diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000..c38d98f
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,24 @@
+# House rule: all text is stored and checked out with LF.
+# (Batch wrappers are the only exception -- cmd.exe needs CRLF.)
+* text=auto eol=lf
+*.bat text eol=crlf
+*.cmd text eol=crlf
+
+# Binaries: never touch line endings, never treat as text
+*.so binary
+*.a binary
+*.dll binary
+*.exe binary
+*.png binary
+*.jpg binary
+*.jpeg binary
+*.ico binary
+*.icns binary
+*.jar binary
+*.zip binary
+*.bin binary
+*.gguf binary
+*.safetensors binary
+*.e0b binary
+# repo-specific addition: the technical report lives at paper/main.pdf
+*.pdf binary
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 9c3a3ad..78cfea2 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -9,6 +9,9 @@ jobs:
   test-macos:
     name: Unit tests (macOS + MLX)
     runs-on: macos-14
+    defaults:
+      run:
+        working-directory: python
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -22,6 +25,9 @@ jobs:
   hygiene:
     name: Repo hygiene (pure Python, no MLX)
     runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: python
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
diff --git a/.gitignore b/.gitignore
index c3b7432..b2587d4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,9 +13,21 @@ dist/
 # Trained adapter artifacts (never committed; see README)
 artifacts/
 
+# Downloaded model checkpoints (never committed; see README quick start).
+# Anchored so the source package python/src/edge0/models/ is not matched.
+/models/
+/python/models/
+
+# Platform build areas materialized by the engine/vendor scripts (never commit):
+# isolated patch-replay worktrees, the pinned llama.cpp checkout, and the
+# local engine-depot bootstrap fallback (see windows/scripts/vendor-build.ps1).
+/wt/
+/vendor/
+.edge0/
+
 # macOS
 .DS_Store
 
 # Bench / runtime logs
 *.log
-scripts/alignment_results.json
+python/scripts/alignment_results.json
diff --git a/NOTICE b/NOTICE
index ddee8a9..2b044cc 100644
--- a/NOTICE
+++ b/NOTICE
@@ -4,7 +4,7 @@ edge0
 This project bundles or adapts the following third-party works:
 
 mlx-lm (MIT License)
-  - `src/edge0/backends/mlx/_impl/qwen3_next.py` is vendored from mlx-lm
+  - `python/src/edge0/backends/mlx/_impl/qwen3_next.py` is vendored from mlx-lm
     0.31.0 (https://github.com/ml-explore/mlx-lm), with additive hooks for
     prerouter routing and layer callbacks.  Non-prerouter paths are kept
     byte-identical to the upstream file.
@@ -12,10 +12,24 @@ mlx-lm (MIT License)
   Copyright (c) 2023-2024 Apple Inc. and the mlx-lm contributors.
 
 bailing_hybrid (Apache License 2.0)
-  - `src/edge0/backends/mlx/_impl/bailing_hybrid.py` is vendored from the
+  - `python/src/edge0/backends/mlx/_impl/bailing_hybrid.py` is vendored from the
     Ling 3.0 MLX implementation (rapid-mlx / Ling-3.0-tiny, Apache-2.0), with
     prerouter hook identifiers renamed for consistency with the framework.
   Copyright (c) the bailing / Ling MLX contributors.
 
+llama.cpp (MIT License)
+  - The Android and Windows engines build their inference core from a
+    pristine upstream checkout of llama.cpp
+    (https://github.com/ggml-org/llama.cpp), pinned at commit 7ab4ee7
+    (release b11100) via `vendor.llama.pin` and materialized into the
+    gitignored `vendor/llama.cpp` by the platform build scripts.
+  - It is never patched in place: all modifications ship as replayable
+    patch bands under `patches/llama.cpp/{common,android,windows}` and are
+    replayed into an isolated build worktree at compile time.
+  - ggml (Apache License 2.0) is developed within the same upstream tree.
+  - Not vendored as files in this repo; built from source by the consumer
+    (see `android/NOTICE` and `windows/NOTICE` for the per-platform detail).
+  Copyright (c) 2023-2024 The ggml authors.
+
 All other files are Copyright (c) 2025 edge0 contributors, licensed under
 the Apache License, Version 2.0 (see LICENSE).
diff --git a/README.md b/README.md
index 5e0240a..c335726 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,8 @@
 
 **An open-source streaming MoE inference framework — SSD expert offload + Recover-LoRA + prerouter routing prediction.**
 
+**Python** · **macOS** · **iOS** · **Android** — one recipe, every device.
+
 [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--35B--A3B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
 [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--8B--A1B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
 [![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--35B--A3B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
@@ -14,16 +16,52 @@
 [![GitHub](https://img.shields.io/badge/GitHub-Edge0--AI%2FEdge0-black?style=for-the-badge&logo=github)](https://github.com/Edge0-AI/Edge0)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue?style=for-the-badge)](LICENSE)
 
-English | [中文](README_zh.md)
+English | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
 
 </div>
 
+## News
+
+- **[2026-09-30]** We released the **edge0 inference engines for four platforms — iOS, macOS, Android and Windows** — so users get the best inference experience across architectures and platforms. The source is open-sourced in this repo ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/) · [`windows/`](windows/)) — see each directory's README for details. The **unified inference framework** follows in **Q4 2026**; see the [Roadmap](#roadmap).
+- **[2026-09-20]** `--prefill-ondemand` lands for `edge0 demo / chat / serve` ([#112](https://github.com/Edge0-AI/Edge0/pull/112)).
+- **[2026-09-16]** Our technical report is on arXiv: [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://arxiv.org/abs/2609.18063).
+- **[2026-09-08]** Initial open-source release of **edge0**, together with both model tiers — [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) and [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) — on Hugging Face and ModelScope.
+
+## About
+
 **edge0** is an open-source streaming MoE inference framework. It
 generalizes the production-proven recipe — **SSD expert offload +
 Recover-LoRA + prerouter routing prediction** — into an extensible
-framework. The backend is isolated by design: the current MLX backend
-runs on Apple Silicon, and additional platforms (CUDA, …) plug into the
-same core abstractions.
+framework that runs large sparse-MoE models on consumer hardware: peak
+memory is bounded by the *active* expert set, not the parameter count.
+
+### Core mechanisms
+
+- **SSD expert offload**: expert weights are streamed from storage on
+  demand; peak memory is bounded by the active set, not the parameter
+  count.
+- **Prerouter**: a trained head predicts expert routing one step
+  ahead, so expert loads overlap the forward pass instead of stalling
+  it — **up to +59%** decode throughput; the gain grows with storage
+  latency, model size, and routed width *K*.
+- **Recover-LoRA**: the int4 base is frozen and LoRA adapters are
+  trained by distillation from the FP teacher, recovering most of the
+  quantization loss at 4-bit (see [Quality](#quality)). Adapters stay
+  unmerged: one read-only base serves multiple adapter sets.
+
+### Platforms
+
+One repo, one recipe, per-platform runtimes:
+
+| Platform | Directory | Stack | Status |
+|---|---|---|---|
+| **Python** (macOS · Apple Silicon) | [`python/`](python/README.md) | Python + MLX | ✅ Available now |
+| **macOS** app & CLI | [`macos/`](macos/README.md) | Rust | ✅ Open-sourced (2026-09-30) |
+| **iOS** app | [`ios/`](ios/README.md) | Swift + MLX Swift | ✅ Open-sourced (2026-09-30) |
+| **Android** app & engine | [`android/`](android/README.md) | Kotlin + native engine | ✅ Open-sourced (2026-09-30) |
+| **Windows** app & engine | [`windows/`](windows/README.md) | C++ + Vulkan | ✅ Open-sourced (2026-09-30) |
+
+### Models
 
 Two model tiers ship with the framework. Each tier is an end-to-end
 release: the released checkpoint, the trained LoRA adapters, and the
@@ -40,32 +78,19 @@ LoRA and prerouter training done for this framework — the adapter files
 are co-located with each checkpoint and load automatically, so
 `edge0 serve <tier>` runs the trained pipeline out of the box.
 
-## Requirements
-
-- **OS / hardware**: the MLX backend runs on macOS with Apple Silicon
-  (M1/M2/M3/M4). The CUDA backend is on the roadmap — no other
-  platforms are supported yet.
-- **Python**: 3.10+ (3.12 recommended).
-- **MLX**: `mlx==0.30.6` / `mlx-metal==0.30.6` with `mlx-lm==0.31.0` (see
-  `pyproject.toml`). Garbled, mixed-language output on Apple A18 / A18 Pro
-  means an older `mlx`: `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
-  ([#8](https://github.com/Edge0-AI/Edge0/issues/8)).
-- **Memory**: ~2.9 GB peak active memory for `edge0-35b`, ~1.0 GB for
-  `edge0-8b` (short contexts; see [Benchmark](#benchmark)). Add
-  headroom for the OS, tokenizer, and long-context KV growth.
-- **Disk**: the 4-bit checkpoints are ~23 GB (`edge0-35b`) and ~4.2 GB
-  (`edge0-8b`); expert weights are mmapped and read on demand, they are
-  not loaded into RAM up front.
-
-## Design
+### Design
 
 - **transformers-style usage**: `AutoModel` / `AutoConfig` / `AutoEngine`
   resolve the tier from the model name;
-- **Backend isolation**: all MLX code lives under `edge0/backends/mlx/`;
-  the core logic (model specs, prerouter, streaming expert pool, server)
-  depends only on the backend facade (`edge0/backends/base.py`), so a new
-  backend implements the same facade (`backends/cuda/` is a reserved
-  slot) with zero changes to core code;
+- **Backend isolation (Python framework)**: within the Python framework,
+  all MLX code lives under `python/src/edge0/backends/mlx/`; the core
+  logic (model specs, prerouter, streaming expert pool, server) depends
+  only on the backend facade (`backends/base.py`), so a new backend
+  implements the same facade (`backends/cuda/` is a reserved slot) with
+  zero changes to core code. The iOS / macOS / Android engines ship
+  platform-native stacks today — bringing every platform under one
+  access layer is exactly what the unified inference framework
+  (see [Roadmap](#roadmap)) will deliver;
 - **Adapters as safetensors**: LoRA and prerouter weights are
   `.safetensors` files with provenance metadata (source, version, owner
   layers), resolved from the model directory or `artifacts/`;
@@ -74,34 +99,79 @@ are co-located with each checkpoint and load automatically, so
   and that model's adapters; upgrading adapters swaps adapter
   files only — the base stays read-only and is never merged.
 
-## Core mechanisms
+### Quality
 
-- **SSD expert offload**: expert weights are streamed from storage on
-  demand; peak memory is bounded by the active set, not the parameter
-  count.
-- **Prerouter**: a trained head predicts expert routing one step
-  ahead, so expert loads overlap the forward pass instead of stalling
-  it — **up to +59%** decode throughput; the gain grows with storage
-  latency, model size, and routed width *K*.
-- **Recover-LoRA**: the int4 base is frozen and LoRA adapters are
-  trained by distillation from the FP teacher, recovering most of the
-  quantization loss at 4-bit (see [Quality](#quality)).  Adapters stay
-  unmerged: one read-only base serves multiple adapter sets.
+All benchmarks were run by us with [OpenCompass](https://github.com/open-compass/opencompass)
+under identical settings and parameters for both the edge0 models (int4 +
+trained adapters + prerouter routing) and the original fp16 base models.
+The loss of the edge0 pipeline is small: **3.9 points on average for
+edge0-35b, 2.8 for edge0-8b** (MMLU-Pro is even above the base). Max 100:
+
+| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **Average** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+### Benchmark
+
+Measured with `python/examples/bench.py` (3.3k-token prompt prefill → 10
+sampled warmup steps → 200 timed sampled decode tokens, 2 runs per tier):
+
+| Tier | Decode speed | Prefill throughput (cold / warm)* | Peak active memory | Test machine |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*Cold = first request after process start (expert weights fault in from
+SSD); warm = subsequent requests (page cache resident). Prefill numbers
+are throughput over a ~3.3k-token prompt (`BENCH_LONG=1`).*
+
+Reproduce:
+
+```bash
+cd python
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+```
 
-## Quick start
+## Getting Started
+
+### Python (macOS · Apple Silicon)
+
+#### Requirements
+
+- **OS / hardware**: the MLX backend runs on macOS with Apple Silicon
+  (M1/M2/M3/M4). The CUDA backend is on the roadmap — no other
+  platforms are supported by the Python framework yet.
+- **Python**: 3.10+ (3.12 recommended).
+- **MLX**: `mlx==0.30.6` / `mlx-metal==0.30.6` with `mlx-lm==0.31.0` (see
+  `python/pyproject.toml`). Garbled, mixed-language output on Apple A18 /
+  A18 Pro means an older `mlx`: `pip install 'mlx==0.30.6'
+  'mlx-metal==0.30.6'` ([#8](https://github.com/Edge0-AI/Edge0/issues/8)).
+- **Memory**: ~2.9 GB peak active memory for `edge0-35b`, ~1.0 GB for
+  `edge0-8b` (short contexts; see [Benchmark](#benchmark)). Add
+  headroom for the OS, tokenizer, and long-context KV growth.
+- **Disk**: the 4-bit checkpoints are ~23 GB (`edge0-35b`) and ~4.2 GB
+  (`edge0-8b`); expert weights are mmapped and read on demand, they are
+  not loaded into RAM up front.
 
-### 1) Install
+#### 1) Install
 
 ```bash
+cd python
 # Python >= 3.10; the MLX backend requires macOS with Apple Silicon
 python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
 ```
 
-### 2) Download a model
+#### 2) Download a model
 
-The two tiers are published on Hugging Face and ModelScope — each repo bundles the
-base checkpoint and the trained LoRA + prerouter adapters in **one
-directory**, so a single download is a ready-to-run model:
+The two tiers are published on Hugging Face and ModelScope — each repo
+bundles the base checkpoint and the trained LoRA + prerouter adapters in
+**one directory**, so a single download is a ready-to-run model:
 
 - [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) (~23 GB) · [ModelScope mirror](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
 - [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) (~4.2 GB) · [ModelScope mirror](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
@@ -125,7 +195,7 @@ models/edge0-35b/
 └── prerouter_edge0_35b.safetensors     # trained prerouter heads
 ```
 
-### 3) Point edge0 at it
+#### 3) Point edge0 at it
 
 Tier names resolve to local directories via environment variables
 (where you put the download is up to you):
@@ -143,7 +213,7 @@ edge0 demo models/edge0-35b
 edge0 serve models/edge0-8b
 ```
 
-### 4) Run
+#### 4) Run
 
 ```bash
 # quick demo
@@ -156,7 +226,7 @@ edge0 serve edge0-35b
 ```bash
 curl http://127.0.0.1:8000/v1/chat/completions \
   -H 'Content-Type: application/json' \
-  -d '{"messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+  -d '{"model":"edge0-35b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
 
 # 5) One-shot chat (pass --max-new to cap length; add --show-thinking to
 #    print the model's reasoning block too)
@@ -165,7 +235,7 @@ edge0 chat edge0-35b --prompt "Explain streaming inference in one sentence."
 
 `python -m edge0 ...` is equivalent to `edge0 ...`.
 
-### Python API
+#### Python API
 
 ```python
 from edge0 import AutoEngine
@@ -182,10 +252,10 @@ print(engine._tok.decode(tokens))
 engine.close()   # release mmaps / expert cache
 ```
 
-`examples/demo.py` is the same minimal walkthrough (`edge0 demo` runs
-this exact path).
+`python/examples/demo.py` is the same minimal walkthrough (`edge0 demo`
+runs this exact path).
 
-### Models and adapters
+#### Models and adapters
 
 - **Checkpoint**: the original model directory (`config.json`,
   `model*.safetensors`, tokenizer). `edge0 serve <dir>` /
@@ -195,72 +265,133 @@ this exact path).
   location automatically:
   - the model directory (recommended): side by side with the base, e.g.
     `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`;
-  - `artifacts/` (repo root, gitignored): convert once from
-    training-side npz exports via `edge0 convert-adapters --npz-dir ...`.
+  - `artifacts/` at the Python project root (gitignored): an optional
+    fallback cache for adapter safetensors not co-located with the model.
 - The published model repos bundle both the base checkpoint and the
   current default adapter release, so `scripts/fetch_models.py` produces
-  a ready-to-run model directory.  Check each model's doc page for its
+  a ready-to-run model directory. Check each model's doc page for its
   adapter provenance (training data, owner-layer layout).
 - Both adapters are required for the prerouter + LoRA pipeline; if a
   file is missing, `edge0` fails with a clear message (or pass
   `--no-prerouter` / `--no-lora` to run the plain base model).
-## Quality
 
-All benchmarks were run by us with [OpenCompass](https://github.com/open-compass/opencompass)
-under identical settings and parameters for both the edge0 models (int4 +
-trained adapters + prerouter routing) and the original fp16 base models.
-The loss of the edge0 pipeline is small: **3.9 points on average for
-edge0-35b, 2.8 for edge0-8b** (MMLU-Pro is even above the base). Max 100:
+#### Documentation
 
-| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
-|---|---:|---:|---:|---:|
-| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
-| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
-| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
-| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
-| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
-| **Average** | **79.2** | **83.2** | **69.9** | **72.7** |
+- [Architecture](docs/architecture.md)
+- [Attention](docs/attention.md) / [MoE](docs/moe.md) / [SSD streaming](docs/streaming.md) / [prerouter](docs/prerouter.md)
+- [Adding a model](docs/adding-a-model.md)
+- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md)
+- Technical report: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](paper/main.pdf))
 
-## Benchmark
+### macOS / iOS / Android / Windows
 
-Measured with `examples/bench.py` (3.3k-token prompt prefill → 10 sampled
-warmup steps → 200 timed sampled decode tokens, 2 runs per tier):
+The four platform engines are open-sourced in this repo — more details
+in each directory's README:
 
-| Tier | Decode speed | Prefill throughput (cold / warm)* | Peak active memory | Test machine |
-|---|---|---|---|---|
-| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
-| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+- **macOS**: local CLI / daemon / desktop app (Rust) — see [`macos/README.md`](macos/README.md)
+- **iOS**: on-device iPhone app (Swift + MLX Swift) — see [`ios/README.md`](ios/README.md)
+- **Android**: on-device app + native engine (Kotlin) — see [`android/README.md`](android/README.md)
+- **Windows**: desktop app + native engine (C++ + Vulkan) — see [`windows/README.md`](windows/README.md)
 
-*Cold = first request after process start (expert weights fault in from
-SSD); warm = subsequent requests (page cache resident). Prefill numbers
-are throughput over a ~3.3k-token prompt (`BENCH_LONG=1`).*
+The **unified inference framework** — one access layer, runtime
+auto-adapting to iOS / macOS / Android / Windows / Python — arrives in
+**Q4 2026**; see the [Roadmap](#roadmap).
 
-Reproduce:
+## Roadmap
 
-```bash
-python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
-python examples/bench.py edge0-8b    # via $EDGE0_8B_MODEL
-```
+### Q4 2026
+
+**Platforms & systems**
+
+- **edge0 unified inference framework** — we will
+  open-source a unified inference framework: **one unified access
+  layer** (a single API across chat / serve / on-device use), with the
+  **runtime automatically adapting to the hardware platform** — iOS,
+  macOS, Android, Windows and Python. It builds on the platform
+  engines already open-sourced in this repo (`ios/` · `macos/` ·
+  `android/` · `windows/`).
+- **CUDA backend** for the Python framework — reserved slot at
+  `python/src/edge0/backends/cuda/`, core code needs zero changes.
+
+**Models & algorithms**
+
+Q4 works two fronts: bringing a next-generation architecture into the
+framework, and turning latent reasoning into a real latency saving rather
+than just an arithmetic one.
+
+- **Next-gen architecture support (Qwen3.8-Flash class)** — run hybrid
+  linear attention (GDN + QSA), gated multi-branch residual, and N-gram
+  embedding topologies on edge0. These designs suit SSD streaming offload
+  naturally: O(1)-state attention keeps long thinking from becoming a
+  KV-cache problem, and lookup-only N-gram tables stream on demand. Goal:
+  the tier runs on a single device and benchmarks within an acceptable
+  gap of the fp16 base.
+- **Latent thinking + batched expert pre-prediction** — make latent
+  reasoning a *latency* saving, not only an arithmetic one. The core
+  engineering problem: move expert routing from per-position to **once
+  per block**, so one prediction covers every position and round of a
+  block and **expert load volume decouples from the reasoning loop
+  count** — plus cross-block prefetch that loads the next block's experts
+  inside the current block's compute window. Progress is measured as
+  **end-to-end thinking-phase time at matched accuracy** (never tokens/s).
+- More model tiers and adapter releases on the existing pipeline.
 
-## Tests
+## Contributing
+
+Contributions are welcome — issues, PRs, benchmark reports and model
+ports all count.
+
+**Python framework** (available now):
 
 ```bash
+cd python
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+
 pytest                 # unit tests (no real weights)
 EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q
-                        # real-weight generation; missing tiers are skipped
+                       # real-weight generation; missing tiers are skipped
 .venv/bin/python scripts/e2e_smoke.py \
   --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
-                        # staged vs exact consistency + generation smoke
+                       # staged vs exact consistency + generation smoke
 scripts/generate_example.py   # full-pipeline API example
-examples/demo.py       # minimal API walkthrough
+examples/demo.py              # minimal API walkthrough
 ```
 
-## Documentation
+CI runs unit tests (macOS + MLX) and a repo-hygiene suite (no hardcoded
+paths, backend-boundary and secret checks) on every PR.
 
-- [Architecture](docs/architecture.md)
-- [Attention](docs/attention.md) / [MoE](docs/moe.md) / [SSD streaming](docs/streaming.md) / [prerouter](docs/prerouter.md)
-- [Adding a model](docs/adding-a-model.md)
-- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md)
+**Platform runtimes** (macOS / iOS / Android / Windows): each platform
+directory ships its own build guide and tests — see the directory
+READMEs.
+
+Workflow: fork → feature branch → PR against `main`. Please keep the
+hygiene suite green and add tests for new behavior.
+
+## Citation
+
+If you find edge0 useful, please cite our technical report:
+
+```bibtex
+@article{lin2026other,
+  title   = {The Other Half of the Memory Wall: Serving 35B MoEs from SSD
+             with Trained Routing Prediction},
+  author  = {Lin, Yu and Wang, Yiming and Cai, Runyuan and Liu, Hanze and
+             Zeng, Xiaodong},
+  journal = {arXiv preprint arXiv:2609.18063},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2609.18063}
+}
+```
+
+## Contact Us
+
+Community and support channels are coming soon — this section will list
+the official ways to reach us:
+
+- **Email**: samuel@edge0.ai
+
+For bugs and feature requests, please use
+[GitHub Issues](https://github.com/Edge0-AI/Edge0/issues).
 
 ## License
 
diff --git a/README_es.md b/README_es.md
new file mode 100644
index 0000000..c5f8dcf
--- /dev/null
+++ b/README_es.md
@@ -0,0 +1,419 @@
+<div align="center">
+
+<img src="assets/20260908-223115.jpg" alt="edge0" width="100%">
+
+# edge0
+
+**Un framework de inferencia MoE en streaming de código abierto — SSD expert offload + Recover-LoRA + predicción de enrutamiento del prerouter.**
+
+**Python** · **macOS** · **iOS** · **Android** — una receta, todos los dispositivos.
+
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--35B--A3B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--8B--A1B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+[![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--35B--A3B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
+[![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--8B--A1B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
+[![arXiv](https://img.shields.io/badge/arXiv-2609.18063-B31B1B?style=for-the-badge&logo=arxiv&logoColor=white)](https://arxiv.org/abs/2609.18063)
+[![GitHub](https://img.shields.io/badge/GitHub-Edge0--AI%2FEdge0-black?style=for-the-badge&logo=github)](https://github.com/Edge0-AI/Edge0)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue?style=for-the-badge)](LICENSE)
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | Español | [Français](README_fr.md)
+
+</div>
+
+## Novedades
+
+- **[2026-09-30]** Lanzamos los **motores de inferencia de edge0 para cuatro plataformas — iOS, macOS, Android y Windows** — de modo que los usuarios obtienen la mejor experiencia de inferencia en todas las arquitecturas y plataformas. El código fuente se ha liberado en este repositorio ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/) · [`windows/`](windows/)) — consulta el README de cada directorio para más detalles. El **framework de inferencia unificado** llegará en **Q4 2026**; consulta la [Hoja de ruta](#hoja-de-ruta).
+- **[2026-09-20]** `--prefill-ondemand` llega a `edge0 demo / chat / serve` ([#112](https://github.com/Edge0-AI/Edge0/pull/112)).
+- **[2026-09-16]** Nuestro informe técnico está en arXiv: [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://arxiv.org/abs/2609.18063).
+- **[2026-09-08]** Primera versión de código abierto de **edge0**, junto con los dos niveles de modelos — [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) y [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) — en Hugging Face y ModelScope.
+
+## Acerca de
+
+**edge0** es un framework de inferencia MoE en streaming de código abierto.
+Generaliza la receta probada en producción — **SSD expert offload +
+Recover-LoRA + predicción de enrutamiento del prerouter** — en un framework
+extensible que ejecuta grandes modelos sparse-MoE en hardware de consumo: la
+memoria pico está limitada por el conjunto de expertos *activos*, no por el
+número de parámetros.
+
+### Mecanismos principales
+
+- **SSD expert offload**: los pesos de los expertos se transmiten desde el
+  almacenamiento bajo demanda; la memoria pico está limitada por el conjunto
+  activo, no por el número de parámetros.
+- **Prerouter**: una cabeza entrenada predice el enrutamiento de expertos un
+  paso por delante, de modo que las cargas de expertos se solapan con el
+  forward pass en lugar de bloquearlo — **hasta +59%** de throughput de
+  decode; la ganancia crece con la latencia del almacenamiento, el tamaño del
+  modelo y el ancho de enrutamiento *K*.
+- **Recover-LoRA**: la base int4 se mantiene congelada y los adaptadores LoRA
+  se entrenan por destilación desde el teacher FP, recuperando la mayor parte
+  de la pérdida de cuantización a 4 bits (consulta [Calidad](#calidad)). Los
+  adaptadores permanecen sin fusionar: una única base de solo lectura sirve a
+  varios conjuntos de adaptadores.
+
+### Plataformas
+
+Un repositorio, una receta, runtimes por plataforma:
+
+| Plataforma | Directorio | Stack | Estado |
+|---|---|---|---|
+| **Python** (macOS · Apple Silicon) | [`python/`](python/README_es.md) | Python + MLX | ✅ Disponible ahora |
+| Aplicación y CLI para **macOS** | [`macos/`](macos/README_es.md) | Rust | ✅ Código abierto (2026-09-30) |
+| Aplicación para **iOS** | [`ios/`](ios/README_es.md) | Swift + MLX Swift | ✅ Código abierto (2026-09-30) |
+| Aplicación y motor para **Android** | [`android/`](android/README_es.md) | Kotlin + motor nativo | ✅ Código abierto (2026-09-30) |
+| Aplicación y motor para **Windows** | [`windows/`](windows/README_es.md) | C++ + Vulkan | ✅ Código abierto (2026-09-30) |
+
+### Modelos
+
+El framework incluye dos niveles de modelos. Cada nivel es una versión de
+extremo a extremo: el checkpoint publicado, los adaptadores LoRA entrenados y
+las cabezas prerouter entrenadas funcionan en conjunto como una sola unidad.
+
+| Nivel | Checkpoint publicado | Perfil de inferencia |
+|---|---|---|
+| `edge0-35b` | [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) · [ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview) | 4 bits, 40 capas, 256 expertos, prerouter K=4 |
+| `edge0-8b` | [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview) | 4 bits, 24 capas, 128 expertos, prerouter K=8 |
+
+Ambos checkpoints se construyen sobre modelos base sparse-MoE abiertos
+(Qwen3.6-35B-A3B y el híbrido bailing Ling 3.0, respectivamente) e incluyen
+el entrenamiento de LoRA y prerouter realizado para este framework — los
+archivos de adaptadores están junto a cada checkpoint y se cargan
+automáticamente, de modo que `edge0 serve <nivel>` ejecuta el pipeline
+entrenado sin ninguna configuración adicional.
+
+### Diseño
+
+- **Uso al estilo transformers**: `AutoModel` / `AutoConfig` / `AutoEngine`
+  resuelven el nivel a partir del nombre del modelo;
+- **Aislamiento de backends (framework de Python)**: dentro del framework de
+  Python, todo el código MLX vive en `python/src/edge0/backends/mlx/`; la
+  lógica principal (especificaciones de modelos, prerouter, pool de expertos
+  en streaming, servidor) depende solo de la fachada del backend
+  (`backends/base.py`), de modo que un nuevo backend implementa la misma
+  fachada (`backends/cuda/` es un espacio reservado) sin ningún cambio en el
+  código principal. Los motores de iOS / macOS / Android incluyen hoy stacks
+  nativos de cada plataforma — llevar todas las plataformas bajo una única
+  capa de acceso es exactamente lo que entregará el framework de inferencia
+  unificado (consulta la [Hoja de ruta](#hoja-de-ruta));
+- **Adaptadores como safetensors**: los pesos de LoRA y del prerouter son
+  archivos `.safetensors` con metadatos de procedencia (fuente, versión, capas
+  propietarias), resueltos desde el directorio del modelo o `artifacts/`;
+- **Modelo + adaptadores en un solo directorio**: un directorio de modelo
+  contiene tanto el checkpoint base (`config.json` / `model*.safetensors` /
+  tokenizer) como los adaptadores de ese modelo; actualizar los adaptadores
+  solo reemplaza los archivos de adaptadores — la base permanece de solo
+  lectura y nunca se fusiona.
+
+### Calidad
+
+Todos los benchmarks fueron ejecutados por nosotros con
+[OpenCompass](https://github.com/open-compass/opencompass) bajo configuraciones
+y parámetros idénticos tanto para los modelos edge0 (int4 + adaptadores
+entrenados + enrutamiento prerouter) como para los modelos base fp16
+originales. La pérdida del pipeline edge0 es pequeña: **3.9 puntos de media
+para edge0-35b, 2.8 para edge0-8b** (MMLU-Pro incluso supera a la base).
+Máximo 100:
+
+| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **Media** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+### Benchmark
+
+Medido con `python/examples/bench.py` (prefill de un prompt de 3.3k tokens →
+10 pasos de calentamiento muestreados → 200 tokens de decode muestreados y
+cronometrados, 2 ejecuciones por nivel):
+
+| Nivel | Velocidad de decode | Throughput de prefill (frío / caliente)* | Memoria activa pico | Máquina de pruebas |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*Frío = primera petición tras iniciar el proceso (los pesos de expertos entran
+por fallos de página desde el SSD); caliente = peticiones posteriores
+(residentes en la caché de páginas). Las cifras de prefill son throughput sobre
+un prompt de ~3.3k tokens (`BENCH_LONG=1`).*
+
+Reproducir:
+
+```bash
+cd python
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+```
+
+## Primeros pasos
+
+### Python (macOS · Apple Silicon)
+
+#### Requisitos
+
+- **SO / hardware**: el backend MLX se ejecuta en macOS con Apple Silicon
+  (M1/M2/M3/M4). El backend CUDA está en la hoja de ruta — el framework de
+  Python aún no admite otras plataformas.
+- **Python**: 3.10+ (se recomienda 3.12).
+- **MLX**: `mlx==0.30.6` / `mlx-metal==0.30.6` con `mlx-lm==0.31.0` (consulta
+  `python/pyproject.toml`). Una salida ilegible o mezclada de idiomas en Apple
+  A18 / A18 Pro indica un `mlx` antiguo: `pip install 'mlx==0.30.6'
+  'mlx-metal==0.30.6'` ([#8](https://github.com/Edge0-AI/Edge0/issues/8)).
+- **Memoria**: ~2.9 GB de memoria activa pico para `edge0-35b`, ~1.0 GB para
+  `edge0-8b` (contextos cortos; consulta [Benchmark](#benchmark)). Añade
+  margen para el SO, el tokenizer y el crecimiento de la KV cache en contextos
+  largos.
+- **Disco**: los checkpoints de 4 bits ocupan ~23 GB (`edge0-35b`) y ~4.2 GB
+  (`edge0-8b`); los pesos de los expertos se mapean con mmap y se leen bajo
+  demanda, no se cargan en RAM por adelantado.
+
+#### 1) Instalación
+
+```bash
+cd python
+# Python >= 3.10; the MLX backend requires macOS with Apple Silicon
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+#### 2) Descargar un modelo
+
+Los dos niveles están publicados en Hugging Face y ModelScope — cada
+repositorio incluye el checkpoint base y los adaptadores entrenados de LoRA +
+prerouter en **un solo directorio**, de modo que una única descarga es un
+modelo listo para ejecutar:
+
+- [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) (~23 GB) · [mirror en ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
+- [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) (~4.2 GB) · [mirror en ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
+
+```bash
+# with the repo's helper (defaults to the two repos above):
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview     --local-dir models/edge0-35b
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview     --local-dir models/edge0-8b
+```
+
+De cualquier forma obtendrás un directorio así:
+
+```
+models/edge0-35b/
+├── config.json, model-*.safetensors, tokenizer files   # base checkpoint
+├── lora_edge0_35b.safetensors          # trained LoRA adapters
+└── prerouter_edge0_35b.safetensors     # trained prerouter heads
+```
+
+#### 3) Apuntar edge0 al modelo
+
+Los nombres de nivel se resuelven a directorios locales mediante variables de
+entorno (dónde guardes la descarga depende de ti):
+
+```bash
+export EDGE0_35B_MODEL=$PWD/models/edge0-35b
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+O sáltate las variables de entorno y pasa el directorio directamente — el
+nivel se detecta automáticamente desde el `config.json` del checkpoint:
+
+```bash
+edge0 demo models/edge0-35b
+edge0 serve models/edge0-8b
+```
+
+#### 4) Ejecutar
+
+```bash
+# quick demo
+edge0 demo edge0-35b
+
+# serve (OpenAI-compatible /v1/chat/completions)
+edge0 serve edge0-35b
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-35b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+
+# 5) One-shot chat (pass --max-new to cap length; add --show-thinking to
+#    print the model's reasoning block too)
+edge0 chat edge0-35b --prompt "Explain streaming inference in one sentence."
+```
+
+`python -m edge0 ...` es equivalente a `edge0 ...`.
+
+#### API de Python
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`python/examples/demo.py` es el mismo recorrido mínimo (`edge0 demo` ejecuta
+exactamente esta ruta).
+
+#### Modelos y adaptadores
+
+- **Checkpoint**: el directorio original del modelo (`config.json`,
+  `model*.safetensors`, tokenizer). `edge0 serve <dir>` /
+  `AutoEngine.from_pretrained(<dir>)` detectan el nivel desde `config.json`.
+- Los **adaptadores** (LoRA + prerouter, safetensors) se resuelven
+  automáticamente desde cualquiera de estas ubicaciones:
+  - el directorio del modelo (recomendado): junto a la base, p. ej.
+    `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`;
+  - `artifacts/` en la raíz del proyecto Python (ignorado por git): una caché
+    alternativa opcional para los safetensors de adaptadores que no estén
+    junto al modelo.
+- Los repositorios de modelos publicados incluyen tanto el checkpoint base
+  como la versión actual por defecto de los adaptadores, de modo que
+  `scripts/fetch_models.py` produce un directorio de modelo listo para
+  ejecutar. Consulta la página de documentación de cada modelo para conocer la
+  procedencia de sus adaptadores (datos de entrenamiento, disposición de las
+  capas propietarias).
+- Ambos adaptadores son necesarios para el pipeline prerouter + LoRA; si falta
+  un archivo, `edge0` falla con un mensaje claro (o pasa `--no-prerouter` /
+  `--no-lora` para ejecutar el modelo base simple).
+
+#### Documentación
+
+- [Arquitectura](docs/architecture.md)
+- [Atención](docs/attention.md) / [MoE](docs/moe.md) / [Streaming SSD](docs/streaming.md) / [prerouter](docs/prerouter.md)
+- [Añadir un modelo](docs/adding-a-model.md)
+- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md)
+- Informe técnico: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](paper/main.pdf))
+
+### macOS / iOS / Android / Windows
+
+Los cuatro motores de plataforma son de código abierto en este repositorio —
+más detalles en el README de cada directorio:
+
+- **macOS**: CLI / demonio / aplicación de escritorio local (Rust) — consulta [`macos/README_es.md`](macos/README_es.md)
+- **iOS**: aplicación para iPhone en el dispositivo (Swift + MLX Swift) — consulta [`ios/README_es.md`](ios/README_es.md)
+- **Android**: aplicación en el dispositivo + motor nativo (Kotlin) — consulta [`android/README_es.md`](android/README_es.md)
+- **Windows**: aplicación de escritorio + motor nativo (C++ + Vulkan) — consulta [`windows/README_es.md`](windows/README_es.md)
+
+El **framework de inferencia unificado** — una única capa de acceso, con el
+runtime adaptándose automáticamente a iOS / macOS / Android / Windows /
+Python — llegará en **Q4 2026**; consulta la [Hoja de ruta](#hoja-de-ruta).
+
+## Hoja de ruta
+
+### Q4 2026
+
+**Plataformas y sistemas**
+
+- **Framework de inferencia unificado de edge0** — liberaremos
+  un framework de inferencia unificado: **una única capa de acceso
+  unificada** (una sola API para chat / serve / uso en el dispositivo), con el
+  **runtime adaptándose automáticamente a la plataforma de hardware** — iOS,
+  macOS, Android, Windows y Python. Se construye sobre los motores de
+  plataforma ya liberados en este repositorio (`ios/` · `macos/` ·
+  `android/` · `windows/`).
+- **Backend CUDA** para el framework de Python — espacio reservado en
+  `python/src/edge0/backends/cuda/`, el código principal no necesita ningún
+  cambio.
+
+**Modelos y algoritmos**
+
+Q4 trabaja en dos frentes: incorporar una arquitectura de próxima generación
+al framework y convertir el razonamiento latente en un ahorro real de latencia
+y no solo aritmético.
+
+- **Soporte de arquitectura de próxima generación (clase Qwen3.8-Flash)** —
+  ejecutar atención lineal híbrida (GDN + QSA), residual multirrama con
+  puertas y topologías de embedding N-gram en edge0. Estos diseños encajan de
+  forma natural con el streaming offload a SSD: la atención de estado O(1)
+  evita que el thinking prolongado se convierta en un problema de KV cache, y
+  las tablas N-gram de solo consulta se transmiten bajo demanda. Objetivo: que
+  el nivel se ejecute en un solo dispositivo y que los benchmarks queden a una
+  distancia aceptable de la base fp16.
+- **Thinking latente + predicción anticipada de expertos por lotes** —
+  convertir el razonamiento latente en un ahorro de *latencia*, no solo
+  aritmético. El problema central de ingeniería: mover el enrutamiento de
+  expertos de por posición a **una vez por bloque**, de modo que una sola
+  predicción cubra cada posición y ronda de un bloque y **el volumen de carga
+  de expertos se desacople del número de bucles de razonamiento** — además de
+  prefetch entre bloques que carga los expertos del siguiente bloque dentro de
+  la ventana de cómputo del bloque actual. El progreso se mide como **tiempo
+  de la fase de thinking de extremo a extremo con la misma precisión** (nunca
+  tokens/s).
+- Más niveles de modelos y versiones de adaptadores sobre el pipeline
+  existente.
+
+## Contribuir
+
+Las contribuciones son bienvenidas — issues, PRs, informes de benchmarks y
+ports de modelos, todo cuenta.
+
+**Framework de Python** (disponible ahora):
+
+```bash
+cd python
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q
+                       # real-weight generation; missing tiers are skipped
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+                       # staged vs exact consistency + generation smoke
+scripts/generate_example.py   # full-pipeline API example
+examples/demo.py              # minimal API walkthrough
+```
+
+CI ejecuta pruebas unitarias (macOS + MLX) y una suite de higiene del
+repositorio (sin rutas hardcodeadas, checks de límites de backend y de
+secretos) en cada PR.
+
+**Runtimes de plataforma** (macOS / iOS / Android / Windows): cada
+directorio de plataforma incluye su propia guía de compilación y pruebas
+— consulta el README de cada directorio.
+
+Flujo de trabajo: fork → rama de feature → PR contra `main`. Mantén la suite
+de higiene en verde y añade pruebas para los nuevos comportamientos.
+
+## Cita
+
+Si edge0 te resulta útil, cita nuestro informe técnico:
+
+```bibtex
+@article{lin2026other,
+  title   = {The Other Half of the Memory Wall: Serving 35B MoEs from SSD
+             with Trained Routing Prediction},
+  author  = {Lin, Yu and Wang, Yiming and Cai, Runyuan and Liu, Hanze and
+             Zeng, Xiaodong},
+  journal = {arXiv preprint arXiv:2609.18063},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2609.18063}
+}
+```
+
+## Contacto
+
+Los canales de comunidad y soporte llegarán pronto — esta sección enumerará
+las vías oficiales para contactarnos:
+
+- **Email**: samuel@edge0.ai
+
+Para bugs y solicitudes de funcionalidades, usa
+[GitHub Issues](https://github.com/Edge0-AI/Edge0/issues).
+
+## Licencia
+
+Apache-2.0, incluido el código de terceros incorporado (consulta
+[NOTICE](NOTICE)).
diff --git a/README_fr.md b/README_fr.md
new file mode 100644
index 0000000..ade6b8d
--- /dev/null
+++ b/README_fr.md
@@ -0,0 +1,430 @@
+<div align="center">
+
+<img src="assets/20260908-223115.jpg" alt="edge0" width="100%">
+
+# edge0
+
+**Un framework open source d'inférence MoE en streaming — SSD expert offload + Recover-LoRA + prédiction de routage prerouter.**
+
+**Python** · **macOS** · **iOS** · **Android** — une seule recette, tous les appareils.
+
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--35B--A3B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--8B--A1B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+[![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--35B--A3B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
+[![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--8B--A1B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
+[![arXiv](https://img.shields.io/badge/arXiv-2609.18063-B31B1B?style=for-the-badge&logo=arxiv&logoColor=white)](https://arxiv.org/abs/2609.18063)
+[![GitHub](https://img.shields.io/badge/GitHub-Edge0--AI%2FEdge0-black?style=for-the-badge&logo=github)](https://github.com/Edge0-AI/Edge0)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue?style=for-the-badge)](LICENSE)
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | Français
+
+</div>
+
+## Actualités
+
+- **[2026-09-30]** Nous avons publié les **moteurs d'inférence edge0 pour quatre plateformes — iOS, macOS, Android et Windows** — afin d'offrir aux utilisateurs la meilleure expérience d'inférence à travers les architectures et les plateformes. Le code source est publié en open source dans ce dépôt ([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/) · [`windows/`](windows/)) — consultez le README de chaque répertoire pour plus de détails. Le **framework d'inférence unifié** suivra au **T4 2026** ; voir la [Feuille de route](#feuille-de-route).
+- **[2026-09-20]** `--prefill-ondemand` arrive pour `edge0 demo / chat / serve` ([#112](https://github.com/Edge0-AI/Edge0/pull/112)).
+- **[2026-09-16]** Notre rapport technique est sur arXiv : [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://arxiv.org/abs/2609.18063).
+- **[2026-09-08]** Première publication open source d'**edge0**, avec les deux gammes de modèles — [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) et [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) — sur Hugging Face et ModelScope.
+
+## À propos
+
+**edge0** est un framework open source d'inférence MoE en streaming. Il
+généralise la recette éprouvée en production — **SSD expert offload +
+Recover-LoRA + prédiction de routage prerouter** — en un framework
+extensible qui exécute de grands modèles MoE creux (sparse-MoE) sur du
+matériel grand public : la mémoire de pointe est limitée par l'ensemble
+*actif* d'experts, et non par le nombre de paramètres.
+
+### Mécanismes fondamentaux
+
+- **SSD expert offload** : les poids des experts sont lus en streaming
+  depuis le stockage à la demande ; la mémoire de pointe est limitée par
+  l'ensemble actif, et non par le nombre de paramètres.
+- **Prerouter** : une tête entraînée prédit le routage des experts un pas
+  à l'avance, de sorte que les chargements d'experts chevauchent la passe
+  avant au lieu de la bloquer — **jusqu'à +59 %** de débit de decode ; le
+  gain augmente avec la latence du stockage, la taille du modèle et la
+  largeur routée *K*.
+- **Recover-LoRA** : la base int4 est gelée et les adaptateurs LoRA sont
+  entraînés par distillation depuis l'enseignant FP, récupérant
+  l'essentiel de la perte de quantification à 4 bits (voir
+  [Qualité](#qualité)). Les adaptateurs restent non fusionnés : une base
+  unique en lecture seule sert plusieurs jeux d'adaptateurs.
+
+### Plateformes
+
+Un dépôt, une recette, des runtimes par plateforme :
+
+| Plateforme | Répertoire | Stack | Statut |
+|---|---|---|---|
+| **Python** (macOS · Apple Silicon) | [`python/`](python/README_fr.md) | Python + MLX | ✅ Disponible maintenant |
+| Application et CLI **macOS** | [`macos/`](macos/README_fr.md) | Rust | ✅ Open source (2026-09-30) |
+| Application **iOS** | [`ios/`](ios/README_fr.md) | Swift + MLX Swift | ✅ Open source (2026-09-30) |
+| Application et moteur **Android** | [`android/`](android/README_fr.md) | Kotlin + moteur natif | ✅ Open source (2026-09-30) |
+| Application et moteur **Windows** | [`windows/`](windows/README_fr.md) | C++ + Vulkan | ✅ Open source (2026-09-30) |
+
+### Modèles
+
+Deux gammes de modèles sont fournies avec le framework. Chaque gamme est
+une publication complète : le checkpoint publié, les adaptateurs LoRA
+entraînés et les têtes prerouter entraînées fonctionnent ensemble comme
+une seule unité.
+
+| Gamme | Checkpoint publié | Profil d'inférence |
+|---|---|---|
+| `edge0-35b` | [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) · [ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview) | 4-bit, 40 couches, 256 experts, prerouter K=4 |
+| `edge0-8b` | [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview) | 4-bit, 24 couches, 128 experts, prerouter K=8 |
+
+Les deux checkpoints sont construits à partir de modèles de base MoE creux
+open source (respectivement Qwen3.6-35B-A3B et l'hybride bailing Ling 3.0)
+et sont fournis avec l'entraînement LoRA et prerouter réalisé pour ce
+framework — les fichiers d'adaptateurs sont colocalisés avec chaque
+checkpoint et se chargent automatiquement, de sorte que
+`edge0 serve <tier>` exécute le pipeline entraîné prêt à l'emploi.
+
+### Conception
+
+- **Utilisation façon transformers** : `AutoModel` / `AutoConfig` / `AutoEngine`
+  résolvent la gamme à partir du nom du modèle ;
+- **Isolation du backend (framework Python)** : dans le framework Python,
+  tout le code MLX se trouve sous `python/src/edge0/backends/mlx/` ; la
+  logique centrale (spécifications des modèles, prerouter, pool d'experts
+  en streaming, serveur) ne dépend que de la façade du backend
+  (`backends/base.py`), de sorte qu'un nouveau backend implémente la même
+  façade (`backends/cuda/` est un emplacement réservé) sans aucune
+  modification du code central. Les moteurs iOS / macOS / Android sont
+  aujourd'hui fournis avec des stacks natives par plateforme — regrouper
+  toutes les plateformes sous une couche d'accès unique est précisément ce
+  que fournira le framework d'inférence unifié (voir la
+  [Feuille de route](#feuille-de-route)) ;
+- **Adaptateurs au format safetensors** : les poids LoRA et prerouter sont
+  des fichiers `.safetensors` avec métadonnées de provenance (source,
+  version, couches propriétaires), résolus depuis le répertoire du modèle
+  ou `artifacts/` ;
+- **Modèle + adaptateurs dans un même répertoire** : un répertoire de
+  modèle contient à la fois le checkpoint de base (`config.json` /
+  `model*.safetensors` / tokenizer) et les adaptateurs de ce modèle ; la
+  mise à niveau des adaptateurs ne remplace que les fichiers d'adaptateurs
+  — la base reste en lecture seule et n'est jamais fusionnée.
+
+### Qualité
+
+Tous les benchmarks ont été exécutés par nos soins avec
+[OpenCompass](https://github.com/open-compass/opencompass), dans des
+conditions et avec des paramètres identiques, pour les modèles edge0
+(int4 + adaptateurs entraînés + routage prerouter) et pour les modèles de
+base fp16 d'origine. La perte du pipeline edge0 est faible : **3,9 points
+en moyenne pour edge0-35b, 2,8 pour edge0-8b** (MMLU-Pro est même
+au-dessus de la base). Maximum 100 :
+
+| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **Moyenne** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+### Benchmark
+
+Mesuré avec `python/examples/bench.py` (prefill d'un prompt de 3,3 k
+tokens → 10 pas d'échauffement échantillonnés → 200 tokens de decode
+échantillonnés et chronométrés, 2 exécutions par gamme) :
+
+| Gamme | Vitesse de decode | Débit de prefill (à froid / à chaud)* | Mémoire active de pointe | Machine de test |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*À froid = première requête après le démarrage du processus (les poids des
+experts arrivent depuis le SSD par défaut de page) ; à chaud = requêtes
+suivantes (page cache résident). Les chiffres de prefill sont des débits
+mesurés sur un prompt d'environ 3,3 k tokens (`BENCH_LONG=1`).*
+
+Reproduire :
+
+```bash
+cd python
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+```
+
+## Démarrage rapide
+
+### Python (macOS · Apple Silicon)
+
+#### Prérequis
+
+- **OS / matériel** : le backend MLX fonctionne sur macOS avec Apple
+  Silicon (M1/M2/M3/M4). Le backend CUDA est sur la feuille de route —
+  aucune autre plateforme n'est encore prise en charge par le framework
+  Python.
+- **Python** : 3.10+ (3.12 recommandé).
+- **MLX** : `mlx==0.30.6` / `mlx-metal==0.30.6` avec `mlx-lm==0.31.0` (voir
+  `python/pyproject.toml`). Une sortie illisible ou mêlant plusieurs
+  langues sur Apple A18 / A18 Pro indique une version plus ancienne de
+  `mlx` : `pip install 'mlx==0.30.6'
+  'mlx-metal==0.30.6'` ([#8](https://github.com/Edge0-AI/Edge0/issues/8)).
+- **Mémoire** : environ 2,9 GB de mémoire active de pointe pour
+  `edge0-35b`, environ 1,0 GB pour `edge0-8b` (contextes courts ; voir
+  [Benchmark](#benchmark)). Prévoyez de la marge pour l'OS, le tokenizer
+  et la croissance de la KV cache en contexte long.
+- **Disque** : les checkpoints 4-bit pèsent environ 23 GB (`edge0-35b`) et
+  environ 4,2 GB (`edge0-8b`) ; les poids des experts sont mmappés et lus
+  à la demande, ils ne sont pas chargés en RAM au démarrage.
+
+#### 1) Installer
+
+```bash
+cd python
+# Python >= 3.10; the MLX backend requires macOS with Apple Silicon
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+#### 2) Télécharger un modèle
+
+Les deux gammes sont publiées sur Hugging Face et ModelScope — chaque
+dépôt regroupe le checkpoint de base et les adaptateurs LoRA + prerouter
+entraînés dans **un seul répertoire**, de sorte qu'un unique
+téléchargement constitue un modèle prêt à l'emploi :
+
+- [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) (~23 GB) · [Miroir ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
+- [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) (~4,2 GB) · [Miroir ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
+
+```bash
+# with the repo's helper (defaults to the two repos above):
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview     --local-dir models/edge0-35b
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview     --local-dir models/edge0-8b
+```
+
+Quelle que soit la méthode, vous obtenez un répertoire comme celui-ci :
+
+```
+models/edge0-35b/
+├── config.json, model-*.safetensors, tokenizer files   # base checkpoint
+├── lora_edge0_35b.safetensors          # trained LoRA adapters
+└── prerouter_edge0_35b.safetensors     # trained prerouter heads
+```
+
+#### 3) Pointer edge0 vers le modèle
+
+Les noms de gamme sont résolus vers des répertoires locaux via des
+variables d'environnement (l'emplacement du téléchargement est libre) :
+
+```bash
+export EDGE0_35B_MODEL=$PWD/models/edge0-35b
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+Ou ignorez complètement les variables d'environnement et passez
+directement le répertoire — la gamme est détectée automatiquement depuis
+le `config.json` du checkpoint :
+
+```bash
+edge0 demo models/edge0-35b
+edge0 serve models/edge0-8b
+```
+
+#### 4) Exécuter
+
+```bash
+# quick demo
+edge0 demo edge0-35b
+
+# serve (OpenAI-compatible /v1/chat/completions)
+edge0 serve edge0-35b
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-35b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+
+# 5) One-shot chat (pass --max-new to cap length; add --show-thinking to
+#    print the model's reasoning block too)
+edge0 chat edge0-35b --prompt "Explain streaming inference in one sentence."
+```
+
+`python -m edge0 ...` est équivalent à `edge0 ...`.
+
+#### API Python
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`python/examples/demo.py` est le même parcours minimal (`edge0 demo`
+exécute exactement ce chemin).
+
+#### Modèles et adaptateurs
+
+- **Checkpoint** : le répertoire du modèle d'origine (`config.json`,
+  `model*.safetensors`, tokenizer). `edge0 serve <dir>` /
+  `AutoEngine.from_pretrained(<dir>)` détectent la gamme depuis
+  `config.json`.
+- Les **adaptateurs** (LoRA + prerouter, safetensors) sont résolus
+  automatiquement depuis l'un ou l'autre emplacement :
+  - le répertoire du modèle (recommandé) : à côté de la base, par exemple
+    `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors` ;
+  - `artifacts/` à la racine du projet Python (ignoré par git) : un cache
+    de secours optionnel pour les safetensors d'adaptateurs non
+    colocalisés avec le modèle.
+- Les dépôts de modèles publiés regroupent à la fois le checkpoint de base
+  et la version par défaut actuelle des adaptateurs, de sorte que
+  `scripts/fetch_models.py` produit un répertoire de modèle prêt à
+  l'emploi. Consultez la page de documentation de chaque modèle pour la
+  provenance de ses adaptateurs (données d'entraînement, disposition des
+  couches propriétaires).
+- Les deux adaptateurs sont requis pour le pipeline prerouter + LoRA ; si
+  un fichier manque, `edge0` échoue avec un message clair (ou passez
+  `--no-prerouter` / `--no-lora` pour exécuter le modèle de base seul).
+
+#### Documentation
+
+- [Architecture](docs/architecture.md)
+- [Attention](docs/attention.md) / [MoE](docs/moe.md) / [Streaming SSD](docs/streaming.md) / [prerouter](docs/prerouter.md)
+- [Ajouter un modèle](docs/adding-a-model.md)
+- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md)
+- Rapport technique : [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](paper/main.pdf))
+
+### macOS / iOS / Android / Windows
+
+Les quatre moteurs par plateforme sont publiés en open source dans ce
+dépôt — plus de détails dans le README de chaque répertoire :
+
+- **macOS** : CLI locale / démon / application de bureau (Rust) — voir [`macos/README_fr.md`](macos/README_fr.md)
+- **iOS** : application iPhone embarquée (Swift + MLX Swift) — voir [`ios/README_fr.md`](ios/README_fr.md)
+- **Android** : application embarquée + moteur natif (Kotlin) — voir [`android/README_fr.md`](android/README_fr.md)
+- **Windows** : application de bureau + moteur natif (C++ + Vulkan) — voir [`windows/README_fr.md`](windows/README_fr.md)
+
+Le **framework d'inférence unifié** — une couche d'accès unique, un
+runtime s'adaptant automatiquement à iOS / macOS / Android / Windows /
+Python — arrivera au **T4 2026** ; voir la
+[Feuille de route](#feuille-de-route).
+
+## Feuille de route
+
+### T4 2026
+
+**Plateformes et systèmes**
+
+- **Framework d'inférence unifié edge0** — nous publierons en open source
+  un framework d'inférence unifié : **une couche d'accès unifiée** (une
+  API unique pour le chat / le serving / l'utilisation embarquée), avec le
+  **runtime s'adaptant automatiquement à la plateforme matérielle** — iOS,
+  macOS, Android, Windows et Python. Il s'appuie sur les moteurs par
+  plateforme déjà publiés en open source dans ce dépôt (`ios/` · `macos/`
+  · `android/` · `windows/`).
+- **Backend CUDA** pour le framework Python — emplacement réservé dans
+  `python/src/edge0/backends/cuda/`, le code central ne nécessite aucune
+  modification.
+
+**Modèles et algorithmes**
+
+Le T4 travaille sur deux fronts : intégrer une architecture de nouvelle
+génération dans le framework, et transformer le raisonnement latent en un
+véritable gain de latence et non seulement arithmétique.
+
+- **Prise en charge des architectures de nouvelle génération (classe
+  Qwen3.8-Flash)** — exécuter sur edge0 l'attention linéaire hybride
+  (GDN + QSA), les résidus multi-branches à portail et les topologies
+  d'embedding N-gram. Ces conceptions se prêtent naturellement au SSD
+  expert offload en streaming : l'attention à état O(1) évite que les
+  longues réflexions ne deviennent un problème de KV cache, et les tables
+  N-gram, consultables uniquement en lecture (lookup-only), se chargent en
+  streaming à la demande. Objectif : que la gamme s'exécute sur un seul
+  appareil et que les benchmarks se situent à un écart acceptable de la
+  base fp16.
+- **Raisonnement latent + pré-prédiction d'experts par lots** — faire du
+  raisonnement latent une économie de *latence*, et pas seulement
+  arithmétique. Le problème d'ingénierie central : faire passer le routage
+  des experts d'une prédiction par position à **une prédiction par bloc**,
+  afin qu'une seule prédiction couvre toutes les positions et tous les
+  tours d'un bloc et que **le volume de chargement des experts soit
+  découplé du nombre de tours de la boucle de raisonnement** — plus un
+  préchargement inter-blocs qui charge les experts du bloc suivant dans la
+  fenêtre de calcul du bloc courant. Les progrès sont mesurés en **temps
+  de phase de réflexion de bout en bout à précision égale** (jamais en
+  tokens/s).
+- Davantage de gammes de modèles et de versions d'adaptateurs sur le
+  pipeline existant.
+
+## Contribution
+
+Les contributions sont les bienvenues — issues, PR, rapports de benchmark
+et portages de modèles, tout compte.
+
+**Framework Python** (disponible maintenant) :
+
+```bash
+cd python
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q
+                       # real-weight generation; missing tiers are skipped
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+                       # staged vs exact consistency + generation smoke
+scripts/generate_example.py   # full-pipeline API example
+examples/demo.py              # minimal API walkthrough
+```
+
+La CI exécute les tests unitaires (macOS + MLX) et une suite d'hygiène du
+dépôt (pas de chemins codés en dur, vérification des frontières de backend
+et des secrets) sur chaque PR.
+
+**Runtimes par plateforme** (macOS / iOS / Android / Windows) : chaque
+répertoire de plateforme dispose de son propre guide de compilation et de
+ses tests — voir le README de chaque répertoire.
+
+Workflow : fork → branche de fonctionnalité → PR vers `main`. Merci de
+garder la suite d'hygiène au vert et d'ajouter des tests pour les nouveaux
+comportements.
+
+## Citation
+
+Si vous trouvez edge0 utile, merci de citer notre rapport technique :
+
+```bibtex
+@article{lin2026other,
+  title   = {The Other Half of the Memory Wall: Serving 35B MoEs from SSD
+             with Trained Routing Prediction},
+  author  = {Lin, Yu and Wang, Yiming and Cai, Runyuan and Liu, Hanze and
+             Zeng, Xiaodong},
+  journal = {arXiv preprint arXiv:2609.18063},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2609.18063}
+}
+```
+
+## Nous contacter
+
+Les canaux communautaires et d'assistance arrivent bientôt — cette section
+listera les moyens officiels de nous joindre :
+
+- **Email** : samuel@edge0.ai
+
+Pour les bugs et les demandes de fonctionnalités, merci d'utiliser
+[GitHub Issues](https://github.com/Edge0-AI/Edge0/issues).
+
+## Licence
+
+Apache-2.0, y compris le code tiers intégré (voir [NOTICE](NOTICE)).
diff --git a/README_ja.md b/README_ja.md
new file mode 100644
index 0000000..d37cae2
--- /dev/null
+++ b/README_ja.md
@@ -0,0 +1,302 @@
+<div align="center">
+
+<img src="assets/20260908-223115.jpg" alt="edge0" width="100%">
+
+# edge0
+
+**オープンソースのストリーミング MoE 推論フレームワーク — SSD expert offload + Recover-LoRA + prerouter ルーティング予測。**
+
+**Python** · **macOS** · **iOS** · **Android** — 1 つのレシピで、すべてのデバイスに対応。
+
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--35B--A3B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--8B--A1B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+[![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--35B--A3B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
+[![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--8B--A1B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
+[![arXiv](https://img.shields.io/badge/arXiv-2609.18063-B31B1B?style=for-the-badge&logo=arxiv&logoColor=white)](https://arxiv.org/abs/2609.18063)
+[![GitHub](https://img.shields.io/badge/GitHub-Edge0--AI%2FEdge0-black?style=for-the-badge&logo=github)](https://github.com/Edge0-AI/Edge0)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue?style=for-the-badge)](LICENSE)
+
+[English](README.md) | [中文](README_zh.md) | 日本語 | [Español](README_es.md) | [Français](README_fr.md)
+
+</div>
+
+## ニュース
+
+- **[2026-09-30]** **4 つのプラットフォーム — iOS、macOS、Android、Windows — 向けの edge0 推論エンジン**をリリースしました。ユーザーはアーキテクチャとプラットフォームをまたいで最適な推論体験を得られます。ソースはこのリポジトリでオープンソース化されています([`ios/`](ios/) · [`macos/`](macos/) · [`android/`](android/) · [`windows/`](windows/))— 詳細は各ディレクトリの README を参照してください。**統合推論フレームワーク**は **2026 年 Q4** に控えています。[ロードマップ](#ロードマップ)を参照。
+- **[2026-09-20]** `edge0 demo / chat / serve` に `--prefill-ondemand` が登場([#112](https://github.com/Edge0-AI/Edge0/pull/112))。
+- **[2026-09-16]** 技術レポートが arXiv に掲載: [The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction](https://arxiv.org/abs/2609.18063)。
+- **[2026-09-08]** **edge0** の初オープンソースリリース。両モデルティア — [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) と [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) — を Hugging Face と ModelScope で同時公開。
+
+## 概要
+
+**edge0** はオープンソースのストリーミング MoE 推論フレームワークです。本番で実証済みのレシピ — **SSD expert offload + Recover-LoRA + prerouter ルーティング予測** — を拡張可能なフレームワークへと一般化し、コンシューマーハードウェアで大規模スパース MoE モデルを実行します。ピークメモリは、パラメータ数ではなく*アクティブ*なエキスパート集合によって上限が決まります。
+
+### コア機構
+
+- **SSD expert offload**: エキスパートの重みは必要に応じてストレージからストリーミングされます。ピークメモリはパラメータ数ではなく、アクティブな集合によって上限が決まります。
+- **Prerouter**: 訓練済みヘッドがエキスパートルーティングを 1 ステップ先に予測するため、エキスパートのロードがフォワードパスを停滞させるのではなくオーバーラップします — デコードスループットは**最大 +59%**。ストレージレイテンシ、モデルサイズ、ルーティング幅 *K* が増すほど効果は大きくなります。
+- **Recover-LoRA**: int4 のベースは凍結され、FP 教師からの蒸留によって LoRA アダプタが訓練されることで、4-bit 時の量子化損失の大部分を回復します([品質](#品質)を参照)。アダプタはマージされません。1 つの読み取り専用ベースが複数のアダプタセットに対応します。
+
+### プラットフォーム
+
+1 つのリポジトリ、1 つのレシピ、プラットフォームごとのランタイム:
+
+| プラットフォーム | ディレクトリ | スタック | ステータス |
+|---|---|---|---|
+| **Python**(macOS · Apple Silicon) | [`python/`](python/README_ja.md) | Python + MLX | ✅ 現在利用可能 |
+| **macOS** アプリ & CLI | [`macos/`](macos/README_ja.md) | Rust | ✅ オープンソース化(2026-09-30) |
+| **iOS** アプリ | [`ios/`](ios/README_ja.md) | Swift + MLX Swift | ✅ オープンソース化(2026-09-30) |
+| **Android** アプリ & エンジン | [`android/`](android/README_ja.md) | Kotlin + ネイティブエンジン | ✅ オープンソース化(2026-09-30) |
+| **Windows** アプリ & エンジン | [`windows/`](windows/README_ja.md) | C++ + Vulkan | ✅ オープンソース化(2026-09-30) |
+
+### モデル
+
+フレームワークには 2 つのモデルティアが付属します。各ティアはエンドツーエンドのリリースであり、公開されるチェックポイント、訓練済み LoRA アダプタ、訓練済み prerouter ヘッドが 1 つの単位として機能します。
+
+| ティア | 公開チェックポイント | 推論プロファイル |
+|---|---|---|
+| `edge0-35b` | [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview) · [ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview) | 4-bit、40 層、256 エキスパート、prerouter K=4 |
+| `edge0-8b` | [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [ModelScope](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview) | 4-bit、24 層、128 エキスパート、prerouter K=8 |
+
+両方のチェックポイントともオープンのスパース MoE ベースモデル(それぞれ Qwen3.6-35B-A3B と Ling 3.0 bailing hybrid)上に構築されており、本フレームワークのために実施された LoRA と prerouter の訓練成果物を同梱しています — アダプタファイルは各チェックポイントと同じ場所に配置され、自動的にロードされるため、`edge0 serve <tier>` はすぐさま訓練済みパイプラインを実行できます。
+
+### 設計
+
+- **transformers スタイルの使い方**: `AutoModel` / `AutoConfig` / `AutoEngine` がモデル名からティアを解決します。
+- **バックエンド分離(Python フレームワーク)**: Python フレームワーク内では、すべての MLX コードは `python/src/edge0/backends/mlx/` 配下に置かれます。コアロジック(モデル仕様、prerouter、ストリーミングエキスパートプール、サーバー)はバックエンドファサード(`backends/base.py`)にのみ依存するため、新しいバックエンドは同じファサードを実装するだけで済み(`backends/cuda/` は予約スロット)、コアコードの変更はゼロです。iOS / macOS / Android エンジンは現在プラットフォームネイティブのスタックで提供されています — すべてのプラットフォームを 1 つのアクセス層の下にまとめることこそ、統合推論フレームワーク([ロードマップ](#ロードマップ)を参照)が実現するものです。
+- **safetensors 形式のアダプタ**: LoRA と prerouter の重みは、来歴メタデータ(source、version、owner layers)付きの `.safetensors` ファイルであり、モデルディレクトリまたは `artifacts/` から解決されます。
+- **モデル + アダプタを 1 つのディレクトリに**: モデルディレクトリには、ベースチェックポイント(`config.json` / `model*.safetensors` / トークナイザー)とそのモデルのアダプタの両方が格納されます。アダプタのアップグレードはアダプタファイルの交換のみ — ベースは読み取り専用のままで、決してマージされません。
+
+### 品質
+
+すべてのベンチマークは、edge0 モデル(int4 + 訓練済みアダプタ + prerouter ルーティング)とオリジナルの fp16 ベースモデルの両方について、同一の設定とパラメータで [OpenCompass](https://github.com/open-compass/opencompass) を使って私たち自身が実行しました。edge0 パイプラインの損失は小さく、**edge0-35b で平均 3.9 ポイント、edge0-8b で 2.8 ポイント**です(MMLU-Pro はベースを上回ります)。満点 100:
+
+| ベンチマーク | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **平均** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+### ベンチマーク
+
+`python/examples/bench.py` で測定(3.3k トークンのプロンプトのプリフィル → サンプリングによるウォームアップ 10 ステップ → 計測付きのサンプリングデコード 200 トークン、ティアごとに 2 回実行):
+
+| ティア | デコード速度 | プリフィルスループット(コールド / ウォーム)* | ピークアクティブメモリ | テストマシン |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*コールド = プロセス起動後の最初のリクエスト(エキスパートの重みが SSD からフォールトインする);ウォーム = 以降のリクエスト(ページキャッシュに常駐)。プリフィルの数値は約 3.3k トークンのプロンプトに対するスループット(`BENCH_LONG=1`)。*
+
+再現手順:
+
+```bash
+cd python
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+```
+
+## はじめに
+
+### Python(macOS · Apple Silicon)
+
+#### 要件
+
+- **OS / ハードウェア**: MLX バックエンドは Apple Silicon(M1/M2/M3/M4)搭載の macOS で動作します。CUDA バックエンドはロードマップ上にあります — Python フレームワークはまだ他のプラットフォームをサポートしていません。
+- **Python**: 3.10+(3.12 推奨)。
+- **MLX**: `mlx==0.30.6` / `mlx-metal==0.30.6` と `mlx-lm==0.31.0`(`python/pyproject.toml` を参照)。Apple A18 / A18 Pro で文字化けした言語混在の出力が発生する場合、古い `mlx` が原因です: `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`([#8](https://github.com/Edge0-AI/Edge0/issues/8))。
+- **メモリ**: `edge0-35b` でピークアクティブメモリ約 2.9 GB、`edge0-8b` で約 1.0 GB(短いコンテキストの場合;[ベンチマーク](#ベンチマーク)を参照)。OS、トークナイザー、長いコンテキストでの KV キャッシュの増加のための余裕を上乗せしてください。
+- **ディスク**: 4-bit チェックポイントは約 23 GB(`edge0-35b`)と約 4.2 GB(`edge0-8b`)です。エキスパートの重みは mmap されてオンデマンドで読み取られ、事前に RAM へロードされることはありません。
+
+#### 1) インストール
+
+```bash
+cd python
+# Python >= 3.10; the MLX backend requires macOS with Apple Silicon
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+#### 2) モデルのダウンロード
+
+2 つのティアは Hugging Face と ModelScope で公開されています — 各リポジトリは、ベースチェックポイントと訓練済み LoRA + prerouter アダプタを**1 つのディレクトリ**にまとめており、1 回のダウンロードですぐ実行できるモデルが揃います:
+
+- [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)(~23 GB)· [ModelScope ミラー](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
+- [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)(~4.2 GB)· [ModelScope ミラー](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
+
+```bash
+# with the repo's helper (defaults to the two repos above):
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview     --local-dir models/edge0-35b
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview     --local-dir models/edge0-8b
+```
+
+どちらの方法でも、次のようなディレクトリが得られます:
+
+```
+models/edge0-35b/
+├── config.json, model-*.safetensors, tokenizer files   # base checkpoint
+├── lora_edge0_35b.safetensors          # trained LoRA adapters
+└── prerouter_edge0_35b.safetensors     # trained prerouter heads
+```
+
+#### 3) edge0 にモデルを指定する
+
+ティア名は環境変数を通じてローカルディレクトリに解決されます(ダウンロード先は自由です):
+
+```bash
+export EDGE0_35B_MODEL=$PWD/models/edge0-35b
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+環境変数を完全に省略して、ディレクトリを直接渡すこともできます — ティアはチェックポイントの `config.json` から自動検出されます:
+
+```bash
+edge0 demo models/edge0-35b
+edge0 serve models/edge0-8b
+```
+
+#### 4) 実行
+
+```bash
+# quick demo
+edge0 demo edge0-35b
+
+# serve (OpenAI-compatible /v1/chat/completions)
+edge0 serve edge0-35b
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-35b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+
+# 5) One-shot chat (pass --max-new to cap length; add --show-thinking to
+#    print the model's reasoning block too)
+edge0 chat edge0-35b --prompt "Explain streaming inference in one sentence."
+```
+
+`python -m edge0 ...` は `edge0 ...` と等価です。
+
+#### Python API
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`python/examples/demo.py` も同じ最小限のウォークスルーです(`edge0 demo` はまさにこのパスを実行します)。
+
+#### モデルとアダプタ
+
+- **チェックポイント**: 元のモデルディレクトリ(`config.json`、`model*.safetensors`、トークナイザー)。`edge0 serve <dir>` / `AutoEngine.from_pretrained(<dir>)` は `config.json` からティアを検出します。
+- **アダプタ**(LoRA + prerouter、safetensors)は、次のいずれかの場所から自動的に解決されます:
+  - モデルディレクトリ(推奨): ベースと同じ場所。例: `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`
+  - Python プロジェクトルートの `artifacts/`(gitignore 済み): モデルと同じ場所にないアダプタ safetensors のための、オプションのフォールバックキャッシュ。
+- 公開されているモデルリポジトリには、ベースチェックポイントと現在のデフォルトのアダプタリリースの両方が同梱されているため、`scripts/fetch_models.py` はすぐ実行できるモデルディレクトリを生成します。アダプタの来歴(訓練データ、owner-layer レイアウト)は各モデルのドキュメントページで確認してください。
+- prerouter + LoRA パイプラインには両方のアダプタが必要です。ファイルが不足している場合、`edge0` は明確なメッセージとともに失敗します(あるいは `--no-prerouter` / `--no-lora` を渡せば、素のベースモデルを実行できます)。
+
+#### ドキュメント
+
+- [アーキテクチャ](docs/architecture.md)
+- [アテンション](docs/attention.md) / [MoE](docs/moe.md) / [SSD ストリーミング](docs/streaming.md) / [prerouter](docs/prerouter.md)
+- [モデルの追加](docs/adding-a-model.md)
+- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md)
+- 技術レポート: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063)([PDF](paper/main.pdf))
+
+### macOS / iOS / Android / Windows
+
+4 つのプラットフォームエンジンはこのリポジトリでオープンソース化されています — 詳細は各ディレクトリの README を参照してください:
+
+- **macOS**: ローカル CLI / デーモン / デスクトップアプリ(Rust)— [`macos/README_ja.md`](macos/README_ja.md) を参照
+- **iOS**: オンデバイス iPhone アプリ(Swift + MLX Swift)— [`ios/README_ja.md`](ios/README_ja.md) を参照
+- **Android**: オンデバイスアプリ + ネイティブエンジン(Kotlin)— [`android/README_ja.md`](android/README_ja.md) を参照
+- **Windows**: デスクトップアプリ + ネイティブエンジン(C++ + Vulkan)— [`windows/README_ja.md`](windows/README_ja.md) を参照
+
+**統合推論フレームワーク** — 1 つのアクセス層、iOS / macOS / Android / Windows / Python へ自動適応するランタイム — は **2026 年 Q4** に登場します。[ロードマップ](#ロードマップ)を参照。
+
+## ロードマップ
+
+### 2026 年 Q4
+
+**プラットフォーム & システム**
+
+- **edge0 統合推論フレームワーク** — 統合推論フレームワークをオープンソース化します: **1 つの統一アクセス層**(chat / serve / オンデバイス利用を横断する単一 API)を備え、**ランタイムがハードウェアプラットフォームへ自動適応**します — iOS、macOS、Android、Windows、Python。このリポジトリですでにオープンソース化されているプラットフォームエンジン(`ios/` · `macos/` · `android/` · `windows/`)の上に構築されます。
+- Python フレームワーク向け **CUDA バックエンド** — `python/src/edge0/backends/cuda/` に予約スロットがあり、コアコードの変更はゼロです。
+
+**モデル & アルゴリズム**
+
+Q4 は 2 つの前線で進めます: 次世代アーキテクチャをフレームワークに取り込むこと、そして潜在推論を単なる算術上の削減ではなく実際のレイテンシ削減に変えることです。
+
+- **次世代アーキテクチャのサポート(Qwen3.8-Flash クラス)** — ハイブリッド線形アテンション(GDN + QSA)、ゲート付きマルチブランチ残差、N-gram 埋め込みトポロジを edge0 で実行します。これらの設計は SSD ストリーミングオフロードと自然に相性が良いものです: O(1) 状態のアテンションにより長い思考が KV キャッシュの問題になるのを防ぎ、ルックアップのみの N-gram テーブルはオンデマンドでストリーミングされます。目標: そのティアが単一デバイスで動作し、ベンチマークが fp16 ベースに対して許容可能な差に収まること。
+- **潜在 thinking + バッチ化されたエキスパート事前予測** — 潜在推論を、単なる算術上の削減ではなく*レイテンシ*の削減にします。コアとなるエンジニアリング課題: エキスパートルーティングを位置ごとから**ブロックごとに 1 回**へ移行し、1 回の予測がブロック内のすべての位置とラウンドをカバーすることで、**エキスパートのロード量を推論ループの回数から切り離す**こと — さらに、現在のブロックの計算ウィンドウ内で次のブロックのエキスパートをロードするクロスブロックプリフェッチも加えます。進捗は**同一精度でのエンドツーエンドの thinking フェーズ時間**で測定します(決して tokens/s では測りません)。
+- 既存パイプラインにおける、さらなるモデルティアとアダプタのリリース。
+
+## コントリビュート
+
+コントリビューションを歓迎します — issue、PR、ベンチマークレポート、モデルポート、すべてが対象です。
+
+**Python フレームワーク**(現在利用可能):
+
+```bash
+cd python
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q
+                       # real-weight generation; missing tiers are skipped
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+                       # staged vs exact consistency + generation smoke
+scripts/generate_example.py   # full-pipeline API example
+examples/demo.py              # minimal API walkthrough
+```
+
+CI はすべての PR で、ユニットテスト(macOS + MLX)とリポジトリ衛生スイート(ハードコードされたパスの禁止、バックエンド境界とシークレットのチェック)を実行します。
+
+**プラットフォームランタイム**(macOS / iOS / Android / Windows): 各プラットフォームのディレクトリには独自のビルドガイドとテストが含まれています——各ディレクトリの README を参照してください。
+
+ワークフロー: fork → フィーチャーブランチ → `main` への PR。衛生スイートをグリーンのまま保ち、新しい動作にはテストを追加してください。
+
+## 引用
+
+edge0 が役に立った場合は、私たちの技術レポートを引用してください:
+
+```bibtex
+@article{lin2026other,
+  title   = {The Other Half of the Memory Wall: Serving 35B MoEs from SSD
+             with Trained Routing Prediction},
+  author  = {Lin, Yu and Wang, Yiming and Cai, Runyuan and Liu, Hanze and
+             Zeng, Xiaodong},
+  journal = {arXiv preprint arXiv:2609.18063},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2609.18063}
+}
+```
+
+## お問い合わせ
+
+コミュニティとサポートのチャネルは近日公開予定です — このセクションに、私たちへの公式の連絡方法を掲載します:
+
+- **メール**: samuel@edge0.ai
+
+バグと機能リクエストには、[GitHub Issues](https://github.com/Edge0-AI/Edge0/issues) をご利用ください。
+
+## ライセンス
+
+Apache-2.0。ベンダー化されたサードパーティコードを含みます([NOTICE](NOTICE) を参照)。
diff --git a/README_zh.md b/README_zh.md
index 09f0d5e..328fb5e 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -6,6 +6,8 @@
 
 **开源流式 MoE 推理框架 —— SSD 专家 offload + Recover-LoRA + prerouter 路由预判**
 
+**Python** · **macOS** · **iOS** · **Android** —— 一套配方，全端落地。
+
 [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--35B--A3B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
 [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Edge0--8B--A1B--preview-yellow?style=for-the-badge)](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
 [![ModelScope](https://img.shields.io/badge/ModelScope-Edge0--35B--A3B--preview-624AFF?style=for-the-badge)](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
@@ -14,13 +16,57 @@
 [![GitHub](https://img.shields.io/badge/GitHub-Edge0--AI%2FEdge0-black?style=for-the-badge&logo=github)](https://github.com/Edge0-AI/Edge0)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue?style=for-the-badge)](LICENSE)
 
-[English](README.md) | 中文
+[English](README.md) | 中文 | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
 
 </div>
 
+## 新闻
+
+- **【2026-09-30】** 我们发布了**四端推理引擎**，适配 iOS、macOS、
+  Android、Windows 平台，让用户在不同架构、不同平台上都有最佳的推理
+  体验。四端源码已开源至本仓库（[`ios/`](ios/) · [`macos/`](macos/) ·
+  [`android/`](android/) · [`windows/`](windows/)），更多细节见各目录
+  README；**统一推理框架**将于 **2026 Q4** 发布，详见[路线图](#路线图)。
+- **【2026-09-20】** `edge0 demo / chat / serve` 新增 `--prefill-ondemand`
+  （[#112](https://github.com/Edge0-AI/Edge0/pull/112)）。
+- **【2026-09-16】** 技术报告上线 arXiv：[The Other Half of the Memory
+  Wall: Serving 35B MoEs from SSD with Trained Routing
+  Prediction](https://arxiv.org/abs/2609.18063)。
+- **【2026-09-08】** **edge0** 首次开源发布，两个模型档位
+  [`Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+  与 [`Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+  同步登陆 Hugging Face 与 ModelScope。
+
+## 关于 edge0
+
 **edge0** 是一个开源的流式 MoE 推理框架：把「SSD 专家 offload +
-Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架。后端隔离设计，
-当前实现 MLX 后端（Apple Silicon），更多平台（CUDA 等）即将接入。
+Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架，让大型稀疏
+MoE 模型跑在消费级硬件上——峰值内存由**激活**专家集而非参数量决定。
+
+### 核心机制
+
+- **SSD 专家 offload**：专家权重按需从存储流式加载，峰值内存由
+  激活集而非参数量决定；
+- **prerouter**：训练头提前一步预测专家路由，专家装载与前向计算重叠
+  而非阻塞——解码吞吐**最高 +59%**，收益随存储延迟、模型规模与路由
+  宽度 *K* 增大；
+- **Recover-LoRA**：冻结 int4 基模，用 FP teacher 蒸馏训练 LoRA，
+  在 4bit 下恢复绝大部分量化损失（见[质量](#质量)）。适配器不合并，
+  一份只读基模服务多套适配器。
+
+### 平台
+
+一个仓库、一套配方、各平台独立 runtime：
+
+| 平台 | 目录 | 技术栈 | 状态 |
+|---|---|---|---|
+| **Python**（macOS · Apple Silicon） | [`python/`](python/README_zh.md) | Python + MLX | ✅ 现已可用 |
+| **macOS** 桌面 App 与 CLI | [`macos/`](macos/README_zh.md) | Rust | ✅ 已开源（2026-09-30） |
+| **iOS** App | [`ios/`](ios/README_zh.md) | Swift + MLX Swift | ✅ 已开源（2026-09-30） |
+| **Android** App 与引擎 | [`android/`](android/README_zh.md) | Kotlin + 原生引擎 | ✅ 已开源（2026-09-30） |
+| **Windows** App 与引擎 | [`windows/`](windows/README_zh.md) | C++ + Vulkan | ✅ 已开源（2026-09-30） |
+
+### 模型
 
 框架随附两个模型档位。每个档位是一个端到端发布：发布的 checkpoint、
 训练好的 LoRA 适配器与训练好的 prerouter 头作为整体协同工作。
@@ -35,60 +81,92 @@ Recover-LoRA + prerouter 路由预判」抽象成可扩展的通用框架。后
 适配器文件与 checkpoint 同目录、自动加载，`edge0 serve <tier>` 开箱即跑
 训练好的完整管线。
 
-## 环境要求
-
-- **系统 / 硬件**：MLX 后端目前仅支持 Apple Silicon 的 macOS
-  （M1/M2/M3/M4）；CUDA 后端在路线图中，其余平台暂不支持。
-- **Python**：3.10+（推荐 3.12）。
-- **MLX**：`mlx==0.30.6` / `mlx-metal==0.30.6`（`mlx-lm==0.31.0`，见
-  `pyproject.toml`）。Apple A18 / A18 Pro 上输出乱码 = mlx 版本旧：
-  `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
-  （[#8](https://github.com/Edge0-AI/Edge0/issues/8)）。
-- **内存**：短上下文下 `edge0-35b` ≈2.9 GB、`edge0-8b` ≈1.0 GB
-  峰值激活内存（见[性能实测](#性能实测)）；另为系统、tokenizer 与
-  长上下文 KV 增长预留余量。
-- **磁盘**：4bit checkpoint 约 23 GB（`edge0-35b`）/ 4.2 GB
-  （`edge0-8b`）；专家权重 mmap 按需读取，不一次性载入内存。
-
-## 设计
+### 设计
 
 - **像 transformers 一样使用**：`AutoModel` / `AutoConfig` / `AutoEngine`
   按模型名自动选类；
-- **后端隔离**：全部 MLX 代码收在 `edge0/backends/mlx/`，核心逻辑
-  （模型 spec / prerouter / 流式专家池 / server）只依赖后端门面
-  （`edge0/backends/base.py` 的 `core` / `nn` 门面），新增后端实现同一门面
-  即可平级接入（`backends/cuda/` 预留插槽），核心代码零改动；
+- **后端隔离（Python 框架）**：在 Python 框架内，全部 MLX 代码收在
+  `python/src/edge0/backends/mlx/`，核心逻辑（模型 spec / prerouter /
+  流式专家池 / server）只依赖后端门面（`backends/base.py` 的 `core` /
+  `nn` 门面），新增后端实现同一门面即可平级接入（`backends/cuda/`
+  预留插槽），核心代码零改动。iOS / macOS / Android 引擎目前是各自的
+  平台原生技术栈——把所有平台收进统一接入层，正是统一推理框架
+  （见[路线图](#路线图)）要交付的内容；
 - **适配器统一为 safetensors**：LoRA 与 prerouter 权重均为带元数据
   （来源、版本、owner 层）的 `.safetensors`，放模型目录或 `artifacts/`
   均可自动解析；
 - **模型 + 适配器同目录布局**：一个模型目录同时放基模（`config.json` /
-  `model*.safetensors` / tokenizer）和该模型的适配器，升级适配器只换适配器文件
-  文件，基模不动、不 merge。
+  `model*.safetensors` / tokenizer）和该模型的适配器，升级适配器只换
+  适配器文件，基模不动、不 merge。
 
-## 核心机制
+### 质量
 
-- **SSD 专家 offload**：专家权重按需从存储流式加载，峰值内存由
-  激活集而非参数量决定；
-- **prerouter**：训练头提前一步预测专家路由，专家装载与前向计算重叠
-  而非阻塞——解码吞吐**最高 +59%**，收益随存储延迟、模型规模与路由
-  宽度 *K* 增大；
-- **Recover-LoRA**：冻结 int4 基模，用 FP teacher 蒸馏训练 LoRA，
-  在 4bit 下恢复绝大部分量化损失（见[质量](#质量)）。适配器不合并，
-  一份只读基模服务多套适配器。
+全部评测由我们使用 [OpenCompass](https://github.com/open-compass/opencompass)、
+在完全相同的设置与参数下对 edge0 模型（int4 + 训练适配器 + prerouter 路由）
+与原 fp16 基座模型测得。edge0 管线的损失很小：**edge0-35b 平均仅落后
+3.9 分、edge0-8b 落后 2.8 分**（MMLU-Pro 甚至反超基座）。满分 100：
+
+| 评测集 | edge0-35b（int4） | Qwen3.6-35B-A3B（fp16） | edge0-8b（int4） | Ling 3.0 tiny（fp16） |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **平均** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+### 性能实测
+
+`python/examples/bench.py` 实测（3.3k token prompt prefill → 10 步采样
+warmup → 200 token 计时段，每档 2 轮）：
+
+| 档位 | 解码速度 | Prefill 吞吐（冷/热）* | 峰值 active 内存 | 测试机器 |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*冷 = 进程启动后首请求（专家权重从 SSD 逐页换入）；热 = 后续请求（页缓存常驻）。Prefill 为 ≈3.3k token 长 prompt 的吞吐（`BENCH_LONG=1`）。*
+
+复现：
+
+```bash
+cd python
+python examples/bench.py edge0-35b    # 经 $EDGE0_35B_MODEL
+python examples/bench.py edge0-8b     # 经 $EDGE0_8B_MODEL
+```
 
 ## 快速开始
 
-### 1) 安装
+### Python（macOS · Apple Silicon）
+
+#### 环境要求
+
+- **系统 / 硬件**：MLX 后端目前仅支持 Apple Silicon 的 macOS
+  （M1/M2/M3/M4）；CUDA 后端在路线图中，Python 框架暂不支持其他平台。
+- **Python**：3.10+（推荐 3.12）。
+- **MLX**：`mlx==0.30.6` / `mlx-metal==0.30.6`（`mlx-lm==0.31.0`，见
+  `python/pyproject.toml`）。Apple A18 / A18 Pro 上输出乱码 = mlx 版本旧：
+  `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
+  （[#8](https://github.com/Edge0-AI/Edge0/issues/8)）。
+- **内存**：短上下文下 `edge0-35b` ≈2.9 GB、`edge0-8b` ≈1.0 GB
+  峰值激活内存（见[性能实测](#性能实测)）；另为系统、tokenizer 与
+  长上下文 KV 增长预留余量。
+- **磁盘**：4bit checkpoint 约 23 GB（`edge0-35b`）/ 4.2 GB
+  （`edge0-8b`）；专家权重 mmap 按需读取，不一次性载入内存。
+
+#### 1) 安装
 
 ```bash
+cd python
 # Python ≥3.10；MLX 后端需 macOS + Apple Silicon
 python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
 ```
 
-### 2) 下载模型
+#### 2) 下载模型
 
-两个档位发布在 Hugging Face 与 ModelScope——每个仓库把基模 checkpoint 与训练好的
-LoRA + prerouter 适配器打包在**同一目录**，一次下载即为可运行的模型：
+两个档位发布在 Hugging Face 与 ModelScope——每个仓库把基模 checkpoint
+与训练好的 LoRA + prerouter 适配器打包在**同一目录**，一次下载即为
+可运行的模型：
 
 - [`Edge0/Edge0-35B-A3B-preview`](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)（约 23 GB） · [ModelScope 镜像](https://www.modelscope.cn/models/Edge0/Edge0-35B-A3B-preview)
 - [`Edge0/Edge0-8B-A1B-preview`](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)（约 4.2 GB） · [ModelScope 镜像](https://www.modelscope.cn/models/Edge0/Edge0-8B-A1B-preview)
@@ -114,7 +192,7 @@ models/edge0-35b/
 └── prerouter_edge0_35b.safetensors     # 训练好的 prerouter 头
 ```
 
-### 3) 指向模型目录
+#### 3) 指向模型目录
 
 档位名经环境变量解析到本地目录（放哪由你决定）：
 
@@ -131,7 +209,7 @@ edge0 demo models/edge0-35b
 edge0 serve models/edge0-8b
 ```
 
-### 4) 运行
+#### 4) 运行
 
 ```bash
 # 快速演示
@@ -144,7 +222,7 @@ edge0 serve edge0-35b
 ```bash
 curl http://127.0.0.1:8000/v1/chat/completions \
   -H 'Content-Type: application/json' \
-  -d '{"messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+  -d '{"model":"edge0-35b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
 
 # 5) 单轮对话（--max-new 限制生成长度；加 --show-thinking 会一并打印思考块）
 edge0 chat edge0-35b --prompt "用一句话解释流式推理。"
@@ -152,7 +230,7 @@ edge0 chat edge0-35b --prompt "用一句话解释流式推理。"
 
 `python -m edge0 ...` 等价于 `edge0 ...`。
 
-### Python API
+#### Python API
 
 ```python
 from edge0 import AutoEngine
@@ -169,9 +247,9 @@ print(engine._tok.decode(tokens))
 engine.close()   # 释放 mmap / 专家缓存
 ```
 
-`examples/demo.py` 就是这条最小路径（`edge0 demo` 内部等价运行）。
+`python/examples/demo.py` 就是这条最小路径（`edge0 demo` 内部等价运行）。
 
-### 模型与适配器
+#### 模型与适配器
 
 - **checkpoint**：原始模型目录（`config.json`、`model*.safetensors`、
   tokenizer）。`edge0 serve <dir>` / `AutoEngine.from_pretrained(<dir>)`
@@ -179,68 +257,122 @@ engine.close()   # 释放 mmap / 专家缓存
 - **适配器**（LoRA + prerouter，safetensors）放两处任一，自动解析：
   - **模型目录内**（推荐）：与基模同目录，如
     `lora_edge0_35b.safetensors` + `prerouter_edge0_35b.safetensors`；
-  - `artifacts/`（仓库根，gitignored）：`edge0 convert-adapters` 从
-    训练侧 npz 一次性转换。
+  - `artifacts/`（Python 项目根，gitignored）：可选的回退缓存，放置未
+    与模型同目录的适配器 safetensors。
 - 发布的模型仓库同时包含基模与当前默认适配器版本，
   `scripts/fetch_models.py` 下载后即为可运行的模型目录。适配器来源
   （训练数据、owner 层分布）见各模型文档页。
 - prerouter + LoRA 两条适配器都必需；缺文件时 `edge0` 会给出明确报错
   （也可加 `--no-prerouter` / `--no-lora` 直接跑裸基模）。
 
-## 质量
+#### 文档
 
-全部评测由我们使用 [OpenCompass](https://github.com/open-compass/opencompass)、
-在完全相同的设置与参数下对 edge0 模型（int4 + 训练适配器 + prerouter 路由）
-与原 fp16 基座模型测得。edge0 管线的损失很小：**edge0-35b 平均仅落后
-3.9 分、edge0-8b 落后 2.8 分**（MMLU-Pro 甚至反超基座）。满分 100：
+- [架构总览](docs/architecture.md)
+- [注意力抽象](docs/attention.md) / [MoE 抽象](docs/moe.md) / [SSD 流式](docs/streaming.md) / [prerouter](docs/prerouter.md)
+- [如何接入新模型](docs/adding-a-model.md)
+- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md)
+- 技术报告：[The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063)（[PDF](paper/main.pdf)）
 
-| 评测集 | edge0-35b（int4） | Qwen3.6-35B-A3B（fp16） | edge0-8b（int4） | Ling 3.0 tiny（fp16） |
-|---|---:|---:|---:|---:|
-| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
-| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
-| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
-| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
-| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
-| **平均** | **79.2** | **83.2** | **69.9** | **72.7** |
+### macOS / iOS / Android / Windows
 
-## 性能实测
+四端推理引擎已开源在本仓库，更多细节见各目录 README：
 
-`examples/bench.py` 实测（3.3k token prompt prefill → 10 步采样 warmup →
-200 token 计时段，每档 2 轮）：
+- **macOS**：本地 CLI / daemon / 桌面 App（Rust）—— 详见
+  [`macos/README_zh.md`](macos/README_zh.md)
+- **iOS**：iPhone 端侧 App（Swift + MLX Swift）—— 详见
+  [`ios/README_zh.md`](ios/README_zh.md)
+- **Android**：端侧 App + 原生引擎（Kotlin）—— 详见
+  [`android/README_zh.md`](android/README_zh.md)
+- **Windows**：桌面 App + 原生引擎（C++ + Vulkan）—— 详见
+  [`windows/README_zh.md`](windows/README_zh.md)
 
-| 档位 | 解码速度 | Prefill 吞吐（冷/热）* | 峰值 active 内存 | 测试机器 |
-|---|---|---|---|---|
-| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
-| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+**统一推理框架**（接入层统一，runtime 自动适配 iOS / macOS / Android /
+Windows / Python）将于 **2026 Q4** 发布，见[路线图](#路线图)。
 
-*冷 = 进程启动后首请求（专家权重从 SSD 逐页换入）；热 = 后续请求（页缓存常驻）。Prefill 为 ≈3.3k token 长 prompt 的吞吐（`BENCH_LONG=1`）。
+## 路线图
 
-复现：
+### 2026 Q4
 
-```bash
-python examples/bench.py edge0-35b    # 经 $EDGE0_35B_MODEL
-python examples/bench.py edge0-8b    # 经 $EDGE0_8B_MODEL
-```
+**平台与系统**
+
+- **edge0 统一推理框架开源** —— 发布统一推理框架：**接入层统一**为一套
+  API（chat / serve / 端侧共用），**runtime 自动适配不同硬件平台**——
+  iOS、macOS、Android、Windows、Python；在本仓库已开源的平台引擎
+  （`ios/` · `macos/` · `android/` · `windows/`）之上统一收编。
+- **CUDA 后端**（Python 框架）—— 插槽已预留在
+  `python/src/edge0/backends/cuda/`，核心代码零改动。
+
+**模型与算法**
+
+Q4 推进两条主线：把下一代架构引入框架，以及把潜在推理（latent
+thinking）变成真正的延迟收益、而不只是算力节省。
+
+- **下一代架构支持（Qwen3.8-Flash 级）** —— 在 edge0 上跑通混合线性
+  注意力（GDN + QSA）、门控多分支残差、N-gram embedding 拓扑。这些
+  设计天然契合 SSD 流式 offload：O(1) 状态的注意力层让长思考不再是
+  KV cache 问题，纯查表的 N-gram 表按需流式加载。目标：该档位在单
+  设备跑通，且基准与 fp16 基座的差距在可接受范围内。
+- **潜在思考 + 批量专家预预测** —— 让潜在推理成为*延迟*收益，而不
+  只是算力收益。核心工程问题：把专家路由从 per-position 改为**每 block
+  一次**，一次预测覆盖一个 block 内所有位置与所有轮次，使**专家加载量
+  与推理循环次数解耦**；配合跨 block 预取，在当前 block 的计算窗口内
+  加载下一个 block 的专家集。进展以**等准确率下的端到端思考阶段耗时**
+  衡量（而非 tokens/s）。
+- 更多模型档位与适配器版本在现有管线上持续发布。
+
+## 贡献
+
+欢迎各种形式的贡献——issue、PR、性能实测报告、新模型接入都算。
 
-## 验证
+**Python 框架**（现已可用）：
 
 ```bash
+cd python
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+
 pytest                 # 单元测试（不含真实权重）
 EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q
-                        # 真实权重生成测试；缺少的档位会明确 skip
+                       # 真实权重生成测试；缺少的档位会明确 skip
 .venv/bin/python scripts/e2e_smoke.py \
   --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
-                        # staged vs exact 一致性 + 生成冒烟
+                       # staged vs exact 一致性 + 生成冒烟
 scripts/generate_example.py   # 完整 API 上手例子（prefill→生成→解码全链路）
-examples/demo.py       # 最小 API walkthrough（edge0 demo 的等价代码）
+examples/demo.py              # 最小 API walkthrough（edge0 demo 的等价代码）
 ```
 
-## 文档
+每个 PR 都会跑 CI：单元测试（macOS + MLX）与仓库卫生检查（无硬编码
+路径、后端边界、无密钥）。
 
-- [架构总览](docs/architecture.md)
-- [注意力抽象](docs/attention.md) / [MoE 抽象](docs/moe.md) / [SSD 流式](docs/streaming.md) / [prerouter](docs/prerouter.md)
-- [如何接入新模型](docs/adding-a-model.md)
-- [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md)
+**平台 runtime**（macOS / iOS / Android / Windows）：各平台目录自带
+构建指南与测试——见各目录 README。
+
+流程：fork → 功能分支 → 向 `main` 发 PR。请保持卫生检查通过，并为新
+行为补测试。
+
+## 引用
+
+如果 edge0 对你有帮助，请引用我们的技术报告：
+
+```bibtex
+@article{lin2026other,
+  title   = {The Other Half of the Memory Wall: Serving 35B MoEs from SSD
+             with Trained Routing Prediction},
+  author  = {Lin, Yu and Wang, Yiming and Cai, Runyuan and Liu, Hanze and
+             Zeng, Xiaodong},
+  journal = {arXiv preprint arXiv:2609.18063},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2609.18063}
+}
+```
+
+## 联系我们
+
+社区与支持渠道即将上线，本节将列出联系我们的官方方式：
+
+- **邮箱**：samuel@edge0.ai
+
+Bug 与功能建议请直接提
+[GitHub Issues](https://github.com/Edge0-AI/Edge0/issues)。
 
 ## License
 
diff --git a/android/.gitignore b/android/.gitignore
new file mode 100644
index 0000000..6249d89
--- /dev/null
+++ b/android/.gitignore
@@ -0,0 +1,16 @@
+*.iml
+.gradle
+/local.properties
+/.idea
+.DS_Store
+/build
+*/build
+*/.cxx
+/captures
+.externalNativeBuild
+.cxx
+local.properties
+# engine build workbench (staging for libs; artifacts are never committed)
+/build-dl
+# model sources and converted gguf (self-service downloads)
+/models
diff --git a/android/NOTICE b/android/NOTICE
new file mode 100644
index 0000000..e7f3546
--- /dev/null
+++ b/android/NOTICE
@@ -0,0 +1,20 @@
+NOTICE - edge0 / android
+================================
+
+Licensed under Apache-2.0 (see the LICENSE at the monorepo root). Besides our
+own code, this project contains or references the following third-party material:
+
+1. llama.cpp - MIT (commit 7ab4ee7, release b11100, pinned)
+   Location: ../vendor/llama.cpp (pristine upstream checkout, never patched in
+   place). Built with ../patches/llama.cpp/{common,android} replayed (20
+   hook-point bands) into the shipped shared libraries (libllama, libggml*).
+   ggml (Apache-2.0) is developed inside the same upstream tree. Not vendored
+   as files; built from source by the consumer (tools/llama/build_vendor_libs.sh).
+
+2. Gradle/Maven dependencies (Apache-2.0 / MIT / BSD)
+   Kotlin, Jetpack Compose, AndroidX (Activity, Room, DataStore, Lifecycle),
+   kotlinx-serialization, mikepenz multiplatform-markdown-renderer-m3 (MIT),
+   JUnit. Resolved at build time; each ships its own license text.
+
+3. Model assets - Edge0/Edge0-8B-A1B-preview and Edge0/Edge0-35B-A3B-preview on
+   Hugging Face; terms per the model cards. Not part of the source distribution.
diff --git a/android/README.md b/android/README.md
new file mode 100644
index 0000000..ee9197c
--- /dev/null
+++ b/android/README.md
@@ -0,0 +1,185 @@
+# edge0-android — Release Build Guide
+
+English | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+This document is for **developers building or evaluating the Android app from source**. It covers: quick start (engine → app → models → test), measured performance on the reference device, and the technical choices behind the stack.
+
+edge0 is published as a **monorepo** — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — whose top level holds the shared engine supply (`vendor.llama.pin` + the scripts-materialized `vendor/llama.cpp`, `patches/llama.cpp/` band-sets) and the platform subprojects (`windows/` = the desktop companion, `android/` = this app). Inference runs on pinned upstream llama.cpp, patched as a replayable patch set, fully on-CPU with ARM-NEON kernels and a demand-paged expert pool.
+
+Third-party attribution: see `NOTICE` in this directory. Model cards: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+
+---
+
+## 1. Quick Start
+
+### 1.1 Prerequisites
+
+| component | requirement |
+|---|---|
+| Device | arm64-v8a, Android 13+ (API 33); Snapdragon 8 Elite class recommended |
+| RAM | 8 GB+ runs 8B; 12–16 GB runs 35B (expert paging, see §3.2) |
+| Toolchain | JDK 17+, Android SDK 35, **NDK r28** (`28.2.13676358`), CMake ≥ 3.21 + Ninja |
+| Python | 3.10+ with `numpy` (model converter, runs from the sibling `windows/tools`) |
+| Disk | ≥ 30 GB free for model sources and converted GGUF builds |
+
+The NDK is **not** part of this repository — install it once (Android Studio:
+*Android SDK → SDK Tools → NDK (Side by side)*, or via CLI):
+
+```bash
+sdkmanager --install "ndk;28.2.13676358"
+```
+
+Then point the build at it with `NDK_DIR` (or an exported `ANDROID_NDK_HOME`):
+`build_vendor_libs.sh` reads the compiler toolchain **and** the `libomp.so`
+runtime it stages (§1.2) from inside that NDK installation; a missing NDK
+fails fast with that hint rather than producing a broken library set.
+
+### 1.2 Build the engine
+
+The native libraries come from the pinned upstream tree with this platform's
+patch bands replayed into an isolated worktree — the vendor tree is **never
+patched in place**:
+
+```bash
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/android
+bash tools/llama/build_vendor_libs.sh
+```
+
+The script replays `../patches/llama.cpp/{common,android}` (6 + 14 bands) onto
+the pinned llama.cpp tree — the pin lives in `../vendor.llama.pin` (currently
+`7ab4ee7`, tag b11100); the tree is **not** a submodule, so on first run it is
+cloned from upstream into `../vendor/llama.cpp` (gitignored; set
+`EDGE0_LLAMA_URL` to use a mirror) and detached at the pin. Replays happen in a
+gitignored consumer worktree, the script
+asserts the golden result-tree hash, and produces the four engine shared
+libraries (plus the NDK `libomp.so` runtime that `libggml-cpu` needs) with
+headers into `build-dl/llama-libs/` (the app's jniLibs staging point).
+`--replay` re-applies the bands after patch changes; mismatched tree ⇒ RED, the
+build refuses to start. This step is required once before the app build —
+the Gradle plugin reads these libraries from `build-dl/llama-libs/`.
+
+### 1.3 Build & install the app
+
+```bash
+./gradlew :app:assembleDebug
+./gradlew :app:installDebug        # or adb install -r app/build/outputs/apk/debug/app-debug.apk
+```
+
+### 1.4 Models
+
+GGUF builds are produced locally from the published checkpoints — everything
+stays under this directory's `models/` (gitignored):
+
+```bash
+huggingface-cli download Edge0/Edge0-8B-A1B-preview --local-dir models/edge0-8b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-8b
+#  → models/edge0-8b-gguf/{edge0-8b.gguf, lora_edge0_8b-gguf.gguf, manifest.json}
+
+huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-35b
+
+bash tools/model/push_models.sh --all    # md5-gated staging onto the device
+```
+
+The converter needs only python3 + numpy (no MLX/torch runtime — "MLX" denotes
+the on-disk checkpoint layout). It runs the r3 repack with numeric parity gates
+and emits a sha256 manifest; a correct conversion reproduces the baseline
+checksums listed in `push_models.sh` byte-for-byte. Models can also be copied
+into `files/models/` with the in-app picker.
+
+### 1.5 Run
+
+Launch **Edge0 Chat**. The title bar shows the active model (8B / 35B), the
+top-right button switches. The composer is send-only; temperature, thinking and
+the system prompt live in the drawer settings. Each reply carries an inline
+metrics line: `tokens · TTFT · prefill t/s · decode t/s · RSS`.
+
+### 1.6 Test it
+
+Instrumented regression (needs both models staged — reinstalling the test APK
+wipes app data, so restage right before the run):
+
+```bash
+./gradlew :app:installDebugAndroidTest
+bash tools/model/push_models.sh --all
+adb shell am instrument -w -e class dev.edge0.runtime.app.LlamaRuntimeTest \
+  dev.edge0.runtime.app.test/androidx.test.runner.AndroidJUnitRunner
+# expected: OK (8 tests), ~7 min on the reference device
+```
+
+Coverage: 8B/35B smoke, 35B↔8B in-process switching, prefix-reuse fidelity,
+thinking on/off × system-prompt quadrants (8B gate + 35B off leak probe), and
+identity retention across multi-turn rendering. Host-side logic tests:
+`./gradlew :app:testDebugUnitTest` (26 tests, no device needed).
+
+---
+
+## 2. Performance
+
+Reference device: **Lenovo TB322FC (Snapdragon 8 Elite, 16 GB RAM)**, shipping
+config, sustained windows (first segments = boost clocks, tail = thermal steady
+state — both reported rather than cherry-picked peaks).
+
+| Model | TTFT (warm turn) | Decode | Prefill | Session RSS |
+|---|---:|---:|---:|---:|
+| **8B** (Q8-class GGUF + LoRA) | ≈ 1.4 s | 29–32 → ~10 t/s over a 480 s window | ~100 t/s | ≈ 250 MB |
+| **35B** (mixed int8 MoE, demand-paged) | ≈ 1.1 s warm (≈ 10–15 s on the very first turn of a new topic — cold expert pool, flash-bound) | 6–9 t/s in-app (9.46 t/s sustained CLI) | ~1 s per incremental turn | pool budget 2–6 GB; resident ≪ file size |
+
+Notes: "first turn of a new topic" pays the cold-pool tax — e.g. a 55-token
+question measured prefill 12.8 s with 25956 expert loads, 73 % of wall time in
+flash I/O wait; that is demand paging working as designed, not a regression.
+Turn two onwards is warm: KV prefix reuse + keepwarm refill bring TTFT to ~1 s.
+Decode decay along a long window is DVFS/thermal behavior on this SoC.
+
+---
+
+## 3. Technical Details
+
+### 3.1 Architecture
+
+```mermaid
+graph TD
+    subgraph App ["Kotlin / Jetpack Compose"]
+        UI[ChatScreen · dark · send-only composer] --> VM[ChatViewModel]
+        VM --> RT[LlamaRuntime<br/>coroutines + Flow events]
+        VM --> DB[(Room · threads & messages)]
+        ST[SettingsStore] --> VM
+    end
+    subgraph Native ["C JNI shell (llama_chat.c)"]
+        SHELL[generate loop · template-aware thinking control<br/>segment-wise history render · UTF-8-safe streaming]
+    end
+    subgraph Engine ["patched llama.cpp @ b11100 · arm64 CPU-only"]
+        LLIB[libllama.so]
+        GCPU[libggml-cpu.so<br/>NEON kernels + moe_pool]
+    end
+    RT -->|JNI| SHELL --> LLIB --> GCPU
+    GCPU -->|demand-paged expert IO| MODELS[GGUF on flash]
+```
+
+### 3.2 Why a 21.7 GB MoE fits on a phone — `moe_pool`
+
+The 35B model activates only a moving subset of its 256-per-layer experts per
+token, so the shipped design pages experts **on demand** from flash instead of
+resident memory (`ggml/src/ggml-cpu/moe_pool.c`, developed as the 14-patch
+android band): copy-in private frames, an equal-slot state machine, background
+IO staging queues tuned to measured UFS throughput, pin/blob/trim controls, a
+turn-end keepwarm refill within a byte budget, and full cross-model reset
+enabling 8B↔35B switching inside one process. With the pool disabled the engine
+resolves expert rows exactly like upstream (NULL resolver ⇒ zero perturbation,
+verified by symbol-set diffing). This is the same mechanism family the desktop
+project implements over NVMe + Vulkan; here it is CPU/NEON by design — GPU
+backends stay out of the shipping config for deterministic numerics and a
+single memory model.
+
+### 3.3 Repo map
+
+```
+app/                     Android app: Compose UI (src/main/java), JNI shell (src/main/cpp),
+                         instrumented + unit tests (src/androidTest, src/test)
+tools/llama/             build_vendor_libs.sh — engine rebuild from the pinned tree + bands
+tools/model/             push_models.sh — md5-gated model staging to devices
+../vendor/llama.cpp/     materialized by the build scripts from vendor.llama.pin (gitignored; never patched in place)
+../patches/llama.cpp/    common(6) + android(14) hook-point bands + ledger README
+../windows/              desktop companion — hosts the MLX→GGUF converter used in §1.4
+```
diff --git a/android/README_es.md b/android/README_es.md
new file mode 100644
index 0000000..ef65486
--- /dev/null
+++ b/android/README_es.md
@@ -0,0 +1,191 @@
+# edge0-android — Guía de compilación de release
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | Español | [Français](README_fr.md)
+
+Este documento está dirigido a **desarrolladores que compilan o evalúan la aplicación de Android desde el código fuente**. Cubre: inicio rápido (motor → aplicación → modelos → pruebas), rendimiento medido en el dispositivo de referencia y las decisiones técnicas detrás del stack.
+
+edge0 se publica como **monorepo** — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — cuyo nivel superior contiene el suministro compartido del motor (`vendor.llama.pin` + el `vendor/llama.cpp` materializado por los scripts, los conjuntos de bandas de `patches/llama.cpp/`) y los subproyectos de plataforma (`windows/` = el complemento de escritorio, `android/` = esta aplicación). La inferencia se ejecuta sobre una versión fijada (pin) de llama.cpp upstream, parcheada como un conjunto de parches reejecutable, totalmente en CPU con kernels ARM-NEON y un pool de expertos paginado bajo demanda.
+
+Atribución de terceros: consulta `NOTICE` en este directorio. Fichas de modelo: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+
+---
+
+## 1. Inicio rápido
+
+### 1.1 Requisitos previos
+
+| componente | requisito |
+|---|---|
+| Dispositivo | arm64-v8a, Android 13+ (API 33); se recomienda un Snapdragon 8 Elite o similar |
+| RAM | 8 GB+ ejecuta 8B; 12–16 GB ejecuta 35B (paginación de expertos, ver §3.2) |
+| Toolchain | JDK 17+, Android SDK 35, **NDK r28** (`28.2.13676358`), CMake ≥ 3.21 + Ninja |
+| Python | 3.10+ con `numpy` (conversor de modelos, se ejecuta desde el `windows/tools` hermano) |
+| Disco | ≥ 30 GB libres para las fuentes de los modelos y las compilaciones GGUF convertidas |
+
+El NDK **no** forma parte de este repositorio — instálalo una vez (Android Studio:
+*Android SDK → SDK Tools → NDK (Side by side)*, o mediante CLI):
+
+```bash
+sdkmanager --install "ndk;28.2.13676358"
+```
+
+Después apunta la compilación hacia él con `NDK_DIR` (o una variable `ANDROID_NDK_HOME`
+exportada): `build_vendor_libs.sh` lee la toolchain del compilador **y** el runtime
+`libomp.so` que prepara (§1.2) desde dentro de esa instalación del NDK; si falta el NDK,
+falla rápidamente con esa indicación en lugar de producir un conjunto de bibliotecas
+defectuoso.
+
+### 1.2 Compilar el motor
+
+Las bibliotecas nativas provienen del árbol upstream fijado con las bandas de parches de
+esta plataforma reejecutadas en un worktree aislado — el árbol del vendor **nunca** se
+parchea in situ:
+
+```bash
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/android
+bash tools/llama/build_vendor_libs.sh
+```
+
+El script reejecuta `../patches/llama.cpp/{common,android}` (6 + 14 bandas) sobre el árbol
+fijado de llama.cpp — el pin está en `../vendor.llama.pin` (actualmente `7ab4ee7`, tag
+b11100); el árbol **no** es un submódulo, así que en la primera ejecución se clona desde
+upstream en `../vendor/llama.cpp` (ignorado por git; establece `EDGE0_LLAMA_URL` para usar
+un mirror) y se hace detach en el pin. Las reejecuciones ocurren en un worktree consumidor
+ignorado por git, el script verifica el hash golden del árbol de resultado y produce las
+cuatro bibliotecas compartidas del motor (más el runtime `libomp.so` del NDK que
+`libggml-cpu` necesita) con las cabeceras en `build-dl/llama-libs/` (el punto de
+preparación de jniLibs de la aplicación). `--replay` vuelve a aplicar las bandas tras
+cambios en los parches; árbol discordante ⇒ ROJO, la compilación se niega a empezar. Este
+paso es necesario una vez antes de compilar la aplicación — el plugin de Gradle lee estas
+bibliotecas desde `build-dl/llama-libs/`.
+
+### 1.3 Compilar e instalar la aplicación
+
+```bash
+./gradlew :app:assembleDebug
+./gradlew :app:installDebug        # or adb install -r app/build/outputs/apk/debug/app-debug.apk
+```
+
+### 1.4 Modelos
+
+Las compilaciones GGUF se producen localmente a partir de los checkpoints publicados —
+todo permanece en el `models/` de este directorio (ignorado por git):
+
+```bash
+huggingface-cli download Edge0/Edge0-8B-A1B-preview --local-dir models/edge0-8b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-8b
+#  → models/edge0-8b-gguf/{edge0-8b.gguf, lora_edge0_8b-gguf.gguf, manifest.json}
+
+huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-35b
+
+bash tools/model/push_models.sh --all    # md5-gated staging onto the device
+```
+
+El conversor solo necesita python3 + numpy (sin runtime MLX/torch — «MLX» denota la
+disposición en disco del checkpoint). Ejecuta el reempaquetado r3 con puertas de paridad
+numérica y emite un manifiesto sha256; una conversión correcta reproduce byte a byte las
+sumas de verificación de referencia listadas en `push_models.sh`. Los modelos también se
+pueden copiar a `files/models/` con el selector de la aplicación.
+
+### 1.5 Ejecución
+
+Inicia **Edge0 Chat**. La barra de título muestra el modelo activo (8B / 35B), el botón
+superior derecho lo cambia. El compositor solo envía; la temperatura, el thinking y el
+system prompt están en los ajustes del panel lateral. Cada respuesta lleva una línea de
+métricas integrada: `tokens · TTFT · prefill t/s · decode t/s · RSS`.
+
+### 1.6 Probarlo
+
+Regresión instrumentada (requiere ambos modelos preparados — reinstalar el APK de prueba
+borra los datos de la aplicación, así que vuelve a prepararlos justo antes de la
+ejecución):
+
+```bash
+./gradlew :app:installDebugAndroidTest
+bash tools/model/push_models.sh --all
+adb shell am instrument -w -e class dev.edge0.runtime.app.LlamaRuntimeTest \
+  dev.edge0.runtime.app.test/androidx.test.runner.AndroidJUnitRunner
+# expected: OK (8 tests), ~7 min on the reference device
+```
+
+Cobertura: smoke de 8B/35B, cambio en el proceso 35B↔8B, fidelidad de la reutilización de
+prefijos, cuadrantes de thinking activado/desactivado × system prompt (puerta 8B + prueba
+de fuga con 35B desactivado) y retención de identidad en el renderizado multiturno.
+Pruebas de lógica del lado del host: `./gradlew :app:testDebugUnitTest` (26 pruebas, sin
+necesidad de dispositivo).
+
+---
+
+## 2. Rendimiento
+
+Dispositivo de referencia: **Lenovo TB322FC (Snapdragon 8 Elite, 16 GB de RAM)**,
+configuración de serie, ventanas sostenidas (los primeros segmentos = relojes de boost, la
+cola = estado estacionario térmico — se reportan ambos en lugar de picos seleccionados a
+dedo).
+
+| Modelo | TTFT (turno caliente) | Decode | Prefill | RSS de sesión |
+|---|---:|---:|---:|---:|
+| **8B** (GGUF clase Q8 + LoRA) | ≈ 1.4 s | 29–32 → ~10 t/s en una ventana de 480 s | ~100 t/s | ≈ 250 MB |
+| **35B** (MoE int8 mixto, paginado bajo demanda) | ≈ 1.1 s caliente (≈ 10–15 s en el primer turno de un tema nuevo — pool de expertos frío, limitado por flash) | 6–9 t/s en la aplicación (9.46 t/s sostenidos en CLI) | ~1 s por turno incremental | presupuesto del pool 2–6 GB; residente ≪ tamaño del archivo |
+
+Notas: el «primer turno de un tema nuevo» paga el costo del pool frío — p. ej., una
+pregunta de 55 tokens midió un prefill de 12.8 s con 25956 cargas de expertos, el 73 % del
+tiempo real en espera de E/S de flash; eso es la paginación bajo demanda funcionando según
+lo diseñado, no una regresión. Desde el segundo turno todo está caliente: la reutilización
+de prefijos KV + el refill keepwarm llevan el TTFT a ~1 s. El descenso del decode a lo
+largo de una ventana larga es comportamiento DVFS/térmico en este SoC.
+
+---
+
+## 3. Detalles técnicos
+
+### 3.1 Arquitectura
+
+```mermaid
+graph TD
+    subgraph App ["Kotlin / Jetpack Compose"]
+        UI[ChatScreen · dark · send-only composer] --> VM[ChatViewModel]
+        VM --> RT[LlamaRuntime<br/>coroutines + Flow events]
+        VM --> DB[(Room · threads & messages)]
+        ST[SettingsStore] --> VM
+    end
+    subgraph Native ["C JNI shell (llama_chat.c)"]
+        SHELL[generate loop · template-aware thinking control<br/>segment-wise history render · UTF-8-safe streaming]
+    end
+    subgraph Engine ["patched llama.cpp @ b11100 · arm64 CPU-only"]
+        LLIB[libllama.so]
+        GCPU[libggml-cpu.so<br/>NEON kernels + moe_pool]
+    end
+    RT -->|JNI| SHELL --> LLIB --> GCPU
+    GCPU -->|demand-paged expert IO| MODELS[GGUF on flash]
+```
+
+### 3.2 Por qué un MoE de 21.7 GB cabe en un teléfono — `moe_pool`
+
+El modelo de 35B solo activa un subconjunto móvil de sus 256 expertos por capa en cada
+token, así que el diseño de serie pagina los expertos **bajo demanda** desde la flash en
+lugar de la memoria residente (`ggml/src/ggml-cpu/moe_pool.c`, desarrollado como la banda
+android de 14 parches): frames privados copy-in, una máquina de estados de slots iguales,
+colas de preparación de E/S en segundo plano ajustadas al throughput UFS medido, controles
+de pin/blob/trim, un refill keepwarm al final del turno dentro de un presupuesto de bytes y
+un reinicio completo entre modelos que permite cambiar de 8B↔35B dentro de un mismo
+proceso. Con el pool desactivado, el motor resuelve las filas de expertos exactamente igual
+que upstream (resolver NULL ⇒ perturbación cero, verificado con diff del conjunto de
+símbolos). Es la misma familia de mecanismos que el proyecto de escritorio implementa sobre
+NVMe + Vulkan; aquí es CPU/NEON por diseño — los backends de GPU quedan fuera de la
+configuración de serie para garantizar una numérica determinista y un único modelo de
+memoria.
+
+### 3.3 Mapa del repositorio
+
+```
+app/                     Android app: Compose UI (src/main/java), JNI shell (src/main/cpp),
+                         instrumented + unit tests (src/androidTest, src/test)
+tools/llama/             build_vendor_libs.sh — engine rebuild from the pinned tree + bands
+tools/model/             push_models.sh — md5-gated model staging to devices
+../vendor/llama.cpp/     materialized by the build scripts from vendor.llama.pin (gitignored; never patched in place)
+../patches/llama.cpp/    common(6) + android(14) hook-point bands + ledger README
+../windows/              desktop companion — hosts the MLX→GGUF converter used in §1.4
+```
diff --git a/android/README_fr.md b/android/README_fr.md
new file mode 100644
index 0000000..5d0a988
--- /dev/null
+++ b/android/README_fr.md
@@ -0,0 +1,200 @@
+# edge0-android — Guide de compilation de la release
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | Français
+
+Ce document s'adresse aux **développeurs qui compilent ou évaluent l'application Android depuis les sources**. Il couvre : le démarrage rapide (moteur → application → modèles → test), les performances mesurées sur l'appareil de référence et les choix techniques derrière la stack.
+
+edge0 est publié sous forme de **monorepo** — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — dont le niveau supérieur contient l'approvisionnement partagé du moteur (`vendor.llama.pin` + le `vendor/llama.cpp` matérialisé par les scripts, les jeux de bandes `patches/llama.cpp/`) et les sous-projets par plateforme (`windows/` = le compagnon de bureau, `android/` = cette application). L'inférence s'exécute sur une version épinglée de llama.cpp upstream, patchée sous forme d'un jeu de patchs rejouable, entièrement sur CPU avec des kernels ARM-NEON et un pool d'experts paginé à la demande.
+
+Attribution tierce : voir `NOTICE` dans ce répertoire. Cartes de modèles : [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+
+---
+
+## 1. Démarrage rapide
+
+### 1.1 Prérequis
+
+| composant | exigence |
+|---|---|
+| Appareil | arm64-v8a, Android 13+ (API 33) ; Snapdragon 8 Elite ou équivalent recommandé |
+| RAM | 8 GB+ pour le 8B ; 12–16 GB pour le 35B (pagination des experts, voir §3.2) |
+| Chaîne d'outils | JDK 17+, Android SDK 35, **NDK r28** (`28.2.13676358`), CMake ≥ 3.21 + Ninja |
+| Python | 3.10+ avec `numpy` (convertisseur de modèles, exécuté depuis le `windows/tools` voisin) |
+| Disque | ≥ 30 GB libres pour les sources des modèles et les compilations GGUF converties |
+
+Le NDK ne fait **pas** partie de ce dépôt — installez-le une fois (Android Studio :
+*Android SDK → SDK Tools → NDK (Side by side)*, ou via la CLI) :
+
+```bash
+sdkmanager --install "ndk;28.2.13676358"
+```
+
+Pointez ensuite la compilation dessus avec `NDK_DIR` (ou un `ANDROID_NDK_HOME`
+exporté) : `build_vendor_libs.sh` lit la chaîne de compilation **et** le runtime
+`libomp.so` qu'il met en place (§1.2) depuis cette installation NDK ; un NDK
+manquant échoue immédiatement avec cette indication plutôt que de produire un
+jeu de bibliothèques cassé.
+
+### 1.2 Compiler le moteur
+
+Les bibliothèques natives proviennent de l'arborescence upstream épinglée, avec
+les bandes de patchs de cette plateforme rejouées dans un worktree isolé —
+l'arborescence vendor n'est **jamais patchée en place** :
+
+```bash
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/android
+bash tools/llama/build_vendor_libs.sh
+```
+
+Le script rejoue `../patches/llama.cpp/{common,android}` (6 + 14 bandes) sur
+l'arborescence llama.cpp épinglée — l'épingle se trouve dans
+`../vendor.llama.pin` (actuellement `7ab4ee7`, tag b11100) ; l'arborescence
+n'est **pas** un sous-module, elle est donc clonée depuis upstream au premier
+lancement dans `../vendor/llama.cpp` (ignoré par git ; définissez
+`EDGE0_LLAMA_URL` pour utiliser un miroir) et détachée sur l'épingle. Les
+rejeux se produisent dans un worktree consommateur ignoré par git, le script
+vérifie le hash de l'arborescence résultat attendue (golden) et produit les
+quatre bibliothèques partagées du moteur (plus le runtime NDK `libomp.so` dont
+`libggml-cpu` a besoin) avec les en-têtes dans `build-dl/llama-libs/` (le point
+de mise en place jniLibs de l'application). `--replay` réapplique les bandes
+après modification des patchs ; arborescence différente ⇒ ROUGE, la
+compilation refuse de démarrer. Cette étape est requise une fois avant la
+compilation de l'application — le plugin Gradle lit ces bibliothèques depuis
+`build-dl/llama-libs/`.
+
+### 1.3 Compiler et installer l'application
+
+```bash
+./gradlew :app:assembleDebug
+./gradlew :app:installDebug        # or adb install -r app/build/outputs/apk/debug/app-debug.apk
+```
+
+### 1.4 Modèles
+
+Les compilations GGUF sont produites localement à partir des checkpoints
+publiés — tout reste sous le `models/` de ce répertoire (ignoré par git) :
+
+```bash
+huggingface-cli download Edge0/Edge0-8B-A1B-preview --local-dir models/edge0-8b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-8b
+#  → models/edge0-8b-gguf/{edge0-8b.gguf, lora_edge0_8b-gguf.gguf, manifest.json}
+
+huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-35b
+
+bash tools/model/push_models.sh --all    # md5-gated staging onto the device
+```
+
+Le convertisseur n'a besoin que de python3 + numpy (pas de runtime MLX/torch —
+« MLX » désigne la disposition sur disque du checkpoint). Il exécute le repack
+r3 avec des portes de parité numérique et émet un manifeste sha256 ; une
+conversion correcte reproduit octet pour octet les checksums de référence
+listées dans `push_models.sh`. Les modèles peuvent aussi être copiés dans
+`files/models/` avec le sélecteur intégré à l'application.
+
+### 1.5 Exécution
+
+Lancez **Edge0 Chat**. La barre de titre affiche le modèle actif (8B / 35B), le
+bouton en haut à droite permet d'en changer. La zone de saisie sert uniquement
+à envoyer ; la température, la réflexion (thinking) et le prompt système se
+trouvent dans les paramètres du tiroir. Chaque réponse comporte une ligne de
+métriques inline : `tokens · TTFT · prefill t/s · decode t/s · RSS`.
+
+### 1.6 Tester
+
+Régression instrumentée (nécessite les deux modèles mis en place —
+réinstaller l'APK de test efface les données de l'application, replacez donc
+les modèles juste avant l'exécution) :
+
+```bash
+./gradlew :app:installDebugAndroidTest
+bash tools/model/push_models.sh --all
+adb shell am instrument -w -e class dev.edge0.runtime.app.LlamaRuntimeTest \
+  dev.edge0.runtime.app.test/androidx.test.runner.AndroidJUnitRunner
+# expected: OK (8 tests), ~7 min on the reference device
+```
+
+Couverture : smoke 8B/35B, commutation 35B↔8B dans le même processus,
+fidélité de la réutilisation de préfixe, quadrants thinking activé/désactivé ×
+prompt système (porte 8B + sonde de fuite 35B désactivée), et rétention
+d'identité à travers le rendu multi-tours. Tests logiques côté hôte :
+`./gradlew :app:testDebugUnitTest` (26 tests, aucun appareil requis).
+
+---
+
+## 2. Performances
+
+Appareil de référence : **Lenovo TB322FC (Snapdragon 8 Elite, 16 GB RAM)**,
+configuration de série, fenêtres soutenues (premiers segments = fréquences
+boost, fin = régime thermique stable — les deux sont rapportés plutôt que des
+pics triés sur le volet).
+
+| Modèle | TTFT (tour à chaud) | Decode | Prefill | RSS de session |
+|---|---:|---:|---:|---:|
+| **8B** (GGUF classe Q8 + LoRA) | ≈ 1.4 s | 29–32 → ~10 t/s sur une fenêtre de 480 s | ~100 t/s | ≈ 250 MB |
+| **35B** (MoE int8 mixte, paginé à la demande) | ≈ 1.1 s à chaud (≈ 10–15 s au tout premier tour d'un nouveau sujet — pool d'experts froid, limité par le flash) | 6–9 t/s dans l'application (9.46 t/s soutenu en CLI) | ~1 s par tour incrémental | budget du pool 2–6 GB ; résident ≪ taille du fichier |
+
+Notes : le « premier tour d'un nouveau sujet » paie le coût du pool froid —
+par exemple, une question de 55 tokens a mesuré un prefill de 12,8 s avec
+25 956 chargements d'experts, 73 % du temps total passé en attente d'E/S
+flash ; c'est la pagination à la demande qui fonctionne comme prévu, pas une
+régression. À partir du deuxième tour, tout est à chaud : la réutilisation du
+préfixe KV + le remplissage keepwarm ramènent le TTFT à ~1 s. La décroissance
+du decode le long d'une longue fenêtre est le comportement DVFS/thermique sur
+ce SoC.
+
+---
+
+## 3. Détails techniques
+
+### 3.1 Architecture
+
+```mermaid
+graph TD
+    subgraph App ["Kotlin / Jetpack Compose"]
+        UI[ChatScreen · dark · send-only composer] --> VM[ChatViewModel]
+        VM --> RT[LlamaRuntime<br/>coroutines + Flow events]
+        VM --> DB[(Room · threads & messages)]
+        ST[SettingsStore] --> VM
+    end
+    subgraph Native ["C JNI shell (llama_chat.c)"]
+        SHELL[generate loop · template-aware thinking control<br/>segment-wise history render · UTF-8-safe streaming]
+    end
+    subgraph Engine ["patched llama.cpp @ b11100 · arm64 CPU-only"]
+        LLIB[libllama.so]
+        GCPU[libggml-cpu.so<br/>NEON kernels + moe_pool]
+    end
+    RT -->|JNI| SHELL --> LLIB --> GCPU
+    GCPU -->|demand-paged expert IO| MODELS[GGUF on flash]
+```
+
+### 3.2 Pourquoi un MoE de 21,7 GB tient sur un téléphone — `moe_pool`
+
+Le modèle 35B n'active qu'un sous-ensemble mobile de ses 256 experts par
+couche et par token, de sorte que la conception livrée pagine les experts **à
+la demande** depuis le flash au lieu de la mémoire résidente
+(`ggml/src/ggml-cpu/moe_pool.c`, développé comme la bande android de
+14 patchs) : cadres privés recopiés à l'entrée (copy-in), une machine à états
+à créneaux égaux, des files d'attente d'E/S en arrière-plan calibrées sur le
+débit UFS mesuré, des contrôles pin/blob/trim, un remplissage keepwarm de fin
+de tour dans une limite d'octets, et une réinitialisation complète
+inter-modèles permettant la commutation 8B↔35B à l'intérieur d'un même
+processus. Pool désactivé, le moteur résout les lignes d'experts exactement
+comme upstream (résolveur NULL ⇒ perturbation nulle, vérifiée par diff des
+jeux de symboles). C'est la même famille de mécanismes que le projet de
+bureau implémente sur NVMe + Vulkan ; ici, c'est CPU/NEON par conception —
+les backends GPU restent hors de la configuration livrée pour garantir un
+comportement numérique déterministe et un modèle de mémoire unique.
+
+### 3.3 Structure du dépôt
+
+```
+app/                     Android app: Compose UI (src/main/java), JNI shell (src/main/cpp),
+                         instrumented + unit tests (src/androidTest, src/test)
+tools/llama/             build_vendor_libs.sh — engine rebuild from the pinned tree + bands
+tools/model/             push_models.sh — md5-gated model staging to devices
+../vendor/llama.cpp/     materialized by the build scripts from vendor.llama.pin (gitignored; never patched in place)
+../patches/llama.cpp/    common(6) + android(14) hook-point bands + ledger README
+../windows/              desktop companion — hosts the MLX→GGUF converter used in §1.4
+```
diff --git a/android/README_ja.md b/android/README_ja.md
new file mode 100644
index 0000000..5f0c93a
--- /dev/null
+++ b/android/README_ja.md
@@ -0,0 +1,140 @@
+# edge0-android — リリースビルドガイド
+
+[English](README.md) | [中文](README_zh.md) | 日本語 | [Español](README_es.md) | [Français](README_fr.md)
+
+このドキュメントは、**ソースから Android アプリをビルドまたは評価する開発者**向けのものです。内容は、クイックスタート(エンジン → アプリ → モデル → テスト)、リファレンスデバイスでの実測パフォーマンス、そしてスタックの背後にある技術的選択です。
+
+edge0 は**モノレポ**として公開されています — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — 最上位階層には共有のエンジンサプライ(`vendor.llama.pin` + スクリプトが実体化する `vendor/llama.cpp`、`patches/llama.cpp/` のバンドセット)とプラットフォームのサブプロジェクト(`windows/` = デスクトップ版コンパニオン、`android/` = 本アプリ)が含まれます。推論はピン留めされた上流の llama.cpp 上で実行され、リプレイ可能なパッチセットとしてパッチが適用され、ARM-NEON カーネルとデマンドページ方式のエキスパートプールにより、完全に CPU 上で動作します。
+
+サードパーティの帰属表示: このディレクトリの `NOTICE` を参照してください。モデルカード: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)。
+
+---
+
+## 1. クイックスタート
+
+### 1.1 前提条件
+
+| コンポーネント | 要件 |
+|---|---|
+| デバイス | arm64-v8a、Android 13+(API 33);Snapdragon 8 Elite クラスを推奨 |
+| RAM | 8 GB+ で 8B が動作;12–16 GB で 35B が動作(エキスパートページング、§3.2 を参照) |
+| ツールチェーン | JDK 17+、Android SDK 35、**NDK r28**(`28.2.13676358`)、CMake ≥ 3.21 + Ninja |
+| Python | 3.10+ と `numpy`(モデルコンバータ;隣接する `windows/tools` から実行される) |
+| ディスク | モデルソースと変換済み GGUF ビルド用に 30 GB 以上の空き容量 |
+
+NDK はこのリポジトリの一部では**ありません** — 一度インストールしてください(Android Studio: *Android SDK → SDK Tools → NDK (Side by side)*、または CLI 経由):
+
+```bash
+sdkmanager --install "ndk;28.2.13676358"
+```
+
+次に `NDK_DIR`(またはエクスポートされた `ANDROID_NDK_HOME`)でビルドにそれを指定します:
+`build_vendor_libs.sh` は、その NDK インストール内からコンパイラツールチェーン**と**、ステージングする `libomp.so` ランタイム(§1.2)を読み取ります。NDK が見つからない場合は、壊れたライブラリセットを生成するのではなく、そのヒントを示して即座に失敗します。
+
+### 1.2 エンジンのビルド
+
+ネイティブライブラリは、ピン留めされた上流ツリーにこのプラットフォームのパッチバンドを隔離されたワークツリーへリプレイして生成されます — ベンダーツリーが**その場でパッチされることは決してありません**:
+
+```bash
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/android
+bash tools/llama/build_vendor_libs.sh
+```
+
+スクリプトは `../patches/llama.cpp/{common,android}`(6 + 14 バンド)をピン留めされた llama.cpp ツリーにリプレイします — ピンは `../vendor.llama.pin` にあります(現在 `7ab4ee7`、tag b11100)。このツリーはサブモジュールでは**ない**ため、初回実行時に上流から `../vendor/llama.cpp` にクローンされ(gitignore 済み;ミラーを使うには `EDGE0_LLAMA_URL` を設定)、ピンの位置で detach されます。リプレイは gitignore されたコンシューマーワークツリー内で行われ、スクリプトはゴールデンな結果ツリーハッシュをアサートし、4 つのエンジン共有ライブラリ(および `libggml-cpu` が必要とする NDK の `libomp.so` ランタイム)をヘッダとともに `build-dl/llama-libs/`(アプリの jniLibs ステージングポイント)に生成します。`--replay` はパッチ変更後にバンドを再適用します。ツリーが不一致の場合は RED となり、ビルドは開始を拒否します。このステップはアプリのビルドの前に一度必要です — Gradle プラグインはこれらのライブラリを `build-dl/llama-libs/` から読み取ります。
+
+### 1.3 アプリのビルド & インストール
+
+```bash
+./gradlew :app:assembleDebug
+./gradlew :app:installDebug        # or adb install -r app/build/outputs/apk/debug/app-debug.apk
+```
+
+### 1.4 モデル
+
+GGUF ビルドは公開チェックポイントからローカルで生成されます — すべてのものはこのディレクトリの `models/` 配下に置かれます(gitignore 済み):
+
+```bash
+huggingface-cli download Edge0/Edge0-8B-A1B-preview --local-dir models/edge0-8b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-8b
+#  → models/edge0-8b-gguf/{edge0-8b.gguf, lora_edge0_8b-gguf.gguf, manifest.json}
+
+huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-35b
+
+bash tools/model/push_models.sh --all    # md5-gated staging onto the device
+```
+
+コンバータが必要とするのは python3 + numpy のみです(MLX/torch ランタイムは不要 — 「MLX」はディスク上のチェックポイントレイアウトを指します)。数値パリティゲートを備えた r3 リパックを実行し、sha256 マニフェストを出力します。正しい変換は、`push_models.sh` に記載されたベースラインチェックサムをバイト単位で再現します。モデルはアプリ内ピッカーで `files/models/` にコピーすることもできます。
+
+### 1.5 実行
+
+**Edge0 Chat** を起動します。タイトルバーはアクティブなモデル(8B / 35B)を表示し、右上のボタンで切り替えます。コンポーザーは送信専用です。temperature、thinking、システムプロンプトはドロワーの設定にあります。各返信にはインラインのメトリクス行が付きます: `tokens · TTFT · prefill t/s · decode t/s · RSS`。
+
+### 1.6 テスト
+
+計装リグレッションテスト(両モデルのステージングが必要 — テスト APK の再インストールはアプリデータを消去するため、実行直前に再ステージングしてください):
+
+```bash
+./gradlew :app:installDebugAndroidTest
+bash tools/model/push_models.sh --all
+adb shell am instrument -w -e class dev.edge0.runtime.app.LlamaRuntimeTest \
+  dev.edge0.runtime.app.test/androidx.test.runner.AndroidJUnitRunner
+# expected: OK (8 tests), ~7 min on the reference device
+```
+
+カバレッジ: 8B/35B スモーク、35B↔8B のプロセス内切り替え、プレフィックス再利用の忠実性、thinking オン/オフ × システムプロンプトの四象限(8B ゲート + 35B オフ時のリークプローブ)、マルチターンレンダリングを通じたアイデンティティの保持。ホスト側のロジックテスト: `./gradlew :app:testDebugUnitTest`(26 テスト、デバイス不要)。
+
+---
+
+## 2. パフォーマンス
+
+リファレンスデバイス: **Lenovo TB322FC(Snapdragon 8 Elite、16 GB RAM)**、出荷時構成、持続ウィンドウ(最初のセグメント = ブーストクロック、テール = 熱的定常状態 — チェリーピックしたピーク値ではなく両方を報告)。
+
+| モデル | TTFT(ウォームターン) | デコード | プリフィル | セッション RSS |
+|---|---:|---:|---:|---:|
+| **8B**(Q8 クラス GGUF + LoRA) | ≈ 1.4 s | 480 秒のウィンドウで 29–32 → ~10 t/s | ~100 t/s | ≈ 250 MB |
+| **35B**(混合 int8 MoE、デマンドページ方式) | ≈ 1.1 s ウォーム(≈ 10–15 s 新しいトピックの最初のターン — コールドエキスパートプール、フラッシュ律速) | アプリ内 6–9 t/s(CLI 持続値 9.46 t/s) | インクリメンタルターンあたり ~1 s | プールバジェット 2–6 GB;常駐 ≪ ファイルサイズ |
+
+注記: 「新しいトピックの最初のターン」はコールドプールのコストを支払います — 例えば 55 トークンの質問で、プリフィル 12.8 s、エキスパートロード 25956 回、実時間の 73 % がフラッシュ I/O 待ちと測定されました。これはデマンドページングが設計通りに機能しているものであり、リグレッションではありません。2 ターン目以降はウォームです: KV プレフィックスの再利用 + keepwarm リフィルにより、TTFT は約 1 s になります。長いウィンドウに沿ったデコードの減衰は、この SoC の DVFS/熱による挙動です。
+
+---
+
+## 3. 技術詳細
+
+### 3.1 アーキテクチャ
+
+```mermaid
+graph TD
+    subgraph App ["Kotlin / Jetpack Compose"]
+        UI[ChatScreen · dark · send-only composer] --> VM[ChatViewModel]
+        VM --> RT[LlamaRuntime<br/>coroutines + Flow events]
+        VM --> DB[(Room · threads & messages)]
+        ST[SettingsStore] --> VM
+    end
+    subgraph Native ["C JNI shell (llama_chat.c)"]
+        SHELL[generate loop · template-aware thinking control<br/>segment-wise history render · UTF-8-safe streaming]
+    end
+    subgraph Engine ["patched llama.cpp @ b11100 · arm64 CPU-only"]
+        LLIB[libllama.so]
+        GCPU[libggml-cpu.so<br/>NEON kernels + moe_pool]
+    end
+    RT -->|JNI| SHELL --> LLIB --> GCPU
+    GCPU -->|demand-paged expert IO| MODELS[GGUF on flash]
+```
+
+### 3.2 21.7 GB の MoE はなぜスマートフォンに収まるのか — `moe_pool`
+
+35B モデルはトークンごとに、レイヤーあたり 256 のエキスパートのうち移動する部分集合のみをアクティブ化します。そのため出荷時の設計では、エキスパートを常駐メモリではなくフラッシュから**オンデマンドで**ページングします(`ggml/src/ggml-cpu/moe_pool.c`、14 パッチの android バンドとして開発): コピーインのプライベートフレーム、等スロットの状態機械、実測 UFS スループットにチューニングされたバックグラウンド IO ステージングキュー、pin/blob/trim 制御、バイトバジェット内でのターン終了時の keepwarm リフィル、そして 1 つのプロセス内での 8B↔35B 切り替えを可能にする完全なクロスモデルリセット。プールを無効にすると、エンジンは上流とまったく同じようにエキスパート行を解決します(NULL リゾルバー ⇒ ゼロ摂動、シンボルセットの diff により検証済み)。これは、デスクトッププロジェクトが NVMe + Vulkan 上で実装しているものと同じ機構ファミリーです。ここでは設計により CPU/NEON です — 決定的な数値計算と単一のメモリモデルのため、GPU バックエンドは出荷時構成から外されています。
+
+### 3.3 リポジトリマップ
+
+```
+app/                     Android app: Compose UI (src/main/java), JNI shell (src/main/cpp),
+                         instrumented + unit tests (src/androidTest, src/test)
+tools/llama/             build_vendor_libs.sh — engine rebuild from the pinned tree + bands
+tools/model/             push_models.sh — md5-gated model staging to devices
+../vendor/llama.cpp/     materialized by the build scripts from vendor.llama.pin (gitignored; never patched in place)
+../patches/llama.cpp/    common(6) + android(14) hook-point bands + ledger README
+../windows/              desktop companion — hosts the MLX→GGUF converter used in §1.4
+```
diff --git a/android/README_zh.md b/android/README_zh.md
new file mode 100644
index 0000000..143b9c8
--- /dev/null
+++ b/android/README_zh.md
@@ -0,0 +1,177 @@
+# edge0-android — 发布构建指南
+
+[English](README.md) | 中文 | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+本文档面向**从源码构建或评估 Android App 的开发者**。内容包括：快速开始（引擎 → App → 模型 → 测试）、参考设备上的性能实测，以及技术栈背后的技术选型。
+
+edge0 以 **monorepo** 形式发布 —— [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) —— 顶层目录存放共享的引擎供给（`vendor.llama.pin` + 由脚本物化的 `vendor/llama.cpp`、`patches/llama.cpp/` 补丁 band 集）与各平台子项目（`windows/` = 桌面伙伴项目，`android/` = 本 App）。推理运行在固定版本（pinned）的上游 llama.cpp 上，以可重放补丁集的方式打补丁，完全在 CPU 上运行，采用 ARM-NEON kernel 与按需分页的专家池。
+
+第三方声明：见本目录下的 `NOTICE`。模型卡：[Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)。
+
+---
+
+## 1. 快速开始
+
+### 1.1 环境要求
+
+| 组件 | 要求 |
+|---|---|
+| 设备 | arm64-v8a，Android 13+（API 33）；推荐骁龙 8 Elite 级别 |
+| 内存 | 8 GB+ 可运行 8B；12–16 GB 可运行 35B（专家分页，见 §3.2） |
+| 工具链 | JDK 17+，Android SDK 35，**NDK r28**（`28.2.13676358`），CMake ≥ 3.21 + Ninja |
+| Python | 3.10+ 且带 `numpy`（模型转换器，从同级目录 `windows/tools` 运行） |
+| 磁盘 | ≥ 30 GB 可用空间，用于模型源文件与转换后的 GGUF 构建产物 |
+
+NDK **不在**本仓库内 —— 只需安装一次（Android Studio：
+*Android SDK → SDK Tools → NDK (Side by side)*，或通过 CLI）：
+
+```bash
+sdkmanager --install "ndk;28.2.13676358"
+```
+
+然后用 `NDK_DIR`（或导出的 `ANDROID_NDK_HOME`）让构建指向它：
+`build_vendor_libs.sh` 会从该 NDK 安装中读取编译器工具链**以及**其暂存
+（§1.2）的 `libomp.so` runtime；若缺少 NDK，会按该提示快速失败，而不是
+产出损坏的库集合。
+
+### 1.2 构建引擎
+
+原生库来自固定版本的上游源码树，本平台的补丁 band 会被重放到一个隔离的
+worktree 中 —— vendor 树**绝不就地打补丁**：
+
+```bash
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/android
+bash tools/llama/build_vendor_libs.sh
+```
+
+脚本把 `../patches/llama.cpp/{common,android}`（6 + 14 个 band）重放到
+固定版本的 llama.cpp 树上 —— pin 记录在 `../vendor.llama.pin`（当前为
+`7ab4ee7`，tag b11100）；该树**不是** submodule，因此首次运行时会从上游
+克隆到 `../vendor/llama.cpp`（gitignored；设置 `EDGE0_LLAMA_URL` 可使用
+镜像）并 detach 到 pin 的提交。重放发生在一个 gitignored 的消费方
+worktree 中，脚本会校验黄金结果树哈希，并产出四个引擎共享库（外加
+`libggml-cpu` 所需的 NDK `libomp.so` runtime）与头文件到
+`build-dl/llama-libs/`（App 的 jniLibs 暂存点）。补丁变更后用 `--replay`
+重新应用 band；树不匹配 ⇒ RED，构建拒绝启动。App 构建前需要执行一次
+本步骤 —— Gradle 插件会从 `build-dl/llama-libs/` 读取这些库。
+
+### 1.3 构建并安装 App
+
+```bash
+./gradlew :app:assembleDebug
+./gradlew :app:installDebug        # or adb install -r app/build/outputs/apk/debug/app-debug.apk
+```
+
+### 1.4 模型
+
+GGUF 构建产物由已发布的 checkpoint 在本地生成 —— 全部内容都留在本目录的
+`models/` 下（gitignored）：
+
+```bash
+huggingface-cli download Edge0/Edge0-8B-A1B-preview --local-dir models/edge0-8b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-8b
+#  → models/edge0-8b-gguf/{edge0-8b.gguf, lora_edge0_8b-gguf.gguf, manifest.json}
+
+huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+python ../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-35b
+
+bash tools/model/push_models.sh --all    # md5-gated staging onto the device
+```
+
+转换器只需 python3 + numpy（无需 MLX/torch runtime —— “MLX” 指的是磁盘上
+的 checkpoint 布局）。它执行带数值一致性校验门的 r3 重打包并输出 sha256
+清单；正确的转换会逐字节复现 `push_models.sh` 中列出的基线校验和。也可以
+通过 App 内的选择器把模型拷贝到 `files/models/`。
+
+### 1.5 运行
+
+启动 **Edge0 Chat**。标题栏显示当前模型（8B / 35B），右上角按钮用于切换。
+输入框只有发送功能；temperature、thinking 与系统提示词在抽屉设置中。每条
+回复附带一行内联指标：`tokens · TTFT · prefill t/s · decode t/s · RSS`。
+
+### 1.6 测试
+
+仪器化回归测试（需要两个模型都已就位 —— 重新安装测试 APK 会清空 App
+数据，因此请在运行前重新推送模型）：
+
+```bash
+./gradlew :app:installDebugAndroidTest
+bash tools/model/push_models.sh --all
+adb shell am instrument -w -e class dev.edge0.runtime.app.LlamaRuntimeTest \
+  dev.edge0.runtime.app.test/androidx.test.runner.AndroidJUnitRunner
+# expected: OK (8 tests), ~7 min on the reference device
+```
+
+覆盖范围：8B/35B 冒烟、35B↔8B 进程内切换、前缀复用保真度、thinking
+开/关 × 系统提示词四象限（8B 开启 + 35B 关闭泄漏探测），以及多轮渲染下的
+身份保持。宿主侧逻辑测试：`./gradlew :app:testDebugUnitTest`（26 个测试，
+无需设备）。
+
+---
+
+## 2. 性能实测
+
+参考设备：**Lenovo TB322FC（骁龙 8 Elite，16 GB RAM）**，出厂配置，持续
+测试窗口（前段 = 加速频率，尾段 = 热稳态 —— 两者都如实报告，而非只挑
+峰值）。
+
+| 模型 | TTFT（热轮次） | 解码 | Prefill | 会话 RSS |
+|---|---:|---:|---:|---:|
+| **8B**（Q8 级 GGUF + LoRA） | ≈ 1.4 s | 480 s 窗口内 29–32 → 约 10 t/s | 约 100 t/s | ≈ 250 MB |
+| **35B**（混合 int8 MoE，按需分页） | 热轮次 ≈ 1.1 s（新话题首轮 ≈ 10–15 s —— 专家池冷启动，受闪存 IO 限制） | App 内 6–9 t/s（CLI 持续 9.46 t/s） | 每个增量轮次约 1 s | 池预算 2–6 GB；常驻 ≪ 文件大小 |
+
+说明：“新话题首轮”要付出冷池代价 —— 例如一个 55 token 的问题实测
+prefill 12.8 s、25956 次专家加载，73% 的墙钟时间耗在闪存 I/O 等待上；
+这是按需分页按设计工作，而非性能回退。第二轮起即为热轮次：KV 前缀复用 +
+keepwarm 回填把 TTFT 降到约 1 s。长窗口内解码速度的衰减是该 SoC 上的
+DVFS/热行为。
+
+---
+
+## 3. 技术细节
+
+### 3.1 架构
+
+```mermaid
+graph TD
+    subgraph App ["Kotlin / Jetpack Compose"]
+        UI[ChatScreen · dark · send-only composer] --> VM[ChatViewModel]
+        VM --> RT[LlamaRuntime<br/>coroutines + Flow events]
+        VM --> DB[(Room · threads & messages)]
+        ST[SettingsStore] --> VM
+    end
+    subgraph Native ["C JNI shell (llama_chat.c)"]
+        SHELL[generate loop · template-aware thinking control<br/>segment-wise history render · UTF-8-safe streaming]
+    end
+    subgraph Engine ["patched llama.cpp @ b11100 · arm64 CPU-only"]
+        LLIB[libllama.so]
+        GCPU[libggml-cpu.so<br/>NEON kernels + moe_pool]
+    end
+    RT -->|JNI| SHELL --> LLIB --> GCPU
+    GCPU -->|demand-paged expert IO| MODELS[GGUF on flash]
+```
+
+### 3.2 为什么 21.7 GB 的 MoE 能装进手机 —— `moe_pool`
+
+35B 模型每个 token 只激活其每层 256 个专家中一个滚动的子集，因此发布版
+设计让专家从闪存**按需**分页，而非驻留内存
+（`ggml/src/ggml-cpu/moe_pool.c`，以 14 个补丁的 android band 开发）：
+copy-in 私有帧、等槽位状态机、按实测 UFS 吞吐调优的后台 IO 暂存队列、
+pin/blob/trim 控制、字节预算内的轮末 keepwarm 回填，以及支持单进程内
+8B↔35B 切换的完整跨模型重置。禁用专家池时，引擎解析专家行的行为与上游
+完全一致（NULL resolver ⇒ 零扰动，经符号集 diff 验证）。这与桌面项目在
+NVMe + Vulkan 上实现的是同一族机制；这里按设计走 CPU/NEON —— GPU 后端
+不进入发布配置，以保证确定性数值与单一内存模型。
+
+### 3.3 仓库结构
+
+```
+app/                     Android app: Compose UI (src/main/java), JNI shell (src/main/cpp),
+                         instrumented + unit tests (src/androidTest, src/test)
+tools/llama/             build_vendor_libs.sh — engine rebuild from the pinned tree + bands
+tools/model/             push_models.sh — md5-gated model staging to devices
+../vendor/llama.cpp/     materialized by the build scripts from vendor.llama.pin (gitignored; never patched in place)
+../patches/llama.cpp/    common(6) + android(14) hook-point bands + ledger README
+../windows/              desktop companion — hosts the MLX→GGUF converter used in §1.4
+```
diff --git a/android/app/build.gradle.kts b/android/app/build.gradle.kts
new file mode 100644
index 0000000..37876ca
--- /dev/null
+++ b/android/app/build.gradle.kts
@@ -0,0 +1,95 @@
+// app/build.gradle.kts - the chat app module.
+// Single APK: native llama shell (shared libs staged via jniLibs) + Compose UI.
+// git SHA provenance: configuration cache forbids external processes at configure time,
+// so the build takes E0_GIT_SHA from the environment:
+// unset falls back to "unknown" (never blocks the build).
+plugins {
+    alias(libs.plugins.android.application)
+    alias(libs.plugins.kotlin.android)
+    alias(libs.plugins.kotlin.compose)
+    alias(libs.plugins.ksp)
+    alias(libs.plugins.kotlinx.serialization)
+}
+
+android {
+    namespace = "dev.edge0.runtime.app"
+    compileSdk = libs.versions.compileSdk.get().toInt()
+
+    defaultConfig {
+        applicationId = "dev.edge0.runtime.app"
+        minSdk = libs.versions.minSdk.get().toInt()
+        targetSdk = libs.versions.compileSdk.get().toInt()
+        versionCode = 1
+        versionName = "0.1.0"
+        ndk { abiFilters += "arm64-v8a" } // arm64 only
+        externalNativeBuild {
+            cmake {
+                abiFilters += "arm64-v8a"
+                arguments += "-DLLAMA_LIBS=" + rootProject.file("build-dl/llama-libs/arm64-v8a").absolutePath
+            }
+        }
+        // instrumented suites follow install -> stage assets -> instrument; the first step is gradle.
+        testInstrumentationRunner = "androidx.test.runner.AndroidJUnitRunner"
+        buildConfigField("String", "GIT_SHA",
+            "\"${System.getenv("E0_GIT_SHA") ?: "unknown"}\"")
+    }
+    buildTypes {
+        debug { isMinifyEnabled = false }
+        release {
+            isMinifyEnabled = false // no minification at this stage
+        }
+    }
+    // llama native shell; .so assembly point is build-dl/llama-libs (gitignored)
+    externalNativeBuild { cmake { path = file("src/main/cpp/CMakeLists.txt"); version = "3.22.1" } }
+    sourceSets {
+        getByName("main") {
+            val ll = rootProject.file("build-dl/llama-libs/arm64-v8a")
+            if (ll.isDirectory) jniLibs.srcDirs(rootProject.file("build-dl/llama-libs"))
+        }
+    }
+
+    buildFeatures {
+        compose = true
+        buildConfig = true
+    }
+    compileOptions {
+        sourceCompatibility = JavaVersion.VERSION_17
+        targetCompatibility = JavaVersion.VERSION_17
+    }
+    kotlinOptions { jvmTarget = "17" }
+    packaging {
+        resources.excludes += "/META-INF/{AL2.0,LGPL2.1}"
+    }
+}
+
+ksp {
+    // Room schema export into app/schemas/ for migration traceability
+    arg("room.schemaLocation", "$projectDir/schemas".toString())
+    arg("room.generateKotlin", "true")
+}
+
+dependencies {
+    implementation(platform(libs.compose.bom))
+    implementation(libs.compose.ui)
+    implementation(libs.compose.ui.tooling.preview)
+    implementation(libs.compose.material3)
+    implementation(libs.compose.material.icons.extended)
+    debugImplementation(libs.compose.ui.tooling)
+    implementation(libs.androidx.activity.compose)
+    implementation(libs.androidx.navigation.compose)
+    implementation(libs.androidx.lifecycle.viewmodel.compose)
+    implementation(libs.androidx.lifecycle.runtime.compose)
+    implementation(libs.androidx.lifecycle.process) // lifecycle hooks for trim mapping
+    implementation(libs.androidx.room.runtime)
+    implementation(libs.androidx.room.ktx)
+    ksp(libs.androidx.room.compiler)
+    implementation(libs.androidx.datastore.preferences)
+    implementation(libs.kotlinx.coroutines.android)
+    implementation(libs.kotlinx.serialization.json)
+    implementation(libs.markdown.renderer.m3) // markdown rendering (tables/links/strikethrough)
+
+    testImplementation("junit:junit:4.13.2") // JVM unit tests
+    androidTestImplementation("junit:junit:4.13.2")
+    androidTestImplementation("androidx.test.ext:junit:1.2.1")
+    androidTestImplementation("androidx.test:runner:1.6.2")
+}
diff --git a/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/1.json b/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/1.json
new file mode 100644
index 0000000..c7705c3
--- /dev/null
+++ b/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/1.json
@@ -0,0 +1,169 @@
+{
+  "formatVersion": 1,
+  "database": {
+    "version": 1,
+    "identityHash": "7e3b71bb48f755556c1020fe2717fb45",
+    "entities": [
+      {
+        "tableName": "threads",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`id` TEXT NOT NULL, `title` TEXT, `modelId` TEXT NOT NULL, `createdAt` INTEGER NOT NULL, `updatedAt` INTEGER NOT NULL, PRIMARY KEY(`id`))",
+        "fields": [
+          {
+            "fieldPath": "id",
+            "columnName": "id",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "title",
+            "columnName": "title",
+            "affinity": "TEXT"
+          },
+          {
+            "fieldPath": "modelId",
+            "columnName": "modelId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "createdAt",
+            "columnName": "createdAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "updatedAt",
+            "columnName": "updatedAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": false,
+          "columnNames": [
+            "id"
+          ]
+        },
+        "indices": [
+          {
+            "name": "index_threads_updatedAt",
+            "unique": false,
+            "columnNames": [
+              "updatedAt"
+            ],
+            "orders": [],
+            "createSql": "CREATE INDEX IF NOT EXISTS `index_threads_updatedAt` ON `${TABLE_NAME}` (`updatedAt`)"
+          }
+        ]
+      },
+      {
+        "tableName": "messages",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`id` TEXT NOT NULL, `threadId` TEXT NOT NULL, `role` TEXT NOT NULL, `content` TEXT NOT NULL, `thinking` TEXT, `thinkingMs` INTEGER, `status` TEXT NOT NULL, `modelId` TEXT NOT NULL, `createdAt` INTEGER NOT NULL, `promptTokens` INTEGER, `newTokens` INTEGER, `decodeTokS` REAL, PRIMARY KEY(`id`), FOREIGN KEY(`threadId`) REFERENCES `threads`(`id`) ON UPDATE NO ACTION ON DELETE CASCADE )",
+        "fields": [
+          {
+            "fieldPath": "id",
+            "columnName": "id",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "threadId",
+            "columnName": "threadId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "role",
+            "columnName": "role",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "content",
+            "columnName": "content",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "thinking",
+            "columnName": "thinking",
+            "affinity": "TEXT"
+          },
+          {
+            "fieldPath": "thinkingMs",
+            "columnName": "thinkingMs",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "status",
+            "columnName": "status",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "modelId",
+            "columnName": "modelId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "createdAt",
+            "columnName": "createdAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "promptTokens",
+            "columnName": "promptTokens",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "newTokens",
+            "columnName": "newTokens",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "decodeTokS",
+            "columnName": "decodeTokS",
+            "affinity": "REAL"
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": false,
+          "columnNames": [
+            "id"
+          ]
+        },
+        "indices": [
+          {
+            "name": "index_messages_threadId_createdAt",
+            "unique": false,
+            "columnNames": [
+              "threadId",
+              "createdAt"
+            ],
+            "orders": [],
+            "createSql": "CREATE INDEX IF NOT EXISTS `index_messages_threadId_createdAt` ON `${TABLE_NAME}` (`threadId`, `createdAt`)"
+          }
+        ],
+        "foreignKeys": [
+          {
+            "table": "threads",
+            "onDelete": "CASCADE",
+            "onUpdate": "NO ACTION",
+            "columns": [
+              "threadId"
+            ],
+            "referencedColumns": [
+              "id"
+            ]
+          }
+        ]
+      }
+    ],
+    "setupQueries": [
+      "CREATE TABLE IF NOT EXISTS room_master_table (id INTEGER PRIMARY KEY,identity_hash TEXT)",
+      "INSERT OR REPLACE INTO room_master_table (id,identity_hash) VALUES(42, '7e3b71bb48f755556c1020fe2717fb45')"
+    ]
+  }
+}
\ No newline at end of file
diff --git a/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/2.json b/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/2.json
new file mode 100644
index 0000000..e39ee62
--- /dev/null
+++ b/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/2.json
@@ -0,0 +1,174 @@
+{
+  "formatVersion": 1,
+  "database": {
+    "version": 2,
+    "identityHash": "95511891effc3c18a7d5a92ebfc02e22",
+    "entities": [
+      {
+        "tableName": "threads",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`id` TEXT NOT NULL, `title` TEXT, `modelId` TEXT NOT NULL, `createdAt` INTEGER NOT NULL, `updatedAt` INTEGER NOT NULL, PRIMARY KEY(`id`))",
+        "fields": [
+          {
+            "fieldPath": "id",
+            "columnName": "id",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "title",
+            "columnName": "title",
+            "affinity": "TEXT"
+          },
+          {
+            "fieldPath": "modelId",
+            "columnName": "modelId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "createdAt",
+            "columnName": "createdAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "updatedAt",
+            "columnName": "updatedAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": false,
+          "columnNames": [
+            "id"
+          ]
+        },
+        "indices": [
+          {
+            "name": "index_threads_updatedAt",
+            "unique": false,
+            "columnNames": [
+              "updatedAt"
+            ],
+            "orders": [],
+            "createSql": "CREATE INDEX IF NOT EXISTS `index_threads_updatedAt` ON `${TABLE_NAME}` (`updatedAt`)"
+          }
+        ]
+      },
+      {
+        "tableName": "messages",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`id` TEXT NOT NULL, `threadId` TEXT NOT NULL, `role` TEXT NOT NULL, `content` TEXT NOT NULL, `thinking` TEXT, `thinkingMs` INTEGER, `status` TEXT NOT NULL, `modelId` TEXT NOT NULL, `createdAt` INTEGER NOT NULL, `promptTokens` INTEGER, `newTokens` INTEGER, `decodeTokS` REAL, `genTokens` BLOB, PRIMARY KEY(`id`), FOREIGN KEY(`threadId`) REFERENCES `threads`(`id`) ON UPDATE NO ACTION ON DELETE CASCADE )",
+        "fields": [
+          {
+            "fieldPath": "id",
+            "columnName": "id",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "threadId",
+            "columnName": "threadId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "role",
+            "columnName": "role",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "content",
+            "columnName": "content",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "thinking",
+            "columnName": "thinking",
+            "affinity": "TEXT"
+          },
+          {
+            "fieldPath": "thinkingMs",
+            "columnName": "thinkingMs",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "status",
+            "columnName": "status",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "modelId",
+            "columnName": "modelId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "createdAt",
+            "columnName": "createdAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "promptTokens",
+            "columnName": "promptTokens",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "newTokens",
+            "columnName": "newTokens",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "decodeTokS",
+            "columnName": "decodeTokS",
+            "affinity": "REAL"
+          },
+          {
+            "fieldPath": "genTokens",
+            "columnName": "genTokens",
+            "affinity": "BLOB"
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": false,
+          "columnNames": [
+            "id"
+          ]
+        },
+        "indices": [
+          {
+            "name": "index_messages_threadId_createdAt",
+            "unique": false,
+            "columnNames": [
+              "threadId",
+              "createdAt"
+            ],
+            "orders": [],
+            "createSql": "CREATE INDEX IF NOT EXISTS `index_messages_threadId_createdAt` ON `${TABLE_NAME}` (`threadId`, `createdAt`)"
+          }
+        ],
+        "foreignKeys": [
+          {
+            "table": "threads",
+            "onDelete": "CASCADE",
+            "onUpdate": "NO ACTION",
+            "columns": [
+              "threadId"
+            ],
+            "referencedColumns": [
+              "id"
+            ]
+          }
+        ]
+      }
+    ],
+    "setupQueries": [
+      "CREATE TABLE IF NOT EXISTS room_master_table (id INTEGER PRIMARY KEY,identity_hash TEXT)",
+      "INSERT OR REPLACE INTO room_master_table (id,identity_hash) VALUES(42, '95511891effc3c18a7d5a92ebfc02e22')"
+    ]
+  }
+}
\ No newline at end of file
diff --git a/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/3.json b/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/3.json
new file mode 100644
index 0000000..e795529
--- /dev/null
+++ b/android/app/schemas/dev.edge0.runtime.app.data.AppDatabase/3.json
@@ -0,0 +1,189 @@
+{
+  "formatVersion": 1,
+  "database": {
+    "version": 3,
+    "identityHash": "bedb08d6046568298f8264b8d94ffd0f",
+    "entities": [
+      {
+        "tableName": "threads",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`id` TEXT NOT NULL, `title` TEXT, `modelId` TEXT NOT NULL, `createdAt` INTEGER NOT NULL, `updatedAt` INTEGER NOT NULL, PRIMARY KEY(`id`))",
+        "fields": [
+          {
+            "fieldPath": "id",
+            "columnName": "id",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "title",
+            "columnName": "title",
+            "affinity": "TEXT"
+          },
+          {
+            "fieldPath": "modelId",
+            "columnName": "modelId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "createdAt",
+            "columnName": "createdAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "updatedAt",
+            "columnName": "updatedAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": false,
+          "columnNames": [
+            "id"
+          ]
+        },
+        "indices": [
+          {
+            "name": "index_threads_updatedAt",
+            "unique": false,
+            "columnNames": [
+              "updatedAt"
+            ],
+            "orders": [],
+            "createSql": "CREATE INDEX IF NOT EXISTS `index_threads_updatedAt` ON `${TABLE_NAME}` (`updatedAt`)"
+          }
+        ]
+      },
+      {
+        "tableName": "messages",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`id` TEXT NOT NULL, `threadId` TEXT NOT NULL, `role` TEXT NOT NULL, `content` TEXT NOT NULL, `thinking` TEXT, `thinkingMs` INTEGER, `status` TEXT NOT NULL, `modelId` TEXT NOT NULL, `createdAt` INTEGER NOT NULL, `promptTokens` INTEGER, `newTokens` INTEGER, `decodeTokS` REAL, `genTokens` BLOB, `ttftMs` INTEGER, `prefillTokS` REAL, `memBytes` INTEGER, PRIMARY KEY(`id`), FOREIGN KEY(`threadId`) REFERENCES `threads`(`id`) ON UPDATE NO ACTION ON DELETE CASCADE )",
+        "fields": [
+          {
+            "fieldPath": "id",
+            "columnName": "id",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "threadId",
+            "columnName": "threadId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "role",
+            "columnName": "role",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "content",
+            "columnName": "content",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "thinking",
+            "columnName": "thinking",
+            "affinity": "TEXT"
+          },
+          {
+            "fieldPath": "thinkingMs",
+            "columnName": "thinkingMs",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "status",
+            "columnName": "status",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "modelId",
+            "columnName": "modelId",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "createdAt",
+            "columnName": "createdAt",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "promptTokens",
+            "columnName": "promptTokens",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "newTokens",
+            "columnName": "newTokens",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "decodeTokS",
+            "columnName": "decodeTokS",
+            "affinity": "REAL"
+          },
+          {
+            "fieldPath": "genTokens",
+            "columnName": "genTokens",
+            "affinity": "BLOB"
+          },
+          {
+            "fieldPath": "ttftMs",
+            "columnName": "ttftMs",
+            "affinity": "INTEGER"
+          },
+          {
+            "fieldPath": "prefillTokS",
+            "columnName": "prefillTokS",
+            "affinity": "REAL"
+          },
+          {
+            "fieldPath": "memBytes",
+            "columnName": "memBytes",
+            "affinity": "INTEGER"
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": false,
+          "columnNames": [
+            "id"
+          ]
+        },
+        "indices": [
+          {
+            "name": "index_messages_threadId_createdAt",
+            "unique": false,
+            "columnNames": [
+              "threadId",
+              "createdAt"
+            ],
+            "orders": [],
+            "createSql": "CREATE INDEX IF NOT EXISTS `index_messages_threadId_createdAt` ON `${TABLE_NAME}` (`threadId`, `createdAt`)"
+          }
+        ],
+        "foreignKeys": [
+          {
+            "table": "threads",
+            "onDelete": "CASCADE",
+            "onUpdate": "NO ACTION",
+            "columns": [
+              "threadId"
+            ],
+            "referencedColumns": [
+              "id"
+            ]
+          }
+        ]
+      }
+    ],
+    "setupQueries": [
+      "CREATE TABLE IF NOT EXISTS room_master_table (id INTEGER PRIMARY KEY,identity_hash TEXT)",
+      "INSERT OR REPLACE INTO room_master_table (id,identity_hash) VALUES(42, 'bedb08d6046568298f8264b8d94ffd0f')"
+    ]
+  }
+}
\ No newline at end of file
diff --git a/android/app/src/androidTest/java/dev/edge0/runtime/app/LlamaRuntimeTest.kt b/android/app/src/androidTest/java/dev/edge0/runtime/app/LlamaRuntimeTest.kt
new file mode 100644
index 0000000..4744267
--- /dev/null
+++ b/android/app/src/androidTest/java/dev/edge0/runtime/app/LlamaRuntimeTest.kt
@@ -0,0 +1,326 @@
+// LlamaRuntimeTest - app-domain smoke suite for the llama engine shell.
+// Prerequisite: stage the models AFTER installing the test APK -
+//   adb shell run-as dev.edge0.runtime.app cp /data/local/tmp/lgguf/edge0-8b.gguf files/models/edge0-8b.gguf
+// instrumented installs wipe app data, so staging must follow installation.
+// Assertions keep a conservative floor (availability over speed; the CLI benchmark
+// owns ceiling checks): >= 8 tokens produced, decode > 1 t/s, non-empty text.
+package dev.edge0.runtime.app
+
+import dev.edge0.runtime.app.runtime.GenEvent
+import dev.edge0.runtime.app.runtime.LlamaRuntime
+import dev.edge0.runtime.engine.ChatMessage
+import dev.edge0.runtime.engine.GenParams
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.runBlocking
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import java.io.File
+
+class LlamaRuntimeTest {
+    @Test
+    fun eightBChatRoundTrip() = runBlocking {
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        // files staged via run-as may carry an SELinux label the instrumentation cannot
+        // read; search candidate roots instead. Production loads through the normal path.
+        val want = androidx.test.platform.app.InstrumentationRegistry.getArguments()
+            .getString("model") ?: "edge0-35b.gguf"
+        val model = sequenceOf(
+            File(ctx.getDir("models", 0), want),
+            File(ctx.filesDir, "models/$want"),
+            File(ctx.cacheDir, "models/$want"),
+            File(ctx.getExternalFilesDir(null) ?: ctx.filesDir, "models/$want"),
+        ).firstOrNull { it.isFile }
+            ?: run {
+                val hit = generateSequence(ctx.dataDir) { f ->
+                    f.listFiles()?.firstOrNull { it.isDirectory && !it.name.startsWith(".") }
+                }.flatMap { (it.walk().firstOrNull { f2 -> f2.name == want }?.let { x -> sequenceOf(x) } ?: emptySequence()) }
+                    .firstOrNull()
+                hit ?: File(ctx.getDir("models", 0), want)
+            }
+        assertTrue("model not found anywhere: ${model.path}", model.isFile)
+
+        val rt = LlamaRuntime()
+        rt.ensureLoaded(model.absolutePath)
+        assertTrue("load failed: ${rt.state.value}", rt.activeModelDir != null)
+
+        val pieces = StringBuilder()
+        var tps = 0.0
+        rt.submit("t1", listOf(ChatMessage("user", "Introduce yourself in one sentence")),
+                  true, GenParams(temperature = 0f, maxNewTokens = 32))
+            .toList().forEach { ev ->
+                when (ev) {
+                    is GenEvent.TextPiece -> pieces.append(ev.s)
+                    is GenEvent.ThinkPiece -> pieces.append(ev.s)
+                    is GenEvent.Done -> tps = ev.metrics?.decodeTokS ?: 0.0
+                    is GenEvent.Failed -> throw AssertionError("gen failed: ${ev.diag}")
+                    else -> Unit
+                }
+            }
+        android.util.Log.i("LlamaRT", "reply=${pieces.take(80)} tps=$tps")
+        // metrics wiring assertions (pool/rss must not be N/A)
+        val rm = rt.lastMetrics.value
+        assertTrue("metrics missing", rm != null)
+        assertTrue("pool metrics missing", rm?.pool != null)
+        if (want.contains("35b")) {
+            assertTrue("35B pool counters not wired", (rm!!.pool!!.hits + rm.pool!!.misses) > 0 && rm.pool!!.slots_total > 0)
+        } else {
+            // the 8B tier is fully resident: counters exist and read zero - that is the correct shape
+            assertTrue("8B should not occupy the pool", rm!!.pool!!.slots_total == 0L)
+        }
+        assertTrue("peak rss N/A", (rm?.peakRssBytes ?: 0L) > 100L * 1024 * 1024)
+        assertTrue("first-token timing anomaly: ${rm?.firstTokenMs}", (rm?.firstTokenMs ?: 0) > 0)
+        assertTrue("empty reply", pieces.isNotBlank())
+        val floor = if (want.contains("35b")) 0.5 else 1.0
+        assertTrue("tps too low: $tps", tps > floor)
+
+        // second turn exercises KV prefix reuse (history diff sync)
+        val r2 = rt.submit("t2", listOf(
+            ChatMessage("user", "Introduce yourself in one sentence"),
+            ChatMessage("assistant", pieces.toString().trim().take(40)),
+            ChatMessage("user", "a bit shorter"),
+        ), true, GenParams(temperature = 0f, maxNewTokens = 24)).toList()
+        assertTrue("round2 no done", r2.any { it is GenEvent.Done })
+        rt.unloadActive()
+    }
+
+    // crash regression: use 8B, switch to 35B, send in a new session - stale LoRA/session
+// state attaching to the new graph used to abort; bidirectional switching must survive.
+    @Test
+    fun switch8bThen35bSurvival() = runBlocking {
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        fun find(n: String) = sequenceOf(File(ctx.getDir("models", 0), n), File(ctx.filesDir, "models/$n"))
+            .firstOrNull { it.isFile }
+        val m8 = find("edge0-8b.gguf") ?: return@runBlocking
+        val rt = LlamaRuntime()
+        fun round(model: File, q: String) = runBlocking {
+            rt.ensureLoaded(model.absolutePath)
+            assertTrue("load fail ${model.name}: ${rt.state.value}", rt.activeModelDir != null)
+            var n = 0
+            rt.submit("sw", listOf(ChatMessage("user", q)), true,
+                      GenParams(temperature = 0f, maxNewTokens = 8))
+                .toList().forEach { ev ->
+                    when (ev) {
+                        is GenEvent.TextPiece -> n += ev.s.length
+                        is GenEvent.ThinkPiece -> n += ev.s.length
+                        is GenEvent.Failed -> throw AssertionError("${model.name}: ${ev.diag}")
+                        else -> Unit
+                    }
+                }
+            assertTrue("no output ${model.name}", n > 2)
+        }
+        round(m8, "hello")
+        val m35 = find("edge0-35b.gguf")
+        if (m35 != null) {
+            round(m35, "hello")          // the original crash point
+            round(m8, "hello")           // switch back (pool/registry reset verified both ways)
+        }
+        rt.unloadActive()
+    }
+
+    // real multi-turn on 35B (long answer, then thanks): log reused/prefill per round to locate forks
+    @Test
+    fun prefixReuseProbe() = runBlocking {
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        val want = androidx.test.platform.app.InstrumentationRegistry.getArguments()
+            .getString("model") ?: "edge0-35b.gguf"
+        val model = sequenceOf(File(ctx.getDir("models", 0), want), File(ctx.filesDir, "models/$want"))
+            .firstOrNull { it.isFile } ?: return@runBlocking
+        val syncMode = androidx.test.platform.app.InstrumentationRegistry.getArguments()
+            .getString("sync") ?: ""
+        if (syncMode.isNotEmpty()) android.system.Os.setenv("E0_POOL_SYNC", "1", true)
+        val rt = LlamaRuntime()
+        rt.ensureLoaded(model.absolutePath)
+        val probeThinking = (androidx.test.platform.app.InstrumentationRegistry.getArguments()
+            .getString("thinking") ?: "off") == "on"
+        val hist = mutableListOf<ChatMessage>()
+        val qs = listOf("Describe three core advantages of Mixture-of-Experts models in detail.", "thanks")
+        for ((i, q) in qs.withIndex()) {
+            hist.add(ChatMessage("user", q))
+            val sb = StringBuilder()
+            rt.submit("probe$i", hist.toList(), probeThinking, GenParams(temperature = 0f, maxNewTokens = if (i == 0) 200 else 20))
+                .toList().forEach { ev ->
+                    when (ev) {
+                        is GenEvent.TextPiece -> sb.append(ev.s)
+                        is GenEvent.ThinkPiece -> sb.append("«T»").append(ev.s)
+                        is GenEvent.Failed -> throw AssertionError(ev.diag)
+                        else -> Unit
+                    }
+                }
+            hist.add(ChatMessage("assistant", sb.toString()))
+            if (i == 0) Thread.sleep(25000)   // simulate reading the reply: the keepwarm battleground
+            android.util.Log.i("ReuseProbe", "round$i head=[${sb.toString().take(70).replace("\n", " | ")}] len=${sb.length} metrics=${rt.lastMetrics.value?.let { "prompt=${it.promptTokens} reused=${it.prefix_reused_tokens} prefill_ms=${it.prefillMs} first_ms=${it.firstTokenMs} pool=${it.pool?.let { q -> "hits=${q.hits} miss=${q.misses} evict=${q.evictions} stall_ms=${q.stall_ns / 1000000} res_MB=${q.resident_bytes / 1048576} pf=${q.prefetch_used}/${q.prefetch_mispredict}" }}" }}")
+        }
+        rt.unloadActive()
+    }
+
+    // thinking quality gate: same greedy prompt, 12 tokens, quadrants of {system,thinking}.
+    // Verdict via logcat ThinkGate: off cells must answer directly, on cells must think.
+    // Regression anchor: custom-system off leaking thinking (the ensure-position fix).
+    @Test
+    fun thinkSwitchQualityGate() = runBlocking {
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        val want = androidx.test.platform.app.InstrumentationRegistry.getArguments()
+            .getString("model") ?: "edge0-8b.gguf"
+        val model = sequenceOf(File(ctx.getDir("models", 0), want), File(ctx.filesDir, "models/$want"))
+            .firstOrNull { it.isFile } ?: return@runBlocking
+        val q = "Explain why the sky is blue in one short sentence."
+        val sp = "You're Edge0, an on-device AI assistant."
+        val heads = ArrayList<String>()
+        for (sys in listOf(false, true)) for (th in listOf(false, true)) {
+            val rt = LlamaRuntime()
+            rt.ensureLoaded(model.absolutePath)
+            val base0 = if (sys) listOf(ChatMessage("system", sp)) else emptyList()
+            val turns = ArrayList<String>()
+            val hist = base0.toMutableList()
+            for (t in 0..1) {  // round two goes through the sidecar path (the empty-answer incident scene)
+                val sb = StringBuilder()
+                hist.add(ChatMessage("user", if (t == 0) q else "thanks, one more short sentence"))
+                rt.submit("g$sys$th$t", hist.toList(), th,
+                          GenParams(temperature = 0f, maxNewTokens = 12))
+                    .toList().forEach { ev ->
+                        when (ev) {
+                            is GenEvent.TextPiece -> sb.append(ev.s)
+                            is GenEvent.ThinkPiece -> sb.append("<T>").append(ev.s)
+                            is GenEvent.Failed -> throw AssertionError(ev.diag)
+                            else -> Unit
+                        }
+                    }
+                turns.add("t" + t + "=" + sb.toString().take(44).replace("\n", " "))
+                hist.add(ChatMessage("assistant", sb.toString()))
+            }
+            heads.add((if (sys) "sys" else "no") + "-th" + (if (th) "on" else "off") + ": " + turns.joinToString(" ~ "))
+            // empty-answer anchor: no quadrant, no round may return empty
+            org.junit.Assert.assertTrue("gate: empty reply in " + heads.last(),
+                turns.all { it.length > 6 })
+            rt.unloadActive()
+        }
+        android.util.Log.i("ThinkGate", heads.joinToString(" ||| "))
+        org.junit.Assert.assertTrue("gate: 4 quadrants", heads.size == 4)
+        org.junit.Assert.assertTrue("gate: off-no-sys empty", heads[0].length > 12)
+    }
+
+    // multi-turn identity probe: chat once (builds sidecar history), then ask who you are.
+    // Single-turn was green; the incident lived in the multi-turn path.
+    @Test
+    fun whoAreYouSysProbe() = runBlocking {
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        val model = File(ctx.filesDir, "models/edge0-8b.gguf")
+        if (!model.isFile) return@runBlocking
+        val rt = LlamaRuntime()
+        rt.ensureLoaded(model.absolutePath)
+        var heads_all = ""
+        val hist = mutableListOf(ChatMessage("system", "You're Edge0, an on-device AI assistant."))
+        for ((i, q) in listOf("What is the capital of France? One word.",
+                              "Who are you? What is your name? Answer in one sentence.").withIndex()) {
+            hist.add(ChatMessage("user", q))
+            val sb = StringBuilder()
+            rt.submit("way$i", hist.toList(), false, GenParams(temperature = 0.2f, maxNewTokens = 60))
+                .toList().forEach { ev ->
+                    when (ev) {
+                        is GenEvent.TextPiece -> sb.append(ev.s)
+                        is GenEvent.ThinkPiece -> sb.append("<T>")
+                        is GenEvent.Failed -> throw AssertionError(ev.diag)
+                        else -> Unit
+                    }
+                }
+            heads_all += sb.toString(); android.util.Log.i("WayProbe", "turn" + i + " len=" + sb.length + " ans=" + sb.toString().take(90))
+            hist.add(ChatMessage("assistant", sb.toString()))
+        }
+        rt.unloadActive()
+        org.junit.Assert.assertTrue("way: identity lost in multi-turn",
+            heads_all.contains("Edge0"))
+    }
+
+    // fork dissecting table: mode=ui replays the stored UI text as the assistant turn;
+// mode=native replays lastAssistantNative ground truth
+    @Test
+    fun reuseForkBisect() = runBlocking {
+        val mode = androidx.test.platform.app.InstrumentationRegistry.getArguments()
+            .getString("mode") ?: "ui"
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        val want = androidx.test.platform.app.InstrumentationRegistry.getArguments().getString("model") ?: "edge0-35b.gguf"
+        val model = sequenceOf(
+            File(ctx.getDir("models", 0), want),
+            File(ctx.filesDir, "models/$want"),
+        ).firstOrNull { it.isFile } ?: return@runBlocking
+        val rt = LlamaRuntime()
+        rt.ensureLoaded(model.absolutePath)
+        val hist = mutableListOf<ChatMessage>()
+        val qs = listOf("Introduce three advantages of Mixture-of-Experts in detail.", "thanks")
+        for ((i, q) in qs.withIndex()) {
+            hist.add(ChatMessage("user", q))
+            val sb = StringBuilder()
+            rt.submit("bs$i", hist.toList(), true, GenParams(temperature = 0f, maxNewTokens = if (i == 0) 160 else 16))
+                .toList().forEach { ev ->
+                    when (ev) {
+                        is GenEvent.TextPiece -> sb.append(ev.s)
+                        is GenEvent.ThinkPiece -> sb.append(ev.s)
+                        is GenEvent.Failed -> throw AssertionError(ev.diag)
+                        else -> Unit
+                    }
+                }
+            val content = if (mode == "native") (rt.lastAssistantNative() ?: sb.toString()) else sb.toString()
+            hist.add(ChatMessage("assistant", content))
+            android.util.Log.i("ReuseProbe", "mode=$mode round$i uiLen=${sb.length} sentLen=${content.length} m=${rt.lastMetrics.value?.let { "prompt=${it.promptTokens} reused=${it.prefix_reused_tokens} prefill_ms=${it.prefillMs}" }}")
+        }
+        rt.unloadActive()
+    }
+
+    // stack-smash regression: long ASCII pieces (>16 bytes) from byte-fallback must not crash
+    @Test
+    fun longAsciiPieceSurvival() = runBlocking {
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        val want = "edge0-8b.gguf"
+        val model = sequenceOf(
+            File(ctx.getDir("models", 0), want),
+            File(ctx.filesDir, "models/$want"),
+        ).firstOrNull { it.isFile } ?: return@runBlocking
+        val rt = LlamaRuntime()
+        rt.ensureLoaded(model.absolutePath)
+        var count = 0
+        rt.submit("t-ascii", listOf(
+            ChatMessage("user", "Repeat exactly, no thinking: AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB 123456789012345678901234567890 CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC"),
+        ), true, GenParams(temperature = 0f, maxNewTokens = 192)).toList()
+            .forEach { ev -> when (ev) {
+                is GenEvent.TextPiece -> count += ev.s.length
+                is GenEvent.ThinkPiece -> count += ev.s.length
+                is GenEvent.Failed -> throw AssertionError(ev.diag)
+                else -> Unit
+            } }
+        android.util.Log.i("LlamaRT", "ascii survival chars=$count")
+        assertTrue("no output", count > 20)
+        rt.unloadActive()
+    }
+
+    // 35B off-mode leak measurement: the empty-pair path is live (parity with the
+    // reference template); this case quantifies residual drift via thinkChars and
+    // driftFirst. Assertions guard gross shape only; drift is read from LeakProbe logs.
+    @Test
+    fun thinkOffLeak35b() = runBlocking {
+        val ctx = androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().targetContext
+        val want = "edge0-35b.gguf"
+        val model = sequenceOf(File(ctx.getDir("models", 0), want), File(ctx.filesDir, "models/$want"))
+            .firstOrNull { it.isFile } ?: return@runBlocking
+        val rt = LlamaRuntime()
+        rt.ensureLoaded(model.absolutePath)
+        val qs = listOf("What is the capital of France? One word.",
+                        "If you pick 3, 3, 8, 8 once each with +,-,*,/ to reach 24, give the expression directly.",
+                        "Name one prime number larger than 90 and smaller than 100.")
+        for ((i, q) in qs.withIndex()) {
+            val sb = StringBuilder(); var thinkChars = 0; var textChars = 0; var firstKind = "?"
+            rt.submit("leak$i", listOf(ChatMessage("user", q)), false,
+                      GenParams(temperature = 0f, maxNewTokens = 48)).toList().forEach { ev ->
+                when (ev) {
+                    is GenEvent.TextPiece -> { if (firstKind == "?") firstKind = "T"; textChars += ev.s.length; sb.append(ev.s) }
+                    is GenEvent.ThinkPiece -> { if (firstKind == "?") firstKind = "K"; thinkChars += ev.s.length; sb.append("<T>").append(ev.s) }
+                    is GenEvent.Failed -> throw AssertionError(ev.diag)
+                    else -> Unit
+                }
+            }
+            android.util.Log.i("LeakProbe", "q$i first=$firstKind think=$thinkChars text=$textChars out=" +
+                sb.toString().take(70).replace("\n", " "))
+            org.junit.Assert.assertTrue("leak probe empty q$i", textChars + thinkChars > 3)
+        }
+        rt.unloadActive()
+    }
+}
diff --git a/android/app/src/main/AndroidManifest.xml b/android/app/src/main/AndroidManifest.xml
new file mode 100644
index 0000000..803a6c4
--- /dev/null
+++ b/android/app/src/main/AndroidManifest.xml
@@ -0,0 +1,27 @@
+<?xml version="1.0" encoding="utf-8"?>
+<manifest xmlns:android="http://schemas.android.com/apk/res/android">
+    <!-- Single APK delivery: no INTERNET permission - the app is fully on-device;
+              models are staged into app-private dirs, no storage permission needed. -->
+    <application
+        android:name=".Edge0App"
+        android:allowBackup="false"
+        android:label="@string/app_name"
+        android:supportsRtl="true"
+        android:theme="@style/Theme.Edge0">
+        <!-- Optional vendor GPU library: on Android 12+ undeclared vendor libs fail to
+                  dlopen even when publicly listed; required=false keeps devices without it
+                  installable and the runtime probe degrades silently. -->
+        <uses-native-library android:name="libOpenCL.so" android:required="false" />
+        <uses-native-library android:name="libOpenCL_adreno.so" android:required="false" />
+        <activity
+            android:name=".MainActivity"
+            android:exported="true"
+            android:configChanges="orientation|screenSize|screenLayout|keyboardHidden|uiMode|density"
+            android:windowSoftInputMode="adjustResize">
+            <intent-filter>
+                <action android:name="android.intent.action.MAIN" />
+                <category android:name="android.intent.category.LAUNCHER" />
+            </intent-filter>
+        </activity>
+    </application>
+</manifest>
diff --git a/android/app/src/main/cpp/CMakeLists.txt b/android/app/src/main/cpp/CMakeLists.txt
new file mode 100644
index 0000000..9e66bb5
--- /dev/null
+++ b/android/app/src/main/cpp/CMakeLists.txt
@@ -0,0 +1,21 @@
+# app/src/main/cpp/CMakeLists.txt - native llama engine shell for the app.
+# Links the five shared libs staged at build-dl/llama-libs/arm64-v8a (built by
+# tools/llama/build_vendor_libs.sh from the pinned supply tree: vendor@7ab4ee7 with
+# the common+android patch bands replayed; see patches/llama.cpp/README).
+cmake_minimum_required(VERSION 3.22.1)
+project(edge0llama C)
+
+if(NOT DEFINED LLAMA_LIBS)
+    set(LLAMA_LIBS "${CMAKE_CURRENT_SOURCE_DIR}/../../../../build-dl/llama-libs/arm64-v8a")
+endif()
+
+add_library(edge0llama SHARED llama_chat.c)
+target_include_directories(edge0llama PRIVATE
+    "${LLAMA_LIBS}/../include"            # llama.h/ggml.h headers (synced at build time)
+    "${LLAMA_LIBS}")
+target_link_libraries(edge0llama
+    "${LLAMA_LIBS}/libllama.so"
+    "${LLAMA_LIBS}/libggml.so"
+    "${LLAMA_LIBS}/libggml-base.so"
+    "${LLAMA_LIBS}/libggml-cpu.so"
+    android log)
diff --git a/android/app/src/main/cpp/llama_chat.c b/android/app/src/main/cpp/llama_chat.c
new file mode 100644
index 0000000..a3ee9c1
--- /dev/null
+++ b/android/app/src/main/cpp/llama_chat.c
@@ -0,0 +1,650 @@
+// llama_chat.c - native llama engine shell for the Edge0 chat app.
+// Session model: static history of turns; every request renders the chat template
+// over the append-only transcript, tokenizes it, and lets KV prefix reuse turn the
+// full render into an incremental prefill. Sampling = classic token_data_array
+// chain (greedy when temperature <= 0). nativeInit carries expert-pool tier config
+// (the 35B tier requires E0_MMAP_NOPREFETCH + E0_NO_REPACK), nativeReset drops KV,
+// and each generate returns a structured metrics line for the Kotlin layer.
+#include <jni.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <time.h>
+#include <android/log.h>
+#include "llama.h"
+
+#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, "E0LLAMA", __VA_ARGS__)
+
+static struct llama_model *g_model = NULL;
+static struct llama_context *g_ctx = NULL;
+static struct llama_adapter_lora *g_lora = NULL;
+
+#define MAX_MSG 128
+static llama_chat_message g_hist[MAX_MSG];
+static char *g_strs[MAX_MSG * 2];
+static int g_nmsg = 0;
+
+#define MAX_TOKS 8192
+static llama_token g_toks[MAX_TOKS];
+static llama_token *g_msg_ids[MAX_MSG];      // assistant turns: stepped token ids (sidecar)
+static int g_msg_nids[MAX_MSG];
+static llama_token g_msg_eog[MAX_MSG];      // end-of-turn token for each history turn
+static llama_token g_out_ids[MAX_TOKS];      // ids stepped this turn (EOS excluded)
+static int g_out_n = 0;
+static llama_token g_out_eog = -1;
+static int g_nprocessed = 0;
+
+static volatile int g_cancel = 0;
+// pool counters forward decl (defined at the file tail; the metrics window inside nativeGenerate needs them)
+extern void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
+                           long long * bytes_loaded, size_t * cap_bytes, size_t * used_bytes);
+extern void moe_pool_extra(long long * stall_ns, long long * slots_used, long long * slots_total,
+                           long long * pf_used, long long * pf_mispredict);
+
+// Thinking control: the built-in BAILING2 renderer emits no think tags and ignores
+// jinja variables - thinking is model-native. The real switch is the template's
+// "detailed thinking on/off" instruction line in the system turn (training contract).
+static int g_thinking = 0;
+// The on direction rewrites in place with equal or shorter length, byte-matching the
+// reference jinja render; returns the corrected render length.
+// Both template families align with the reference (real-template) render:
+//   8B Bailing V3: off = instruction line at the end of the system turn (the built-in
+//     renderer self-injects it when no system turn exists); assistant heads get the open
+//     tag for on, an open+close empty pair for off.
+//   35B Qwen2/CHATML: off = empty pair after the generation head only - the instruction
+//     line is a foreign signal for this family (verified against the embedded template).
+
+static int32_t bailing_flip_on(char * r, int32_t len) {   // V3: off -> on, shrink 1 byte left
+    const char * needle = "detailed thinking off";
+    for (int32_t i = 0; i + 21 <= len; i++) {
+        if (memcmp(r + i, needle, 21) == 0) {
+            r[i + 19] = 'n';
+            memmove(r + i + 20, r + i + 21, (size_t) (len - i - 20));
+            return len - 1;
+        }
+    }
+    return len;
+}
+
+// Reference-template evidence: with enable_thinking=false the jinja template emits NO
+// instruction line - just the empty pair after the generation head. Carrying the
+// "detailed thinking off" text into a fabricated system turn is a foreign signal for
+// the qwen family (it invites the model to mimic reasoning prose in the body), so
+// off-mode here is a pure erase (also catches residue from user-provided system text).
+static int32_t qwen_pure_erase_directive(char * r, int32_t len) {
+    char needle[24];
+    memcpy(needle, "detailed thin", 13);
+    memcpy(needle + 13, "king off", 9);
+    needle[22] = 0;
+    for (;;) {
+        char * f = strstr(r, needle);
+        if (!f) return (int32_t) strlen(r);
+        int32_t st = (int32_t) (f - r), en = st + 22;
+        if (st > 0 && r[st - 1] == '\n') st--;
+        else if (r[en] == '\n') en++;
+        memmove(r + st, r + en, (size_t) (len - en + 1));
+        len -= (en - st);
+    }
+}
+
+static int g_tpl_imstart = -1;
+
+// Bailing fallback: when a custom system turn exists the built-in renderer skips the
+// instruction line entirely and the off signal is lost. The V3 contract inlines the
+// line at the end of the SYSTEM turn content, before its role_end. Insert there;
+// direction from g_thinking; only acts when no "detailed thinking" text exists.
+static void bail_ensure_directive(char * r, int32_t * plen, int32_t cap) {
+    if (g_tpl_imstart == 1) return;
+    if (strstr(r, "detailed thinking")) return;   // built-in already injected it, or the user wrote their own
+    char * sy = strstr(r, "<role>SYSTEM</role>");
+    if (!sy) return;
+    char * re = strstr(sy + 19, "<|" "role_end" "|>");
+    if (!re) return;
+    char frag[32];
+    int fl = snprintf(frag, sizeof(frag), "\ndetailed thinking %s", g_thinking ? "on" : "off");
+    if (fl <= 0 || *plen + fl + 1 > cap) return;
+    memmove(re + fl, re, (size_t) (*plen - (re - r) + 1));
+    memcpy(re, frag, (size_t) fl);
+    *plen += fl;
+}
+
+// Product decision for the 35B family: keep the template-conformant empty-pair form
+// and do not seed the continuation; sampling-side forcing measured unreliable.
+
+/* Literal of the built-in BAILING2 default head for history without a system turn:
+   segment-wise rendering would inject it between segments - strip it to keep one SYSTEM. */
+static const char FAKE_SYS_HEAD[] =
+    "<" "role" ">" "SYSTEM" "<" "/" "role" ">" "detailed thinking off" "<|" "role_end" "|>";
+static int32_t bail_strip_fake_head(char * r, int32_t len) {
+    const int fl = (int) sizeof(FAKE_SYS_HEAD) - 1;
+    if (len >= fl && memcmp(r, FAKE_SYS_HEAD, (size_t) fl) == 0) {
+        memmove(r, r + fl, (size_t) (len - fl + 1));
+        return len - fl;
+    }
+    return len;
+}
+
+// Golden-path head form (instruction line alone is a half-fix against the reference
+// render): after every bare assistant head emit newline + (on ? open : open+close pair).
+static void bailing_think_heads(char * r, int32_t * plen, int32_t cap) {
+    const char * head = "<role>ASSISTANT</role>";
+    const int hl = 22;
+    static const char tso[] = "\n" "<" "th" "ink" ">";
+    static const char tsof[] = "\n" "<" "th" "ink" ">" "</" "th" "ink" ">";
+    const char * inj = g_thinking ? tso : tsof;
+    const int il = (int) strlen(inj);   /* tso 8B / tsof 16B, self-adaptive */
+    int32_t pos = 0;
+    for (;;) {
+        char * h = strstr(r + pos, head);
+        if (!h) break;
+        int32_t at = (int32_t) (h - r) + hl;
+        // already annotated (tag follows the newline) - skip to prevent double injection
+        if (memcmp(r + at, inj, (size_t) il) == 0) { pos = at + il; continue; }
+        if (*plen + il + 1 > cap) return;
+        memmove(r + at + il, r + at, (size_t) (*plen - at + 1));
+        memcpy(r + at, inj, (size_t) il);
+        *plen += il;
+        pos = at + il;
+    }
+}
+
+static void qwen_inject_empty_pair(char * r, char const * head, int32_t cap) {
+    size_t l = strlen(r), hl = strlen(head);
+    if (l == 0 || hl == 0 || l + 32 > (size_t) cap) return;
+    static const char pair[] = "<th" "ink>\n\n</th" "ink>\n\n";
+    /* Training form = reference template: exactly one newline after the assistant
+       head, then the empty pair. The old append style added an extra newline (a
+       blank-paragraph continuation cue) and the model answered with mimic'd
+       "thinking process" prose - this overlay keeps it token-identical. */
+    if (l >= hl && memcmp(r + l - hl, head, hl) == 0 && !strstr(r + l - hl, "<th" "ink>")) {
+        memcpy(r + (l - 1 - hl), pair, sizeof(pair) - 1);   /* keep head+newline; overlay empty pair right after it */
+        r[(l - 1 - hl) + sizeof(pair) - 1] = 0;                 /* the overlay straddles the old terminator: rewrite it */
+    }
+}
+
+
+static char g_reply_tail[8192];      // tail 8KB of the last reply (cross-check for history sync; longer = TOOLONG)
+static size_t g_rlen = 0;
+static llama_token_data *g_td = NULL;
+static int g_reset_next = 0;
+
+// With byte-fallback vocabularies one token may be half a UTF-8 codepoint; feeding an
+// illegal continuation byte to NewStringUTF aborts under CheckJNI. Emit only complete
+// codepoints; carry the partial tail into the next piece.
+static char g_pend[8];
+static int g_pend_len = 0;
+
+static void emit_valid_utf8(JNIEnv *env, jobject sink, jmethodID invoke,
+                            const char *buf, int32_t len) {
+    char comb[16];
+    int n = g_pend_len + len;
+    if (n > (int)sizeof(comb)) n = (int)sizeof(comb);
+    memcpy(comb, g_pend, g_pend_len);
+    memcpy(comb + g_pend_len, buf, n - g_pend_len);   // clamp every memory op to the truncated n, never the raw len
+    int cut = n;
+    for (int k = 1; k <= 3 && k <= n; k++) {
+        unsigned char b = (unsigned char)comb[n - k];
+        if ((b & 0xC0) != 0x80) {
+            int expect = (b < 0x80) ? 1 : (b < 0xE0) ? 2 : (b < 0xF0) ? 3 : 4;
+            if (k < expect) cut = n - k;
+            break;
+        }
+    }
+    g_pend_len = n - cut;
+    if (g_pend_len) memcpy(g_pend, comb + cut, g_pend_len);
+    if (cut > 0) {
+        char safe[16]; memcpy(safe, comb, cut); safe[cut] = 0;
+        jstring js = (*env)->NewStringUTF(env, safe);
+        (*env)->CallObjectMethod(env, sink, invoke, js);
+        (*env)->DeleteLocalRef(env, js);
+    }
+}
+
+static void flush_pend(void) { g_pend_len = 0; }
+
+static long long now_ms(void) {
+    struct timespec ts; clock_gettime(CLOCK_MONOTONIC, &ts);
+    return (long long) ts.tv_sec * 1000 + ts.tv_nsec / 1000000;
+}
+
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeInit(
+        JNIEnv *env, jclass cls, jstring jmodel, jstring jlora, jstring jheads,
+        jint n_ctx, jint pool_mb, jint blob_mb, jint pred, jint io_threads) {
+    (void) cls;
+    // idempotent self-clear: model/LoRA/pool registry/history from a previous session
+    // must never bleed into a new load (an 8B adapter on a 35B graph aborts on shape asserts)
+    extern void moe_pool_reset_all(void);
+    if (g_model || g_ctx || g_lora) {
+        if (g_ctx) { llama_free(g_ctx); g_ctx = NULL; }
+        if (g_lora) { llama_adapter_lora_free(g_lora); g_lora = NULL; }
+        if (g_model) { llama_model_free(g_model); g_model = NULL; }
+        for (int i = 0; i < g_nmsg * 2; i++) { free(g_strs[i]); g_strs[i] = NULL; }
+        for (int i = 0; i < g_nmsg; i++) { free(g_msg_ids[i]); g_msg_ids[i] = NULL; g_msg_nids[i] = 0; }
+        g_nmsg = 0; g_nprocessed = 0;
+    }
+    moe_pool_reset_all();
+    const char *model = (*env)->GetStringUTFChars(env, jmodel, NULL);
+    const char *lora = (*env)->GetStringUTFChars(env, jlora, NULL);
+    const char *heads = (*env)->GetStringUTFChars(env, jheads, NULL);
+    char err[512];
+    // pool env must be set before model load (moe_pool_init reads env at register time);
+    // 8B passes 0/0 = zero perturbation.
+    setenv("E0_MMAP_NOPREFETCH", "1", 1);   // verified freeze-avoidance on 35B; harmless for 8B
+    setenv("E0_NO_REPACK", "1", 1);
+    { char b[32];
+      snprintf(b, sizeof(b), "%d", pool_mb);    setenv("E0_POOL_MB", b, 1);
+      snprintf(b, sizeof(b), "%d", blob_mb);    setenv("E0_BLOB_MB", b, 1);
+      snprintf(b, sizeof(b), "%d", io_threads); setenv("E0_IO_THREADS", b, 1);
+      setenv("E0_MEM_FLOOR_MB", "2000", 1); }
+    if (heads && heads[0]) setenv("E0_HEADS", heads, 1);
+    if (blob_mb > 0) setenv("E0_KEEPWARM_MB", "512", 1);   // 35B tier: turn-end keepwarm refill (moot without a pool)
+    if (pred) setenv("E0_PREROUTER", "1", 1); else unsetenv("E0_PREROUTER");
+    llama_backend_init();
+    struct llama_model_params mp = llama_model_default_params();
+    mp.n_gpu_layers = 0;
+    g_model = llama_model_load_from_file(model, mp);
+    if (!g_model) {
+        snprintf(err, sizeof(err), "ERR load: %s", model);
+        goto ret;
+    }
+    if (lora && lora[0]) {
+        g_lora = llama_adapter_lora_init(g_model, lora);
+        if (!g_lora) { snprintf(err, sizeof(err), "ERR lora init: %s", lora); goto ret; }
+    }
+    {
+        struct llama_context_params cp = llama_context_default_params();
+        cp.n_ctx = (uint32_t) n_ctx;
+        cp.n_batch = 512;
+        g_ctx = llama_init_from_model(g_model, cp);
+        if (!g_ctx) { snprintf(err, sizeof(err), "ERR ctx init"); goto ret; }
+        if (g_lora) {
+        float scale = 1.0f;   // canonical form: file alpha 32 + explicit scale 1.0
+            if (llama_set_adapters_lora(g_ctx, &g_lora, 1, &scale) != 0) {
+                snprintf(err, sizeof(err), "ERR lora set"); goto ret;
+            }
+        }
+    }
+    g_reset_next = 1;
+    snprintf(err, sizeof(err), "OK %s pool=%d blob=%d pred=%d", llama_version(), pool_mb, blob_mb, pred);
+ret:
+    (*env)->ReleaseStringUTFChars(env, jmodel, model);
+    (*env)->ReleaseStringUTFChars(env, jlora, lora);
+    (*env)->ReleaseStringUTFChars(env, jheads, heads);
+    return (*env)->NewStringUTF(env, err);
+}
+
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeReset(
+        JNIEnv *env, jclass cls) {
+    (void) cls;
+    // session reset: clear history and force full KV replay (thread switch / history fork)
+    for (int i = 0; i < g_nmsg * 2; i++) { free(g_strs[i]); g_strs[i] = NULL; }
+    for (int i = 0; i < g_nmsg; i++) { free(g_msg_ids[i]); g_msg_ids[i] = NULL; g_msg_nids[i] = 0; }
+    g_nmsg = 0;
+    g_reset_next = 1;
+    return (*env)->NewStringUTF(env, "OK");
+}
+
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeAddRole(
+        JNIEnv *env, jclass cls, jstring jrole, jstring jtext) {
+    (void) cls;
+    const char *role = (*env)->GetStringUTFChars(env, jrole, NULL);
+    const char *text = (*env)->GetStringUTFChars(env, jtext, NULL);
+    int full = g_nmsg >= MAX_MSG;
+    if (!full) {
+        char *r = strdup(role); char *t = strdup(text);
+        g_strs[g_nmsg * 2] = r; g_strs[g_nmsg * 2 + 1] = t;
+        g_hist[g_nmsg].role = r; g_hist[g_nmsg].content = t;
+        g_nmsg++;
+    }
+    (*env)->ReleaseStringUTFChars(env, jrole, role);
+    (*env)->ReleaseStringUTFChars(env, jtext, text);
+    return (*env)->NewStringUTF(env, full ? "FULL" : "OK");
+}
+
+// assistant turn enters history via sidecar: ids big-endian 4B per element; text render is skipped (stitched in nativeGenerate)
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeAddRoleIds(
+        JNIEnv *env, jclass cls, jstring jrole, jbyteArray jids, jint eog) {
+    (void) cls;
+    const char *role = (*env)->GetStringUTFChars(env, jrole, NULL);
+    const int nb = jids ? (*env)->GetArrayLength(env, jids) : 0;
+    int full = (g_nmsg >= MAX_MSG) || (nb % 4) != 0;
+    if (!full) {
+        llama_token *ids = NULL;
+        if (nb > 0) {
+            jbyte *bd = (*env)->GetByteArrayElements(env, jids, NULL);
+            ids = (llama_token *) malloc(sizeof(llama_token) * (nb / 4));
+            for (int i = 0; i < nb / 4; i++) {
+                unsigned char *q = (unsigned char *) (bd + i * 4);
+                ids[i] = (llama_token) ((q[0] << 24) | (q[1] << 16) | (q[2] << 8) | q[3]);
+            }
+            (*env)->ReleaseByteArrayElements(env, jids, bd, JNI_ABORT);
+        }
+        char *r = strdup(role);
+        g_strs[g_nmsg * 2] = r; g_strs[g_nmsg * 2 + 1] = NULL;
+        g_hist[g_nmsg].role = r; g_hist[g_nmsg].content = "";
+        g_msg_ids[g_nmsg] = ids; g_msg_nids[g_nmsg] = nb / 4; g_msg_eog[g_nmsg] = (llama_token) eog;
+        g_nmsg++;
+    }
+    (*env)->ReleaseStringUTFChars(env, jrole, role);
+    return (*env)->NewStringUTF(env, full ? "FULL" : "OK");
+}
+
+// ids stepped this turn (big-endian 4B) plus the end-of-turn token
+JNIEXPORT jbyteArray JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeLastAssistantIds(JNIEnv *env, jclass cls) {
+    (void) cls;
+    jbyteArray out = (*env)->NewByteArray(env, g_out_n * 4);
+    if (out && g_out_n > 0) {
+        jbyte *bd = (*env)->GetByteArrayElements(env, out, NULL);
+        for (int i = 0; i < g_out_n; i++) {
+            unsigned char *q = (unsigned char *) (bd + i * 4);
+            q[0] = (unsigned char) (g_out_ids[i] >> 24); q[1] = (unsigned char) (g_out_ids[i] >> 16);
+            q[2] = (unsigned char) (g_out_ids[i] >> 8);  q[3] = (unsigned char) g_out_ids[i];
+        }
+        (*env)->ReleaseByteArrayElements(env, out, bd, 0);
+    }
+    return out;
+}
+JNIEXPORT jint JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeLastEog(JNIEnv *env, jclass cls) {
+    (void) env; (void) cls; return (jint) g_out_eog;
+}
+
+#include <pthread.h>
+static void * kw_thread(void * arg) { (void) arg; extern void moe_pool_keepwarm(int); moe_pool_keepwarm(0); return NULL; }
+
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeGenerate(
+        JNIEnv *env, jclass cls, jobject sink, jint max_tokens, jfloat temp, jboolean thinking_on) {
+    // Toggling thinking changes the rendered instruction line, which the UI-text
+    // history comparison cannot see - trusting positions would keep dirty KV context.
+    // Authoritative checkpoint: on a thinking change with live context, force a full rebuild.
+    if (g_ctx && g_thinking != (thinking_on ? 1 : 0) && g_nprocessed > 0) g_reset_next = 1;
+    g_thinking = thinking_on ? 1 : 0;
+    char out[512];
+    out[0] = 0;
+    if (!g_ctx) { snprintf(out, sizeof(out), "ERR not init"); return (*env)->NewStringUTF(env, out); }
+        {
+        const char *t0 = llama_model_chat_template(g_model, NULL);
+        g_tpl_imstart = (t0 && strstr(t0, "<|im_start|>assistant")) ? 1 : 0;
+    }
+    const struct llama_vocab *vocab = llama_model_get_vocab(g_model);
+    const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+
+    static char render[65536];
+    const char *tmpl = llama_model_chat_template(g_model, NULL); // BAILING2 -> built-in renderer
+    int has_ids = 0;
+    for (int i = 0; i < g_nmsg; i++) if (g_msg_nids[i] > 0) { has_ids = 1; break; }
+    int32_t n_tok = 0;
+    if (!has_ids) {
+        int32_t need = llama_chat_apply_template(tmpl, g_hist, (size_t) g_nmsg, true, render, sizeof(render));
+        if (need <= 0 || (size_t) need >= sizeof(render)) {
+            snprintf(out, sizeof(out), "ERR template need=%d", need);
+            return (*env)->NewStringUTF(env, out);
+        }
+        if (!g_thinking) {
+            if (g_tpl_imstart == 1) {
+                // qwen off = pure empty pair (reference-template verified): no instruction line, no fabricated system turn
+                need = qwen_pure_erase_directive(render, need);
+            }
+        } else {
+            if (g_tpl_imstart != 1) need = bailing_flip_on(render, need);
+        }
+        bail_ensure_directive(render, &need, (int32_t) sizeof(render));
+        if (g_tpl_imstart != 1) bailing_think_heads(render, &need, (int32_t) sizeof(render));
+        qwen_inject_empty_pair(render, "<|im_start|>assistant\n", (int32_t) sizeof(render));
+        n_tok = llama_tokenize(vocab, render, need, g_toks, MAX_TOKS, false, true);
+        if (n_tok <= 0) { snprintf(out, sizeof(out), "ERR tokenize %d", n_tok); return (*env)->NewStringUTF(env, out); }
+    } else {
+        // sidecar stitching: text segments render+tokenize, id segments paste verbatim
+        // (kills BPE re-encode boundary drift). Every text segment carries the generation
+        // prompt; id turns = stepped body + that turn's end-of-generation token.
+        int32_t w = 0; int ii = 0; int err = 0;
+        while (ii < g_nmsg) {
+            if (g_msg_nids[ii] > 0) {
+                if (w + g_msg_nids[ii] + 1 >= MAX_TOKS - 8) { err = 2; break; }
+                memcpy(&g_toks[w], g_msg_ids[ii], sizeof(llama_token) * (size_t) g_msg_nids[ii]);
+                w += g_msg_nids[ii];
+                if (g_msg_eog[ii] >= 0) g_toks[w++] = g_msg_eog[ii];
+                ii++;
+                continue;
+            }
+            int j = ii; while (j < g_nmsg && g_msg_nids[j] == 0) j++;
+            int32_t need = llama_chat_apply_template(tmpl, &g_hist[ii], (size_t) (j - ii), true, render, sizeof(render));
+            if (need <= 0 || (size_t) need >= sizeof(render)) { err = 1; break; }
+            /* fake SYSTEM head stripping: the built-in renderer injects its default
+               system block into every non-first segment start; strip or the prompt
+               grows a second SYSTEM block and the model loses its identity */
+            if (g_tpl_imstart != 1 && ii > 0) need = bail_strip_fake_head(render, need);
+            if (!g_thinking) {
+                if (g_tpl_imstart == 1) {
+                    // qwen off = pure empty pair (reference-template verified): no instruction line, no fabricated system turn
+                    need = qwen_pure_erase_directive(render, need);
+                }
+            } else {
+                if (g_tpl_imstart != 1) need = bailing_flip_on(render, need);
+            }
+            // run every segment unconditionally: gating this on the thinking branch or
+            // the last segment lost the instruction line from turn two onward (the
+            // empty-answer bug). Families without a SYSTEM in the segment self-check to a no-op.
+            if (g_tpl_imstart != 1) bail_ensure_directive(render, &need, (int32_t) sizeof(render));
+            if (g_tpl_imstart != 1) bailing_think_heads(render, &need, (int32_t) sizeof(render));
+            if (g_tpl_imstart == 1 && j >= g_nmsg) qwen_inject_empty_pair(render, "<|im_start>|assistant\n", (int32_t) sizeof(render));
+            int32_t room = MAX_TOKS - 8 - w;
+            int32_t tn = llama_tokenize(vocab, render, need, g_toks + w, room, false, true);
+            if (tn <= 0) { err = 1; break; }
+            w += tn; ii = j;
+        }
+        if (err) { snprintf(out, sizeof(out), "ERR sidecar concat %d", err); return (*env)->NewStringUTF(env, out); }
+        n_tok = w;
+    }
+    if (n_tok >= MAX_TOKS - 8) { snprintf(out, sizeof(out), "ERR ctx too long"); return (*env)->NewStringUTF(env, out); }
+
+    {   // render forensics: single-line dump of the tail (newlines escaped) for empty-pair / migration checks
+        char dbg[512]; int di = 0;
+        int rl = (int) strlen(render);
+        int tl = rl > 110 ? 110 : rl;
+        const char * t = render + rl - tl;
+        for (int z = 0; z < tl && di < 380; z++) {
+            if (t[z] == '\n') { dbg[di++] = '\\'; dbg[di++] = 'n'; }
+            else dbg[di++] = t[z];
+        }
+        dbg[di] = 0;
+        LOGI("REND tail=[%s]", dbg);
+        {   /* head forensics: SYSTEM segment identity + instruction-line shape */
+            char dh[260]; int dj = 0;
+            for (int z = 0; z < (int) strlen(render) && z < 130 && dj < 250; z++) {
+                char ch = render[z];
+                if (ch == '\n') { dh[dj++] = '\\'; dh[dj++] = 'n'; }
+                else dh[dj++] = ch;
+            }
+            dh[dj] = 0;
+            LOGI("REND head=[%s]", dh);
+        }
+    }
+    int start = g_reset_next ? 0 : g_nprocessed;
+    if (start > n_tok) start = 0;
+    if (g_reset_next) {
+        llama_memory_seq_rm(llama_get_memory(g_ctx), 0, -1, -1);
+        g_reset_next = 0;
+    }
+
+    long long ph0=0, pm0=0, pl0=0, pe0=0, pb0=0, pst0=0;
+    size_t pcap0=0, pused0=0;
+    moe_pool_stats(&ph0, &pm0, &pl0, &pe0, (long long*)&pb0, &pcap0, &pused0);
+    { long long su,st,pu,pm; moe_pool_extra(&pst0,&su,&st,&pu,&pm); }
+    jmethodID invoke = (*env)->GetMethodID(env, (*env)->GetObjectClass(env, sink),
+                                           "invoke", "(Ljava/lang/Object;)Ljava/lang/Object;");
+    long long prefill_ms;
+    {
+        long long t0 = now_ms();
+        // small incremental batches decode token-by-token: MoE gather cost does not
+        // amortize on tiny batches (measured 7x slower than decode); <=32 new tokens
+        // take the decode-like path, long batches keep the 512 window.
+        const int32_t inc = n_tok - start;
+        const int32_t chunk_max = (inc > 0 && inc <= 32) ? 1 : 512;
+        for (int i = start; i < n_tok; i += chunk_max) {
+            int chunk = (n_tok - i < chunk_max) ? n_tok - i : chunk_max;
+            struct llama_batch b = llama_batch_get_one(g_toks + i, chunk);
+            if (llama_decode(g_ctx, b) != 0) { snprintf(out, sizeof(out), "ERR prefill@%d", i); return (*env)->NewStringUTF(env, out); }
+            if (g_cancel) { g_cancel = 0; snprintf(out, sizeof(out), "CANCELLED"); return (*env)->NewStringUTF(env, out); }
+        }
+        prefill_ms = now_ms() - t0;
+        { long long h, m, l, e, b; long long su, stq, pu, pm; size_t cap, used; long long st1;
+          moe_pool_stats(&h, &m, &l, &e, &b, &cap, &used);
+          moe_pool_extra(&st1, &su, &stq, &pu, &pm);
+          char dbg[256];
+          snprintf(dbg, sizeof(dbg),
+                   "PRE fill hit=%lld miss=%lld evict=%lld stall_ms=%lld loads=%lld",
+                   h - ph0, m - pm0, e - pe0, (st1 - pst0) / 1000000, l - pl0);
+          LOGI("%s", dbg); }
+    }
+    int reused = start;
+    g_nprocessed = n_tok;
+
+    {   // vocab width changes across models: the token buffer must realloc (8B-sized buffer fed a 35B vocab = OOB write)
+        static size_t g_td_cap = 0;
+        if (g_td_cap < (size_t) n_vocab) {
+            free(g_td);
+            g_td = (llama_token_data *) malloc(sizeof(llama_token_data) * (size_t) n_vocab);
+            g_td_cap = g_td ? (size_t) n_vocab : 0;
+        }
+    }
+    struct llama_sampler *smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());
+    if (temp <= 0.0f) llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
+    else {
+        llama_sampler_chain_add(smpl, llama_sampler_init_temp(temp));
+        llama_sampler_chain_add(smpl, llama_sampler_init_top_k(40));
+        llama_sampler_chain_add(smpl, llama_sampler_init_dist(0xC0FFEEu));
+    }
+    char reply[32768]; size_t rlen = 0; reply[0] = 0;
+    g_rlen = 0;
+    flush_pend();
+    char piece[256];
+    long long t0 = now_ms();
+    int gen = 0;
+    long long first_ms = 0;
+    g_out_n = 0; g_out_eog = -1;
+    double tps = 0.0;
+    llama_token next = -1;
+    while (gen < max_tokens && !g_cancel) {
+        const float *lg = llama_get_logits_ith(g_ctx, -1);
+        if (!lg) { snprintf(out, sizeof(out), "ERR logits@%d", gen); break; }
+        for (int32_t j = 0; j < n_vocab; j++) { g_td[j].id = j; g_td[j].logit = lg[j]; g_td[j].p = 1.0f; }
+        llama_token_data_array cur = { .data = g_td, .size = (size_t) n_vocab, .selected = -1, .sorted = false };
+        llama_sampler_apply(smpl, &cur);
+        next = cur.data[cur.selected].id;
+        if (llama_vocab_is_eog(vocab, next)) { g_out_eog = next; break; }
+        int32_t plen = llama_token_to_piece(vocab, next, piece, sizeof(piece) - 1, 0, false);
+        if (plen > 0) {
+            piece[plen] = 0;
+            if (rlen + plen < sizeof(reply)) { memcpy(reply + rlen, piece, plen); rlen += plen; }
+            if (g_rlen + plen < sizeof(g_reply_tail)) { memcpy(g_reply_tail + g_rlen, piece, plen); g_rlen += plen; }
+            else g_rlen = sizeof(g_reply_tail) + 1;
+            emit_valid_utf8(env, sink, invoke, piece, plen);
+        }
+        if (g_out_n < MAX_TOKS) g_out_ids[g_out_n++] = next;
+        gen++;
+        if (gen == 1) first_ms = now_ms() - t0;
+        struct llama_batch b = llama_batch_get_one(&next, 1);
+        if (llama_decode(g_ctx, b) != 0) { snprintf(out, sizeof(out), "ERR decode@%d", gen); break; }
+        g_nprocessed++;
+        tps = gen * 1000.0 / (double) (now_ms() - t0 + 1);
+    }
+    flush_pend();
+    while (rlen > 0 && (((unsigned char)reply[rlen - 1] & 0xC0) == 0x80)) rlen--;
+    if (rlen > 0) { unsigned char b = (unsigned char)reply[rlen - 1];
+        int expect = (b < 0x80) ? 1 : (b < 0xE0) ? 2 : (b < 0xF0) ? 3 : 4;
+        if (expect > 1) rlen--; }
+    reply[rlen] = 0;
+    if (g_cancel) { g_cancel = 0;
+        snprintf(out, sizeof(out), "CANCELLED gen=%d tps=%.2f prefill_ms=%lld first_ms=%lld prompt=%d reused=%d",
+                 gen, tps, prefill_ms, first_ms, n_tok, reused); }
+    else snprintf(out, sizeof(out), "OK gen=%d tps=%.2f prefill_ms=%lld first_ms=%lld prompt=%d reused=%d",
+                  gen, tps, prefill_ms, first_ms, n_tok, reused);
+    LOGI("%s", out);
+
+    if (gen > 0 && getenv("E0_KEEPWARM_MB")) { pthread_t kt; if (pthread_create(&kt, NULL, kw_thread, NULL) == 0) pthread_detach(kt); }
+    {   // assistant turn recorded with stepped ids (sidecar); the text reply stays in logs only
+        char *ar = strdup("assistant");
+        g_strs[g_nmsg * 2] = ar; g_strs[g_nmsg * 2 + 1] = NULL;
+        if (g_nmsg < MAX_MSG) {
+            llama_token *ids = NULL;
+            if (g_out_n > 0) { ids = (llama_token *) malloc(sizeof(llama_token) * (size_t) g_out_n); memcpy(ids, g_out_ids, sizeof(llama_token) * (size_t) g_out_n); }
+            g_hist[g_nmsg].role = ar; g_hist[g_nmsg].content = "";
+            g_msg_ids[g_nmsg] = ids; g_msg_nids[g_nmsg] = ids ? g_out_n : 0; g_msg_eog[g_nmsg] = g_out_eog;
+            g_nmsg++;
+        }
+    }
+    (void) reply;
+    llama_sampler_free(smpl);
+    return (*env)->NewStringUTF(env, out);
+}
+
+// pool counters + process RSS export (the metrics line under each reply: hit/stall/peak-rss).
+// moe_pool_* symbols live in libggml-cpu.so (linked explicitly at build time).
+extern void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
+                           long long * bytes_loaded, size_t * cap_bytes, size_t * used_bytes);
+extern void moe_pool_extra(long long * stall_ns, long long * slots_used, long long * slots_total,
+                           long long * pf_used, long long * pf_mispredict);
+
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeStats(JNIEnv *env, jclass cls) {
+    (void) cls;
+    long long h=0, m=0, l=0, e=0, bl=0, st=0, su=0, stot=0, pu=0, pm=0;
+    size_t cap=0, used=0;
+    moe_pool_stats(&h, &m, &l, &e, &bl, &cap, &used);
+    moe_pool_extra(&st, &su, &stot, &pu, &pm);
+    long vmhwm = 0, vmrss = 0;
+    FILE * f = fopen("/proc/self/status", "r");
+    if (f) {
+        char ln[256];
+        while (fgets(ln, sizeof(ln), f)) {
+            if (!strncmp(ln, "VmHWM:", 6)) vmhwm = strtol(ln + 6, NULL, 10);
+            else if (!strncmp(ln, "VmRSS:", 6)) vmrss = strtol(ln + 6, NULL, 10);
+        }
+        fclose(f);
+    }
+    char buf[512];
+    snprintf(buf, sizeof(buf),
+        "hits=%lld misses=%lld loads=%lld evicts=%lld flash=%lld resident=%lld cap=%lld\n" 
+        "stall_ns=%lld slots=%lld/%lld pf_used=%lld pf_miss=%lld vmhwm_kb=%ld vmrss_kb=%ld",
+        h, m, l, e, bl, (long long)used, (long long)cap, st, su, stot, pu, pm, vmhwm, vmrss);
+    return (*env)->NewStringUTF(env, buf);
+}
+
+// tail 8KB of the last assistant reply in native form (when the UI shows equivalent text, replaying it keeps full prefix reuse; longer = TOOLONG)
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeLastAssistant(JNIEnv *env, jclass cls) {
+    (void) cls;
+    if (g_rlen > sizeof(g_reply_tail)) return (*env)->NewStringUTF(env, "TOOLONG");
+    char tmp[8200];
+    size_t n = g_rlen; if (n > sizeof(g_reply_tail) - 1) n = sizeof(g_reply_tail) - 1;
+    memcpy(tmp, g_reply_tail, n); tmp[n] = 0;
+    return (*env)->NewStringUTF(env, tmp);
+}
+
+JNIEXPORT void JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeCancel(JNIEnv *env, jclass cls) {
+    (void) env; (void) cls; g_cancel = 1;
+}
+
+// unload/free: real teardown. Order is law: sampler first, then adapter before model
+// (freeing the model under an attached adapter crashed at 452s once) - plus all-static reset.
+JNIEXPORT jstring JNICALL
+Java_dev_edge0_runtime_app_runtime_LlamaNative_nativeFree(JNIEnv *env, jclass cls) {
+    (void) cls;
+    if (g_ctx) { llama_free(g_ctx); g_ctx = NULL; }
+    if (g_lora) { llama_adapter_lora_free(g_lora); g_lora = NULL; }
+    if (g_model) { llama_model_free(g_model); g_model = NULL; }
+    for (int i = 0; i < g_nmsg * 2; i++) { free(g_strs[i]); g_strs[i] = NULL; }
+    for (int i = 0; i < g_nmsg; i++) { free(g_msg_ids[i]); g_msg_ids[i] = NULL; g_msg_nids[i] = 0; }
+    g_nmsg = 0; g_nprocessed = 0; g_reset_next = 0;
+    { extern void moe_pool_reset_all(void); moe_pool_reset_all(); }
+    return (*env)->NewStringUTF(env, "OK");
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/Edge0App.kt b/android/app/src/main/java/dev/edge0/runtime/app/Edge0App.kt
new file mode 100644
index 0000000..7e57d95
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/Edge0App.kt
@@ -0,0 +1,79 @@
+// Edge0App.kt - process-level container (manual DI, no framework) plus:
+// cold-start cleanup of orphaned streaming rows, and onTrimMemory -> engine pressure relief
+// (logcat probe tag=Edge0Trim).
+package dev.edge0.runtime.app
+
+import android.app.Application
+import android.content.ComponentCallbacks2
+import android.content.res.Configuration
+import android.util.Log
+import dev.edge0.runtime.app.data.AppDatabase
+import dev.edge0.runtime.app.data.ChatRepository
+import dev.edge0.runtime.app.data.RoomChatRepository
+import dev.edge0.runtime.app.data.SettingsStore
+import dev.edge0.runtime.app.runtime.LlamaRuntime
+import dev.edge0.runtime.app.runtime.Runtime
+import dev.edge0.runtime.app.runtime.ThermalSource
+import java.io.File
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.launch
+
+class AppContainer(private val app: Application) {
+    private val llamaRuntime by lazy { LlamaRuntime() }
+    val runtime: Runtime by lazy { llamaRuntime }
+    val settings = SettingsStore(app)
+    val db by lazy { AppDatabase.create(app) }
+    val repository: ChatRepository by lazy { RoomChatRepository(db) }
+    val thermal by lazy { ThermalSource(app) }
+    /** Model import roots: internal files/models (primary - on some OEM builds
+     *  shell-pushed files are invisible to the app uid, so device-internal copy is the
+     *  reliable channel) plus external files/models; both roots are scanned. */
+    val modelRoots: List<File?>
+        get() = listOf(File(app.filesDir, "models"), app.getExternalFilesDir("models"))
+}
+
+class Edge0App : Application() {
+    lateinit var container: AppContainer
+        private set
+
+    override fun onCreate() {
+        super.onCreate()
+        container = AppContainer(this)
+        // cold start: mark rows left in streaming by a killed process as interrupted (visible, not lost)
+        CoroutineScope(SupervisorJob() + Dispatchers.IO).launch {
+            runCatching { container.repository.markStreamingInterrupted() }
+        }
+        // ComponentCallbacks2 registered directly on the Application - lifecycle observers
+        // (2.8) do not carry onTrimMemory; this is the official hook
+        registerComponentCallbacks(TrimBridge(container))
+    }
+}
+
+/** Pressure mapping: RUNNING_LOW/BACKGROUND levels -> trim(0.5); COMPLETE -> cancel then unload. */
+class TrimBridge(private val container: AppContainer) : ComponentCallbacks2 {
+    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Main.immediate)
+
+    override fun onConfigurationChanged(newConfig: Configuration) {}
+    @Deprecated("compat") override fun onLowMemory() {}
+
+    override fun onTrimMemory(level: Int) {
+        // level bands: RUNNING_MODERATE=5/LOW=10/CRITICAL=15; UI_HIDDEN=16;
+        // BACKGROUND=20, MODERATE=40, COMPLETE=80
+        when {
+            level == android.content.ComponentCallbacks2.TRIM_MEMORY_COMPLETE -> {
+                Log.i(TAG, "onTrimMemory level=$level action=cancel+unload")
+                container.runtime.cancelActive()
+                scope.launch { runCatching { container.runtime.unloadActive() } }
+            }
+            level in 5..19 || level in 20..79 -> {
+                Log.i(TAG, "onTrimMemory level=$level action=trim(0.5)")
+                scope.launch { container.runtime.trim(0.5f) }
+            }
+            else -> Log.i(TAG, "onTrimMemory level=$level action=none")
+        }
+    }
+
+    companion object { const val TAG = "Edge0Trim" }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/MainActivity.kt b/android/app/src/main/java/dev/edge0/runtime/app/MainActivity.kt
new file mode 100644
index 0000000..4c0e6bd
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/MainActivity.kt
@@ -0,0 +1,27 @@
+// MainActivity.kt - single-Activity host (edge-to-edge; Compose-drawn; theme follows settings).
+package dev.edge0.runtime.app
+
+import android.os.Bundle
+import androidx.activity.ComponentActivity
+import androidx.activity.compose.setContent
+import androidx.activity.enableEdgeToEdge
+import androidx.compose.runtime.getValue
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import dev.edge0.runtime.app.data.AppSettings
+import dev.edge0.runtime.app.ui.AppRoot
+import dev.edge0.runtime.app.ui.theme.Edge0Theme
+
+class MainActivity : ComponentActivity() {
+    override fun onCreate(savedInstanceState: Bundle?) {
+        enableEdgeToEdge()
+        super.onCreate(savedInstanceState)
+        val container = (application as Edge0App).container
+        setContent {
+            val settings by container.settings.flow
+                .collectAsStateWithLifecycle(initialValue = AppSettings())
+            Edge0Theme(mode = settings.themeMode, fontScale = settings.fontScale) {
+                AppRoot()
+            }
+        }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/data/AppDatabase.kt b/android/app/src/main/java/dev/edge0/runtime/app/data/AppDatabase.kt
new file mode 100644
index 0000000..3395588
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/data/AppDatabase.kt
@@ -0,0 +1,38 @@
+// data/AppDatabase.kt
+package dev.edge0.runtime.app.data
+
+import android.content.Context
+import androidx.room.Database
+import androidx.room.Room
+import androidx.room.RoomDatabase
+import androidx.room.migration.Migration
+import androidx.sqlite.db.SupportSQLiteDatabase
+
+// v2: assistant token-id sidecar - messages.genTokens BLOB (additive ALTER).
+private val MIGRATION_1_2 = object : Migration(1, 2) {
+    override fun migrate(db: SupportSQLiteDatabase) {
+        db.execSQL("ALTER TABLE messages ADD COLUMN genTokens BLOB NULL")
+    }
+}
+
+// v3: per-reply inline stats line - ttftMs/prefillTokS/memBytes columns.
+private val MIGRATION_2_3 = object : Migration(2, 3) {
+    override fun migrate(db: SupportSQLiteDatabase) {
+        db.execSQL("ALTER TABLE messages ADD COLUMN ttftMs INTEGER NULL")
+        db.execSQL("ALTER TABLE messages ADD COLUMN prefillTokS REAL NULL")
+        db.execSQL("ALTER TABLE messages ADD COLUMN memBytes INTEGER NULL")
+    }
+}
+
+@Database(entities = [ThreadEntity::class, MessageEntity::class], version = 3,
+          exportSchema = true)
+abstract class AppDatabase : RoomDatabase() {
+    abstract fun threads(): ThreadDao
+    abstract fun messages(): MessageDao
+
+    companion object {
+        fun create(context: Context): AppDatabase =
+            Room.databaseBuilder(context, AppDatabase::class.java, "edge0-chat.db")
+                .addMigrations(MIGRATION_1_2, MIGRATION_2_3).build()
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/data/ChatRepository.kt b/android/app/src/main/java/dev/edge0/runtime/app/data/ChatRepository.kt
new file mode 100644
index 0000000..af1cf62
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/data/ChatRepository.kt
@@ -0,0 +1,139 @@
+// data/ChatRepository.kt - seam between UI and persistence. A fake implementation
+// can be injected so UI walkthroughs run without the engine or disk.
+package dev.edge0.runtime.app.data
+
+import dev.edge0.runtime.app.runtime.StoredMessage
+import dev.edge0.runtime.app.runtime.decodeIds
+import java.util.UUID
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.map
+
+/** Message as the UI sees it (includes persistence status). */
+data class MessageView(
+    val id: String,
+    val role: String,
+    val content: String,
+    val thinking: String?,
+    val thinkingMs: Long?,
+    val status: String,
+    val createdAt: Long,
+    val newTokens: Int?,
+    val decodeTokS: Double?,
+    val ttftMs: Long? = null,
+    val prefillTokS: Double? = null,
+    val memBytes: Long? = null,
+)
+
+data class ThreadSummary(val id: String, val title: String?, val updatedAt: Long,
+                         val modelId: String)
+
+interface ChatRepository {
+    fun observeThreads(): Flow<List<ThreadSummary>>
+    fun observeMessages(threadId: String): Flow<List<MessageView>>
+    /** Creates the thread only with the first message (no empty entries). Returns threadId. */
+    suspend fun startThreadWithUserMessage(modelId: String, userText: String): String
+    /** Appends a user message to an existing thread and bumps updatedAt. */
+    suspend fun addUserMessage(threadId: String, userText: String)
+    /** Inserts a streaming placeholder assistant message, returns its id (finalize overwrites). */
+    suspend fun beginAssistant(threadId: String, modelId: String): String
+    suspend fun finalizeAssistant(id: String, content: String, thinking: String?,
+                                  thinkingMs: Long?, status: String,
+                                  promptTokens: Int?, newTokens: Int?, decodeTokS: Double?,
+                                  genTokens: ByteArray? = null,
+                                  ttftMs: Long? = null, prefillTokS: Double? = null,
+                                  memBytes: Long? = null)
+    suspend fun renameThread(id: String, title: String)
+    suspend fun deleteThread(id: String)
+    /** Editing a user message truncates everything after it (incl. old answer). */
+    suspend fun editUserMessage(id: String, newText: String)
+    suspend fun deleteMessage(id: String)
+    suspend fun markStreamingInterrupted()
+    /** Thread history (old to new, without the trailing user turn the caller appends). */
+    suspend fun history(threadId: String): List<StoredMessage>
+    /** Full view with id/timestamps (used by regenerate to locate the last assistant row). */
+    suspend fun historyWithIds(threadId: String): List<MessageView>
+    /** Deletes all messages at/after the given timestamp (regenerate truncation). */
+    suspend fun deleteFrom(threadId: String, fromTs: Long)
+}
+
+class RoomChatRepository(private val db: AppDatabase) : ChatRepository {
+    override fun observeThreads(): Flow<List<ThreadSummary>> =
+        db.threads().observeAll().map { l -> l.map { ThreadSummary(it.id, it.title, it.updatedAt, it.modelId) } }
+
+    override fun observeMessages(threadId: String): Flow<List<MessageView>> =
+        db.messages().observeFor(threadId).map { l ->
+            l.map { MessageView(it.id, it.role, it.content, it.thinking, it.thinkingMs,
+                                it.status, it.createdAt, it.newTokens, it.decodeTokS,
+                                it.ttftMs, it.prefillTokS, it.memBytes) }
+        }
+
+    override suspend fun startThreadWithUserMessage(modelId: String, userText: String): String {
+        val now = System.currentTimeMillis()
+        val tid = UUID.randomUUID().toString()
+        db.threads().upsert(ThreadEntity(tid, userText.take(60), modelId, now, now))
+        db.messages().insert(MessageEntity(UUID.randomUUID().toString(), tid, "user",
+                                           userText, null, null, MsgStatus.OK, modelId, now,
+                                           null, null, null))
+        return tid
+    }
+
+    override suspend fun beginAssistant(threadId: String, modelId: String): String {
+        val id = UUID.randomUUID().toString()
+        db.messages().insert(MessageEntity(id, threadId, "assistant", "", null, null,
+                                           MsgStatus.STREAMING, modelId,
+                                           System.currentTimeMillis(), null, null, null))
+        return id
+    }
+
+    override suspend fun addUserMessage(threadId: String, userText: String) {
+        val now = System.currentTimeMillis()
+        db.messages().insert(MessageEntity(UUID.randomUUID().toString(), threadId, "user",
+                                           userText, null, null, MsgStatus.OK, "", now,
+                                           null, null, null))
+        db.threads().touch(threadId, null, now)
+    }
+
+    override suspend fun finalizeAssistant(id: String, content: String, thinking: String?,
+                                           thinkingMs: Long?, status: String,
+                                           promptTokens: Int?, newTokens: Int?,
+                                           decodeTokS: Double?, genTokens: ByteArray?,
+                                           ttftMs: Long?, prefillTokS: Double?, memBytes: Long?) {
+        db.messages().finalizeAssistant(id, content, thinking, thinkingMs, status,
+                                        promptTokens, newTokens, decodeTokS, genTokens,
+                                        ttftMs, prefillTokS, memBytes)
+    }
+
+    override suspend fun renameThread(id: String, title: String) =
+        db.threads().touch(id, title, System.currentTimeMillis())
+
+    override suspend fun deleteThread(id: String) = db.threads().delete(id)
+
+    override suspend fun editUserMessage(id: String, newText: String) {
+        val msg = db.messages().byId(id) ?: return
+        db.messages().update(msg.copy(content = newText))
+        db.messages().deleteFrom(msg.threadId, msg.createdAt + 1) // truncate everything after (incl. old answer)
+        // editing the first user message renames the thread (title = first user message)
+        if (db.messages().olderUserCount(msg.threadId, msg.createdAt) == 0) {
+            db.threads().touch(msg.threadId, newText.take(60), System.currentTimeMillis())
+        }
+    }
+
+    override suspend fun deleteMessage(id: String) = db.messages().deleteById(id)
+
+    override suspend fun markStreamingInterrupted() = db.messages().markStreamingInterrupted()
+
+    override suspend fun history(threadId: String): List<StoredMessage> =
+        db.messages().observeFor(threadId).first()
+            .map { StoredMessage(it.role, it.content, it.thinking, decodeIds(it.genTokens)) }
+
+    override suspend fun historyWithIds(threadId: String): List<MessageView> =
+        db.messages().observeFor(threadId).first()
+            .map { MessageView(it.id, it.role, it.content, it.thinking, it.thinkingMs,
+                               it.status, it.createdAt, it.newTokens, it.decodeTokS,
+                               it.ttftMs, it.prefillTokS, it.memBytes) }
+
+    override suspend fun deleteFrom(threadId: String, fromTs: Long) =
+        db.messages().deleteFrom(threadId, fromTs)
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/data/Daos.kt b/android/app/src/main/java/dev/edge0/runtime/app/data/Daos.kt
new file mode 100644
index 0000000..5b2d9a5
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/data/Daos.kt
@@ -0,0 +1,63 @@
+// data/Daos.kt - thread/message DAOs (flat list ordered by updatedAt desc).
+package dev.edge0.runtime.app.data
+
+import androidx.room.Dao
+import androidx.room.Delete
+import androidx.room.Insert
+import androidx.room.OnConflictStrategy
+import androidx.room.Query
+import androidx.room.Update
+import kotlinx.coroutines.flow.Flow
+
+@Dao
+interface ThreadDao {
+    @Query("SELECT * FROM threads ORDER BY updatedAt DESC")
+    fun observeAll(): Flow<List<ThreadEntity>>
+
+    @Insert(onConflict = OnConflictStrategy.REPLACE)
+    suspend fun upsert(thread: ThreadEntity)
+
+    @Query("UPDATE threads SET title = COALESCE(:title, title), updatedAt = :ts WHERE id = :id")
+    suspend fun touch(id: String, title: String?, ts: Long)
+
+    @Query("DELETE FROM threads WHERE id = :id")
+    suspend fun delete(id: String)
+}
+
+@Dao
+interface MessageDao {
+    @Query("SELECT * FROM messages WHERE threadId = :threadId ORDER BY createdAt ASC")
+    fun observeFor(threadId: String): Flow<List<MessageEntity>>
+
+    @Insert
+    suspend fun insert(msg: MessageEntity)
+
+    @Update
+    suspend fun update(msg: MessageEntity)
+
+    @Query("UPDATE messages SET content=:content, thinking=:thinking, thinkingMs=:thinkingMs, " +
+           "status=:status, promptTokens=:pt, newTokens=:nt, decodeTokS=:dts, " +
+           "genTokens=:genTokens, ttftMs=:ttft, prefillTokS=:pfs, memBytes=:mem WHERE id=:id")
+    suspend fun finalizeAssistant(id: String, content: String, thinking: String?,
+                                  thinkingMs: Long?, status: String,
+                                  pt: Int?, nt: Int?, dts: Double?, genTokens: ByteArray?,
+                                  ttft: Long?, pfs: Double?, mem: Long?)
+
+    /** Cold-start sweep: rows left streaming by a killed process become interrupted. */
+    @Query("UPDATE messages SET status='interrupted' WHERE status='streaming'")
+    suspend fun markStreamingInterrupted()
+
+    /** Edit/regenerate truncation: delete every message at/after the given created-at. */
+    @Query("DELETE FROM messages WHERE threadId=:threadId AND createdAt >= :fromTs")
+    suspend fun deleteFrom(threadId: String, fromTs: Long)
+
+    @Query("SELECT * FROM messages WHERE id = :id")
+    suspend fun byId(id: String): MessageEntity?
+
+    @Query("SELECT COUNT(*) FROM messages WHERE threadId=:threadId AND createdAt < :ts " +
+           "AND role='user'")
+    suspend fun olderUserCount(threadId: String, ts: Long): Int
+
+    @Query("DELETE FROM messages WHERE id = :id")
+    suspend fun deleteById(id: String)
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/data/Entities.kt b/android/app/src/main/java/dev/edge0/runtime/app/data/Entities.kt
new file mode 100644
index 0000000..d2da789
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/data/Entities.kt
@@ -0,0 +1,52 @@
+// data/Entities.kt - minimal two-table Room schema. Tagged raw text is not stored:
+// thinking/body are separate columns; history reconstruction goes through
+// HistoryWindow (tag pairs are re-added at render time).
+package dev.edge0.runtime.app.data
+
+import androidx.room.Entity
+import androidx.room.ForeignKey
+import androidx.room.Index
+import androidx.room.PrimaryKey
+
+@Entity(tableName = "threads", indices = [Index("updatedAt")])
+data class ThreadEntity(
+    @PrimaryKey val id: String,
+    val title: String?,            // truncated first user message (no AI summary)
+    val modelId: String,
+    val createdAt: Long,
+    val updatedAt: Long,
+)
+
+/** status: ok | cancelled | error | interrupted (killed process leaves interrupted) */
+@Entity(
+    tableName = "messages",
+    foreignKeys = [ForeignKey(entity = ThreadEntity::class, parentColumns = ["id"],
+                              childColumns = ["threadId"], onDelete = ForeignKey.CASCADE)],
+    indices = [Index("threadId", "createdAt")],
+)
+data class MessageEntity(
+    @PrimaryKey val id: String,
+    val threadId: String,
+    val role: String,              // user | assistant
+    val content: String,           // tag-free body text
+    val thinking: String?,         // extracted reasoning (null = no thinking this turn)
+    val thinkingMs: Long?,
+    val status: String,            // "streaming" while generating, overwritten at terminal state
+    val modelId: String,
+    val createdAt: Long,
+    val promptTokens: Int?,
+    val newTokens: Int?,
+    val decodeTokS: Double?,
+    val genTokens: ByteArray? = null,  // sidecar: token ids stepped this turn (int32 LE BLOB)
+    val ttftMs: Long? = null,          // v3: perceived time-to-first-token, wall clock ms
+    val prefillTokS: Double? = null,   // v3: prefill tok/s
+    val memBytes: Long? = null,        // v3: expert-pool resident bytes (excludes cacheable mmap)
+)
+
+object MsgStatus {
+    const val STREAMING = "streaming"
+    const val OK = "ok"
+    const val CANCELLED = "cancelled"
+    const val ERROR = "error"
+    const val INTERRUPTED = "interrupted"
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/data/FakeChatRepository.kt b/android/app/src/main/java/dev/edge0/runtime/app/data/FakeChatRepository.kt
new file mode 100644
index 0000000..ad05eae
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/data/FakeChatRepository.kt
@@ -0,0 +1,99 @@
+// data/FakeChatRepository.kt - in-memory implementation for UI walkthroughs (no engine, no disk).
+package dev.edge0.runtime.app.data
+
+import dev.edge0.runtime.app.runtime.StoredMessage
+import dev.edge0.runtime.app.runtime.decodeIds
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.map
+
+class FakeChatRepository : ChatRepository {
+    data class Msg(val id: String, val threadId: String, val role: String, var content: String,
+                   var thinking: String?, var status: String, val createdAt: Long,
+                   var genTokens: ByteArray? = null)
+
+    private val threads = MutableStateFlow<List<ThreadSummary>>(emptyList())
+    private val msgs = MutableStateFlow<List<Msg>>(emptyList())
+    private var seq = 0
+    private fun nid() = "fx${seq++}"
+
+    override fun observeThreads(): Flow<List<ThreadSummary>> = threads
+    override fun observeMessages(threadId: String): Flow<List<MessageView>> =
+        msgs.map { l -> l.filter { it.threadId == threadId }.sortedBy { it.createdAt }
+            .map { MessageView(it.id, it.role, it.content, it.thinking, null, it.status,
+                               it.createdAt, null, null) } }
+
+    override suspend fun startThreadWithUserMessage(modelId: String, userText: String): String {
+        val tid = nid()
+        val now = System.currentTimeMillis()
+        msgs.value += Msg(tid, tid, "user", userText, null, MsgStatus.OK, now)
+        threads.value = listOf(ThreadSummary(tid, userText.take(60), now, modelId)) + threads.value
+        return tid
+    }
+
+    override suspend fun beginAssistant(threadId: String, modelId: String): String {
+        val id = nid()
+        msgs.value += Msg(id, threadId, "assistant", "", null, MsgStatus.STREAMING,
+                          System.currentTimeMillis())
+        return id
+    }
+
+    override suspend fun addUserMessage(threadId: String, userText: String) {
+        val now = System.currentTimeMillis()
+        msgs.value += Msg(nid(), threadId, "user", userText, null, MsgStatus.OK, now)
+        threads.value = threads.value.map {
+            if (it.id == threadId) it.copy(updatedAt = now) else it
+        }
+    }
+
+    override suspend fun finalizeAssistant(id: String, content: String, thinking: String?,
+                                           thinkingMs: Long?, status: String,
+                                           promptTokens: Int?, newTokens: Int?,
+                                           decodeTokS: Double?, genTokens: ByteArray?,
+                                           ttftMs: Long?, prefillTokS: Double?, memBytes: Long?) {
+        msgs.value = msgs.value.map {
+            if (it.id == id) it.copy(content = content, thinking = thinking, status = status,
+                                     genTokens = genTokens)
+            else it
+        }
+    }
+
+    override suspend fun renameThread(id: String, title: String) {
+        threads.value = threads.value.map { if (it.id == id) it.copy(title = title) else it }
+    }
+
+    override suspend fun deleteThread(id: String) {
+        threads.value = threads.value.filter { it.id != id }
+        msgs.value = msgs.value.filter { it.threadId != id }
+    }
+
+    override suspend fun editUserMessage(id: String, newText: String) {
+        val target = msgs.value.firstOrNull { it.id == id } ?: return
+        msgs.value = msgs.value.filter { it.createdAt < target.createdAt } +
+            target.copy(content = newText)
+    }
+
+    override suspend fun deleteMessage(id: String) {
+        msgs.value = msgs.value.filter { it.id != id }
+    }
+
+    override suspend fun markStreamingInterrupted() {
+        msgs.value = msgs.value.map {
+            if (it.status == MsgStatus.STREAMING) it.copy(status = MsgStatus.INTERRUPTED) else it
+        }
+    }
+
+    override suspend fun history(threadId: String): List<StoredMessage> =
+        msgs.value.filter { it.threadId == threadId && it.status != MsgStatus.STREAMING }
+            .sortedBy { it.createdAt }
+            .map { StoredMessage(it.role, it.content, it.thinking, decodeIds(it.genTokens)) }
+
+    override suspend fun historyWithIds(threadId: String): List<MessageView> =
+        msgs.value.filter { it.threadId == threadId }.sortedBy { it.createdAt }
+            .map { MessageView(it.id, it.role, it.content, it.thinking, null, it.status,
+                               it.createdAt, null, null) }
+
+    override suspend fun deleteFrom(threadId: String, fromTs: Long) {
+        msgs.value = msgs.value.filter { !(it.threadId == threadId && it.createdAt >= fromTs) }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/data/SettingsStore.kt b/android/app/src/main/java/dev/edge0/runtime/app/data/SettingsStore.kt
new file mode 100644
index 0000000..3fc1595
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/data/SettingsStore.kt
@@ -0,0 +1,97 @@
+// data/SettingsStore.kt - DataStore preferences snapshot (appearance / assistant / sampling / toggles).
+// Defaults follow the upstream sampling quartet; installs without a stored key
+// fall back to the current defaults.
+package dev.edge0.runtime.app.data
+
+import android.content.Context
+import androidx.datastore.core.DataStore
+import androidx.datastore.preferences.core.Preferences
+import androidx.datastore.preferences.core.booleanPreferencesKey
+import androidx.datastore.preferences.core.edit
+import androidx.datastore.preferences.core.floatPreferencesKey
+import androidx.datastore.preferences.core.intPreferencesKey
+import androidx.datastore.preferences.core.longPreferencesKey
+import androidx.datastore.preferences.core.stringPreferencesKey
+import androidx.datastore.preferences.preferencesDataStore
+import dev.edge0.runtime.app.ui.theme.ThemeMode
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.map
+
+private val Context.dataStore: DataStore<Preferences> by preferencesDataStore("edge0_settings")
+
+data class AppSettings(
+    val themeMode: ThemeMode = ThemeMode.DARK,   // dark by default (does not follow system)
+    val fontScale: Float = 1.0f,          // 0.85..1.3
+    val systemPrompt: String = "You're Edge0, an on-device AI assistant.",
+    val useSystemPrompt: Boolean = true,    // explicit toggle; empty text also means off
+    val temperature: Float = 0.2f,        // precise default (control lives in Settings)
+    val topK: Int = 64,
+    val topP: Float = 0.95f,
+    val repetitionPenalty: Float = 1.1f,  // aligned with upstream; 1.0 = off
+    val maxNewTokens: Int = 0,            // 0 = app default 1024 (reasoning models need headroom)
+    val seed: Long = 0x20260915L,
+    val enableThinking: Boolean = false,    // master thinking toggle: drives both the
+    //   model-side template instruction and UI block visibility;
+    //   storage key kept as show_thinking for backward compatibility
+    val metricsOverlay: Boolean = false,  // off by default
+    val activeModelDir: String = "",
+)
+
+object SettingsKeys {
+    val THEME = stringPreferencesKey("theme_mode")
+    val FONT_SCALE = floatPreferencesKey("font_scale")
+    val SYSTEM_PROMPT = stringPreferencesKey("system_prompt")
+    val USE_SYS_PROMPT = booleanPreferencesKey("use_system_prompt")
+    val TEMPERATURE = floatPreferencesKey("temperature")
+    val TOP_K = intPreferencesKey("top_k")
+    val TOP_P = floatPreferencesKey("top_p")
+    val REP_PENALTY = floatPreferencesKey("repetition_penalty")
+    val MAX_NEW = intPreferencesKey("max_new_tokens")
+    val SEED = longPreferencesKey("seed")
+    val SHOW_THINKING = booleanPreferencesKey("show_thinking")
+    val OVERLAY = booleanPreferencesKey("metrics_overlay")
+    val ACTIVE_MODEL = stringPreferencesKey("active_model_dir")
+}
+
+class SettingsStore(private val context: Context) {
+
+    val flow: Flow<AppSettings> = context.dataStore.data.map { p ->
+        val d = AppSettings()
+        AppSettings(
+            themeMode = p[SettingsKeys.THEME]?.let { runCatching { ThemeMode.valueOf(it) }.getOrNull() } ?: d.themeMode,
+            fontScale = p[SettingsKeys.FONT_SCALE] ?: d.fontScale,
+            systemPrompt = p[SettingsKeys.SYSTEM_PROMPT] ?: d.systemPrompt,
+            useSystemPrompt = p[SettingsKeys.USE_SYS_PROMPT] ?: d.useSystemPrompt,
+            temperature = p[SettingsKeys.TEMPERATURE] ?: d.temperature,
+            topK = p[SettingsKeys.TOP_K] ?: d.topK,
+            topP = p[SettingsKeys.TOP_P] ?: d.topP,
+            repetitionPenalty = p[SettingsKeys.REP_PENALTY] ?: d.repetitionPenalty,
+            maxNewTokens = p[SettingsKeys.MAX_NEW] ?: d.maxNewTokens,
+            seed = p[SettingsKeys.SEED] ?: d.seed,
+            enableThinking = p[SettingsKeys.SHOW_THINKING] ?: d.enableThinking,
+            metricsOverlay = p[SettingsKeys.OVERLAY] ?: d.metricsOverlay,
+            activeModelDir = p[SettingsKeys.ACTIVE_MODEL] ?: d.activeModelDir,
+        )
+    }
+
+    suspend fun update(block: (AppSettings) -> AppSettings) {
+        val current = runCatching { flow.first() }.getOrElse { AppSettings() }
+        val next = block(current)
+        context.dataStore.edit { p ->
+            p[SettingsKeys.THEME] = next.themeMode.name
+            p[SettingsKeys.FONT_SCALE] = next.fontScale
+            p[SettingsKeys.SYSTEM_PROMPT] = next.systemPrompt
+            p[SettingsKeys.USE_SYS_PROMPT] = next.useSystemPrompt
+            p[SettingsKeys.TEMPERATURE] = next.temperature
+            p[SettingsKeys.TOP_K] = next.topK
+            p[SettingsKeys.TOP_P] = next.topP
+            p[SettingsKeys.REP_PENALTY] = next.repetitionPenalty
+            p[SettingsKeys.MAX_NEW] = next.maxNewTokens
+            p[SettingsKeys.SEED] = next.seed
+            p[SettingsKeys.SHOW_THINKING] = next.enableThinking
+            p[SettingsKeys.OVERLAY] = next.metricsOverlay
+            p[SettingsKeys.ACTIVE_MODEL] = next.activeModelDir
+        }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/GenState.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/GenState.kt
new file mode 100644
index 0000000..bec25d0
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/GenState.kt
@@ -0,0 +1,26 @@
+// runtime/GenState.kt - generation state machine and events (typed, single-slot).
+package dev.edge0.runtime.app.runtime
+
+import dev.edge0.runtime.engine.Status
+
+sealed interface GenState {
+    data object Idle : GenState
+    data class Loading(val modelDir: String) : GenState
+    data class Streaming(val threadId: String) : GenState
+    data object Cancelling : GenState // settling gate: cancel sent, old job not yet home (new submits refused)
+    data class Error(val status: Status, val diag: String) : GenState
+}
+
+sealed interface GenEvent {
+    data class ThinkPiece(val s: String) : GenEvent
+    data class TextPiece(val s: String) : GenEvent
+    // genIds: token ids actually stepped this turn (sidecar; absent on cancel/failure)
+    data class Done(val metrics: RequestMetrics?, val genIds: IntArray = IntArray(0)) : GenEvent
+    data class Cancelled(val metrics: RequestMetrics?) : GenEvent
+    data class Failed(val status: Status, val diag: String) : GenEvent
+}
+
+fun Seg.toEvent(): GenEvent = when (this) {
+    is Seg.Text -> GenEvent.TextPiece(s)
+    is Seg.Think -> GenEvent.ThinkPiece(s)
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/HistoryWindow.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/HistoryWindow.kt
new file mode 100644
index 0000000..758b686
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/HistoryWindow.kt
@@ -0,0 +1,69 @@
+// runtime/HistoryWindow.kt - multi-turn context assembly with a character budget.
+// The budget is conservative (mixed CJK/Latin ~1.5 chars/token against a 4096-token
+// context); the hard overflow gate is the engine itself (INVALID_ARG with a
+// "prompt tokens=" diagnostic) and callers retry with a tighter budget.
+// Assistant history re-enters the prompt wrapped in think tag pairs at render time.
+package dev.edge0.runtime.app.runtime
+
+import dev.edge0.runtime.engine.ChatMessage
+
+/** Minimal persisted-message shape (Room entities map to this).
+ *  genIds = assistant token-id sidecar: lives and dies with its row. */
+data class StoredMessage(val role: String, val content: String, val thinking: String? = null,
+                         val genIds: IntArray? = null)
+
+/** Sidecar serialization: IntArray to int32-LE BLOB (Room column); empty becomes null. */
+fun encodeIds(ids: IntArray?): ByteArray? {
+    if (ids == null || ids.isEmpty()) return null
+    val bb = java.nio.ByteBuffer.allocate(ids.size * 4).order(java.nio.ByteOrder.LITTLE_ENDIAN)
+    ids.forEach { bb.putInt(it) }
+    return bb.array()
+}
+
+/** Sidecar deserialization (length not a multiple of 4, or empty, falls back to text). */
+fun decodeIds(b: ByteArray?): IntArray? {
+    if (b == null || b.isEmpty() || b.size % 4 != 0) return null
+    val bb = java.nio.ByteBuffer.wrap(b).order(java.nio.ByteOrder.LITTLE_ENDIAN)
+    return IntArray(b.size / 4) { bb.int }
+}
+
+object HistoryWindow {
+    const val DEFAULT_BUDGET_CHARS = 10000
+
+    /** turns old to new (user/assistant alternating, last is the user question); system first when non-blank. */
+    fun build(systemPrompt: String?, turns: List<StoredMessage>,
+              budgetChars: Int = DEFAULT_BUDGET_CHARS): List<ChatMessage> {
+        val msgs = ArrayList<ChatMessage>(turns.size + 1)
+        var used = 0
+        if (!systemPrompt.isNullOrBlank()) {
+            msgs.add(ChatMessage("system", systemPrompt))
+            used += systemPrompt.length
+        }
+        // poison-turn filter: blank assistant turns without ids must not reach the
+        // model - replaying "I answered nothing" derails every later turn. Dropped from
+        // the model-facing context only; the UI record is untouched.
+        val clean = turns.filterNot { it.role == "assistant" && it.content.isBlank() && it.genIds?.isEmpty() != false }
+        val picked = ArrayList<StoredMessage>()
+        for (t in clean.asReversed()) {
+            val cost = t.content.length + (t.thinking?.length ?: 0) + 12
+            // the current question is always kept, even over budget; engine overflow diagnostic triggers retry
+            if (picked.isNotEmpty() && used + cost > budgetChars) break
+            picked.add(t)
+            used += cost
+        }
+        picked.reverse()
+        for (t in picked) {
+            // Keep the engine transcript exact while the UI still stores and
+            // renders thinking separately. The native renderer handles the
+            // family-specific newline/empty-pair form.
+            val engineContent = if (t.role == "assistant" && !t.thinking.isNullOrEmpty()) {
+                "<think>" + t.thinking + "</think>" + t.content
+            } else t.content
+            msgs.add(ChatMessage(t.role, engineContent, t.genIds))
+        }
+        return msgs
+    }
+
+    /** Tightened budget for the overflow retry (caller catches the INVALID_ARG diagnostic). */
+    fun tighterBudget(budgetChars: Int): Int = (budgetChars / 2).coerceAtLeast(400)
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/LlamaNative.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/LlamaNative.kt
new file mode 100644
index 0000000..098108c
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/LlamaNative.kt
@@ -0,0 +1,25 @@
+// LlamaNative.kt - JNI declaration surface for llama_chat.c.
+// external fun symbols bind hard to the package name
+// (dev_edge0_runtime_app_runtime_LlamaNative_*): moving packages requires
+// mirroring the rename on the C side. Top-level object members keep stable
+// exported names (no @JvmName needed).
+package dev.edge0.runtime.app.runtime
+
+object LlamaNative {
+    init {
+        // packaged via APK jniLibs (five shared libs from tools/llama/build_vendor_libs.sh);
+// DT_NEEDED resolves libllama -> libggml* automatically
+        System.loadLibrary("edge0llama")
+    }
+
+    external fun nativeInit(model: String, lora: String, heads: String,
+                            nCtx: Int, poolMb: Int, blobMb: Int, pred: Int, ioThreads: Int): String
+    external fun nativeReset(): String
+    external fun nativeAddRole(role: String, text: String): String
+    external fun nativeGenerate(sink: (String) -> Unit, maxTokens: Int, temp: Float,
+                           thinkingOn: Boolean): String
+    external fun nativeCancel()
+    external fun nativeStats(): String
+    external fun nativeLastAssistant(): String?
+    external fun nativeFree(): String
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/LlamaRuntime.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/LlamaRuntime.kt
new file mode 100644
index 0000000..780372a
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/LlamaRuntime.kt
@@ -0,0 +1,264 @@
+// runtime/LlamaRuntime.kt - the Runtime implementation backed by llama.cpp + GGUF.
+// Native side = LlamaNative / llama_chat.c; shared libs are staged into
+// build-dl/llama-libs at build time (gitignored, produced by
+// tools/llama/build_vendor_libs.sh from the pinned engine tree).
+// Auto tiering (mirrors the validated CLI configuration): model >= 12GB takes the
+//   35B tier - expert pool 6144MB + blob 3072MB + n_ctx 4096, demand-only
+//   (prefetch measured net-negative on-device, so E0_PREROUTER is never set);
+//   otherwise the 8B tier - fully resident, pool/blob disabled, n_ctx 8192,
+//   companion LoRA auto-discovered (lora_*.gguf).
+// LoRA companion applies to the 8B tier only: the 35B build bakes its LoRA in,
+// and attaching an adapter on top fails initialization (verified on-device).
+// History contract: native is append-only; this class keeps a mirror and does a
+// prefix diff - any fork triggers a full rebuild (thread switch / message edit).
+package dev.edge0.runtime.app.runtime
+
+import dev.edge0.runtime.engine.ChatMessage
+import dev.edge0.runtime.engine.GenParams
+import dev.edge0.runtime.engine.Status
+import dev.edge0.runtime.app.runtime.PoolMetrics
+import java.io.File
+import java.util.concurrent.Executors
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+
+class LlamaRuntime : Runtime {
+    override val state = MutableStateFlow<GenState>(GenState.Idle)
+    override val lastMetrics = MutableStateFlow<RequestMetrics?>(null)
+
+    @Volatile private var _dir: String? = null
+    override val activeModelDir: String? get() = _dir
+
+    private val slotMutex = Mutex()
+    private val genExecutor = Executors.newSingleThreadExecutor { r ->
+        Thread(r, "llama-gen").apply { isDaemon = true }
+    }
+    private val genCtx: CoroutineDispatcher = genExecutor.asCoroutineDispatcher()
+
+    private val mirror = ArrayList<Pair<String, String>>()   // (role, content) kept in sync with native
+
+    companion object {
+        const val TAG = "LlamaRuntime"
+        const val BIG_MODEL_BYTES = 12L shl 30   // tier threshold (midpoint of 8B=5.0GB / 35B=21.7GB)
+    }
+
+    private fun runCatchingNative(block: () -> String): String? =
+        try { block() } catch (t: Throwable) {
+            // never swallow native failures into a generic "so load failed" -
+            // the real cause must surface in the return value
+            "ERR native: ${t.javaClass.simpleName}: ${t.message?.take(160)}"
+        }
+
+    override suspend fun ensureLoaded(path: String) {
+        slotMutex.withLock {
+            if (_dir == path && state.value !is GenState.Error) return
+            val model = resolveGguf(path)
+                ?: run {
+                    state.value = GenState.Error(Status.FORMAT_VERSION, "model.gguf not found: $path")
+                    return
+                }
+            state.value = GenState.Loading(path)
+            if (_dir != null) {   // model swap: unload first (never two loads resident; nativeInit self-clears as backstop)
+                runCatching { LlamaNative.nativeFree() }
+                _dir = null
+            }
+            val big = model.length() >= BIG_MODEL_BYTES
+            val dir = model.parentFile
+            val lora = if (big) "" else dir?.listFiles { f ->
+                f.name.startsWith("lora_") && f.name.endsWith(".gguf")
+            }?.firstOrNull()?.absolutePath ?: ""
+            val heads = dir?.let { File(it, "e0_heads.bin") }?.takeIf { it.exists() }?.absolutePath ?: ""
+            val r = withContext(Dispatchers.IO) {
+                runCatchingNative {
+                    LlamaNative.nativeInit(model.absolutePath, lora,
+                        if (big) heads else "",          // heads ship with the 35B tier (idle unless E0_PREROUTER is enabled)
+                        if (big) 4096 else 8192,
+                        if (big) 6144 else 0,
+                        if (big) 3072 else 0,
+                        0,                               // prefetch: measured net-negative, permanently off
+                        4)
+                }
+            }
+            if (r == null || !r.startsWith("OK")) {
+                state.value = GenState.Error(Status.IO, (r ?: "libedge0llama.so failed to load"))
+                return
+            }
+            mirror.clear()
+            _dir = path
+            state.value = GenState.Idle
+        }
+    }
+
+    /** path may be a .gguf file, or a directory containing model.gguf. */
+    private fun resolveGguf(path: String): File? {
+        val f = File(path)
+        return when {
+            f.isFile && path.endsWith(".gguf") -> f
+            f.isDirectory -> File(f, "model.gguf").takeIf { it.exists() }
+                ?: f.listFiles { x -> x.name.endsWith(".gguf") && !x.name.startsWith("lora_") }?.firstOrNull()
+            else -> null
+        }
+    }
+
+    override suspend fun unloadActive() {
+        slotMutex.withLock {
+            if (_dir == null) return
+            withContext(Dispatchers.IO) { runCatchingNative { LlamaNative.nativeFree() } }
+            mirror.clear()
+            _dir = null
+            if (state.value !is GenState.Error) state.value = GenState.Idle
+        }
+    }
+
+    override suspend fun trim(ratio: Float): Boolean {
+        // TODO: expose moe_pool trim through native; the pool currently self-sheds via E0_MEM_FLOOR_MB
+        return true
+    }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    override fun submit(threadId: String, messages: List<ChatMessage>,
+                        thinkingOn: Boolean, params: GenParams): Flow<GenEvent> =
+        callbackFlow {
+            if (_dir == null) {
+                trySend(GenEvent.Failed(Status.INVALID_ARG, "no active llama session")); close(); return@callbackFlow
+            }
+            if (!state.compareAndSet(GenState.Idle, GenState.Streaming(threadId))) {
+                trySend(GenEvent.Failed(Status.INVALID_ARG, "generation slot busy")); close(); return@callbackFlow
+            }
+            launch(genCtx) {
+                val splitter = ThinkSplitter(initialThink = thinkingOn)
+                val tSend = System.nanoTime()
+                var tFirst = 0L
+                val uiText = StringBuilder()
+                val sink: (String) -> Unit = { piece ->
+                    if (piece.isNotEmpty()) {
+                        if (tFirst == 0L) tFirst = System.nanoTime()
+                        uiText.append(piece)
+                        splitter.feed(piece).forEach { trySend(it.toEvent()) }
+                    }
+                }
+                syncHistory(messages)
+                val res = runCatchingNative {
+                    LlamaNative.nativeGenerate(sink, params.maxNewTokens.coerceAtLeast(1),
+                        params.temperature, thinkingOn)
+                } ?: "ERR native"
+                splitter.flush().forEach { trySend(it.toEvent()) }
+                // native already appended this turn with stepped ids (sidecar); mirror keeps index parity
+                if (res.startsWith("OK")) mirror.add("assistant" to uiText.toString())
+                val term: GenEvent = when {
+                    res.startsWith("OK") || res.startsWith("CANCELLED") -> {
+                        val m = parseMetrics(res)
+                        val mm = m?.copy(
+                            // user-perceived TTFT measured here (native first_ms is post-prefill sampling time)
+                            firstTokenMs = if (tFirst > 0) (tFirst - tSend) / 1_000_000 else m.firstTokenMs,
+                            peakRssBytes = readPeakRss(),
+                            pool = poolMetrics(),
+                            gpu = GpuMetrics(probed = false, ok = false, active = false,
+                                device = "cpu(llama-route)", why = "CPU-only (shipping configuration)"),
+                        )
+                        if (mm != null) lastMetrics.value = mm
+                        if (res.startsWith("OK")) GenEvent.Done(mm) else GenEvent.Cancelled(mm)
+                    }
+                    else -> GenEvent.Failed(Status.IO, res)
+                }
+                trySend(term)
+                if (state.value !is GenState.Error) state.value = GenState.Idle
+                close()
+            }
+            awaitClose { }
+        }
+
+    private fun syncHistory(msgs: List<ChatMessage>) {
+        // Prefix match: user turns compare verbatim (UI text feeds native losslessly);
+        // assistant turns match by ROLE ONLY - the ground truth lives on the native side
+        // as stepped ids (sidecar), while UI text is a lossy reassembly with think tags
+        // stripped. Verbatim comparison there would always fake a fork and replay the
+        // whole history every turn (once cost 60s of dead waiting). Non-contiguous hits
+        // must not count toward the common prefix either.
+        // Known edge: editing/regenerating history falls back to full rebuild.
+        fun ok(i: Int) = i < mirror.size && mirror[i].first == msgs[i].role &&
+            (msgs[i].role != "user" || mirror[i].second == msgs[i].content)
+        var common = 0
+        while (common < msgs.size && ok(common)) common++
+        if (common == mirror.size) {
+            msgs.drop(common).forEach { m ->
+                val content = m.content
+                LlamaNative.nativeAddRole(m.role, content)
+                mirror.add(m.role to content)
+            }
+        } else {
+            LlamaNative.nativeReset()
+            mirror.clear()
+            msgs.forEach { m ->
+                val content = m.content
+                LlamaNative.nativeAddRole(m.role, content)
+                mirror.add(m.role to content)
+            }
+        }
+    }
+
+    internal fun parseMetrics(res: String): RequestMetrics? {
+        val kv = Regex("""(\w+)=(\S+)""").findAll(res).associate { it.groupValues[1] to it.groupValues[2] }
+        val gen = kv["gen"]?.toIntOrNull() ?: return null
+        val tps = kv["tps"]?.toDoubleOrNull() ?: 0.0
+        val pf = kv["prefill_ms"]?.toLongOrNull() ?: 0L
+        val first = kv["first_ms"]?.toLongOrNull() ?: 0L
+        val prompt = kv["prompt"]?.toIntOrNull() ?: 0
+        val reused = kv["reused"]?.toIntOrNull() ?: 0
+        return RequestMetrics(
+            promptTokens = prompt,
+            prefix_reused = reused > 0,
+            prefix_reused_tokens = reused,
+            newTokens = gen,
+            prefillMs = pf,
+            firstTokenMs = first,
+            decodeMs = (if (gen > 0) (gen * 1000.0 / (tps + 1e-6)).toLong() else 0L) + first,
+            prefillTokS = if (pf > 0) (prompt - reused) * 1000.0 / pf else 0.0,
+            decodeTokS = tps,
+            peakRssBytes = 0,
+            chatTurns = 1,
+        )
+    }
+
+    private fun poolMetrics(): PoolMetrics? {
+        val raw = runCatching { LlamaNative.nativeStats() }.getOrNull() ?: return null
+        if (!raw.contains("hits=")) return null
+        val kv = raw.split(" ", "\n").mapNotNull {
+            val i = it.indexOf('='); if (i <= 0) null else it.substring(0, i) to it.substring(i + 1)
+        }.toMap()
+        fun L(k: String) = (kv[k]?.toLongOrNull() ?: 0L)
+        val slots = (kv["slots"] ?: "").split("/")
+        return PoolMetrics(
+            resident_bytes = L("resident"),
+            hits = L("hits"), misses = L("misses"),
+            prefetch_used = L("pf_used"), prefetch_mispredict = L("pf_miss"),
+            evictions = L("evicts"), stall_ns = L("stall_ns"),
+            flash_bytes = L("flash"),
+            slots_used = slots.firstOrNull()?.toLongOrNull() ?: 0L,
+            slots_total = slots.getOrNull(1)?.toLongOrNull() ?: 0L,
+        )
+    }
+
+    private fun readPeakRss(): Long {
+        val raw = runCatching { LlamaNative.nativeStats() }.getOrNull() ?: return 0
+        return Regex("vmhwm_kb=(\\d+)").find(raw)?.groupValues?.get(1)?.toLongOrNull()?.times(1024) ?: 0
+    }
+
+    /** Raw native text of the last assistant turn (ground truth for history sync). */
+    fun lastAssistantNative(): String? = runCatching { LlamaNative.nativeLastAssistant() }.getOrNull()
+
+    override fun cancelActive() {
+        runCatching { LlamaNative.nativeCancel() }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/MetricsMapper.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/MetricsMapper.kt
new file mode 100644
index 0000000..4e04782
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/MetricsMapper.kt
@@ -0,0 +1,99 @@
+// runtime/MetricsMapper.kt - typed mapping of engine metrics (append-only schema;
+// unknown fields ignored, missing fields default to the N/A display semantics).
+package dev.edge0.runtime.app.runtime
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.json.Json
+
+@Serializable
+data class RouteMetrics(
+    val mode: String = "",        // predicted|exact (audit echo, prevents track mixing)
+    val feature: String = "",     // executed|teacher|n/a
+    @SerialName("dec_steps") val decSteps: Long = 0,
+    @SerialName("pred_used") val predUsed: Long = 0,
+    @SerialName("route_set_size") val routeSetSize: Int = 0,
+    val agree: Double? = null,
+)
+
+@Serializable
+data class GpuMetrics(
+    val probed: Boolean = false,
+    val ok: Boolean = false,
+    val active: Boolean = false,  // whether this request used the GPU (prefill window)
+    val policy: Int = 0,
+    val device: String? = null,
+    val why: String = "",         // probe conclusion echo (ok / dlopen:.. / policy-cpu ...)
+    @SerialName("cache_hit") val cacheHit: Boolean = false,
+)
+
+@Serializable
+data class ConfigEcho(
+    @SerialName("clip_value") val clipValue: Double = 0.0, // per-layer sanitize clamp (0 = off)
+    @SerialName("bundle_mode") val bundleMode: String = "",
+    @SerialName("production_top_k") val productionTopK: Int = 0,
+)
+
+@Serializable
+data class ParamsEcho(
+    @SerialName("repetition_penalty") val repPenalty: Double = 0.0,
+    @SerialName("first_token_greedy") val firstTokenGreedy: Int = 0,
+)
+
+@Serializable
+data class PoolMetrics(
+    val resident_bytes: Long = 0,
+    val hits: Long = 0,
+    val misses: Long = 0,
+    val prefetch_used: Long = 0,
+    val prefetch_mispredict: Long = 0,
+    val evictions: Long = 0,
+    val stall_ns: Long = 0,
+    val flash_bytes: Long = 0,
+    val slots_used: Long = 0,
+    val slots_total: Long = 0,
+)
+
+@Serializable
+data class RequestMetrics(
+    val schema: Int = 0,
+    @SerialName("prompt_tokens") val promptTokens: Int = 0,
+    val prefix_reused: Boolean = false,
+    val prefix_reused_tokens: Int = 0,
+    @SerialName("new_tokens") val newTokens: Int = 0,
+    @SerialName("prefill_ms") val prefillMs: Long = 0,
+    @SerialName("first_token_ms") val firstTokenMs: Long = 0,
+    @SerialName("decode_ms") val decodeMs: Long = 0,
+    @SerialName("prefill_tok_s") val prefillTokS: Double = 0.0,
+    @SerialName("decode_tok_s") val decodeTokS: Double = 0.0,
+    @SerialName("peak_rss_bytes") val peakRssBytes: Long = 0,
+    @SerialName("chat_turns") val chatTurns: Int = 0,
+    @SerialName("thinking_on") val thinkingOn: Boolean = true,
+    val cancelled: Boolean = false,
+    val error: String = "",
+    // additive echo section: missing fields on older engines fall back to N/A defaults
+    val config: ConfigEcho? = null,
+    val params: ParamsEcho? = null,
+    val route: RouteMetrics? = null,
+    val pool: PoolMetrics? = null,
+    val gpu: GpuMetrics? = null,
+) {
+    val hitRate: Double? get() = pool?.let {
+        val t = it.hits + it.misses
+        if (t > 0) it.hits.toDouble() / t else null
+    }
+}
+
+object MetricsMapper {
+    private val json = Json { ignoreUnknownKeys = true; isLenient = true }
+
+    /** Parse failure / empty input returns null (UI shows N/A, never crashes). */
+    fun parse(raw: String?): RequestMetrics? {
+        if (raw.isNullOrBlank()) return null
+        return try {
+            json.decodeFromString(RequestMetrics.serializer(), raw)
+        } catch (_: Exception) {
+            null
+        }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/ModelInventory.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/ModelInventory.kt
new file mode 100644
index 0000000..8773832
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/ModelInventory.kt
@@ -0,0 +1,51 @@
+// runtime/ModelInventory.kt - model discovery for the import/models screen.
+// Scans the app's model roots (internal files/models first; external files/models
+// as a secondary root - on some OEM Android builds shell-pushed files under
+// Android/data/<pkg> are invisible to the app uid, so in-app copy to the
+// internal dir is the primary staging channel). Only flat *.gguf files are
+// entries; lora_*.gguf companions are picked up by the runtime, not listed.
+package dev.edge0.runtime.app.runtime
+
+import java.io.File
+
+data class ModelEntry(
+    val name: String,
+    val dir: File,                 // the .gguf file itself
+    val totalBytes: Long,          // file size (requirement yardstick)
+    val dirBytes: Long,            // same, kept for the UI's size column
+    val valid: Boolean,
+    val hint: String?,             // reason shown when invalid
+    val mode: String = "llama-gguf",
+)
+
+object ModelInventory {
+
+    fun scanRoots(roots: List<File?>): List<ModelEntry> =
+        roots.filterNotNull().flatMap { scan(it) }
+             .distinctBy { it.name }
+             .sortedBy { it.name }
+
+    fun scan(root: File): List<ModelEntry> =
+        (root.listFiles() ?: emptyArray()).filter {
+            it.isFile && it.name.endsWith(".gguf") && !it.name.startsWith("lora_")
+        }.map { f ->
+            val ok = f.length() > 0
+            ModelEntry(f.name.removeSuffix(".gguf"), f, f.length(), f.length(), ok,
+                if (ok) null else "empty file")
+        }
+
+    /** Free space on the volume backing the primary model root (UI warning gate). */
+    fun volumeAvailBytes(root: File?): Long =
+        root?.let { runCatching { android.os.StatFs(it.absolutePath).availableBytes }
+                     .getOrDefault(0L) } ?: 0L
+
+    /** Import instructions shown on the models screen. */
+    fun importTemplate(pkg: String): String = listOf(
+        "Place model files flat under files/models/ of this app.",
+        "Option 1 - in-app: use the picker to copy from device storage.",
+        "Option 2 - adb (debug builds):",
+        "  adb push <file> /data/local/tmp/",
+        "  adb shell run-as $pkg cp /data/local/tmp/<file> files/models/<file>",
+        "Expected names: edge0-8b.gguf (plus lora_edge0_8b-gguf.gguf), edge0-35b.gguf.",
+    ).joinToString("\n")
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/Runtime.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/Runtime.kt
new file mode 100644
index 0000000..c3c5a05
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/Runtime.kt
@@ -0,0 +1,19 @@
+// runtime/Runtime.kt - the runtime facade contract consumed by ViewModel/screens.
+package dev.edge0.runtime.app.runtime
+
+import dev.edge0.runtime.engine.ChatMessage
+import dev.edge0.runtime.engine.GenParams
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+
+interface Runtime {
+    val state: MutableStateFlow<GenState>
+    val lastMetrics: MutableStateFlow<RequestMetrics?>
+    val activeModelDir: String?
+    suspend fun ensureLoaded(path: String)
+    suspend fun unloadActive()
+    suspend fun trim(ratio: Float): Boolean
+    fun submit(threadId: String, messages: List<ChatMessage>,
+               thinkingOn: Boolean, params: GenParams): Flow<GenEvent>
+    fun cancelActive()
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/ThermalSource.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/ThermalSource.kt
new file mode 100644
index 0000000..c4c4db9
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/ThermalSource.kt
@@ -0,0 +1,66 @@
+// runtime/ThermalSource.kt - dual-source thermal telemetry:
+// PowerManager.currentThermalStatus (API 29+ enum) plus /sys/class/thermal
+// nodes (non-root readability verified per node; unreadable ones degrade to N/A).
+package dev.edge0.runtime.app.runtime
+
+import android.content.Context
+import android.os.Build
+import android.os.PowerManager
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.flowOn
+import java.io.File
+
+data class ThermalReading(
+    val statusLevel: Int?,         // PowerManager level (null when unavailable)
+    val statusName: String?,
+    val topCpuGpuMilli: Int?,      // max across compute-thermal zones (cpu/gpuss/ddr/qfprom)
+    val tsEpochMs: Long,
+)
+
+class ThermalSource(context: Context) {
+    private val power = context.getSystemService(PowerManager::class.java)
+
+    fun read(): ThermalReading {
+        var lvl: Int? = null
+        var name: String? = null
+        try {
+            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
+                lvl = power.currentThermalStatus
+                name = when (lvl) {
+                    PowerManager.THERMAL_STATUS_NONE -> "NONE"
+                    PowerManager.THERMAL_STATUS_LIGHT -> "LIGHT"
+                    PowerManager.THERMAL_STATUS_MODERATE -> "MODERATE"
+                    PowerManager.THERMAL_STATUS_SEVERE -> "SEVERE"
+                    PowerManager.THERMAL_STATUS_CRITICAL -> "CRITICAL"
+                    PowerManager.THERMAL_STATUS_EMERGENCY -> "EMERGENCY"
+                    PowerManager.THERMAL_STATUS_SHUTDOWN -> "SHUTDOWN"
+                    else -> "S$lvl"
+                }
+            }
+        } catch (_: Exception) { /* degrade to N/A */ }
+        var max = -1
+        try {
+            val zones = File("/sys/class/thermal").listFiles() ?: emptyArray()
+            for (z in zones) {
+                if (!z.name.startsWith("thermal_zone")) continue
+                val type = try { File(z, "type").readText().trim() } catch (_: Exception) { continue }
+                if (!(type.contains("cpu", true) || type.contains("gpuss", true) ||
+                        type.contains("ddr", true) || type.contains("qfprom", true))) continue
+                val milli = try { File(z, "temp").readText().trim().toIntOrNull() } catch (_: Exception) { null }
+                if (milli != null && milli > max) max = milli
+            }
+        } catch (_: Exception) { /* unreadable: null */ }
+        return ThermalReading(lvl, name, if (max >= 0) max else null, System.currentTimeMillis())
+    }
+
+    /** 1 Hz sampling flow (collected only while the overlay is visible; unsubscribes when hidden). */
+    fun flow(): Flow<ThermalReading> = flow {
+        while (true) {
+            emit(read())
+            delay(1000)
+        }
+    }.flowOn(Dispatchers.IO)
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/runtime/ThinkSplitter.kt b/android/app/src/main/java/dev/edge0/runtime/app/runtime/ThinkSplitter.kt
new file mode 100644
index 0000000..0812705
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/runtime/ThinkSplitter.kt
@@ -0,0 +1,86 @@
+// runtime/ThinkSplitter.kt - streaming piece splitter that demuxes engine output
+// into body/thinking increments. The engine never strips think tags (template purity
+// contract), so parsing happens here. Tag literals are assembled in two pieces to
+// survive tooling that strips special-token lookalikes; verify bytes on disk after edits.
+package dev.edge0.runtime.app.runtime
+
+sealed interface Seg {
+    data class Text(val s: String) : Seg
+    data class Think(val s: String) : Seg
+}
+
+class ThinkSplitter(initialThink: Boolean = false) {
+    companion object {
+        const val OPEN = "<th" + "ink>"
+        const val CLOSE = "</th" + "ink>"
+        private const val MIN_OVERLAP_KEEP = 1
+    }
+
+    // when thinking_on the opening tag is injected at the prompt tail by the
+    // template, so the stream starts inside the thinking region - callers must pass
+    // initialThink=thinkingOn or thinking text leaks wholesale into the body.
+    private var inThink = initialThink
+    private var carry = "" // sole retained state: suffix of the watched tag that is a real prefix of it
+
+    /** Whether we are currently inside a thinking segment (drives the UI timer). */
+    val thinkingNow: Boolean get() = inThink
+
+    fun feed(piece: String): List<Seg> {
+        if (piece.isEmpty() && carry.isEmpty()) return emptyList()
+        val s = carry + piece
+        carry = ""
+        val out = ArrayList<Seg>(4)
+        var i = 0
+        while (i < s.length || i == 0) {
+            val tag = if (inThink) CLOSE else OPEN
+            val idx = s.indexOf(tag, i)
+            if (idx >= 0) {
+                if (idx > i) emit(out, s.substring(i, idx))
+                inThink = !inThink
+                i = idx + tag.length
+                if (i >= s.length) return out
+            } else {
+                val keep = overlapSuffixLen(s, tag)
+                if (keep > 0 && i <= s.length - keep) {
+                    if (s.length - i - keep > 0) emit(out, s.substring(i, s.length - keep))
+                    carry = s.substring(s.length - keep)
+                } else if (i < s.length) {
+                    emit(out, s.substring(i))
+                }
+                return out
+            }
+        }
+        return out
+    }
+
+    /** End of stream: flush the carried prefix per current mode (unclosed thinking is tolerated). */
+    fun flush(): List<Seg> {
+        if (carry.isEmpty()) return emptyList()
+        val out = listOf(segOf(carry))
+        carry = ""
+        return out
+    }
+
+    // maximal overlap of the tail with a real prefix of tag (length <= tag.size-1); cross-chunk tags live only here
+    private fun overlapSuffixLen(s: String, tag: String): Int {
+        val max = minOf(tag.length - 1, s.length)
+        for (k in max downTo MIN_OVERLAP_KEEP) {
+            if (s.endsWith(tag.substring(0, k))) return k
+        }
+        return 0
+    }
+
+    private fun emit(out: MutableList<Seg>, text: String) {
+        val seg = segOf(text)
+        // merge consecutive same-kind segments (keep event count minimal after flips within one feed)
+        val last = out.lastOrNull()
+        out.add(when {
+            last is Seg.Text && seg is Seg.Text -> { out.removeAt(out.size - 1); Seg.Text(last.s + seg.s) }
+            last is Seg.Think && seg is Seg.Think -> { out.removeAt(out.size - 1); Seg.Think(last.s + seg.s) }
+            else -> seg
+        })
+    }
+
+    private fun segOf(text: String): Seg =
+        if (inThink) Seg.Think(text) else Seg.Text(text)
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/AppRoot.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/AppRoot.kt
new file mode 100644
index 0000000..2a16ac5
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/AppRoot.kt
@@ -0,0 +1,206 @@
+// ui/AppRoot.kt - DrawerHost + NavHost wiring.
+// Side sheet (288dp): New Chat / thread list (updatedAt desc, busy pulse dot,
+// long-press delete menu) / footer with model management and settings.
+package dev.edge0.runtime.app.ui
+
+import androidx.compose.foundation.background
+import androidx.compose.foundation.clickable
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.PaddingValues
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.layout.width
+import androidx.compose.foundation.lazy.LazyColumn
+import androidx.compose.foundation.lazy.items
+import androidx.compose.foundation.shape.CircleShape
+import androidx.compose.material.icons.Icons
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material.icons.outlined.Close
+import androidx.compose.ui.draw.clip
+import androidx.compose.material.icons.outlined.Settings
+import androidx.compose.animation.core.RepeatMode
+import androidx.compose.animation.core.animateFloat
+import androidx.compose.animation.core.infiniteRepeatable
+import androidx.compose.animation.core.rememberInfiniteTransition
+import androidx.compose.animation.core.tween
+import androidx.compose.material3.DrawerValue
+import androidx.compose.material3.Icon
+import androidx.compose.material3.IconButton
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.ModalDrawerSheet
+import androidx.compose.material3.ModalNavigationDrawer
+import androidx.compose.material3.NavigationDrawerItem
+import androidx.compose.material3.Scaffold
+import androidx.compose.material3.Text
+import androidx.compose.material3.rememberDrawerState
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.draw.alpha
+import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.text.style.TextOverflow
+import androidx.compose.ui.unit.dp
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import androidx.lifecycle.viewmodel.compose.viewModel
+import androidx.navigation.compose.NavHost
+import androidx.navigation.compose.composable
+import androidx.navigation.compose.rememberNavController
+import dev.edge0.runtime.app.Edge0App
+import dev.edge0.runtime.app.data.AppSettings
+import dev.edge0.runtime.app.data.ThreadSummary
+import dev.edge0.runtime.app.runtime.GenState
+import dev.edge0.runtime.app.ui.chat.ChatScreen
+import dev.edge0.runtime.app.ui.chat.ChatViewModel
+import dev.edge0.runtime.app.ui.models.ModelsScreen
+import dev.edge0.runtime.app.ui.settings.SettingsScreen
+import kotlinx.coroutines.launch
+
+object Routes {
+    const val CHAT = "chat"
+    const val MODELS = "models"
+    const val SETTINGS = "settings"
+}
+
+@Composable
+fun AppRoot() {
+    val app = LocalContext.current.applicationContext as Edge0App
+    val container = app.container
+    val factory = remember {
+        object : ViewModelProvider.Factory {
+            override fun <T : androidx.lifecycle.ViewModel> create(
+                modelClass: Class<T>): T =
+                ChatViewModel(null, container.repository, container.runtime,
+                              container.settings) as T
+        }
+    }
+    val vm: ChatViewModel = viewModel(factory = factory)
+    val drawerState = rememberDrawerState(DrawerValue.Closed)
+    val nav = rememberNavController()
+    val scope = rememberCoroutineScope()
+    val threads by container.repository.observeThreads()
+        .collectAsStateWithLifecycle(initialValue = emptyList())
+    val genState by vm.genState.collectAsStateWithLifecycle()
+    val settings by container.settings.flow
+        .collectAsStateWithLifecycle(initialValue = AppSettings())
+
+    val closeDrawer: () -> Unit = { scope.launch { drawerState.close() } }
+
+    ModalNavigationDrawer(drawerState = drawerState, drawerContent = {
+        ModalDrawerSheet(modifier = Modifier.width(288.dp),
+                         drawerContainerColor = MaterialTheme.colorScheme.background) {
+            Column(Modifier.fillMaxSize()) {
+                // title row: name left, X closes the drawer (thread/model entries live in the top bar)
+                Row(Modifier.fillMaxWidth().padding(start = 20.dp, end = 8.dp,
+                                                    top = 28.dp, bottom = 14.dp),
+                    verticalAlignment = Alignment.CenterVertically) {
+                    Text("Edge0 Chat", style = MaterialTheme.typography.titleLarge,
+                         modifier = Modifier.weight(1f))
+                    IconButton(onClick = { closeDrawer() }) {
+                        Icon(Icons.Outlined.Close, "Close drawer",
+                             tint = MaterialTheme.colorScheme.onSurface)
+                    }
+                }
+                Box(Modifier.fillMaxWidth().height(1.dp)
+                    .background(MaterialTheme.colorScheme.outline))
+                Text("History", style = MaterialTheme.typography.labelLarge,
+                     color = MaterialTheme.colorScheme.onSurfaceVariant,
+                     modifier = Modifier.padding(start = 20.dp, top = 16.dp, bottom = 6.dp))
+                LazyColumn(Modifier.weight(1f),
+                           contentPadding = PaddingValues(horizontal = 10.dp)) {
+                    items(threads, key = { it.id }) { t ->
+                        ThreadItem(
+                            t = t,
+                            selected = vm.currentThreadId() == t.id,
+                            busyHere = genState is GenState.Streaming &&
+                                vm.currentThreadId() == t.id,
+                            onOpen = { vm.openThread(t.id); nav.navigate(Routes.CHAT)
+                                      closeDrawer() },
+                            onDelete = { vm.deleteThread(t.id) },
+                        )
+                    }
+                }
+                NavigationDrawerItem(label = { Text("Settings") },
+                    icon = { Icon(Icons.Outlined.Settings, null) }, selected = false,
+                    onClick = { nav.navigate(Routes.SETTINGS); closeDrawer() })
+                Spacer(Modifier.height(8.dp))
+            }
+        }
+    }) {
+        Scaffold { insets ->
+            NavHost(navController = nav, startDestination = Routes.CHAT,
+                    modifier = Modifier.padding(insets)) {
+                composable(Routes.CHAT) {
+                    Box(Modifier.fillMaxSize()) {
+                        ChatScreen(
+                            vm = vm,
+                            modelLabel = settings.activeModelDir.substringAfterLast('/'),
+                            onOpenDrawer = { scope.launch { drawerState.open() } },
+                            onOpenModels = { nav.navigate(Routes.MODELS) },
+                            onNewChat = { vm.newChat(); nav.navigate(Routes.CHAT) },
+                        )
+                    }
+                }
+                composable(Routes.MODELS) {
+                    ModelsScreen(onBack = { nav.popBackStack(Routes.CHAT, false) })
+                }
+                composable(Routes.SETTINGS) {
+                    SettingsScreen(onBack = { nav.popBackStack(Routes.CHAT, false) })
+                }
+            }
+        }
+    }
+}
+
+@Composable
+private fun ThreadItem(t: ThreadSummary, selected: Boolean, busyHere: Boolean,
+                       onOpen: () -> Unit, onDelete: () -> Unit) {
+    // row = rounded card (selected gets surfaceVariant backing); trailing X deletes
+    // immediately without confirmation; rename was retired by product decision.
+    Row(Modifier.fillMaxWidth().padding(vertical = 2.dp)
+            .clip(RoundedCornerShape(12.dp))
+            .background(if (selected) MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.55f)
+                        else androidx.compose.ui.graphics.Color.Transparent)
+            .clickable(onClick = onOpen)
+            .padding(start = 14.dp, end = 4.dp, top = 13.dp, bottom = 13.dp),
+        verticalAlignment = Alignment.CenterVertically) {
+        if (busyHere) {
+            val tr = rememberInfiniteTransition(label = "pulse")
+            val a by tr.animateFloat(0.25f, 1f,
+                infiniteRepeatable(tween(700), RepeatMode.Reverse), label = "p")
+            Box(Modifier.size(6.dp).background(
+                MaterialTheme.colorScheme.primary, CircleShape).alpha(a))
+            Spacer(Modifier.width(8.dp))
+        }
+        Text(t.title ?: "New chat", maxLines = 1, overflow = TextOverflow.Ellipsis,
+             style = MaterialTheme.typography.bodyMedium,
+             modifier = Modifier.weight(1f))
+        IconButton(onClick = onDelete, modifier = Modifier.size(36.dp)) {
+            Icon(Icons.Outlined.Close, "Delete chat",
+                 tint = MaterialTheme.colorScheme.onSurfaceVariant,
+                 modifier = Modifier.size(17.dp))
+        }
+    }
+}
+
+@Composable
+private fun PlaceholderScreen(name: String) {
+    Column(Modifier.fillMaxSize(), horizontalAlignment = Alignment.CenterHorizontally,
+        verticalArrangement = Arrangement.Center) {
+        Text(name, style = MaterialTheme.typography.bodyLarge)
+        Text("S4-Scaffold-OK", style = MaterialTheme.typography.bodyLarge,
+             color = MaterialTheme.colorScheme.error)
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ActionRow.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ActionRow.kt
new file mode 100644
index 0000000..8e82166
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ActionRow.kt
@@ -0,0 +1,81 @@
+// ui/chat/ActionRow.kt - per-message action row: copy / edit / delete
+// (+ regenerate on the last assistant turn) and a tok/s badge; hidden while streaming.
+package dev.edge0.runtime.app.ui.chat
+
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.layout.width
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.outlined.ContentCopy
+import androidx.compose.material.icons.outlined.Delete
+import androidx.compose.material.icons.outlined.Edit
+import androidx.compose.material.icons.outlined.Refresh
+import androidx.compose.material3.Icon
+import androidx.compose.material3.IconButton
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.platform.LocalClipboardManager
+import androidx.compose.ui.text.AnnotatedString
+import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
+import dev.edge0.runtime.app.ui.theme.LocalFontScale
+
+@Composable
+fun ActionRow(
+    text: String,
+    canEdit: Boolean,
+    canRegenerate: Boolean,
+    onEdit: () -> Unit,
+    onDelete: () -> Unit,
+    onRegenerate: () -> Unit,
+    modifier: Modifier = Modifier,
+    endAligned: Boolean = false,
+) {
+    val clip = LocalClipboardManager.current
+    Row(modifier.padding(vertical = 2.dp), verticalAlignment = Alignment.CenterVertically) {
+        if (endAligned) Spacer(Modifier.weight(1f))
+        IconButton(onClick = { clip.setText(AnnotatedString(text)) },
+                   modifier = Modifier.width(36.dp)) {
+            Icon(Icons.Outlined.ContentCopy, "Copy", tint = MaterialTheme.colorScheme.onSurfaceVariant,
+                 modifier = Modifier.size(16.dp))
+        }
+        if (canEdit) IconButton(onClick = onEdit, modifier = Modifier.width(36.dp)) {
+            Icon(Icons.Outlined.Edit, "Edit", tint = MaterialTheme.colorScheme.onSurfaceVariant,
+                 modifier = Modifier.size(16.dp))
+        }
+        IconButton(onClick = onDelete, modifier = Modifier.width(36.dp)) {
+            Icon(Icons.Outlined.Delete, "Delete", tint = MaterialTheme.colorScheme.onSurfaceVariant,
+                 modifier = Modifier.size(16.dp))
+        }
+        if (canRegenerate) IconButton(onClick = onRegenerate, modifier = Modifier.width(36.dp)) {
+            Icon(Icons.Outlined.Refresh, "Regenerate", tint = MaterialTheme.colorScheme.onSurfaceVariant,
+                 modifier = Modifier.size(16.dp))
+        }
+        if (!endAligned) Spacer(Modifier.weight(1f))
+    }
+}
+
+/** Inline stats line under each reply (Atomic style). Memory segment reports
+ *  resident pool bytes only (cacheable mmap views excluded). Missing fields elide. */
+@Composable
+fun StatsLine(newTokens: Int?, ttftMs: Long?, prefillTokS: Double?,
+              decodeTokS: Double?, memBytes: Long?,
+              modifier: Modifier = Modifier) {
+    val parts = ArrayList<String>(5)
+    if (newTokens != null && newTokens > 0) parts.add("$newTokens tokens")
+    if (ttftMs != null && ttftMs > 0) parts.add("TTFT %.1fs".format(ttftMs / 1000.0))
+    if (prefillTokS != null && prefillTokS > 0.05) parts.add("prefill %.1f tok/s".format(prefillTokS))
+    if (decodeTokS != null && decodeTokS > 0) parts.add("decode %.2f tok/s".format(decodeTokS))
+    if (memBytes != null && memBytes > 0) parts.add("pool %.0f MB".format(memBytes / 1048576.0))
+    if (parts.isEmpty()) return
+    Text(parts.joinToString("  ·  "),
+         style = MaterialTheme.typography.bodySmall.copy(
+            fontSize = 11.5.sp * LocalFontScale.current,
+            color = MaterialTheme.colorScheme.onSurfaceVariant),
+         modifier = modifier.padding(top = 6.dp))
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ChatScreen.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ChatScreen.kt
new file mode 100644
index 0000000..5f6e516
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ChatScreen.kt
@@ -0,0 +1,367 @@
+// ui/chat/ChatScreen.kt - the chat page: top bar (drawer button left, context usage
+// right), centered message column (max 768dp) sticking to bottom, user bubbles on the
+// right, full-width assistant markdown + reasoning, action row on the last message,
+// Working... shimmer, destructive error card, rounded composer card.
+package dev.edge0.runtime.app.ui.chat
+
+import androidx.compose.foundation.BorderStroke
+import androidx.compose.foundation.clickable
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.imePadding
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.layout.width
+import androidx.compose.foundation.layout.widthIn
+import androidx.compose.foundation.lazy.LazyColumn
+import androidx.compose.foundation.lazy.items
+import androidx.compose.foundation.lazy.rememberLazyListState
+import androidx.compose.foundation.shape.CircleShape
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.foundation.text.BasicTextField
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.outlined.ArrowDownward
+import androidx.compose.material.icons.outlined.Close
+import androidx.compose.material.icons.outlined.Edit
+import androidx.compose.material.icons.outlined.Menu
+import androidx.compose.material.icons.outlined.SwapHoriz
+import androidx.compose.foundation.background
+import androidx.compose.ui.graphics.Color
+import androidx.compose.material.icons.outlined.Warning
+import androidx.compose.material3.FloatingActionButton
+import androidx.compose.material3.Icon
+import androidx.compose.material3.IconButton
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Surface
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.derivedStateOf
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberCoroutineScope
+import kotlinx.coroutines.launch
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.draw.clip
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.text.font.FontWeight
+import androidx.compose.ui.text.style.TextOverflow
+import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import dev.edge0.runtime.app.data.MessageView
+import dev.edge0.runtime.app.data.MsgStatus
+import dev.edge0.runtime.app.ui.markdown.MarkdownText
+import dev.edge0.runtime.app.ui.theme.LocalFontScale
+
+@Composable
+fun ChatScreen(
+    vm: ChatViewModel,
+    modelLabel: String,
+    onOpenDrawer: () -> Unit,
+    onOpenModels: () -> Unit,
+    onNewChat: () -> Unit,
+) {
+    val draft by vm.draft.collectAsStateWithLifecycle()
+    val streaming by vm.streaming.collectAsStateWithLifecycle()
+    val genState by vm.genState.collectAsStateWithLifecycle()
+    val banner by vm.banner.collectAsStateWithLifecycle()
+    val messages by vm.messages.collectAsStateWithLifecycle()
+
+    val busy = genState is dev.edge0.runtime.app.runtime.GenState.Streaming ||
+        genState is dev.edge0.runtime.app.runtime.GenState.Cancelling || streaming != null
+    val loading = genState is dev.edge0.runtime.app.runtime.GenState.Loading
+
+    val listState = rememberLazyListState()
+    val scrollScope = rememberCoroutineScope()
+    val atBottom by remember {
+        derivedStateOf {
+            val info = listState.layoutInfo
+            val last = info.visibleItemsInfo.lastOrNull()
+            last == null || last.index >= info.totalItemsCount - 1
+        }
+    }
+    // stick to bottom: streaming deltas and new messages follow; a user scroll up detaches without interrupting generation
+    LaunchedEffect(messages.size, streaming?.text, streaming?.thinking) {
+        if (atBottom) listState.scrollToItem(Int.MAX_VALUE / 2)
+    }
+
+    Column(Modifier.fillMaxSize().imePadding()) {
+        // -- top bar --
+        Row(Modifier.fillMaxWidth().height(52.dp).padding(horizontal = 8.dp),
+            verticalAlignment = Alignment.CenterVertically) {
+            IconButton(onClick = onOpenDrawer) {
+                Icon(Icons.Outlined.Menu, "Chats")
+            }
+            Spacer(Modifier.weight(1f))
+            // centered title: perceived model name + On-device Inference subtitle
+            Column(horizontalAlignment = Alignment.CenterHorizontally) {
+                val short = modelLabel.lowercase()
+                val title = when {
+                    "35b" in short -> "Edge0 35B"
+                    "8b" in short -> "Edge0 8B"
+                    else -> "Edge0 Chat"
+                }
+                Text(title, style = MaterialTheme.typography.titleMedium.copy(
+                    fontWeight = FontWeight.Bold, fontSize = 16.sp * LocalFontScale.current))
+                Text("On-device Inference", style = MaterialTheme.typography.labelSmall.copy(
+                    fontSize = 11.sp * LocalFontScale.current,
+                    color = MaterialTheme.colorScheme.onSurfaceVariant))
+            }
+            Spacer(Modifier.weight(1f))
+            // new-thread button (lifted out of the drawer, beside the model switch)
+            Box(Modifier.size(40.dp), contentAlignment = Alignment.Center) {
+                Surface(shape = CircleShape,
+                        color = MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.5f),
+                        modifier = Modifier.size(40.dp).clickable(onClick = onNewChat)) {}
+                Icon(Icons.Outlined.Edit, "New chat",
+                     modifier = Modifier.size(20.dp), tint = MaterialTheme.colorScheme.onSurface)
+            }
+            Spacer(Modifier.width(12.dp))
+            // top-right model switch (opens the models page)
+            Box(Modifier.size(40.dp), contentAlignment = Alignment.Center) {
+                Surface(shape = CircleShape,
+                        color = MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.5f),
+                        modifier = Modifier.size(40.dp).clickable(onClick = onOpenModels)) {}
+                Icon(Icons.Outlined.SwapHoriz, "Switch model ($modelLabel)",
+                     modifier = Modifier.size(20.dp),
+                     tint = MaterialTheme.colorScheme.onSurface)
+                val dot = when {
+                    loading -> MaterialTheme.colorScheme.tertiary
+                    busy -> MaterialTheme.colorScheme.primary
+                    modelLabel.isNotBlank() -> Color(0xFF22C55E)
+                    else -> MaterialTheme.colorScheme.outline
+                }
+                Box(Modifier.align(Alignment.TopEnd).padding(3.dp).size(8.dp)
+                    .background(dot, CircleShape))
+            }
+            Spacer(Modifier.width(8.dp))
+        }
+
+        // -- message column --
+        Box(Modifier.weight(1f), contentAlignment = Alignment.TopCenter) {
+            LazyColumn(state = listState, modifier = Modifier.fillMaxSize(),
+                       contentPadding = androidx.compose.foundation.layout.PaddingValues(
+                           start = 16.dp, end = 16.dp, top = 4.dp, bottom = 8.dp),
+                       horizontalAlignment = Alignment.CenterHorizontally) {
+                if (messages.isEmpty() && streaming == null) {
+                    item {
+                        WelcomePage(onPick = { pr -> vm.draft.value = pr; vm.sendMessage() })
+                    }
+                }
+                items(messages.size) { i ->
+                    val m = messages[i]
+                    MessageRow(
+                        m = m,
+                        isLastAssistant = m.role == "assistant" &&
+                            i == messages.indexOfLast { it.role == "assistant" },
+                        showActions = !busy && m.status != MsgStatus.STREAMING,
+                        onEdit = { text -> vm.editUserMessage(m.id, text) },
+                        onDelete = { vm.deleteMessage(m.id) },
+                        onRegenerate = { vm.regenerate() },
+                    )
+                }
+                streaming?.let { s ->
+                    item {
+                        AssistantBody(
+                            body = s.text, thinking = s.thinking,
+                            reasoning = ReasoningState(
+                                active = s.thinkingActive, startedAt = s.thinkStartedMs,
+                                endedAt = s.thinkEndedMs),
+                            prefilling = s.prefilling,
+                            status = null, isStreaming = true,
+                        )
+                    }
+                }
+                item { Spacer(Modifier.height(2.dp)) }
+            }
+            if (!atBottom) {
+                FloatingActionButton(
+                    onClick = { scrollScope.launch { listState.animateScrollToItem(Int.MAX_VALUE / 2) } },
+                    modifier = Modifier.align(Alignment.BottomEnd).padding(12.dp)
+                        .size(36.dp),
+                    shape = CircleShape, containerColor = MaterialTheme.colorScheme.surfaceVariant,
+                ) { Icon(Icons.Outlined.ArrowDownward, "Scroll to bottom",
+                        modifier = Modifier.size(18.dp)) }
+            }
+        }
+
+        // -- error / guidance banner --
+        banner?.let { msg ->
+            Surface(Modifier.fillMaxWidth().padding(horizontal = 12.dp, vertical = 4.dp),
+                    shape = RoundedCornerShape(10.dp),
+                    color = MaterialTheme.colorScheme.errorContainer.copy(alpha = 0.35f),
+                    border = BorderStroke(1.dp, MaterialTheme.colorScheme.error.copy(alpha = 0.4f))
+            ) {
+                Row(Modifier.padding(10.dp), verticalAlignment = Alignment.CenterVertically) {
+                    Icon(Icons.Outlined.Warning, null, tint = MaterialTheme.colorScheme.error,
+                         modifier = Modifier.size(16.dp))
+                    Spacer(Modifier.width(8.dp))
+                    Text(msg, modifier = Modifier.weight(1f), maxLines = 3,
+                         overflow = TextOverflow.Ellipsis,
+                         style = MaterialTheme.typography.bodySmall)
+                    IconButton(onClick = { vm.banner.value = null }) {
+                        Icon(Icons.Outlined.Close, "Dismiss", modifier = Modifier.size(14.dp))
+                    }
+                }
+            }
+        }
+
+        // -- composer (input card + send button; other controls moved to Settings/top bar) --
+        Composer(
+            value = draft, onValueChange = { vm.draft.value = it },
+            busy = busy,
+            enabledSend = draft.isNotBlank() && !loading && modelLabel.isNotBlank(),
+            onSend = { vm.sendMessage() }, onStop = { vm.cancel() },
+        )
+    }
+}
+
+@Composable
+private fun MessageRow(m: MessageView, isLastAssistant: Boolean, showActions: Boolean,
+                       onEdit: (String) -> Unit, onDelete: () -> Unit, onRegenerate: () -> Unit) {
+    if (m.role == "user") {
+        var editing by remember(m.id) { mutableStateOf(false) }
+        var buf by remember(m.id) { mutableStateOf(m.content) }
+        Column(Modifier.fillMaxWidth().padding(top = 10.dp),
+               horizontalAlignment = Alignment.End) {
+            if (editing) {
+                Surface(Modifier.fillMaxWidth(0.92f), shape = RoundedCornerShape(10.dp),
+                        border = BorderStroke(1.dp, MaterialTheme.colorScheme.outline)) {
+                    Column(Modifier.padding(8.dp)) {
+                        BasicTextField(value = buf, onValueChange = { buf = it },
+                                       textStyle = MaterialTheme.typography.bodyLarge.copy(
+                                           color = MaterialTheme.colorScheme.onSurface),
+                                       modifier = Modifier.fillMaxWidth().height(88.dp))
+                        Row(horizontalArrangement = Arrangement.End,
+                            modifier = Modifier.fillMaxWidth()) {
+                            Text("Cancel", color = MaterialTheme.colorScheme.onSurfaceVariant,
+                                 modifier = Modifier.padding(8.dp)
+                                     .clickableText { editing = false })
+                            Text("Save", color = MaterialTheme.colorScheme.primary,
+                                 modifier = Modifier.padding(8.dp).clickableText {
+                                     if (buf.isNotBlank()) onEdit(buf.trim())
+                                     editing = false
+                                 })
+                        }
+                    }
+                }
+            } else {
+                Surface(shape = RoundedCornerShape(18.dp),
+                        color = MaterialTheme.colorScheme.secondary,
+                        modifier = Modifier.widthIn(max = 560.dp)) {
+                    Text(m.content, Modifier.padding(horizontal = 12.dp, vertical = 8.dp),
+                         style = MaterialTheme.typography.bodyLarge.copy(
+                            fontSize = 14.sp * LocalFontScale.current))
+                }
+                if (showActions) {
+                    ActionRow(text = m.content, canEdit = true, canRegenerate = false,
+                              onEdit = { editing = true }, onDelete = onDelete,
+                              onRegenerate = {}, endAligned = true)
+                }
+            }
+        }
+    } else {
+        AssistantBody(
+            body = m.content, thinking = m.thinking ?: "",
+            // persisted row: startedAt at assistant-row creation, endedAt plus thinking duration
+            reasoning = ReasoningState(active = false, startedAt = m.createdAt,
+                                       endedAt = m.thinkingMs?.let { m.createdAt + it }),
+            prefilling = false,
+            status = m.status, isStreaming = false,
+            extra = {
+                if (showActions) Row {
+                    Column(Modifier.weight(1f)) {
+                        StatsLine(newTokens = m.newTokens, ttftMs = m.ttftMs,
+                                  prefillTokS = m.prefillTokS, decodeTokS = m.decodeTokS,
+                                  memBytes = m.memBytes)
+                        ActionRow(text = m.content, canEdit = false,
+                                  canRegenerate = isLastAssistant,
+                                  onEdit = {}, onDelete = onDelete,
+                                  onRegenerate = onRegenerate)
+                    }
+                }
+            },
+        )
+    }
+}
+
+private data class ReasoningState(val active: Boolean, val startedAt: Long?, val endedAt: Long?)
+
+@Composable
+private fun AssistantBody(body: String, thinking: String, reasoning: ReasoningState,
+                          prefilling: Boolean, status: String?, isStreaming: Boolean,
+                          extra: (@Composable () -> Unit)? = null) {
+    Column(Modifier.fillMaxWidth().widthIn(max = 768.dp).padding(top = 10.dp)) {
+        if (thinking.isNotEmpty()) {
+            ReasoningBlock(thinking = thinking, active = reasoning.active,
+                           startedAtMs = reasoning.startedAt, endedAtMs = reasoning.endedAt,
+                           nowMs = System.currentTimeMillis())
+        }
+        if (prefilling && isStreaming) {
+            WorkingShimmer("Working...")
+        }
+        if (body.isNotEmpty()) {
+            MarkdownText(body, Modifier.fillMaxWidth())
+        }
+        status?.let { st ->
+            val chip = when (st) {
+                MsgStatus.CANCELLED -> "Cancelled"
+                MsgStatus.INTERRUPTED -> "Interrupted"
+                MsgStatus.ERROR -> "Error"
+                else -> null
+            }
+            if (chip != null) {
+                Text(chip, style = MaterialTheme.typography.labelSmall.copy(
+                    color = MaterialTheme.colorScheme.error),
+                     modifier = Modifier.padding(top = 2.dp))
+            }
+        }
+        extra?.invoke()
+    }
+}
+
+// lightweight text button
+private fun Modifier.clickableText(onClick: () -> Unit): Modifier =
+    this.clickable(onClick = onClick)
+
+// starter prompts (three fixed questions; display truncates with ellipsis, send uses full text)
+private val PRESET_PROMPTS = listOf(
+    "Will the water level rise when ice floating in it melts? Explain in two sentences.",
+    "Write a sci-fi story with a twist in exactly six words.",
+    "Make 24 using 3, 3, 8, and 8 exactly once each, with only basic arithmetic and parentheses.",
+)
+
+@Composable
+private fun WelcomePage(onPick: (String) -> Unit) {
+    Column(Modifier.fillMaxWidth().padding(top = 110.dp),
+           horizontalAlignment = Alignment.CenterHorizontally) {
+        Text("edge0", style = MaterialTheme.typography.displaySmall.copy(
+            fontWeight = FontWeight.Bold, fontSize = 34.sp * LocalFontScale.current))
+        Spacer(Modifier.height(10.dp))
+        Text("Private AI, running locally on your device",
+             style = MaterialTheme.typography.bodyMedium.copy(
+                color = MaterialTheme.colorScheme.onSurfaceVariant))
+        Spacer(Modifier.height(30.dp))
+        PRESET_PROMPTS.forEach { pr ->
+            Surface(Modifier.padding(vertical = 5.dp).widthIn(max = 320.dp)
+                    .clip(RoundedCornerShape(20.dp))
+                    .clickable { onPick(pr) },
+                    shape = RoundedCornerShape(20.dp),
+                    color = MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.5f)) {
+                Text(pr, Modifier.padding(horizontal = 16.dp, vertical = 10.dp),
+                     maxLines = 1, overflow = TextOverflow.Ellipsis,
+                     style = MaterialTheme.typography.bodyMedium.copy(
+                        fontSize = 13.5.sp * LocalFontScale.current))
+            }
+        }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ChatViewModel.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ChatViewModel.kt
new file mode 100644
index 0000000..d8aaea8
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ChatViewModel.kt
@@ -0,0 +1,220 @@
+// ui/chat/ChatViewModel.kt - chat page orchestration: send -> load -> collect
+// generation events -> persist terminal state. Engine cancellation goes exclusively
+// through runtime.cancelActive (no fake interruption in the VM). The engine
+// overflow diagnostic (INVALID_ARG "prompt tokens=") triggers one retry with a
+// tighter budget. Both Done and Cancelled finalize to storage (a cancelled bubble
+// keeps its partial content).
+package dev.edge0.runtime.app.ui.chat
+
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.viewModelScope
+import dev.edge0.runtime.app.data.AppSettings
+import dev.edge0.runtime.app.data.ChatRepository
+import dev.edge0.runtime.app.data.MessageView
+import dev.edge0.runtime.app.data.MsgStatus
+import dev.edge0.runtime.app.data.SettingsStore
+import dev.edge0.runtime.app.runtime.Runtime
+import dev.edge0.runtime.app.runtime.GenEvent
+import dev.edge0.runtime.app.runtime.GenState
+import dev.edge0.runtime.app.runtime.encodeIds
+import dev.edge0.runtime.app.runtime.HistoryWindow
+import dev.edge0.runtime.app.runtime.RequestMetrics
+import dev.edge0.runtime.engine.GenParams
+import dev.edge0.runtime.engine.Status
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+
+data class StreamingDraft(
+    val text: String = "",
+    val thinking: String = "",
+    val prefilling: Boolean = true,   // before the first event: show the Working... shimmer
+    val thinkingActive: Boolean = false,
+    val thinkStartedMs: Long? = null,
+    val thinkEndedMs: Long? = null,
+)
+
+class ChatViewModel(
+    initialThreadId: String?,
+    private val repo: ChatRepository,
+    private val runtime: Runtime,
+    private val settingsStore: SettingsStore,
+) : ViewModel() {
+
+    val draft = MutableStateFlow("")
+    val streaming = MutableStateFlow<StreamingDraft?>(null)
+    val banner = MutableStateFlow<String?>(null)
+    val genState: StateFlow<GenState> = runtime.state
+    val lastMetrics: StateFlow<RequestMetrics?> = runtime.lastMetrics
+
+    private val activeThreadId = MutableStateFlow(initialThreadId)
+    fun currentThreadId(): String? = activeThreadId.value
+
+    /** Message stream of the current thread (re-follows on thread switch). */
+    @OptIn(ExperimentalCoroutinesApi::class)
+    val messages: StateFlow<List<MessageView>> = activeThreadId
+        .flatMapLatest { tid -> if (tid == null) flowOf(emptyList()) else repo.observeMessages(tid) }
+        .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5_000), emptyList())
+
+    fun newChat() {
+        activeThreadId.value = null
+        draft.value = ""
+        banner.value = null
+    }
+
+    fun openThread(id: String) {
+        if (runtime.state.value is GenState.Streaming) return // no thread switch mid-generation (single slot)
+        activeThreadId.value = id
+        banner.value = null
+    }
+
+    fun sendMessage() {
+        val userText = draft.value.trim()
+        if (userText.isEmpty()) return
+        val st = runtime.state.value
+        if (st is GenState.Streaming || st is GenState.Cancelling || streaming.value != null) return
+        draft.value = ""
+        viewModelScope.launch {
+            val settings = runCatching { settingsStore.flow.first() }.getOrNull() ?: return@launch
+            val existing = activeThreadId.value
+            val tid = existing ?: repo.startThreadWithUserMessage(
+                runtime.activeModelDir ?: settings.activeModelDir, userText).also {
+                activeThreadId.value = it
+            }
+            if (existing != null) repo.addUserMessage(tid, userText)
+
+            if (runtime.activeModelDir == null) {
+                if (settings.activeModelDir.isBlank()) {
+                    banner.value = "No model loaded: import one into the app (adb push) via Models, then load it"
+                    return@launch
+                }
+                try {
+                    runtime.ensureLoaded(settings.activeModelDir)
+                } catch (e: Exception) {
+                    banner.value = "Load failed: ${e.message}"
+                    return@launch
+                }
+            }
+            runGenerate(tid, settings, budget = HistoryWindow.DEFAULT_BUDGET_CHARS,
+                        allowRetry = true)
+        }
+    }
+
+    /** One generation round (with the overflow-trim retry). The last history entry is the new user turn. */
+    private suspend fun runGenerate(tid: String, settings: AppSettings,
+                                    budget: Int, allowRetry: Boolean) {
+        val history = repo.history(tid)
+        val msgs = HistoryWindow.build(
+            settings.systemPrompt.takeIf { settings.useSystemPrompt }?.ifBlank { null },
+            history, budget)
+        val mid = repo.beginAssistant(tid, runtime.activeModelDir ?: "")
+        streaming.value = StreamingDraft()
+        val text = StringBuilder()
+        val think = StringBuilder()
+        var thinkStart = 0L
+        var thinkMs: Long? = null
+        var terminal: GenEvent? = null
+        val params = settings.toGenParams()
+        runtime.submit(tid, msgs, settings.enableThinking, params).collect { ev ->
+            when (ev) {
+                is GenEvent.ThinkPiece -> {
+                    if (think.isEmpty()) thinkStart = System.currentTimeMillis()
+                    think.append(ev.s)
+                    streaming.value = (streaming.value ?: StreamingDraft()).copy(
+                        prefilling = false, thinkingActive = true,
+                        thinking = think.toString(), thinkStartedMs = thinkStart)
+                }
+                is GenEvent.TextPiece -> {
+                    if (streaming.value?.thinkingActive == true && thinkMs == null) {
+                        thinkMs = System.currentTimeMillis() - thinkStart
+                    }
+                    text.append(ev.s)
+                    streaming.value = (streaming.value ?: StreamingDraft()).copy(
+                        prefilling = false, thinkingActive = false,
+                        text = text.toString(), thinking = think.toString(),
+                        thinkEndedMs = thinkMs)
+                }
+                else -> terminal = ev
+            }
+        }
+        streaming.value = null
+        val metrics = (terminal as? GenEvent.Done)?.metrics
+            ?: (terminal as? GenEvent.Cancelled)?.metrics
+        // timing fallback: stopped/cancelled turns never flip THINK to TEXT -
+        // freeze the clock at the terminal moment. Unclosed thinking is kept whole in
+        // the thinking region (same semantics as the upstream split helper).
+        var thinkMsFinal = thinkMs
+        if (thinkMsFinal == null && think.isNotEmpty() && thinkStart > 0) {
+            thinkMsFinal = System.currentTimeMillis() - thinkStart
+        }
+        val status = when (terminal) {
+            is GenEvent.Done -> MsgStatus.OK
+            is GenEvent.Cancelled -> MsgStatus.CANCELLED
+            else -> MsgStatus.ERROR
+        }
+        // the sidecar persists only on successful turns (null on cancel/failure)
+        repo.finalizeAssistant(mid, text.toString(), think.toString().ifBlank { null },
+                               thinkMsFinal, status, metrics?.promptTokens, metrics?.newTokens,
+                               metrics?.decodeTokS,
+                               (terminal as? GenEvent.Done)?.genIds?.takeIf { it.isNotEmpty() }
+                                   ?.let { encodeIds(it) },
+                               ttftMs = metrics?.firstTokenMs?.takeIf { it > 0 },
+                               prefillTokS = metrics?.prefillTokS?.takeIf { it > 0.05 },
+                               // report resident pool usage only (cacheable mmap views excluded)
+                               memBytes = metrics?.pool?.resident_bytes?.takeIf { it > 0 })
+        val fail = terminal as? GenEvent.Failed
+        if (fail != null) {
+            if (allowRetry && fail.status == Status.INVALID_ARG &&
+                fail.diag.contains("prompt tokens=")) {
+                runGenerate(tid, settings,
+                            HistoryWindow.tighterBudget(budget), allowRetry = false)
+            } else {
+                banner.value = "Generation failed[${fail.status}]: ${fail.diag}"
+            }
+        }
+    }
+
+    fun cancel() { runtime.cancelActive() }
+
+    fun regenerate() {
+        val st = runtime.state.value
+        if (st !is GenState.Idle || streaming.value != null) return
+        viewModelScope.launch {
+            val tid = activeThreadId.value ?: return@launch
+            val rows = repo.historyWithIds(tid)
+            val lastAssistant = rows.lastOrNull { it.role == "assistant" } ?: return@launch
+            repo.deleteFrom(tid, lastAssistant.createdAt)
+            val settings = settingsStore.flow.first()
+            runGenerate(tid, settings, HistoryWindow.DEFAULT_BUDGET_CHARS, allowRetry = true)
+        }
+    }
+
+    fun deleteMessage(id: String) = viewModelScope.launch { repo.deleteMessage(id) }
+
+    fun editUserMessage(id: String, newText: String) = viewModelScope.launch {
+        repo.editUserMessage(id, newText)
+    }
+
+    fun renameThread(id: String, title: String) =
+        viewModelScope.launch { repo.renameThread(id, title) }
+
+    fun deleteThread(id: String) = viewModelScope.launch {
+        repo.deleteThread(id)
+        if (activeThreadId.value == id) newChat()
+    }
+}
+
+private fun AppSettings.toGenParams() =
+    GenParams(temperature = temperature, topK = topK, topP = topP, seed = seed,
+              // reasoning models need more than the 256-token engine default (observed truncation on a "hello" turn)
+              maxNewTokens = if (maxNewTokens > 0) maxNewTokens else 1024,
+              // repetition penalty follows the settings panel; first-token greedy is fixed upstream behavior
+              repetitionPenalty = repetitionPenalty,
+              firstTokenGreedy = true)
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/Composer.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/Composer.kt
new file mode 100644
index 0000000..f6eebc1
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/Composer.kt
@@ -0,0 +1,102 @@
+// ui/chat/Composer.kt - rounded input card plus send button, nothing else.
+// Temperature/thinking live in Settings; model switching in the top bar.
+// IME-safe: Enter always inserts a newline; sending is button-only.
+package dev.edge0.runtime.app.ui.chat
+
+import androidx.compose.foundation.background
+import androidx.compose.foundation.clickable
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.defaultMinSize
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.heightIn
+import androidx.compose.foundation.layout.imePadding
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.layout.width
+import androidx.compose.foundation.rememberScrollState
+import androidx.compose.foundation.shape.CircleShape
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.foundation.text.BasicTextField
+import androidx.compose.foundation.verticalScroll
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.filled.ArrowUpward
+import androidx.compose.material.icons.filled.Stop
+import androidx.compose.material3.Icon
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Surface
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.SolidColor
+import androidx.compose.ui.text.TextStyle
+import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
+import dev.edge0.runtime.app.ui.theme.LocalFontScale
+
+@Composable
+fun Composer(
+    value: String,
+    onValueChange: (String) -> Unit,
+    busy: Boolean,           // streaming/cancelling shows the stop button
+    enabledSend: Boolean,
+    onSend: () -> Unit,
+    onStop: () -> Unit,
+    modifier: Modifier = Modifier,
+) {
+    Row(modifier.fillMaxWidth().imePadding()
+        .padding(horizontal = 12.dp, vertical = 8.dp),
+        verticalAlignment = Alignment.Bottom,
+        horizontalArrangement = Arrangement.spacedBy(10.dp)) {
+        Surface(Modifier.weight(1f),
+                shape = RoundedCornerShape(26.dp),
+                color = MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.45f),
+                border = null) {
+            BasicTextField(
+                value = value, onValueChange = onValueChange,
+                textStyle = TextStyle(
+                    color = MaterialTheme.colorScheme.onSurface,
+                    fontSize = 15.sp * LocalFontScale.current,
+                    lineHeight = 21.sp * LocalFontScale.current),
+                cursorBrush = SolidColor(MaterialTheme.colorScheme.primary),
+                modifier = Modifier.fillMaxWidth()
+                    .defaultMinSize(minHeight = 26.dp)
+                    .heightIn(max = 168.dp) // ~8 lines
+                    .verticalScroll(rememberScrollState())
+                    .padding(horizontal = 16.dp, vertical = 12.dp),
+                decorationBox = { inner ->
+                    Box(contentAlignment = Alignment.CenterStart) {
+                        if (value.isEmpty()) {
+                            Text("Ask anything...",
+                                 color = MaterialTheme.colorScheme.onSurfaceVariant,
+                                 fontSize = 15.sp * LocalFontScale.current)
+                        }
+                        inner()
+                    }
+                })
+        }
+        Spacer(Modifier.width(0.dp))
+        val canSend = busy || enabledSend
+        val bg = when {
+            busy -> MaterialTheme.colorScheme.error
+            canSend -> MaterialTheme.colorScheme.primary
+            else -> MaterialTheme.colorScheme.surfaceVariant
+        }
+        val fg = when {
+            busy -> Color.White
+            canSend -> MaterialTheme.colorScheme.onPrimary
+            else -> MaterialTheme.colorScheme.onSurfaceVariant
+        }
+        Box(Modifier.size(48.dp).background(bg, CircleShape)
+            .clickable(enabled = canSend) { if (busy) onStop() else onSend() },
+            contentAlignment = Alignment.Center) {
+            Icon(if (busy) Icons.Filled.Stop else Icons.Filled.ArrowUpward,
+                 contentDescription = if (busy) "Stop" else "Send",
+                 tint = fg, modifier = Modifier.size(22.dp))
+        }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ReasoningBlock.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ReasoningBlock.kt
new file mode 100644
index 0000000..ae6b445
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/ReasoningBlock.kt
@@ -0,0 +1,67 @@
+// ui/chat/ReasoningBlock.kt - collapsible reasoning card: shows
+// "Thinking for Ns..." while streaming, "Thought for Ns" when done; collapsed by default.
+package dev.edge0.runtime.app.ui.chat
+
+import androidx.compose.animation.AnimatedVisibility
+import androidx.compose.foundation.clickable
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.automirrored.filled.KeyboardArrowRight
+import androidx.compose.material3.Icon
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Surface
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.text.font.FontStyle
+import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
+import dev.edge0.runtime.app.ui.theme.LocalFontScale
+
+@Composable
+fun ReasoningBlock(thinking: String, active: Boolean, startedAtMs: Long?, endedAtMs: Long?,
+                   nowMs: Long, modifier: Modifier = Modifier) {
+    if (thinking.isEmpty() && !active) return
+    var expanded by remember { mutableStateOf(false) }
+    val header = if (active) {
+        "Thinking for %.0fs...".format((nowMs - (startedAtMs ?: nowMs)) / 1000.0)
+    } else if (startedAtMs == null && endedAtMs == null) {
+        "Thought"   // legacy rows without a timing basis: qualitative label only
+    } else {
+        "Thought for %.0fs".format(((endedAtMs ?: nowMs) - (startedAtMs ?: endedAtMs!!)) / 1000.0)
+    }
+    Surface(
+        modifier = modifier.fillMaxWidth().padding(vertical = 4.dp),
+        shape = RoundedCornerShape(10.dp),
+        color = MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.5f),
+    ) {
+        Column {
+            Row(Modifier.fillMaxWidth().clickable { expanded = !expanded }
+                .padding(horizontal = 12.dp, vertical = 8.dp),
+                verticalAlignment = Alignment.CenterVertically) {
+                Icon(Icons.AutoMirrored.Filled.KeyboardArrowRight, null,
+                     modifier = Modifier.padding(end = 6.dp),
+                     tint = MaterialTheme.colorScheme.onSurfaceVariant)
+                Text(header, style = MaterialTheme.typography.bodyMedium.copy(
+                    color = MaterialTheme.colorScheme.onSurfaceVariant,
+                    fontSize = 13.sp * LocalFontScale.current,
+                    fontStyle = FontStyle.Italic))
+            }
+            AnimatedVisibility(visible = expanded) {
+                Text(thinking, style = MaterialTheme.typography.bodyMedium.copy(
+                    color = MaterialTheme.colorScheme.onSurfaceVariant,
+                    fontSize = 13.sp * LocalFontScale.current),
+                    modifier = Modifier.padding(start = 30.dp, end = 12.dp, bottom = 10.dp))
+            }
+        }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/Shimmer.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/Shimmer.kt
new file mode 100644
index 0000000..b14e5fd
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/chat/Shimmer.kt
@@ -0,0 +1,36 @@
+// ui/chat/Shimmer.kt - the Working... waiting state: gradient sweep text (opacity pulse variant).
+package dev.edge0.runtime.app.ui.chat
+
+import androidx.compose.animation.core.LinearEasing
+import androidx.compose.animation.core.RepeatMode
+import androidx.compose.animation.core.animateFloat
+import androidx.compose.animation.core.infiniteRepeatable
+import androidx.compose.animation.core.rememberInfiniteTransition
+import androidx.compose.animation.core.tween
+import androidx.compose.foundation.text.BasicText
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.draw.alpha
+import androidx.compose.ui.text.TextStyle
+import androidx.compose.ui.text.font.FontStyle
+import androidx.compose.ui.unit.sp
+import androidx.compose.material3.MaterialTheme
+import dev.edge0.runtime.app.ui.theme.LocalFontScale
+
+/** Waiting indicator before the first token (prefilling=true). */
+@Composable
+fun WorkingShimmer(label: String = "Working...", modifier: Modifier = Modifier) {
+    val t = rememberInfiniteTransition(label = "shimmer")
+    val a by t.animateFloat(0.35f, 0.9f,
+        infiniteRepeatable(tween(900, easing = LinearEasing), RepeatMode.Reverse),
+        label = "a")
+    BasicText(
+        label,
+        modifier = modifier.alpha(a),
+        style = TextStyle(
+            color = MaterialTheme.colorScheme.onSurfaceVariant,
+            fontSize = 13.sp * LocalFontScale.current,
+            fontStyle = FontStyle.Italic),
+    )
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/markdown/MarkdownText.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/markdown/MarkdownText.kt
new file mode 100644
index 0000000..a1d0acc
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/markdown/MarkdownText.kt
@@ -0,0 +1,48 @@
+// ui/markdown/MarkdownText.kt - thin rendering adapter.
+// Markdown is rendered by mikepenz multiplatform-markdown-renderer-m3 (tables,
+// links, strikethrough included). Pinned to 0.27.0: newer releases require
+// compileSdk >= 36 / Kotlin >= 2.1, while this project stays on AGP 8.7.3 (35)
+// + Kotlin 2.0.21, so 0.27.0 is the last compatible line. Colors come from
+// MaterialTheme, keeping the monochrome palette intact.
+//
+// LaTeX fallback: the renderer recognizes $$..$$ / $..$ math nodes but ships no
+// math renderer, so model math output used to render as blank space. mathFallback
+// degrades math to code style before parsing - raw LaTeX stays readable in the
+// monochrome theme without pulling a typesetting engine into the app.
+package dev.edge0.runtime.app.ui.markdown
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.mikepenz.markdown.m3.Markdown
+import com.mikepenz.markdown.m3.markdownColor
+import com.mikepenz.markdown.m3.markdownTypography
+
+@Composable
+fun MarkdownText(src: String, modifier: Modifier = Modifier) {
+    Markdown(mathFallback(src), markdownColor(), markdownTypography(), modifier)
+}
+
+/**
+ * Math degradation: the renderer recognizes math nodes but has no typesetter.
+ * Fenced code blocks break when the math sits inside a list item (the fence
+ * gets swallowed by the list paragraph), so every math form becomes inline
+ * code with its LaTeX backslashes escaped - readable anywhere in the tree.
+ */
+internal fun mathFallback(src: String): String {
+    if ('$' !in src && "\\(" !in src && "\\[" !in src) return src
+    var s = BLOCK_DOLLAR.replace(src) { inlineMath(it.groupValues[1]) }
+    s = BLOCK_BRACKET.replace(s) { inlineMath(it.groupValues[1]) }
+    s = INLINE_DOLLAR.replace(s) { inlineMath(it.groupValues[1]) }
+    s = INLINE_PAREN.replace(s) { inlineMath(it.groupValues[1]) }
+    return s
+}
+
+private fun inlineMath(body: String): String {
+    val oneLine = body.trim().replace(Regex("\\s+"), " ").replace("\\", "\\\\")
+    return if (oneLine.length > 160) oneLine.take(160) else "`$oneLine`"
+}
+
+private val BLOCK_DOLLAR = Regex("""\$\$\s*([\s\S]+?)\s*\$\$""")
+private val BLOCK_BRACKET = Regex("""\\\[\s*([\s\S]+?)\s*\\\]""")
+private val INLINE_DOLLAR = Regex("""\$([^\s$](?:[^$\n]{0,88}[^\s$])?)\$""")
+private val INLINE_PAREN = Regex("""\\\(([^$\n]{1,90}?)\\\)""")
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/models/ModelsScreen.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/models/ModelsScreen.kt
new file mode 100644
index 0000000..e27f658
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/models/ModelsScreen.kt
@@ -0,0 +1,200 @@
+// ui/models/ModelsScreen.kt - model management: scans files/models, shows
+// size/integrity/load state, load|unload buttons (disabled while generating), a disk
+// headroom bar (warn under 1.2x requirement) and an import-instructions card.
+package dev.edge0.runtime.app.ui.models
+
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.rememberScrollState
+import androidx.compose.foundation.shape.CircleShape
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.foundation.verticalScroll
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.automirrored.outlined.ArrowBack
+import androidx.compose.material.icons.outlined.Warning
+import androidx.compose.material3.Button
+import androidx.compose.material3.CircularProgressIndicator
+import androidx.compose.material3.Icon
+import androidx.compose.material3.IconButton
+import androidx.compose.material3.LinearProgressIndicator
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.OutlinedButton
+import androidx.compose.material3.Surface
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.text.font.FontFamily
+import androidx.compose.ui.text.style.TextOverflow
+import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import dev.edge0.runtime.app.Edge0App
+import dev.edge0.runtime.app.runtime.GenState
+import dev.edge0.runtime.app.runtime.ModelEntry
+import dev.edge0.runtime.app.runtime.ModelInventory
+import dev.edge0.runtime.app.ui.theme.LocalFontScale
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+
+private fun fmtGiB(b: Long): String = "%.2f GiB".format(b / 1073741824.0)
+
+@Composable
+fun ModelsScreen(onBack: () -> Unit) {
+    val app = LocalContext.current.applicationContext as Edge0App
+    val container = app.container
+    val scope = rememberCoroutineScope()
+    var entries by remember { mutableStateOf<List<ModelEntry>>(emptyList()) }
+    var avail by remember { mutableStateOf(-1L) }
+    var note by remember { mutableStateOf<String?>(null) }
+    val genState by container.runtime.state.collectAsStateWithLifecycle()
+    var refreshTick by remember { mutableStateOf(0) }
+
+    LaunchedEffect(refreshTick) {
+        withContext(Dispatchers.IO) {
+            entries = ModelInventory.scanRoots(container.modelRoots)
+            avail = ModelInventory.volumeAvailBytes(app.filesDir)
+        }
+    }
+
+    Column(Modifier.fillMaxSize().verticalScroll(rememberScrollState())) {
+        Row(Modifier.fillMaxWidth().padding(horizontal = 4.dp, vertical = 4.dp),
+            verticalAlignment = Alignment.CenterVertically) {
+            IconButton(onClick = onBack) { Icon(Icons.AutoMirrored.Outlined.ArrowBack, "Back") }
+            Text("Models", style = MaterialTheme.typography.titleMedium)
+            Spacer(Modifier.weight(1f))
+            OutlinedButton(onClick = { refreshTick++ }, modifier = Modifier.padding(end = 12.dp)) {
+                Text("Refresh")
+            }
+        }
+        // disk headroom bar (available vs required x 1.2)
+        if (avail >= 0) {
+            val maxReq = entries.filter { it.valid }.maxOfOrNull { it.totalBytes } ?: 0L
+            val warn = maxReq > 0 && avail < maxReq * 1.2
+            Surface(Modifier.fillMaxWidth().padding(horizontal = 12.dp, vertical = 4.dp),
+                    shape = RoundedCornerShape(10.dp),
+                    color = MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.4f)) {
+                Column(Modifier.padding(12.dp)) {
+                    Text("Storage: ${if (avail > 0) fmtGiB(avail) else "N/A"} free" +
+                        if (maxReq > 0) " / largest model ${fmtGiB(maxReq)}, ${fmtGiB((maxReq * 1.2).toLong())} recommended" else "",
+                         style = MaterialTheme.typography.bodySmall)
+                    if (warn) {
+                        Text("Free space below 1.2x requirement: loading may fail midway (engine reports out-of-storage)",
+                             color = MaterialTheme.colorScheme.error,
+                             style = MaterialTheme.typography.bodySmall)
+                    }
+                }
+            }
+        }
+        note?.let {
+            Text(it, color = MaterialTheme.colorScheme.error,
+                 style = MaterialTheme.typography.bodySmall,
+                 modifier = Modifier.padding(horizontal = 16.dp, vertical = 4.dp))
+        }
+        if (entries.isEmpty()) {
+            Text("No models found. Import via adb to:",
+                 modifier = Modifier.padding(16.dp), style = MaterialTheme.typography.bodyMedium)
+            Surface(Modifier.fillMaxWidth().padding(horizontal = 12.dp),
+                    shape = RoundedCornerShape(10.dp),
+                    color = MaterialTheme.colorScheme.surfaceVariant.copy(alpha = 0.3f)) {
+                Text(ModelInventory.importTemplate(app.packageName),
+                     fontFamily = FontFamily.Monospace,
+                     style = MaterialTheme.typography.bodySmall.copy(
+                        fontSize = 11.sp * LocalFontScale.current),
+                     modifier = Modifier.padding(12.dp))
+            }
+            Text("Tap Refresh (top right) after importing. Note: uninstalling the app wipes its data " +
+                 "(Android policy); re-import assets with the command above if they disappear.",
+                 style = MaterialTheme.typography.bodySmall,
+                 color = MaterialTheme.colorScheme.onSurfaceVariant,
+                 modifier = Modifier.padding(16.dp))
+        }
+        for (e in entries) {
+            ModelRow(e,
+                     active = container.runtime.activeModelDir == e.dir.absolutePath,
+                     busy = genState is GenState.Streaming || genState is GenState.Cancelling,
+                     loading = genState is GenState.Loading,
+                     onLoad = {
+                         scope.launch {
+                             note = null
+                             try {
+                                 container.runtime.ensureLoaded(e.dir.absolutePath)
+                                 container.settings.update { s -> s.copy(activeModelDir = e.dir.absolutePath) }
+                             } catch (ex: Exception) {
+                                 android.util.Log.e("Edge0Load", "ensureLoaded failed", ex)
+                                 note = "Load failed: ${ex.message} (if the app was just updated, re-push assets to the import path)"
+                             }
+                         }
+                     },
+                     onUnload = {
+                         scope.launch {
+                             note = null
+                             try { container.runtime.unloadActive() } catch (ex: Exception) {
+                                 note = "Unload failed: ${ex.message}"
+                             }
+                         }
+                     })
+            Spacer(Modifier.height(8.dp))
+        }
+    }
+}
+
+@Composable
+private fun ModelRow(e: ModelEntry, active: Boolean, busy: Boolean, loading: Boolean,
+                     onLoad: () -> Unit, onUnload: () -> Unit) {
+    Surface(Modifier.fillMaxWidth().padding(horizontal = 12.dp),
+            shape = RoundedCornerShape(10.dp),
+            color = MaterialTheme.colorScheme.surface,
+            border = androidx.compose.foundation.BorderStroke(
+                1.dp, if (active) MaterialTheme.colorScheme.primary
+                      else MaterialTheme.colorScheme.outline)) {
+        Row(Modifier.padding(12.dp), verticalAlignment = Alignment.CenterVertically) {
+            Column(Modifier.weight(1f)) {
+                Row(verticalAlignment = Alignment.CenterVertically) {
+                    Text(e.name, maxLines = 1, overflow = TextOverflow.Ellipsis,
+                         style = MaterialTheme.typography.titleSmall)
+                    Spacer(Modifier.size(6.dp))
+                    Text("e0b", style = MaterialTheme.typography.labelSmall,
+                         modifier = Modifier.padding(horizontal = 4.dp))
+                }
+                Text("${fmtGiB(if (e.dirBytes > 0) e.dirBytes else e.totalBytes)}" +
+                     if (!e.valid) " · incomplete (${e.hint})" else "",
+                     style = MaterialTheme.typography.bodySmall,
+                     color = if (e.valid) MaterialTheme.colorScheme.onSurfaceVariant
+                             else MaterialTheme.colorScheme.error)
+                Text(if (active) "● Loaded" else if (loading) "◌ Loading..." else "○ Not loaded",
+                     style = MaterialTheme.typography.bodySmall.copy(
+                        fontSize = 11.sp * LocalFontScale.current),
+                     color = if (active) MaterialTheme.colorScheme.primary
+                             else MaterialTheme.colorScheme.onSurfaceVariant)
+            }
+            if (!e.valid) {
+                Icon(Icons.Outlined.Warning, "Broken assets", tint = MaterialTheme.colorScheme.error)
+                return@Row
+            }
+            if (active) {
+                OutlinedButton(onClick = onUnload, enabled = !busy && !loading) { Text("Unload") }
+            } else {
+                Button(onClick = onLoad, enabled = !busy && !loading) {
+                    if (loading) CircularProgressIndicator(Modifier.size(14.dp))
+                    else Text("Load")
+                }
+            }
+        }
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/settings/SettingsScreen.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/settings/SettingsScreen.kt
new file mode 100644
index 0000000..490b513
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/settings/SettingsScreen.kt
@@ -0,0 +1,167 @@
+// ui/settings/SettingsScreen.kt - settings in a single-column card layout:
+// appearance (theme / font scale), assistant (system prompt + thinking), sampling
+// parameters, about. Writes go through SettingsStore.update.
+package dev.edge0.runtime.app.ui.settings
+
+import androidx.compose.foundation.BorderStroke
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.ColumnScope
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.rememberScrollState
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.foundation.verticalScroll
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.automirrored.outlined.ArrowBack
+import androidx.compose.material3.Icon
+import androidx.compose.material3.IconButton
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.OutlinedTextField
+import androidx.compose.material3.SegmentedButton
+import androidx.compose.material3.SegmentedButtonDefaults
+import androidx.compose.material3.SingleChoiceSegmentedButtonRow
+import androidx.compose.material3.Slider
+import androidx.compose.material3.Surface
+import androidx.compose.material3.Switch
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.text.font.FontWeight
+import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import dev.edge0.runtime.app.Edge0App
+import dev.edge0.runtime.app.data.AppSettings
+import dev.edge0.runtime.app.ui.theme.LocalFontScale
+import dev.edge0.runtime.app.ui.theme.ThemeMode
+import kotlinx.coroutines.launch
+
+@Composable
+fun SettingsScreen(onBack: () -> Unit) {
+    val app = LocalContext.current.applicationContext as Edge0App
+    val store = app.container.settings
+    val s by store.flow.collectAsStateWithLifecycle(initialValue = AppSettings())
+    val scope = rememberCoroutineScope()
+    fun patch(f: (AppSettings) -> AppSettings) = scope.launch { store.update(f) }
+
+    Column(Modifier.fillMaxSize().verticalScroll(rememberScrollState())) {
+        Row(Modifier.fillMaxWidth().padding(horizontal = 4.dp, vertical = 4.dp),
+            verticalAlignment = Alignment.CenterVertically) {
+            IconButton(onClick = onBack) { Icon(Icons.AutoMirrored.Outlined.ArrowBack, "Back") }
+            Text("Settings", style = MaterialTheme.typography.titleMedium)
+        }
+
+        Group("Appearance") {
+            Row(Modifier.fillMaxWidth().padding(vertical = 6.dp),
+                verticalAlignment = Alignment.CenterVertically) {
+                Text("Theme", modifier = Modifier.weight(1f))
+                SingleChoiceSegmentedButtonRow {
+                    val opts = listOf(ThemeMode.SYSTEM to "System", ThemeMode.LIGHT to "Light",
+                                      ThemeMode.DARK to "Dark")
+                    opts.forEachIndexed { i, (mode, label) ->
+                        SegmentedButton(
+                            selected = s.themeMode == mode,
+                            shape = SegmentedButtonDefaults.itemShape(i, opts.size),
+                            onClick = { patch { cur -> cur.copy(themeMode = mode) } },
+                        ) { Text(label, fontSize = 12.sp) }
+                    }
+                }
+            }
+            Text("Font size ${"%.2f".format(s.fontScale)}", style = MaterialTheme.typography.bodySmall)
+            Slider(value = s.fontScale, onValueChange = { v ->
+                scope.launch { store.update { cur -> cur.copy(fontScale = v) } } },
+                valueRange = 0.85f..1.3f)
+        }
+
+        Group("Assistant") {
+            Row(Modifier.fillMaxWidth().padding(vertical = 4.dp),
+                verticalAlignment = Alignment.CenterVertically) {
+                Text("Use system prompt", modifier = Modifier.weight(1f))
+                Switch(checked = s.useSystemPrompt,
+                       onCheckedChange = { c -> patch { cur -> cur.copy(useSystemPrompt = c) } })
+            }
+            if (s.useSystemPrompt) {
+                OutlinedTextField(value = s.systemPrompt,
+                    onValueChange = { v -> patch { cur -> cur.copy(systemPrompt = v) } },
+                    label = { Text("System prompt") },
+                    minLines = 3, maxLines = 6, modifier = Modifier.fillMaxWidth())
+            }
+            Row(Modifier.fillMaxWidth().padding(vertical = 4.dp),
+                verticalAlignment = Alignment.CenterVertically) {
+                Text("Enable thinking", modifier = Modifier.weight(1f))
+                Switch(checked = s.enableThinking,
+                       onCheckedChange = { c -> patch { cur -> cur.copy(enableThinking = c) } })
+            }
+        }
+
+        Group("Sampling") {
+            // defaults mirror the upstream sampling quartet
+            ParamSlider("repetition penalty (1.0=off)", s.repetitionPenalty, 1f, 1.5f) { v ->
+                patch { cur -> cur.copy(repetitionPenalty = v) } }
+            ParamSlider("temperature (0=greedy)", s.temperature, 0f, 2f) { v ->
+                patch { cur -> cur.copy(temperature = v) } }
+            ParamSliderInt("top-k (0=off)", s.topK, 0, 100) { v ->
+                patch { cur -> cur.copy(topK = v) } }
+            ParamSlider("top-p (>=1=off)", s.topP, 0.1f, 1f) { v ->
+                patch { cur -> cur.copy(topP = v) } }
+            ParamSliderInt("max new tokens (0=default 1024)", s.maxNewTokens, 0, 2048) { v ->
+                patch { cur -> cur.copy(maxNewTokens = v) } }
+            OutlinedTextField(value = s.seed.toString(),
+                onValueChange = { v -> v.toLongOrNull()?.let { seed ->
+                    patch { cur -> cur.copy(seed = seed) } } },
+                label = { Text("seed") }, singleLine = true,
+                modifier = Modifier.fillMaxWidth().padding(top = 6.dp))
+        }
+
+        Group("About") {
+            Text("llama.cpp (pinned build) · " +
+                 "app ${dev.edge0.runtime.app.BuildConfig.VERSION_NAME}" +
+                 "+${dev.edge0.runtime.app.BuildConfig.GIT_SHA}",
+                 style = MaterialTheme.typography.bodySmall)
+            Text("Runtime and UI ship in a single APK",
+                 style = MaterialTheme.typography.bodySmall,
+                 color = MaterialTheme.colorScheme.onSurfaceVariant)
+        }
+        Spacer(Modifier.height(24.dp))
+    }
+}
+
+@Composable
+private fun Group(title: String, content: @Composable ColumnScope.() -> Unit) {
+    Column(Modifier.fillMaxWidth().padding(horizontal = 12.dp, vertical = 6.dp)) {
+        Text(title.uppercase(), style = MaterialTheme.typography.labelMedium.copy(
+            fontWeight = FontWeight.SemiBold, fontSize = 11.sp * LocalFontScale.current,
+            color = MaterialTheme.colorScheme.onSurfaceVariant),
+            modifier = Modifier.padding(start = 4.dp, bottom = 4.dp))
+        Surface(shape = RoundedCornerShape(10.dp),
+                color = MaterialTheme.colorScheme.surface,
+                border = BorderStroke(1.dp, MaterialTheme.colorScheme.outline)) {
+            Column(Modifier.padding(12.dp), content = content)
+        }
+    }
+}
+
+@Composable
+private fun ParamSlider(label: String, value: Float, min: Float, max: Float,
+                        onValue: (Float) -> Unit) {
+    Text("$label = ${"%.2f".format(value)}", style = MaterialTheme.typography.bodySmall)
+    Slider(value = value.coerceIn(min, max), onValueChange = onValue, valueRange = min..max,
+           modifier = Modifier.fillMaxWidth())
+}
+
+@Composable
+private fun ParamSliderInt(label: String, value: Int, min: Int, max: Int,
+                           onValue: (Int) -> Unit) {
+    Text("$label = $value", style = MaterialTheme.typography.bodySmall)
+    Slider(value = value.toFloat().coerceIn(min.toFloat(), max.toFloat()),
+           onValueChange = { onValue(it.toInt()) }, valueRange = min.toFloat()..max.toFloat(),
+           modifier = Modifier.fillMaxWidth())
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/theme/Color.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/theme/Color.kt
new file mode 100644
index 0000000..db7e6d3
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/theme/Color.kt
@@ -0,0 +1,41 @@
+// ui/theme/Color.kt - palette converted from the OKLCH design values to sRGB
+// (OKLab -> linear sRGB -> gamma, per Björn Ottosson's transform). The original
+// oklch() values are kept in paired comments for traceability. The palette is
+// grayscale; blue appears only in active states such as download progress.
+package dev.edge0.runtime.app.ui.theme
+
+import androidx.compose.ui.graphics.Color
+
+// ---- Light (index.css :root)----
+val LightBackground = Color(0xFFFFFFFF)        // oklch(1 0 0)
+val LightForeground = Color(0xFF0A0A0A)        // oklch(0.145 0 0)
+val LightPrimary = Color(0xFF171717)           // oklch(0.205 0 0) send-button background
+val LightOnPrimary = Color(0xFFFAFAFA)         // primary foreground, shadcn default white
+val LightSecondary = Color(0xFFF5F5F5)         // oklch(0.97 0 0) user bubble background
+val LightOnSecondary = Color(0xFF0A0A0A)       // secondary-foreground
+val LightMutedForeground = Color(0xFF737373)   // oklch(0.556 0 0) secondary text / action row
+val LightBorder = Color(0xFFE5E5E5)            // oklch(0.922 0 0)
+val LightSidebar = Color(0xFFF2F2F2)           // oklch(0.96 0 0)
+val LightDestructive = Color(0xFFE7000B)       // oklch(0.577 0.245 27.325)
+
+// ---- Dark (index.css .dark)----
+val DarkBackground = Color(0xFF121212)         // oklch(0.18 0 0)
+val DarkForeground = Color(0xFFFAFAFA)         // oklch(0.985 0 0)
+val DarkPrimary = Color(0xFFE5E5E5)            // oklch(0.922 0 0) dark primary = light gray, dark glyph on it
+val DarkOnPrimary = Color(0xFF171717)
+val DarkSecondary = Color(0xFF383838)          // oklch(0.34 0 0)
+val DarkOnSecondary = Color(0xFFFAFAFA)
+val DarkMutedForeground = Color(0xFFB7B7B7)    // oklch(0.78 0 0)
+val DarkBorder = Color(0x24FFFFFF)             // oklch(1 0 0 / 14%)
+val DarkSidebar = Color(0xFF1B1B1B)            // oklch(0.22 0 0)
+val DarkDestructive = Color(0xFFFF6467)        // oklch(0.704 0.191 22.216)
+val DarkRing = Color(0xFFA4A4A4)               // oklch(0.72 0 0)
+
+// layout constants (design spec to dp)
+object Dimens {
+    const val RadiusCard = 10f    // --radius: 0.625rem ≈ 10dp
+    const val RadiusComposer = 24f // rounded-3xl (composer card)
+    const val RadiusBubble = 6f    // rounded-md (bubble)
+    const val DrawerWidth = 288f   // Sheet 18rem (sidebar.tsx SIDEBAR_WIDTH_MOBILE)
+    const val ContentMaxWidth = 768f // max-w-3xl message column (auto-centers on tablets)
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/app/ui/theme/Theme.kt b/android/app/src/main/java/dev/edge0/runtime/app/ui/theme/Theme.kt
new file mode 100644
index 0000000..d8b2b6b
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/app/ui/theme/Theme.kt
@@ -0,0 +1,91 @@
+// ui/theme/Theme.kt - three theme modes (light/dark/system) plus a global font scale.
+// LocalFontScale is the Compose equivalent of the design system's font-size-base
+// multiplier; screens apply it explicitly (Material3 typography baseline is 14sp).
+package dev.edge0.runtime.app.ui.theme
+
+import androidx.compose.foundation.isSystemInDarkTheme
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Shapes
+import androidx.compose.material3.Typography
+import androidx.compose.material3.darkColorScheme
+import androidx.compose.material3.lightColorScheme
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.compositionLocalOf
+import androidx.compose.ui.text.TextStyle
+import androidx.compose.ui.text.font.FontFamily
+import androidx.compose.ui.text.font.FontWeight
+import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
+
+// whole-app text-sm baseline (system sans fallback; custom font files out of scope for v1)
+private val Edge0Typography = Typography(
+    bodyLarge = TextStyle(fontFamily = FontFamily.Default, fontSize = 14.sp, lineHeight = 20.sp),
+    bodyMedium = TextStyle(fontFamily = FontFamily.Default, fontSize = 14.sp, lineHeight = 20.sp),
+    titleLarge = TextStyle(fontFamily = FontFamily.Default, fontSize = 18.sp,
+                           fontWeight = FontWeight.Medium),
+)
+
+private val LightScheme = lightColorScheme(
+    primary = LightPrimary,
+    onPrimary = LightOnPrimary,
+    secondary = LightSecondary,
+    onSecondary = LightOnSecondary,
+    background = LightBackground,
+    onBackground = LightForeground,
+    surface = LightBackground,
+    onSurface = LightForeground,
+    surfaceVariant = LightSecondary,
+    outline = LightBorder,
+    error = LightDestructive,
+)
+
+private val DarkScheme = darkColorScheme(
+    primary = DarkPrimary,
+    onPrimary = DarkOnPrimary,
+    secondary = DarkSecondary,
+    onSecondary = DarkOnSecondary,
+    background = DarkBackground,
+    onBackground = DarkForeground,
+    surface = DarkBackground,
+    onSurface = DarkForeground,
+    surfaceVariant = DarkSecondary,
+    outline = DarkBorder,
+    error = DarkDestructive,
+)
+
+/** Radius ladder: md = bubble / lg = card (10dp) / 2xl = composer (24dp). */
+private val Edge0Shapes = Shapes(
+    extraSmall = RoundedCornerShape(4.dp),
+    small = RoundedCornerShape(6.dp),   // user bubble
+    medium = RoundedCornerShape(10.dp), // cards
+    large = RoundedCornerShape(16.dp),
+    extraLarge = RoundedCornerShape(24.dp), // composer
+)
+
+/** Font scale (0.85..1.3, driven by the settings slider, persisted in DataStore). */
+val LocalFontScale = compositionLocalOf { 1.0f }
+
+/** Theme mode: SYSTEM follows the platform setting. */
+enum class ThemeMode { SYSTEM, LIGHT, DARK }
+
+@Composable
+fun Edge0Theme(
+    mode: ThemeMode = ThemeMode.SYSTEM,
+    fontScale: Float = 1.0f,
+    content: @Composable () -> Unit,
+) {
+    val dark = when (mode) {
+        ThemeMode.SYSTEM -> isSystemInDarkTheme()
+        ThemeMode.LIGHT -> false
+        ThemeMode.DARK -> true
+    }
+    androidx.compose.runtime.CompositionLocalProvider(LocalFontScale provides fontScale) {
+        MaterialTheme(
+            colorScheme = if (dark) DarkScheme else LightScheme,
+            typography = Edge0Typography,
+            shapes = Edge0Shapes,
+            content = content,
+        )
+    }
+}
diff --git a/android/app/src/main/java/dev/edge0/runtime/engine/Types.kt b/android/app/src/main/java/dev/edge0/runtime/engine/Types.kt
new file mode 100644
index 0000000..9c55074
--- /dev/null
+++ b/android/app/src/main/java/dev/edge0/runtime/engine/Types.kt
@@ -0,0 +1,36 @@
+// Shared data types of the runtime layer.
+package dev.edge0.runtime.engine
+
+enum class Status(val code: Int) {
+    OK(0),
+    INSUFFICIENT_STORAGE(1),
+    FORMAT_VERSION(2),
+    ASSET_CORRUPT(3),
+    MEM_BUDGET(4),
+    IO(5),
+    CANCELLED(6),
+    SESSION_EXISTS(7),
+    INVALID_ARG(8),
+    NOT_IMPLEMENTED(9),
+    INTERNAL(100),
+    ;
+    companion object { fun of(v: Int): Status = entries.firstOrNull { it.code == v } ?: INTERNAL }
+}
+
+class EngineException(val status: Status, message: String) : Exception("[$status] $message")
+
+data class GenParams(
+    val temperature: Float = 0f,      // <= 0 selects greedy decoding (deterministic)
+    val topK: Int = 0,
+    val topP: Float = 1f,
+    val seed: Long = 0x20260915L,
+    val maxNewTokens: Int = 0,
+    val repetitionPenalty: Float = 0f,    // 0 or 1.0 = off; > 0 and != 1 applies penalty
+    val firstTokenGreedy: Boolean = false, // true = argmax for the first token
+)
+
+/** A chat turn. role is one of system/user/assistant; at most one system turn, first.
+ *  genIds = assistant token-id sidecar: the exact token ids this assistant turn
+ *  stepped through. When present the engine replays the turn from ids (the text
+ *  is for the UI only), keeping the prompt byte-identical across turns. */
+data class ChatMessage(val role: String, val content: String, val genIds: IntArray? = null)
diff --git a/android/app/src/main/res/values/strings.xml b/android/app/src/main/res/values/strings.xml
new file mode 100644
index 0000000..56cb79d
--- /dev/null
+++ b/android/app/src/main/res/values/strings.xml
@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="utf-8"?>
+<resources>
+    <string name="app_name">Edge0 Chat</string>
+</resources>
diff --git a/android/app/src/main/res/values/themes.xml b/android/app/src/main/res/values/themes.xml
new file mode 100644
index 0000000..461b543
--- /dev/null
+++ b/android/app/src/main/res/values/themes.xml
@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="utf-8"?>
+<resources>
+    <!-- Compose draws the full UI; this theme only sets the window background
+      (edge-to-edge is enabled in the Activity). -->
+    <style name="Theme.Edge0" parent="android:Theme.DeviceDefault.NoActionBar">
+        <item name="android:windowBackground">@android:color/transparent</item>
+    </style>
+</resources>
diff --git a/android/app/src/test/java/dev/edge0/runtime/app/runtime/HistoryWindowTest.kt b/android/app/src/test/java/dev/edge0/runtime/app/runtime/HistoryWindowTest.kt
new file mode 100644
index 0000000..d739be8
--- /dev/null
+++ b/android/app/src/test/java/dev/edge0/runtime/app/runtime/HistoryWindowTest.kt
@@ -0,0 +1,74 @@
+// HistoryWindowTest - budget trimming, ordering, assistant replay shape.
+package dev.edge0.runtime.app.runtime
+
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+class HistoryWindowTest {
+    private val O = ThinkSplitter.OPEN
+    private val C = ThinkSplitter.CLOSE
+
+    @Test fun tailWithinBudgetKeptInOrder() {
+        val turns = listOf(
+            StoredMessage("user", "q1"), StoredMessage("assistant", "a1"),
+            StoredMessage("user", "q2"), StoredMessage("assistant", "a2"),
+            StoredMessage("user", "q3"),
+        )
+        val msgs = HistoryWindow.build("sys", turns, budgetChars = 10000)
+        assertEquals(listOf("system", "user", "assistant", "user", "assistant", "user"),
+            msgs.map { it.role })
+        assertEquals("sys", msgs[0].content)
+        assertEquals("q3", msgs.last().content)
+    }
+
+    @Test fun budgetTrimsOldestNotLast() {
+        val big = "W".repeat(100)
+        val turns = listOf(
+            StoredMessage("user", big), StoredMessage("assistant", big),
+            StoredMessage("user", "lastQ"),
+        )
+        val msgs = HistoryWindow.build(null, turns, budgetChars = 60)
+        // the last turn always survives; older turns beyond budget are trimmed
+        assertEquals(1, msgs.size)
+        assertEquals("lastQ", msgs[0].content)
+    }
+
+    @Test fun assistantThinkingReplayedIntoEngineTranscript() {
+        // Revisited under the sidecar regime: thinking MUST enter the context -
+        // KV prefix reuse requires bit-identical replay of the live tokens.
+        val turns = listOf(
+            StoredMessage("user", "q1"),
+            StoredMessage("assistant", "body", thinking = "reason"),
+            StoredMessage("user", "q2"),
+        )
+        val msgs = HistoryWindow.build(null, turns)
+        assertEquals(O + "reason" + C + "body", msgs[1].content)
+    }
+
+    @Test fun sidecarIdsRideAlongReplay() {
+        val ids = intArrayOf(5, 42, 17)
+        val turns = listOf(
+            StoredMessage("user", "q"),
+            StoredMessage("assistant", "body", genIds = ids),
+            StoredMessage("user", "q2"),
+        )
+        val msgs = HistoryWindow.build(null, turns)
+        assertTrue(msgs[1].genIds!!.contentEquals(ids))
+        assertEquals("body", msgs[1].content) // no thinking column -> no tags (ids are authoritative)
+    }
+
+    @Test fun assistantWithoutThinkingPlainBody() {
+        val turns = listOf(
+            StoredMessage("user", "q1"),
+            StoredMessage("assistant", "plainbody"),
+        )
+        val msgs = HistoryWindow.build(null, turns)
+        assertEquals("plainbody", msgs[1].content)
+    }
+
+    @Test fun tighterBudgetFloor() {
+        assertEquals(400, HistoryWindow.tighterBudget(700))
+        assertTrue(HistoryWindow.tighterBudget(300) >= 400)
+    }
+}
diff --git a/android/app/src/test/java/dev/edge0/runtime/app/runtime/LlamaRuntimeLogicTest.kt b/android/app/src/test/java/dev/edge0/runtime/app/runtime/LlamaRuntimeLogicTest.kt
new file mode 100644
index 0000000..2233415
--- /dev/null
+++ b/android/app/src/test/java/dev/edge0/runtime/app/runtime/LlamaRuntimeLogicTest.kt
@@ -0,0 +1,58 @@
+// LlamaRuntimeLogicTest - pure-logic unit tests (host JVM, no device).
+// Device-facing behavior is covered by the instrumented LlamaRuntimeTest suite.
+package dev.edge0.runtime.app.runtime
+
+import dev.edge0.runtime.engine.ChatMessage
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+class LlamaRuntimeLogicTest {
+
+    @Test fun metricsParseRoundTrip() {
+        val rt = LlamaRuntime()
+        val m = rt.parseMetrics(
+            "OK gen=64 tps=18.19 prefill_ms=3087 first_ms=182 prompt=17 reused=12")
+        checkNotNull(m)
+        assertEquals(64, m.newTokens)
+        assertEquals(18.19, m.decodeTokS, 1e-6)
+        assertEquals(3087, m.prefillMs)
+        assertEquals(17, m.promptTokens)
+        assertEquals(12, m.prefix_reused_tokens)
+        assertTrue(m.prefix_reused)
+        // effective prefill rate = new input tokens (17-12=5) over prefill time
+        assertEquals(5 * 1000.0 / 3087, m.prefillTokS, 1e-6)
+    }
+
+    @Test fun metricsParseCancelledAndGarbage() {
+        val rt = LlamaRuntime()
+        checkNotNull(rt.parseMetrics("CANCELLED gen=8 tps=2.00 prefill_ms=100 first_ms=10 prompt=5 reused=0"))
+        assertNull(rt.parseMetrics("ERR template need=-1"))
+    }
+
+    @Test fun historyWindowSystemFirstAndLastTurnKept() {
+        val long = "x".repeat(400)
+        val turns = (1..30).map { i ->
+            if (i % 2 == 1) StoredMessage("user", "q$i $long")
+            else StoredMessage("assistant", "a$i $long")
+        } + StoredMessage("user", "final question " + long)
+        val out = HistoryWindow.build("You are Edge0.", turns, budgetChars = 1000)
+        assertEquals("system", out.first().role)                      // system always prepended
+        assertEquals("user", out.last().role)                         // current question survives
+        assertTrue(out.size < turns.size + 1)                         // budget clipped history
+    }
+
+    @Test fun historyWindowDropsBlankAssistantPoisonTurns() {
+        // regression sentinel (empty-answer incident): blank assistant turns without
+        // genIds must never reach the model - the model would "see itself" answering empty
+        val turns = listOf(
+            StoredMessage("user", "hi"),
+            StoredMessage("assistant", ""),
+            StoredMessage("user", "who are you"),
+        )
+        val out = HistoryWindow.build(null, turns)
+        assertTrue(out.none { it.role == "assistant" && it.content.isBlank() })
+        assertEquals(2, out.size)
+    }
+}
diff --git a/android/app/src/test/java/dev/edge0/runtime/app/runtime/MetricsMapperTest.kt b/android/app/src/test/java/dev/edge0/runtime/app/runtime/MetricsMapperTest.kt
new file mode 100644
index 0000000..1af086d
--- /dev/null
+++ b/android/app/src/test/java/dev/edge0/runtime/app/runtime/MetricsMapperTest.kt
@@ -0,0 +1,63 @@
+// MetricsMapperTest - real-device schema1 sample capture plus tolerance checks.
+package dev.edge0.runtime.app.runtime
+
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNotNull
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+class MetricsMapperTest {
+
+    private fun resource(): String =
+        javaClass.classLoader!!.getResourceAsStream("metrics_real.json")!!
+            .readBytes().toString(Charsets.UTF_8)
+
+    @Test fun realDeviceSampleMaps() {
+        val m = MetricsMapper.parse(resource())
+        assertNotNull(m)
+        m!!
+        assertEquals(1, m.schema)
+        assertEquals(23, m.promptTokens)
+        assertEquals(false, m.prefix_reused)
+        assertEquals(0, m.prefix_reused_tokens)
+        assertEquals(8, m.newTokens)
+        assertEquals(1, m.chatTurns)
+        assertTrue(m.thinkingOn)
+        assertEquals(1_366_167_552L, m.peakRssBytes)
+        assertTrue(m.decodeTokS > 4.0 && m.decodeTokS < 4.4)
+        assertEquals("predicted", m.route?.mode)
+        assertEquals("executed", m.route?.feature)
+        assertNull(m.route?.agree) // JSON null → Kotlin null
+        assertEquals(635L, m.pool?.hits)
+        assertEquals(2095L, m.pool?.misses)
+        assertTrue(m.gpu?.active == true)
+        assertTrue(m.gpu?.device?.contains("Adreno") == true)
+        assertEquals("ok", m.gpu?.why)        // backend-audit row data source
+        assertTrue(m.gpu?.cacheHit == true)   // kernel cache hit (CLI-domain precedent)
+        val hr = m.hitRate!!
+        assertTrue(hr > 0.22 && hr < 0.24) // 635/2730
+    }
+
+    @Test fun missingFieldsDefaultNotThrow() {
+        val m = MetricsMapper.parse("""{"schema": 1}""")
+        assertNotNull(m)
+        assertEquals(0, m!!.newTokens)
+        assertEquals(false, m.prefix_reused)
+        assertEquals(0, m.prefix_reused_tokens)
+        assertNull(m.route)
+        assertNull(m.hitRate)
+    }
+
+    @Test fun unknownFieldsIgnored() {
+        val m = MetricsMapper.parse("""{"schema":1,"brand_new_field":{"a":[1,2]},"decode_tok_s":3.0}""")
+        assertNotNull(m)
+        assertEquals(3.0, m!!.decodeTokS, 1e-9)
+    }
+
+    @Test fun corruptInputReturnsNull() {
+        assertNull(MetricsMapper.parse("not json at all"))
+        assertNull(MetricsMapper.parse(""))
+        assertNull(MetricsMapper.parse(null))
+    }
+}
diff --git a/android/app/src/test/java/dev/edge0/runtime/app/runtime/ThinkSplitterTest.kt b/android/app/src/test/java/dev/edge0/runtime/app/runtime/ThinkSplitterTest.kt
new file mode 100644
index 0000000..e62b77e
--- /dev/null
+++ b/android/app/src/test/java/dev/edge0/runtime/app/runtime/ThinkSplitterTest.kt
@@ -0,0 +1,110 @@
+// ThinkSplitterTest - incremental tag-pair splitting table (cross-chunk tag cases).
+// Tag literals are always assembled from constants (tooling strips special-looking
+// literals; the rule applies to tests too).
+// Delivery semantics: same-kind segments merge within one feed; across feeds the
+// UI aggregates.
+package dev.edge0.runtime.app.runtime
+
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+class ThinkSplitterTest {
+    private val O = ThinkSplitter.OPEN
+    private val C = ThinkSplitter.CLOSE
+
+    private fun run(pieces: List<String>): List<Seg> {
+        val sp = ThinkSplitter()
+        val out = ArrayList<Seg>()
+        for (p in pieces) out.addAll(sp.feed(p))
+        out.addAll(sp.flush())
+        return foldAdjacent(out)
+    }
+
+    /** Fold adjacent same-kind segments (normalize the cross-feed freedom). */
+    private fun foldAdjacent(segs: List<Seg>): List<Seg> {
+        val out = ArrayList<Seg>()
+        for (s in segs) {
+            val last = out.lastOrNull()
+            if (last != null && last::class == s::class) {
+                out[out.size - 1] = when (s) {
+                    is Seg.Text -> Seg.Text((last as Seg.Text).s + s.s)
+                    is Seg.Think -> Seg.Think((last as Seg.Think).s + s.s)
+                }
+            } else out.add(s)
+        }
+        return out
+    }
+
+    @Test fun promptInjectedOpenStartsInThink() {
+        // real-device scenario: the OPEN tag lives in the prompt, the stream starts inside thinking
+        val sp = ThinkSplitter(initialThink = true)
+        val out = ArrayList<Seg>()
+        out.addAll(sp.feed("plan text"))
+        out.addAll(sp.feed("$C" + "answer"))
+        out.addAll(sp.flush())
+        assertEquals(listOf<Seg>(Seg.Think("plan text"), Seg.Text("answer")), foldAdjacent(out))
+    }
+
+    @Test fun wholePairInOneFeed() {
+        assertEquals(listOf<Seg>(Seg.Think("abc"), Seg.Text("def")),
+            run(listOf("$O" + "abc$C" + "def")))
+    }
+
+    @Test fun openTagSplitAcrossPieces() {
+        // open tag split across three chunks + close tag across two (folded result must be whole)
+        assertEquals(listOf<Seg>(Seg.Text("hi "), Seg.Think("thinkingbody"), Seg.Text("tail")),
+            run(listOf("hi <", "th", "ink>thinking", "body</t", "h", "ink>tail")))
+    }
+
+    @Test fun plainTextNoTags() {
+        assertEquals(listOf<Seg>(Seg.Text("purebody.")), run(listOf("pure", "body.")))
+    }
+
+    @Test fun fakePrefixNotTag() {
+        // body containing a fake "<th" prefix not followed by the rest -> stays body
+        assertEquals(listOf<Seg>(Seg.Text("a <th b")), run(listOf("a <th", " b")))
+    }
+
+    @Test fun nestedOpenInsideThinkIsPlainThink() {
+        assertEquals(listOf<Seg>(Seg.Think("  again $O same again"), Seg.Text("post")),
+            run(listOf("$O  again $O same again$C" + "post")))
+    }
+
+    @Test fun unclosedToEosIsAllThink() {
+        assertEquals(listOf<Seg>(Seg.Think("never closed")), run(listOf("$O" + "never closed")))
+    }
+
+    @Test fun emptyPairEmitsNoThinkSegment() {
+        assertEquals(listOf<Seg>(Seg.Text("x")), run(listOf("$O$C" + "x")))
+    }
+
+    @Test fun trailingFakeOpenFlushesAsText() {
+        // EOS with a half-tag carried: flush delivers under the current mode (TEXT)
+        assertEquals(listOf<Seg>(Seg.Text("before<th")), run(listOf("before<th")))
+    }
+
+    @Test fun thinkThenMoreThinkBlocks() {
+        assertEquals(
+            listOf<Seg>(Seg.Think("a"), Seg.Text("b"), Seg.Think("c"), Seg.Text("d")),
+            run(listOf("$O" + "a$C" + "b$O" + "c$C" + "d")))
+    }
+
+    @Test fun charByCharFeedMatchesOneShot() {
+        val sp1 = ThinkSplitter()
+        val src = "$O" + "multi step" + "thinking$C" + "reply $O" + "again$C" + "end"
+        val oneShot = foldAdjacent(sp1.feed(src) + sp1.flush())
+        val sp2 = ThinkSplitter()
+        val per = ArrayList<Seg>()
+        for (ch in src) { per.addAll(sp2.feed(ch.toString())) }
+        per.addAll(sp2.flush())
+        assertEquals(oneShot, foldAdjacent(per))
+    }
+
+    @Test fun thinkingNowTracksMode() {
+        val sp = ThinkSplitter()
+        sp.feed("$O" + "x")
+        assertEquals(true, sp.thinkingNow)
+        sp.feed("y$C" + "z")
+        assertEquals(false, sp.thinkingNow)
+    }
+}
diff --git a/android/app/src/test/java/dev/edge0/runtime/app/ui/markdown/MathFallbackTest.kt b/android/app/src/test/java/dev/edge0/runtime/app/ui/markdown/MathFallbackTest.kt
new file mode 100644
index 0000000..735a61a
--- /dev/null
+++ b/android/app/src/test/java/dev/edge0/runtime/app/ui/markdown/MathFallbackTest.kt
@@ -0,0 +1,49 @@
+// MathFallbackTest - LaTeX degradation preprocessing for the markdown adapter.
+// Guards the blank-math-render incident: $$..$$ math nodes were accepted by the
+// renderer but drawn as empty space; mathFallback makes them visible as inline
+// code (fenced blocks break inside list items, which is where models put math).
+package dev.edge0.runtime.app.ui.markdown
+
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+class MathFallbackTest {
+    private val D = "\$" // a dollar sign without triggering string interpolation
+
+    @Test fun blockDollarBecomesInlineCode() {
+        val src = "1. Divide 8 by 3:\n\n" + D + D + "\\frac{8}{3}" + D + D + "\n\ndone"
+        val out = mathFallback(src)
+        assertTrue("math content visible", "frac" in out)
+        assertTrue("math delimiters gone", D !in out)
+        assertTrue("no fence noise", "```" !in out)
+    }
+
+    @Test fun inlineDollarBecomesCode() {
+        val out = mathFallback("where " + D + "x^2 + y" + D + " holds")
+        assertTrue("backtick-wrapped", out.contains("`x^2 + y`"))
+    }
+
+    @Test fun moneyTextUntouched() {
+        val src = "costs " + D + "5 and " + D + "6 total"
+        assertEquals(src, mathFallback(src))
+    }
+
+    @Test fun plainTextPassesThrough() {
+        assertEquals("no math here at all", mathFallback("no math here at all"))
+    }
+
+    @Test fun parenAndBracketForms() {
+        assertTrue("`a+b`" in mathFallback("do \\(a+b\\) now"))
+        val br = mathFallback("\\[ \\frac{1}{2} \\]")
+        assertTrue("bracket body visible", "frac" in br)
+        assertTrue("delimiters gone", "\\[" !in br)
+    }
+
+    @Test fun multilineBlockCollapsesToOneLine() {
+        val src = D + D + "\n\\sum_{i=1}^{n} i\n" + D + D
+        val out = mathFallback(src)
+        assertTrue("body present", "sum_{i=1}^{n} i" in out)
+        assertTrue("single line", !out.trim().contains("\n"))
+    }
+}
diff --git a/android/app/src/test/resources/metrics_real.json b/android/app/src/test/resources/metrics_real.json
new file mode 100644
index 0000000..0e7ffdb
--- /dev/null
+++ b/android/app/src/test/resources/metrics_real.json
@@ -0,0 +1 @@
+{"cancelled": false, "chat_turns": 1, "config": {"compute_threads": 1, "io_threads": 4, "max_context_tokens": 4096, "pool_budget_bytes": 100663296, "scrub_percent": 0}, "decode_ms": 1908, "decode_tok_s": 4.1928721174004195, "error": "", "first_token_ms": 2725, "generate_seq": 1, "gpu": {"active": true, "cache_hit": true, "cus": 12, "device": "QUALCOMM Adreno(TM) 830", "driver": "OpenCL 3.0 QUALCOMM build: 0800.72 Compiler E031.47.18.49", "mem_mib": 7576, "ms": 1415.627034187317, "n_submit": 7944, "n_sync": 207, "ok": true, "policy": 0, "probed": true, "why": "ok"}, "int8_effective": false, "layer_track": 0, "new_tokens": 8, "params": {"max_new_tokens": 8, "seed": 539363605, "temperature": 0.0, "top_k": 0, "top_p": 1.0}, "peak_rss_bytes": 1366167552, "pool": {"claim_ns": 0, "claim_ns_per_tok": 0.0, "evictions": 4029, "flash_bytes": 5429182464, "hits": 635, "io_failures": 0, "lookup_ns": 0, "lookup_ns_per_tok": 0.0, "misses": 2095, "prefetch_dropped": 0, "prefetch_issued": 4216, "prefetch_wasted": 2220, "resident_bytes": 100564992, "slots_total": 62, "slots_used": 62, "stall_ns": 637740209, "sync_rescues": 0}, "prefill_cold_s": 8.443465491923641, "prefill_ms": 2724, "prefill_tok_s": 8.443465491923641, "prefill_warm_s": 0.0, "prompt_tokens": 23, "route": {"agree": null, "dec_steps": 8, "feature": "executed", "mode": "predicted", "pred_used": 128}, "schema": 1, "thinking_on": true}
\ No newline at end of file
diff --git a/android/build.gradle.kts b/android/build.gradle.kts
new file mode 100644
index 0000000..97a679c
--- /dev/null
+++ b/android/build.gradle.kts
@@ -0,0 +1,9 @@
+plugins {
+    alias(libs.plugins.android.library) apply false
+    alias(libs.plugins.android.application) apply false
+    alias(libs.plugins.kotlin.android) apply false
+    // compose compiler plugin version must equal Kotlin (2.0.21); Room via KSP; metrics via kotlinx-serialization
+    alias(libs.plugins.kotlin.compose) apply false
+    alias(libs.plugins.ksp) apply false
+    alias(libs.plugins.kotlinx.serialization) apply false
+}
diff --git a/android/gradle.properties b/android/gradle.properties
new file mode 100644
index 0000000..d703336
--- /dev/null
+++ b/android/gradle.properties
@@ -0,0 +1,5 @@
+org.gradle.jvmargs=-Xmx4g -Dfile.encoding=UTF-8
+org.gradle.caching=true
+org.gradle.configuration-cache=true
+android.useAndroidX=true
+android.nonTransitiveRClass=true
diff --git a/android/gradle/libs.versions.toml b/android/gradle/libs.versions.toml
new file mode 100644
index 0000000..07b528d
--- /dev/null
+++ b/android/gradle/libs.versions.toml
@@ -0,0 +1,47 @@
+[versions]
+agp = "8.7.3"
+kotlin = "2.0.21"
+# aligned with the target platform constraints in the docs
+minSdk = "33"
+compileSdk = "35"
+ndk = "27.2.12479018"
+# pinned, verified combination; compose compiler plugin version must equal Kotlin
+composeBom = "2024.12.01"
+activityCompose = "1.9.3"
+navigationCompose = "2.8.5"
+lifecycle = "2.8.7"
+room = "2.7.0"
+ksp = "2.0.21-1.0.28"
+datastore = "1.1.1"
+coroutines = "1.9.0"
+serializationJson = "1.7.3"
+markdownRenderer = "0.27.0"   # last release binary-compatible with Kotlin 2.0.21
+
+[libraries]
+# Compose (versions governed by the BOM)
+compose-bom = { group = "androidx.compose", name = "compose-bom", version.ref = "composeBom" }
+compose-ui = { group = "androidx.compose.ui", name = "ui" }
+compose-ui-tooling = { group = "androidx.compose.ui", name = "ui-tooling" }
+compose-ui-tooling-preview = { group = "androidx.compose.ui", name = "ui-tooling-preview" }
+compose-material3 = { group = "androidx.compose.material3", name = "material3" }
+compose-material-icons-extended = { group = "androidx.compose.material", name = "material-icons-extended" }
+androidx-activity-compose = { group = "androidx.activity", name = "activity-compose", version.ref = "activityCompose" }
+androidx-navigation-compose = { group = "androidx.navigation", name = "navigation-compose", version.ref = "navigationCompose" }
+androidx-lifecycle-viewmodel-compose = { group = "androidx.lifecycle", name = "lifecycle-viewmodel-compose", version.ref = "lifecycle" }
+androidx-lifecycle-runtime-compose = { group = "androidx.lifecycle", name = "lifecycle-runtime-compose", version.ref = "lifecycle" }
+androidx-lifecycle-process = { group = "androidx.lifecycle", name = "lifecycle-process", version.ref = "lifecycle" }
+androidx-room-runtime = { group = "androidx.room", name = "room-runtime", version.ref = "room" }
+androidx-room-ktx = { group = "androidx.room", name = "room-ktx", version.ref = "room" }
+androidx-room-compiler = { group = "androidx.room", name = "room-compiler", version.ref = "room" }
+androidx-datastore-preferences = { group = "androidx.datastore", name = "datastore-preferences", version.ref = "datastore" }
+kotlinx-coroutines-android = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-android", version.ref = "coroutines" }
+kotlinx-serialization-json = { group = "org.jetbrains.kotlinx", name = "kotlinx-serialization-json", version.ref = "serializationJson" }
+markdown-renderer-m3 = { group = "com.mikepenz", name = "multiplatform-markdown-renderer-m3", version.ref = "markdownRenderer" }
+
+[plugins]
+android-library = { id = "com.android.library", version.ref = "agp" }
+android-application = { id = "com.android.application", version.ref = "agp" }
+kotlin-android = { id = "org.jetbrains.kotlin.android", version.ref = "kotlin" }
+kotlin-compose = { id = "org.jetbrains.kotlin.plugin.compose", version.ref = "kotlin" }
+ksp = { id = "com.google.devtools.ksp", version.ref = "ksp" }
+kotlinx-serialization = { id = "org.jetbrains.kotlin.plugin.serialization", version.ref = "kotlin" }
diff --git a/android/gradle/wrapper/gradle-wrapper.jar b/android/gradle/wrapper/gradle-wrapper.jar
new file mode 100644
index 0000000..a4b76b9
Binary files /dev/null and b/android/gradle/wrapper/gradle-wrapper.jar differ
diff --git a/android/gradle/wrapper/gradle-wrapper.properties b/android/gradle/wrapper/gradle-wrapper.properties
new file mode 100644
index 0000000..e2847c8
--- /dev/null
+++ b/android/gradle/wrapper/gradle-wrapper.properties
@@ -0,0 +1,7 @@
+distributionBase=GRADLE_USER_HOME
+distributionPath=wrapper/dists
+distributionUrl=https\://services.gradle.org/distributions/gradle-8.11.1-bin.zip
+networkTimeout=10000
+validateDistributionUrl=true
+zipStoreBase=GRADLE_USER_HOME
+zipStorePath=wrapper/dists
diff --git a/android/gradlew b/android/gradlew
new file mode 100644
index 0000000..d95bf61
--- /dev/null
+++ b/android/gradlew
@@ -0,0 +1,252 @@
+#!/bin/sh
+
+#
+# Copyright © 2015-2021 the original authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+##############################################################################
+#
+#   Gradle start up script for POSIX generated by Gradle.
+#
+#   Important for running:
+#
+#   (1) You need a POSIX-compliant shell to run this script. If your /bin/sh is
+#       noncompliant, but you have some other compliant shell such as ksh or
+#       bash, then to run this script, type that shell name before the whole
+#       command line, like:
+#
+#           ksh Gradle
+#
+#       Busybox and similar reduced shells will NOT work, because this script
+#       requires all of these POSIX shell features:
+#         * functions;
+#         * expansions «$var», «${var}», «${var:-default}», «${var+SET}»,
+#           «${var#prefix}», «${var%suffix}», and «$( cmd )»;
+#         * compound commands having a testable exit status, especially «case»;
+#         * various built-in commands including «command», «set», and «ulimit».
+#
+#   Important for patching:
+#
+#   (2) This script targets any POSIX shell, so it avoids extensions provided
+#       by Bash, Ksh, etc; in particular arrays are avoided.
+#
+#       The "traditional" practice of packing multiple parameters into a
+#       space-separated string is a well documented source of bugs and security
+#       problems, so this is (mostly) avoided, by progressively accumulating
+#       options in "$@", and eventually passing that to Java.
+#
+#       Where the inherited environment variables (DEFAULT_JVM_OPTS, JAVA_OPTS,
+#       and GRADLE_OPTS) rely on word-splitting, this is performed explicitly;
+#       see the in-line comments for details.
+#
+#       There are tweaks for specific operating systems such as AIX, CygWin,
+#       Darwin, MinGW, and NonStop.
+#
+#   (3) This script is generated from the Groovy template
+#       https://github.com/gradle/gradle/blob/HEAD/platforms/jvm/plugins-application/src/main/resources/org/gradle/api/internal/plugins/unixStartScript.txt
+#       within the Gradle project.
+#
+#       You can find Gradle at https://github.com/gradle/gradle/.
+#
+##############################################################################
+
+# Attempt to set APP_HOME
+
+# Resolve links: $0 may be a link
+app_path=$0
+
+# Need this for daisy-chained symlinks.
+while
+    APP_HOME=${app_path%"${app_path##*/}"}  # leaves a trailing /; empty if no leading path
+    [ -h "$app_path" ]
+do
+    ls=$( ls -ld "$app_path" )
+    link=${ls#*' -> '}
+    case $link in             #(
+      /*)   app_path=$link ;; #(
+      *)    app_path=$APP_HOME$link ;;
+    esac
+done
+
+# This is normally unused
+# shellcheck disable=SC2034
+APP_BASE_NAME=${0##*/}
+# Discard cd standard output in case $CDPATH is set (https://github.com/gradle/gradle/issues/25036)
+APP_HOME=$( cd -P "${APP_HOME:-./}" > /dev/null && printf '%s
+' "$PWD" ) || exit
+
+# Use the maximum available, or set MAX_FD != -1 to use that value.
+MAX_FD=maximum
+
+warn () {
+    echo "$*"
+} >&2
+
+die () {
+    echo
+    echo "$*"
+    echo
+    exit 1
+} >&2
+
+# OS specific support (must be 'true' or 'false').
+cygwin=false
+msys=false
+darwin=false
+nonstop=false
+case "$( uname )" in                #(
+  CYGWIN* )         cygwin=true  ;; #(
+  Darwin* )         darwin=true  ;; #(
+  MSYS* | MINGW* )  msys=true    ;; #(
+  NONSTOP* )        nonstop=true ;;
+esac
+
+CLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar
+
+
+# Determine the Java command to use to start the JVM.
+if [ -n "$JAVA_HOME" ] ; then
+    if [ -x "$JAVA_HOME/jre/sh/java" ] ; then
+        # IBM's JDK on AIX uses strange locations for the executables
+        JAVACMD=$JAVA_HOME/jre/sh/java
+    else
+        JAVACMD=$JAVA_HOME/bin/java
+    fi
+    if [ ! -x "$JAVACMD" ] ; then
+        die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME
+
+Please set the JAVA_HOME variable in your environment to match the
+location of your Java installation."
+    fi
+else
+    JAVACMD=java
+    if ! command -v java >/dev/null 2>&1
+    then
+        die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+
+Please set the JAVA_HOME variable in your environment to match the
+location of your Java installation."
+    fi
+fi
+
+# Increase the maximum file descriptors if we can.
+if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then
+    case $MAX_FD in #(
+      max*)
+        # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked.
+        # shellcheck disable=SC2039,SC3045
+        MAX_FD=$( ulimit -H -n ) ||
+            warn "Could not query maximum file descriptor limit"
+    esac
+    case $MAX_FD in  #(
+      '' | soft) :;; #(
+      *)
+        # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked.
+        # shellcheck disable=SC2039,SC3045
+        ulimit -n "$MAX_FD" ||
+            warn "Could not set maximum file descriptor limit to $MAX_FD"
+    esac
+fi
+
+# Collect all arguments for the java command, stacking in reverse order:
+#   * args from the command line
+#   * the main class name
+#   * -classpath
+#   * -D...appname settings
+#   * --module-path (only if needed)
+#   * DEFAULT_JVM_OPTS, JAVA_OPTS, and GRADLE_OPTS environment variables.
+
+# For Cygwin or MSYS, switch paths to Windows format before running java
+if "$cygwin" || "$msys" ; then
+    APP_HOME=$( cygpath --path --mixed "$APP_HOME" )
+    CLASSPATH=$( cygpath --path --mixed "$CLASSPATH" )
+
+    JAVACMD=$( cygpath --unix "$JAVACMD" )
+
+    # Now convert the arguments - kludge to limit ourselves to /bin/sh
+    for arg do
+        if
+            case $arg in                                #(
+              -*)   false ;;                            # don't mess with options #(
+              /?*)  t=${arg#/} t=/${t%%/*}              # looks like a POSIX filepath
+                    [ -e "$t" ] ;;                      #(
+              *)    false ;;
+            esac
+        then
+            arg=$( cygpath --path --ignore --mixed "$arg" )
+        fi
+        # Roll the args list around exactly as many times as the number of
+        # args, so each arg winds up back in the position where it started, but
+        # possibly modified.
+        #
+        # NB: a `for` loop captures its iteration list before it begins, so
+        # changing the positional parameters here affects neither the number of
+        # iterations, nor the values presented in `arg`.
+        shift                   # remove old arg
+        set -- "$@" "$arg"      # push replacement arg
+    done
+fi
+
+
+# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+DEFAULT_JVM_OPTS='-Dfile.encoding=UTF-8 "-Xmx64m" "-Xms64m"'
+
+# Collect all arguments for the java command:
+#   * DEFAULT_JVM_OPTS, JAVA_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments,
+#     and any embedded shellness will be escaped.
+#   * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be
+#     treated as '${Hostname}' itself on the command line.
+
+set -- \
+        "-Dorg.gradle.appname=$APP_BASE_NAME" \
+        -classpath "$CLASSPATH" \
+        org.gradle.wrapper.GradleWrapperMain \
+        "$@"
+
+# Stop when "xargs" is not available.
+if ! command -v xargs >/dev/null 2>&1
+then
+    die "xargs is not available"
+fi
+
+# Use "xargs" to parse quoted args.
+#
+# With -n1 it outputs one arg per line, with the quotes and backslashes removed.
+#
+# In Bash we could simply go:
+#
+#   readarray ARGS < <( xargs -n1 <<<"$var" ) &&
+#   set -- "${ARGS[@]}" "$@"
+#
+# but POSIX shell has neither arrays nor command substitution, so instead we
+# post-process each arg (as a line of input to sed) to backslash-escape any
+# character that might be a shell metacharacter, then use eval to reverse
+# that process (while maintaining the separation between arguments), and wrap
+# the whole thing up as a single "set" statement.
+#
+# This will of course break if any of these variables contains a newline or
+# an unmatched quote.
+#
+
+eval "set -- $(
+        printf '%s\n' "$DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS" |
+        xargs -n1 |
+        sed ' s~[^-[:alnum:]+,./:=@_]~\\&~g; ' |
+        tr '\n' ' '
+    )" '"$@"'
+
+exec "$JAVACMD" "$@"
diff --git a/android/gradlew.bat b/android/gradlew.bat
new file mode 100755
index 0000000..640d686
--- /dev/null
+++ b/android/gradlew.bat
@@ -0,0 +1,94 @@
+@rem
+@rem Copyright 2015 the original author or authors.
+@rem
+@rem Licensed under the Apache License, Version 2.0 (the "License");
+@rem you may not use this file except in compliance with the License.
+@rem You may obtain a copy of the License at
+@rem
+@rem      https://www.apache.org/licenses/LICENSE-2.0
+@rem
+@rem Unless required by applicable law or agreed to in writing, software
+@rem distributed under the License is distributed on an "AS IS" BASIS,
+@rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+@rem See the License for the specific language governing permissions and
+@rem limitations under the License.
+@rem
+@rem SPDX-License-Identifier: Apache-2.0
+@rem
+
+@if "%DEBUG%"=="" @echo off
+@rem ##########################################################################
+@rem
+@rem  Gradle startup script for Windows
+@rem
+@rem ##########################################################################
+
+@rem Set local scope for the variables with windows NT shell
+if "%OS%"=="Windows_NT" setlocal
+
+set DIRNAME=%~dp0
+if "%DIRNAME%"=="" set DIRNAME=.
+@rem This is normally unused
+set APP_BASE_NAME=%~n0
+set APP_HOME=%DIRNAME%
+
+@rem Resolve any "." and ".." in APP_HOME to make it shorter.
+for %%i in ("%APP_HOME%") do set APP_HOME=%%~fi
+
+@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+set DEFAULT_JVM_OPTS=-Dfile.encoding=UTF-8 "-Xmx64m" "-Xms64m"
+
+@rem Find java.exe
+if defined JAVA_HOME goto findJavaFromJavaHome
+
+set JAVA_EXE=java.exe
+%JAVA_EXE% -version >NUL 2>&1
+if %ERRORLEVEL% equ 0 goto execute
+
+echo. 1>&2
+echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 1>&2
+echo. 1>&2
+echo Please set the JAVA_HOME variable in your environment to match the 1>&2
+echo location of your Java installation. 1>&2
+
+goto fail
+
+:findJavaFromJavaHome
+set JAVA_HOME=%JAVA_HOME:"=%
+set JAVA_EXE=%JAVA_HOME%/bin/java.exe
+
+if exist "%JAVA_EXE%" goto execute
+
+echo. 1>&2
+echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 1>&2
+echo. 1>&2
+echo Please set the JAVA_HOME variable in your environment to match the 1>&2
+echo location of your Java installation. 1>&2
+
+goto fail
+
+:execute
+@rem Setup the command line
+
+set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
+
+
+@rem Execute Gradle
+"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %*
+
+:end
+@rem End local scope for the variables with windows NT shell
+if %ERRORLEVEL% equ 0 goto mainEnd
+
+:fail
+rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
+rem the _cmd.exe /c_ return code!
+set EXIT_CODE=%ERRORLEVEL%
+if %EXIT_CODE% equ 0 set EXIT_CODE=1
+if not ""=="%GRADLE_EXIT_CONSOLE%" exit %EXIT_CODE%
+exit /b %EXIT_CODE%
+
+:mainEnd
+if "%OS%"=="Windows_NT" endlocal
+
+:omega
diff --git a/android/settings.gradle.kts b/android/settings.gradle.kts
new file mode 100644
index 0000000..408eaee
--- /dev/null
+++ b/android/settings.gradle.kts
@@ -0,0 +1,18 @@
+pluginManagement {
+    repositories {
+        google()
+        mavenCentral()
+        gradlePluginPortal()
+    }
+}
+
+dependencyResolutionManagement {
+    repositoriesMode.set(RepositoriesMode.FAIL_ON_PROJECT_REPOS)
+    repositories {
+        google()
+        mavenCentral()
+    }
+}
+
+rootProject.name = "edge0-runtime-android"
+include(":app")  // single APK: native engine + Compose UI
diff --git a/android/tools/llama/build_vendor_libs.sh b/android/tools/llama/build_vendor_libs.sh
new file mode 100644
index 0000000..49a9006
--- /dev/null
+++ b/android/tools/llama/build_vendor_libs.sh
@@ -0,0 +1,63 @@
+#!/usr/bin/env bash
+# build_vendor_libs.sh - rebuild the four shared libraries from the pinned supply tree.
+# Supply topology (patches/llama.cpp/README in the depot repo):
+#   vendor@pin (vendor.llama.pin at the depot root; materialized by clone on first run -
+#   no submodule) -> wt/and worktree with common(6)+android(14) bands replayed in order
+#   -> golden tree hash. Set EDGE0_LLAMA_URL to clone from a mirror instead of GitHub.
+# Usage: bash tools/llama/build_vendor_libs.sh [--replay]  (--replay = re-apply bands + verify)
+set -euo pipefail
+REPO=$(cd "$(dirname "$0")/../.." && pwd)
+DEPOT=${EDGE0_DEPOT:-$(cd "$REPO/.." && pwd -W)}   # depot root (this repo lives at <depot>/android)
+WT=$DEPOT/wt/and
+[ -f "$DEPOT/vendor.llama.pin" ] || { echo "ABORT: vendor.llama.pin missing at depot root ($DEPOT)"; exit 1; }
+PIN=$(cut -d' ' -f1 "$DEPOT/vendor.llama.pin")
+VENDOR=$DEPOT/vendor/llama.cpp
+GOLD_TREE=e974be50ba5c184bf0ba9a4a26adbbae4235c2ab
+NDK=${NDK_DIR:-${ANDROID_NDK_HOME:-}}
+[ -n "$NDK" ] || { echo "ABORT: set NDK_DIR (or ANDROID_NDK_HOME) to the Android NDK r28 path"; exit 1; }
+LL=$REPO/build-dl/llama-libs
+
+# materialize the pinned vendor tree when absent (submodule-free supply)
+if [ ! -d "$VENDOR/.git" ]; then
+    echo "[0/3] materializing vendor: clone llama.cpp at pinned $PIN -> $VENDOR"
+    mkdir -p "$DEPOT/vendor"
+    git clone --filter=blob:none "${EDGE0_LLAMA_URL:-https://github.com/ggml-org/llama.cpp}" "$VENDOR"
+fi
+[ "$(git -C "$VENDOR" rev-parse HEAD)" = "$PIN" ] || git -C "$VENDOR" checkout --detach "$PIN"
+[ -z "$(git -C "$VENDOR" status --porcelain)" ] || { echo "ABORT: vendor tree dirty - the pristine supply is never patched in place; restore with: git -C $VENDOR checkout --detach $PIN"; exit 1; }
+
+# self-heal: a fresh clone has no consumer worktree yet - create it at the pin,
+# and treat the bare tree as needing the band replay (otherwise the tree gate
+# would compare an unpatched pin against the golden patched tree -> false RED)
+FRESH_WT=0
+if [ ! -d "$WT" ]; then
+    git -C "$VENDOR" worktree add --detach "$WT" "$PIN"
+    FRESH_WT=1
+fi
+
+if [ "${1:-}" = "--replay" ] || [ "$FRESH_WT" = "1" ]; then
+    git -C "$WT" am --abort 2>/dev/null || true
+    git -C "$WT" reset --hard "$PIN"
+    for band in common android; do
+        for f in $(ls "$DEPOT/patches/llama.cpp/$band/"0*.patch | sort); do
+            git -C "$WT" am --3way "$f"
+        done
+    done
+fi
+T=$(git -C "$WT" rev-parse 'HEAD^{tree}')
+[ "$T" = "$GOLD_TREE" ] || { echo "ABORT: tree mismatch got=$T want=$GOLD_TREE"; exit 1; }
+echo "tree GREEN $T"
+
+cmake -S "$WT" -B "$WT/build-android-cpu" -G Ninja \
+    -DCMAKE_TOOLCHAIN_FILE="$NDK/build/cmake/android.toolchain.cmake" \
+    -DANDROID_ABI=arm64-v8a -DANDROID_PLATFORM=android-28 -DCMAKE_BUILD_TYPE=Release \
+    -DBUILD_SHARED_LIBS=ON -DLLAMA_BUILD_COMMON=ON -DGGML_OPENCL=OFF -DGGML_NATIVE=OFF
+ninja -C "$WT/build-android-cpu" llama ggml ggml-base ggml-cpu
+
+mkdir -p "$LL/arm64-v8a" "$LL/include"
+cp "$WT"/build-android-cpu/bin/{libllama.so,libggml.so,libggml-base.so,libggml-cpu.so} "$LL/arm64-v8a/"
+# libggml-cpu DT_NEEDED libomp.so (NDK OpenMP runtime) - ship it alongside,
+# otherwise the app fails dlopen at launch on a fresh build.
+cp "$NDK"/toolchains/llvm/prebuilt/*/lib/clang/*/lib/linux/aarch64/libomp.so "$LL/arm64-v8a/"
+cp "$WT"/include/*.h "$WT"/ggml/include/*.h "$LL/include/"
+echo "OK: four libs from wt/and @$(git -C "$WT" rev-parse --short HEAD)"
diff --git a/android/tools/model/push_models.sh b/android/tools/model/push_models.sh
new file mode 100644
index 0000000..53258a2
--- /dev/null
+++ b/android/tools/model/push_models.sh
@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+# push_models.sh - stage the three GGUF artifacts onto the device files/models/.
+# Usage: bash tools/model/push_models.sh --8b | --35b | --all
+#   Default sources are this repo's models/edge0-{8b,35b}-gguf/ (converter output).
+#   SRC_DIR=<dir> overrides; ADB_SERIAL selects the device (default: first online).
+# Triple-hop discipline: host md5 -> adb push -> run-as copy -> on-device md5.
+# Any 8-hex mismatch aborts. The 21.7GB artifact takes 2-6 min; reruns are idempotent.
+set -euo pipefail
+REPO=$(cd "$(dirname "$0")/../.." && pwd)
+# Self-contained layout: sources and converted artifacts live under models/ (gitignored).
+# The converter (../windows/tools/convert_mlx_to_gguf.py --dir models/edge0-{tier})
+# writes models/edge0-{tier}-gguf/; names pass through unchanged.
+# SRC_DIR=<dir> overrides with a single flat directory.
+src_for() {  # $1 = artifact name -> host source path (convention dirs unless SRC_DIR set)
+  local dir
+  case "$1" in edge0-8b.gguf|lora_edge0_8b-gguf.gguf) dir=${SRC:-$REPO/models/edge0-8b-gguf} ;;
+                 *)                                     dir=${SRC:-$REPO/models/edge0-35b-gguf} ;;
+  esac
+  echo "$dir/$1"
+}
+DEV=${ADB_SERIAL:-}
+PKG=dev.edge0.runtime.app
+T=/data/local/tmp/lgguf
+
+# artifact -> official md5 prefix baseline (converter convention names, 2026-09)
+md5_of() { case "$1" in
+  edge0-8b.gguf)         echo b5f6021f ;;
+  lora_edge0_8b-gguf.gguf) echo e73c92e3 ;;
+  edge0-35b.gguf)        echo 7d2c2e8c ;;
+  *) echo "?" ;; esac; }
+
+FILES_8B=(edge0-8b.gguf lora_edge0_8b-gguf.gguf)
+FILES_35B=(edge0-35b.gguf)
+case "${1:-}" in
+  --8b)  FILES=("${FILES_8B[@]}") ;;
+  --35b) FILES=("${FILES_35B[@]}") ;;
+  --all) FILES=("${FILES_8B[@]}" "${FILES_35B[@]}") ;;
+  *) echo "usage: push_models.sh --8b|--35b|--all" >&2; exit 2 ;;
+esac
+
+ADB=(adb ${DEV:+-s "$DEV"})
+"${ADB[@]}" get-state >/dev/null || { echo "ABORT: no device ($DEV)" >&2; exit 1; }
+"${ADB[@]}" shell "mkdir -p $T"
+for f in "${FILES[@]}"; do
+  want=$(md5_of "$f"); src="$(src_for "$f")"
+  [ -f "$src" ] || { echo "ABORT: missing $src (download + run converter, or place artifact)"; exit 1; }
+  got=$(md5sum "$src" | cut -c1-8)
+  [ "$got" = "$want" ] || { echo "ABORT: $f host md5=$got want=$want"; exit 1; }
+  echo "push $f ($(du -h "$src" | cut -f1)) ..."
+  MSYS_NO_PATHCONV=1 "${ADB[@]}" push "$(cygpath -w "$src" 2>/dev/null || echo "$src")" "$T/$f" >/dev/null
+  # copy into the app-private dir via run-as (non-root path) + on-device verification
+  "${ADB[@]}" shell "run-as $PKG mkdir -p files/models; run-as $PKG cp $T/$f files/models/$f"
+  dev=$(MSYS_NO_PATHCONV=1 "${ADB[@]}" shell "run-as $PKG md5sum files/models/$f" | tr -d '\r' | cut -c1-8)
+  [ "$dev" = "$want" ] || { echo "ABORT: $f device md5=$dev want=$want (truncated transfer? re-run)"; exit 1; }
+  echo "OK $f md5=$dev"
+done
+echo "staged -> $PKG/files/models"
diff --git a/docs/adding-a-model.md b/docs/adding-a-model.md
index 6af69b4..be304bb 100644
--- a/docs/adding-a-model.md
+++ b/docs/adding-a-model.md
@@ -42,7 +42,7 @@ The subclass only needs to implement the class method `_defaults(model_dir) -> C
 
 Key points:
 
-- The repository root ships `scripts/convert_adapters_legacy.py`, which one-shot converts the training npz exports into safetensors artifacts under `artifacts/`; `ModelConfig.artifact(name)` (`models/base.py:28–30`) returns the absolute paths of those artifacts, and the adapter uses them to fill in the LoRA / prerouter weight paths.
+- `ModelConfig.artifact(name)` (`models/base.py`) resolves an adapter's absolute path — from the model directory first, falling back to the gitignored `artifacts/` cache — and the adapter uses it to fill in the LoRA / prerouter weight paths.
 - LoRA overrides go through `resolve_lora` (`models/base.py:88–95`): a bare model name resolves to that tier's artifacts, `"model_dir"` means "keep the training weights in place", and an empty string disables it.
 
 ## Step 3: `build_model` / `build_engine`
@@ -105,7 +105,7 @@ The lookup order of `_resolve_name` (registry.py:42–56): explicit `name` → t
 
 - `config.json` exists, and its `model_type` is registered (or covered by your `TYPE_ALIASES`);
 - the weights are in safetensors format, with key prefixes matching `moe_spec.key_template` (e.g. `language_model.model.layers.N.mlp.switch_mlp`);
-- the LoRA / prerouter weights are `.safetensors` files carrying metadata (see the comment at the top of `models/base.py`: they are converted one-shot from the training npz by `convert_adapters_legacy.py`).
+- the LoRA / prerouter weights are `.safetensors` files carrying metadata (see the comment at the top of `models/base.py`).
 
 The directory basename is the last-resort resolution fallback (`_model_type_from_dir`), so aligning the directory name with the registered name is recommended, but not required.
 
diff --git a/docs/architecture.md b/docs/architecture.md
index 3a07c42..c5adc22 100644
--- a/docs/architecture.md
+++ b/docs/architecture.md
@@ -19,12 +19,10 @@ box: `edge0-35b` (Qwen3.6-35B-A3B, K=4) and `edge0-8b` (Ling 3.0 hybrid, K=8).
    future CUDA backend (`backends/cuda/`, selected via the `EDGE0_BACKEND`
    environment variable);
 3. **Adapters unified on safetensors**: LoRA and prerouter weights are both
-   `.safetensors` files carrying metadata; legacy npz training exports are
-   converted by a one-off migration script and then deprecated;
+   `.safetensors` files carrying metadata;
 4. **Unified terminology**: the pre-routing head is always called the
-   **prerouter** — zero legacy-term residue in code and docs. The legacy npz
-   exports are migrated by `scripts/convert_adapters_legacy.py`, which
-   normalizes their key namespace to `layers.<N>.<part>.weight`.
+   **prerouter** — zero legacy-term residue in code and docs. Prerouter
+   weight keys are normalized to `layers.<N>.<part>.weight`.
 
 ## Layered Structure
 
diff --git a/docs/models/edge0-35b.md b/docs/models/edge0-35b.md
index 99d611b..e6c12af 100644
--- a/docs/models/edge0-35b.md
+++ b/docs/models/edge0-35b.md
@@ -22,7 +22,7 @@ The performance profile is based on benchmarks of the current release adapter ve
 | Prefill chunk | 2048 |
 | Hot window | 4 |
 | Streaming prefetch history | on (`prefetch_history=True`) |
-| Serving port | 8085 |
+| Serving port | 8000 |
 | Measured throughput | 14.9–17.7 tok/s (M4 Pro) |
 | Measured peak activation memory | ≈ 2.9 GiB |
 
@@ -42,7 +42,7 @@ This tier uses the `LayerOptions.staged_k4()` preset:
 ### Serving via the CLI
 
 ```bash
-edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8085
+edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8000
 ```
 
 Optional arguments:
@@ -99,7 +99,7 @@ cfg = AutoConfig.from_pretrained("/path/to/checkpoint", name="edge0-35b")
 ### Non-streaming chat
 
 ```bash
-curl -s http://127.0.0.1:8085/v1/chat/completions \
+curl -s http://127.0.0.1:8000/v1/chat/completions \
   -H 'Content-Type: application/json' \
   -d '{
     "model": "edge0-35b",
@@ -131,7 +131,7 @@ Optional request fields: `model`, `messages` (with `role`/`content`; content sup
 ### Streaming chat (requires Flask)
 
 ```bash
-curl -N http://127.0.0.1:8085/v1/chat/completions \
+curl -N http://127.0.0.1:8000/v1/chat/completions \
   -H 'Content-Type: application/json' \
   -d '{
     "model": "edge0-35b",
@@ -152,7 +152,7 @@ from edge0 import AutoEngine
 engine = AutoEngine.from_pretrained(
     "/path/to/checkpoint",
     name="edge0-35b",
-    port=9090,                    # override the default port 8085
+    port=9090,                    # override the default port 8000
     target_tok_s=14.0,            # override the acceptance throughput target
     prerouter=None,               # disable the prerouter
     lora="",                      # disable LoRA
diff --git a/docs/models/edge0-8b.md b/docs/models/edge0-8b.md
index 55c2db0..a71943f 100644
--- a/docs/models/edge0-8b.md
+++ b/docs/models/edge0-8b.md
@@ -22,7 +22,7 @@ The performance profile is based on benchmarks of the current release adapter ve
 | Prefill chunk | 2048 |
 | Hot window | 1 |
 | Streaming prefetch history | on (`prefetch_history=True`) |
-| Serving port | 8083 |
+| Serving port | 8000 |
 | Measured throughput | 23.9–25.3 tok/s (M4 Pro) |
 | Measured peak activation memory | ≈ 1.0 GiB |
 
@@ -41,7 +41,7 @@ This tier uses the `LayerOptions.prod_k8()` preset (aligned with the reference d
 ### Serving via the CLI
 
 ```bash
-edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8083
+edge0 serve /path/to/checkpoint --host 127.0.0.1 --port 8000
 ```
 
 Optional arguments:
@@ -99,7 +99,7 @@ cfg = AutoConfig.from_pretrained("/path/to/checkpoint", name="edge0-8b")
 ### Non-streaming chat
 
 ```bash
-curl -s http://127.0.0.1:8083/v1/chat/completions \
+curl -s http://127.0.0.1:8000/v1/chat/completions \
   -H 'Content-Type: application/json' \
   -d '{
     "model": "edge0-8b",
@@ -131,7 +131,7 @@ Optional request fields: `model`, `messages` (with `role`/`content`; content sup
 ### Streaming chat (requires Flask)
 
 ```bash
-curl -N http://127.0.0.1:8083/v1/chat/completions \
+curl -N http://127.0.0.1:8000/v1/chat/completions \
   -H 'Content-Type: application/json' \
   -d '{
     "model": "edge0-8b",
@@ -152,7 +152,7 @@ from edge0 import AutoEngine
 engine = AutoEngine.from_pretrained(
     "/path/to/checkpoint",
     name="edge0-8b",
-    port=9083,                    # override the default port 8083
+    port=9083,                    # override the default port 8000
     target_tok_s=35.0,            # override the acceptance throughput target
     prerouter=None,               # disable the prerouter
     lora="",                      # disable LoRA
diff --git a/docs/prerouter.md b/docs/prerouter.md
index 460a1da..8744fa4 100644
--- a/docs/prerouter.md
+++ b/docs/prerouter.md
@@ -73,7 +73,7 @@ per `start_layer`/`owners`. Family differences are expressed through three hooks
 | weight file | `prerouter_edge0_35b.safetensors` | `prerouter_edge0_8b.safetensors` |
 
 Both weight files are resolved from the model directory by default, with
-`artifacts/` (repo root, gitignored) as the fallback — see the README's
+`artifacts/` (Python project root, gitignored) as the fallback — see the README's
 "Models and adapters" section.
 
 ## Tests and regression
diff --git a/ios/Config/App.xcconfig b/ios/Config/App.xcconfig
new file mode 100644
index 0000000..234f81e
--- /dev/null
+++ b/ios/Config/App.xcconfig
@@ -0,0 +1 @@
+#include? "Local.xcconfig"
diff --git a/ios/Config/Local.xcconfig.example b/ios/Config/Local.xcconfig.example
new file mode 100644
index 0000000..80c6e5d
--- /dev/null
+++ b/ios/Config/Local.xcconfig.example
@@ -0,0 +1,3 @@
+// Copy this file to Config/Local.xcconfig and set your Apple Developer Team ID.
+// Config/Local.xcconfig is excluded by .gitignore.
+DEVELOPMENT_TEAM =
diff --git a/ios/Edge0PhoneProbe.xcodeproj/project.pbxproj b/ios/Edge0PhoneProbe.xcodeproj/project.pbxproj
new file mode 100644
index 0000000..4bdcaac
--- /dev/null
+++ b/ios/Edge0PhoneProbe.xcodeproj/project.pbxproj
@@ -0,0 +1,259 @@
+// !$*UTF8*$!
+{
+	archiveVersion = 1;
+	classes = {
+	};
+	objectVersion = 71;
+	objects = {
+
+/* Begin PBXBuildFile section */
+		A00000000000000000000001 /* Edge0PhoneProbeApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = A00000000000000000000011 /* Edge0PhoneProbeApp.swift */; };
+		A00000000000000000000002 /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = A00000000000000000000012 /* ContentView.swift */; };
+		A00000000000000000000004 /* Edge0Core in Frameworks */ = {isa = PBXBuildFile; productRef = A00000000000000000000071 /* Edge0Core */; };
+		A00000000000000000000005 /* Edge0MLX in Frameworks */ = {isa = PBXBuildFile; productRef = A00000000000000000000072 /* Edge0MLX */; };
+		A00000000000000000000006 /* Models in Resources */ = {isa = PBXBuildFile; fileRef = A00000000000000000000014 /* Models */; };
+/* End PBXBuildFile section */
+
+/* Begin PBXFileReference section */
+		A00000000000000000000003 /* Edge0Phone.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = Edge0Phone.app; sourceTree = BUILT_PRODUCTS_DIR; };
+		A00000000000000000000011 /* Edge0PhoneProbeApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Edge0PhoneProbeApp.swift; sourceTree = "<group>"; };
+		A00000000000000000000012 /* ContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = "<group>"; };
+		A00000000000000000000013 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
+		A00000000000000000000014 /* Models */ = {isa = PBXFileReference; lastKnownFileType = folder; path = Models; sourceTree = "<group>"; };
+		A00000000000000000000015 /* App.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; path = App.xcconfig; sourceTree = "<group>"; };
+/* End PBXFileReference section */
+
+/* Begin PBXFrameworksBuildPhase section */
+		A00000000000000000000021 /* Frameworks */ = {
+			isa = PBXFrameworksBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				A00000000000000000000004 /* Edge0Core in Frameworks */,
+				A00000000000000000000005 /* Edge0MLX in Frameworks */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXFrameworksBuildPhase section */
+
+/* Begin PBXGroup section */
+		A00000000000000000000022 = {
+			isa = PBXGroup;
+			children = (
+				A00000000000000000000033 /* Config */,
+				A00000000000000000000014 /* Models */,
+				A00000000000000000000031 /* Edge0PhoneProbe */,
+				A00000000000000000000032 /* Products */,
+			);
+			sourceTree = "<group>";
+		};
+		A00000000000000000000031 /* Edge0PhoneProbe */ = {
+			isa = PBXGroup;
+			children = (
+				A00000000000000000000011 /* Edge0PhoneProbeApp.swift */,
+				A00000000000000000000012 /* ContentView.swift */,
+				A00000000000000000000013 /* Info.plist */,
+			);
+			path = Edge0PhoneProbe;
+			sourceTree = "<group>";
+		};
+		A00000000000000000000032 /* Products */ = {
+			isa = PBXGroup;
+			children = (
+				A00000000000000000000003 /* Edge0Phone.app */,
+			);
+			name = Products;
+			sourceTree = "<group>";
+		};
+		A00000000000000000000033 /* Config */ = {
+			isa = PBXGroup;
+			children = (
+				A00000000000000000000015 /* App.xcconfig */,
+			);
+			path = Config;
+			sourceTree = "<group>";
+		};
+/* End PBXGroup section */
+
+/* Begin PBXNativeTarget section */
+		A00000000000000000000041 /* Edge0Phone */ = {
+			isa = PBXNativeTarget;
+			buildConfigurationList = A00000000000000000000051 /* Build configuration list for PBXNativeTarget "Edge0Phone" */;
+			buildPhases = (
+				A00000000000000000000061 /* Sources */,
+				A00000000000000000000021 /* Frameworks */,
+				A00000000000000000000062 /* Resources */,
+			);
+			buildRules = (
+			);
+			dependencies = (
+			);
+			name = Edge0Phone;
+			packageProductDependencies = (
+				A00000000000000000000071 /* Edge0Core */,
+				A00000000000000000000072 /* Edge0MLX */,
+			);
+			productName = Edge0Phone;
+			productReference = A00000000000000000000003 /* Edge0Phone.app */;
+			productType = "com.apple.product-type.application";
+		};
+/* End PBXNativeTarget section */
+
+/* Begin PBXProject section */
+		A00000000000000000000042 /* Project object */ = {
+			isa = PBXProject;
+			attributes = {
+				BuildIndependentTargetsInParallel = 1;
+				LastSwiftUpdateCheck = 2610;
+				LastUpgradeCheck = 2610;
+				TargetAttributes = {
+					A00000000000000000000041 = {
+						CreatedOnToolsVersion = 26.6;
+					};
+				};
+			};
+			buildConfigurationList = A00000000000000000000052 /* Build configuration list for PBXProject "Edge0PhoneProbe" */;
+			compatibilityVersion = "Xcode 16.0";
+			developmentRegion = en;
+			hasScannedForEncodings = 0;
+			knownRegions = (
+				en,
+				Base,
+			);
+			mainGroup = A00000000000000000000022;
+			packageReferences = (
+				A00000000000000000000081 /* XCLocalSwiftPackageReference "." */,
+			);
+			productRefGroup = A00000000000000000000032 /* Products */;
+			projectDirPath = "";
+			projectRoot = "";
+			targets = (
+				A00000000000000000000041 /* Edge0Phone */,
+			);
+		};
+/* End PBXProject section */
+
+/* Begin PBXResourcesBuildPhase section */
+		A00000000000000000000062 /* Resources */ = {
+			isa = PBXResourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				A00000000000000000000006 /* Models in Resources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXResourcesBuildPhase section */
+
+/* Begin PBXSourcesBuildPhase section */
+		A00000000000000000000061 /* Sources */ = {
+			isa = PBXSourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				A00000000000000000000001 /* Edge0PhoneProbeApp.swift in Sources */,
+				A00000000000000000000002 /* ContentView.swift in Sources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXSourcesBuildPhase section */
+
+/* Begin XCBuildConfiguration section */
+		A00000000000000000000091 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			baseConfigurationReference = A00000000000000000000015 /* App.xcconfig */;
+			buildSettings = {
+				CODE_SIGN_ENTITLEMENTS = Edge0PhoneProbe/Edge0PhoneProbe.entitlements;
+				CODE_SIGN_STYLE = Automatic;
+				DEVELOPMENT_TEAM = X756QY2CTK;
+				GENERATE_INFOPLIST_FILE = NO;
+				INFOPLIST_FILE = Edge0PhoneProbe/Info.plist;
+				IPHONEOS_DEPLOYMENT_TARGET = 17.0;
+				PRODUCT_BUNDLE_IDENTIFIER = com.example.edge0phone;
+				PRODUCT_NAME = "$(TARGET_NAME)";
+				SDKROOT = iphoneos;
+				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator";
+				SWIFT_VERSION = 6.0;
+				TARGETED_DEVICE_FAMILY = 1;
+			};
+			name = Debug;
+		};
+		A00000000000000000000092 /* Release */ = {
+			isa = XCBuildConfiguration;
+			baseConfigurationReference = A00000000000000000000015 /* App.xcconfig */;
+			buildSettings = {
+				CODE_SIGN_ENTITLEMENTS = Edge0PhoneProbe/Edge0PhoneProbe.entitlements;
+				CODE_SIGN_STYLE = Automatic;
+				DEVELOPMENT_TEAM = X756QY2CTK;
+				GENERATE_INFOPLIST_FILE = NO;
+				INFOPLIST_FILE = Edge0PhoneProbe/Info.plist;
+				IPHONEOS_DEPLOYMENT_TARGET = 17.0;
+				PRODUCT_BUNDLE_IDENTIFIER = com.example.edge0phone;
+				PRODUCT_NAME = "$(TARGET_NAME)";
+				SDKROOT = iphoneos;
+				SUPPORTED_PLATFORMS = "iphoneos iphonesimulator";
+				SWIFT_VERSION = 6.0;
+				TARGETED_DEVICE_FAMILY = 1;
+			};
+			name = Release;
+		};
+		A00000000000000000000093 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
+				CLANG_ENABLE_MODULES = YES;
+			};
+			name = Debug;
+		};
+		A00000000000000000000094 /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
+				CLANG_ENABLE_MODULES = YES;
+			};
+			name = Release;
+		};
+/* End XCBuildConfiguration section */
+
+/* Begin XCConfigurationList section */
+		A00000000000000000000051 /* Build configuration list for PBXNativeTarget "Edge0Phone" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				A00000000000000000000091 /* Debug */,
+				A00000000000000000000092 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+		A00000000000000000000052 /* Build configuration list for PBXProject "Edge0PhoneProbe" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				A00000000000000000000093 /* Debug */,
+				A00000000000000000000094 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+/* End XCConfigurationList section */
+
+/* Begin XCLocalSwiftPackageReference section */
+		A00000000000000000000081 /* XCLocalSwiftPackageReference "." */ = {
+			isa = XCLocalSwiftPackageReference;
+			relativePath = .;
+		};
+/* End XCLocalSwiftPackageReference section */
+
+/* Begin XCSwiftPackageProductDependency section */
+		A00000000000000000000071 /* Edge0Core */ = {
+			isa = XCSwiftPackageProductDependency;
+			package = A00000000000000000000081 /* XCLocalSwiftPackageReference "." */;
+			productName = Edge0Core;
+		};
+		A00000000000000000000072 /* Edge0MLX */ = {
+			isa = XCSwiftPackageProductDependency;
+			package = A00000000000000000000081 /* XCLocalSwiftPackageReference "." */;
+			productName = Edge0MLX;
+		};
+/* End XCSwiftPackageProductDependency section */
+	};
+	rootObject = A00000000000000000000042 /* Project object */;
+}
diff --git a/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/contents.xcworkspacedata b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/contents.xcworkspacedata
new file mode 100644
index 0000000..919434a
--- /dev/null
+++ b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/contents.xcworkspacedata
@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<Workspace
+   version = "1.0">
+   <FileRef
+      location = "self:">
+   </FileRef>
+</Workspace>
diff --git a/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
new file mode 100644
index 0000000..e1dd047
--- /dev/null
+++ b/ios/Edge0PhoneProbe.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
@@ -0,0 +1,33 @@
+{
+  "originHash" : "bb8f886c2113b757a4f1caf03c3e9422df8f501d7a701d5dfdb178ee6043c4eb",
+  "pins" : [
+    {
+      "identity" : "mlx-swift",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ml-explore/mlx-swift",
+      "state" : {
+        "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021",
+        "version" : "0.31.6"
+      }
+    },
+    {
+      "identity" : "swift-argument-parser",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-argument-parser",
+      "state" : {
+        "revision" : "6a52f3251125d74daf04fcbd5e6f08a75d074382",
+        "version" : "1.8.2"
+      }
+    },
+    {
+      "identity" : "swift-numerics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-numerics",
+      "state" : {
+        "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2",
+        "version" : "1.1.1"
+      }
+    }
+  ],
+  "version" : 3
+}
diff --git a/ios/Edge0PhoneProbe.xcodeproj/xcshareddata/xcschemes/Edge0Phone.xcscheme b/ios/Edge0PhoneProbe.xcodeproj/xcshareddata/xcschemes/Edge0Phone.xcscheme
new file mode 100644
index 0000000..0136b32
--- /dev/null
+++ b/ios/Edge0PhoneProbe.xcodeproj/xcshareddata/xcschemes/Edge0Phone.xcscheme
@@ -0,0 +1,76 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<Scheme
+   LastUpgradeVersion = "2610"
+   version = "1.7">
+   <BuildAction
+      parallelizeBuildables = "YES"
+      buildImplicitDependencies = "YES">
+      <BuildActionEntries>
+         <BuildActionEntry
+            buildForTesting = "YES"
+            buildForRunning = "YES"
+            buildForProfiling = "YES"
+            buildForArchiving = "YES"
+            buildForAnalyzing = "YES">
+            <BuildableReference
+               BuildableIdentifier = "primary"
+               BlueprintIdentifier = "A00000000000000000000041"
+               BuildableName = "Edge0Phone.app"
+               BlueprintName = "Edge0Phone"
+               ReferencedContainer = "container:Edge0PhoneProbe.xcodeproj">
+            </BuildableReference>
+         </BuildActionEntry>
+      </BuildActionEntries>
+   </BuildAction>
+   <TestAction
+      buildConfiguration = "Debug"
+      selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
+      selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
+      shouldUseLaunchSchemeArgsEnv = "YES">
+   </TestAction>
+   <LaunchAction
+      buildConfiguration = "Debug"
+      selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
+      selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
+      launchStyle = "0"
+      useCustomWorkingDirectory = "NO"
+      ignoresPersistentStateOnLaunch = "NO"
+      debugDocumentVersioning = "YES"
+      debugServiceExtension = "internal"
+      allowLocationSimulation = "YES">
+      <BuildableProductRunnable
+         runnableDebuggingMode = "0">
+         <BuildableReference
+            BuildableIdentifier = "primary"
+            BlueprintIdentifier = "A00000000000000000000041"
+            BuildableName = "Edge0Phone.app"
+            BlueprintName = "Edge0Phone"
+            ReferencedContainer = "container:Edge0PhoneProbe.xcodeproj">
+         </BuildableReference>
+      </BuildableProductRunnable>
+   </LaunchAction>
+   <ProfileAction
+      buildConfiguration = "Release"
+      shouldUseLaunchSchemeArgsEnv = "YES"
+      savedToolIdentifier = ""
+      useCustomWorkingDirectory = "NO"
+      debugDocumentVersioning = "YES">
+      <BuildableProductRunnable
+         runnableDebuggingMode = "0">
+         <BuildableReference
+            BuildableIdentifier = "primary"
+            BlueprintIdentifier = "A00000000000000000000041"
+            BuildableName = "Edge0Phone.app"
+            BlueprintName = "Edge0Phone"
+            ReferencedContainer = "container:Edge0PhoneProbe.xcodeproj">
+         </BuildableReference>
+      </BuildableProductRunnable>
+   </ProfileAction>
+   <AnalyzeAction
+      buildConfiguration = "Debug">
+   </AnalyzeAction>
+   <ArchiveAction
+      buildConfiguration = "Release"
+      revealArchiveInOrganizer = "YES">
+   </ArchiveAction>
+</Scheme>
diff --git a/ios/Edge0PhoneProbe/ContentView.swift b/ios/Edge0PhoneProbe/ContentView.swift
new file mode 100644
index 0000000..4e33ad3
--- /dev/null
+++ b/ios/Edge0PhoneProbe/ContentView.swift
@@ -0,0 +1,552 @@
+import Edge0MLX
+import Darwin
+import Foundation
+import SwiftUI
+
+private struct ChatMessage: Identifiable, Sendable {
+    enum Role: Sendable { case user, assistant }
+    let id: UUID
+    let role: Role
+    var text: String
+    var detail: String?
+
+    init(id: UUID = UUID(), role: Role, text: String, detail: String?) {
+        self.id = id; self.role = role; self.text = text; self.detail = detail
+    }
+}
+
+private enum LocalModel: String, CaseIterable, Identifiable, Sendable {
+    case edge8
+    case edge35b
+
+    var id: String { rawValue }
+    var title: String {
+        switch self {
+        case .edge8: "Edge0 8B"
+        case .edge35b: "Edge0 35B"
+        }
+    }
+    var folderName: String {
+        self == .edge8 ? "Edge0-8B-A1B-preview" : Edge0ChatEngine35B.modelFolderName
+    }
+}
+
+private actor ChatRuntime {
+    private enum LoadedEngine {
+        case edge8(Edge0ChatEngine)
+        case edge35b(Edge0ChatEngine35B)
+    }
+    private var engine: LoadedEngine?
+
+    func load(_ model: LocalModel, modelURL: URL) throws {
+        // Do not keep both models alive while switching; 35B is memory constrained.
+        engine = nil
+        switch model {
+        case .edge8:
+            engine = .edge8(try Edge0ChatEngine(modelURL: modelURL))
+        case .edge35b:
+            engine = .edge35b(try Edge0ChatEngine35B(modelURL: modelURL))
+        }
+    }
+
+    func reply(to text: String,
+               maxTokens: Int? = nil,
+               thinking: Bool,
+               onText: @escaping @Sendable (String) -> Void) async throws -> Edge0GenerationResult {
+        let shouldContinue = { @Sendable in
+            !withUnsafeCurrentTask { $0?.isCancelled ?? false }
+        }
+        switch engine {
+        case .edge8(let engine):
+            if let maxTokens {
+                return try engine.reply(
+                    to: text, maxTokens: maxTokens, thinking: thinking,
+                    onText: onText, shouldContinue: shouldContinue)
+            }
+            return try engine.reply(
+                to: text, thinking: thinking, onText: onText, shouldContinue: shouldContinue)
+        case .edge35b(let engine):
+            if let maxTokens {
+                return try await engine.reply(
+                    to: text, maxTokens: maxTokens, thinking: thinking,
+                    onText: onText, shouldContinue: shouldContinue)
+            }
+            return try await engine.reply(
+                to: text, thinking: thinking, onText: onText, shouldContinue: shouldContinue)
+        case nil:
+            throw RuntimeFailure.modelNotLoaded
+        }
+    }
+
+    func reset() {
+        switch engine {
+        case .edge8(let engine): engine.reset()
+        case .edge35b(let engine): engine.reset()
+        case nil: break
+        }
+    }
+
+    func unload() { engine = nil }
+
+    private enum RuntimeFailure: LocalizedError {
+        case modelNotLoaded
+        var errorDescription: String? { "Model is not loaded" }
+    }
+}
+
+@MainActor
+private final class ChatViewModel: ObservableObject {
+    enum Phase: Equatable { case choosing, ready, preparing, generating, failed(String) }
+
+    @Published var input = ""
+    @Published private(set) var messages: [ChatMessage] = []
+    @Published private(set) var phase: Phase = .choosing
+    @Published private(set) var selectedModel: LocalModel?
+
+    private let runtime = ChatRuntime()
+    private var generationTask: Task<Void, Never>?
+
+    var isRunning: Bool { phase == .preparing || phase == .generating }
+    /// Off unless the process is launched with `--thinking`. `--no-thinking` forces it off.
+    private let thinkingEnabled: Bool = {
+        let args = CommandLine.arguments
+        if args.contains("--no-thinking") { return false }
+        return args.contains("--thinking")
+    }()
+    var statusText: String? {
+        switch phase {
+        case .choosing: nil
+        case .ready: nil
+        case .preparing: "Loading \(selectedModel?.title ?? "local model")…"
+        case .generating: "\(selectedModel?.title ?? "Model") is generating…"
+        case .failed(let message): message
+        }
+    }
+
+    func isInstalled(_ model: LocalModel) -> Bool {
+        let url = resolveModelURL(model)
+        switch model {
+        case .edge8:
+            return FileManager.default.fileExists(
+                atPath: url.appendingPathComponent("model.safetensors").path)
+        case .edge35b:
+            return FileManager.default.fileExists(
+                atPath: url.appendingPathComponent("model.safetensors.index.json").path)
+                && FileManager.default.fileExists(
+                    atPath: url.appendingPathComponent("tokenizer.bin").path)
+                && FileManager.default.fileExists(
+                    atPath: url.appendingPathComponent("experts-L39.bin").path)
+        }
+    }
+
+    func choose(_ model: LocalModel) {
+        guard !isRunning else { return }
+        selectedModel = model
+        messages.removeAll()
+        input = ""
+        phase = .preparing
+        let url = resolveModelURL(model)
+        Task {
+            do {
+                try await runtime.load(model, modelURL: url)
+                if phase == .preparing { phase = .ready }
+            } catch {
+                phase = .failed(Self.friendlyMessage(for: error))
+            }
+        }
+    }
+
+    func chooseAnotherModel() {
+        guard !isRunning else { return }
+        messages.removeAll()
+        input = ""
+        selectedModel = nil
+        phase = .choosing
+        Task { await runtime.unload() }
+    }
+
+    func useSuggestion(_ text: String, maxTokens: Int? = nil) {
+        input = text
+        send(maxTokens: maxTokens)
+    }
+
+    func send(maxTokens: Int? = nil) {
+        let text = input.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard !text.isEmpty, !isRunning else { return }
+        input = ""
+        messages.append(ChatMessage(role: .user, text: text, detail: nil))
+        let assistantID = UUID()
+        messages.append(ChatMessage(id: assistantID, role: .assistant,
+                                    text: "", detail: nil))
+        phase = .generating
+        generationTask = Task {
+            do {
+                let result = try await runtime.reply(
+                    to: text,
+                    maxTokens: maxTokens,
+                    thinking: thinkingEnabled,
+                    onText: { [weak self] partial in
+                        Task { @MainActor in
+                            guard let self,
+                                  let index = self.messages.firstIndex(where: { $0.id == assistantID })
+                            else { return }
+                            self.messages[index].text = partial
+                        }
+                    })
+                let detail = String(
+                    format: "%d tokens · TTFT %.1fs · prefill %.1f tok/s · decode %.2f tok/s · peak %.0f MB",
+                    result.generatedTokenCount, result.timeToFirstTokenSeconds,
+                    result.prefillTokensPerSecond,
+                    result.decodeTokensPerSecond,
+                    Double(result.peakMemoryBytes) / 1_048_576
+                )
+                if let index = messages.firstIndex(where: { $0.id == assistantID }) {
+                    messages[index].text = result.text.isEmpty ? "(The model returned an empty response)" : result.text
+                    messages[index].detail = detail
+                }
+                print("[generation] \(detail)")
+                if CommandLine.arguments.contains("--bench-second"),
+                   CommandLine.arguments.contains("--bench-35b") {
+                    print("[bench] second turn")
+                    let second = try await runtime.reply(
+                        to: "Summarize it in one sentence.",
+                        maxTokens: 64,
+                        thinking: thinkingEnabled,
+                        onText: { _ in })
+                    let secondDetail = String(
+                        format: "%d tokens · TTFT %.1fs · prefill %.1f tok/s · decode %.2f tok/s · peak %.0f MB",
+                        second.generatedTokenCount, second.timeToFirstTokenSeconds,
+                        second.prefillTokensPerSecond,
+                        second.decodeTokensPerSecond,
+                        Double(second.peakMemoryBytes) / 1_048_576
+                    )
+                    print("[generation turn2] \(secondDetail)")
+                }
+                if CommandLine.arguments.contains("--bench-8b")
+                    || CommandLine.arguments.contains("--bench-35b") {
+                    var info = task_vm_info_data_t()
+                    var count = mach_msg_type_number_t(
+                        MemoryLayout<task_vm_info_data_t>.stride / MemoryLayout<natural_t>.stride)
+                    let status = withUnsafeMutablePointer(to: &info) {
+                        $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
+                            task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count)
+                        }
+                    }
+                    print("[bench memory] endFootprintMiB=\(status == KERN_SUCCESS ? Double(info.phys_footprint) / 1_048_576 : -1) thermal=\(ProcessInfo.processInfo.thermalState.rawValue)")
+                    print("[bench timing] ttftSeconds=\(result.timeToFirstTokenSeconds) prefillSeconds=\(result.prefillSeconds) prefillTokPerSec=\(result.prefillTokensPerSecond) decodeTokPerSec=\(result.decodeTokensPerSecond)")
+                }
+                phase = .ready
+                if CommandLine.arguments.contains("--bench-8b")
+                    || CommandLine.arguments.contains("--bench-35b")
+                    || CommandLine.arguments.contains("--smoke-35b") {
+                    fflush(nil)
+                    exit(0)
+                }
+            } catch is CancellationError {
+                phase = .ready
+            } catch {
+                phase = .failed(Self.friendlyMessage(for: error))
+                if CommandLine.arguments.contains("--bench-8b")
+                    || CommandLine.arguments.contains("--bench-35b")
+                    || CommandLine.arguments.contains("--smoke-35b") {
+                    fflush(stdout)
+                    exit(1)
+                }
+            }
+            generationTask = nil
+        }
+    }
+
+    func stop() { generationTask?.cancel() }
+
+    func newConversation() {
+        guard !isRunning, selectedModel != nil else { return }
+        messages.removeAll()
+        input = ""
+        phase = .ready
+        Task { await runtime.reset() }
+    }
+
+    private func resolveModelURL(_ model: LocalModel) -> URL {
+        let models = Bundle.main.resourceURL?
+            .appendingPathComponent("Models", isDirectory: true)
+            ?? Bundle.main.bundleURL.appendingPathComponent("Models", isDirectory: true)
+        return models.appendingPathComponent(model.folderName, isDirectory: true)
+    }
+
+    private static func friendlyMessage(for error: Error) -> String {
+        let description = error.localizedDescription
+        if description.contains("doesn’t exist") || description.contains("No such file")
+            || description.contains("incomplete") {
+            return "Model not included in this build"
+        }
+        return "Run failed: \(description)"
+    }
+}
+
+struct ContentView: View {
+    @StateObject private var chat = ChatViewModel()
+    @FocusState private var inputFocused: Bool
+
+    private let suggestions = [
+        "What is artificial intelligence?",
+        "How do LLMs understand context?",
+        "Tell me about renewable energy",
+        "Write a short poem about autumn",
+    ]
+
+    var body: some View {
+        ZStack {
+            Color.black.ignoresSafeArea()
+            if chat.selectedModel == nil {
+                modelPicker
+            } else {
+                VStack(spacing: 0) {
+                    header
+                    conversation
+                }
+            }
+        }
+        .preferredColorScheme(.dark)
+        .safeAreaInset(edge: .bottom, spacing: 0) {
+            if chat.selectedModel != nil { composer }
+        }
+        .task {
+            // Keeps the normal launch picker intact while allowing a repeatable
+            // on-device smoke/benchmark run from `devicectl`.
+            let smoke = CommandLine.arguments.contains("--smoke-35b")
+            let bench8 = CommandLine.arguments.contains("--bench-8b")
+            let bench35 = CommandLine.arguments.contains("--bench-35b")
+            let benchmark = bench8 || bench35
+            let target: LocalModel =
+                bench8 || CommandLine.arguments.contains("--model=8b") ? .edge8 : .edge35b
+            if chat.selectedModel == nil,
+               (CommandLine.arguments.contains("--model=35b")
+                || CommandLine.arguments.contains("--model=8b")
+                || smoke || benchmark),
+               chat.isInstalled(target) {
+                print("[bench] loading \(target.title)")
+                print("[bench] thermalAtLaunch=\(ProcessInfo.processInfo.thermalState.rawValue)")
+                chat.choose(target)
+                if smoke || benchmark {
+                    while chat.phase == .preparing {
+                        try? await Task.sleep(for: .milliseconds(100))
+                    }
+                    if case .failed(let message) = chat.phase {
+                        print("[bench] load failed: \(message)")
+                    } else if chat.phase == .ready {
+                        let long = CommandLine.arguments.contains("--bench-long")
+                        let sentence = "Streaming inference overlaps disk bandwidth with compute. Each mixture-of-experts layer activates only a few experts."
+                        let prompt = smoke
+                            ? "Reply with this only: 35B is running"
+                            : (long
+                                ? Array(repeating: sentence, count: 80).joined()
+                                : "What is artificial intelligence?")
+                        chat.useSuggestion(prompt, maxTokens: long ? 8 : (benchmark ? 64 : nil))
+                    }
+                }
+            }
+        }
+    }
+
+    private var modelPicker: some View {
+        ScrollView {
+            VStack(alignment: .leading, spacing: 22) {
+                Spacer(minLength: 54)
+                Text("Choose a local model")
+                    .font(.system(size: 34, weight: .bold, design: .rounded))
+
+                ForEach(LocalModel.allCases) { model in
+                    let installed = chat.isInstalled(model)
+                    Button { chat.choose(model) } label: {
+                        HStack(spacing: 16) {
+                            Image(systemName: model == .edge8 ? "bolt.fill" : "brain.head.profile")
+                                .font(.system(size: 23, weight: .semibold))
+                                .frame(width: 48, height: 48)
+                                .background(.white.opacity(0.09), in: RoundedRectangle(cornerRadius: 14))
+                            VStack(alignment: .leading, spacing: 5) {
+                                Text(model.title).font(.headline)
+                                if !installed {
+                                    Text("Not included")
+                                        .font(.caption2)
+                                        .foregroundStyle(.orange)
+                                }
+                            }
+                            Spacer()
+                            Image(systemName: "chevron.right")
+                                .foregroundStyle(.secondary)
+                        }
+                        .padding(18)
+                        .background(.white.opacity(0.07), in: RoundedRectangle(cornerRadius: 22))
+                    }
+                    .buttonStyle(.plain)
+                    .disabled(!installed)
+                    .opacity(installed ? 1 : 0.62)
+                }
+
+                Spacer(minLength: 40)
+            }
+            .padding(.horizontal, 22)
+            .frame(maxWidth: 620)
+            .frame(maxWidth: .infinity)
+        }
+    }
+
+    private var header: some View {
+        HStack(spacing: 10) {
+            Text(chat.selectedModel?.title ?? "Edge0")
+                .font(.headline.weight(.semibold))
+            Spacer()
+            Button(action: chat.chooseAnotherModel) {
+                Image(systemName: "arrow.left.arrow.right")
+                    .font(.system(size: 15, weight: .semibold))
+                    .frame(width: 38, height: 38)
+                    .background(.white.opacity(0.07), in: Circle())
+            }
+            .foregroundStyle(.white)
+            .disabled(chat.isRunning)
+            .accessibilityLabel("Switch model")
+            Button(action: chat.newConversation) {
+                Image(systemName: "square.and.pencil")
+                    .font(.system(size: 17, weight: .semibold))
+                    .frame(width: 38, height: 38)
+                    .background(.white.opacity(0.07), in: Circle())
+            }
+            .foregroundStyle(.white)
+            .disabled(chat.isRunning)
+            .accessibilityLabel("New chat")
+        }
+        .padding(.horizontal, 20)
+        .padding(.top, 8)
+        .padding(.bottom, 10)
+    }
+
+    @ViewBuilder
+    private var conversation: some View {
+        if chat.messages.isEmpty {
+            welcome
+        } else {
+            ScrollViewReader { proxy in
+                ScrollView {
+                    LazyVStack(spacing: 18) {
+                        ForEach(chat.messages) { message in
+                            messageBubble(message).id(message.id)
+                        }
+                        if let status = chat.statusText { statusRow(status).id("status") }
+                    }
+                    .padding(.horizontal, 18)
+                    .padding(.top, 16)
+                    .padding(.bottom, 24)
+                }
+                .scrollDismissesKeyboard(.interactively)
+                .onChange(of: chat.messages.count) {
+                    withAnimation { proxy.scrollTo(chat.messages.last?.id, anchor: .bottom) }
+                }
+                .onChange(of: chat.phase) {
+                    withAnimation { proxy.scrollTo("status", anchor: .bottom) }
+                }
+            }
+        }
+    }
+
+    private var welcome: some View {
+        ScrollView {
+            VStack(spacing: 18) {
+                Spacer(minLength: 90)
+                Text("edge0")
+                    .font(.system(size: 58, weight: .medium, design: .rounded))
+                    .tracking(-3)
+                Text("Private AI, running locally on your iPhone")
+                    .font(.subheadline)
+                    .foregroundStyle(.secondary)
+                    .multilineTextAlignment(.center)
+                VStack(spacing: 10) {
+                    ForEach(suggestions, id: \.self) { suggestion in
+                        Button(suggestion) { chat.useSuggestion(suggestion) }
+                            .font(.subheadline.weight(.medium))
+                            .foregroundStyle(.white.opacity(0.88))
+                            .padding(.horizontal, 16)
+                            .padding(.vertical, 11)
+                            .background(.white.opacity(0.07), in: Capsule())
+                    }
+                }
+                .padding(.top, 8)
+                if let status = chat.statusText { statusRow(status) }
+                Spacer(minLength: 120)
+            }
+            .frame(maxWidth: .infinity)
+            .padding(.horizontal, 24)
+        }
+    }
+
+    private func messageBubble(_ message: ChatMessage) -> some View {
+        VStack(alignment: message.role == .user ? .trailing : .leading, spacing: 6) {
+            Text(message.text)
+                .font(.body)
+                .foregroundStyle(.white.opacity(0.94))
+                .textSelection(.enabled)
+                .padding(.horizontal, 15)
+                .padding(.vertical, 12)
+                .background(
+                    message.role == .user ? Color.white.opacity(0.14) : Color.white.opacity(0.07),
+                    in: RoundedRectangle(cornerRadius: 19, style: .continuous)
+                )
+            if let detail = message.detail {
+                Text(detail)
+                    .font(.caption2.monospacedDigit())
+                    .foregroundStyle(.secondary)
+                    .padding(.horizontal, 6)
+            }
+        }
+        .frame(maxWidth: .infinity, alignment: message.role == .user ? .trailing : .leading)
+    }
+
+    private func statusRow(_ text: String) -> some View {
+        HStack(spacing: 10) {
+            if chat.isRunning { ProgressView().controlSize(.small) }
+            Text(text)
+                .font(.footnote)
+                .foregroundStyle(chat.isRunning ? Color.secondary : Color.red.opacity(0.9))
+        }
+        .frame(maxWidth: .infinity, alignment: .leading)
+        .padding(.horizontal, 6)
+    }
+
+    private var composer: some View {
+        HStack(alignment: .bottom, spacing: 10) {
+            TextField("Ask anything…", text: $chat.input, axis: .vertical)
+                .lineLimit(1...5)
+                .focused($inputFocused)
+                .submitLabel(.send)
+                .onSubmit { chat.send() }
+                .padding(.horizontal, 16)
+                .padding(.vertical, 13)
+                .background(.white.opacity(0.08), in: RoundedRectangle(cornerRadius: 22, style: .continuous))
+
+            Button {
+                if chat.isRunning, chat.phase == .generating {
+                    chat.stop()
+                } else {
+                    chat.send()
+                }
+                inputFocused = false
+            } label: {
+                Image(systemName: chat.phase == .generating ? "stop.fill" : "arrow.up")
+                    .font(.system(size: 17, weight: .bold))
+                    .foregroundStyle(.black)
+                    .frame(width: 45, height: 45)
+                    .background(.white, in: Circle())
+            }
+            .disabled(chat.phase == .preparing ||
+                      chat.phase == .choosing ||
+                      (chat.phase != .generating &&
+                       chat.input.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty))
+            .opacity(chat.phase == .preparing ? 0.5 : 1)
+            .accessibilityLabel(chat.phase == .generating ? "Stop" : "Send")
+        }
+        .padding(.horizontal, 16)
+        .padding(.top, 10)
+        .padding(.bottom, 8)
+        .background(.ultraThinMaterial)
+    }
+}
diff --git a/ios/Edge0PhoneProbe/Edge0PhoneProbe.entitlements b/ios/Edge0PhoneProbe/Edge0PhoneProbe.entitlements
new file mode 100644
index 0000000..6631ffa
--- /dev/null
+++ b/ios/Edge0PhoneProbe/Edge0PhoneProbe.entitlements
@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+</dict>
+</plist>
diff --git a/ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift b/ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift
new file mode 100644
index 0000000..0423c3d
--- /dev/null
+++ b/ios/Edge0PhoneProbe/Edge0PhoneProbeApp.swift
@@ -0,0 +1,10 @@
+import SwiftUI
+
+@main
+struct Edge0PhoneProbeApp: App {
+    var body: some Scene {
+        WindowGroup {
+            ContentView()
+        }
+    }
+}
diff --git a/ios/Edge0PhoneProbe/Info.plist b/ios/Edge0PhoneProbe/Info.plist
new file mode 100644
index 0000000..91ac1d1
--- /dev/null
+++ b/ios/Edge0PhoneProbe/Info.plist
@@ -0,0 +1,35 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>CFBundleDevelopmentRegion</key>
+	<string>$(DEVELOPMENT_LANGUAGE)</string>
+	<key>CFBundleExecutable</key>
+	<string>$(EXECUTABLE_NAME)</string>
+	<key>CFBundleIdentifier</key>
+	<string>$(PRODUCT_BUNDLE_IDENTIFIER)</string>
+	<key>CFBundleInfoDictionaryVersion</key>
+	<string>6.0</string>
+	<key>CFBundleName</key>
+	<string>$(PRODUCT_NAME)</string>
+	<key>CFBundlePackageType</key>
+	<string>$(PRODUCT_BUNDLE_PACKAGE_TYPE)</string>
+	<key>CFBundleShortVersionString</key>
+	<string>1.0</string>
+	<key>CFBundleVersion</key>
+	<string>1</string>
+	<key>LSRequiresIPhoneOS</key>
+	<true/>
+	<key>LSSupportsOpeningDocumentsInPlace</key>
+	<true/>
+	<key>UIApplicationSceneManifest</key>
+	<dict>
+		<key>UIApplicationSupportsMultipleScenes</key>
+		<false/>
+	</dict>
+	<key>UIFileSharingEnabled</key>
+	<true/>
+	<key>UILaunchScreen</key>
+	<dict/>
+</dict>
+</plist>
diff --git a/ios/Models/.gitkeep b/ios/Models/.gitkeep
new file mode 100644
index 0000000..e69de29
diff --git a/ios/Package.resolved b/ios/Package.resolved
new file mode 100644
index 0000000..66a2fa4
--- /dev/null
+++ b/ios/Package.resolved
@@ -0,0 +1,33 @@
+{
+  "originHash" : "14c443ca8e8c4734ac842efe15d1c40af0856e85f2a4579f1ff7b75f93b2ccc7",
+  "pins" : [
+    {
+      "identity" : "mlx-swift",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ml-explore/mlx-swift",
+      "state" : {
+        "revision" : "0bb916c67f4b9e5c682cbe02a42c701c93ab5021",
+        "version" : "0.31.6"
+      }
+    },
+    {
+      "identity" : "swift-argument-parser",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-argument-parser",
+      "state" : {
+        "revision" : "6a52f3251125d74daf04fcbd5e6f08a75d074382",
+        "version" : "1.8.2"
+      }
+    },
+    {
+      "identity" : "swift-numerics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-numerics",
+      "state" : {
+        "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2",
+        "version" : "1.1.1"
+      }
+    }
+  ],
+  "version" : 3
+}
diff --git a/ios/Package.swift b/ios/Package.swift
new file mode 100644
index 0000000..4672abb
--- /dev/null
+++ b/ios/Package.swift
@@ -0,0 +1,33 @@
+// swift-tools-version: 6.1
+import PackageDescription
+
+let package = Package(
+    name: "Edge0Phone",
+    platforms: [
+        .macOS(.v14),
+        .iOS(.v17),
+    ],
+    products: [
+        .library(name: "Edge0Core", targets: ["Edge0Core"]),
+        .library(name: "Edge0MLX", targets: ["Edge0MLX"]),
+    ],
+    dependencies: [
+        .package(
+            url: "https://github.com/ml-explore/mlx-swift",
+            .upToNextMinor(from: "0.31.4")
+        ),
+    ],
+    targets: [
+        .target(name: "Edge0Core"),
+        .target(
+            name: "Edge0MLX",
+            dependencies: [
+                "Edge0Core",
+                .product(name: "MLX", package: "mlx-swift"),
+                .product(name: "MLXFast", package: "mlx-swift"),
+            ]
+        ),
+        .testTarget(name: "Edge0MLXTests", dependencies: ["Edge0MLX", "Edge0Core"], resources: [.copy("Fixtures")]),
+        .testTarget(name: "Edge0CoreTests", dependencies: ["Edge0Core"]),
+    ]
+)
diff --git a/ios/README.md b/ios/README.md
new file mode 100644
index 0000000..1d3d24e
--- /dev/null
+++ b/ios/README.md
@@ -0,0 +1,111 @@
+# Edge0Phone
+
+English | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+On-device chat for iPhone. The app can run Edge0 8B, Edge0 35B, or both, and only the model you select is loaded. Model weights are not in this folder.
+
+## Quick Start
+
+You need an Apple Silicon Mac, Xcode, and an iPhone on iOS 17 or later with Developer Mode on. Python 3.10 or later is only needed to download weights and to convert 35B. If Xcode reports a missing Metal compiler, run `xcodebuild -downloadComponent MetalToolchain`.
+
+The `Models/` directory is copied into the app at build time. Put a model there before you build if you want that button to be selectable. An empty `Models/` directory still builds; the missing model shows **Not included**.
+
+### Download a model
+
+```sh
+python3 -m venv .venv-tools
+.venv-tools/bin/python -m pip install --upgrade pip huggingface_hub numpy
+```
+
+8B can be packed as downloaded. It is about 5 GB, so the installed app is about that size.
+
+```sh
+.venv-tools/bin/hf download Edge0/Edge0-8B-A1B-preview \
+  --local-dir Models/Edge0-8B-A1B-preview
+```
+
+35B is a layout change, not a change to the weights. The published checkpoint stores each expert as nine scattered reads across a 19 GB file. On an iPhone that is too slow, so the tools below rewrite each expert into one sequential read and write one file per layer. Keep the original download outside `Models/`. Only `Models/repacked/` is what the app should contain. A 35B build is about 20 GB.
+
+```sh
+mkdir -p checkpoints
+.venv-tools/bin/hf download Edge0/Edge0-35B-A3B-preview \
+  --local-dir checkpoints/Edge0-35B-A3B-preview
+
+SOURCE=checkpoints/Edge0-35B-A3B-preview
+OUTPUT=Models/repacked
+
+.venv-tools/bin/python tools/repack_experts.py pack "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py resident "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py verify "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/convert_tokenizer.py "$SOURCE" "$OUTPUT/tokenizer.bin"
+.venv-tools/bin/python tools/convert_pregate.py \
+  "$SOURCE/prerouter_edge0_35b.safetensors" \
+  "$OUTPUT/prerouter-stacked.safetensors"
+cp "$SOURCE/lora_edge0_35b.safetensors" "$OUTPUT/"
+```
+
+`Models/repacked/` must contain `resident.safetensors`, `tokenizer.bin`, and `experts-L00.bin` through `experts-L39.bin`. You can pack either model, or both.
+
+Checkpoints:
+
+- [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+- [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+
+### Build and install
+
+`com.example.edge0phone` in the project is a placeholder bundle identifier. Before building or signing, replace it with a unique bundle identifier that you control. In Xcode, open `Edge0PhoneProbe.xcodeproj`, select the **Edge0Phone** target, then change **Signing & Capabilities → Bundle Identifier**. Choose your Apple Developer Team, connect the iPhone, and Run.
+
+To set the team from a file instead, copy `Config/Local.xcconfig.example` to `Config/Local.xcconfig` and fill in `DEVELOPMENT_TEAM`.
+
+The same build from the command line:
+
+```sh
+DEVICE_ID=<device id from `xcrun devicectl list devices`>
+TEAM_ID=<your Apple Developer Team ID>
+BUNDLE_ID=com.yourname.edge0phone # replace with your own unique bundle identifier
+
+xcodebuild -project Edge0PhoneProbe.xcodeproj \
+  -scheme Edge0Phone \
+  -destination "id=$DEVICE_ID" \
+  -derivedDataPath .build-device \
+  -allowProvisioningUpdates \
+  -skipPackagePluginValidation \
+  CODE_SIGN_STYLE=Automatic \
+  DEVELOPMENT_TEAM="$TEAM_ID" \
+  PRODUCT_BUNDLE_IDENTIFIER="$BUNDLE_ID" \
+  build
+
+xcrun devicectl device install app \
+  --device "$DEVICE_ID" \
+  .build-device/Build/Products/Debug-iphoneos/Edge0Phone.app
+```
+
+The weights are inside that `.app`. There is no separate copy onto the phone.
+
+### Test
+
+`scripts/test.sh` builds the MLX Metal library with Xcode, then runs the unit tests. The first run resolves the MLX Swift package.
+
+```sh
+scripts/test.sh
+EDGE0_TEST_GPU=1 scripts/test.sh
+```
+
+### Use the app
+
+Launch the app and choose **Edge0 8B** or **Edge0 35B**. A model that was not in `Models/` when you built shows **Not included** and cannot be selected. Type a question, or tap one of the suggestions. The line under a reply shows the token count, time to first token, prefill, decode, and peak memory. Switching models unloads the one that is loaded. New chat clears the conversation.
+
+## Performance
+
+| Model | Device | OS | Prefill | TTFT | Decode |
+| --- | --- | --- | --- | --- | --- |
+| 8B | iPhone 16 Pro | iOS 26.6.2 | 7.2 tok/s | 3.6 s | 10.9 tok/s |
+| 35B | iPhone 16 Pro | iOS 26.6.2 | 4.9 tok/s | 2.1 s | 6.4 tok/s |
+
+## Technical Details
+
+The app is SwiftUI. `Edge0MLX` runs the model. `Edge0Core` reads the checkpoint, tokenizer, and routing index. Math goes through MLX Swift on Metal. The deployment target is iOS 17.
+
+Edge0 8B is 24 layers, alternating Kimi Delta Attention and multi-head latent attention, Top-8, INT4. Prompt prefill runs a whole layer at a time, in chunks of up to 2048 tokens. The first token is greedy. Later tokens use temperature 0.7, top-k 64, top-p 0.95, and a repetition penalty of 1.1.
+
+Edge0 35B is 40 layers, 256 experts per layer, Top-4. Resident weights and the per-layer expert files are read separately. The prefix cache is stored in Application Support, because the app bundle is read-only. The two models are never loaded at the same time.
diff --git a/ios/README_es.md b/ios/README_es.md
new file mode 100644
index 0000000..3eb3ed8
--- /dev/null
+++ b/ios/README_es.md
@@ -0,0 +1,111 @@
+# Edge0Phone
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | Español | [Français](README_fr.md)
+
+Chat en el dispositivo para iPhone. La aplicación puede ejecutar Edge0 8B, Edge0 35B o ambos, y solo se carga el modelo que selecciones. Los pesos del modelo no están en esta carpeta.
+
+## Inicio rápido
+
+Necesitas un Mac con Apple Silicon, Xcode y un iPhone con iOS 17 o posterior y el modo de desarrollador activado. Python 3.10 o posterior solo se necesita para descargar los pesos y para convertir el 35B. Si Xcode indica que falta el compilador de Metal, ejecuta `xcodebuild -downloadComponent MetalToolchain`.
+
+El directorio `Models/` se copia dentro de la aplicación al compilar. Coloca ahí un modelo antes de compilar si quieres que ese botón sea seleccionable. Un directorio `Models/` vacío compila igualmente; el modelo ausente muestra **No incluido**.
+
+### Descargar un modelo
+
+```sh
+python3 -m venv .venv-tools
+.venv-tools/bin/python -m pip install --upgrade pip huggingface_hub numpy
+```
+
+El 8B se puede empaquetar tal como se descarga. Ocupa unos 5 GB, así que la aplicación instalada tiene aproximadamente ese tamaño.
+
+```sh
+.venv-tools/bin/hf download Edge0/Edge0-8B-A1B-preview \
+  --local-dir Models/Edge0-8B-A1B-preview
+```
+
+El 35B es un cambio de disposición (layout), no un cambio en los pesos. El checkpoint publicado almacena cada experto como nueve lecturas dispersas a lo largo de un archivo de 19 GB. En un iPhone eso es demasiado lento, así que las herramientas de abajo reescriben cada experto en una sola lectura secuencial y escriben un archivo por capa. Conserva la descarga original fuera de `Models/`. Solo `Models/repacked/` es lo que debe contener la aplicación. Una compilación del 35B ocupa unos 20 GB.
+
+```sh
+mkdir -p checkpoints
+.venv-tools/bin/hf download Edge0/Edge0-35B-A3B-preview \
+  --local-dir checkpoints/Edge0-35B-A3B-preview
+
+SOURCE=checkpoints/Edge0-35B-A3B-preview
+OUTPUT=Models/repacked
+
+.venv-tools/bin/python tools/repack_experts.py pack "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py resident "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py verify "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/convert_tokenizer.py "$SOURCE" "$OUTPUT/tokenizer.bin"
+.venv-tools/bin/python tools/convert_pregate.py \
+  "$SOURCE/prerouter_edge0_35b.safetensors" \
+  "$OUTPUT/prerouter-stacked.safetensors"
+cp "$SOURCE/lora_edge0_35b.safetensors" "$OUTPUT/"
+```
+
+`Models/repacked/` debe contener `resident.safetensors`, `tokenizer.bin` y desde `experts-L00.bin` hasta `experts-L39.bin`. Puedes empaquetar cualquiera de los dos modelos, o ambos.
+
+Checkpoints:
+
+- [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+- [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+
+### Compilación e instalación
+
+`com.example.edge0phone` en el proyecto es un identificador de bundle de ejemplo. Antes de compilar o firmar, reemplázalo por un identificador de bundle único que controles. En Xcode, abre `Edge0PhoneProbe.xcodeproj`, selecciona el target **Edge0Phone** y cambia **Signing & Capabilities → Bundle Identifier**. Elige tu Apple Developer Team, conecta el iPhone y pulsa Run.
+
+Para configurar el equipo desde un archivo, copia `Config/Local.xcconfig.example` a `Config/Local.xcconfig` y completa `DEVELOPMENT_TEAM`.
+
+La misma compilación desde la línea de comandos:
+
+```sh
+DEVICE_ID=<device id from `xcrun devicectl list devices`>
+TEAM_ID=<your Apple Developer Team ID>
+BUNDLE_ID=com.yourname.edge0phone # replace with your own unique bundle identifier
+
+xcodebuild -project Edge0PhoneProbe.xcodeproj \
+  -scheme Edge0Phone \
+  -destination "id=$DEVICE_ID" \
+  -derivedDataPath .build-device \
+  -allowProvisioningUpdates \
+  -skipPackagePluginValidation \
+  CODE_SIGN_STYLE=Automatic \
+  DEVELOPMENT_TEAM="$TEAM_ID" \
+  PRODUCT_BUNDLE_IDENTIFIER="$BUNDLE_ID" \
+  build
+
+xcrun devicectl device install app \
+  --device "$DEVICE_ID" \
+  .build-device/Build/Products/Debug-iphoneos/Edge0Phone.app
+```
+
+Los pesos van dentro de ese `.app`. No hay una copia independiente al teléfono.
+
+### Pruebas
+
+`scripts/test.sh` compila la biblioteca MLX Metal con Xcode y después ejecuta las pruebas unitarias. La primera ejecución resuelve el paquete MLX Swift.
+
+```sh
+scripts/test.sh
+EDGE0_TEST_GPU=1 scripts/test.sh
+```
+
+### Uso de la aplicación
+
+Abre la aplicación y elige **Edge0 8B** o **Edge0 35B**. Un modelo que no estaba en `Models/` al compilar muestra **No incluido** y no se puede seleccionar. Escribe una pregunta o toca una de las sugerencias. La línea bajo una respuesta muestra el recuento de tokens, el tiempo hasta el primer token, prefill, decode y memoria pico. Al cambiar de modelo se libera el que estaba cargado. Nuevo chat borra la conversación.
+
+## Rendimiento
+
+| Modelo | Dispositivo | SO | Prefill | TTFT | Decode |
+| --- | --- | --- | --- | --- | --- |
+| 8B | iPhone 16 Pro | iOS 26.6.2 | 7.2 tok/s | 3.6 s | 10.9 tok/s |
+| 35B | iPhone 16 Pro | iOS 26.6.2 | 4.9 tok/s | 2.1 s | 6.4 tok/s |
+
+## Detalles técnicos
+
+La aplicación es SwiftUI. `Edge0MLX` ejecuta el modelo. `Edge0Core` lee el checkpoint, el tokenizer y el índice de enrutamiento. Las fórmulas matemáticas pasan por MLX Swift sobre Metal. El objetivo de despliegue es iOS 17.
+
+Edge0 8B tiene 24 capas, alternando Kimi Delta Attention y atención latente multicabeza, Top-8, INT4. El prefill del prompt ejecuta una capa completa cada vez, en fragmentos de hasta 2048 tokens. El primer token es greedy. Los tokens siguientes usan temperatura 0.7, top-k 64, top-p 0.95 y una penalización de repetición de 1.1.
+
+Edge0 35B tiene 40 capas, 256 expertos por capa, Top-4. Los pesos residentes y los archivos de expertos por capa se leen por separado. La caché de prefijos se guarda en Application Support, porque el bundle de la aplicación es de solo lectura. Los dos modelos nunca se cargan al mismo tiempo.
diff --git a/ios/README_fr.md b/ios/README_fr.md
new file mode 100644
index 0000000..e93c38c
--- /dev/null
+++ b/ios/README_fr.md
@@ -0,0 +1,111 @@
+# Edge0Phone
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | Français
+
+Chat embarqué pour iPhone. L'application peut exécuter Edge0 8B, Edge0 35B, ou les deux, et seul le modèle que vous sélectionnez est chargé. Les poids des modèles ne sont pas dans ce dossier.
+
+## Démarrage rapide
+
+Vous avez besoin d'un Mac Apple Silicon, de Xcode et d'un iPhone sous iOS 17 ou ultérieur avec le mode développeur activé. Python 3.10 ou ultérieur n'est nécessaire que pour télécharger les poids et pour convertir le 35B. Si Xcode signale un compilateur Metal manquant, exécutez `xcodebuild -downloadComponent MetalToolchain`.
+
+Le répertoire `Models/` est copié dans l'application au moment de la compilation. Placez-y un modèle avant de compiler si vous voulez que ce bouton soit sélectionnable. Un répertoire `Models/` vide compile quand même ; le modèle manquant affiche **Not included**.
+
+### Télécharger un modèle
+
+```sh
+python3 -m venv .venv-tools
+.venv-tools/bin/python -m pip install --upgrade pip huggingface_hub numpy
+```
+
+Le 8B peut être empaqueté tel que téléchargé. Il pèse environ 5 GB, l'application installée a donc à peu près cette taille.
+
+```sh
+.venv-tools/bin/hf download Edge0/Edge0-8B-A1B-preview \
+  --local-dir Models/Edge0-8B-A1B-preview
+```
+
+Le 35B est un changement de disposition, pas une modification des poids. Le checkpoint publié stocke chaque expert sous forme de neuf lectures dispersées dans un fichier de 19 GB. Sur un iPhone, c'est trop lent, donc les outils ci-dessous réécrivent chaque expert en une seule lecture séquentielle et écrivent un fichier par couche. Conservez le téléchargement d'origine hors de `Models/`. Seul `Models/repacked/` doit contenir ce que l'application doit embarquer. Une compilation 35B pèse environ 20 GB.
+
+```sh
+mkdir -p checkpoints
+.venv-tools/bin/hf download Edge0/Edge0-35B-A3B-preview \
+  --local-dir checkpoints/Edge0-35B-A3B-preview
+
+SOURCE=checkpoints/Edge0-35B-A3B-preview
+OUTPUT=Models/repacked
+
+.venv-tools/bin/python tools/repack_experts.py pack "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py resident "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py verify "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/convert_tokenizer.py "$SOURCE" "$OUTPUT/tokenizer.bin"
+.venv-tools/bin/python tools/convert_pregate.py \
+  "$SOURCE/prerouter_edge0_35b.safetensors" \
+  "$OUTPUT/prerouter-stacked.safetensors"
+cp "$SOURCE/lora_edge0_35b.safetensors" "$OUTPUT/"
+```
+
+`Models/repacked/` doit contenir `resident.safetensors`, `tokenizer.bin`, ainsi que `experts-L00.bin` à `experts-L39.bin`. Vous pouvez empaqueter l'un ou l'autre modèle, ou les deux.
+
+Checkpoints :
+
+- [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+- [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+
+### Compilation et installation
+
+`com.example.edge0phone` dans le projet est un identifiant de bundle fictif. Avant de compiler ou de signer, remplacez-le par un identifiant de bundle unique que vous contrôlez. Dans Xcode, ouvrez `Edge0PhoneProbe.xcodeproj`, sélectionnez la cible **Edge0Phone**, puis modifiez **Signing & Capabilities → Bundle Identifier**. Choisissez votre équipe Apple Developer, connectez l'iPhone et lancez (Run).
+
+Pour définir l'équipe depuis un fichier, copiez `Config/Local.xcconfig.example` vers `Config/Local.xcconfig` et renseignez `DEVELOPMENT_TEAM`.
+
+La même compilation depuis la ligne de commande :
+
+```sh
+DEVICE_ID=<device id from `xcrun devicectl list devices`>
+TEAM_ID=<your Apple Developer Team ID>
+BUNDLE_ID=com.yourname.edge0phone # replace with your own unique bundle identifier
+
+xcodebuild -project Edge0PhoneProbe.xcodeproj \
+  -scheme Edge0Phone \
+  -destination "id=$DEVICE_ID" \
+  -derivedDataPath .build-device \
+  -allowProvisioningUpdates \
+  -skipPackagePluginValidation \
+  CODE_SIGN_STYLE=Automatic \
+  DEVELOPMENT_TEAM="$TEAM_ID" \
+  PRODUCT_BUNDLE_IDENTIFIER="$BUNDLE_ID" \
+  build
+
+xcrun devicectl device install app \
+  --device "$DEVICE_ID" \
+  .build-device/Build/Products/Debug-iphoneos/Edge0Phone.app
+```
+
+Les poids se trouvent à l'intérieur de ce `.app`. Il n'y a pas de copie séparée vers le téléphone.
+
+### Tests
+
+`scripts/test.sh` compile la bibliothèque MLX Metal avec Xcode, puis exécute les tests unitaires. La première exécution résout le paquet MLX Swift.
+
+```sh
+scripts/test.sh
+EDGE0_TEST_GPU=1 scripts/test.sh
+```
+
+### Utiliser l'application
+
+Lancez l'application et choisissez **Edge0 8B** ou **Edge0 35B**. Un modèle qui n'était pas dans `Models/` au moment de la compilation affiche **Not included** et ne peut pas être sélectionné. Saisissez une question ou touchez l'une des suggestions. La ligne sous une réponse affiche le nombre de tokens, le temps avant le premier token, le prefill, le decode et la mémoire de pointe. Changer de modèle décharge celui qui est chargé. New chat efface la conversation.
+
+## Performances
+
+| Modèle | Appareil | OS | Prefill | TTFT | Decode |
+| --- | --- | --- | --- | --- | --- |
+| 8B | iPhone 16 Pro | iOS 26.6.2 | 7.2 tok/s | 3.6 s | 10.9 tok/s |
+| 35B | iPhone 16 Pro | iOS 26.6.2 | 4.9 tok/s | 2.1 s | 6.4 tok/s |
+
+## Détails techniques
+
+L'application est en SwiftUI. `Edge0MLX` exécute le modèle. `Edge0Core` lit le checkpoint, le tokenizer et l'index de routage. Les calculs mathématiques passent par MLX Swift sur Metal. La cible de déploiement est iOS 17.
+
+Edge0 8B comporte 24 couches, alternant Kimi Delta Attention et attention latente multi-têtes, Top-8, INT4. Le prefill du prompt exécute une couche entière à la fois, par blocs de 2048 tokens au plus. Le premier token est glouton (greedy). Les tokens suivants utilisent une température de 0,7, top-k 64, top-p 0,95 et une pénalité de répétition de 1,1.
+
+Edge0 35B comporte 40 couches, 256 experts par couche, Top-4. Les poids résidents et les fichiers d'experts par couche sont lus séparément. Le cache de préfixes est stocké dans Application Support, car le bundle de l'application est en lecture seule. Les deux modèles ne sont jamais chargés en même temps.
diff --git a/ios/README_ja.md b/ios/README_ja.md
new file mode 100644
index 0000000..7f5f2a0
--- /dev/null
+++ b/ios/README_ja.md
@@ -0,0 +1,111 @@
+# Edge0Phone
+
+[English](README.md) | [中文](README_zh.md) | 日本語 | [Español](README_es.md) | [Français](README_fr.md)
+
+iPhone 向けオンデバイスチャット。アプリは Edge0 8B、Edge0 35B、またはその両方を実行でき、選択したモデルのみがロードされます。モデルの重みはこのフォルダには含まれていません。
+
+## クイックスタート
+
+Apple Silicon Mac、Xcode、そして Developer Mode をオンにした iOS 17 以降の iPhone が必要です。Python 3.10 以降は、重みのダウンロードと 35B の変換にのみ必要です。Xcode が Metal コンパイラの不足を報告する場合は、`xcodebuild -downloadComponent MetalToolchain` を実行してください。
+
+`Models/` ディレクトリはビルド時にアプリへコピーされます。ボタンを選択可能にしたい場合は、ビルド前にモデルをそこに置いてください。空の `Models/` ディレクトリでもビルドは成功します。存在しないモデルは **Not included** と表示されます。
+
+### モデルのダウンロード
+
+```sh
+python3 -m venv .venv-tools
+.venv-tools/bin/python -m pip install --upgrade pip huggingface_hub numpy
+```
+
+8B はダウンロードしたままの状態でパックできます。約 5 GB なので、インストールされるアプリもほぼそのサイズになります。
+
+```sh
+.venv-tools/bin/hf download Edge0/Edge0-8B-A1B-preview \
+  --local-dir Models/Edge0-8B-A1B-preview
+```
+
+35B は重みの変更ではなくレイアウトの変更です。公開されているチェックポイントは、各エキスパートを 19 GB のファイル全体に散らばった 9 回の読み取りとして格納しています。iPhone ではこれは遅すぎるため、以下のツールは各エキスパートを 1 回のシーケンシャル読み取りに書き換え、レイヤーごとに 1 ファイルを書き出します。元のダウンロードは `Models/` の外に保管してください。アプリに含めるべきものは `Models/repacked/` だけです。35B ビルドは約 20 GB になります。
+
+```sh
+mkdir -p checkpoints
+.venv-tools/bin/hf download Edge0/Edge0-35B-A3B-preview \
+  --local-dir checkpoints/Edge0-35B-A3B-preview
+
+SOURCE=checkpoints/Edge0-35B-A3B-preview
+OUTPUT=Models/repacked
+
+.venv-tools/bin/python tools/repack_experts.py pack "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py resident "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py verify "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/convert_tokenizer.py "$SOURCE" "$OUTPUT/tokenizer.bin"
+.venv-tools/bin/python tools/convert_pregate.py \
+  "$SOURCE/prerouter_edge0_35b.safetensors" \
+  "$OUTPUT/prerouter-stacked.safetensors"
+cp "$SOURCE/lora_edge0_35b.safetensors" "$OUTPUT/"
+```
+
+`Models/repacked/` には `resident.safetensors`、`tokenizer.bin`、および `experts-L00.bin` から `experts-L39.bin` が含まれている必要があります。どちらか一方のモデル、または両方をパックできます。
+
+チェックポイント:
+
+- [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+- [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+
+### ビルドとインストール
+
+プロジェクト内の `com.example.edge0phone` はプレースホルダのバンドル識別子です。ビルドや署名の前に、自分が管理する一意のバンドル識別子に置き換えてください。Xcode で `Edge0PhoneProbe.xcodeproj` を開き、**Edge0Phone** ターゲットを選択して、**Signing & Capabilities → Bundle Identifier** を変更します。Apple Developer Team を選び、iPhone を接続して Run します。
+
+代わりにファイルからチームを設定するには、`Config/Local.xcconfig.example` を `Config/Local.xcconfig` にコピーして `DEVELOPMENT_TEAM` を記入します。
+
+コマンドラインからの同じビルド:
+
+```sh
+DEVICE_ID=<device id from `xcrun devicectl list devices`>
+TEAM_ID=<your Apple Developer Team ID>
+BUNDLE_ID=com.yourname.edge0phone # replace with your own unique bundle identifier
+
+xcodebuild -project Edge0PhoneProbe.xcodeproj \
+  -scheme Edge0Phone \
+  -destination "id=$DEVICE_ID" \
+  -derivedDataPath .build-device \
+  -allowProvisioningUpdates \
+  -skipPackagePluginValidation \
+  CODE_SIGN_STYLE=Automatic \
+  DEVELOPMENT_TEAM="$TEAM_ID" \
+  PRODUCT_BUNDLE_IDENTIFIER="$BUNDLE_ID" \
+  build
+
+xcrun devicectl device install app \
+  --device "$DEVICE_ID" \
+  .build-device/Build/Products/Debug-iphoneos/Edge0Phone.app
+```
+
+重みはその `.app` の中に含まれています。スマートフォンへの個別のコピーはありません。
+
+### テスト
+
+`scripts/test.sh` は Xcode で MLX Metal ライブラリをビルドし、その後ユニットテストを実行します。初回実行では MLX Swift パッケージを解決します。
+
+```sh
+scripts/test.sh
+EDGE0_TEST_GPU=1 scripts/test.sh
+```
+
+### アプリの使い方
+
+アプリを起動し、**Edge0 8B** または **Edge0 35B** を選択します。ビルド時に `Models/` になかったモデルは **Not included** と表示され、選択できません。質問を入力するか、提案の 1 つをタップします。返信の下の行には、トークン数、最初のトークンまでの時間、プリフィル、デコード、ピークメモリが表示されます。モデルを切り替えると、ロード済みのモデルがアンロードされます。New chat は会話をクリアします。
+
+## パフォーマンス
+
+| モデル | デバイス | OS | プリフィル | TTFT | デコード |
+| --- | --- | --- | --- | --- | --- |
+| 8B | iPhone 16 Pro | iOS 26.6.2 | 7.2 tok/s | 3.6 s | 10.9 tok/s |
+| 35B | iPhone 16 Pro | iOS 26.6.2 | 4.9 tok/s | 2.1 s | 6.4 tok/s |
+
+## 技術詳細
+
+アプリは SwiftUI です。`Edge0MLX` がモデルを実行します。`Edge0Core` はチェックポイント、トークナイザー、ルーティングインデックスを読み取ります。計算は Metal 上の MLX Swift を通じて行われます。デプロイメントターゲットは iOS 17 です。
+
+Edge0 8B は 24 層で、Kimi Delta Attention とマルチヘッド潜在アテンションが交互に配置され、Top-8、INT4 です。プロンプトのプリフィルは 1 度に 1 層全体を、最大 2048 トークンのチャンクで実行します。最初のトークンは greedy です。以降のトークンは temperature 0.7、top-k 64、top-p 0.95、repetition penalty 1.1 を使用します。
+
+Edge0 35B は 40 層、レイヤーごとに 256 エキスパート、Top-4 です。常駐重みとレイヤーごとのエキスパートファイルは別々に読み取られます。プレフィックスキャッシュは Application Support に保存されます — アプリバンドルは読み取り専用だからです。2 つのモデルが同時にロードされることはありません。
diff --git a/ios/README_zh.md b/ios/README_zh.md
new file mode 100644
index 0000000..4b20f87
--- /dev/null
+++ b/ios/README_zh.md
@@ -0,0 +1,111 @@
+# Edge0Phone
+
+[English](README.md) | 中文 | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+面向 iPhone 的端侧对话 App。可运行 Edge0 8B、Edge0 35B 或两者兼备，且只加载你选中的模型。模型权重不在本目录内。
+
+## 快速开始
+
+你需要一台 Apple Silicon Mac、Xcode，以及一部开启开发者模式、运行 iOS 17 或更高版本的 iPhone。Python 3.10+ 仅用于下载权重与转换 35B。若 Xcode 提示缺少 Metal 编译器，运行 `xcodebuild -downloadComponent MetalToolchain`。
+
+`Models/` 目录会在构建时拷贝进 App。若希望对应按钮可选，请在构建前把模型放入其中。`Models/` 为空也能构建；缺失的模型会显示 **Not included**。
+
+### 下载模型
+
+```sh
+python3 -m venv .venv-tools
+.venv-tools/bin/python -m pip install --upgrade pip huggingface_hub numpy
+```
+
+8B 下载后即可直接打包。它约 5 GB，因此安装后的 App 也约为该体积。
+
+```sh
+.venv-tools/bin/hf download Edge0/Edge0-8B-A1B-preview \
+  --local-dir Models/Edge0-8B-A1B-preview
+```
+
+35B 只是布局变化，权重本身不变。发布的 checkpoint 把每个专家存为一个 19 GB 文件中的九次分散读取。这在 iPhone 上太慢，因此下面的工具会把每个专家重写为一次顺序读取，并按层各写一个文件。原始下载请放在 `Models/` 之外。App 只应包含 `Models/repacked/`。35B 构建约 20 GB。
+
+```sh
+mkdir -p checkpoints
+.venv-tools/bin/hf download Edge0/Edge0-35B-A3B-preview \
+  --local-dir checkpoints/Edge0-35B-A3B-preview
+
+SOURCE=checkpoints/Edge0-35B-A3B-preview
+OUTPUT=Models/repacked
+
+.venv-tools/bin/python tools/repack_experts.py pack "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py resident "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/repack_experts.py verify "$SOURCE" "$OUTPUT"
+.venv-tools/bin/python tools/convert_tokenizer.py "$SOURCE" "$OUTPUT/tokenizer.bin"
+.venv-tools/bin/python tools/convert_pregate.py \
+  "$SOURCE/prerouter_edge0_35b.safetensors" \
+  "$OUTPUT/prerouter-stacked.safetensors"
+cp "$SOURCE/lora_edge0_35b.safetensors" "$OUTPUT/"
+```
+
+`Models/repacked/` 必须包含 `resident.safetensors`、`tokenizer.bin` 以及 `experts-L00.bin` 至 `experts-L39.bin`。可以只打包其中一个模型，也可以两个都打包。
+
+Checkpoint：
+
+- [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+- [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)
+
+### 构建与安装
+
+项目中的 `com.example.edge0phone` 是占位 bundle identifier。构建或签名前，请替换为你自己掌控的唯一 bundle identifier。在 Xcode 中打开 `Edge0PhoneProbe.xcodeproj`，选择 **Edge0Phone** target，然后修改 **Signing & Capabilities → Bundle Identifier**。选择你的 Apple Developer Team，连接 iPhone 并运行。
+
+也可以改用文件方式设置 team：把 `Config/Local.xcconfig.example` 复制为 `Config/Local.xcconfig`，填入 `DEVELOPMENT_TEAM`。
+
+命令行执行同样的构建：
+
+```sh
+DEVICE_ID=<device id from `xcrun devicectl list devices`>
+TEAM_ID=<your Apple Developer Team ID>
+BUNDLE_ID=com.yourname.edge0phone # replace with your own unique bundle identifier
+
+xcodebuild -project Edge0PhoneProbe.xcodeproj \
+  -scheme Edge0Phone \
+  -destination "id=$DEVICE_ID" \
+  -derivedDataPath .build-device \
+  -allowProvisioningUpdates \
+  -skipPackagePluginValidation \
+  CODE_SIGN_STYLE=Automatic \
+  DEVELOPMENT_TEAM="$TEAM_ID" \
+  PRODUCT_BUNDLE_IDENTIFIER="$BUNDLE_ID" \
+  build
+
+xcrun devicectl device install app \
+  --device "$DEVICE_ID" \
+  .build-device/Build/Products/Debug-iphoneos/Edge0Phone.app
+```
+
+权重就在该 `.app` 内部，不会另行拷贝到手机上。
+
+### 测试
+
+`scripts/test.sh` 会用 Xcode 构建 MLX Metal 库，然后运行单元测试。首次运行会解析 MLX Swift 包。
+
+```sh
+scripts/test.sh
+EDGE0_TEST_GPU=1 scripts/test.sh
+```
+
+### 使用 App
+
+启动 App，选择 **Edge0 8B** 或 **Edge0 35B**。构建时不在 `Models/` 中的模型会显示 **Not included** 且不可选择。输入问题，或点按一条建议。回复下方一行显示 token 数、首 token 时间、prefill、decode 与峰值内存。切换模型会卸载当前已加载的模型。新对话会清空当前会话。
+
+## 性能实测
+
+| 模型 | 设备 | 系统 | Prefill | TTFT | 解码 |
+| --- | --- | --- | --- | --- | --- |
+| 8B | iPhone 16 Pro | iOS 26.6.2 | 7.2 tok/s | 3.6 s | 10.9 tok/s |
+| 35B | iPhone 16 Pro | iOS 26.6.2 | 4.9 tok/s | 2.1 s | 6.4 tok/s |
+
+## 技术细节
+
+App 采用 SwiftUI。`Edge0MLX` 负责运行模型。`Edge0Core` 读取 checkpoint、tokenizer 与路由索引。数学运算经由 Metal 上的 MLX Swift 完成。部署目标为 iOS 17。
+
+Edge0 8B 为 24 层，Kimi Delta Attention 与多头潜在注意力（multi-head latent attention）交替，Top-8，INT4。Prompt prefill 以整层为单位运行，每块最多 2048 token。首 token 采用贪心解码，后续 token 使用 temperature 0.7、top-k 64、top-p 0.95 与 1.1 的重复惩罚。
+
+Edge0 35B 为 40 层，每层 256 专家，Top-4。常驻权重与各层专家文件分开读取。前缀缓存存放在 Application Support 中，因为 App bundle 是只读的。两个模型绝不会同时加载。
diff --git a/ios/Sources/Edge0Core/Edge0Configuration8B.swift b/ios/Sources/Edge0Core/Edge0Configuration8B.swift
new file mode 100644
index 0000000..2047695
--- /dev/null
+++ b/ios/Sources/Edge0Core/Edge0Configuration8B.swift
@@ -0,0 +1,99 @@
+import Foundation
+
+/// Minimal configuration surface needed for the Edge0 8B port.
+/// Mirrors Edge0/Edge0-8B-A1B-preview config.json.
+public struct Edge0Configuration8B: Codable, Sendable, Equatable {
+    public let hiddenSize: Int
+    public let hiddenLayers: Int
+    public let intermediateSize: Int
+    public let attentionHeads: Int
+    public let keyValueHeads: Int
+    public let headDim: Int
+    public let vocabularySize: Int
+    public let maxPositionEmbeddings: Int
+    public let rmsNormEps: Float
+
+    public let layerGroupSize: Int
+    public let firstKDenseReplace: Int
+    public let shortConvKernelSize: Int
+    public let kdaSafeGate: Bool
+    public let kdaLowerBound: Float
+
+    public let qLoraRank: Int
+    public let kvLoraRank: Int
+    public let qkNopeHeadDim: Int
+    public let qkRopeHeadDim: Int
+    public let vHeadDim: Int
+    public let ropeTheta: Float
+    public let ropeInterleave: Bool
+
+    public let numExperts: Int
+    public let expertsPerToken: Int
+    public let numSharedExperts: Int
+    public let moeIntermediateSize: Int
+    public let sharedExpertIntermediateSize: Int
+    public let nGroup: Int
+    public let topkGroup: Int
+    public let normTopkProb: Bool
+    public let routedScalingFactor: Float
+    public let routerHasExpertBias: Bool
+
+    public let quantization: Quantization
+
+    public struct Quantization: Codable, Sendable, Equatable {
+        public let groupSize: Int
+        public let bits: Int
+        public let mode: String
+
+        enum CodingKeys: String, CodingKey {
+            case groupSize = "group_size"
+            case bits
+            case mode
+        }
+    }
+
+    enum CodingKeys: String, CodingKey {
+        case hiddenSize = "hidden_size"
+        case hiddenLayers = "num_hidden_layers"
+        case intermediateSize = "intermediate_size"
+        case attentionHeads = "num_attention_heads"
+        case keyValueHeads = "num_key_value_heads"
+        case headDim = "head_dim"
+        case vocabularySize = "vocab_size"
+        case maxPositionEmbeddings = "max_position_embeddings"
+        case rmsNormEps = "rms_norm_eps"
+        case layerGroupSize = "layer_group_size"
+        case firstKDenseReplace = "first_k_dense_replace"
+        case shortConvKernelSize = "short_conv_kernel_size"
+        case kdaSafeGate = "kda_safe_gate"
+        case kdaLowerBound = "kda_lower_bound"
+        case qLoraRank = "q_lora_rank"
+        case kvLoraRank = "kv_lora_rank"
+        case qkNopeHeadDim = "qk_nope_head_dim"
+        case qkRopeHeadDim = "qk_rope_head_dim"
+        case vHeadDim = "v_head_dim"
+        case ropeTheta = "rope_theta"
+        case ropeInterleave = "rope_interleave"
+        case numExperts = "num_experts"
+        case expertsPerToken = "num_experts_per_tok"
+        case numSharedExperts = "num_shared_experts"
+        case moeIntermediateSize = "moe_intermediate_size"
+        case sharedExpertIntermediateSize = "moe_shared_expert_intermediate_size"
+        case nGroup = "n_group"
+        case topkGroup = "topk_group"
+        case normTopkProb = "norm_topk_prob"
+        case routedScalingFactor = "routed_scaling_factor"
+        case routerHasExpertBias = "moe_router_enable_expert_bias"
+        case quantization
+    }
+
+    public static func load(from url: URL) throws -> Self {
+        let data = try Data(contentsOf: url)
+        return try JSONDecoder().decode(Self.self, from: data)
+    }
+
+    public func isMLALayer(_ index: Int) -> Bool {
+        let full = hiddenLayers / layerGroupSize * layerGroupSize
+        return (index + 1) % layerGroupSize == 0 || index >= full
+    }
+}
diff --git a/ios/Sources/Edge0Core/Edge0Tokenizer8B.swift b/ios/Sources/Edge0Core/Edge0Tokenizer8B.swift
new file mode 100644
index 0000000..dd1a9eb
--- /dev/null
+++ b/ios/Sources/Edge0Core/Edge0Tokenizer8B.swift
@@ -0,0 +1,186 @@
+import Foundation
+
+public enum Edge0TokenizerError8B: Error, LocalizedError {
+    case invalid(String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .invalid(let message): message
+        }
+    }
+}
+
+/// Native Swift implementation of the checkpoint's NFC + regex + ByteLevel BPE pipeline.
+/// It reads `tokenizer.json` from the model folder to use the checkpoint vocabulary and
+/// merge table directly.
+public final class Edge0Tokenizer8B: @unchecked Sendable {
+    private struct FileFormat: Decodable {
+        struct AddedToken: Decodable {
+            let id: Int
+            let content: String
+        }
+
+        struct Model: Decodable {
+            let vocab: [String: Int]
+            let merges: [[String]]
+        }
+
+        let added_tokens: [AddedToken]
+        let model: Model
+    }
+
+    private struct Pair: Hashable {
+        let left: String
+        let right: String
+    }
+
+    private static let splitPattern = #"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"#
+
+    private let vocabulary: [String: Int]
+    private let tokensByID: [Int: String]
+    private let mergeRanks: [Pair: Int]
+    private let addedTokens: [(content: String, id: Int)]
+    private let byteEncoder: [UInt8: String]
+    private let byteDecoder: [Character: UInt8]
+    private let splitRegex: NSRegularExpression
+
+    public let endOfTurnTokenID: Int
+
+    public init(contentsOf url: URL) throws {
+        let payload = try JSONDecoder().decode(FileFormat.self, from: Data(contentsOf: url))
+        vocabulary = payload.model.vocab
+        tokensByID = Dictionary(uniqueKeysWithValues: payload.model.vocab.map { ($0.value, $0.key) })
+        mergeRanks = Dictionary(uniqueKeysWithValues: payload.model.merges.enumerated().compactMap { rank, merge in
+            guard merge.count == 2 else { return nil }
+            return (Pair(left: merge[0], right: merge[1]), rank)
+        })
+        addedTokens = payload.added_tokens
+            .map { ($0.content, $0.id) }
+            .sorted { lhs, rhs in lhs.content.count > rhs.content.count }
+        guard let eos = payload.added_tokens.first(where: { $0.content == "<|role_end|>" })?.id else {
+            throw Edge0TokenizerError8B.invalid("tokenizer.json is missing <|role_end|>")
+        }
+        endOfTurnTokenID = eos
+
+        let maps = Self.makeByteMaps()
+        byteEncoder = maps.encoder
+        byteDecoder = maps.decoder
+        splitRegex = try NSRegularExpression(pattern: Self.splitPattern)
+    }
+
+    public func encode(_ text: String) throws -> [Int] {
+        let normalized = text.precomposedStringWithCanonicalMapping
+        var ids: [Int] = []
+        var plain = ""
+        var cursor = normalized.startIndex
+
+        func flushPlain() throws {
+            guard !plain.isEmpty else { return }
+            ids.append(contentsOf: try encodePlain(plain))
+            plain.removeAll(keepingCapacity: true)
+        }
+
+        while cursor < normalized.endIndex {
+            let suffix = normalized[cursor...]
+            if let token = addedTokens.first(where: { suffix.hasPrefix($0.content) }) {
+                try flushPlain()
+                ids.append(token.id)
+                cursor = normalized.index(cursor, offsetBy: token.content.count)
+            } else {
+                let next = normalized.index(after: cursor)
+                plain.append(contentsOf: normalized[cursor..<next])
+                cursor = next
+            }
+        }
+        try flushPlain()
+        return ids
+    }
+
+    public func decode(_ ids: [Int], skipSpecialTokens: Bool = false) throws -> String {
+        var bytes: [UInt8] = []
+        for id in ids {
+            guard let token = tokensByID[id] ?? addedTokens.first(where: { $0.id == id })?.content else {
+                throw Edge0TokenizerError8B.invalid("Unknown token ID \(id)")
+            }
+            if skipSpecialTokens, addedTokens.contains(where: { $0.id == id }) { continue }
+            for character in token {
+                if let byte = byteDecoder[character] {
+                    bytes.append(byte)
+                } else {
+                    bytes.append(contentsOf: String(character).utf8)
+                }
+            }
+        }
+        return String(decoding: bytes, as: UTF8.self)
+    }
+
+    private func encodePlain(_ text: String) throws -> [Int] {
+        let range = NSRange(text.startIndex..<text.endIndex, in: text)
+        let matches = splitRegex.matches(in: text, range: range)
+        var ids: [Int] = []
+        for match in matches {
+            guard let swiftRange = Range(match.range, in: text) else { continue }
+            let byteLevel = text[swiftRange].utf8.map { byteEncoder[$0]! }.joined()
+            for piece in bpe(byteLevel) {
+                guard let id = vocabulary[piece] else {
+                    throw Edge0TokenizerError8B.invalid("BPE vocabulary is missing a merged token")
+                }
+                ids.append(id)
+            }
+        }
+        return ids
+    }
+
+    private func bpe(_ token: String) -> [String] {
+        var pieces = token.map(String.init)
+        while pieces.count > 1 {
+            var selected: Pair?
+            var selectedRank = Int.max
+            for index in 0..<(pieces.count - 1) {
+                let pair = Pair(left: pieces[index], right: pieces[index + 1])
+                if let rank = mergeRanks[pair], rank < selectedRank {
+                    selected = pair
+                    selectedRank = rank
+                }
+            }
+            guard let selected else { break }
+
+            var merged: [String] = []
+            var index = 0
+            while index < pieces.count {
+                if index + 1 < pieces.count,
+                   pieces[index] == selected.left,
+                   pieces[index + 1] == selected.right {
+                    merged.append(selected.left + selected.right)
+                    index += 2
+                } else {
+                    merged.append(pieces[index])
+                    index += 1
+                }
+            }
+            pieces = merged
+        }
+        return pieces
+    }
+
+    private static func makeByteMaps() -> (encoder: [UInt8: String], decoder: [Character: UInt8]) {
+        var bytes = Array(UInt8(33)...UInt8(126))
+        bytes += Array(UInt8(161)...UInt8(172))
+        bytes += Array(UInt8(174)...UInt8(255))
+        var scalars = bytes.map(Int.init)
+        var extra = 0
+        for value in 0...255 where !bytes.contains(UInt8(value)) {
+            bytes.append(UInt8(value))
+            scalars.append(256 + extra)
+            extra += 1
+        }
+        var encoder: [UInt8: String] = [:]
+        var decoder: [Character: UInt8] = [:]
+        for (byte, scalarValue) in zip(bytes, scalars) {
+            let character = Character(String(UnicodeScalar(scalarValue)!))
+            encoder[byte] = String(character)
+            decoder[character] = byte
+        }
+        return (encoder, decoder)
+    }
+}
diff --git a/ios/Sources/Edge0Core/ExpertTensorStore.swift b/ios/Sources/Edge0Core/ExpertTensorStore.swift
new file mode 100644
index 0000000..ae7412f
--- /dev/null
+++ b/ios/Sources/Edge0Core/ExpertTensorStore.swift
@@ -0,0 +1,60 @@
+import Foundation
+
+public enum ExpertProjection: String, CaseIterable, Sendable {
+    case up = "up_proj"
+    case gate = "gate_proj"
+    case down = "down_proj"
+}
+
+public enum QuantPart: String, CaseIterable, Sendable {
+    case weight
+    case scales
+    case biases
+}
+
+public struct ExpertTensorKey: Hashable, Sendable {
+    public let layer: Int
+    public let expert: Int
+    public let projection: ExpertProjection
+    public let part: QuantPart
+}
+
+public struct ExpertTensorSlice: @unchecked Sendable {
+    public let descriptor: TensorDescriptor
+    public let pointer: UnsafeMutableRawPointer
+}
+
+/// Index + mmap layer for the Edge0 checkpoint. No MLX dependency here.
+/// The MLX adapter can turn these slices into arrays only when an expert is selected.
+public final class ExpertTensorStore: @unchecked Sendable {
+    public let index: SafetensorsIndex
+    public let mappedFile: MappedFile
+    public let expertCount: Int
+
+    public init(modelURL: URL, expertCount: Int = 128) throws {
+        self.index = try SafetensorsIndex(fileURL: modelURL)
+        self.mappedFile = try MappedFile(url: modelURL)
+        self.expertCount = expertCount
+    }
+
+    public func stackedTensorName(layer: Int, projection: ExpertProjection, part: QuantPart) -> String {
+        "model.layers.\(layer).mlp.experts.\(projection.rawValue).\(part.rawValue)"
+    }
+
+    public func slice(
+        layer: Int,
+        expert: Int,
+        projection: ExpertProjection,
+        part: QuantPart
+    ) throws -> ExpertTensorSlice {
+        precondition((0..<expertCount).contains(expert))
+        let name = stackedTensorName(layer: layer, projection: projection, part: part)
+        let stacked = try index.tensor(named: name)
+        guard stacked.shape.first == expertCount else {
+            throw SafetensorsError.invalidExpertAxis(name)
+        }
+        let descriptor = try stacked.axisZeroSlice(expert)
+        let pointer = try mappedFile.pointer(to: descriptor.byteRange)
+        return ExpertTensorSlice(descriptor: descriptor, pointer: pointer)
+    }
+}
diff --git a/ios/Sources/Edge0Core/GroupedExpertRouter.swift b/ios/Sources/Edge0Core/GroupedExpertRouter.swift
new file mode 100644
index 0000000..21a2560
--- /dev/null
+++ b/ios/Sources/Edge0Core/GroupedExpertRouter.swift
@@ -0,0 +1,58 @@
+import Foundation
+
+/// CPU reference for Edge0 8B's noaux_tc grouped router.
+/// This is intentionally simple and testable; the production MLX version lives in Edge0MLX.
+public enum GroupedExpertRouter {
+    public struct Selection: Equatable, Sendable {
+        public let indices: [Int]
+        public let weights: [Float]
+    }
+
+    public static func select(
+        logits: [Float],
+        expertBias: [Float],
+        topK: Int = 8,
+        nGroup: Int = 8,
+        topkGroup: Int = 4,
+        normalize: Bool = true,
+        routedScale: Float = 2.5
+    ) -> Selection {
+        precondition(nGroup > 0)
+        precondition(logits.count / nGroup >= 2)
+        precondition(topK <= topkGroup * (logits.count / nGroup))
+        precondition(logits.count == expertBias.count)
+        precondition(logits.count % nGroup == 0)
+        precondition(topK > 0 && topK <= logits.count)
+        precondition(topkGroup > 0 && topkGroup <= nGroup)
+
+        let raw = logits.map { 1 / (1 + Foundation.exp(-$0)) }
+        var selection = zip(raw, expertBias).map(+)
+        let perGroup = logits.count / nGroup
+
+        let groupScores: [(group: Int, score: Float)] = (0..<nGroup).map { group in
+            let start = group * perGroup
+            let end = start + perGroup
+            let bestTwo = selection[start..<end].sorted(by: >).prefix(2)
+            return (group, bestTwo.reduce(0, +))
+        }
+        let surviving = Set(groupScores.sorted { $0.score > $1.score }.prefix(topkGroup).map(\.group))
+
+        for group in 0..<nGroup where !surviving.contains(group) {
+            let start = group * perGroup
+            for i in start..<(start + perGroup) { selection[i] = -.infinity }
+        }
+
+        let indices = selection.indices.sorted {
+            if selection[$0] == selection[$1] { return $0 < $1 }
+            return selection[$0] > selection[$1]
+        }.prefix(topK).map { $0 }
+
+        var weights = indices.map { raw[$0] }
+        if normalize {
+            let denominator = weights.reduce(0, +) + 1e-20
+            weights = weights.map { $0 / denominator }
+        }
+        weights = weights.map { $0 * routedScale }
+        return Selection(indices: indices, weights: weights)
+    }
+}
diff --git a/ios/Sources/Edge0Core/MappedFile.swift b/ios/Sources/Edge0Core/MappedFile.swift
new file mode 100644
index 0000000..51d1853
--- /dev/null
+++ b/ios/Sources/Edge0Core/MappedFile.swift
@@ -0,0 +1,123 @@
+import Foundation
+#if canImport(Darwin)
+import Darwin
+#elseif canImport(Glibc)
+import Glibc
+#endif
+
+public enum MappedFileError: Error, LocalizedError {
+    case openFailed(String)
+    case statFailed(String)
+    case mapFailed(String)
+    case outOfBounds(Range<Int>)
+
+    public var errorDescription: String? {
+        switch self {
+        case .openFailed(let p): return "Could not open \(p)."
+        case .statFailed(let p): return "Could not stat \(p)."
+        case .mapFailed(let p): return "Could not mmap \(p)."
+        case .outOfBounds(let r): return "Mapped byte range out of bounds: \(r)."
+        }
+    }
+}
+
+/// Read-only virtual-memory mapping of a model file.
+/// mmap does not make the whole file resident; pages are faulted in on demand.
+public final class MappedFile: @unchecked Sendable {
+    public let url: URL
+    public let count: Int
+    private let fd: Int32
+    private let base: UnsafeMutableRawPointer
+
+    public init(url: URL) throws {
+        self.url = url
+        let path = url.path
+        let descriptor = open(path, O_RDONLY)
+        guard descriptor >= 0 else { throw MappedFileError.openFailed(path) }
+
+        var st = stat()
+        guard fstat(descriptor, &st) == 0 else {
+            close(descriptor)
+            throw MappedFileError.statFailed(path)
+        }
+        let size = Int(st.st_size)
+        guard size > 0 else {
+            close(descriptor)
+            throw MappedFileError.statFailed(path)
+        }
+
+        let mapping = mmap(nil, size, PROT_READ, MAP_PRIVATE, descriptor, 0)
+        guard mapping != MAP_FAILED, let mapping else {
+            close(descriptor)
+            throw MappedFileError.mapFailed(path)
+        }
+
+        self.fd = descriptor
+        self.count = size
+        self.base = mapping
+    }
+
+    deinit {
+        munmap(base, count)
+        close(fd)
+    }
+
+    public func pointer(to range: Range<Int>) throws -> UnsafeMutableRawPointer {
+        guard range.lowerBound >= 0, range.upperBound <= count else {
+            throw MappedFileError.outOfBounds(range)
+        }
+        return base.advanced(by: range.lowerBound)
+    }
+
+    public func bytes(in range: Range<Int>) throws -> UnsafeRawBufferPointer {
+        let ptr = try pointer(to: range)
+        return UnsafeRawBufferPointer(start: UnsafeRawPointer(ptr), count: range.count)
+    }
+
+    /// Advisory only; ignored if the OS declines it.
+    public func adviseSequential() {
+        #if canImport(Darwin)
+        _ = madvise(base, count, MADV_SEQUENTIAL)
+        #endif
+    }
+
+    /// Official startup prewarm hint. Advisory only; iOS remains free to
+    /// limit page-cache residency under memory pressure.
+    public func adviseWillNeed() {
+        #if canImport(Darwin)
+        _ = madvise(base, count, MADV_WILLNEED)
+        #endif
+    }
+
+    /// Advisory readahead for a checkpoint subrange. Streaming experts use
+    /// this before copying their tensor slices so page faults can be serviced
+    /// in bulk rather than serially on the decode critical path.
+    public func adviseWillNeed(range: Range<Int>) {
+        #if canImport(Darwin)
+        guard range.lowerBound >= 0, range.upperBound <= count,
+              !range.isEmpty else { return }
+        let pageSize = max(4096, Int(sysconf(_SC_PAGESIZE)))
+        let start = range.lowerBound / pageSize * pageSize
+        let end = min(
+            count,
+            ((range.upperBound + pageSize - 1) / pageSize) * pageSize)
+        _ = madvise(base.advanced(by: start), end - start, MADV_WILLNEED)
+        #endif
+    }
+
+    /// Official `SafetensorsMmap.seq_read`: force a real sequential pass over
+    /// the checkpoint. `madvise` alone does not reliably fault every page in.
+    public func sequentialRead(chunkSize: Int = 1 << 24) {
+        let size = max(4096, chunkSize)
+        let buffer = UnsafeMutableRawPointer.allocate(
+            byteCount: size, alignment: MemoryLayout<UInt64>.alignment)
+        defer { buffer.deallocate() }
+        var offset = 0
+        while offset < count {
+            let requested = min(size, count - offset)
+            let bytesRead = pread(fd, buffer, requested, off_t(offset))
+            guard bytesRead > 0 else { break }
+            offset += bytesRead
+        }
+    }
+}
diff --git a/ios/Sources/Edge0Core/SafetensorsIndex.swift b/ios/Sources/Edge0Core/SafetensorsIndex.swift
new file mode 100644
index 0000000..be08109
--- /dev/null
+++ b/ios/Sources/Edge0Core/SafetensorsIndex.swift
@@ -0,0 +1,152 @@
+import Foundation
+
+public enum SafetensorsError: Error, LocalizedError {
+    case fileTooSmall
+    case invalidHeaderLength(UInt64)
+    case invalidHeader
+    case malformedTensor(String)
+    case unsupportedDType(String)
+    case tensorNotFound(String)
+    case invalidExpertAxis(String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .fileTooSmall: return "Safetensors file is too small."
+        case .invalidHeaderLength(let n): return "Invalid safetensors header length: \(n)."
+        case .invalidHeader: return "Safetensors JSON header is invalid."
+        case .malformedTensor(let name): return "Malformed tensor entry: \(name)."
+        case .unsupportedDType(let dtype): return "Unsupported safetensors dtype: \(dtype)."
+        case .tensorNotFound(let name): return "Tensor not found: \(name)."
+        case .invalidExpertAxis(let name): return "Tensor has no valid expert axis: \(name)."
+        }
+    }
+}
+
+public enum SafetensorsDType: String, Sendable, Codable {
+    case bool = "BOOL"
+    case u8 = "U8"
+    case i8 = "I8"
+    case i16 = "I16"
+    case u16 = "U16"
+    case i32 = "I32"
+    case u32 = "U32"
+    case i64 = "I64"
+    case u64 = "U64"
+    case f16 = "F16"
+    case bf16 = "BF16"
+    case f32 = "F32"
+    case f64 = "F64"
+
+    public var byteWidth: Int {
+        switch self {
+        case .bool, .u8, .i8: 1
+        case .i16, .u16, .f16, .bf16: 2
+        case .i32, .u32, .f32: 4
+        case .i64, .u64, .f64: 8
+        }
+    }
+}
+
+public struct TensorDescriptor: Sendable, Equatable {
+    public let name: String
+    public let dtype: SafetensorsDType
+    public let shape: [Int]
+    /// Absolute byte range within the safetensors file.
+    public let byteRange: Range<Int>
+
+    public var byteCount: Int { byteRange.count }
+
+    /// Returns a contiguous slice along axis 0. Edge0 stores routed experts stacked on axis 0.
+    public func axisZeroSlice(_ index: Int) throws -> TensorDescriptor {
+        guard let count = shape.first, count > 0, index >= 0, index < count else {
+            throw SafetensorsError.invalidExpertAxis(name)
+        }
+        guard byteCount % count == 0 else {
+            throw SafetensorsError.invalidExpertAxis(name)
+        }
+        let stride = byteCount / count
+        let start = byteRange.lowerBound + index * stride
+        return TensorDescriptor(
+            name: "\(name)[\(index)]",
+            dtype: dtype,
+            shape: Array(shape.dropFirst()),
+            byteRange: start..<(start + stride)
+        )
+    }
+}
+
+public struct SafetensorsIndex: Sendable {
+    public let fileURL: URL
+    public let headerLength: Int
+    public let dataOffset: Int
+    public let tensors: [String: TensorDescriptor]
+
+    public init(fileURL: URL) throws {
+        self.fileURL = fileURL
+        let handle = try FileHandle(forReadingFrom: fileURL)
+        defer { try? handle.close() }
+
+        let prefix = try handle.read(upToCount: 8) ?? Data()
+        guard prefix.count == 8 else { throw SafetensorsError.fileTooSmall }
+        let headerLength64 = prefix.withUnsafeBytes { raw -> UInt64 in
+            raw.loadUnaligned(as: UInt64.self).littleEndian
+        }
+        guard headerLength64 <= UInt64(Int.max) else {
+            throw SafetensorsError.invalidHeaderLength(headerLength64)
+        }
+        let headerLength = Int(headerLength64)
+        guard headerLength > 0 else { throw SafetensorsError.invalidHeaderLength(headerLength64) }
+
+        let headerData = try handle.read(upToCount: headerLength) ?? Data()
+        guard headerData.count == headerLength else { throw SafetensorsError.fileTooSmall }
+        guard let json = try JSONSerialization.jsonObject(with: headerData) as? [String: Any] else {
+            throw SafetensorsError.invalidHeader
+        }
+
+        let base = 8 + headerLength
+        var parsed: [String: TensorDescriptor] = [:]
+        parsed.reserveCapacity(json.count)
+
+        for (name, value) in json where name != "__metadata__" {
+            guard
+                let entry = value as? [String: Any],
+                let dtypeString = entry["dtype"] as? String,
+                let dtype = SafetensorsDType(rawValue: dtypeString),
+                let shapeNumbers = entry["shape"] as? [NSNumber],
+                let offsets = entry["data_offsets"] as? [NSNumber],
+                offsets.count == 2
+            else {
+                if let entry = value as? [String: Any], let dtypeString = entry["dtype"] as? String,
+                   SafetensorsDType(rawValue: dtypeString) == nil {
+                    throw SafetensorsError.unsupportedDType(dtypeString)
+                }
+                throw SafetensorsError.malformedTensor(name)
+            }
+
+            let shape = shapeNumbers.map(\.intValue)
+            let relativeStart = offsets[0].intValue
+            let relativeEnd = offsets[1].intValue
+            guard relativeStart >= 0, relativeEnd >= relativeStart else {
+                throw SafetensorsError.malformedTensor(name)
+            }
+
+            parsed[name] = TensorDescriptor(
+                name: name,
+                dtype: dtype,
+                shape: shape,
+                byteRange: (base + relativeStart)..<(base + relativeEnd)
+            )
+        }
+
+        self.headerLength = headerLength
+        self.dataOffset = base
+        self.tensors = parsed
+    }
+
+    public subscript(_ name: String) -> TensorDescriptor? { tensors[name] }
+
+    public func tensor(named name: String) throws -> TensorDescriptor {
+        guard let tensor = tensors[name] else { throw SafetensorsError.tensorNotFound(name) }
+        return tensor
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/ChatTemplate.swift b/ios/Sources/Edge0MLX/Edge035B/ChatTemplate.swift
new file mode 100644
index 0000000..7cd6966
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/ChatTemplate.swift
@@ -0,0 +1,122 @@
+import Foundation
+
+/// The text-only slice of this checkpoint's chat template.
+///
+/// The shipped `chat_template.jinja` is 7,764 characters of Jinja with macros, image and
+/// video handling, tool calls and vision counters. None of that is reachable from a
+/// text-only app, and running a Jinja interpreter on device to reach a fixed string would
+/// be a large amount of machinery for no behaviour.
+///
+/// What the text path actually renders is regular:
+///
+/// ```
+/// <|im_start|>{role}\n{content}<|im_end|>\n     per message
+/// <|im_start|>assistant\n<think>\n              generation prompt (thinking on)
+/// <|im_start|>assistant\n                       generation prompt (thinking off)
+/// ```
+///
+/// **Scope is the point, not a shortcut.** Vision and tool-call messages are out — and
+/// they fail loudly rather than rendering something that looks plausible, because a
+/// silently dropped image would produce a coherent answer about nothing.
+///
+/// The trailing `<think>\n` is the checkpoint's default: replies open inside a reasoning
+/// block. Passing `thinking: false` omits it so the model writes the answer directly —
+/// fewer tokens before anything useful appears, at the cost of losing that reasoning pass.
+enum ChatTemplate {
+
+    struct Message {
+        let role: String
+        let content: String
+    }
+
+    enum Failure: LocalizedError {
+        case unsupportedRole(String)
+
+        var errorDescription: String? {
+            switch self {
+            case .unsupportedRole(let role):
+                return "chat template supports system/user/assistant only, not '\(role)'"
+            }
+        }
+    }
+
+    static let supportedRoles: Set<String> = ["system", "user", "assistant"]
+    static let userOpening = "<|im_start|>user\n"
+
+    /// The opening tag every assistant turn is generated under.
+    static func assistantOpening(thinking: Bool) -> String {
+        thinking
+            ? "<|im_start|>assistant\n<think>\n"
+            : "<|im_start|>assistant\n"
+    }
+
+    static func render(
+        _ messages: [Message],
+        addGenerationPrompt: Bool = true,
+        thinking: Bool = true
+    ) throws -> String {
+        var out = ""
+        for message in messages {
+            guard supportedRoles.contains(message.role) else {
+                throw Failure.unsupportedRole(message.role)
+            }
+            // **Past assistant turns are reopened with the same thinking mode.**
+            // Generation runs under `assistantOpening`, so a thinking-on reply begins
+            // inside a reasoning block and contains its `</think>`. Rendering that text
+            // back without the opening tag put a dangling close tag in the history —
+            // well-formed-looking output describing a prompt the model never saw.
+            //
+            // It also makes replay exact, which is what lets a conversation's state be
+            // extended instead of rebuilt.
+            if message.role == "assistant" {
+                out += assistantOpening(thinking: thinking) + "\(message.content)<|im_end|>\n"
+            } else {
+                out += message.role == "user"
+                    ? userOpening + "\(message.content)<|im_end|>\n"
+                    : "<|im_start|>\(message.role)\n\(message.content)<|im_end|>\n"
+            }
+        }
+        if addGenerationPrompt {
+            out += assistantOpening(thinking: thinking)
+        }
+        return out
+    }
+
+    /// What must be fed to continue a conversation whose state already ends inside the
+    /// previous assistant turn.
+    ///
+    /// This is the append path: it closes the open assistant turn, adds the new question,
+    /// and reopens for the answer. Nothing before it is re-rendered, so nothing before it
+    /// is recomputed.
+    static func continuation(user: String, thinking: Bool = true) -> String {
+        "<|im_end|>\n" + userOpening + "\(user)<|im_end|>\n" + assistantOpening(thinking: thinking)
+    }
+}
+
+/// When to stop generating.
+///
+/// Two conditions, and both are needed: `eos_token_id` is a **list** here — 248046
+/// (`<|im_end|>`) and 248044 (`<|endoftext|>`) — so checking only the first lets the
+/// model run past the end of its turn. The length bound is the backstop for the case
+/// where neither is ever emitted, which is what a mis-specified prompt produces.
+struct StopCondition {
+    let endTokens: Set<Int32>
+    let maximumTokens: Int
+
+    init(endTokens: [Int32] = [248046, 248044], maximumTokens: Int = 256) {
+        self.endTokens = Set(endTokens)
+        self.maximumTokens = maximumTokens
+    }
+
+    enum Outcome: Equatable {
+        case keepGoing
+        case hitEndToken(Int32)
+        case hitLimit
+    }
+
+    func evaluate(token: Int32, produced: Int) -> Outcome {
+        if endTokens.contains(token) { return .hitEndToken(token) }
+        if produced >= maximumTokens { return .hitLimit }
+        return .keepGoing
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/CoreMLVocabularyHead.swift b/ios/Sources/Edge0MLX/Edge035B/CoreMLVocabularyHead.swift
new file mode 100644
index 0000000..084e5c9
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/CoreMLVocabularyHead.swift
@@ -0,0 +1,60 @@
+import CoreML
+import Foundation
+import MLX
+
+/// Optional Core ML vocabulary projection used by the 35B path.
+///
+/// The model is discovered beside the streamed checkpoint rather than bundled with the
+/// app: it is a derived 273 MiB asset and the existing MLX path remains the fallback
+/// when it is absent or fails parity. Argmax stays in the graph, so only one Int32 token
+/// crosses back from Core ML.
+final class CoreMLVocabularyHead: @unchecked Sendable {
+    static let directoryName = "lm-head.mlmodelc"
+    static let width = 2_048
+
+    private let model: MLModel
+
+    init(directory: URL) throws {
+        let url = directory.appendingPathComponent(Self.directoryName, isDirectory: true)
+        let configuration = MLModelConfiguration()
+        // Keep Core ML off the GPU so its projection does not compete with the MLX/Metal
+        // decoder. Unsupported operations can fall back to the CPU.
+        configuration.computeUnits = .cpuAndNeuralEngine
+        model = try MLModel(contentsOf: url, configuration: configuration)
+    }
+
+    func predict(_ hidden: MLXArray) async throws -> Int32 {
+        let values = hidden.asType(.float16).asArray(Float16.self)
+        guard values.count == Self.width else {
+            throw Failure.badHiddenSize(values.count)
+        }
+
+        let input = try MLMultiArray(
+            shape: [1, NSNumber(value: Self.width)], dataType: .float16)
+        values.withUnsafeBytes { source in
+            input.dataPointer.copyMemory(from: source.baseAddress!, byteCount: source.count)
+        }
+        let provider = try MLDictionaryFeatureProvider(dictionary: [
+            "hidden_state": MLFeatureValue(multiArray: input)
+        ])
+        let output = try await model.prediction(from: provider)
+        guard let token = output.featureValue(for: "token")?.multiArrayValue,
+              token.count == 1
+        else { throw Failure.missingToken }
+        return token[0].int32Value
+    }
+
+    private enum Failure: LocalizedError {
+        case badHiddenSize(Int)
+        case missingToken
+
+        var errorDescription: String? {
+            switch self {
+            case .badHiddenSize(let count):
+                "Core ML lm_head expected \(CoreMLVocabularyHead.width) values, got \(count)"
+            case .missingToken:
+                "Core ML lm_head did not return token"
+            }
+        }
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/Edge0ChatEngine35B.swift b/ios/Sources/Edge0MLX/Edge035B/Edge0ChatEngine35B.swift
new file mode 100644
index 0000000..518c58c
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/Edge0ChatEngine35B.swift
@@ -0,0 +1,284 @@
+import Darwin
+import Foundation
+import MLX
+
+/// Product-facing wrapper around the streaming Edge0 35B implementation.
+///
+/// This engine deliberately lives beside `Edge0ChatEngine`: selecting 35B creates only
+/// this model, while selecting 8B continues through the existing engine unchanged.
+public final class Edge0ChatEngine35B: @unchecked Sendable {
+    public static let modelFolderName = "repacked"
+
+    private let directory: URL
+    private let model: Edge0Model35B
+    private let tokenizer: Tokenizer
+    private let instructions: String
+    private let systemPrefix: [Int32]
+    private let state = Edge0Model35B.State()
+
+    private var processed: [Int32] = []
+    private var cachedPrefix: PrefixCache?
+
+    public init(
+        modelURL: URL,
+        instructions: String = "You are a helpful assistant.",
+        progress: (String) -> Void = { _ in }
+    ) throws {
+        directory = modelURL
+        self.instructions = instructions
+
+        // Configure MLX memory use for the selected 35B runtime.
+        let highMem = CommandLine.arguments.contains("--high-mem-35b")
+        Memory.cacheLimit = (highMem ? 1536 : 384) * 1_048_576
+        print(String(
+            format: "[mem] cacheLimit=%dMiB highMem=%@ available=%dMiB",
+            Memory.cacheLimit / 1_048_576,
+            highMem ? "true" : "false",
+            os_proc_available_memory() / 1_048_576))
+        progress("mapping Edge0 35B weights")
+        var configuration = Edge0Model35B.Configuration()
+        configuration.round6 = CommandLine.arguments.contains("--round6")
+        if configuration.round6 {
+            configuration.activeExperts = 2
+        }
+        configuration.retainAllResidentWeights = highMem
+        model = try Edge0Model35B(directory: modelURL, configuration: configuration)
+        guard model.availableLayers == model.configuration.layers else {
+            throw Edge035BFailure.incompleteExperts(
+                found: model.availableLayers, expected: model.configuration.layers)
+        }
+        progress("loading Edge0 35B tokenizer")
+        tokenizer = try Tokenizer(url: modelURL.appendingPathComponent("tokenizer.bin"))
+
+        if instructions.isEmpty {
+            systemPrefix = []
+        } else {
+            let text = try ChatTemplate.render(
+                [.init(role: "system", content: instructions)],
+                addGenerationPrompt: false)
+            // The first user role marker is fixed too. Persisting its state saves those
+            // tokens on every new conversation; `PrefixCache.matches` still verifies the
+            // encoded prefix against the complete prompt before it is trusted.
+            systemPrefix = tokenizer.encode(text + ChatTemplate.userOpening)
+        }
+        cachedPrefix = PrefixCache.read(from: Self.prefixURL(
+            directory: modelURL, fingerprint: model.configuration.fingerprint))
+        try model.prepareForFirstTurn()
+        progress("Edge0 35B ready")
+    }
+
+    public func reset() {
+        state.reset()
+        processed.removeAll(keepingCapacity: true)
+    }
+
+    public func reply(
+        to userText: String,
+        maxTokens: Int = 512,
+        thinking: Bool = true,
+        onText: @escaping @Sendable (String) -> Void = { _ in },
+        shouldContinue: @escaping @Sendable () -> Bool = { true }
+    ) async throws -> Edge0GenerationResult {
+        let prompt = userText.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard !prompt.isEmpty else { throw Edge035BFailure.invalid("Message cannot be empty") }
+        guard maxTokens > 0 else { throw Edge035BFailure.invalid("maxTokens must be positive") }
+
+        let continuing = !processed.isEmpty
+        let ids: [Int32]
+        if continuing {
+            ids = tokenizer.encode(ChatTemplate.continuation(user: prompt, thinking: thinking))
+        } else {
+            var messages: [ChatTemplate.Message] = []
+            if !instructions.isEmpty {
+                messages.append(.init(role: "system", content: instructions))
+            }
+            messages.append(.init(role: "user", content: prompt))
+            ids = tokenizer.encode(try ChatTemplate.render(messages, thinking: thinking))
+        }
+        guard !ids.isEmpty else { throw Edge035BFailure.invalid("Tokenizer returned an empty prompt") }
+        print("[35b] thinking=\(thinking ? "on" : "off")")
+
+        GPU.resetPeakMemory()
+        let started = Date()
+        let thermalAtStart = ProcessInfo.processInfo.thermalState
+        var start = 0
+        if !continuing {
+            state.reset()
+            processed.removeAll(keepingCapacity: true)
+            if let cache = cachedPrefix, cache.matches(ids) {
+                restore(cache.state, into: state)
+                processed = cache.tokens
+                start = cache.tokens.count
+            }
+        }
+
+        var next: Int32 = 0
+        // Later turns already have attention weights resident. Keep that cache on
+        // for a continued prefill so those projections are not copied again.
+        // The first prefill still copies each of them once; retaining earlier
+        // only adds dirty pages with nothing to reuse inside that single pass.
+        if continuing { model.beginDecodeCaching() }
+        model.profile.reset()
+        // An older cache may contain only the system message. Restore it, compute just
+        // the newly-added user opening, then atomically replace it with the longer cache.
+        let capturePoint = !continuing && start < systemPrefix.count
+            && ids.starts(with: systemPrefix) ? systemPrefix.count : nil
+        var index = start
+        while index < ids.count {
+            guard shouldContinue() else { return try cancelled() }
+            var size = min(Edge0Model35B.maximumBatch, ids.count - index)
+            if let capturePoint, index < capturePoint {
+                size = min(size, capturePoint - index)
+            }
+            let slice = Array(ids[index ..< index + size])
+            let isLast = index + size == ids.count
+            let logits = try model.step(
+                tokens: MLXArray(slice, [1, size]), state: state, needsLogits: isLast)
+            if isLast, let logits { next = try await model.greedy(logits) }
+            processed.append(contentsOf: slice)
+            index += size
+
+            if let capturePoint, index == capturePoint {
+                let captured = PrefixCache(tokens: systemPrefix, state: state.snapshot())
+                cachedPrefix = captured
+                captured.write(to: Self.prefixURL(
+                    directory: directory, fingerprint: model.configuration.fingerprint))
+            }
+        }
+        let prefillFinished = Date()
+        let prefillSeconds = prefillFinished.timeIntervalSince(started)
+        let prefillTokens = ids.count - start
+        print(String(
+            format: "[35b prefill] tokens=%d total=%.2fs rate=%.2f tok/s %@",
+            prefillTokens,
+            prefillSeconds,
+            Double(prefillTokens) / max(0.000_001, prefillSeconds),
+            model.profile.report))
+        model.beginDecodeCaching()
+        model.profile.reset()
+
+        let stop = StopCondition(maximumTokens: maxTokens)
+        var raw = ""
+        var shown = ""
+        var produced = 0
+        var firstTokenAt: Date?
+        var decodeStartedAt: Date?
+
+        while true {
+            switch stop.evaluate(token: next, produced: produced) {
+            case .hitEndToken, .hitLimit:
+                let finished = Date()
+                let decodeSeconds = decodeStartedAt.map {
+                    finished.timeIntervalSince($0)
+                } ?? 0
+                let footprint = Self.physicalFootprintBytes()
+                print(String(
+                    format: "[35b decode] tokens=%d total=%.2fs rate=%.2f tok/s thermal=%@→%@ mlxPeak=%.0fMB footprint=%.0fMB available=%dMB %@",
+                    max(0, produced - 1),
+                    decodeSeconds,
+                    Double(max(0, produced - 1)) / max(0.000_001, decodeSeconds),
+                    Self.thermalName(thermalAtStart),
+                    Self.thermalName(ProcessInfo.processInfo.thermalState),
+                    Double(Memory.peakMemory) / 1_048_576,
+                    Double(footprint) / 1_048_576,
+                    os_proc_available_memory() / 1_048_576,
+                    model.profile.report))
+                shown = thinking
+                    ? ThinkingBlock.settled(raw)
+                    : ThinkingBlock.settledReply(raw)
+                if !shown.isEmpty { onText(shown) }
+                return Edge0GenerationResult(
+                    text: shown,
+                    generatedTokenCount: produced,
+                    elapsedSeconds: finished.timeIntervalSince(started),
+                    prefillSeconds: prefillSeconds,
+                    prefillTokensPerSecond: Double(prefillTokens) /
+                        max(0.000_001, prefillSeconds),
+                    timeToFirstTokenSeconds: firstTokenAt?.timeIntervalSince(started)
+                        ?? finished.timeIntervalSince(started),
+                    decodeTokensPerSecond: decodeStartedAt.map {
+                        Double(max(0, produced - 1)) /
+                            max(0.000_001, finished.timeIntervalSince($0))
+                    } ?? 0,
+                    peakMemoryBytes: max(Memory.peakMemory, footprint))
+            case .keepGoing:
+                break
+            }
+
+            guard shouldContinue() else { return try cancelled() }
+            if firstTokenAt == nil {
+                firstTokenAt = Date()
+                decodeStartedAt = firstTokenAt
+            }
+            raw += tokenizer.decode([next])
+            let visible = thinking
+                ? ThinkingBlock.visible(in: raw)
+                : ThinkingBlock.visibleReply(in: raw)
+            if let visible, visible != shown {
+                shown = visible
+                onText(shown)
+            }
+            produced += 1
+
+            let logits = try model.step(tokens: MLXArray([next], [1, 1]), state: state)
+            processed.append(next)
+            next = try await model.greedy(logits)
+        }
+    }
+
+    private func cancelled<T>() throws -> T {
+        reset()
+        throw CancellationError()
+    }
+
+    private func restore(_ source: Edge0Model35B.State, into target: Edge0Model35B.State) {
+        let snapshot = source.snapshot()
+        target.convolution = snapshot.convolution
+        target.recurrent = snapshot.recurrent
+        target.keys = snapshot.keys
+        target.values = snapshot.values
+        target.offset = snapshot.offset
+    }
+
+    private static func prefixURL(directory: URL, fingerprint: String) -> URL {
+        PrefixCache.fileURL(in: directory, fingerprint: fingerprint)
+    }
+
+    private static func thermalName(_ state: ProcessInfo.ThermalState) -> String {
+        switch state {
+        case .nominal: "nominal"
+        case .fair: "fair"
+        case .serious: "serious"
+        case .critical: "critical"
+        @unknown default: "unknown"
+        }
+    }
+
+    /// Process physical footprint (jetsam-relevant), not just MLX allocator peak.
+    private static func physicalFootprintBytes() -> Int {
+        var info = task_vm_info_data_t()
+        var count = mach_msg_type_number_t(
+            MemoryLayout<task_vm_info_data_t>.stride / MemoryLayout<natural_t>.stride)
+        let result = withUnsafeMutablePointer(to: &info) {
+            $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
+                task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count)
+            }
+        }
+        guard result == KERN_SUCCESS else { return 0 }
+        return Int(info.phys_footprint)
+    }
+}
+
+private enum Edge035BFailure: LocalizedError {
+    case incompleteExperts(found: Int, expected: Int)
+    case invalid(String)
+
+    var errorDescription: String? {
+        switch self {
+        case .incompleteExperts(let found, let expected):
+            return "Edge0 35B expert files are incomplete (found \(found), expected \(expected))"
+        case .invalid(let message):
+            return message
+        }
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/Edge0Model35B.swift b/ios/Sources/Edge0MLX/Edge035B/Edge0Model35B.swift
new file mode 100644
index 0000000..1b1669a
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/Edge0Model35B.swift
@@ -0,0 +1,1012 @@
+import Foundation
+import MLX
+import MLXFast
+
+/// Streaming Edge0 35B model implementation used by the iOS chat engine.
+///
+/// This implementation is kept independent of the Python reference so parity checks can
+/// compare separate code paths. It is not main-actor isolated because model steps perform
+/// synchronous GPU and file I/O. The chat engine serializes access to one model instance.
+final class Edge0Model35B: @unchecked Sendable {
+
+    /// How many tokens may go through one forward pass.
+    ///
+    /// Sets the expert buffer ceiling: a batch of T tokens can in the worst case route to
+    /// T*K distinct experts. The loader grows staging storage only to the union actually
+    /// observed, so this ceiling no longer reserves the worst case up front.
+    ///
+    /// Storage is sized for the observed expert union rather than the worst-case T*K.
+    /// Top-4 routing makes the batch size a tradeoff between repeated base-weight reads
+    /// and the memory required for the active graph and expert union.
+    /// Tokens per prefill forward. Decode ignores this.
+    ///
+    /// A full-layer read is about 8 s of flash for 40 layers, independent of
+    /// how many tokens share it. 512 is the first size at which that fixed
+    /// cost can land near 60 tok/s before GPU time; the Gated DeltaNet scan
+    /// is what will jetsam first if 512 is too large.
+    static let maximumBatch = 512
+    /// Below this, prefill keeps the decode-style expert union. A full-layer
+    /// read is 432 MiB × 40 and only pays for itself on a long chunk.
+    static let fullLayerPrefillMinimum = 64
+
+    struct Configuration {
+        var layers = 40
+        var hidden = 2048
+        var moeIntermediate = 512
+        var expertsPerLayer = 256
+        // The public Edge0-35B-A3B-preview adapter and prerouter were trained with K=4.
+        // This must stay aligned with the shipped artifacts: changing K changes both
+        // the MoE output and the routing history consumed by the next-step heads.
+        var activeExperts = 4
+        /// Round-6 pregate + LoRA, trained at K=2. The two files are a pair.
+        var round6 = false
+        var groupSize = 64
+        var bits = 4
+        var fullAttentionInterval = 4
+        var vocabulary = 248_320
+
+        /// When true, keep every non-expert resident tensor after first materialisation.
+        /// The default retains only selected tensors to limit active memory use.
+        var retainAllResidentWeights = false
+
+        /// Source used to prefetch expert reads for the next decode step.
+        var routingGuess: RoutingGuess = .trainedHeads
+
+        /// Number of recent routing steps used by `.recentRouting`.
+        var routingLookback = 1
+
+        enum RoutingGuess {
+            /// Do not prefetch expert reads.
+            case none
+            /// Prefetch the experts selected by each layer on the previous step.
+            case recentRouting
+            /// The trained pregate heads, when their file is present. Predicts the
+            /// *next* layer a *token* ahead, so it is the only source that can cover a
+            /// layer before that layer has run — and it replaces the router to do it.
+            case trainedHeads
+        }
+
+        /// Identifies this configuration in caches whose contents depend on it.
+        var fingerprint: String { "full" }
+
+        /// `(index + 1) % interval == 0` are full attention; the rest are linear.
+        func isLinear(_ layer: Int) -> Bool {
+            (layer + 1) % fullAttentionInterval != 0
+        }
+    }
+
+    enum Failure: LocalizedError {
+        case weightsMissing(URL)
+        case tensorMissing(String)
+        case layerFailed(Int)
+
+        var errorDescription: String? {
+            switch self {
+            case .weightsMissing(let url):
+                return "no model at \(url.lastPathComponent)"
+            case .tensorMissing(let name):
+                return "missing tensor \(name)"
+            case .layerFailed(let layer):
+                return "layer \(layer) could not run"
+            }
+        }
+    }
+
+    /// State carried between decode steps.
+    ///
+    /// Two kinds, and **they are not interchangeable**. Linear layers hold a fixed-size
+    /// `[1, heads, dim, dim]` float32 matrix updated in place: it does not grow with
+    /// context and cannot be pruned. Full-attention layers hold a KV cache that does
+    /// grow. Any context policy that drops old turns shortens the second and leaves the
+    /// first untouched — a model that appears to forget while still carrying everything.
+    final class State {
+        var convolution: [Int: MLXArray] = [:]
+        var recurrent: [Int: MLXArray] = [:]
+        var keys: [Int: MLXArray] = [:]
+        var values: [Int: MLXArray] = [:]
+        var offset = 0
+
+        /// Routing carried between decode steps. Per-conversation, not per-model: two
+        /// conversations share one set of heads but must never share the routing
+        /// history that feeds them.
+        let pregate = PregateHeads.State()
+
+        /// Layer → the experts it ran, most recent step first, capped at the lookback.
+        /// Previous-step router output, used by `.recentRouting`.
+        var recentRouting: [Int: [[Int]]] = [:]
+
+        /// Layer → what was asked for on its behalf. Exists so the hit rate is a
+        /// measurement of what was *issued* rather than a restatement of the guess:
+        /// `recentRouting` is overwritten as each layer runs, so comparing against it
+        /// would compare a layer's routing with its own.
+        var prefetched: [Int: [Int]] = [:]
+
+        var byteCount: Int {
+            [convolution, recurrent, keys, values].reduce(0) { total, table in
+                total + table.values.reduce(0) { $0 + $1.size * $1.itemSize }
+            }
+        }
+
+        func reset() {
+            convolution.removeAll(); recurrent.removeAll()
+            keys.removeAll(); values.removeAll()
+            offset = 0
+            pregate.reset()
+            recentRouting.removeAll(); prefetched.removeAll()
+        }
+    }
+
+    let configuration: Configuration
+    private let weights: SafetensorsBundle
+    private let experts: ExpertLoader
+    private let lora: SafetensorsShard?
+    private let directory: URL
+    private let coreMLHead: CoreMLVocabularyHead?
+    private var coreMLHeadEnabled: Bool
+    private var coreMLHeadValidated = false
+
+    init(directory: URL, configuration: Configuration = Configuration()) throws {
+        let index = directory.appendingPathComponent("model.safetensors.index.json")
+        guard FileManager.default.fileExists(atPath: index.path) else {
+            throw Failure.weightsMissing(directory)
+        }
+        self.directory = directory
+        self.configuration = configuration
+        self.weights = try SafetensorsBundle(directory: directory)
+        self.experts = try ExpertLoader(
+            directory: directory, layers: configuration.layers,
+            slots: Edge0Model35B.maximumBatch * configuration.activeExperts,
+            hidden: configuration.hidden,
+            moeIntermediate: configuration.moeIntermediate,
+            groupSize: configuration.groupSize)
+        self.experts.profile = self.profile
+        let round6LoRA = directory.appendingPathComponent("lora-round6.safetensors")
+        let officialLoRA = directory.appendingPathComponent("lora_edge0_35b.safetensors")
+        let legacyLoRA = directory.appendingPathComponent("lora.safetensors")
+        let loraURL: URL
+        if configuration.round6 {
+            loraURL = round6LoRA
+        } else if FileManager.default.fileExists(atPath: officialLoRA.path) {
+            loraURL = officialLoRA
+        } else {
+            loraURL = legacyLoRA
+        }
+        self.lora = try? SafetensorsShard(url: loraURL)
+        if configuration.round6 {
+            print("[round6] K=\(configuration.activeExperts) lora=\(self.lora == nil ? "missing" : "loaded")")
+        }
+        let compiledHead = directory.appendingPathComponent(
+            CoreMLVocabularyHead.directoryName, isDirectory: true)
+        if CommandLine.arguments.contains("--ane-35b"),
+           FileManager.default.fileExists(atPath: compiledHead.path) {
+            do {
+                self.coreMLHead = try CoreMLVocabularyHead(directory: directory)
+                self.coreMLHeadEnabled = true
+                print("[coreml-head] loaded with computeUnits=cpuAndNeuralEngine")
+            } catch {
+                self.coreMLHead = nil
+                self.coreMLHeadEnabled = false
+                print("[coreml-head] unavailable: \(error)")
+            }
+        } else {
+            self.coreMLHead = nil
+            self.coreMLHeadEnabled = false
+            print("[coreml-head] disabled or absent")
+        }
+        let free = (try? directory.resourceValues(
+            forKeys: [.volumeAvailableCapacityForImportantUsageKey]))?
+            .volumeAvailableCapacityForImportantUsage
+        print("[weights] free=\(free.map { String($0 / 1_048_576) } ?? "?")MiB")
+
+        // Pregate weights are optional. Report an invalid present file so a fallback to
+        // the normal router is visible during diagnostics.
+        let pregateURL = PregateHeads.fileURL(in: directory, round6: configuration.round6)
+        if configuration.routingGuess == .trainedHeads,
+           FileManager.default.fileExists(atPath: pregateURL.path) {
+            do {
+                let heads = try PregateHeads(
+                    url: pregateURL, hidden: configuration.hidden,
+                    experts: RepackedExperts.expertsPerLayer)
+                self.pregate = heads
+                print("[pregate] heads loaded: owners \(heads.owners.first ?? -1)..\(heads.owners.last ?? -1) n=\(heads.owners.count)")
+                if let report = heads.selfCheck(
+                    fixture: directory.appendingPathComponent("pregate-fixture.safetensors")) {
+                    print(report)
+                }
+            } catch {
+                print("[pregate] present but unusable: \(error)")
+                self.pregate = nil
+            }
+        } else {
+            self.pregate = nil
+        }
+    }
+
+    /// The trained routing heads, when the file is present. See `PregateHeads`.
+    private let pregate: PregateHeads?
+
+    /// See `StepProfile` for why only the synchronous phases are timed.
+    let profile = StepProfile()
+
+    var loraLoaded: Bool { lora != nil }
+    var availableLayers: Int { experts.availableLayers.count }
+
+    // MARK: - Weights
+
+    /// Selected tensors retained after first use to reduce repeated materialization.
+    private var retainedWeights: [String: MLXArray] = [:]
+    private var retainAttentionWeights = false
+
+    /// First-turn prefill is latency-sensitive and sees every attention weight only
+    /// once, so retaining them there adds pressure without reuse. Decode immediately
+    /// reuses them; enable the cache at that boundary and keep it warm for later turns.
+    func beginDecodeCaching() {
+        retainAttentionWeights = true
+    }
+
+    /// Prepare selected attention weights and issue a discarded token before generation
+    /// so initial Metal setup is outside the first reply's latency measurement.
+    func prepareForFirstTurn() throws {
+        beginDecodeCaching()
+        let started = Date()
+        var count = 0
+        for name in weights.index.keys where
+            name.hasPrefix("language_model.lm_head.")
+            || name.contains(".linear_attn.")
+            || name.contains(".self_attn.") {
+            if weight(name) != nil { count += 1 }
+        }
+        print(String(
+            format: "[weights] first-turn cache tensors=%d %.2fs",
+            count, Date().timeIntervalSince(started)))
+
+        let gpu = Date()
+        let scratch = State()
+        _ = try step(tokens: MLXArray([Int32(0)], [1, 1]), state: scratch)
+        profile.reset()
+        print(String(format: "[weights] gpu warm %.2fs", Date().timeIntervalSince(gpu)))
+    }
+
+    /// Materialize a tensor from the mapped checkpoint. Retain selected tensors according
+    /// to the configured memory policy; other tensors are released after use.
+    private func weight(_ name: String) -> MLXArray? {
+        if let retained = retainedWeights[name] { return retained }
+        guard let entry = weights.entry(name), let pointer = weights.pointer(to: name)
+        else { return nil }
+        let raw = UnsafeRawBufferPointer(start: pointer, count: entry.byteCount)
+        profile.recordWeight(bytes: entry.byteCount)
+        let made = profile.measure(.weightCopy) { () -> MLXArray? in
+            switch entry.dtype {
+            case "U32", "I32": return MLXArray(raw, entry.shape, type: UInt32.self)
+            case "BF16":
+                return MLXArray(raw, entry.shape, type: UInt16.self).view(dtype: .bfloat16)
+            case "F32": return MLXArray(raw, entry.shape, type: Float.self)
+            default: return nil
+            }
+        }
+        let shouldRetain = configuration.retainAllResidentWeights
+            || (!coreMLHeadEnabled && name.hasPrefix("language_model.lm_head."))
+            || (retainAttentionWeights
+                && (name.contains(".linear_attn.") || name.contains(".self_attn.")))
+        if shouldRetain, let made {
+            retainedWeights[name] = made
+        }
+        return made
+    }
+
+    /// Apply the base quantized linear layer and its optional LoRA update.
+    private func linear(_ x: MLXArray, _ stem: String, bits: Int) -> MLXArray? {
+        guard let w = weight("\(stem).weight"),
+              let s = weight("\(stem).scales"),
+              let b = weight("\(stem).biases")
+        else { return nil }
+        let base = quantizedMatmul(
+            x.asType(.bfloat16), w, scales: s, biases: b, transpose: true,
+            groupSize: configuration.groupSize, bits: bits, mode: .affine)
+
+        guard let prepared = preparedLoRA(stem) else { return base }
+        let delta = matmul(matmul(x.asType(.float16), prepared.a), prepared.b)
+        return base + delta.asType(base.dtype)
+    }
+
+    /// LoRA matrices are converted, transposed, scaled, and cached on first use.
+    private struct PreparedLoRA {
+        let a: MLXArray
+        let b: MLXArray
+    }
+    private var loraCache: [String: PreparedLoRA] = [:]
+
+    private func preparedLoRA(_ stem: String) -> PreparedLoRA? {
+        if let cached = loraCache[stem] { return cached }
+        guard let lora,
+              let aEntry = lora.entries["\(stem).lora_A"],
+              let bEntry = lora.entries["\(stem).lora_B"],
+              let aPointer = lora.pointer(to: "\(stem).lora_A"),
+              let bPointer = lora.pointer(to: "\(stem).lora_B")
+        else { return nil }
+
+        func loadAdapter(_ entry: SafetensorsShard.Entry,
+                         _ pointer: UnsafeRawPointer) -> MLXArray? {
+            let raw = UnsafeRawBufferPointer(start: pointer, count: entry.byteCount)
+            switch entry.dtype {
+            case "F16": return MLXArray(raw, entry.shape, type: Float16.self)
+            case "F32": return MLXArray(raw, entry.shape, type: Float.self)
+            default: return nil
+            }
+        }
+        guard let a = loadAdapter(aEntry, aPointer),
+              let bb = loadAdapter(bEntry, bPointer)
+        else { return nil }
+        // float16 throughout, as the Python implementation does.
+        let prepared = PreparedLoRA(
+            a: a.asType(.float16).transposed(),
+            b: (Float16(2.0) * bb.asType(.float16)).transposed())
+        // Materialised here so the graph that produced them is not replayed per token.
+        eval(prepared.a, prepared.b)
+        loraCache[stem] = prepared
+        return prepared
+    }
+
+    private func norm(_ x: MLXArray, _ name: String) -> MLXArray? {
+        guard let w = weight(name) else { return nil }
+        return MLXFast.rmsNorm(x.asType(.bfloat16), weight: w, eps: 1e-6)
+    }
+
+    private func rmsNormalise(_ x: MLXArray) -> MLXArray {
+        let squares = mean(x.asType(.float32) * x.asType(.float32), axis: -1, keepDims: true)
+        return (x.asType(.float32) * rsqrt(squares + 1e-6)).asType(x.dtype)
+    }
+
+    // MARK: - Layers
+
+    /// Copy just the rows named, straight out of the mapping.
+    ///
+    /// `weight(_:)` builds an `MLXArray` over a whole tensor, and building one **copies**
+    /// — mlx-c routes it to `allocator::malloc` plus `std::copy`. For a table indexed one
+    /// row at a time that is the wrong shape by five orders of magnitude.
+    private func rows(_ name: String, _ indices: [Int32]) -> MLXArray? {
+        guard let entry = weights.entry(name), let base = weights.pointer(to: name),
+              entry.shape.count == 2, entry.shape[0] > 0
+        else { return nil }
+        let columns = entry.shape[1]
+        let rowBytes = entry.byteCount / entry.shape[0]
+        var gathered = [UInt8](repeating: 0, count: rowBytes * indices.count)
+        gathered.withUnsafeMutableBytes { destination in
+            for (slot, row) in indices.enumerated() {
+                guard row >= 0, Int(row) < entry.shape[0] else { continue }
+                memcpy(destination.baseAddress!.advanced(by: slot * rowBytes),
+                       base.advanced(by: Int(row) * rowBytes), rowBytes)
+            }
+        }
+        let shape = [indices.count, columns]
+        return gathered.withUnsafeBytes { raw -> MLXArray? in
+            switch entry.dtype {
+            case "U32", "I32": return MLXArray(raw, shape, type: UInt32.self)
+            case "BF16": return MLXArray(raw, shape, type: UInt16.self).view(dtype: .bfloat16)
+            case "F32": return MLXArray(raw, shape, type: Float.self)
+            default: return nil
+            }
+        }
+    }
+
+    /// Gather only the embedding rows needed by the current token batch.
+    private func embed(_ tokens: MLXArray) -> MLXArray? {
+        let stem = "language_model.model.embed_tokens"
+        let ids = tokens.reshaped([-1]).asArray(Int32.self)
+        guard let w = rows("\(stem).weight", ids),
+              let s = rows("\(stem).scales", ids),
+              let b = rows("\(stem).biases", ids)
+        else { return nil }
+        // Dequantize only the gathered rows rather than the full embedding table.
+        let out = dequantized(w, scales: s, biases: b,
+                              groupSize: configuration.groupSize, bits: configuration.bits,
+                              mode: .affine)
+        return out.reshaped(tokens.shape + [out.dim(-1)])
+    }
+
+    /// Grouped-query attention with partial RoPE and an output gate.
+    ///
+    /// The gate is the second half of `q_proj`; RoPE covers 64 of each 256-dimensional
+    /// head and uses a base of 10,000,000.
+    private func fullAttention(_ x: MLXArray, layer: Int, state: State) -> MLXArray? {
+        let stem = "language_model.model.layers.\(layer).self_attn"
+        let heads = 16, kvHeads = 2, headDim = 256
+        let (batch, length) = (x.dim(0), x.dim(1))
+
+        guard let qOut = linear(x, "\(stem).q_proj", bits: 4),
+              let kOut = linear(x, "\(stem).k_proj", bits: 4),
+              let vOut = linear(x, "\(stem).v_proj", bits: 4),
+              let qNorm = weight("\(stem).q_norm.weight"),
+              let kNorm = weight("\(stem).k_norm.weight")
+        else { return nil }
+
+        let split = qOut.reshaped([batch, length, heads, 2 * headDim])
+        var queries = split[.ellipsis, 0 ..< headDim]
+        let gate = split[.ellipsis, headDim ..< (2 * headDim)]
+            .reshaped([batch, length, heads * headDim])
+
+        var keys = kOut.reshaped([batch, length, kvHeads, headDim])
+        var values = vOut.reshaped([batch, length, kvHeads, headDim])
+
+        queries = MLXFast.rmsNorm(queries, weight: qNorm, eps: 1e-6).transposed(0, 2, 1, 3)
+        keys = MLXFast.rmsNorm(keys, weight: kNorm, eps: 1e-6).transposed(0, 2, 1, 3)
+        values = values.transposed(0, 2, 1, 3)
+
+        let offset = state.offset
+        func rope(_ a: MLXArray) -> MLXArray {
+            MLXFast.RoPE(a, dimensions: 64, traditional: false,
+                         base: 10_000_000, scale: 1, offset: offset)
+        }
+        queries = rope(queries)
+        keys = rope(keys)
+
+        if let pastKeys = state.keys[layer], let pastValues = state.values[layer] {
+            keys = concatenated([pastKeys, keys], axis: 2)
+            values = concatenated([pastValues, values], axis: 2)
+        }
+        state.keys[layer] = keys
+        state.values[layer] = values
+
+        // A single query attending to everything before it needs no mask; passing a
+        // square one sized for the prompt is how a working prefill breaks decode.
+        let total = keys.dim(2)
+        var mask: MLXArray? = nil
+        if length > 1 {
+            var entries = [Float](repeating: 0, count: length * total)
+            for row in 0 ..< length {
+                for column in (offset + row + 1) ..< total {
+                    entries[row * total + column] = -Float.greatestFiniteMagnitude
+                }
+            }
+            mask = MLXArray(entries, [length, total]).asType(queries.dtype)
+        }
+
+        let attended = MLXFast.scaledDotProductAttention(
+            queries: queries, keys: keys, values: values, scale: 1.0 / 16.0, mask: mask)
+        let merged = attended.transposed(0, 2, 1, 3)
+            .reshaped([batch, length, heads * headDim])
+        return linear(merged * sigmoid(gate), "\(stem).o_proj", bits: 4)
+    }
+
+    /// Gated DeltaNet: depthwise causal convolution, then a gated delta-rule recurrence
+    /// over a float32 state.
+    ///
+    /// `q` and `k` carry 16 heads against `v`'s 32 and are matched by **consecutive**
+    /// duplication — h0, h0, h1, h1 — not by tiling. Both run; one is correct.
+    private func linearAttention(_ x: MLXArray, layer: Int, state: State) -> MLXArray? {
+        let stem = "language_model.model.layers.\(layer).linear_attn"
+        let kHeads = 16, vHeads = 32, headDim = 128
+        let keyDim = kHeads * headDim, valueDim = vHeads * headDim
+        let convDim = keyDim * 2 + valueDim, kernel = 4
+        let (batch, length) = (x.dim(0), x.dim(1))
+
+        guard let qkv = linear(x, "\(stem).in_proj_qkv", bits: 4),
+              let zRaw = linear(x, "\(stem).in_proj_z", bits: 4),
+              let aRaw = linear(x, "\(stem).in_proj_a", bits: 4),
+              let bRaw = linear(x, "\(stem).in_proj_b", bits: 4),
+              let convWeight = weight("\(stem).conv1d.weight"),
+              let aLog = weight("\(stem).A_log"),
+              let dtBias = weight("\(stem).dt_bias"),
+              let normWeight = weight("\(stem).norm.weight")
+        else { return nil }
+
+        let history = state.convolution[layer]?.asType(qkv.dtype)
+            ?? MLXArray.zeros([batch, kernel - 1, convDim]).asType(qkv.dtype)
+        let tail = concatenated([history, qkv], axis: 1)
+
+        // `conv1d` accumulates internally in float32. Preserve that precision for
+        // numerically consistent recurrent state updates.
+        let convolved = conv1d(tail, convWeight, groups: convDim)
+        let activated = convolved * sigmoid(convolved)
+        state.convolution[layer] = tail[0..., (tail.dim(1) - (kernel - 1))...]
+
+        var q = activated[.ellipsis, 0 ..< keyDim].reshaped([batch, length, kHeads, headDim])
+        var k = activated[.ellipsis, keyDim ..< (2 * keyDim)]
+            .reshaped([batch, length, kHeads, headDim])
+        let v = activated[.ellipsis, (2 * keyDim) ..< convDim]
+            .reshaped([batch, length, vHeads, headDim])
+
+        // q is scaled by Dk⁻¹ and k by Dk^-0.5 — not the same factor.
+        let inverseScale = Float(pow(Double(headDim), -0.5))
+        q = rmsNormalise(q) * (inverseScale * inverseScale)
+        k = rmsNormalise(k) * inverseScale
+        q = repeated(q, count: vHeads / kHeads, axis: -2)
+        k = repeated(k, count: vHeads / kHeads, axis: -2)
+
+        let beta = sigmoid(bRaw.asType(.float32))
+        let shifted = aRaw.asType(.float32) + dtBias.asType(.float32)
+        let g = MLX.exp(-MLX.exp(aLog.asType(.float32)) * MLX.log(MLX.exp(shifted) + 1))
+
+        var recurrent = state.recurrent[layer]?.asType(.float32)
+            ?? MLXArray.zeros([batch, vHeads, headDim, headDim], dtype: .float32)
+        var outputs: [MLXArray] = []
+        // Use matrix contractions for the recurrent updates to avoid materializing large
+        // broadcast-multiply intermediates:
+        //
+        //     (recurrent * kt[..., None, :]).sum(-1)     2 MiB temporary, reduced to 16 KB
+        //     matmul(recurrent, kt[..., None])           16 KB, straight out of a GEMV
+        //
+        // GEMV accumulation order can differ from elementwise reduction in the last bits.
+        for t in 0 ..< length {
+            let qt = q[0..., t].asType(.float32).expandedDimensions(axis: -1)
+            let kt = k[0..., t].asType(.float32).expandedDimensions(axis: -1)
+            let vt = v[0..., t].asType(.float32)
+            let decay = g[0..., t].expandedDimensions(axis: -1).expandedDimensions(axis: -1)
+
+            recurrent = recurrent * decay
+            let memory = matmul(recurrent, kt).squeezed(axis: -1)
+            let delta = (vt - memory) * beta[0..., t].expandedDimensions(axis: -1)
+            // The outer product has to materialise 2 MiB whichever way it is written, so
+            // this one stays as it was.
+            recurrent = recurrent
+                + kt.swappedAxes(-1, -2) * delta.expandedDimensions(axis: -1)
+            outputs.append(matmul(recurrent, qt).squeezed(axis: -1))
+        }
+        state.recurrent[layer] = recurrent
+
+        let stacked = MLX.stacked(outputs, axis: 1).asType(x.dtype)
+        let z = zRaw.reshaped([batch, length, vHeads, headDim])
+        // RMSNormGated, in float32 as trained.
+        let normalised = MLXFast.rmsNorm(stacked, weight: normWeight, eps: 1e-6)
+            .asType(.float32)
+        let gated = (sigmoid(z.asType(.float32)) * z.asType(.float32) * normalised)
+            .asType(x.dtype)
+        return linear(gated.reshaped([batch, length, valueDim]), "\(stem).out_proj", bits: 4)
+    }
+
+    /// Router, K streamed experts, and the shared expert.
+    ///
+    /// Apply softmax over all expert logits before top-k selection, then renormalize the
+    /// selected scores.
+    /// Routed experts for **one or more** tokens.
+    ///
+    /// The single-token form this replaces treated the K experts as the batch axis, every
+    /// one of them seeing the same input. With T tokens there are T*K (token, expert)
+    /// pairs and each needs its own input row, so the batch axis becomes T*K and the
+    /// expert for each row is named by `rhsIndices`.
+    ///
+    /// For a batch, read the union of experts selected by its tokens. The union may be
+    /// smaller than T*K when tokens share expert selections.
+    private func moe(_ x: MLXArray, layer: Int, state: State) -> MLXArray? {
+        let prefix = "language_model.model.layers.\(layer).mlp"
+        let k = configuration.activeExperts
+        let tokens = x.dim(1)
+
+        if tokens >= Self.fullLayerPrefillMinimum {
+            return moeFullLayer(x, layer: layer, state: state, prefix: prefix, k: k, tokens: tokens)
+        }
+
+        let ids: [Int]
+        var scores: MLXArray
+
+        // Predictions replace the normal gate during single-token decode. Prefill uses
+        // the normal gate because its multi-token routing shape differs.
+        if tokens == 1, let predicted = state.pregate.predicted[layer],
+           let predictedScores = state.pregate.scores[layer], predicted.count == k {
+            ids = predicted
+            // Rebuilt here, from k floats. The alternative — keeping the slice of the
+            // staging array — would hold that array and its graph across the step.
+            scores = MLXArray(predictedScores, [1, 1, k])
+        } else {
+            guard let logits = linear(x, "\(prefix).gate", bits: 8) else { return nil }
+            let gates = softmax(logits.asType(.float32), axis: -1)
+            let chosen = argPartition(-gates, kth: k - 1, axis: -1)[.ellipsis, 0 ..< k]
+            scores = takeAlong(gates, chosen, axis: -1)
+            scores = scores / scores.sum(axis: -1, keepDims: true)
+
+            // Expert IDs must reach the CPU before their file reads can be issued.
+            let kind: StepProfile.Phase =
+                configuration.isLinear(layer) ? .syncLinear : .syncFull
+            ids = profile.measure(kind) { () -> [Int] in
+                eval(chosen, scores)
+                return chosen.reshaped([-1]).asArray(Int32.self).map { Int($0) }
+            }
+        }
+        guard ids.count == tokens * k else { return nil }
+
+        // Captured for this layer's own head, which runs at the step boundary and
+        // predicts for `layer + 1`. The last position is the right one for both: a
+        // prefill's final position is what the first decode step continues from.
+        //
+        // Store the routes actually used, including predictions, as inputs for the next
+        // pregate evaluation. Tensor reshaping is deferred until after graph evaluation.
+        if let pregate, pregate.owners.contains(layer) {
+            state.pregate.moeInput[layer] = x
+            state.pregate.executed[layer] = ids.suffix(k).map { Int32($0) }
+        }
+
+        // Distinct experts in first-appearance order, and where each pair's expert landed.
+        var rowOf: [Int: Int] = [:]
+        var union: [Int] = []
+        for id in ids where rowOf[id] == nil {
+            rowOf[id] = union.count
+            union.append(id)
+        }
+        profile.recordRouting(union: union.count, pairs: ids.count)
+
+        // How much of what this layer needs was already asked for. **This is the number
+        // that bounds the whole scheme** — the read can only be hidden for blocks the
+        // guess anticipated — and it is free to collect, because both sets are already
+        // on the CPU. Recorded before the routing is overwritten below.
+        if configuration.routingGuess != .none {
+            let guessed = state.prefetched[layer] ?? []
+            profile.recordPrefetch(
+                hits: union.count { guessed.contains($0) }, total: union.count)
+        }
+        // What executed, for the next step to guess from. The union rather than the
+        // last position's pair: during prefill that is strictly more informative, and
+        // during decode they are the same thing.
+        var history = state.recentRouting[layer] ?? []
+        history.insert(union, at: 0)
+        state.recentRouting[layer] = Array(history.prefix(configuration.routingLookback))
+
+        guard let pool = try? experts.load(layer: layer, experts: union),
+              let rows = try? ids.map({ id -> Int32 in
+                  guard let row = rowOf[id] else { throw Failure.layerFailed(layer) }
+                  return Int32(row)
+              })
+        else { return nil }
+
+        // **The single-token form is kept separate, and it is not premature.** Generalising
+        // this for batching added a `reshaped` on top of a broadcast — a stride-0 view that
+        // a reshape has to materialise — and decode runs this on every token of every
+        // layer. Batching is used only by prefill; making decode pay for it was a
+        // regression, not a simplification.
+        let expanded = x.reshaped([tokens, 1, configuration.hidden]).asType(.bfloat16)
+        let inputs = tokens == 1
+            ? broadcast(expanded, to: [k, 1, configuration.hidden])
+            : broadcast(expanded, to: [tokens, k, configuration.hidden])
+                .reshaped([tokens * k, 1, configuration.hidden])
+        let indices = MLXArray(rows)
+
+        func project(_ input: MLXArray, _ offset: Int) -> MLXArray {
+            gatherQuantizedMM(
+                input, pool[offset], scales: pool[offset + 1], biases: pool[offset + 2],
+                rhsIndices: indices, transpose: true,
+                groupSize: configuration.groupSize, bits: configuration.bits,
+                // **An assertion about the input, not a request**, so it is true only
+                // when it is true. With one token the union is built in first-appearance
+                // order from k distinct experts, making the rows exactly `0 ..< k` and
+                // ascending; with more, rows repeat and are unordered, and claiming
+                // otherwise would silently gather the wrong experts.
+                mode: .affine, sortedIndices: tokens == 1)
+        }
+        let gate = project(inputs, ExpertLoader.gate)
+        let up = project(inputs, ExpertLoader.up)
+        let routed = project(gate * sigmoid(gate) * up, ExpertLoader.down)
+        let weighted = routed
+            * scores.reshaped([tokens * k, 1, 1]).asType(routed.dtype)
+        let combined = weighted
+            .reshaped([tokens, k, configuration.hidden])
+            .sum(axis: 1)
+
+        // Compute the shared expert after routed expert loading to avoid competing for
+        // memory bandwidth during the reads.
+        guard let sharedGate = linear(x, "\(prefix).shared_expert_gate", bits: 8),
+              let sg = linear(x, "\(prefix).shared_expert.gate_proj", bits: 4),
+              let su = linear(x, "\(prefix).shared_expert.up_proj", bits: 4),
+              let shared = linear(sg * sigmoid(sg) * su,
+                                  "\(prefix).shared_expert.down_proj", bits: 4)
+        else { return nil }
+        return combined.reshaped(shared.shape) + sigmoid(sharedGate) * shared
+    }
+
+    /// Prefill for a long chunk: the layer file is read once, and expert ids
+    /// index that stack directly. The ids still come back to the CPU afterwards
+    /// so the pregate feature sees what actually ran; the read no longer waits
+    /// on that sync.
+    private func moeFullLayer(_ x: MLXArray, layer: Int, state: State,
+                              prefix: String, k: Int, tokens: Int) -> MLXArray? {
+        guard let pool = try? experts.loadFullLayer(layer: layer),
+              let logits = linear(x, "\(prefix).gate", bits: 8)
+        else { return nil }
+        let gates = softmax(logits.asType(.float32), axis: -1)
+        let chosen = argPartition(-gates, kth: k - 1, axis: -1)[.ellipsis, 0 ..< k]
+        var scores = takeAlong(gates, chosen, axis: -1)
+        scores = scores / scores.sum(axis: -1, keepDims: true)
+
+        let kind: StepProfile.Phase =
+            configuration.isLinear(layer) ? .syncLinear : .syncFull
+        let ids: [Int] = profile.measure(kind) {
+            eval(chosen, scores)
+            return chosen.reshaped([-1]).asArray(Int32.self).map { Int($0) }
+        }
+        guard ids.count == tokens * k else { return nil }
+        if let pregate, pregate.owners.contains(layer) {
+            state.pregate.moeInput[layer] = x
+            state.pregate.executed[layer] = ids.suffix(k).map { Int32($0) }
+        }
+        var rowOf: [Int: Int] = [:]
+        var union: [Int] = []
+        for id in ids where rowOf[id] == nil {
+            rowOf[id] = union.count
+            union.append(id)
+        }
+        profile.recordRouting(union: union.count, pairs: ids.count)
+        var history = state.recentRouting[layer] ?? []
+        history.insert(union, at: 0)
+        state.recentRouting[layer] = Array(history.prefix(configuration.routingLookback))
+
+        let expanded = x.reshaped([tokens, 1, configuration.hidden]).asType(.bfloat16)
+        let inputs = broadcast(expanded, to: [tokens, k, configuration.hidden])
+            .reshaped([tokens * k, 1, configuration.hidden])
+        // Row `e` of the full stack is expert `e`, so the router index is the gather index.
+        let indices = chosen.reshaped([tokens * k])
+        func project(_ input: MLXArray, _ offset: Int) -> MLXArray {
+            gatherQuantizedMM(
+                input, pool[offset], scales: pool[offset + 1], biases: pool[offset + 2],
+                rhsIndices: indices, transpose: true,
+                groupSize: configuration.groupSize, bits: configuration.bits,
+                mode: .affine, sortedIndices: false)
+        }
+        let gate = project(inputs, ExpertLoader.gate)
+        let up = project(inputs, ExpertLoader.up)
+        let routed = project(gate * sigmoid(gate) * up, ExpertLoader.down)
+        let weighted = routed * scores.reshaped([tokens * k, 1, 1]).asType(routed.dtype)
+        let combined = weighted.reshaped([tokens, k, configuration.hidden]).sum(axis: 1)
+        guard let sharedGate = linear(x, "\(prefix).shared_expert_gate", bits: 8),
+              let sg = linear(x, "\(prefix).shared_expert.gate_proj", bits: 4),
+              let su = linear(x, "\(prefix).shared_expert.up_proj", bits: 4),
+              let shared = linear(sg * sigmoid(sg) * su,
+                                  "\(prefix).shared_expert.down_proj", bits: 4)
+        else { return nil }
+        return combined.reshaped(shared.shape) + sigmoid(sharedGate) * shared
+    }
+
+    /// Run all routing heads as one batch and return the expert IDs for the next step.
+    ///
+    /// Returns the expert set per consuming layer, which is what a prefetch needs.
+    @discardableResult
+    private func stagePregate(state: State) -> [Int: [Int]] {
+        guard let pregate else { return [:] }
+        let k = configuration.activeExperts
+
+        var inputs: [MLXArray] = []
+        var executed: [[Int32]] = [], previous: [[Int32]] = []
+        inputs.reserveCapacity(pregate.owners.count)
+        for owner in pregate.owners {
+            // A pruned or otherwise absent owner would desynchronise the stack from
+            // `owners`, so the whole staging is skipped rather than silently shifted.
+            //
+            // Skipping has to *clear* the predictions, not just decline to write new
+            // ones. `advance()` deliberately keeps `predicted` alive across the step
+            // boundary — that is how a prediction made at step N reaches step N+1 — so
+            // leaving it untouched here would hand step N+1 the routing computed for
+            // step N, one token stale, with nothing in any log to say so.
+            guard let input = state.pregate.moeInput[owner],
+                  let ran = state.pregate.executed[owner]
+            else {
+                state.pregate.predicted = [:]
+                state.pregate.scores = [:]
+                state.pregate.advance()
+                return [:]
+            }
+            // Sliced here rather than at capture time: by now `eval(h)` has run, so this
+            // reshapes a materialised array instead of adding a node to a live graph.
+            // The last position is the right one for a prefill chunk too — it is what
+            // the first decode step continues from.
+            let vector = input.ndim == 1
+                ? input
+                : (input.dim(1) == 1 ? input.reshaped([-1]) : input[0, -1])
+            inputs.append(vector)
+            executed.append(ran)
+            // Use a zero row when no previous-token routing is available.
+            previous.append(state.pregate.previous[owner] ?? [])
+        }
+
+        let logits = pregate.logits(
+            inputs: stacked(inputs, axis: 0),
+            executed: pregate.oneHotRows(executed),
+            previous: pregate.oneHotRows(previous))
+
+        // Apply the same softmax, top-k selection, and score normalization as the router.
+        let gates = softmax(logits, axis: -1)
+        let chosen = argPartition(-gates, kth: k - 1, axis: -1)[.ellipsis, 0 ..< k]
+        var weights = takeAlong(gates, chosen, axis: -1)
+        weights = weights / weights.sum(axis: -1, keepDims: true)
+
+        // The one barrier, and it brings back both halves. Pulling the weights across
+        // too costs 66 floats and means nothing MLX-shaped is carried into the next
+        // step — a view into `weights` would have kept the array, and its graph, alive
+        // for every one of the thirty-three layers that held one.
+        let (flat, weighting) = profile.measure(.pregateStage) { () -> ([Int32], [Float]) in
+            eval(chosen, weights)
+            return (chosen.reshaped([-1]).asArray(Int32.self),
+                    weights.reshaped([-1]).asArray(Float.self))
+        }
+
+        var sets: [Int: [Int]] = [:]
+        for (index, consumer) in pregate.consumers.enumerated() {
+            let range = (index * k) ..< ((index + 1) * k)
+            let slice = flat[range].map { Int($0) }
+            state.pregate.predicted[consumer] = slice
+            state.pregate.scores[consumer] = Array(weighting[range])
+            sets[consumer] = Array(Set(slice)).sorted()
+        }
+
+        state.pregate.advance()
+        return sets
+    }
+
+    // MARK: - Forward
+
+    /// Run one model step and evaluate the resulting graph once at the end.
+    func step(tokens: MLXArray, state: State) throws -> MLXArray {
+        try step(tokens: tokens, state: state, needsLogits: true)!
+    }
+
+    /// One step, with the option of **not** projecting to the vocabulary.
+    ///
+    /// During prefill, intermediate positions do not need vocabulary logits. This option
+    /// skips that projection when the result will not be consumed.
+    func step(tokens: MLXArray, state: State, needsLogits: Bool) throws -> MLXArray? {
+        guard var h = embed(tokens) else { throw Failure.tensorMissing("embed_tokens") }
+        // Checkpoint weights use bfloat16; keep activations in the same format.
+        h = h.asType(.bfloat16)
+
+        let batchedPrefill = tokens.dim(1) >= Self.fullLayerPrefillMinimum
+        for layer in 0 ..< configuration.layers {
+            if batchedPrefill && layer % 10 == 0 {
+                print("[35b layer] \(layer) T=\(tokens.dim(1))")
+            }
+            let stem = "language_model.model.layers.\(layer)"
+            guard let normed = norm(h, "\(stem).input_layernorm.weight") else {
+                throw Failure.layerFailed(layer)
+            }
+            let attended = configuration.isLinear(layer)
+                ? linearAttention(normed, layer: layer, state: state)
+                : fullAttention(normed, layer: layer, state: state)
+            guard let attended else { throw Failure.layerFailed(layer) }
+            h = h + attended
+
+            guard let post = norm(h, "\(stem).post_attention_layernorm.weight"),
+                  let mixed = moe(post, layer: layer, state: state)
+            else { throw Failure.layerFailed(layer) }
+            h = h + mixed
+            // A full-layer expert stack is 432 MiB. Evaluating here drops it
+            // before the next layer's read; holding all forty would be 17 GB.
+            if batchedPrefill {
+                eval(h)
+                if let held = state.pregate.moeInput[layer] {
+                    let last = held[0, -1]
+                    eval(last)
+                    state.pregate.moeInput[layer] = last
+                }
+            }
+        }
+        eval(h)
+
+        // Run routing predictions after graph evaluation; they consume hidden states
+        // produced by this step and provide expert IDs for the next step.
+        switch configuration.routingGuess {
+        case .none:
+            break
+        case .trainedHeads:
+            if tokens.dim(1) > 1 && !needsLogits {
+                // Batched prefill never consumes the one-token-ahead head predictions.
+                // Reuse the real routing as a read hint and avoid loading/running 132 MiB
+                // of heads until the final prefill batch, whose prediction is consumed
+                // by the first decode step.
+                state.prefetched = state.recentRouting.mapValues { Array(Set($0.joined())) }
+                for (layer, chosen) in state.prefetched {
+                    experts.prefetch(layer: layer, experts: chosen)
+                }
+                break
+            }
+            state.prefetched = [:]
+            for (layer, chosen) in stagePregate(state: state) {
+                state.prefetched[layer] = chosen
+                experts.prefetch(layer: layer, experts: chosen)
+            }
+            // The official heads cover consumers 7...39. Preserve their exact routing,
+            // but prefetch the uncovered first layers from their last real routing.
+            // This is advice only: misses still run the router and read the right files.
+            for layer in 0 ..< configuration.layers where state.prefetched[layer] == nil {
+                guard let chosen = state.recentRouting[layer]?.first else { continue }
+                state.prefetched[layer] = chosen
+                experts.prefetch(layer: layer, experts: chosen)
+            }
+        case .recentRouting:
+            // Snapshotted, not read live: `recentRouting` is rewritten layer by layer as
+            // the next step runs, so scoring the hit rate against it would compare a
+            // layer's routing with its own and report a meaningless 100%.
+            state.prefetched = state.recentRouting.mapValues { Array(Set($0.joined())) }
+            for (layer, chosen) in state.prefetched {
+                experts.prefetch(layer: layer, experts: chosen)
+            }
+        }
+
+        state.offset += tokens.dim(1)
+        profile.countStep()
+        guard needsLogits else { return nil }
+        guard let final = norm(h, "language_model.model.norm.weight") else {
+            throw Failure.tensorMissing("model.norm")
+        }
+        if coreMLHeadEnabled {
+            // Core ML consumes only the final position. This is 4 KiB instead of the
+            // 248,320-element logits vector, and avoids retaining MLX's 286 MiB head.
+            let last = final[0, -1].reshaped([1, configuration.hidden])
+            eval(last)
+            return last
+        }
+        guard let logits = linear(final, "language_model.lm_head", bits: 4) else {
+            throw Failure.tensorMissing("lm_head")
+        }
+        return logits
+    }
+
+    /// Argmax over the last position's logits.
+    func greedy(_ output: MLXArray) async throws -> Int32 {
+        if coreMLHeadEnabled, let coreMLHead,
+           output.dim(output.ndim - 1) == configuration.hidden {
+            let coreMLToken = try await profile.measure(.vocabularyHead) {
+                try await coreMLHead.predict(output)
+            }
+            if !coreMLHeadValidated {
+                guard let logits = linear(output, "language_model.lm_head", bits: 4) else {
+                    throw Failure.tensorMissing("lm_head")
+                }
+                let mlxToken = Self.mlxGreedy(logits)
+                if mlxToken != coreMLToken {
+                    coreMLHeadEnabled = false
+                    print("[coreml-head] parity FAIL coreml=\(coreMLToken) mlx=\(mlxToken); falling back")
+                    return mlxToken
+                }
+                coreMLHeadValidated = true
+                retainedWeights = retainedWeights.filter {
+                    !$0.key.hasPrefix("language_model.lm_head.")
+                }
+                print("[coreml-head] parity PASS token=\(coreMLToken)")
+            }
+            return coreMLToken
+        }
+        return Self.mlxGreedy(output)
+    }
+
+    private static func mlxGreedy(_ logits: MLXArray) -> Int32 {
+        // `-1`, not `0`: with a batched forward the interesting logits are the last
+        // position's. Identical for a single token, which is why this was never wrong.
+        let final = logits.ndim == 2 ? logits[0] : logits[0, -1]
+        let choice = argMax(final.asType(.float32), axis: -1)
+        eval(choice)
+        return choice.item(Int32.self)
+    }
+
+    /// Greedy decode. Prompt and generation take the same path — one token at a time —
+    /// which costs a little on prefill and buys a single code path.
+    func generate(promptIds: [Int32], stop: StopCondition,
+                  state: State? = nil,
+                  onToken: @escaping (Int32) -> Bool) throws -> [Int32]
+    {
+        let state = state ?? State()
+        var produced: [Int32] = []
+        var next: Int32 = 0
+
+        for (index, token) in promptIds.enumerated() {
+            let logits = try step(tokens: MLXArray([token], [1, 1]), state: state)
+            if index == promptIds.count - 1 {
+                let choice = argMax(logits[0, 0].asType(.float32), axis: -1)
+                eval(choice)
+                next = choice.item(Int32.self)
+            }
+        }
+
+        while true {
+            switch stop.evaluate(token: next, produced: produced.count) {
+            case .hitEndToken, .hitLimit:
+                return produced
+            case .keepGoing:
+                break
+            }
+            produced.append(next)
+            // The callback decides whether to continue, so cancellation is the caller's
+            // to express rather than something this loop has to be told about.
+            guard onToken(next) else { return produced }
+
+            let logits = try step(tokens: MLXArray([next], [1, 1]), state: state)
+            let choice = argMax(logits[0, 0].asType(.float32), axis: -1)
+            eval(choice)
+            next = choice.item(Int32.self)
+        }
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/ExpertLoader.swift b/ios/Sources/Edge0MLX/Edge035B/ExpertLoader.swift
new file mode 100644
index 0000000..1687f20
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/ExpertLoader.swift
@@ -0,0 +1,241 @@
+import Foundation
+import MLX
+
+/// Loads routed experts from the per-layer files and builds the stacked tensors used by
+/// `gatherQuantizedMM`. Reads for the selected experts are issued concurrently.
+final class ExpertLoader: @unchecked Sendable {
+
+    enum Failure: LocalizedError {
+        case noLayers(URL)
+        case unknownLayer(Int, available: Int)
+        case expertOutOfRange(Int)
+        case wrongSlotCount(asked: Int, configured: Int)
+        case readFailed(layer: Int, expert: Int)
+
+        var errorDescription: String? {
+            switch self {
+            case .noLayers(let url):
+                return "no intact layer files in \(url.lastPathComponent)"
+            case .unknownLayer(let layer, let available):
+                return "layer \(layer) is not loaded (\(available) layers available)"
+            case .expertOutOfRange(let expert):
+                return "expert \(expert) outside 0..<\(RepackedExperts.expertsPerLayer)"
+            case .wrongSlotCount(let asked, let configured):
+                return "asked for \(asked) experts, loader is built for \(configured)"
+            case .readFailed(let layer, let expert):
+                return "read failed for layer \(layer) expert \(expert)"
+            }
+        }
+    }
+
+    /// Where each tensor sits in the stacked pool handed to `gatherQuantizedMM`:
+    /// weight, scales, biases for gate, then up, then down.
+    static let gate = 0, up = 3, down = 6
+
+    private let experts: RepackedExperts
+    private let slots: Int
+    private var buffers: [UnsafeMutableRawPointer] = []
+
+
+    private let shapes: [(weight: [Int], scale: [Int])]
+
+    /// Set by the model so reads and array construction land in the same accounting as
+    /// the router barrier. Optional so the loader stays usable without it.
+    var profile: StepProfile?
+
+    var availableLayers: [Int] { experts.availableLayers }
+    var rejectedFiles: [(layer: Int, bytes: Int)] { experts.rejected }
+
+    init(directory: URL, layers: Int, slots: Int, hidden: Int,
+         moeIntermediate: Int, groupSize: Int) throws
+    {
+        self.experts = RepackedExperts(directory: directory, expectedLayers: layers)
+        guard !experts.availableLayers.isEmpty else { throw Failure.noLayers(directory) }
+
+        self.slots = slots
+
+        self.shapes = [
+            (weight: [moeIntermediate, hidden / 8], scale: [moeIntermediate, hidden / groupSize]),
+            (weight: [moeIntermediate, hidden / 8], scale: [moeIntermediate, hidden / groupSize]),
+            (weight: [hidden, moeIntermediate / 8], scale: [hidden, moeIntermediate / groupSize]),
+        ]
+    }
+
+    /// Nine tensors: gate weight/scales/biases, then up, then down.
+    ///
+    /// Throws rather than returning something plausible. A silently wrong expert
+    /// produces fluent nonsense several layers later, which is far more expensive to
+    /// find than a thrown error here.
+    deinit {
+        buffers.forEach { $0.deallocate() }
+    }
+
+    /// Ask the kernel to prefetch the selected expert blocks into the file cache.
+    func prefetch(layer: Int, experts chosen: [Int]) {
+        for expert in chosen where expert >= 0 && expert < RepackedExperts.expertsPerLayer {
+            experts.prefetch(layer: layer, expert: expert)
+        }
+    }
+
+    /// Load **1 to `slots`** distinct experts.
+    ///
+    /// Was fixed at exactly `slots`, which is what a single token needs at K=2. A batch
+    /// of T tokens routes to the *union* of their choices — anywhere from K (all tokens
+    /// agree) to T*K (none do) — so `slots` became the ceiling rather than the count.
+    func load(layer: Int, experts chosen: [Int]) throws -> [MLXArray] {
+        guard !chosen.isEmpty, chosen.count <= slots else {
+            throw Failure.wrongSlotCount(asked: chosen.count, configured: slots)
+        }
+        let count = chosen.count
+        guard experts.availableLayers.contains(layer) else {
+            throw Failure.unknownLayer(layer, available: experts.availableLayers.count)
+        }
+        for expert in chosen where expert < 0 || expert >= RepackedExperts.expertsPerLayer {
+            throw Failure.expertOutOfRange(expert)
+        }
+
+        // A larger prefill batch needs a larger *possible* expert union, but reserving
+        // T*K blocks eagerly makes the worst case permanent. Grow to the union we have
+        // actually observed and reuse that high-water mark on later layers/turns.
+        while buffers.count < count {
+            buffers.append(UnsafeMutableRawPointer.allocate(
+                byteCount: RepackedExperts.blockBytes,
+                alignment: RepackedExperts.pageSize))
+        }
+
+        // Read each expert block into reusable aligned buffers, then stack its tensors.
+        let sources = Array(buffers.prefix(count))
+        let missing = Array(0 ..< count)
+
+        var failures = [Bool](repeating: false, count: missing.count)
+        let buffers = self.buffers
+        let experts = self.experts
+        let read = {
+            failures.withUnsafeMutableBufferPointer { flags in
+                let flagged = flags
+                let slots = missing
+                DispatchQueue.concurrentPerform(iterations: slots.count) { index in
+                    let slot = slots[index]
+                    flagged[index] = !experts.read(
+                        layer: layer, expert: chosen[slot], into: buffers[slot])
+                }
+            }
+        }
+        if let profile { profile.measure(.expertRead, read) } else { read() }
+        profile?.recordRead(bytes: missing.count * RepackedExperts.blockBytes,
+                            issued: missing.count)
+        if let index = failures.firstIndex(of: true) {
+            throw Failure.readFailed(layer: layer, expert: chosen[missing[index]])
+        }
+
+        // `stack` is where `MLXArray(raw:)` runs, and that is `malloc` + `std::copy`
+        // in mlx-c — eager, so timing it needs no forced evaluation.
+        var pool: [MLXArray] = []
+        pool.reserveCapacity(9)
+        let shapes = self.shapes
+        let build: () -> Void = {
+        for (projection, parts) in [
+            (0, (RepackedExperts.Part.gateWeight, RepackedExperts.Part.gateScales,
+                 RepackedExperts.Part.gateBiases)),
+            (1, (.upWeight, .upScales, .upBiases)),
+            (2, (.downWeight, .downScales, .downBiases)),
+        ] {
+            let shape = shapes[projection]
+            pool.append(self.stack(parts.0, shape.weight, bf16: false, from: sources))
+            pool.append(self.stack(parts.1, shape.scale, bf16: true, from: sources))
+            pool.append(self.stack(parts.2, shape.scale, bf16: true, from: sources))
+        }
+        }
+        if let profile { profile.measure(.expertStack, build) } else { build() }
+        return pool
+    }
+
+    /// Every expert in the layer, stacked as `[256, …]` in expert-id order.
+    ///
+    /// The file is one sequential 432 MiB read. Parts are then packed into the
+    /// nine tensors `gatherQuantizedMM` already understands, so a prefill chunk
+    /// can gather by expert id without a per-expert `pread` or a CPU copy of
+    /// the index list. The file buffer is released before return; MLX has
+    /// copied the nine tensors. Callers must `eval` and drop those tensors
+    /// before the next layer, or forty of them become 17 GB.
+    func loadFullLayer(layer: Int) throws -> [MLXArray] {
+        guard experts.availableLayers.contains(layer) else {
+            throw Failure.unknownLayer(layer, available: experts.availableLayers.count)
+        }
+        let file = UnsafeMutableRawPointer.allocate(
+            byteCount: RepackedExperts.layerBytes, alignment: RepackedExperts.pageSize)
+        defer { file.deallocate() }
+
+        let read: () -> Bool = { self.experts.readLayer(layer, into: file) }
+        let ok = profile?.measure(.expertRead, read) ?? read()
+        guard ok else { throw Failure.readFailed(layer: layer, expert: -1) }
+        profile?.recordRead(bytes: RepackedExperts.layerBytes, issued: 1)
+
+        let shapes = self.shapes
+        let parts: [(RepackedExperts.Part, RepackedExperts.Part, RepackedExperts.Part)] = [
+            (.gateWeight, .gateScales, .gateBiases),
+            (.upWeight, .upScales, .upBiases),
+            (.downWeight, .downScales, .downBiases),
+        ]
+        var pool: [MLXArray] = []
+        pool.reserveCapacity(9)
+        let build: () -> Void = {
+            for (projection, triple) in parts.enumerated() {
+                let shape = shapes[projection]
+                for (part, bf16, partShape) in [
+                    (triple.0, false, shape.weight),
+                    (triple.1, true, shape.scale),
+                    (triple.2, true, shape.scale),
+                ] {
+                    pool.append(self.stackedPart(
+                        part, partShape, bf16: bf16, file: file))
+                }
+            }
+        }
+        if let profile { profile.measure(.expertStack, build) } else { build() }
+        return pool
+    }
+
+    /// Expert id `e` lives at byte `e * block` in the layer file. Gather one
+    /// part from all 256 blocks into a contiguous `[256, …]` tensor.
+    private func stackedPart(_ part: RepackedExperts.Part, _ shape: [Int],
+                             bf16: Bool, file: UnsafeMutableRawPointer) -> MLXArray {
+        let count = RepackedExperts.expertsPerLayer
+        let bytes = part.byteCount
+        let packed = UnsafeMutableRawPointer.allocate(
+            byteCount: count * bytes, alignment: 64)
+        defer { packed.deallocate() }
+        for expert in 0 ..< count {
+            packed.advanced(by: expert * bytes).copyMemory(
+                from: file.advanced(by: expert * RepackedExperts.blockBytes + part.offset),
+                byteCount: bytes)
+        }
+        let raw = UnsafeRawBufferPointer(start: packed, count: count * bytes)
+        let stackedShape = [count] + shape
+        return bf16
+            ? MLXArray(raw, stackedShape, type: UInt16.self).view(dtype: .bfloat16)
+            : MLXArray(raw, stackedShape, type: UInt32.self)
+    }
+
+    /// One part across all K slots, stacked along a new leading axis.
+    ///
+    /// bf16 is reached by reinterpreting the bits, never by converting values. An early
+    /// version widened each half in a Swift loop — roughly 200,000 iterations per layer
+    /// — and that single artifact produced the spike's original 0.5 tok/s figure.
+    private func stack(_ part: RepackedExperts.Part, _ shape: [Int], bf16: Bool,
+                       from sources: [UnsafeMutableRawPointer]) -> MLXArray
+    {
+        // Sources are the staging buffers `load` has just filled. `MLXArray(raw:)`
+        // copies eagerly, so they are free to be overwritten the moment this returns.
+        let slices = sources.map { buffer -> MLXArray in
+            let raw = UnsafeRawBufferPointer(
+                start: buffer.advanced(by: part.offset), count: part.byteCount)
+            return bf16
+                ? MLXArray(raw, shape, type: UInt16.self).view(dtype: .bfloat16)
+                : MLXArray(raw, shape, type: UInt32.self)
+        }
+        return slices.count == 1
+            ? slices[0].expandedDimensions(axis: 0)
+            : MLX.stacked(slices, axis: 0)
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/PrefixCache.swift b/ios/Sources/Edge0MLX/Edge035B/PrefixCache.swift
new file mode 100644
index 0000000..25b980e
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/PrefixCache.swift
@@ -0,0 +1,192 @@
+import Foundation
+import MLX
+
+/// Decode state captured after a fixed run of prompt tokens, so those tokens are never
+/// processed twice.
+///
+/// **This is the single largest win available on time-to-first-character.** A measured
+/// prefill runs at 258 ms/token, and of a 96-token prompt the great majority is the
+/// system instructions — byte-identical on every request, in every conversation, for the
+/// life of the install. Recomputing it is the bulk of a ~25 second wait.
+///
+/// ## Why it is unusually cheap here
+///
+/// Measured on device: **63 MiB**, against a 2.31 GB peak and a ~3.5 GB jetsam line.
+///
+/// Sixty of those MiB are the thirty DeltaNet layers' recurrent state, whose size **does
+/// not depend on how many tokens produced it** — caching a 96-token prefix and a
+/// 10,000-token prefix cost the same. Only the ten full-attention layers contribute a
+/// per-token term, at 2 KiB each.
+///
+/// The property that makes those thirty layers awkward elsewhere — the state is a fixed
+/// matrix updated in place, so dropping history does not shrink it and `replaceHistory`
+/// cannot prune cheaply — is exactly the property that makes them nearly free to cache.
+/// A plain forty-layer transformer of these dimensions would need 80 KiB per token.
+///
+/// ## Why the prefix is verified rather than assumed
+///
+/// The obvious implementation trusts that encoding the system message alone yields the
+/// same tokens as encoding it as the head of the full prompt. Byte-level BPE gives no
+/// such guarantee in general: merges are free to cross a boundary the renderer considers
+/// structural. So the cached token run is stored alongside the state and checked against
+/// the real prompt on every use. A mismatch costs an array comparison and falls back to
+/// prefilling normally — it can never produce a wrong answer.
+///
+/// This also makes invalidation automatic. Editing the system instructions changes the
+/// tokens, the stored run stops matching, and the stale file is simply ignored.
+struct PrefixCache {
+
+    /// The exact tokens this state was produced by.
+    let tokens: [Int32]
+    let state: Edge0Model35B.State
+
+    /// Whether `ids` begins with this cache's tokens, and so can skip them.
+    func matches(_ ids: [Int32]) -> Bool {
+        // An empty run is rejected explicitly. `prefix(0)` equals `[]` for every input,
+        // so without this a cache holding nothing would claim to match everything —
+        // currently unreachable, since neither capture nor `read` can produce one, but
+        // it is the kind of vacuous truth that stops being unreachable quietly.
+        guard !tokens.isEmpty else { return false }
+
+        // A prompt *equal* to the prefix is rejected: replaying the last token is what
+        // produces the logits generation starts from, so there has to be one left.
+        return ids.count > tokens.count && Array(ids.prefix(tokens.count)) == tokens
+    }
+
+    // MARK: - Persistence
+
+    /// Where the snapshot lives between launches.
+    ///
+    /// **Persisting matters more than caching in memory.** An in-process cache only helps
+    /// the second message onward; the first one after launch — the one that decides
+    /// whether the app feels usable at all — still pays in full.
+    ///
+    /// **Deliberately not beside the weights.** `Documents/repacked` is created by
+    /// `devicectl` when the 19 GB is pushed, and the app cannot write into it:
+    ///
+    ///     free=6063MiB dirExists=true
+    ///     smallWrite=Code=513 "You don't have permission to save the file"
+    ///
+    /// Six gigabytes free and a directory that plainly exists, so this is ownership, not
+    /// space. The same cause explains a puzzle from the transfer work — partial files that
+    /// **neither** `devicectl` **nor** the app could delete. Anything the app needs to
+    /// write belongs somewhere the app made.
+    ///
+    /// Application Support rather than Caches: this is derived and rebuildable, but
+    /// rebuilding costs twenty-five seconds of the user's attention, which is not what
+    /// Caches is for.
+    /// - Parameter fingerprint: identifies the model configuration. **Required, because
+    ///   `matches` compares tokens and nothing else.** A snapshot taken by the full
+    ///   forty-layer model is meaningless to a pruned one — same prompt, different
+    ///   function — and restoring it would produce confident nonsense with no error
+    ///   anywhere. Separate files mean a configuration change simply misses the cache.
+    static func fileURL(in directory: URL, fingerprint: String) -> URL {
+        let manager = FileManager.default
+        guard let support = manager.urls(
+            for: .applicationSupportDirectory, in: .userDomainMask).first
+        else {
+            return directory.appendingPathComponent(
+                "prefix-cache-\(fingerprint).safetensors")
+        }
+        try? manager.createDirectory(at: support, withIntermediateDirectories: true)
+        return support.appendingPathComponent("prefix-cache-\(fingerprint).safetensors")
+    }
+
+    /// Flattened for safetensors, which stores a flat name→array map.
+    ///
+    /// The token run is written as an array too, so the file is self-describing: nothing
+    /// outside it needs to stay in sync for the verification above to work.
+    private func flattened() -> [String: MLXArray] {
+        // **Everything is forced contiguous before it is written.** safetensors stores a
+        // plain row-major buffer and cannot represent a strided view, so saving one
+        // fails. The convolution state is exactly that — `state.convolution[layer]` is a
+        // tail slice of a larger array — and the failure was silent: the write was
+        // wrapped in `try?`, so the only symptom was a cache that worked within a session
+        // and never survived a launch.
+        var out: [String: MLXArray] = ["tokens": MLXArray(tokens)]
+        for (layer, value) in state.convolution { out["conv.\(layer)"] = contiguous(value) }
+        for (layer, value) in state.recurrent { out["rec.\(layer)"] = contiguous(value) }
+        for (layer, value) in state.keys { out["k.\(layer)"] = contiguous(value) }
+        for (layer, value) in state.values { out["v.\(layer)"] = contiguous(value) }
+        return out
+    }
+
+    /// Write the snapshot. Cache persistence is best-effort; a missing cache affects
+    /// startup speed but not generation correctness.
+    func write(to url: URL) {
+        do {
+            let arrays = flattened()
+            eval(Array(arrays.values))
+            try save(arrays: arrays, url: url)
+        } catch {
+            let directory = url.deletingLastPathComponent()
+            let free = (try? directory.resourceValues(
+                forKeys: [.volumeAvailableCapacityForImportantUsageKey]))?
+                .volumeAvailableCapacityForImportantUsage
+            let probe = directory.appendingPathComponent(".write-probe")
+            var probeResult = "ok"
+            do {
+                try Data([0]).write(to: probe)
+                try? FileManager.default.removeItem(at: probe)
+            } catch {
+                probeResult = "\(error)"
+            }
+            print("""
+                [prefix-cache] write failed: \(error)
+                [prefix-cache] free=\(free.map { "\($0 / 1_048_576)MiB" } ?? "unknown") \
+                dirExists=\(FileManager.default.fileExists(atPath: directory.path)) \
+                smallWrite=\(probeResult)
+                """)
+        }
+    }
+
+    /// Read a snapshot back, or `nil` if there is nothing usable there.
+    ///
+    /// Every failure mode — absent file, truncated write, a format from an older build —
+    /// resolves to `nil` and a normal prefill. A cache is an optimisation; it is never
+    /// worth failing a request over.
+    static func read(from url: URL) -> PrefixCache? {
+        guard FileManager.default.fileExists(atPath: url.path),
+              let arrays = try? loadArrays(url: url),
+              let tokenArray = arrays["tokens"]
+        else { return nil }
+
+        let state = Edge0Model35B.State()
+        for (name, value) in arrays where name != "tokens" {
+            let parts = name.split(separator: ".")
+            guard parts.count == 2, let layer = Int(parts[1]) else { continue }
+            switch parts[0] {
+            case "conv": state.convolution[layer] = value
+            case "rec": state.recurrent[layer] = value
+            case "k": state.keys[layer] = value
+            case "v": state.values[layer] = value
+            default: break
+            }
+        }
+        let tokens = tokenArray.asArray(Int32.self)
+        guard !tokens.isEmpty else { return nil }
+        state.offset = tokens.count
+        return PrefixCache(tokens: tokens, state: state)
+    }
+}
+
+extension Edge0Model35B.State {
+
+    /// A copy that generation can advance without disturbing the original.
+    ///
+    /// The four tables are Swift dictionaries — value types — so copying them is a real
+    /// snapshot, and the `MLXArray`s they hold are safe to share because **every write in
+    /// the model reassigns rather than mutating** (`state.recurrent[layer] = ...`, all
+    /// four verified at their call sites). If that ever changes to an in-place update,
+    /// this becomes silent corruption rather than a compile error, which is why it is
+    /// written down here.
+    func snapshot() -> Edge0Model35B.State {
+        let copy = Edge0Model35B.State()
+        copy.convolution = convolution
+        copy.recurrent = recurrent
+        copy.keys = keys
+        copy.values = values
+        copy.offset = offset
+        return copy
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/PregateHeads.swift b/ios/Sources/Edge0MLX/Edge035B/PregateHeads.swift
new file mode 100644
index 0000000..388847c
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/PregateHeads.swift
@@ -0,0 +1,220 @@
+import Foundation
+import MLX
+
+/// Trained routing heads that predict expert selections for upcoming decode steps.
+///
+/// During decode, these heads replace the router for their configured consumer layers;
+/// the model's real gate is not evaluated for those layers. The heads and their LoRA
+/// adapter are a trained pair and must be used together. Layers without a prediction use
+/// the normal router.
+///
+/// A head on layer N predicts routing for layer N+1 on the next token. This layer and
+/// token offset allows expert reads to be prefetched before the corresponding forward
+/// pass. Features contain the current hidden state and one-hot encodings of the current
+/// and previous executed routes.
+final class PregateHeads: @unchecked Sendable {
+
+    enum Failure: LocalizedError {
+        case tensorMissing(String)
+        case wrongShape(String, expected: [Int], found: [Int])
+
+        var errorDescription: String? {
+            switch self {
+            case .tensorMissing(let name): return "pregate file has no \(name)"
+            case .wrongShape(let name, let expected, let found):
+                return "pregate \(name) is \(found), expected \(expected)"
+            }
+        }
+    }
+
+    /// Layers whose head runs. Each predicts for `owner + 1`.
+    let owners: [Int]
+    /// Which layers consume a prediction, in the order their heads appear.
+    let consumers: [Int]
+    let hidden: Int
+    let experts: Int
+
+    /// Keep the file mapped and materialize arrays only during a pregate forward pass.
+    private let shard: SafetensorsShard
+    private let middle: Int
+
+    /// Where the file lives, beside the weights. Optional by construction: a build
+    /// without it, the normal router is used.
+    static func fileURL(in directory: URL, round6: Bool = false) -> URL {
+        if round6 {
+            return directory.appendingPathComponent("pregate-round6.safetensors")
+        }
+        let official = directory.appendingPathComponent("prerouter-stacked.safetensors")
+        if FileManager.default.fileExists(atPath: official.path) { return official }
+        return directory.appendingPathComponent("pregate-round6.safetensors")
+    }
+
+    init(url: URL, hidden: Int, experts: Int) throws {
+        let shard = try SafetensorsShard(url: url)
+        self.shard = shard
+        self.hidden = hidden
+        self.experts = experts
+        let inputWidth = hidden + 2 * experts
+
+        // `nil` in the expected shape means "whatever the file says" — the hidden width
+        // of the heads is the file's to decide, and everything else is pinned by the
+        // model it has to plug into.
+        func validate(_ name: String, expecting: [Int?]) throws -> SafetensorsShard.Entry {
+            guard let entry = shard.entries[name] else {
+                throw Failure.tensorMissing(name)
+            }
+            let matches = entry.shape.count == expecting.count
+                && zip(entry.shape, expecting).allSatisfy { $1 == nil || $0 == $1 }
+            guard matches else {
+                throw Failure.wrongShape(
+                    name, expected: expecting.map { $0 ?? -1 }, found: entry.shape)
+            }
+            guard entry.dtype == "F16" else {
+                throw Failure.wrongShape(name, expected: expecting.map { $0 ?? -1 },
+                                         found: entry.shape)
+            }
+            return entry
+        }
+
+        guard let ownersEntry = shard.entries["pregate.owners"],
+              let ownersBase = shard.pointer(to: "pregate.owners")
+        else { throw Failure.tensorMissing("pregate.owners") }
+        let count = ownersEntry.shape.first ?? 0
+        owners = (0 ..< count).map {
+            Int(ownersBase.load(fromByteOffset: $0 * 4, as: Int32.self))
+        }
+        consumers = owners.map { $0 + 1 }
+
+        let fc1 = try validate("pregate.fc1", expecting: [count, inputWidth, nil])
+        middle = fc1.shape[2]
+        _ = try validate("pregate.fc2", expecting: [count, middle, experts])
+        _ = try validate("pregate.linear_init", expecting: [count, inputWidth, experts])
+    }
+
+    private func array(_ name: String) -> MLXArray {
+        guard let entry = shard.entries[name], let base = shard.pointer(to: name) else {
+            preconditionFailure("validated pregate tensor disappeared: \(name)")
+        }
+        return MLXArray(
+            UnsafeRawBufferPointer(start: base, count: entry.byteCount),
+            entry.shape, type: Float16.self)
+    }
+
+    /// One-hot rows for every owner layer at once, built on the CPU.
+    ///
+    /// **Deliberately not `identity[indices].sum(axis: 0)` per layer.** That is the
+    /// obvious phrasing and it was measured: thirty-three lazy MLX nodes, each pinning
+    /// its inputs until the step boundary evaluated them, contributed to a footprint
+    /// that rose from ~600 MiB to 1535 MiB. On this device that is not merely untidy —
+    /// the expert blocks live in *clean file-backed* pages, so anything the process
+    /// holds dirty is taken out of the page cache those reads depend on. `read` doubled
+    /// to 134 ms and the weight copy rate halved, and together they ate a 70 ms saving.
+    ///
+    /// The routing is already on the CPU — `moe` has just called `asArray` on it — so
+    /// the whole `[n, experts]` block is a memset and `n * k` stores, then one array.
+    /// No graph, nothing retained, nothing to evaluate.
+    func oneHotRows(_ perLayer: [[Int32]]) -> MLXArray {
+        var flat = [Float16](repeating: 0, count: perLayer.count * experts)
+        for (row, chosen) in perLayer.enumerated() {
+            for expert in chosen where expert >= 0 && Int(expert) < experts {
+                // Accumulate so duplicate IDs retain their multiplicity.
+                flat[row * experts + Int(expert)] += 1
+            }
+        }
+        return MLXArray(flat, [perLayer.count, experts])
+    }
+
+    /// Run every head as one batch.
+    ///
+    /// - Parameters:
+    ///   - inputs: `[n, hidden]`, the owner layers' MoE inputs this token, in owner order.
+    ///   - executed: `[n, experts]`, the one-hot rows those layers actually ran, from
+    ///     `oneHotRows`.
+    ///   - previous: `[n, experts]`, the previous token's routing or a zero row.
+    /// - Returns: `[n, experts]` float32 logits for the consuming layers.
+    func logits(inputs: MLXArray, executed: MLXArray, previous: MLXArray) -> MLXArray {
+        let fc1 = array("pregate.fc1")
+        let fc2 = array("pregate.fc2")
+        let linearInit = array("pregate.linear_init")
+        let features = concatenated(
+            [inputs.asType(.float16), executed, previous], axis: -1)
+            .expandedDimensions(axis: 1)                          // [n, 1, in]
+        let inner = matmul(features, fc1)                         // [n, 1, mid]
+        // Exact erf GELU, as trained (torch's default). `geluApproximate` would be a
+        // different function evaluated on 33 × 512 values a step, and there is no
+        // reference to check it against; the tanh form is not what these weights saw.
+        let activated = 0.5 * inner * (1 + erf(inner / Float(2).squareRoot()))
+        let out = matmul(features, linearInit) + matmul(activated, fc2)
+        return out.squeezed(axis: 1).asType(.float32)             // [n, experts]
+    }
+
+    /// What has to survive from one decode step to the next.
+    ///
+    /// Per-conversation routing history used by the trained heads.
+    final class State {
+        /// Owner layer to its MoE input for the current token. Stored as an MLX array so
+        /// conversion to CPU does not introduce an additional synchronization point.
+        var moeInput: [Int: MLXArray] = [:]
+        /// Owner layer → the experts it actually ran this token.
+        var executed: [Int: [Int32]] = [:]
+        /// The previous token's executed experts, or absent on the first token.
+        var previous: [Int: [Int32]] = [:]
+        /// Consuming layer → the experts it will use next step, and their weights.
+        ///
+        /// Both plain Swift. They cross a step boundary, and an `MLXArray` that does
+        /// that carries whatever graph produced it into the next step.
+        var predicted: [Int: [Int]] = [:]
+        var scores: [Int: [Float]] = [:]
+
+        /// Advance routing history at the end of a decode step.
+        func advance() {
+            previous = executed
+            executed = [:]
+            moeInput = [:]
+        }
+
+        func reset() {
+            moeInput = [:]; executed = [:]; previous = [:]
+            predicted = [:]; scores = [:]
+        }
+    }
+
+    // MARK: - Optional numerical self-check
+
+    /// Compare the heads' output against an optional fixture containing inputs, expected
+    /// logits, and top-2 expert IDs. Returns `nil` when the fixture is unavailable.
+    func selfCheck(fixture url: URL) -> String? {
+        guard let shard = try? SafetensorsShard(url: url),
+              let feats = shard.entries["feats"], let featsBase = shard.pointer(to: "feats"),
+              let ref = shard.entries["logits"], let refBase = shard.pointer(to: "logits"),
+              let top = shard.entries["top2"], let topBase = shard.pointer(to: "top2")
+        else { return nil }
+
+        let features = MLXArray(
+            UnsafeRawBufferPointer(start: featsBase, count: feats.byteCount),
+            feats.shape, type: Float16.self)
+        let expected = MLXArray(
+            UnsafeRawBufferPointer(start: refBase, count: ref.byteCount),
+            ref.shape, type: Float32.self)
+        let expectedTop = MLXArray(
+            UnsafeRawBufferPointer(start: topBase, count: top.byteCount),
+            top.shape, type: Int32.self)
+
+        // The fixture holds the concatenated 2560-wide features; split them back so the
+        // public entry point is what gets exercised, not a shortcut around it.
+        let n = feats.shape[0]
+        let inputs = features[0..., 0 ..< hidden]
+        let executed = features[0..., hidden ..< (hidden + experts)]
+        let previous = features[0..., (hidden + experts)...]
+        let got = logits(inputs: inputs, executed: executed, previous: previous)
+
+        let difference = abs(got - expected)
+        let gotTop = argSort(-got, axis: -1)[0..., 0 ..< 2]
+        let agree = (gotTop .== expectedTop).all(axis: -1).sum()
+        eval(difference, agree)
+        return String(
+            format: "[pregate] self-check: max|Δ|=%.4f mean|Δ|=%.5f top-2 agree %d/%d",
+            difference.max().item(Float.self), difference.mean().item(Float.self),
+            agree.item(Int32.self), Int32(n))
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/RepackedExperts.swift b/ios/Sources/Edge0MLX/Edge035B/RepackedExperts.swift
new file mode 100644
index 0000000..a42de63
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/RepackedExperts.swift
@@ -0,0 +1,162 @@
+import Foundation
+
+/// The repacked layout: one file per layer, one expert per contiguous block.
+///
+/// Each block holds an expert's nine tensors back to back, 1,769,472 bytes — exactly
+/// 108 pages of 16 KB — so with data starting at offset 0 of its layer's file, **every
+/// block is page-aligned**. Reading one expert is one sequential 1.69 MB read instead
+/// of the nine scattered reads the original safetensors layout forces.
+///
+/// Each file is checked against the expected length before it is opened so truncated
+/// or incomplete files cannot be used as model weights.
+final class RepackedExperts {
+
+    /// Byte layout inside one expert's block, in the order the packer wrote them.
+    enum Part: Int, CaseIterable {
+        case gateWeight, gateScales, gateBiases
+        case upWeight, upScales, upBiases
+        case downWeight, downScales, downBiases
+
+        var byteCount: Int {
+            switch self {
+            case .gateWeight, .upWeight, .downWeight: return 524_288
+            default: return 32_768
+            }
+        }
+
+        var offset: Int {
+            Part.allCases.prefix(rawValue).reduce(0) { $0 + $1.byteCount }
+        }
+    }
+
+    static let blockBytes = 1_769_472
+    static let expertsPerLayer = 256
+    static let layerBytes = blockBytes * expertsPerLayer   // 452,984,832
+    static let pageSize = 16384                            // blockBytes is 108 of these
+
+    private struct Opened {
+        let size: Int
+        let descriptor: Int32
+    }
+
+    // Keep descriptors, not 40 whole-file mappings. Mapping every 432 MiB layer
+    // consumes ~17 GiB of virtual address space and fails after 11 layers when a
+    // development profile lacks Extended Virtual Addressing. The hot path already
+    // uses pread, so those mappings were only an optional prefetch mechanism.
+    private var layers: [Int: Opened] = [:]
+
+    /// Layers whose file was present but the wrong length, with the length found.
+    private(set) var rejected: [(layer: Int, bytes: Int)] = []
+    private(set) var missing: [Int] = []
+
+    var availableLayers: [Int] { layers.keys.sorted() }
+
+    init(directory: URL, expectedLayers: Int) {
+        for layer in 0 ..< expectedLayers {
+            let url = directory.appendingPathComponent(String(format: "experts-L%02d.bin", layer))
+
+            guard let attributes = try? FileManager.default.attributesOfItem(atPath: url.path),
+                  let size = (attributes[.size] as? NSNumber)?.intValue
+            else {
+                missing.append(layer)
+                continue
+            }
+            guard size == Self.layerBytes else {
+                rejected.append((layer, size))
+                continue
+            }
+
+            let descriptor = open(url.path, O_RDONLY)
+            guard descriptor >= 0 else {
+                missing.append(layer)
+                continue
+            }
+            layers[layer] = Opened(size: size, descriptor: descriptor)
+        }
+    }
+
+    deinit {
+        for opened in layers.values {
+            close(opened.descriptor)
+        }
+    }
+
+    /// Ask the kernel for an expert's whole block in one go.
+    func prefetch(layer: Int, expert: Int) {
+        guard let opened = layers[layer], expert >= 0, expert < Self.expertsPerLayer
+        else { return }
+
+        // `F_RDADVISE` is the descriptor equivalent of the old mmap + MADV_WILLNEED
+        // path: it starts asynchronous readahead without reserving 17 GiB of virtual
+        // address space. The eventual `pread` remains the source of truth, so ignored
+        // or failed advice can affect only latency, never model output.
+        var advice = radvisory(
+            ra_offset: off_t(expert * Self.blockBytes),
+            ra_count: Int32(Self.blockBytes))
+        _ = withUnsafeMutablePointer(to: &advice) {
+            fcntl(opened.descriptor, F_RDADVISE, UnsafeMutableRawPointer($0))
+        }
+    }
+
+    /// Remove files that are present but have the wrong length, returning any failures.
+    static func removeTruncated(in directory: URL, expectedLayers: Int)
+        -> (removed: [(Int, Int)], failed: [(Int, String)])
+    {
+        var removed: [(Int, Int)] = []
+        var failed: [(Int, String)] = []
+        for layer in 0 ..< expectedLayers {
+            let url = directory.appendingPathComponent(String(format: "experts-L%02d.bin", layer))
+            guard let attributes = try? FileManager.default.attributesOfItem(atPath: url.path),
+                  let size = (attributes[.size] as? NSNumber)?.intValue,
+                  size != layerBytes
+            else { continue }
+            do {
+                try FileManager.default.removeItem(at: url)
+                removed.append((layer, size))
+            } catch {
+                failed.append((layer, error.localizedDescription))
+            }
+        }
+        return (removed, failed)
+    }
+
+    /// Read one expert's whole block with a single `pread`.
+    ///
+    /// The mmap path reaches the same bytes through page faults, which the kernel
+    /// services in its own units; this asks for all 1,769,472 bytes in one call and
+    /// lets it issue one large sequential read. Worth measuring against the fault path
+    /// now that a block is contiguous — 768 MB/s is still only about a third of what
+    /// the device can do.
+    /// One sequential read of the whole layer file (432 MiB, 256 expert blocks).
+    func readLayer(_ layer: Int, into buffer: UnsafeMutableRawPointer) -> Bool {
+        guard let opened = layers[layer] else { return false }
+        var moved = 0
+        while moved < Self.layerBytes {
+            let got = pread(
+                opened.descriptor,
+                buffer.advanced(by: moved),
+                Self.layerBytes - moved,
+                off_t(moved))
+            if got <= 0 { return false }
+            moved += got
+        }
+        return true
+    }
+
+    func read(layer: Int, expert: Int, into buffer: UnsafeMutableRawPointer) -> Bool {
+        guard let opened = layers[layer], expert >= 0, expert < Self.expertsPerLayer
+        else { return false }
+
+        var moved = 0
+        while moved < Self.blockBytes {
+            let got = pread(
+                opened.descriptor,
+                buffer.advanced(by: moved),
+                Self.blockBytes - moved,
+                off_t(expert * Self.blockBytes + moved))
+            if got <= 0 { return false }
+            moved += got
+        }
+        return true
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/Safetensors.swift b/ios/Sources/Edge0MLX/Edge035B/Safetensors.swift
new file mode 100644
index 0000000..c48b554
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/Safetensors.swift
@@ -0,0 +1,264 @@
+import Foundation
+
+/// Read-only mmap over one safetensors shard, plus the header that says where each
+/// tensor's bytes begin.
+///
+/// A shard is: 8 bytes of little-endian header length, that many bytes of JSON, then
+/// the payload. Every tensor's `data_offsets` are relative to the start of the
+/// payload, so a tensor — or, for a stacked expert tensor, one expert's slice of it
+/// — is a byte range and nothing more. That is the premise the whole streaming
+/// design rests on: **no parsing, no dequantisation, no copy on the read path.**
+///
+/// Mapped `MAP_SHARED` + read-only so the pages stay clean and file-backed. Dirtying
+/// them would turn the 19 GB into the app's own memory and defeat the exercise.
+final class SafetensorsShard {
+
+    struct Entry {
+        let dtype: String
+        let shape: [Int]
+        /// Absolute offset in the file, payload base already added.
+        let offset: Int
+        let byteCount: Int
+    }
+
+    let url: URL
+    private let descriptor: Int32
+    private let base: UnsafeRawPointer
+    private let mappedSize: Int
+    private(set) var entries: [String: Entry] = [:]
+
+    enum Failure: LocalizedError {
+        case cannotOpen(String)
+        case cannotMap(String)
+        case badHeader(String)
+        case truncated(String, needs: Int, has: Int)
+
+        var errorDescription: String? {
+            switch self {
+            case .cannotOpen(let p): return "cannot open \(p)"
+            case .cannotMap(let p): return "cannot mmap \(p)"
+            case .badHeader(let p): return "bad safetensors header in \(p)"
+            case .truncated(let p, let needs, let has):
+                return "\(p) is truncated: header describes \(needs) bytes, file has \(has)"
+            }
+        }
+    }
+
+    init(url: URL) throws {
+        self.url = url
+
+        descriptor = open(url.path, O_RDONLY)
+        guard descriptor >= 0 else { throw Failure.cannotOpen(url.lastPathComponent) }
+
+        var status = stat()
+        guard fstat(descriptor, &status) == 0, status.st_size > 16 else {
+            close(descriptor)
+            throw Failure.cannotOpen(url.lastPathComponent)
+        }
+        mappedSize = Int(status.st_size)
+
+        guard let mapped = mmap(nil, mappedSize, PROT_READ, MAP_SHARED, descriptor, 0),
+              mapped != MAP_FAILED
+        else {
+            close(descriptor)
+            throw Failure.cannotMap(url.lastPathComponent)
+        }
+        base = UnsafeRawPointer(mapped)
+
+        // Read advice is applied at the call site because access patterns vary by tensor.
+
+        let headerLength = Int(base.loadUnaligned(as: UInt64.self))
+        guard headerLength > 0, headerLength + 8 <= mappedSize else {
+            munmap(mapped, mappedSize)
+            close(descriptor)
+            throw Failure.badHeader(url.lastPathComponent)
+        }
+
+        let payloadBase = 8 + headerLength
+        let headerData = Data(bytes: base.advanced(by: 8), count: headerLength)
+
+        guard let json = try JSONSerialization.jsonObject(with: headerData) as? [String: Any]
+        else {
+            munmap(mapped, mappedSize)
+            close(descriptor)
+            throw Failure.badHeader(url.lastPathComponent)
+        }
+
+        for (name, value) in json {
+            guard name != "__metadata__",
+                  let fields = value as? [String: Any],
+                  let dtype = fields["dtype"] as? String,
+                  let shape = fields["shape"] as? [Int],
+                  let offsets = fields["data_offsets"] as? [Int],
+                  offsets.count == 2
+            else { continue }
+
+            entries[name] = Entry(
+                dtype: dtype,
+                shape: shape,
+                offset: payloadBase + offsets[0],
+                byteCount: offsets[1] - offsets[0]
+            )
+        }
+
+        // Does the header describe more bytes than the file actually holds?
+        //
+        // The device-side size check in the transfer script compares against what
+        // `devicectl` reports, which is three significant figures — enough to catch a
+        // file that arrived at 133 MB of 432, useless for catching one that is short by
+        // a few megabytes. This catches that case, cheaply, at the only point where it
+        // can still be reported as a load failure rather than as quiet nonsense several
+        // layers into inference.
+        if let end = entries.values.map({ $0.offset + $0.byteCount }).max(), end > mappedSize {
+            munmap(mapped, mappedSize)
+            close(descriptor)
+            throw Failure.truncated(url.lastPathComponent, needs: end, has: mappedSize)
+        }
+    }
+
+    deinit {
+        munmap(UnsafeMutableRawPointer(mutating: base), mappedSize)
+        close(descriptor)
+    }
+
+    /// Raw pointer to a whole tensor. Does not touch the pages.
+    func pointer(to name: String) -> UnsafeRawPointer? {
+        guard let entry = entries[name] else { return nil }
+        return base.advanced(by: entry.offset)
+    }
+
+    /// Raw pointer to one slice along the leading axis of a stacked tensor.
+    ///
+    /// Expert weights are stored as `[num_experts, out, in]`, so expert *e* is one
+    /// contiguous run — `rowBytes` apart from its neighbour. This single line is why
+    /// streaming is possible at all: pulling one expert is an offset, not a load.
+    func pointer(to name: String, slice index: Int, of leadingDimension: Int) -> UnsafeRawPointer? {
+        guard let entry = entries[name], leadingDimension > 0 else { return nil }
+        let rowBytes = entry.byteCount / leadingDimension
+        guard index >= 0, index < leadingDimension else { return nil }
+        return base.advanced(by: entry.offset + index * rowBytes)
+    }
+
+    func sliceByteCount(of name: String, leadingDimension: Int) -> Int? {
+        guard let entry = entries[name], leadingDimension > 0 else { return nil }
+        return entry.byteCount / leadingDimension
+    }
+
+    /// Reset the kernel's read-ahead policy for the whole mapping.
+    func advise(_ advice: Int32) {
+        madvise(UnsafeMutableRawPointer(mutating: base), mappedSize, advice)
+    }
+
+    /// Ask the kernel to start pulling one expert's slice in.
+    ///
+    /// `MADV_WILLNEED` is advice, not a read: it returns once the I/O is queued rather
+    /// than once the pages are there. Issued immediately before the copy it can only
+    /// overlap within the slice; issued a step ahead — which is what pregate makes
+    /// possible — it has a whole layer of compute to land in.
+    func prefetch(_ name: String, slice index: Int, of leadingDimension: Int) {
+        guard let entry = entries[name], leadingDimension > 0,
+              index >= 0, index < leadingDimension
+        else { return }
+        let rowBytes = entry.byteCount / leadingDimension
+        madvise(
+            UnsafeMutableRawPointer(mutating: base).advanced(by: entry.offset + index * rowBytes),
+            rowBytes, MADV_WILLNEED)
+    }
+
+    /// Ask the kernel to bring one whole tensor in.
+    ///
+    /// Request that the kernel warm the tensor's clean, file-backed pages.
+    /// Returns the number of bytes requested.
+    @discardableResult
+    func warm(_ name: String) -> Int {
+        guard let entry = entries[name] else { return 0 }
+        madvise(UnsafeMutableRawPointer(mutating: base).advanced(by: entry.offset),
+                entry.byteCount, MADV_WILLNEED)
+        return entry.byteCount
+    }
+}
+
+/// The four shards plus the index that says which shard holds which tensor.
+final class SafetensorsBundle {
+
+    private var shards: [String: SafetensorsShard] = [:]
+    private var shardForTensor: [String: String] = [:]
+
+    let directory: URL
+
+    init(directory: URL) throws {
+        self.directory = directory
+
+        let indexURL = directory.appendingPathComponent("model.safetensors.index.json")
+        let indexData = try Data(contentsOf: indexURL)
+        guard let root = try JSONSerialization.jsonObject(with: indexData) as? [String: Any],
+              let map = root["weight_map"] as? [String: String]
+        else {
+            throw SafetensorsShard.Failure.badHeader("model.safetensors.index.json")
+        }
+        shardForTensor = map
+
+        for file in Set(map.values) {
+            shards[file] = try SafetensorsShard(url: directory.appendingPathComponent(file))
+        }
+    }
+
+    var mappedShardCount: Int { shards.count }
+
+    /// Every tensor name in the index, and which shard holds it.
+    var index: [String: String] { shardForTensor }
+
+    func entry(_ name: String) -> SafetensorsShard.Entry? {
+        guard let file = shardForTensor[name] else { return nil }
+        return shards[file]?.entries[name]
+    }
+
+    func pointer(to name: String) -> UnsafeRawPointer? {
+        guard let file = shardForTensor[name] else { return nil }
+        return shards[file]?.pointer(to: name)
+    }
+
+    func pointer(to name: String, slice index: Int, of leadingDimension: Int) -> UnsafeRawPointer? {
+        guard let file = shardForTensor[name] else { return nil }
+        return shards[file]?.pointer(to: name, slice: index, of: leadingDimension)
+    }
+
+    func sliceByteCount(of name: String, leadingDimension: Int) -> Int? {
+        guard let file = shardForTensor[name] else { return nil }
+        return shards[file]?.sliceByteCount(of: name, leadingDimension: leadingDimension)
+    }
+
+    func advise(_ advice: Int32) {
+        for shard in shards.values { shard.advise(advice) }
+    }
+
+    /// Bring in every tensor the decode path reads on **every** step.
+    ///
+    /// Deliberately not "the whole mapping": these shards hold the 256 routed experts per
+    /// layer as well, nineteen gigabytes of data that only two-of-256 of is touched per
+    /// layer per token. Asking for that would evict everything worth having.
+    ///
+    /// Two exclusions, both by measurement rather than taste:
+    ///
+    /// - `switch_mlp` — the routed experts, streamed on demand and prefetched from their
+    ///   own repacked files by their actual routing.
+    /// - `embed_tokens` — 273 MiB of table from which `rows(_:_:)` copies one row per
+    ///   token. Warming it would be 273 MiB of page cache for 1,152 bytes of need.
+    ///
+    /// What remains is what `weight(_:)` copies every step: 20.4 MiB per layer of
+    /// attention and shared-expert projections, plus `lm_head`.
+    @discardableResult
+    func warmResidentWeights() -> Int {
+        var total = 0
+        for (name, file) in shardForTensor
+        where !name.contains("switch_mlp") && !name.contains("embed_tokens") {
+            total += shards[file]?.warm(name) ?? 0
+        }
+        return total
+    }
+
+    func prefetch(_ name: String, slice index: Int, of leadingDimension: Int) {
+        guard let file = shardForTensor[name] else { return }
+        shards[file]?.prefetch(name, slice: index, of: leadingDimension)
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/StepProfile.swift b/ios/Sources/Edge0MLX/Edge035B/StepProfile.swift
new file mode 100644
index 0000000..19e8fa6
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/StepProfile.swift
@@ -0,0 +1,148 @@
+import Foundation
+
+/// Tracks synchronous portions of model steps without forcing additional MLX evaluation.
+/// GPU operations remain lazy, so timing them separately would change execution order.
+final class StepProfile {
+
+    private(set) var syncLinear: Double = 0
+    private(set) var syncFull: Double = 0
+    var routerSync: Double { syncLinear + syncFull }
+    private(set) var expertRead: Double = 0
+    private(set) var expertStack: Double = 0
+
+    /// Time spent constructing MLX arrays from mapped weight bytes.
+    private(set) var weightCopy: Double = 0
+    private(set) var weightBytes: Int = 0
+    private(set) var expertBytes: Int = 0
+
+    /// Number of expert read batches and individual reads issued.
+    private(set) var readBatches = 0
+    private(set) var readsIssued = 0
+
+    /// Number of unique experts requested compared with the naive T*K upper bound.
+    private(set) var unionTotal = 0
+    private(set) var pairsTotal = 0
+
+    /// How often the "same experts as last token" guess was right.
+    ///
+    /// The whole prefetch scheme rests on this. It costs nothing when wrong, so any
+    /// positive rate is upside — but it bounds the gain: the read can only be hidden for
+    /// the fraction of blocks the guess anticipated.
+    private(set) var prefetchHits = 0
+    private(set) var prefetchTotal = 0
+
+    /// Cumulative cost of the batched head forward plus its single barrier.
+    private(set) var pregateStage = 0.0
+    /// Core ML vocabulary projection plus its 4 KiB MLX-to-Core ML handoff.
+    private(set) var vocabularyHead = 0.0
+
+    private var steps = 0
+
+    func reset() {
+        syncLinear = 0; syncFull = 0; expertRead = 0; expertStack = 0; expertBytes = 0
+        weightCopy = 0; weightBytes = 0; pregateStage = 0; vocabularyHead = 0
+        prefetchHits = 0; prefetchTotal = 0
+        readBatches = 0; readsIssued = 0; steps = 0
+        unionTotal = 0; pairsTotal = 0
+    }
+
+    func countStep() { steps += 1 }
+
+    /// Timing categories for synchronous portions of a model step.
+    enum Phase {
+        case syncLinear, syncFull, expertRead, expertStack, weightCopy
+        /// Synchronous work performed by the trained routing heads.
+        case pregateStage, vocabularyHead
+    }
+
+    /// Time a synchronous block. Returns the block's value so call sites stay expressions.
+    @inline(__always)
+    func measure<T>(_ phase: Phase, _ body: () throws -> T) rethrows -> T {
+        let start = DispatchTime.now().uptimeNanoseconds
+        let value = try body()
+        let elapsed = Double(DispatchTime.now().uptimeNanoseconds - start) / 1_000_000
+        switch phase {
+        case .syncLinear: syncLinear += elapsed
+        case .syncFull: syncFull += elapsed
+        case .expertRead: expertRead += elapsed
+        case .expertStack: expertStack += elapsed
+        case .weightCopy: weightCopy += elapsed
+        case .pregateStage: pregateStage += elapsed
+        case .vocabularyHead: vocabularyHead += elapsed
+        }
+        return value
+    }
+
+    @inline(__always)
+    func measure<T>(_ phase: Phase, _ body: () async throws -> T) async rethrows -> T {
+        let start = DispatchTime.now().uptimeNanoseconds
+        let value = try await body()
+        let elapsed = Double(DispatchTime.now().uptimeNanoseconds - start) / 1_000_000
+        switch phase {
+        case .syncLinear: syncLinear += elapsed
+        case .syncFull: syncFull += elapsed
+        case .expertRead: expertRead += elapsed
+        case .expertStack: expertStack += elapsed
+        case .weightCopy: weightCopy += elapsed
+        case .pregateStage: pregateStage += elapsed
+        case .vocabularyHead: vocabularyHead += elapsed
+        }
+        return value
+    }
+
+    func recordRouting(union: Int, pairs: Int) {
+        unionTotal += union
+        pairsTotal += pairs
+    }
+
+    func recordPrefetch(hits: Int, total: Int) {
+        prefetchHits += hits
+        prefetchTotal += total
+    }
+
+    func recordWeight(bytes: Int) { weightBytes += bytes }
+
+    func recordRead(bytes: Int, issued: Int) {
+        expertBytes += bytes
+        readBatches += 1
+        readsIssued += issued
+    }
+
+    /// Per-step medians are not available without keeping every sample; per-step means
+    /// are, and the question here is where a quarter-second goes, not its distribution.
+    var report: String {
+        guard steps > 0 else { return "no steps" }
+        let n = Double(steps)
+        func ms(_ total: Double) -> String { String(format: "%.0f", total / n) }
+        let throughput = expertRead > 0
+            ? String(format: "%.0f", Double(expertBytes) / (expertRead / 1000) / 1_000_000)
+            : "-"
+        let depth = readBatches > 0
+            ? String(format: "%.2f", Double(readsIssued) / Double(readBatches))
+            : "-"
+        let copyRate = weightCopy > 0
+            ? String(format: "%.0f", Double(weightBytes) / (weightCopy / 1000) / 1_000_000)
+            : "-"
+        let hitRate = prefetchTotal > 0
+            ? String(format: "%.0f%%",
+                     Double(prefetchHits) / Double(prefetchTotal) * 100)
+            : "-"
+        let sharing = pairsTotal > 0
+            ? String(format: "%.2f", Double(unionTotal) / Double(pairsTotal))
+            : "-"
+        return "syncLinear=\(ms(syncLinear))ms syncFull=\(ms(syncFull))ms "
+            + "read=\(ms(expertRead))ms "
+            + "stack=\(ms(expertStack))ms readMBps=\(throughput) depth=\(depth) "
+            + "union/pairs=\(sharing) "
+            + "weightCopy=\(ms(weightCopy))ms wMBps=\(copyRate)"
+            // Only when the heads are installed, so a baseline build's line is
+            // byte-identical to the one every earlier measurement was read from.
+            + (pregateStage > 0 ? " pregate=\(ms(pregateStage))ms" : "")
+            + (vocabularyHead > 0 ? " coremlHead=\(ms(vocabularyHead))ms" : "")
+            // `hitRate` was computed and dropped on the floor by an earlier version of
+            // this line. It is the number that bounds the whole prefetch scheme — the
+            // read can only be hidden for the fraction of blocks the guess anticipated
+            // — so leaving it unprinted made the one decision it informs unanswerable.
+            + (prefetchTotal > 0 ? " prefetch=\(hitRate)" : "")
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/ThinkingBlock.swift b/ios/Sources/Edge0MLX/Edge035B/ThinkingBlock.swift
new file mode 100644
index 0000000..5d1f377
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/ThinkingBlock.swift
@@ -0,0 +1,57 @@
+import Foundation
+
+/// Separates the model's reasoning from its reply.
+///
+/// This checkpoint is opened with `<|im_start|>assistant\n<think>\n`, so **every answer
+/// begins inside a reasoning block** and the first thing it writes is working-out, not a
+/// reply. The closing `</think>` is generated; the opening tag never is, because it came
+/// from the prompt.
+///
+/// That shape leaks straight through a stream that treats decoded text as the answer —
+/// the user sees the model's notes to itself, followed by a bare `</think>`, followed by
+/// the reply.
+///
+/// **The cost of hiding it is that nothing appears until the model stops reasoning.** At
+/// a few tokens per second that is a real wait with an empty bubble at the end of it; the
+/// generating indicator is what carries that interval. Showing the reasoning instead was
+/// the alternative, and it is worse for the ordinary case of wanting an answer.
+enum ThinkingBlock {
+
+    static let terminator = "</think>"
+
+    /// The part of `raw` the user should see, or `nil` while the model is still reasoning.
+    ///
+    /// Leading whitespace after the terminator is dropped: the model writes a blank line
+    /// or two before the reply proper, which would otherwise open every answer with a gap.
+    static func visible(in raw: String) -> String? {
+        guard let end = raw.range(of: terminator) else { return nil }
+        return String(raw[end.upperBound...])
+            .trimmingCharacters(in: .whitespacesAndNewlines)
+    }
+
+    /// The non-thinking prompt normally starts directly in the answer, but some
+    /// generations still emit the checkpoint's literal thinking wrapper. Hide that
+    /// block while it is incomplete, then expose only the text after its terminator.
+    static func visibleReply(in raw: String) -> String? {
+        guard raw.contains("<think>") || raw.contains(terminator) else { return raw }
+        return visible(in: raw)
+    }
+
+    /// What to show once generation has finished.
+    ///
+    /// **Falls back to the whole text when the block never closed.** A reply truncated by
+    /// the token limit mid-reasoning has no terminator, and showing nothing at all would
+    /// present a completed generation as an empty answer. Raw reasoning is poor output;
+    /// an empty bubble is a bug report.
+    static func settled(_ raw: String) -> String {
+        visible(in: raw) ?? raw.trimmingCharacters(in: .whitespacesAndNewlines)
+    }
+
+    /// Final text for a prompt that requested a direct answer.
+    static func settledReply(_ raw: String) -> String {
+        guard raw.contains("<think>") || raw.contains(terminator) else {
+            return raw.trimmingCharacters(in: .whitespacesAndNewlines)
+        }
+        return visible(in: raw) ?? ""
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift b/ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift
new file mode 100644
index 0000000..8020dee
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge035B/Tokenizer.swift
@@ -0,0 +1,267 @@
+import Foundation
+
+/// Byte-level BPE, the GPT-4 variant this checkpoint uses.
+///
+/// The pipeline, in order, because getting the order wrong produces a tokenizer that
+/// works on English and mis-splits everything else:
+///
+/// 1. **NFC** normalisation
+/// 2. **Split** on the pre-tokenizer regex, `Isolated` — each match becomes its own piece
+/// 3. **ByteLevel** — every UTF-8 byte maps to one printable character, so BPE never sees
+///    a raw byte and no `unk` token is needed
+/// 4. **BPE** — repeatedly merge the adjacent pair with the lowest rank
+///
+/// Loaded from the compact binary produced by `tools/convert_tokenizer.py`: 5.6 MiB of
+/// flat arrays rather than 19 MiB of JSON, with merges stored as **token id pairs** so
+/// the merge loop never looks anything up by string.
+final class Tokenizer {
+
+    enum Failure: LocalizedError {
+        case cannotOpen(String)
+        case badHeader(String)
+        case truncated(String, needs: Int, has: Int)
+
+        var errorDescription: String? {
+            switch self {
+            case .cannotOpen(let p): return "cannot open \(p)"
+            case .badHeader(let p): return "bad tokenizer header in \(p)"
+            case .truncated(let p, let needs, let has):
+                return "\(p) is truncated: needs \(needs) bytes, has \(has)"
+            }
+        }
+    }
+
+    private let tokens: [String]                  // id -> token text (byte-level encoded)
+    private let ids: [String: Int32]              // token text -> id
+    private let ranks: [UInt64: Int32]            // packed (left, right) -> merge rank
+    private let specials: [String: Int32]
+    private let specialTexts: [Int32: String]
+    private let splitter: NSRegularExpression
+
+    /// Byte ↔ printable-character tables, the GPT-2 mapping.
+    ///
+    /// Bytes that are already printable map to themselves; the rest are shifted into an
+    /// unused Unicode range. This is what lets a BPE vocabulary of *characters* represent
+    /// arbitrary bytes without an unknown token.
+    private static let byteToUnicode: [Character] = {
+        var mapping: [Int] = []
+        var table = [Character](repeating: " ", count: 256)
+        for value in UInt8(ascii: "!") ... UInt8(ascii: "~") { mapping.append(Int(value)) }
+        for value in 0xA1 ... 0xAC { mapping.append(value) }
+        for value in 0xAE ... 0xFF { mapping.append(value) }
+
+        var extra = 0
+        for byte in 0 ..< 256 {
+            if mapping.contains(byte) {
+                table[byte] = Character(UnicodeScalar(byte)!)
+            } else {
+                table[byte] = Character(UnicodeScalar(256 + extra)!)
+                extra += 1
+            }
+        }
+        return table
+    }()
+
+    private static let unicodeToByte: [Character: UInt8] = {
+        var inverse: [Character: UInt8] = [:]
+        for byte in 0 ..< 256 { inverse[byteToUnicode[byte]] = UInt8(byte) }
+        return inverse
+    }()
+
+    init(url: URL) throws {
+        guard let data = try? Data(contentsOf: url, options: .mappedIfSafe),
+              data.count > 8
+        else { throw Failure.cannotOpen(url.lastPathComponent) }
+
+        let headerLength = Int(data.withUnsafeBytes { $0.loadUnaligned(as: UInt64.self) })
+        guard headerLength > 0, 8 + headerLength <= data.count,
+              let header = try? JSONSerialization.jsonObject(
+                with: data[8 ..< (8 + headerLength)]) as? [String: Any],
+              let vocabSize = header["vocab_size"] as? Int,
+              let mergeCount = header["merge_count"] as? Int,
+              let specialCount = header["special_count"] as? Int,
+              let pattern = header["regex"] as? String
+        else { throw Failure.badHeader(url.lastPathComponent) }
+
+        // Sections are fixed-width and laid out in a known order, so each is a slice at a
+        // running offset. Checked against the file length before any of it is read — a
+        // truncated tokenizer that silently loses its tail would encode most text
+        // correctly and fail only on rare tokens.
+        var cursor = 8 + headerLength
+        func take(_ count: Int) throws -> Data {
+            guard cursor + count <= data.count else {
+                throw Failure.truncated(url.lastPathComponent,
+                                        needs: cursor + count, has: data.count)
+            }
+            defer { cursor += count }
+            return data[cursor ..< (cursor + count)]
+        }
+
+        let offsets = try take((vocabSize + 1) * 4).toUInt32Array()
+        let blob = try take(Int(offsets[vocabSize]))
+        let mergeFlat = try take(mergeCount * 8).toUInt32Array()
+        let specialIds = try take(specialCount * 4).toUInt32Array()
+        let specialOffsets = try take((specialCount + 1) * 4).toUInt32Array()
+        let specialBlob = try take(Int(specialOffsets[specialCount]))
+
+        var tokens = [String](repeating: "", count: vocabSize)
+        var ids: [String: Int32] = [:]
+        ids.reserveCapacity(vocabSize)
+        blob.withUnsafeBytes { raw in
+            let base = raw.bindMemory(to: UInt8.self).baseAddress!
+            for index in 0 ..< vocabSize {
+                let start = Int(offsets[index]), end = Int(offsets[index + 1])
+                let text = String(decoding: UnsafeBufferPointer(
+                    start: base + start, count: end - start), as: UTF8.self)
+                tokens[index] = text
+                ids[text] = Int32(index)
+            }
+        }
+        self.tokens = tokens
+        self.ids = ids
+
+        var ranks: [UInt64: Int32] = [:]
+        ranks.reserveCapacity(mergeCount)
+        for rank in 0 ..< mergeCount {
+            let left = UInt64(mergeFlat[rank * 2])
+            let right = UInt64(mergeFlat[rank * 2 + 1])
+            ranks[(left << 32) | right] = Int32(rank)
+        }
+        self.ranks = ranks
+
+        var specials: [String: Int32] = [:]
+        var specialTexts: [Int32: String] = [:]
+        specialBlob.withUnsafeBytes { raw in
+            let base = raw.bindMemory(to: UInt8.self).baseAddress!
+            for index in 0 ..< specialCount {
+                let start = Int(specialOffsets[index]), end = Int(specialOffsets[index + 1])
+                let text = String(decoding: UnsafeBufferPointer(
+                    start: base + start, count: end - start), as: UTF8.self)
+                specials[text] = Int32(specialIds[index])
+                specialTexts[Int32(specialIds[index])] = text
+            }
+        }
+        self.specials = specials
+        self.specialTexts = specialTexts
+
+        // Carried across from the source file rather than rewritten, because this pattern
+        // is where CJK and emoji splitting is decided.
+        self.splitter = try NSRegularExpression(pattern: pattern, options: [])
+    }
+
+    var vocabularySize: Int { tokens.count }
+    var specialTokenCount: Int { specials.count }
+
+    func id(of special: String) -> Int32? { specials[special] }
+
+    // MARK: - Encode
+
+    func encode(_ text: String) -> [Int32] {
+        // Special tokens are matched before anything else and never split — a chat
+        // template's `<|im_start|>` reaching the BPE loop would come back as a handful of
+        // ordinary pieces that happen to spell it.
+        var output: [Int32] = []
+        for (chunk, special) in splitOnSpecials(text) {
+            if let special {
+                output.append(special)
+            } else {
+                output.append(contentsOf: encodeOrdinary(chunk))
+            }
+        }
+        return output
+    }
+
+    private func splitOnSpecials(_ text: String) -> [(String, Int32?)] {
+        guard !specials.isEmpty else { return [(text, nil)] }
+        var pieces: [(String, Int32?)] = []
+        var remainder = Substring(text)
+
+        while !remainder.isEmpty {
+            var earliest: (Range<Substring.Index>, Int32)? = nil
+            for (token, id) in specials {
+                if let range = remainder.range(of: token),
+                   earliest == nil || range.lowerBound < earliest!.0.lowerBound {
+                    earliest = (range, id)
+                }
+            }
+            guard let (range, id) = earliest else { break }
+            if range.lowerBound > remainder.startIndex {
+                pieces.append((String(remainder[remainder.startIndex ..< range.lowerBound]), nil))
+            }
+            pieces.append(("", id))
+            remainder = remainder[range.upperBound...]
+        }
+        if !remainder.isEmpty { pieces.append((String(remainder), nil)) }
+        return pieces
+    }
+
+    private func encodeOrdinary(_ text: String) -> [Int32] {
+        let normalised = text.precomposedStringWithCanonicalMapping   // NFC
+        var output: [Int32] = []
+
+        let full = NSRange(normalised.startIndex ..< normalised.endIndex, in: normalised)
+        splitter.enumerateMatches(in: normalised, options: [], range: full) { match, _, _ in
+            guard let match, let range = Range(match.range, in: normalised) else { return }
+            let piece = String(normalised[range])
+
+            // Byte level: every UTF-8 byte becomes one printable character.
+            let mapped = String(piece.utf8.map { Self.byteToUnicode[Int($0)] })
+            output.append(contentsOf: bpe(mapped))
+        }
+        return output
+    }
+
+    /// Merge the lowest-ranked adjacent pair until none remains.
+    private func bpe(_ text: String) -> [Int32] {
+        if let whole = ids[text] { return [whole] }
+
+        var parts = text.map { String($0) }.compactMap { ids[$0] }
+        guard parts.count > 1 else { return parts }
+
+        while parts.count > 1 {
+            var bestRank = Int32.max
+            var bestIndex = -1
+            for index in 0 ..< (parts.count - 1) {
+                let key = (UInt64(UInt32(bitPattern: parts[index])) << 32)
+                    | UInt64(UInt32(bitPattern: parts[index + 1]))
+                if let rank = ranks[key], rank < bestRank {
+                    bestRank = rank
+                    bestIndex = index
+                }
+            }
+            guard bestIndex >= 0 else { break }
+
+            let merged = tokens[Int(parts[bestIndex])] + tokens[Int(parts[bestIndex + 1])]
+            guard let id = ids[merged] else { break }
+            parts.replaceSubrange(bestIndex ... (bestIndex + 1), with: [id])
+        }
+        return parts
+    }
+
+    // MARK: - Decode
+
+    func decode(_ sequence: [Int32]) -> String {
+        var bytes: [UInt8] = []
+        for id in sequence {
+            if let special = specialTexts[id] {
+                bytes.append(contentsOf: Array(special.utf8))
+                continue
+            }
+            guard id >= 0, Int(id) < tokens.count else { continue }
+            for character in tokens[Int(id)] {
+                if let byte = Self.unicodeToByte[character] { bytes.append(byte) }
+            }
+        }
+        return String(decoding: bytes, as: UTF8.self)
+    }
+}
+
+private extension Data {
+    func toUInt32Array() -> [UInt32] {
+        withUnsafeBytes { raw in
+            Array(UnsafeBufferPointer(
+                start: raw.bindMemory(to: UInt32.self).baseAddress!,
+                count: count / 4))
+        }
+    }
+}
diff --git a/ios/Sources/Edge0MLX/Edge0ChatEngine.swift b/ios/Sources/Edge0MLX/Edge0ChatEngine.swift
new file mode 100644
index 0000000..c5515da
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge0ChatEngine.swift
@@ -0,0 +1,249 @@
+import Edge0Core
+import Foundation
+import MLX
+
+public struct Edge0GenerationResult: Sendable {
+    public let text: String
+    public let generatedTokenCount: Int
+    public let elapsedSeconds: TimeInterval
+    public let prefillSeconds: TimeInterval
+    public let prefillTokensPerSecond: Double
+    public let timeToFirstTokenSeconds: TimeInterval
+    public let decodeTokensPerSecond: Double
+    public let peakMemoryBytes: Int
+}
+
+private enum Edge0ChatTemplate8B {
+    static func firstTurn(_ text: String, thinking: Bool) -> String {
+        "<role>SYSTEM</role>detailed thinking \(thinking ? "on" : "off")<|role_end|>" +
+        "<role>HUMAN</role>\(text)<|role_end|>" + assistantPrefix(thinking: thinking)
+    }
+
+    static func nextTurn(_ text: String, thinking: Bool) -> String {
+        "<|role_end|><role>HUMAN</role>\(text)<|role_end|>" +
+        assistantPrefix(thinking: thinking)
+    }
+
+    private static func assistantPrefix(thinking: Bool) -> String {
+        "<role>ASSISTANT</role>\n" + (thinking ? "<think>" : "<think></think>")
+    }
+}
+
+/// Product-facing, stateful chat entry point. The model math and cache lifetime
+/// remain owned by StreamingEdge0Model8B; this type only mirrors the official
+/// chat template and official sampling loop.
+public final class Edge0ChatEngine: @unchecked Sendable {
+    private let tokenizer: Edge0Tokenizer8B
+    private let model: StreamingEdge0Model8B
+    private var hasConversationContext = false
+    private var tokenCount = 0
+
+    /// Short prompts do not amortize loading all 128 experts per layer.
+    private static let fullLayerPrefillMinimum = 128
+
+    private static var runtimeDevice: Device {
+        #if targetEnvironment(simulator)
+        .cpu
+        #else
+        .gpu
+        #endif
+    }
+
+    public init(modelURL: URL, progress: (String) -> Void = { _ in }) throws {
+        let built = try Device.withDefaultDevice(Self.runtimeDevice) {
+            let configuration = try Edge0Configuration8B.load(from: modelURL.appendingPathComponent("config.json"))
+            let tokenizer = try Edge0Tokenizer8B(contentsOf: modelURL.appendingPathComponent("tokenizer.json"))
+            let store = try ExpertTensorStore(
+                modelURL: modelURL.appendingPathComponent("model.safetensors"),
+                expertCount: configuration.numExperts
+            )
+            let model = try StreamingEdge0Model8B(
+                configuration: configuration, store: store,
+                loraWeightsURL: modelURL.appendingPathComponent(
+                    "lora_edge0_8b.safetensors"),
+                prerouterWeightsURL: modelURL.appendingPathComponent(
+                    "prerouter_edge0_8b.safetensors"),
+                progress: progress)
+            progress("prewarming official decode path")
+            try model.prewarm()
+            return (tokenizer, model)
+        }
+        tokenizer = built.0
+        model = built.1
+    }
+
+    public func reset() {
+        model.reset()
+        hasConversationContext = false
+        tokenCount = 0
+    }
+
+    public func reply(
+        to userText: String,
+        maxTokens: Int = 2048,
+        thinking: Bool = false,
+        seed: UInt64? = nil,
+        onText: @escaping @Sendable (String) -> Void = { _ in },
+        shouldContinue: @escaping @Sendable () -> Bool = { true }
+    ) throws -> Edge0GenerationResult {
+        try Device.withDefaultDevice(Self.runtimeDevice) {
+            try generateReply(to: userText, maxTokens: maxTokens,
+                              thinking: thinking, seed: seed,
+                              onText: onText, shouldContinue: shouldContinue)
+        }
+    }
+
+    private func generateReply(
+        to userText: String, maxTokens: Int, thinking: Bool, seed: UInt64?,
+        onText: @escaping @Sendable (String) -> Void,
+        shouldContinue: @escaping @Sendable () -> Bool
+    ) throws -> Edge0GenerationResult {
+        let trimmed = userText.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard !trimmed.isEmpty else { throw M1Error.invalid("Message cannot be empty") }
+        guard maxTokens > 0 else { throw M1Error.invalid("maxTokens must be positive") }
+
+        let prompt = hasConversationContext
+            ? Edge0ChatTemplate8B.nextTurn(trimmed, thinking: thinking)
+            : Edge0ChatTemplate8B.firstTurn(trimmed, thinking: thinking)
+        let promptIDs = try tokenizer.encode(prompt)
+        guard !promptIDs.isEmpty else { throw M1Error.invalid("Tokenizer returned an empty prompt") }
+
+        let started = Date()
+        guard shouldContinue() else {
+            model.reset(); hasConversationContext = false; tokenCount = 0
+            throw CancellationError()
+        }
+        var output: StreamingModelOutput?
+        // Full-layer prefill loads every expert in a layer. That cost is almost
+        // independent of length, so a short chat prompt pays for the whole 4 GB
+        // expert set. Below this cutoff, walk the same top-8 decode path.
+        let streamingPrefill = promptIDs.count < Self.fullLayerPrefillMinimum
+        if streamingPrefill {
+            for id in promptIDs {
+                guard shouldContinue() else {
+                    model.reset(); hasConversationContext = false; tokenCount = 0
+                    throw CancellationError()
+                }
+                output = try model(tokenID: id)
+            }
+        } else {
+            let prefillChunk = 2048
+            for start in stride(from: 0, to: promptIDs.count, by: prefillChunk) {
+                guard shouldContinue() else {
+                    model.reset(); hasConversationContext = false; tokenCount = 0
+                    throw CancellationError()
+                }
+                let end = min(start + prefillChunk, promptIDs.count)
+                let chunk = Array(promptIDs[start..<end])
+                if chunk.count > 1 {
+                    output = try model.prefill(tokenIDs: chunk)
+                } else if let id = chunk.first {
+                    output = try model(tokenID: id)
+                }
+            }
+        }
+        let prefillFinished = Date()
+        tokenCount += promptIDs.count
+
+        var generated: [Int] = []
+        var samplingHistory = promptIDs
+        var peakMemory = output?.peakMemory ?? 0
+        var firstTokenAt: Date?
+        var decodeStartedAt: Date?
+        while generated.count < maxTokens, let current = output {
+            if !shouldContinue() {
+                if generated.isEmpty {
+                    model.reset(); hasConversationContext = false; tokenCount = 0
+                    throw CancellationError()
+                }
+                break
+            }
+            let next = try sample(current.logits, history: samplingHistory,
+                                  firstToken: generated.isEmpty, seed: seed.map {
+                                      $0 &+ UInt64(generated.count)
+                                  })
+            if next == tokenizer.endOfTurnTokenID { break }
+            generated.append(next)
+            samplingHistory.append(next)
+            onText(try tokenizer.decode(generated))
+            if firstTokenAt == nil {
+                firstTokenAt = Date()
+                decodeStartedAt = firstTokenAt
+            }
+            if generated.count == maxTokens { break }
+            output = try model(tokenID: next)
+            peakMemory = max(peakMemory, output?.peakMemory ?? 0)
+            tokenCount += 1
+        }
+        hasConversationContext = true
+
+        let finished = Date()
+        let firstTokenSeconds = firstTokenAt?.timeIntervalSince(started)
+            ?? finished.timeIntervalSince(started)
+        let decodeTokenCount = max(0, generated.count - 1)
+        let decodeSeconds = decodeStartedAt.map { finished.timeIntervalSince($0) } ?? 0
+        let prefillSeconds = prefillFinished.timeIntervalSince(started)
+        let decodeTokensPerSecond = decodeSeconds > 0
+            ? Double(decodeTokenCount) / decodeSeconds : 0
+        print(String(
+            format: "[8b prefill] path=%@ tokens=%d total=%.2fs rate=%.2f tok/s",
+            streamingPrefill ? "streaming-top8" : "full-layer",
+            promptIDs.count,
+            prefillSeconds,
+            Double(promptIDs.count) / max(0.000_001, prefillSeconds)))
+        print(String(
+            format: "[8b decode] tokens=%d total=%.2fs rate=%.2f tok/s peak=%.0f MB",
+            decodeTokenCount,
+            decodeSeconds,
+            decodeTokensPerSecond,
+            Double(peakMemory) / 1_048_576))
+
+        return Edge0GenerationResult(
+            text: try tokenizer.decode(generated),
+            generatedTokenCount: generated.count,
+            elapsedSeconds: finished.timeIntervalSince(started),
+            prefillSeconds: prefillSeconds,
+            prefillTokensPerSecond: Double(promptIDs.count) /
+                max(0.000_001, prefillSeconds),
+            timeToFirstTokenSeconds: firstTokenSeconds,
+            decodeTokensPerSecond: decodeTokensPerSecond,
+            peakMemoryBytes: peakMemory
+        )
+    }
+
+    /// Official Edge0 policy: first token greedy, then temperature 0.7,
+    /// top-k 64, top-p 0.95 and HF-style repetition penalty 1.1.
+    private func sample(_ input: MLXArray, history: [Int],
+                        firstToken: Bool, seed: UInt64?) throws -> Int {
+        let logits = input.flattened().asType(.float32)
+        if firstToken { return argMax(logits).item(Int32.self).intValue }
+
+        if !history.isEmpty {
+            let unique = Array(Set(history)).sorted()
+            let ids = MLXArray(unique.map(Int32.init))
+            let values = logits[ids]
+            logits[ids] = which(values .> 0, values / 1.1, values * 1.1)
+        }
+        let scaled = logits / 0.7
+        let k = min(64, scaled.dim(0))
+        let candidateIDs = argPartition(-scaled, kth: k - 1)[..<k]
+        let candidateValues = takeAlong(scaled, candidateIDs, axis: 0)
+        let order = argSort(-candidateValues)
+        let rankedIDs = takeAlong(candidateIDs, order, axis: 0)
+        let rankedValues = takeAlong(candidateValues, order, axis: 0)
+        let cumulative = cumsum(softmax(rankedValues), axis: 0)
+        let retained = (cumulative .<= 0.95).asType(.int32).sum()
+        let retainedCount = maximum(retained, MLXArray(Int32(1)))
+        let thresholdIndex = (retainedCount - 1).reshaped([1])
+        let threshold = takeAlong(rankedValues, thresholdIndex, axis: 0)
+        let masked = which(rankedValues .< threshold,
+                           MLXArray(-Float.infinity), rankedValues)
+        let local = MLXRandom.categorical(
+            masked, key: seed.map(MLXRandom.key)).item(UInt32.self)
+        return rankedIDs[Int(local)].item(Int32.self).intValue
+    }
+}
+
+private extension Int32 {
+    var intValue: Int { Int(self) }
+}
diff --git a/ios/Sources/Edge0MLX/Edge0GroupedRouter8B.swift b/ios/Sources/Edge0MLX/Edge0GroupedRouter8B.swift
new file mode 100644
index 0000000..8cc65d7
--- /dev/null
+++ b/ios/Sources/Edge0MLX/Edge0GroupedRouter8B.swift
@@ -0,0 +1,53 @@
+import MLX
+
+/// MLX implementation of Edge0/Ling's SIGMOID_GROUP / noaux_tc routing law.
+/// Mirrors Edge0 8B gate: expert bias affects selection only; routed
+/// weights come from the unbiased sigmoid scores.
+public enum Edge0GroupedRouter8B {
+    public static func select(
+        logits: MLXArray,
+        expertBias: MLXArray,
+        topK: Int = 8,
+        nGroup: Int = 8,
+        topkGroup: Int = 4,
+        normalize: Bool = true,
+        routedScale: Float = 2.5
+    ) -> (indices: MLXArray, weights: MLXArray) {
+        precondition(nGroup > 0)
+        precondition(logits.shape == expertBias.shape || expertBias.shape == [logits.dim(-1)])
+        precondition(logits.dim(-1) / nGroup >= 2)
+        precondition(topK <= topkGroup * (logits.dim(-1) / nGroup))
+        let raw = sigmoid(logits.asType(.float32))
+        var selection = raw + expertBias
+        let experts = selection.dim(-1)
+        precondition(experts % nGroup == 0)
+        precondition(topK > 0 && topK <= experts)
+        precondition(topkGroup > 0 && topkGroup <= nGroup)
+
+        let groupsToDrop = nGroup - topkGroup
+        if groupsToDrop > 0 {
+            let grouped = unflatten(selection, axis: -1, shape: [nGroup, -1])
+            let groupScores = top(grouped, k: 2, axis: -1).sum(axis: -1, keepDims: true)
+            let dropped = argPartition(
+                groupScores,
+                kth: groupsToDrop - 1,
+                axis: -2
+            )[.ellipsis, ..<groupsToDrop, 0...]
+
+            selection = putAlong(
+                grouped,
+                stopGradient(dropped),
+                values: MLXArray(-Float.infinity),
+                axis: -2
+            )
+            selection = flattened(selection, start: -2, end: -1)
+        }
+
+        let indices = argPartition(-selection, kth: topK - 1, axis: -1)[.ellipsis, ..<topK]
+        var weights = takeAlong(raw, indices, axis: -1)
+        if normalize {
+            weights = weights / (weights.sum(axis: -1, keepDims: true) + 1e-20)
+        }
+        return (indices, weights * routedScale)
+    }
+}
diff --git a/ios/Sources/Edge0MLX/LingPrerouter.swift b/ios/Sources/Edge0MLX/LingPrerouter.swift
new file mode 100644
index 0000000..8e75e37
--- /dev/null
+++ b/ios/Sources/Edge0MLX/LingPrerouter.swift
@@ -0,0 +1,108 @@
+import Edge0Core
+import Foundation
+import MLX
+
+struct LingPrerouterFeature {
+    let input: MLXArray
+    let routeIndices: MLXArray
+}
+
+/// Language-level port of `LingPrerouterStager` for Edge0-8B.
+/// Owners 7...22 predict the routed set and weights for consumers 8...23.
+final class LingPrerouter {
+    private let c: Edge0Configuration8B
+    private let owners = Array(7...22)
+    private let fc1: MLXArray
+    private let fc2: MLXArray
+    private let linear: MLXArray
+    private var previousIndices: [Int: MLXArray] = [:]
+
+    init(configuration c: Edge0Configuration8B, weightsURL: URL) throws {
+        self.c = c
+        let store = try ExpertTensorStore(modelURL: weightsURL,
+                                          expertCount: c.numExperts)
+        let ownerIDs = Array(7...22)
+        func weights(_ suffix: String) throws -> MLXArray {
+            try stacked(ownerIDs.map {
+                try store.mlxArrayCopying(named: "layers.\($0).\(suffix)")
+            }, axis: 0)
+        }
+        fc1 = try weights("fc1.weight")
+        fc2 = try weights("fc2.weight")
+        linear = try weights("linear_init.weight")
+        eval(fc1, fc2, linear)
+    }
+
+    func predict(features: [Int: LingPrerouterFeature]) throws
+        -> [Int: PredictedMoERoute] {
+        var activeOwners: [Int] = []
+        var inputs: [MLXArray] = []
+        var currentIndices: [Int: MLXArray] = [:]
+
+        for owner in owners {
+            guard let feature = features[owner] else { continue }
+            let rows = feature.routeIndices.dim(0)
+            guard rows > 0 else { continue }
+            let current = row(feature.routeIndices, at: rows - 1)
+            let previous = previousIndices[owner]
+                ?? (rows > 1 ? row(feature.routeIndices, at: rows - 2) : nil)
+            currentIndices[owner] = current
+
+            let tokenInput = lastToken(feature.input)
+            let combined = concatenated([
+                tokenInput,
+                oneHot(current),
+                oneHot(previous),
+            ], axis: -1)
+            activeOwners.append(owner)
+            inputs.append(combined)
+        }
+        guard activeOwners == owners else { return [:] }
+
+        // Official `head_batch`: 16 owners share three batched matmuls rather
+        // than issuing 48 tiny per-head GPU operations.
+        let x = stacked(inputs, axis: 0).asType(fc1.dtype)
+        let hidden = matmul(x, fc1.transposed(0, 2, 1))
+        let gelu = 0.5 * hidden * (1 + erf(hidden / Float(2).squareRoot()))
+        let logits = matmul(gelu, fc2.transposed(0, 2, 1))
+            + matmul(x, linear.transposed(0, 2, 1))
+        let route = Edge0GroupedRouter8B.select(
+            logits: logits, expertBias: MLXArray.zeros([c.numExperts]),
+            topK: c.expertsPerToken, nGroup: c.nGroup,
+            topkGroup: c.topkGroup, normalize: c.normTopkProb,
+            routedScale: c.routedScalingFactor)
+        eval(route.indices, route.weights)
+        let allIDs = route.indices.asArray(Int32.self).map(Int.init)
+        let allWeights = route.weights.asArray(Float.self)
+
+        var predictions: [Int: PredictedMoERoute] = [:]
+        for (position, owner) in owners.enumerated() {
+            let start = position * c.expertsPerToken
+            let end = start + c.expertsPerToken
+            predictions[owner + 1] = PredictedMoERoute(
+                expertIndices: Array(allIDs[start..<end]),
+                expertWeights: Array(allWeights[start..<end]))
+        }
+        previousIndices.merge(currentIndices) { _, new in new }
+        return predictions
+    }
+
+    private func oneHot(_ ids: MLXArray?) -> MLXArray {
+        var values = MLXArray.zeros([c.numExperts])
+        guard let ids else { return values.reshaped([1, c.numExperts]) }
+        values[ids.flattened()] = MLXArray.ones([c.expertsPerToken])
+        return values.reshaped([1, c.numExperts])
+    }
+
+    private func row(_ values: MLXArray, at index: Int) -> MLXArray {
+        let parts = split(values, indices: [index, index + 1], axis: 0)
+        return parts[1]
+    }
+
+    private func lastToken(_ input: MLXArray) -> MLXArray {
+        guard input.dim(0) > 1 else { return input }
+        return split(input, indices: [input.dim(0) - 1], axis: 0)[1]
+    }
+
+    func reset() { previousIndices.removeAll(keepingCapacity: true) }
+}
diff --git a/ios/Sources/Edge0MLX/MLXExpertTensorStore.swift b/ios/Sources/Edge0MLX/MLXExpertTensorStore.swift
new file mode 100644
index 0000000..3e7962a
--- /dev/null
+++ b/ios/Sources/Edge0MLX/MLXExpertTensorStore.swift
@@ -0,0 +1,56 @@
+import Edge0Core
+import Foundation
+import MLX
+
+public enum Edge0MLXError: Error {
+    case unsupportedDType(SafetensorsDType)
+}
+
+public extension ExpertTensorStore {
+    /// Mirrors the upstream streaming loader's targeted `MADV_WILLNEED`
+    /// pass. It retains no MLX arrays and is safe to issue speculatively.
+    func adviseExpertsWillNeed(layer: Int, experts: [Int]) {
+        for expert in Set(experts) where (0..<expertCount).contains(expert) {
+            for projection in ExpertProjection.allCases {
+                for part in QuantPart.allCases {
+                    let name = stackedTensorName(
+                        layer: layer, projection: projection, part: part)
+                    guard let stacked = index[name],
+                          let slice = try? stacked.axisZeroSlice(expert) else {
+                        continue
+                    }
+                    mappedFile.adviseWillNeed(range: slice.byteRange)
+                }
+            }
+        }
+    }
+
+    /// Correctness-first bridge from an mmap-backed expert slice into MLX.
+    ///
+    /// Only the requested expert tensor is copied. The 4.5 GB checkpoint remains
+    /// memory-mapped and non-resident as a whole. This avoids
+    /// MLXArray(rawPointer:) here: MLX documents that Metal requires compatible
+    /// backing; a normal file mmap is deliberately copied into MLX-owned storage.
+    func mlxArrayCopying(
+        layer: Int,
+        expert: Int,
+        projection: ExpertProjection,
+        part: QuantPart
+    ) throws -> MLXArray {
+        let slice = try slice(layer: layer, expert: expert, projection: projection, part: part)
+        return try copyTensor(descriptor: slice.descriptor, pointer: slice.pointer)
+    }
+
+    /// Copies one row from an arbitrary axis-0-contiguous checkpoint tensor.
+    /// Used by the quantized embedding so a token lookup never materializes
+    /// the full 157k x 1536 embedding table.
+    func mlxAxisZeroSliceCopying(named name: String, index row: Int) throws -> MLXArray {
+        let tensor = try index.tensor(named: name)
+        guard tensor.shape.first.map({ row >= 0 && row < $0 }) == true else {
+            throw M1Error.invalid("Axis-zero index out of range for \(name)")
+        }
+        let slice = try tensor.axisZeroSlice(row)
+        let pointer = try mappedFile.pointer(to: slice.byteRange)
+        return try copyTensor(descriptor: slice, pointer: pointer)
+    }
+}
diff --git a/ios/Sources/Edge0MLX/OfficialGatedDelta.swift b/ios/Sources/Edge0MLX/OfficialGatedDelta.swift
new file mode 100644
index 0000000..acaa296
--- /dev/null
+++ b/ios/Sources/Edge0MLX/OfficialGatedDelta.swift
@@ -0,0 +1,114 @@
+import MLX
+
+/// Language-level Swift port of mlx-lm's vector-gated delta Metal kernel used
+/// by the official Edge0 8B KDA prefill path. The Metal body, grid and
+/// template parameters intentionally retain the upstream structure.
+public enum OfficialGatedDelta {
+    private static let vectorKernel = MLXFast.metalKernel(
+        name: "edge0_gated_delta_step_vec",
+        inputNames: ["q", "k", "v", "g", "beta", "state_in", "T"],
+        outputNames: ["y", "state_out"],
+        source: """
+            auto n = thread_position_in_grid.z;
+            auto b_idx = n / Hv;
+            auto hv_idx = n % Hv;
+            auto hk_idx = hv_idx / (Hv / Hk);
+            constexpr int n_per_t = Dk / 32;
+
+            // q, k: [B, T, Hk, Dk]
+            auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+            auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+            // v, y: [B, T, Hv, Dv]
+            auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+            y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+            auto dk_idx = thread_position_in_threadgroup.x;
+            auto dv_idx = thread_position_in_grid.y;
+
+            // state_in, state_out: [B, Hv, Dv, Dk]
+            auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+            auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+            float state[n_per_t];
+            for (int i = 0; i < n_per_t; ++i) {
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = static_cast<float>(i_state[s_idx]);
+            }
+
+            // g: [B, T, Hv, Dk]
+            auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;
+            auto beta_ = beta + b_idx * T * Hv;
+
+            for (int t = 0; t < T; ++t) {
+              float kv_mem = 0.0f;
+              for (int i = 0; i < n_per_t; ++i) {
+                auto s_idx = n_per_t * dk_idx + i;
+                state[i] = state[i] * g_[s_idx];
+                kv_mem += state[i] * k_[s_idx];
+              }
+              kv_mem = simd_sum(kv_mem);
+
+              auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+
+              float out = 0.0f;
+              for (int i = 0; i < n_per_t; ++i) {
+                auto s_idx = n_per_t * dk_idx + i;
+                state[i] = state[i] + k_[s_idx] * delta;
+                out += state[i] * q_[s_idx];
+              }
+              out = simd_sum(out);
+              if (thread_index_in_simdgroup == 0) {
+                y[dv_idx] = static_cast<InT>(out);
+              }
+
+              q_ += Hk * Dk;
+              k_ += Hk * Dk;
+              v_ += Hv * Dv;
+              y += Hv * Dv;
+              g_ += Hv * Dk;
+              beta_ += Hv;
+            }
+            for (int i = 0; i < n_per_t; ++i) {
+              auto s_idx = n_per_t * dk_idx + i;
+              o_state[s_idx] = static_cast<InT>(state[i]);
+            }
+        """)
+
+    public static func callAsFunction(q: MLXArray, k: MLXArray, v: MLXArray,
+                                      decay: MLXArray, beta: MLXArray,
+                                      state: MLXArray) throws
+        -> (output: MLXArray, state: MLXArray) {
+        guard q.ndim == 4, k.shape == q.shape, v.ndim == 4,
+              decay.shape == q.shape,
+              beta.shape == [q.dim(0), q.dim(1), v.dim(2)],
+              q.dim(2) > 0, v.dim(2) % q.dim(2) == 0,
+              q.dim(3) > 0, q.dim(3) % 32 == 0,
+              state.shape == [q.dim(0), v.dim(2), v.dim(3), q.dim(3)],
+              q.dtype == k.dtype, q.dtype == v.dtype,
+              q.dtype == decay.dtype, q.dtype == state.dtype else {
+            throw M1Error.invalid("Invalid official gated-delta tensor shapes")
+        }
+        let batch = q.dim(0)
+        let tokens = q.dim(1)
+        let keyHeads = q.dim(2)
+        let keyDimension = q.dim(3)
+        let valueHeads = v.dim(2)
+        let valueDimension = v.dim(3)
+        let outputs = vectorKernel(
+            [q, k, v, decay, beta, state, tokens],
+            template: [
+                ("InT", q.dtype), ("Dk", keyDimension),
+                ("Dv", valueDimension), ("Hk", keyHeads),
+                ("Hv", valueHeads),
+            ],
+            grid: (32, valueDimension, batch * valueHeads),
+            threadGroup: (32, 4, 1),
+            outputShapes: [
+                [batch, tokens, valueHeads, valueDimension],
+                state.shape,
+            ],
+            outputDTypes: [q.dtype, q.dtype])
+        return (outputs[0], outputs[1])
+    }
+}
diff --git a/ios/Sources/Edge0MLX/ParallelLoRA.swift b/ios/Sources/Edge0MLX/ParallelLoRA.swift
new file mode 100644
index 0000000..ab9c34d
--- /dev/null
+++ b/ios/Sources/Edge0MLX/ParallelLoRA.swift
@@ -0,0 +1,62 @@
+import Edge0Core
+import MLX
+
+/// Parallel, unmerged LoRA used by the official Edge0 deployment.
+/// Base INT4 weights stay untouched; fp16 rank-16 adapters are added at run time.
+struct ParallelLinear {
+    let base: QuantizedExpertLinear
+    private let loraA: MLXArray?
+    private let loraB: MLXArray?
+    private let scale: Float
+
+    init(base: QuantizedExpertLinear, loraStore: ExpertTensorStore? = nil,
+         prefix: String? = nil, rank: Int = 16, alpha: Float = 32) throws {
+        self.base = base
+        scale = alpha / Float(rank)
+        if let loraStore, let prefix {
+            let a = try loraStore.mlxArrayCopying(named: prefix + ".lora_A")
+            let b = try loraStore.mlxArrayCopying(named: prefix + ".lora_B")
+            guard a.shape == [rank, base.inputSize],
+                  b.shape == [base.outputSize, rank] else {
+                throw M1Error.invalid("LoRA dimensions disagree for \(prefix)")
+            }
+            loraA = a
+            loraB = b
+            eval(a, b)
+        } else {
+            loraA = nil
+            loraB = nil
+        }
+    }
+
+    var inputSize: Int { base.inputSize }
+    var outputSize: Int { base.outputSize }
+
+    func callAsFunction(_ input: MLXArray) throws -> MLXArray {
+        let output = try base(input)
+        guard let loraA, let loraB else { return output }
+        let x = input.asType(loraA.dtype)
+        let delta = matmul(matmul(x, loraA.T), loraB.T) * scale
+        return output + delta.asType(output.dtype)
+    }
+}
+
+struct ParallelExpert {
+    let up: ParallelLinear
+    let gate: ParallelLinear
+    let down: ParallelLinear
+
+    init(base: QuantizedExpert, loraStore: ExpertTensorStore?, prefix: String) throws {
+        up = try ParallelLinear(base: base.up, loraStore: loraStore,
+                                prefix: prefix + ".up_proj")
+        gate = try ParallelLinear(base: base.gate, loraStore: loraStore,
+                                  prefix: prefix + ".gate_proj")
+        down = try ParallelLinear(base: base.down, loraStore: loraStore,
+                                  prefix: prefix + ".down_proj")
+    }
+
+    func callAsFunction(_ input: MLXArray) throws -> MLXArray {
+        let g = try gate(input)
+        return try down((g * sigmoid(g)) * up(input))
+    }
+}
diff --git a/ios/Sources/Edge0MLX/QuantizedExpert.swift b/ios/Sources/Edge0MLX/QuantizedExpert.swift
new file mode 100644
index 0000000..84c6b60
--- /dev/null
+++ b/ios/Sources/Edge0MLX/QuantizedExpert.swift
@@ -0,0 +1,227 @@
+import Edge0Core
+import Foundation
+import MLX
+
+public enum M1Error: Error, LocalizedError {
+    case invalid(String)
+    public var errorDescription: String? {
+        switch self { case .invalid(let message): return message }
+    }
+}
+
+/// Affine INT4: W[o,i] = nibble[o,i] * scale[o,i/64] + bias[o,i/64].
+/// Production affine INT4 linear. Coefficients retain their checkpoint dtype
+/// and the output follows the activation dtype, matching upstream MLX.
+public struct QuantizedExpertLinear {
+    public let weight: MLXArray
+    public let scales: MLXArray
+    public let biases: MLXArray
+    public let inputSize: Int
+    public let outputSize: Int
+
+    public init(weight: MLXArray, scales: MLXArray, biases: MLXArray) throws {
+        guard weight.ndim == 2, weight.dtype == .uint32,
+              weight.dim(0) > 0, weight.dim(1) > 0, weight.dim(1) % 8 == 0,
+              scales.shape == [weight.dim(0), weight.dim(1) / 8],
+              biases.shape == scales.shape,
+              [.float32, .float16, .bfloat16].contains(scales.dtype),
+              [.float32, .float16, .bfloat16].contains(biases.dtype) else {
+            throw M1Error.invalid("Expected affine INT4 [out,in/8] U32 and floating [out,in/64] coefficients")
+        }
+        self.weight = weight
+        self.scales = scales
+        self.biases = biases
+        inputSize = weight.dim(1) * 8
+        outputSize = weight.dim(0)
+    }
+
+    public func callAsFunction(_ input: MLXArray) throws -> MLXArray {
+        guard input.ndim == 2, input.dim(0) > 0, input.dim(1) == inputSize,
+              [.float32, .float16, .bfloat16].contains(input.dtype) else {
+            throw M1Error.invalid("Linear input must be floating [tokens,\(inputSize)]")
+        }
+        return quantizedMM(input, weight, scales: scales,
+                           biases: biases, transpose: true, groupSize: 64, bits: 4, mode: .affine)
+    }
+}
+
+public struct QuantizedExpert {
+    public let up: QuantizedExpertLinear
+    public let gate: QuantizedExpertLinear
+    public let down: QuantizedExpertLinear
+
+    public init(up: QuantizedExpertLinear, gate: QuantizedExpertLinear,
+                down: QuantizedExpertLinear) throws {
+        guard up.inputSize == gate.inputSize, up.outputSize == gate.outputSize,
+              down.inputSize == up.outputSize, down.outputSize == up.inputSize else {
+            throw M1Error.invalid("Inconsistent SwiGLU projection dimensions")
+        }
+        self.up = up; self.gate = gate; self.down = down
+    }
+
+    public func callAsFunction(_ input: MLXArray) throws -> MLXArray {
+        let g = try gate(input)
+        return try down((g * sigmoid(g)) * up(input))
+    }
+}
+
+/// Official Edge0 whole-layer prefill representation: every routed expert for
+/// one projection remains stacked on axis zero and is selected by
+/// `gatherQuantizedMM`. Unlike the one-expert correctness path, checkpoint
+/// BF16 scales/biases stay BF16, matching `StreamingSwitchGLU.load_full_layer`.
+public struct QuantizedExpertStackedLinear {
+    public let weight: MLXArray
+    public let scales: MLXArray
+    public let biases: MLXArray
+    public let expertCount: Int
+    public let inputSize: Int
+    public let outputSize: Int
+
+    public init(weight: MLXArray, scales: MLXArray, biases: MLXArray) throws {
+        guard weight.ndim == 3, weight.dtype == .uint32,
+              weight.dim(0) > 0, weight.dim(1) > 0,
+              weight.dim(2) > 0, weight.dim(2) % 8 == 0,
+              scales.shape == [weight.dim(0), weight.dim(1), weight.dim(2) / 8],
+              biases.shape == scales.shape,
+              [.float32, .float16, .bfloat16].contains(scales.dtype),
+              [.float32, .float16, .bfloat16].contains(biases.dtype) else {
+            throw M1Error.invalid("Expected stacked affine INT4 expert tensors")
+        }
+        self.weight = weight
+        self.scales = scales
+        self.biases = biases
+        expertCount = weight.dim(0)
+        inputSize = weight.dim(2) * 8
+        outputSize = weight.dim(1)
+    }
+
+    public func callAsFunction(_ input: MLXArray, expertIndices: MLXArray,
+                               sortedIndices: Bool = true) throws -> MLXArray {
+        guard input.ndim >= 2, input.dim(-1) == inputSize,
+              [.uint8, .uint16, .uint32, .uint64,
+               .int8, .int16, .int32, .int64].contains(expertIndices.dtype),
+              [.float32, .float16, .bfloat16].contains(input.dtype) else {
+            throw M1Error.invalid("Invalid stacked expert input or indices")
+        }
+        return gatherQuantizedMM(
+            input, weight, scales: scales, biases: biases,
+            rhsIndices: expertIndices, transpose: true,
+            groupSize: 64, bits: 4, mode: .affine,
+            sortedIndices: sortedIndices)
+    }
+}
+
+public struct QuantizedExpertStack {
+    public let up: QuantizedExpertStackedLinear
+    public let gate: QuantizedExpertStackedLinear
+    public let down: QuantizedExpertStackedLinear
+
+    public init(up: QuantizedExpertStackedLinear, gate: QuantizedExpertStackedLinear,
+                down: QuantizedExpertStackedLinear) throws {
+        guard up.expertCount == gate.expertCount,
+              up.expertCount == down.expertCount,
+              up.inputSize == gate.inputSize,
+              up.outputSize == gate.outputSize,
+              down.inputSize == up.outputSize,
+              down.outputSize == up.inputSize else {
+            throw M1Error.invalid("Inconsistent stacked SwiGLU dimensions")
+        }
+        self.up = up
+        self.gate = gate
+        self.down = down
+    }
+
+    /// Official staged-slot assembly: selected expert bundles are stacked in
+    /// prediction order, then consumed by gathered quantized matmuls.
+    public init(experts: [QuantizedExpert]) throws {
+        guard !experts.isEmpty else {
+            throw M1Error.invalid("Cannot stack an empty expert selection")
+        }
+        func projection(_ keyPath: KeyPath<QuantizedExpert, QuantizedExpertLinear>) throws
+            -> QuantizedExpertStackedLinear {
+            let values = experts.map { $0[keyPath: keyPath] }
+            return try QuantizedExpertStackedLinear(
+                weight: stacked(values.map(\.weight), axis: 0),
+                scales: stacked(values.map(\.scales), axis: 0),
+                biases: stacked(values.map(\.biases), axis: 0))
+        }
+        try self.init(up: projection(\.up), gate: projection(\.gate),
+                      down: projection(\.down))
+    }
+}
+
+public extension ExpertTensorStore {
+    /// Copies one named tensor, used for the small router and shared expert.
+    func mlxArrayCopying(named name: String) throws -> MLXArray {
+        let descriptor = try index.tensor(named: name)
+        let pointer = try mappedFile.pointer(to: descriptor.byteRange)
+        return try copyTensor(descriptor: descriptor, pointer: pointer)
+    }
+
+    func loadExpert(layer: Int, expert: Int) throws -> QuantizedExpert {
+        guard (0..<expertCount).contains(expert) else { throw M1Error.invalid("Expert out of range") }
+        func projection(_ p: ExpertProjection) throws -> QuantizedExpertLinear {
+            try QuantizedExpertLinear(
+                weight: mlxArrayCopying(layer: layer, expert: expert, projection: p, part: .weight),
+                scales: mlxArrayCopying(layer: layer, expert: expert, projection: p, part: .scales),
+                biases: mlxArrayCopying(layer: layer, expert: expert, projection: p, part: .biases))
+        }
+        return try QuantizedExpert(up: projection(.up), gate: projection(.gate), down: projection(.down))
+    }
+
+    func loadSharedExpert(layer: Int) throws -> QuantizedExpert {
+        func projection(_ p: ExpertProjection) throws -> QuantizedExpertLinear {
+            let prefix = "model.layers.\(layer).mlp.shared_experts.\(p.rawValue)"
+            return try QuantizedExpertLinear(weight: mlxArrayCopying(named: prefix + ".weight"),
+                scales: mlxArrayCopying(named: prefix + ".scales"),
+                biases: mlxArrayCopying(named: prefix + ".biases"))
+        }
+        return try QuantizedExpert(up: projection(.up), gate: projection(.gate), down: projection(.down))
+    }
+
+    func loadDenseMLP(layer: Int) throws -> QuantizedExpert {
+        func projection(_ p: ExpertProjection) throws -> QuantizedExpertLinear {
+            let prefix = "model.layers.\(layer).mlp.\(p.rawValue)"
+            return try QuantizedExpertLinear(weight: mlxArrayCopying(named: prefix + ".weight"),
+                scales: mlxArrayCopying(named: prefix + ".scales"),
+                biases: mlxArrayCopying(named: prefix + ".biases"))
+        }
+        return try QuantizedExpert(up: projection(.up), gate: projection(.gate), down: projection(.down))
+    }
+
+    /// Direct translation of official `StreamingSwitchGLU.load_full_layer` for
+    /// Edge0-8B's separate gate/up/down layout: nine whole stacked tensors.
+    func loadFullExpertLayer(layer: Int) throws -> QuantizedExpertStack {
+        func projection(_ p: ExpertProjection) throws -> QuantizedExpertStackedLinear {
+            let prefix = "model.layers.\(layer).mlp.experts.\(p.rawValue)."
+            return try QuantizedExpertStackedLinear(
+                weight: mlxArrayCopying(named: prefix + "weight"),
+                scales: mlxArrayCopying(named: prefix + "scales"),
+                biases: mlxArrayCopying(named: prefix + "biases"))
+        }
+        return try QuantizedExpertStack(up: projection(.up), gate: projection(.gate),
+                                        down: projection(.down))
+    }
+}
+
+func copyTensor(descriptor: TensorDescriptor, pointer: UnsafeMutableRawPointer) throws -> MLXArray {
+    let dtype: DType
+    switch descriptor.dtype {
+    case .u32: dtype = .uint32
+    case .bf16: dtype = .bfloat16
+    case .f16: dtype = .float16
+    case .f32: dtype = .float32
+    default: throw Edge0MLXError.unsupportedDType(descriptor.dtype)
+    }
+    var elements = 1
+    for dimension in descriptor.shape {
+        let product = elements.multipliedReportingOverflow(by: dimension)
+        guard dimension > 0, !product.overflow else { throw M1Error.invalid("Invalid tensor shape") }
+        elements = product.partialValue
+    }
+    let size = elements.multipliedReportingOverflow(by: descriptor.dtype.byteWidth)
+    guard !size.overflow, size.partialValue == descriptor.byteCount else {
+        throw M1Error.invalid("Tensor shape/byte count mismatch: \(descriptor.name)")
+    }
+    return MLXArray(Data(bytes: pointer, count: descriptor.byteCount), descriptor.shape, dtype: dtype)
+}
diff --git a/ios/Sources/Edge0MLX/StreamingEdge0Model8B.swift b/ios/Sources/Edge0MLX/StreamingEdge0Model8B.swift
new file mode 100644
index 0000000..06a6bfd
--- /dev/null
+++ b/ios/Sources/Edge0MLX/StreamingEdge0Model8B.swift
@@ -0,0 +1,295 @@
+import Edge0Core
+import Foundation
+import MLX
+
+public struct StreamingModelOutput {
+    public let logits: MLXArray
+    public let activeMemory: Int
+    public let peakMemory: Int
+    public let cacheMemory: Int
+}
+
+private enum StreamingLayer {
+    case kda(StreamingKDADecoderBlock, KDAStateCache)
+    case mla(StreamingMLADecoderBlock, MLAKVCache)
+
+    func callAsFunction(_ hidden: MLXArray,
+                        predictedRoute: PredictedMoERoute? = nil) throws -> MLXArray {
+        switch self {
+        case .kda(let block, let cache):
+            return try block(hidden, cache: cache, predictedRoute: predictedRoute).hidden
+        case .mla(let block, let cache):
+            return try block(hidden, position: cache.count, cache: cache,
+                             predictedRoute: predictedRoute).hidden
+        }
+    }
+
+    var prerouterFeature: LingPrerouterFeature? {
+        switch self {
+        case .kda(let block, _):
+            guard let input = block.prerouterInput,
+                  let indices = block.prerouterIndices else { return nil }
+            return LingPrerouterFeature(input: input, routeIndices: indices)
+        case .mla(let block, _):
+            guard let input = block.prerouterInput,
+                  let indices = block.prerouterIndices else { return nil }
+            return LingPrerouterFeature(input: input, routeIndices: indices)
+        }
+    }
+
+    func stage(experts: [Int]) throws {
+        switch self {
+        case .kda(let block, _): try block.stage(experts: experts)
+        case .mla(let block, _): try block.stage(experts: experts)
+        }
+    }
+
+    func reset() {
+        switch self {
+        case .kda(let block, let cache):
+            cache.reset(); block.resetStaging()
+        case .mla(let block, let cache):
+            cache.reset(); block.resetStaging()
+        }
+    }
+
+    func prefill(_ hidden: MLXArray) throws -> MLXArray {
+        switch self {
+        case .kda(let block, let cache):
+            return try block.prefill(hidden, cache: cache).hidden
+        case .mla(let block, let cache):
+            return try block.prefill(
+                hidden, position: cache.count, cache: cache).hidden
+        }
+    }
+
+    func loadFullExpertLayer() throws {
+        switch self {
+        case .kda(let block, _): try block.loadFullExpertLayer()
+        case .mla(let block, _): try block.loadFullExpertLayer()
+        }
+    }
+
+    func clearFullExpertLayer() {
+        switch self {
+        case .kda(let block, _): block.clearFullExpertLayer()
+        case .mla(let block, _): block.clearFullExpertLayer()
+        }
+    }
+
+    var moeStats: StreamingMoEStats? {
+        switch self {
+        case .kda(let block, _): block.moeStats
+        case .mla(let block, _): block.moeStats
+        }
+    }
+}
+
+public struct StreamingRuntimeStats: Sendable {
+    public let layers: [StreamingMoEStats]
+    public let cacheHits: Int
+    public let cacheMisses: Int
+    public let cachedExperts: Int
+}
+
+/// One-token exact-streaming Edge0 model.
+///
+/// Resident tensors are limited to the non-routed backbone, routers, and shared
+/// experts. Routed experts remain in the checkpoint mmap and are copied only for
+/// the selected top-8, matching the official repository's memory contract. MLX's
+/// reusable free-buffer cache is capped at 256 MiB, also matching the official
+/// engine default.
+public final class StreamingEdge0Model8B {
+    private let configuration: Edge0Configuration8B
+    private let store: ExpertTensorStore
+    private let layers: [StreamingLayer]
+    private let finalNorm: MLXArray
+    private let lmHead: QuantizedExpertLinear
+    private let prerouter: LingPrerouter?
+    private let bundleCache: ExpertBundleCache
+    private var predictedRoutes: [Int: PredictedMoERoute] = [:]
+
+    public init(configuration c: Edge0Configuration8B, store: ExpertTensorStore,
+                cacheLimitBytes: Int = 256 * 1024 * 1024,
+                loraWeightsURL: URL? = nil,
+                prerouterWeightsURL: URL? = nil,
+                progress: (String) -> Void = { _ in }) throws {
+        guard c.quantization.bits == 4, c.quantization.groupSize == 64,
+              c.quantization.mode == "affine" else {
+            throw M1Error.invalid("Full model requires affine INT4/group-64 weights")
+        }
+        Memory.cacheLimit = cacheLimitBytes
+        configuration = c
+        self.store = store
+        finalNorm = try store.mlxArrayCopying(named: "model.norm.weight")
+        lmHead = try QuantizedExpertLinear(
+            weight: store.mlxArrayCopying(named: "lm_head.weight"),
+            scales: store.mlxArrayCopying(named: "lm_head.scales"),
+            biases: store.mlxArrayCopying(named: "lm_head.biases"))
+        guard finalNorm.shape == [c.hiddenSize], lmHead.inputSize == c.hiddenSize,
+              lmHead.outputSize == c.vocabularySize else {
+            throw M1Error.invalid("Final norm or lm-head dimensions disagree with config")
+        }
+        progress("final norm and lm-head loaded")
+
+        let bundleCache = ExpertBundleCache(capacity: 64)
+        self.bundleCache = bundleCache
+        let loraStore = try loraWeightsURL.map {
+            try ExpertTensorStore(modelURL: $0, expertCount: c.numExperts)
+        }
+        var built: [StreamingLayer] = []
+        built.reserveCapacity(c.hiddenLayers)
+        for layer in 0..<c.hiddenLayers {
+            if c.isMLALayer(layer) {
+                let block = try StreamingMLADecoderBlock(
+                    configuration: c, store: store, layer: layer,
+                    loraStore: loraStore,
+                    bundleCache: bundleCache)
+                built.append(.mla(block, MLAKVCache()))
+            } else {
+                let block = try StreamingKDADecoderBlock(
+                    configuration: c, store: store, layer: layer,
+                    loraStore: loraStore,
+                    bundleCache: bundleCache)
+                built.append(.kda(block, KDAStateCache()))
+            }
+            progress("layer \(layer) resident backbone loaded")
+        }
+        layers = built
+        if let prerouterWeightsURL {
+            prerouter = try LingPrerouter(configuration: c,
+                                         weightsURL: prerouterWeightsURL)
+            progress("Ling prerouter loaded")
+        } else {
+            prerouter = nil
+        }
+        eval(finalNorm)
+        progress("24-layer model initialized")
+    }
+
+    public func callAsFunction(tokenID: Int,
+                               afterLayer: (Int, MLXArray) -> Void = { _, _ in }) throws -> StreamingModelOutput {
+        let c = configuration
+        guard tokenID >= 0, tokenID < c.vocabularySize else {
+            throw M1Error.invalid("Token ID out of vocabulary range")
+        }
+        var hidden = try embedding(tokenID: tokenID)
+        var features: [Int: LingPrerouterFeature] = [:]
+        for (index, layer) in layers.enumerated() {
+            hidden = try layer(hidden, predictedRoute: predictedRoutes[index])
+            if (7...22).contains(index), let feature = layer.prerouterFeature {
+                features[index] = feature
+            }
+            afterLayer(index, hidden)
+        }
+        hidden = MLXFast.rmsNorm(hidden, weight: finalNorm, eps: c.rmsNormEps)
+        let logits = try lmHead(hidden)
+        eval(logits)
+        try updateStagedRoutes(features)
+        let memory = Memory.snapshot()
+        return StreamingModelOutput(logits: logits,
+                                    activeMemory: memory.activeMemory,
+                                    peakMemory: memory.peakMemory,
+                                    cacheMemory: memory.cacheMemory)
+    }
+
+    /// Official chunked prefill path: each sparse layer's complete expert
+    /// stack is loaded immediately before that layer and dropped when the
+    /// next layer begins. KDA and MLA process the full token chunk at once.
+    public func prefill(tokenIDs: [Int],
+                        afterLayer: (Int, MLXArray) -> Void = { _, _ in }) throws
+        -> StreamingModelOutput {
+        let c = configuration
+        guard tokenIDs.count > 1,
+              tokenIDs.allSatisfy({ (0..<c.vocabularySize).contains($0) }) else {
+            throw M1Error.invalid("Chunked prefill requires at least two valid token IDs")
+        }
+        defer { layers.forEach { $0.clearFullExpertLayer() } }
+
+        var hidden = try embedding(tokenIDs: tokenIDs)
+        var features: [Int: LingPrerouterFeature] = [:]
+        for (index, layer) in layers.enumerated() {
+            if index > 0 { layers[index - 1].clearFullExpertLayer() }
+            try layer.loadFullExpertLayer()
+            hidden = try layer.prefill(hidden)
+            asyncEval(hidden)
+            if (7...22).contains(index), let feature = layer.prerouterFeature {
+                features[index] = feature
+            }
+            afterLayer(index, hidden)
+        }
+        hidden = MLXFast.rmsNorm(hidden, weight: finalNorm, eps: c.rmsNormEps)
+        let last = split(hidden, indices: [hidden.dim(0) - 1], axis: 0)[1]
+        let logits = try lmHead(last)
+        eval(logits)
+        try updateStagedRoutes(features)
+        let memory = Memory.snapshot()
+        return StreamingModelOutput(logits: logits,
+                                    activeMemory: memory.activeMemory,
+                                    peakMemory: memory.peakMemory,
+                                    cacheMemory: memory.cacheMemory)
+    }
+
+    private func embedding(tokenID: Int) throws -> MLXArray {
+        let prefix = "model.word_embeddings"
+        let weight = try store.mlxAxisZeroSliceCopying(named: prefix + ".weight", index: tokenID)
+            .expandedDimensions(axis: 0)
+        let scales = try store.mlxAxisZeroSliceCopying(named: prefix + ".scales", index: tokenID)
+            .expandedDimensions(axis: 0)
+        let biases = try store.mlxAxisZeroSliceCopying(named: prefix + ".biases", index: tokenID)
+            .expandedDimensions(axis: 0)
+        let hidden = dequantized(weight, scales: scales, biases: biases,
+                                 groupSize: 64, bits: 4, mode: .affine,
+                                 dtype: .bfloat16)
+        guard hidden.shape == [1, configuration.hiddenSize] else {
+            throw M1Error.invalid("Embedding row has an unexpected shape")
+        }
+        return hidden
+    }
+
+    /// Official step boundary: owner layers 7...22 predict consumers 8...23,
+    /// then their bundles are staged before the next token begins.
+    private func updateStagedRoutes(_ features: [Int: LingPrerouterFeature]) throws {
+        guard let prerouter else { return }
+        let next = try prerouter.predict(features: features)
+        for (consumer, route) in next {
+            try layers[consumer].stage(experts: route.expertIndices)
+        }
+        predictedRoutes = next
+    }
+
+    private func embedding(tokenIDs: [Int]) throws -> MLXArray {
+        try concatenated(tokenIDs.map { try embedding(tokenID: $0) }, axis: 0)
+    }
+
+    public func reset() {
+        layers.forEach { $0.reset() }
+        predictedRoutes.removeAll(keepingCapacity: true)
+        prerouter?.reset()
+    }
+
+    public var stats: StreamingRuntimeStats {
+        let cache = bundleCache.stats
+        return StreamingRuntimeStats(
+            layers: layers.compactMap(\.moeStats), cacheHits: cache.hits,
+            cacheMisses: cache.misses, cachedExperts: cache.size)
+    }
+
+    /// Official Ling startup warm-up: request sequential pages, execute a
+    /// tiny prompt plus one decode step, then clear only sequence state.
+    ///
+    /// The four-token warm-up uses the top-8 decode path. Full-layer prefill
+    /// of four tokens loads every expert and is the opposite of a warm-up.
+    public func prewarm(tokenID: Int = 0) throws {
+        store.mappedFile.adviseWillNeed()
+        store.mappedFile.adviseSequential()
+        store.mappedFile.sequentialRead()
+        reset()
+        for _ in 0..<4 {
+            _ = try callAsFunction(tokenID: tokenID)
+        }
+        _ = try callAsFunction(tokenID: tokenID)
+        reset()
+    }
+
+}
diff --git a/ios/Sources/Edge0MLX/StreamingKDADecoderBlock.swift b/ios/Sources/Edge0MLX/StreamingKDADecoderBlock.swift
new file mode 100644
index 0000000..3934a85
--- /dev/null
+++ b/ios/Sources/Edge0MLX/StreamingKDADecoderBlock.swift
@@ -0,0 +1,285 @@
+import Edge0Core
+import MLX
+
+public final class KDAStateCache {
+    fileprivate var qConv: MLXArray?
+    fileprivate var kConv: MLXArray?
+    fileprivate var vConv: MLXArray?
+    fileprivate var recurrent: MLXArray?
+    public fileprivate(set) var count = 0
+
+    public init() {}
+
+    func reset() {
+        qConv = nil; kConv = nil; vConv = nil; recurrent = nil; count = 0
+    }
+}
+
+/// Correctness-first one-token Kimi Delta Attention decoder layer.
+/// The recurrent update is expressed with MLX array operations; a fused Metal
+/// kernel can replace it after numerical parity is established.
+public final class StreamingKDADecoderBlock {
+    private let c: Edge0Configuration8B
+    private let layer: Int
+    private let inputNorm: MLXArray
+    private let postAttentionNorm: MLXArray
+    private let qConvWeight: MLXArray
+    private let kConvWeight: MLXArray
+    private let vConvWeight: MLXArray
+    private let aLog: MLXArray
+    private let dtBias: MLXArray
+    private let outputNorm: MLXArray
+    private let qProj: ParallelLinear
+    private let kProj: ParallelLinear
+    private let vProj: ParallelLinear
+    private let fProj: ParallelLinear
+    private let gProj: ParallelLinear
+    private let bProj: ParallelLinear
+    private let outputProj: ParallelLinear
+    private let moe: StreamingMoE?
+    private let denseMLP: ParallelExpert?
+    public private(set) var prerouterInput: MLXArray?
+    public private(set) var prerouterIndices: MLXArray?
+
+    public init(configuration c: Edge0Configuration8B, store: ExpertTensorStore, layer: Int,
+                loraStore: ExpertTensorStore? = nil,
+                bundleCache: ExpertBundleCache? = nil,
+                progress: (String) -> Void = { _ in }) throws {
+        guard !c.isMLALayer(layer), layer >= 0, layer < c.hiddenLayers,
+              c.kdaSafeGate, c.shortConvKernelSize == 4,
+              c.headDim * c.attentionHeads == 2048 else {
+            throw M1Error.invalid("KDA decoder requires a KDA layer with the preview safe-gate configuration")
+        }
+        self.c = c
+        self.layer = layer
+        let base = "model.layers.\(layer)"
+        let attention = base + ".attention"
+        inputNorm = try store.mlxArrayCopying(named: base + ".input_layernorm.weight")
+        postAttentionNorm = try store.mlxArrayCopying(named: base + ".post_attention_layernorm.weight")
+        qConvWeight = try store.mlxArrayCopying(named: attention + ".q_conv1d.weight")
+        kConvWeight = try store.mlxArrayCopying(named: attention + ".k_conv1d.weight")
+        vConvWeight = try store.mlxArrayCopying(named: attention + ".v_conv1d.weight")
+        aLog = try store.mlxArrayCopying(named: attention + ".A_log")
+        dtBias = try store.mlxArrayCopying(named: attention + ".dt_bias")
+        outputNorm = try store.mlxArrayCopying(named: attention + ".o_norm.weight")
+        progress("KDA state tensors loaded")
+        qProj = try Self.linear(store, attention + ".q_proj", loraStore)
+        kProj = try Self.linear(store, attention + ".k_proj", loraStore)
+        vProj = try Self.linear(store, attention + ".v_proj", loraStore)
+        fProj = try Self.linear(store, attention + ".f_proj", loraStore)
+        gProj = try Self.linear(store, attention + ".g_proj", loraStore)
+        bProj = try Self.linear(store, attention + ".b_proj", loraStore)
+        outputProj = try Self.linear(store, attention + ".o_proj", loraStore)
+        if layer < c.firstKDenseReplace {
+            denseMLP = try ParallelExpert(
+                base: store.loadDenseMLP(layer: layer), loraStore: loraStore,
+                prefix: base + ".mlp")
+            moe = nil
+            progress("dense MLP loaded")
+        } else {
+            moe = try StreamingMoE(configuration: c, store: store, layer: layer,
+                                   bundleCache: bundleCache)
+            denseMLP = nil
+            progress("sparse MoE loaded")
+        }
+        let projection = c.attentionHeads * c.headDim
+        guard [qProj, kProj, vProj, fProj, gProj].allSatisfy({
+            $0.inputSize == c.hiddenSize && $0.outputSize == projection
+        }), bProj.inputSize == c.hiddenSize, bProj.outputSize == c.attentionHeads,
+              outputProj.inputSize == projection, outputProj.outputSize == c.hiddenSize,
+              qConvWeight.shape == [projection, 1, c.shortConvKernelSize],
+              kConvWeight.shape == qConvWeight.shape, vConvWeight.shape == qConvWeight.shape,
+              aLog.shape == [c.attentionHeads], dtBias.shape == [projection],
+              outputNorm.shape == [c.headDim] else {
+            throw M1Error.invalid("M4 KDA tensor dimensions disagree with config")
+        }
+        eval(inputNorm, postAttentionNorm, qConvWeight, kConvWeight, vConvWeight,
+             aLog, dtBias, outputNorm)
+        progress("KDA block initialized")
+    }
+
+    private static func linear(_ store: ExpertTensorStore, _ prefix: String,
+                               _ loraStore: ExpertTensorStore?) throws -> ParallelLinear {
+        let base = try QuantizedExpertLinear(weight: store.mlxArrayCopying(named: prefix + ".weight"),
+            scales: store.mlxArrayCopying(named: prefix + ".scales"),
+            biases: store.mlxArrayCopying(named: prefix + ".biases"))
+        return try ParallelLinear(base: base, loraStore: loraStore, prefix: prefix)
+    }
+
+    public func callAsFunction(_ input: MLXArray, cache: KDAStateCache,
+                               predictedRoute: PredictedMoERoute? = nil) throws -> DecoderBlockOutput {
+        guard input.shape == [1, c.hiddenSize] else {
+            throw M1Error.invalid("M4 accepts one token [1,\(c.hiddenSize)]")
+        }
+        let x = input
+        let normalized = MLXFast.rmsNorm(x, weight: inputNorm, eps: c.rmsNormEps)
+        let attention = try kda(normalized, cache: cache)
+        let residual = x + attention
+        let moeInput = MLXFast.rmsNorm(residual, weight: postAttentionNorm, eps: c.rmsNormEps)
+        let feedForward: MLXArray
+        let expertIndices: [Int]
+        let expertWeights: [Float]
+        if let moe {
+            let output = if let predictedRoute {
+                try moe.predicted(moeInput, route: predictedRoute)
+            } else {
+                try moe(moeInput)
+            }
+            feedForward = output.hidden
+            expertIndices = output.expertIndices
+            expertWeights = output.expertWeights
+            prerouterInput = moeInput
+            prerouterIndices = output.routeIndices
+        } else if let denseMLP {
+            feedForward = try denseMLP(moeInput)
+            expertIndices = []
+            expertWeights = []
+            prerouterInput = nil
+            prerouterIndices = nil
+        } else {
+            throw M1Error.invalid("KDA decoder has no feed-forward block")
+        }
+        let hidden = residual + feedForward
+        cache.count += 1
+        return DecoderBlockOutput(hidden: hidden, expertIndices: expertIndices,
+                                  expertWeights: expertWeights)
+    }
+
+    /// Official multi-token Edge0 8B KDA path used by chunked prefill.
+    public func prefill(_ input: MLXArray, cache: KDAStateCache) throws -> DecoderBlockOutput {
+        guard input.ndim == 2, input.dim(0) > 1, input.dim(1) == c.hiddenSize else {
+            throw M1Error.invalid("KDA prefill accepts [T,\(c.hiddenSize)], T > 1")
+        }
+        let x = input
+        let normalized = MLXFast.rmsNorm(x, weight: inputNorm, eps: c.rmsNormEps)
+        let attention = try kdaPrefill(normalized, cache: cache)
+        let residual = x + attention
+        let moeInput = MLXFast.rmsNorm(residual, weight: postAttentionNorm, eps: c.rmsNormEps)
+        let feedForward: MLXArray
+        let expertIndices: [Int]
+        let expertWeights: [Float]
+        if let moe {
+            let output = try moe.fullLayerPrefill(moeInput)
+            feedForward = output.hidden
+            expertIndices = output.expertIndices
+            expertWeights = output.expertWeights
+            prerouterInput = moeInput
+            prerouterIndices = output.routeIndices
+        } else if let denseMLP {
+            feedForward = try denseMLP(moeInput)
+            expertIndices = []
+            expertWeights = []
+        } else {
+            throw M1Error.invalid("KDA decoder has no feed-forward block")
+        }
+        let hidden = residual + feedForward
+        cache.count += input.dim(0)
+        return DecoderBlockOutput(hidden: hidden, expertIndices: expertIndices,
+                                  expertWeights: expertWeights)
+    }
+
+    public func loadFullExpertLayer() throws { try moe?.loadFullLayer() }
+    public func clearFullExpertLayer() { moe?.clearFullLayer() }
+    public func stage(experts: [Int]) throws { try moe?.stage(experts: experts) }
+    public func resetStaging() { moe?.resetStaging() }
+    public var moeStats: StreamingMoEStats? { moe?.stats }
+
+    private func convStep(_ x: MLXArray, weight: MLXArray,
+                          state: inout MLXArray?) -> MLXArray {
+        let channels = c.attentionHeads * c.headDim
+        let history = state ?? MLXArray.zeros(
+            [c.shortConvKernelSize - 1, channels], dtype: x.dtype)
+        let window = concatenated([history, x], axis: 0)
+        let kernel = weight.reshaped([channels, c.shortConvKernelSize]).T
+        let output = (window * kernel).sum(axis: 0, keepDims: true)
+        state = split(window, indices: [1], axis: 0)[1]
+        return output * sigmoid(output)
+    }
+
+    private func kda(_ x: MLXArray, cache: KDAStateCache) throws -> MLXArray {
+        let heads = c.attentionHeads
+        let dimension = c.headDim
+        let qConv = convStep(try qProj(x), weight: qConvWeight, state: &cache.qConv)
+        let kConv = convStep(try kProj(x), weight: kConvWeight, state: &cache.kConv)
+        let vConv = convStep(try vProj(x), weight: vConvWeight, state: &cache.vConv)
+        var q = qConv.reshaped([1, heads, dimension]).asType(.float32)
+        var k = kConv.reshaped([1, heads, dimension]).asType(.float32)
+        let v = vConv.reshaped([1, heads, dimension]).asType(.float32)
+        let qNorm = sqrt((q * q).sum(axis: -1, keepDims: true)) + 1e-6
+        let kNorm = sqrt((k * k).sum(axis: -1, keepDims: true)) + 1e-6
+        q = q * (1.0 / Float(dimension).squareRoot()) / qNorm
+        k = k / kNorm
+
+        let f = try fProj(x).reshaped([1, heads, dimension]).asType(.float32)
+        let gateInput = exp(aLog.asType(.float32)).reshaped([1, heads, 1])
+            * (f + dtBias.asType(.float32).reshaped([1, heads, dimension]))
+        let decay = exp(Float(c.kdaLowerBound) * sigmoid(gateInput)).expandedDimensions(axis: 2)
+        let beta = sigmoid(try bProj(x).asType(.float32)).expandedDimensions(axis: -1)
+        var state = cache.recurrent ?? MLXArray.zeros([1, heads, dimension, dimension])
+        state = state * decay
+        let kRows = k.expandedDimensions(axis: 2)
+        let memory = (state * kRows).sum(axis: -1)
+        let delta = (v - memory) * beta
+        state = state + kRows * delta.expandedDimensions(axis: -1)
+        var output = (state * q.expandedDimensions(axis: 2)).sum(axis: -1)
+        cache.recurrent = state
+
+        output = MLXFast.rmsNorm(
+            output.asType(x.dtype), weight: outputNorm, eps: c.rmsNormEps)
+        let outputGate = sigmoid(try gProj(x).reshaped([1, heads, dimension]))
+        return try outputProj((output * outputGate).reshaped([1, heads * dimension]))
+    }
+
+    private func convPrefill(_ x: MLXArray, weight: MLXArray,
+                             state: inout MLXArray?) -> MLXArray {
+        let channels = c.attentionHeads * c.headDim
+        let history = state ?? MLXArray.zeros(
+            [c.shortConvKernelSize - 1, channels], dtype: x.dtype)
+        let convInput = concatenated([history, x], axis: 0)
+        let kernel = weight.transposed(0, 2, 1)
+        let output = conv1d(
+            convInput.expandedDimensions(axis: 0), kernel,
+            groups: channels).squeezed(axis: 0)
+        let stateStart = convInput.dim(0) - (c.shortConvKernelSize - 1)
+        state = split(convInput, indices: [stateStart], axis: 0)[1]
+        return output * sigmoid(output)
+    }
+
+    private func kdaPrefill(_ x: MLXArray, cache: KDAStateCache) throws -> MLXArray {
+        let tokens = x.dim(0)
+        let heads = c.attentionHeads
+        let dimension = c.headDim
+        let qConv = convPrefill(try qProj(x), weight: qConvWeight, state: &cache.qConv)
+        let kConv = convPrefill(try kProj(x), weight: kConvWeight, state: &cache.kConv)
+        let vConv = convPrefill(try vProj(x), weight: vConvWeight, state: &cache.vConv)
+        var q = qConv.reshaped([1, tokens, heads, dimension]).asType(.float32)
+        var k = kConv.reshaped([1, tokens, heads, dimension]).asType(.float32)
+        let v = vConv.reshaped([1, tokens, heads, dimension]).asType(.float32)
+        let qNorm = sqrt((q * q).sum(axis: -1, keepDims: true)) + 1e-6
+        let kNorm = sqrt((k * k).sum(axis: -1, keepDims: true)) + 1e-6
+        q = q * (1.0 / Float(dimension).squareRoot()) / qNorm
+        k = k / kNorm
+
+        let f = try fProj(x).reshaped([1, tokens, heads, dimension])
+            .asType(.float32)
+        let gateInput = exp(aLog.asType(.float32)).reshaped([1, 1, heads, 1])
+            * (f + dtBias.asType(.float32)
+                .reshaped([1, 1, heads, dimension]))
+        let decay = exp(Float(c.kdaLowerBound) * sigmoid(gateInput))
+        let beta = sigmoid(try bProj(x).asType(.float32))
+            .reshaped([1, tokens, heads])
+        let initialState = cache.recurrent
+            ?? MLXArray.zeros([1, heads, dimension, dimension])
+        let update = try OfficialGatedDelta.callAsFunction(
+            q: q, k: k, v: v, decay: decay, beta: beta,
+            state: initialState)
+        cache.recurrent = update.state
+        var output = MLXFast.rmsNorm(
+            update.output.asType(x.dtype), weight: outputNorm,
+            eps: c.rmsNormEps)
+        let outputGate = sigmoid(
+            try gProj(x).reshaped([1, tokens, heads, dimension]))
+        output = output * outputGate
+        return try outputProj(output.reshaped([tokens, heads * dimension]))
+    }
+}
diff --git a/ios/Sources/Edge0MLX/StreamingMLADecoderBlock.swift b/ios/Sources/Edge0MLX/StreamingMLADecoderBlock.swift
new file mode 100644
index 0000000..9a8c613
--- /dev/null
+++ b/ios/Sources/Edge0MLX/StreamingMLADecoderBlock.swift
@@ -0,0 +1,278 @@
+import Edge0Core
+import Foundation
+import MLX
+
+public struct DecoderBlockOutput {
+    public let hidden: MLXArray
+    public let expertIndices: [Int]
+    public let expertWeights: [Float]
+}
+
+public final class MLAKVCache {
+    private var keys: MLXArray?
+    private var values: MLXArray?
+    public private(set) var count = 0
+
+    public init() {}
+
+    func reset() { keys = nil; values = nil; count = 0 }
+
+    fileprivate func append(keys newKeys: MLXArray, values newValues: MLXArray) -> (MLXArray, MLXArray) {
+        let allKeys: MLXArray
+        let allValues: MLXArray
+        if let keys, let values {
+            allKeys = concatenated([keys, newKeys], axis: 2)
+            allValues = concatenated([values, newValues], axis: 2)
+        } else {
+            allKeys = newKeys
+            allValues = newValues
+        }
+        keys = allKeys
+        values = allValues
+        count += newKeys.dim(2)
+        return (allKeys, allValues)
+    }
+}
+
+/// Correctness-first implementation of one Edge0 MLA decoder layer for a
+/// single decode token. The current token attends to itself; `position`
+/// exercises the checkpoint's interleaved RoPE convention. KV persistence is
+/// intentionally the next milestone.
+public final class StreamingMLADecoderBlock {
+    private let configuration: Edge0Configuration8B
+    private let layer: Int
+    private let inputNorm: MLXArray
+    private let postAttentionNorm: MLXArray
+    private let qANorm: MLXArray
+    private let kvANorm: MLXArray
+    private let qA: ParallelLinear
+    private let qB: ParallelLinear
+    private let kvA: ParallelLinear
+    private let kvB: ParallelLinear
+    private let gate: ParallelLinear
+    private let dense: ParallelLinear
+    private let moe: StreamingMoE
+    public private(set) var prerouterInput: MLXArray?
+    public private(set) var prerouterIndices: MLXArray?
+
+    public init(configuration c: Edge0Configuration8B, store: ExpertTensorStore, layer: Int,
+                loraStore: ExpertTensorStore? = nil,
+                bundleCache: ExpertBundleCache? = nil,
+                progress: (String) -> Void = { _ in }) throws {
+        let queryHeadDim = c.qkNopeHeadDim + c.qkRopeHeadDim
+        guard c.isMLALayer(layer), c.ropeInterleave,
+              c.qkRopeHeadDim > 0, c.qkRopeHeadDim % 2 == 0,
+              c.attentionHeads == c.keyValueHeads else {
+            throw M1Error.invalid("M2 requires an interleaved-RoPE MLA layer with matching Q/KV heads")
+        }
+        configuration = c
+        self.layer = layer
+        let base = "model.layers.\(layer)"
+        inputNorm = try store.mlxArrayCopying(named: base + ".input_layernorm.weight")
+        postAttentionNorm = try store.mlxArrayCopying(named: base + ".post_attention_layernorm.weight")
+        qANorm = try store.mlxArrayCopying(named: base + ".attention.q_a_layernorm.weight")
+        kvANorm = try store.mlxArrayCopying(named: base + ".attention.kv_a_layernorm.weight")
+        progress("norm tensors loaded")
+        qA = try Self.linear(store, base + ".attention.q_a_proj", nil)
+        progress("q_a loaded")
+        qB = try Self.linear(store, base + ".attention.q_b_proj", loraStore)
+        progress("q_b loaded")
+        kvA = try Self.linear(store, base + ".attention.kv_a_proj_with_mqa", nil)
+        progress("kv_a loaded")
+        kvB = try Self.linear(store, base + ".attention.kv_b_proj", loraStore)
+        progress("kv_b loaded")
+        gate = try Self.linear(store, base + ".attention.g_proj", loraStore)
+        progress("gate loaded")
+        dense = try Self.linear(store, base + ".attention.dense", loraStore)
+        progress("dense loaded")
+        moe = try StreamingMoE(configuration: c, store: store, layer: layer,
+                               bundleCache: bundleCache)
+        progress("moe loaded")
+        guard qA.inputSize == c.hiddenSize, qA.outputSize == c.qLoraRank,
+              qB.inputSize == c.qLoraRank, qB.outputSize == c.attentionHeads * queryHeadDim,
+              kvA.inputSize == c.hiddenSize, kvA.outputSize == c.kvLoraRank + c.qkRopeHeadDim,
+              kvB.inputSize == c.kvLoraRank,
+              kvB.outputSize == c.attentionHeads * (c.qkNopeHeadDim + c.vHeadDim),
+              gate.outputSize == c.attentionHeads,
+              dense.inputSize == c.attentionHeads * c.vHeadDim,
+              dense.outputSize == c.hiddenSize else {
+            throw M1Error.invalid("M2 attention tensor dimensions disagree with config")
+        }
+        eval(inputNorm, postAttentionNorm, qANorm, kvANorm)
+        progress("norm tensors evaluated")
+    }
+
+    private static func linear(_ store: ExpertTensorStore, _ prefix: String,
+                               _ loraStore: ExpertTensorStore?) throws -> ParallelLinear {
+        let base = try QuantizedExpertLinear(
+            weight: store.mlxArrayCopying(named: prefix + ".weight"),
+            scales: store.mlxArrayCopying(named: prefix + ".scales"),
+            biases: store.mlxArrayCopying(named: prefix + ".biases")
+        )
+        return try ParallelLinear(base: base, loraStore: loraStore, prefix: prefix)
+    }
+
+    public func callAsFunction(_ input: MLXArray, position: Int,
+                               cache: MLAKVCache? = nil,
+                               predictedRoute: PredictedMoERoute? = nil) throws -> DecoderBlockOutput {
+        let c = configuration
+        guard input.shape == [1, c.hiddenSize], position >= 0,
+              cache == nil || cache?.count == position else {
+            throw M1Error.invalid("M2 accepts one token [1,\(c.hiddenSize)] and a nonnegative position")
+        }
+        let normalized = MLXFast.rmsNorm(input, weight: inputNorm, eps: c.rmsNormEps)
+        let attention = try selfAttention(normalized, position: position, cache: cache)
+        let residual = input + attention
+        let moeInput = MLXFast.rmsNorm(residual, weight: postAttentionNorm, eps: c.rmsNormEps)
+        let moeOutput = if let predictedRoute {
+            try moe.predicted(moeInput, route: predictedRoute)
+        } else {
+            try moe(moeInput)
+        }
+        prerouterInput = moeInput
+        prerouterIndices = moeOutput.routeIndices
+        let hidden = residual + moeOutput.hidden
+        return DecoderBlockOutput(hidden: hidden, expertIndices: moeOutput.expertIndices,
+                                  expertWeights: moeOutput.expertWeights)
+    }
+
+    /// Official multi-token Edge0 8B MLA path used by chunked prefill.
+    public func prefill(_ input: MLXArray, position: Int,
+                        cache: MLAKVCache) throws -> DecoderBlockOutput {
+        let c = configuration
+        guard input.ndim == 2, input.dim(0) > 1, input.dim(1) == c.hiddenSize,
+              position >= 0, cache.count == position else {
+            throw M1Error.invalid("MLA prefill accepts [T,\(c.hiddenSize)], T > 1")
+        }
+        let x = input
+        let normalized = MLXFast.rmsNorm(x, weight: inputNorm, eps: c.rmsNormEps)
+        let attention = try selfAttentionPrefill(
+            normalized, position: position, cache: cache)
+        let residual = x + attention
+        let moeInput = MLXFast.rmsNorm(
+            residual, weight: postAttentionNorm, eps: c.rmsNormEps)
+        let moeOutput = try moe.fullLayerPrefill(moeInput)
+        prerouterInput = moeInput
+        prerouterIndices = moeOutput.routeIndices
+        let hidden = residual + moeOutput.hidden
+        return DecoderBlockOutput(hidden: hidden,
+                                  expertIndices: moeOutput.expertIndices,
+                                  expertWeights: moeOutput.expertWeights)
+    }
+
+    public func loadFullExpertLayer() throws { try moe.loadFullLayer() }
+    public func clearFullExpertLayer() { moe.clearFullLayer() }
+    public func stage(experts: [Int]) throws { try moe.stage(experts: experts) }
+    public func resetStaging() { moe.resetStaging() }
+    public var moeStats: StreamingMoEStats { moe.stats }
+
+    private func selfAttention(_ x: MLXArray, position: Int, cache: MLAKVCache?) throws -> MLXArray {
+        let c = configuration
+        let queryHeadDim = c.qkNopeHeadDim + c.qkRopeHeadDim
+        let qLatent = MLXFast.rmsNorm(try qA(x), weight: qANorm, eps: c.rmsNormEps)
+        let q = try qB(qLatent).reshaped([1, 1, c.attentionHeads, queryHeadDim])
+            .transposed(0, 2, 1, 3)
+        let qParts = split(q, indices: [c.qkNopeHeadDim], axis: -1)
+
+        let compressedParts = split(try kvA(x), indices: [c.kvLoraRank], axis: -1)
+        let kvLatent = MLXFast.rmsNorm(compressedParts[0], weight: kvANorm, eps: c.rmsNormEps)
+        let kv = try kvB(kvLatent)
+            .reshaped([1, 1, c.attentionHeads, c.qkNopeHeadDim + c.vHeadDim])
+            .transposed(0, 2, 1, 3)
+        let kvParts = split(kv, indices: [c.qkNopeHeadDim], axis: -1)
+        var kRope = compressedParts[1].reshaped([1, 1, 1, c.qkRopeHeadDim])
+
+        let qRope = interleavedRope(qParts[1], position: position)
+        kRope = interleavedRope(kRope, position: position)
+        kRope = broadcast(kRope, to: [1, c.attentionHeads, 1, c.qkRopeHeadDim])
+        let queries = concatenated([qParts[0], qRope], axis: -1)
+        var keys = concatenated([kvParts[0], kRope], axis: -1)
+        var values = kvParts[1]
+        if let cache {
+            (keys, values) = cache.append(keys: keys, values: values)
+        }
+        var output = MLXFast.scaledDotProductAttention(
+            queries: queries, keys: keys, values: values,
+            scale: 1.0 / Float(queryHeadDim).squareRoot(), mask: .none
+        )
+        output = output.transposed(0, 2, 1, 3).reshaped([1, c.attentionHeads, c.vHeadDim])
+        let gates = sigmoid(try gate(x)).reshaped([1, c.attentionHeads, 1])
+        output = (output * gates).reshaped([1, c.attentionHeads * c.vHeadDim])
+        return try dense(output)
+    }
+
+    private func interleavedRope(_ x: MLXArray, position: Int) -> MLXArray {
+        interleavedRope(x, startPosition: position, tokenCount: 1)
+    }
+
+    private func selfAttentionPrefill(_ x: MLXArray, position: Int,
+                                      cache: MLAKVCache) throws -> MLXArray {
+        let c = configuration
+        let tokens = x.dim(0)
+        let queryHeadDim = c.qkNopeHeadDim + c.qkRopeHeadDim
+        let qLatent = MLXFast.rmsNorm(
+            try qA(x), weight: qANorm, eps: c.rmsNormEps)
+        let q = try qB(qLatent)
+            .reshaped([1, tokens, c.attentionHeads, queryHeadDim])
+            .transposed(0, 2, 1, 3)
+        let qParts = split(q, indices: [c.qkNopeHeadDim], axis: -1)
+
+        let compressedParts = split(
+            try kvA(x), indices: [c.kvLoraRank], axis: -1)
+        let kvLatent = MLXFast.rmsNorm(
+            compressedParts[0], weight: kvANorm, eps: c.rmsNormEps)
+        let kv = try kvB(kvLatent)
+            .reshaped([1, tokens, c.attentionHeads,
+                       c.qkNopeHeadDim + c.vHeadDim])
+            .transposed(0, 2, 1, 3)
+        let kvParts = split(kv, indices: [c.qkNopeHeadDim], axis: -1)
+        var kRope = compressedParts[1]
+            .reshaped([1, 1, tokens, c.qkRopeHeadDim])
+
+        let qRope = interleavedRope(
+            qParts[1], startPosition: position, tokenCount: tokens)
+        kRope = interleavedRope(
+            kRope, startPosition: position, tokenCount: tokens)
+        kRope = broadcast(
+            kRope,
+            to: [1, c.attentionHeads, tokens, c.qkRopeHeadDim])
+        let queries = concatenated([qParts[0], qRope], axis: -1)
+        let newKeys = concatenated([kvParts[0], kRope], axis: -1)
+        let (keys, values) = cache.append(keys: newKeys, values: kvParts[1])
+        var output = MLXFast.scaledDotProductAttention(
+            queries: queries, keys: keys, values: values,
+            scale: 1.0 / Float(queryHeadDim).squareRoot(), mask: .causal)
+        output = output.transposed(0, 2, 1, 3)
+            .reshaped([tokens, c.attentionHeads, c.vHeadDim])
+        let gates = sigmoid(try gate(x))
+            .reshaped([tokens, c.attentionHeads, 1])
+        output = (output * gates)
+            .reshaped([tokens, c.attentionHeads * c.vHeadDim])
+        return try dense(output)
+    }
+
+    private func interleavedRope(_ x: MLXArray, startPosition: Int,
+                                 tokenCount: Int) -> MLXArray {
+        let pairs = configuration.qkRopeHeadDim / 2
+        let angles = (0..<tokenCount).flatMap { token in
+            (0..<pairs).map { i in
+                Float(startPosition + token)
+                    / pow(configuration.ropeTheta,
+                          Float(2 * i) / Float(configuration.qkRopeHeadDim))
+            }
+        }
+        let cosines = MLXArray(
+            angles.map { cos($0) }, [1, 1, tokenCount, pairs, 1])
+        let sines = MLXArray(
+            angles.map { sin($0) }, [1, 1, tokenCount, pairs, 1])
+        let pairParts = split(
+            x.reshaped([1, x.dim(1), tokenCount, pairs, 2]),
+            indices: [1], axis: -1)
+        let even = pairParts[0]
+        let odd = pairParts[1]
+        return concatenated(
+            [even * cosines - odd * sines, even * sines + odd * cosines],
+            axis: -1)
+            .reshaped(x.shape)
+    }
+}
diff --git a/ios/Sources/Edge0MLX/StreamingMoE.swift b/ios/Sources/Edge0MLX/StreamingMoE.swift
new file mode 100644
index 0000000..201cb7b
--- /dev/null
+++ b/ios/Sources/Edge0MLX/StreamingMoE.swift
@@ -0,0 +1,475 @@
+import Edge0Core
+import Foundation
+import MLX
+
+public struct MoEOutput {
+    public let hidden: MLXArray
+    public let expertIndices: [Int]
+    public let expertWeights: [Float]
+    public let routeIndices: MLXArray
+}
+
+public struct PredictedMoERoute: Sendable {
+    public let expertIndices: [Int]
+    public let expertWeights: [Float]
+
+    public init(expertIndices: [Int], expertWeights: [Float]) {
+        self.expertIndices = expertIndices
+        self.expertWeights = expertWeights
+    }
+}
+
+/// Official-style whole-model expert-bundle LRU. The production profile uses
+/// 64 slots globally, not per layer.
+public final class ExpertBundleCache: @unchecked Sendable {
+    private let capacity: Int
+    private var values: [String: QuantizedExpert] = [:]
+    private var order: [String] = []
+    private let lock = NSLock()
+    private var hitCount = 0
+    private var missCount = 0
+
+    public init(capacity: Int = 64) { self.capacity = max(0, capacity) }
+
+    func get(layer: Int, expert: Int) -> QuantizedExpert? {
+        lock.lock(); defer { lock.unlock() }
+        let key = "\(layer):\(expert)"
+        guard let value = values[key] else {
+            missCount += 1
+            return nil
+        }
+        hitCount += 1
+        order.removeAll { $0 == key }
+        order.append(key)
+        return value
+    }
+
+    func put(_ value: QuantizedExpert, layer: Int, expert: Int) {
+        lock.lock(); defer { lock.unlock() }
+        guard capacity > 0 else { return }
+        let key = "\(layer):\(expert)"
+        values[key] = value
+        order.removeAll { $0 == key }
+        order.append(key)
+        while order.count > capacity {
+            values.removeValue(forKey: order.removeFirst())
+        }
+    }
+
+    public var stats: (hits: Int, misses: Int, size: Int) {
+        lock.lock(); defer { lock.unlock() }
+        return (hitCount, missCount, values.count)
+    }
+}
+
+public struct StreamingMoEStats: Sendable {
+    public let calls: Int
+    public let loads: Int
+    public let hits: Int
+    public let stagedUsed: Int
+    public let stagedFallback: Int
+    public let stageBuilds: Int
+    public let stageCacheHits: Int
+    public let stageWallSeconds: TimeInterval
+}
+
+private final class ExpertLoadSlots: @unchecked Sendable {
+    private let lock = NSLock()
+    private var values: [QuantizedExpert?]
+    private var firstError: Error?
+
+    init(count: Int) {
+        values = Array(repeating: nil, count: count)
+    }
+
+    func store(_ result: Result<QuantizedExpert, Error>, at index: Int) {
+        lock.lock(); defer { lock.unlock() }
+        switch result {
+        case .success(let expert): values[index] = expert
+        case .failure(let error): if firstError == nil { firstError = error }
+        }
+    }
+
+    func finish() throws -> [QuantizedExpert] {
+        lock.lock(); defer { lock.unlock() }
+        if let firstError { throw firstError }
+        guard values.allSatisfy({ $0 != nil }) else {
+            throw M1Error.invalid("Parallel expert load did not fill every slot")
+        }
+        return values.map { $0! }
+    }
+}
+
+/// One decode token's feed-forward block, including the unconditional shared expert.
+/// It includes the official gathered exact path, global cache and prerouter staging paths.
+public final class StreamingMoE: @unchecked Sendable {
+    private let configuration: Edge0Configuration8B
+    private let store: ExpertTensorStore
+    private let layer: Int
+    private let routerWeight: MLXArray
+    private let expertBias: MLXArray
+    private let shared: QuantizedExpert
+    private let bundleCache: ExpertBundleCache?
+    private var fullLayer: QuantizedExpertStack?
+    private var stagedBundles: [Int: QuantizedExpert] = [:]
+    private var stagedStack: QuantizedExpertStack?
+    private var stagedOrder: [Int] = []
+    private let stagingLock = NSLock()
+    private var pendingOrder: [Int] = []
+    private var pendingBundles: [QuantizedExpert]?
+    private var pendingError: Error?
+    private var calls = 0
+    private var loads = 0
+    private var hits = 0
+    private var stagedUsed = 0
+    private var stagedFallback = 0
+    private var stageBuilds = 0
+    private var stageCacheHits = 0
+    private var stageWall: TimeInterval = 0
+    private let compiledStagedMath: @Sendable ([MLXArray]) -> [MLXArray]
+
+    public init(configuration c: Edge0Configuration8B, store: ExpertTensorStore,
+                layer: Int, bundleCache: ExpertBundleCache? = nil) throws {
+        guard c.quantization.bits == 4, c.quantization.groupSize == 64,
+              c.quantization.mode == "affine", c.numSharedExperts == 1,
+              layer >= c.firstKDenseReplace, layer < c.hiddenLayers,
+              c.numExperts == store.expertCount, c.nGroup > 0,
+              c.numExperts % c.nGroup == 0, c.numExperts / c.nGroup >= 2,
+              c.topkGroup > 0, c.topkGroup <= c.nGroup,
+              c.expertsPerToken > 0,
+              c.expertsPerToken <= c.topkGroup * (c.numExperts / c.nGroup) else {
+            throw M1Error.invalid("M1 requires an affine 4-bit/group-64 MoE layer with one shared expert and valid grouped routing")
+        }
+        configuration = c; self.store = store; self.layer = layer
+        self.bundleCache = bundleCache
+        compiledStagedMath = compile { arrays in
+            let x = arrays[0], local = arrays[1]
+            let up = gatherQuantizedMM(
+                x, arrays[2], scales: arrays[3], biases: arrays[4],
+                rhsIndices: local, transpose: true,
+                groupSize: 64, bits: 4, mode: .affine,
+                sortedIndices: false)
+            let gate = gatherQuantizedMM(
+                x, arrays[5], scales: arrays[6], biases: arrays[7],
+                rhsIndices: local, transpose: true,
+                groupSize: 64, bits: 4, mode: .affine,
+                sortedIndices: false)
+            let activated = (gate * sigmoid(gate)) * up
+            return [gatherQuantizedMM(
+                activated, arrays[8], scales: arrays[9], biases: arrays[10],
+                rhsIndices: local, transpose: true,
+                groupSize: 64, bits: 4, mode: .affine,
+                sortedIndices: false)]
+        }
+        routerWeight = try store.mlxArrayCopying(named: "model.layers.\(layer).mlp.gate.weight").asType(.float32)
+        expertBias = c.routerHasExpertBias
+            ? try store.mlxArrayCopying(named: "model.layers.\(layer).mlp.gate.expert_bias").asType(.float32)
+            : MLXArray.zeros([c.numExperts])
+        shared = try store.loadSharedExpert(layer: layer)
+        guard routerWeight.shape == [c.numExperts, c.hiddenSize], expertBias.shape == [c.numExperts],
+              shared.up.inputSize == c.hiddenSize,
+              shared.up.outputSize == c.sharedExpertIntermediateSize else {
+            throw M1Error.invalid("Router/shared expert dimensions disagree with config")
+        }
+        eval(routerWeight, expertBias)
+    }
+
+    public func callAsFunction(_ input: MLXArray) throws -> MoEOutput {
+        stagingLock.lock(); calls += 1; stagingLock.unlock()
+        let c = configuration
+        guard input.shape == [1, c.hiddenSize],
+              [.float32, .float16, .bfloat16].contains(input.dtype) else {
+            throw M1Error.invalid("M1 MoE accepts exactly one floating token [1,\(c.hiddenSize)]")
+        }
+        let x = input
+        let route = teacherRoute(x.asType(.float32))
+        eval(route.indices, route.weights)
+        let indices = route.indices.asArray(Int32.self).map(Int.init)
+        let weights = route.weights.asArray(Float.self)
+        let result = try executeGathered(x, indices: indices, weights: weights)
+        return MoEOutput(hidden: result, expertIndices: indices,
+                         expertWeights: weights, routeIndices: route.indices)
+    }
+
+    /// Consumer-layer path from the official Ling prerouter. Current
+    /// production configuration feeds the executed top-k to the next head.
+    public func predicted(_ input: MLXArray, route: PredictedMoERoute) throws -> MoEOutput {
+        stagingLock.lock(); calls += 1; stagingLock.unlock()
+        try activatePendingStage(for: route.expertIndices)
+        let x = input
+        let hidden: MLXArray
+        if Set(route.expertIndices) == Set(stagedOrder), let stagedStack {
+            stagingLock.lock(); stagedUsed += route.expertIndices.count; stagingLock.unlock()
+            hidden = try executeStaged(
+                x, stack: stagedStack, indices: route.expertIndices,
+                weights: route.expertWeights)
+        } else {
+            stagingLock.lock(); stagedFallback += 1; stagingLock.unlock()
+            hidden = try executeGathered(
+                x, indices: route.expertIndices, weights: route.expertWeights)
+        }
+        // Current upstream Edge0-8B uses feature_topk="executed": the
+        // prerouter's selected set feeds the next head directly. Do not
+        // recompute the original gate on consumer layers.
+        let executed = MLXArray(
+            route.expertIndices.map(Int32.init),
+            [1, route.expertIndices.count])
+        return MoEOutput(hidden: hidden, expertIndices: route.expertIndices,
+                         expertWeights: route.expertWeights,
+                         routeIndices: executed)
+    }
+
+    /// Official `prod_k8` next-state submission (`staged_sync=True`): fill the
+    /// NEXT set synchronously and promote it only at the next consume boundary.
+    public func stage(experts: [Int]) throws {
+        let ordered = Array(Set(experts)).sorted()
+        let started = Date()
+        var cacheHits = 0
+        var bundles = [QuantizedExpert?](repeating: nil, count: ordered.count)
+        var missingIDs: [Int] = []
+        var missingPositions: [Int] = []
+        store.adviseExpertsWillNeed(layer: layer, experts: ordered)
+        for (position, id) in ordered.enumerated() {
+            if let existing = stagedBundles[id] {
+                cacheHits += 1
+                bundles[position] = existing
+            } else if let cached = bundleCache?.get(layer: layer, expert: id) {
+                cacheHits += 1
+                bundles[position] = cached
+            } else {
+                missingIDs.append(id)
+                missingPositions.append(position)
+            }
+        }
+        let loaded = try loadExpertsFromStore(missingIDs)
+        for ((id, position), expert) in zip(
+            zip(missingIDs, missingPositions), loaded) {
+            bundles[position] = expert
+            bundleCache?.put(expert, layer: layer, expert: id)
+        }
+        let resolved = bundles.map { $0! }
+        let loadCount = missingIDs.count
+        stagingLock.lock()
+        pendingOrder = ordered
+        pendingBundles = resolved
+        pendingError = nil
+        stageBuilds += loadCount
+        stageCacheHits += cacheHits
+        loads += loadCount
+        hits += cacheHits
+        stageWall += Date().timeIntervalSince(started)
+        stagingLock.unlock()
+    }
+
+    private func activatePendingStage(for expected: [Int]) throws {
+        stagingLock.lock()
+        defer { stagingLock.unlock() }
+        if let pendingError { throw pendingError }
+        guard pendingOrder == Array(Set(expected)).sorted(),
+              let bundles = pendingBundles else { return }
+        if pendingOrder == stagedOrder, stagedStack != nil {
+            // Upstream asm-cache behavior: an unchanged staged set reuses
+            // the existing nine stacked graph nodes; only route weights vary.
+            pendingBundles = nil
+            pendingOrder = []
+            return
+        }
+        stagedBundles = Dictionary(uniqueKeysWithValues: zip(pendingOrder, bundles))
+        stagedStack = try QuantizedExpertStack(experts: bundles)
+        stagedOrder = pendingOrder
+        pendingBundles = nil
+        pendingOrder = []
+    }
+
+    public func resetStaging() {
+        stagingLock.lock()
+        pendingBundles = nil
+        pendingError = nil
+        pendingOrder = []
+        stagedBundles.removeAll(keepingCapacity: true)
+        stagedStack = nil
+        stagedOrder = []
+        stagingLock.unlock()
+        clearFullLayer()
+    }
+
+    private func teacherRoute(_ x: MLXArray) -> (indices: MLXArray, weights: MLXArray) {
+        let logits = matmul(x, routerWeight.T)
+        return Edge0GroupedRouter8B.select(
+            logits: logits, expertBias: expertBias,
+            topK: configuration.expertsPerToken, nGroup: configuration.nGroup,
+            topkGroup: configuration.topkGroup,
+            normalize: configuration.normTopkProb,
+            routedScale: configuration.routedScalingFactor)
+    }
+
+    private func resolveExpert(_ id: Int) throws -> QuantizedExpert {
+        if let staged = stagedBundles[id] {
+            stagingLock.lock(); hits += 1; stagingLock.unlock()
+            return staged
+        }
+        if let cached = bundleCache?.get(layer: layer, expert: id) {
+            stagingLock.lock(); hits += 1; stagingLock.unlock()
+            return cached
+        }
+        let expert = try store.loadExpert(layer: layer, expert: id)
+        bundleCache?.put(expert, layer: layer, expert: id)
+        stagingLock.lock(); loads += 1; stagingLock.unlock()
+        return expert
+    }
+
+    /// Match the upstream loader pool: resolve independent expert bundles in
+    /// parallel after issuing page-cache readahead for their tensor ranges.
+    private func resolveExperts(_ ids: [Int]) throws -> [QuantizedExpert] {
+        guard !ids.isEmpty else { return [] }
+        store.adviseExpertsWillNeed(layer: layer, experts: ids)
+        let slots = ExpertLoadSlots(count: ids.count)
+        DispatchQueue.concurrentPerform(iterations: ids.count) { index in
+            slots.store(Result { try self.resolveExpert(ids[index]) }, at: index)
+        }
+        return try slots.finish()
+    }
+
+    private func loadExpertsFromStore(_ ids: [Int]) throws -> [QuantizedExpert] {
+        guard !ids.isEmpty else { return [] }
+        let slots = ExpertLoadSlots(count: ids.count)
+        DispatchQueue.concurrentPerform(iterations: ids.count) { index in
+            slots.store(
+                Result { try self.store.loadExpert(layer: self.layer,
+                                                   expert: ids[index]) },
+                at: index)
+        }
+        return try slots.finish()
+    }
+
+    public var stats: StreamingMoEStats {
+        stagingLock.lock(); defer { stagingLock.unlock() }
+        return StreamingMoEStats(
+            calls: calls, loads: loads, hits: hits, stagedUsed: stagedUsed,
+            stagedFallback: stagedFallback, stageBuilds: stageBuilds,
+            stageCacheHits: stageCacheHits, stageWallSeconds: stageWall)
+    }
+
+    /// Upstream exact path: stack the unique routed bundles and execute the
+    /// same compiled gathered-QMM graph used by fixed-slot decode.
+    private func executeGathered(_ x: MLXArray, indices: [Int],
+                                 weights: [Float]) throws -> MLXArray {
+        let c = configuration
+        guard indices.count == weights.count else {
+            throw M1Error.invalid("Expert IDs and weights disagree")
+        }
+        let unique = Array(Set(indices)).sorted()
+        let experts = try resolveExperts(unique).map { expert in
+            guard expert.up.inputSize == c.hiddenSize,
+                  expert.up.outputSize == c.moeIntermediateSize else {
+                throw M1Error.invalid("Routed expert dimensions disagree with config")
+            }
+            return expert
+        }
+        let localByID = Dictionary(uniqueKeysWithValues:
+            unique.enumerated().map { ($0.element, $0.offset) })
+        return try executeStacked(
+            x, stack: QuantizedExpertStack(experts: experts),
+            localIndices: indices.map { localByID[$0]! }, weights: weights)
+    }
+
+    /// Direct Swift translation of the official fixed staged-slot math. The
+    /// selected K experts map into canonical staged slots and run through three
+    /// gathered INT4 matrix multiplications instead of K independent graphs.
+    private func executeStaged(_ x: MLXArray, stack: QuantizedExpertStack,
+                               indices: [Int],
+                               weights: [Float]) throws -> MLXArray {
+        let k = indices.count
+        guard k == weights.count, stagedOrder.count == stack.up.expertCount else {
+            throw M1Error.invalid("Staged expert stack and route disagree")
+        }
+        let slotByID = Dictionary(uniqueKeysWithValues:
+            stagedOrder.enumerated().map { ($0.element, $0.offset) })
+        return try executeStacked(
+            x, stack: stack, localIndices: indices.map { slotByID[$0]! },
+            weights: weights)
+    }
+
+    private func executeStacked(_ x: MLXArray, stack: QuantizedExpertStack,
+                                localIndices: [Int], weights: [Float]) throws
+        -> MLXArray {
+        let k = localIndices.count
+        guard k == weights.count, !localIndices.isEmpty,
+              localIndices.allSatisfy({
+                  (0..<stack.up.expertCount).contains($0)
+              }) else {
+            throw M1Error.invalid("Stacked expert indices and route disagree")
+        }
+        let local = MLXArray(localIndices.map(Int32.init))
+        let routedInput = broadcast(x, to: [k, configuration.hiddenSize])
+            .expandedDimensions(axis: 1)
+        let outputs = compiledStagedMath([
+            routedInput, local,
+            stack.up.weight, stack.up.scales, stack.up.biases,
+            stack.gate.weight, stack.gate.scales, stack.gate.biases,
+            stack.down.weight, stack.down.scales, stack.down.biases,
+        ])[0].squeezed(axis: 1)
+        let routedWeights = MLXArray(weights, [k, 1]).asType(outputs.dtype)
+        let routed = (outputs * routedWeights).sum(axis: 0, keepDims: true)
+        return try routed + shared(x)
+    }
+
+    /// Official whole-layer prefill lifecycle (`load_full_layer`).
+    public func loadFullLayer() throws {
+        if fullLayer == nil {
+            fullLayer = try store.loadFullExpertLayer(layer: layer)
+        }
+    }
+
+    /// Official whole-layer prefill lifecycle (`clear_full_layer`).
+    public func clearFullLayer() {
+        fullLayer = nil
+    }
+
+    public var hasFullLayer: Bool { fullLayer != nil }
+
+    /// Official full-layer `StreamingSwitchGLU.__call__` branch for a flat
+    /// token batch `[T, H]`. The sorted gather/unsort sequence is kept in the
+    /// same order as mlx-lm's `_gather_sort` and `_scatter_unsort`.
+    public func fullLayerPrefill(_ input: MLXArray) throws -> MoEOutput {
+        let c = configuration
+        guard input.ndim == 2, input.dim(0) > 1, input.dim(1) == c.hiddenSize,
+              let fullLayer else {
+            throw M1Error.invalid("Full-layer prefill requires loaded weights and [T,H], T > 1")
+        }
+        let x = input
+        let logits = matmul(x.asType(.float32), routerWeight.T)
+        let route = Edge0GroupedRouter8B.select(
+            logits: logits, expertBias: expertBias,
+            topK: c.expertsPerToken, nGroup: c.nGroup,
+            topkGroup: c.topkGroup, normalize: c.normTopkProb,
+            routedScale: c.routedScalingFactor)
+
+        let tokenCount = input.dim(0)
+        let topK = c.expertsPerToken
+        let flatIndices = route.indices.flattened()
+        let order = argSort(flatIndices)
+        let inverseOrder = argSort(order)
+        let tokenRows = order.floorDivide(topK)
+        let sortedInput = x.take(tokenRows, axis: 0).expandedDimensions(axis: 1)
+        let sortedExperts = flatIndices.take(order, axis: 0)
+
+        let up = try fullLayer.up(sortedInput, expertIndices: sortedExperts)
+        let gate = try fullLayer.gate(sortedInput, expertIndices: sortedExperts)
+        let activated = (gate * sigmoid(gate)) * up
+        let sortedOutput = try fullLayer.down(activated, expertIndices: sortedExperts)
+        let expertOutput = sortedOutput.take(inverseOrder, axis: 0)
+            .reshaped([tokenCount, topK, c.hiddenSize])
+        let routedWeights = route.weights.asType(expertOutput.dtype)
+            .expandedDimensions(axis: -1)
+        let routed = (expertOutput * routedWeights)
+            .sum(axis: 1)
+        let hidden = try routed + shared(x)
+        eval(hidden, route.indices, route.weights)
+        return MoEOutput(
+            hidden: hidden,
+            expertIndices: route.indices.asArray(Int32.self).map(Int.init),
+            expertWeights: route.weights.asArray(Float.self),
+            routeIndices: route.indices)
+    }
+}
diff --git a/ios/Tests/Edge0CoreTests/Edge0Tokenizer8BTests.swift b/ios/Tests/Edge0CoreTests/Edge0Tokenizer8BTests.swift
new file mode 100644
index 0000000..fa6800e
--- /dev/null
+++ b/ios/Tests/Edge0CoreTests/Edge0Tokenizer8BTests.swift
@@ -0,0 +1,12 @@
+import Edge0Core
+import Foundation
+import Testing
+
+@Test func edge0Tokenizer8BMatchesOfficialHiPrompt() throws {
+    let environment = ProcessInfo.processInfo.environment
+    guard let modelPath = environment["EDGE0_MODEL"] else { return }
+    let tokenizer = try Edge0Tokenizer8B(contentsOf: URL(fileURLWithPath: modelPath).appendingPathComponent("tokenizer.json"))
+    let prompt = "<role>SYSTEM</role>detailed thinking off<|role_end|><role>HUMAN</role>Hi<|role_end|><role>ASSISTANT</role>\n<think></think>"
+    #expect(try tokenizer.encode(prompt) == [157151, 90827, 157152, 14136, 5381, 6350, 928, 156895, 157151, 39, 116171, 157152, 10754, 156895, 157151, 8469, 7342, 5468, 157152, 198, 156903, 156904])
+    #expect(try tokenizer.decode([10754, 0]) == "Hi!")
+}
diff --git a/ios/Tests/Edge0CoreTests/RouterReferenceTests.swift b/ios/Tests/Edge0CoreTests/RouterReferenceTests.swift
new file mode 100644
index 0000000..2d2a292
--- /dev/null
+++ b/ios/Tests/Edge0CoreTests/RouterReferenceTests.swift
@@ -0,0 +1,19 @@
+import Edge0Core
+import Testing
+
+@Test func groupedRouterDropsWeakGroupsAndNormalizes() {
+    // 8 experts, 4 groups, 2 experts/group. Keep 2 groups, choose top 2 experts.
+    let logits: [Float] = [-4, -3, 5, 4, 1, 0, 3, 2]
+    let bias = Array(repeating: Float(0), count: 8)
+    let result = GroupedExpertRouter.select(
+        logits: logits,
+        expertBias: bias,
+        topK: 2,
+        nGroup: 4,
+        topkGroup: 2,
+        normalize: true,
+        routedScale: 2.5
+    )
+    #expect(Set(result.indices) == Set([2, 3]))
+    #expect(abs(result.weights.reduce(0, +) - 2.5) < 0.0001)
+}
diff --git a/ios/Tests/Edge0CoreTests/SafetensorsIndexTests.swift b/ios/Tests/Edge0CoreTests/SafetensorsIndexTests.swift
new file mode 100644
index 0000000..139656d
--- /dev/null
+++ b/ios/Tests/Edge0CoreTests/SafetensorsIndexTests.swift
@@ -0,0 +1,32 @@
+import Edge0Core
+import Foundation
+import Testing
+
+@Test func parsesSafetensorAndSlicesAxisZero() throws {
+    let tmp = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString + ".safetensors")
+    defer { try? FileManager.default.removeItem(at: tmp) }
+
+    let headerObject: [String: Any] = [
+        "x": [
+            "dtype": "U32",
+            "shape": [2, 2],
+            "data_offsets": [0, 16],
+        ]
+    ]
+    var header = try JSONSerialization.data(withJSONObject: headerObject, options: [.sortedKeys])
+    // Safetensors permits JSON whitespace padding.
+    while header.count % 8 != 0 { header.append(0x20) }
+    var headerLength = UInt64(header.count).littleEndian
+    var file = withUnsafeBytes(of: &headerLength) { Data($0) }
+    file.append(header)
+    file.append(Data(repeating: 0xAB, count: 16))
+    try file.write(to: tmp)
+
+    let index = try SafetensorsIndex(fileURL: tmp)
+    let tensor = try index.tensor(named: "x")
+    #expect(tensor.shape == [2, 2])
+    #expect(tensor.byteCount == 16)
+    let second = try tensor.axisZeroSlice(1)
+    #expect(second.shape == [2])
+    #expect(second.byteCount == 8)
+}
diff --git a/ios/Tests/Edge0MLXTests/Edge035BChatTests.swift b/ios/Tests/Edge0MLXTests/Edge035BChatTests.swift
new file mode 100644
index 0000000..1b52e7d
--- /dev/null
+++ b/ios/Tests/Edge0MLXTests/Edge035BChatTests.swift
@@ -0,0 +1,45 @@
+@testable import Edge0MLX
+import Testing
+
+@Test func edge35bFirstTurnTemplateMatchesCheckpointContract() throws {
+    let rendered = try ChatTemplate.render([
+        .init(role: "system", content: "helpful"),
+        .init(role: "user", content: "hello"),
+    ])
+    #expect(rendered == "<|im_start|>system\nhelpful<|im_end|>\n" +
+        "<|im_start|>user\nhello<|im_end|>\n" +
+        "<|im_start|>assistant\n<think>\n")
+}
+
+@Test func edge35bFirstTurnTemplateCanDisableThinking() throws {
+    let rendered = try ChatTemplate.render([
+        .init(role: "system", content: "helpful"),
+        .init(role: "user", content: "hello"),
+    ], thinking: false)
+    #expect(rendered == "<|im_start|>system\nhelpful<|im_end|>\n" +
+        "<|im_start|>user\nhello<|im_end|>\n" +
+        "<|im_start|>assistant\n")
+}
+
+@Test func edge35bContinuationDoesNotReplayHistory() {
+    #expect(ChatTemplate.continuation(user: "again") ==
+        "<|im_end|>\n<|im_start|>user\nagain<|im_end|>\n" +
+        "<|im_start|>assistant\n<think>\n")
+    #expect(ChatTemplate.continuation(user: "again", thinking: false) ==
+        "<|im_end|>\n<|im_start|>user\nagain<|im_end|>\n" +
+        "<|im_start|>assistant\n")
+}
+
+@Test func edge35bStopsOnBothCheckpointEndTokens() {
+    let stop = StopCondition(maximumTokens: 3)
+    #expect(stop.evaluate(token: 248046, produced: 0) == .hitEndToken(248046))
+    #expect(stop.evaluate(token: 248044, produced: 0) == .hitEndToken(248044))
+    #expect(stop.evaluate(token: 42, produced: 3) == .hitLimit)
+    #expect(stop.evaluate(token: 42, produced: 2) == .keepGoing)
+}
+
+@Test func edge35bHidesReasoningBlock() {
+    #expect(ThinkingBlock.visible(in: "private notes") == nil)
+    #expect(ThinkingBlock.visible(in: "private</think>\n\nanswer") == "answer")
+    #expect(ThinkingBlock.settled("unfinished reasoning") == "unfinished reasoning")
+}
diff --git a/ios/Tests/Edge0MLXTests/Fixtures/config.json b/ios/Tests/Edge0MLXTests/Fixtures/config.json
new file mode 100644
index 0000000..903835b
--- /dev/null
+++ b/ios/Tests/Edge0MLXTests/Fixtures/config.json
@@ -0,0 +1,105 @@
+{
+  "attention_dropout": 0.0,
+  "dtype": "bfloat16",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 156895,
+  "expert_swiglu_limit_list": null,
+  "first_k_dense_replace": 1,
+  "gated_attention_proj_granularity_type": "head_wise",
+  "group_norm_size": 1,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 4608,
+  "kda_lower_bound": -5,
+  "kda_safe_gate": true,
+  "kv_lora_rank": 512,
+  "layer_group_size": 4,
+  "linear_silu": true,
+  "max_position_embeddings": 131072,
+  "max_window_layers": 20,
+  "moe_intermediate_size": 512,
+  "moe_router_enable_expert_bias": true,
+  "moe_shared_expert_intermediate_size": 512,
+  "mtp_loss_scaling_factor": 0,
+  "mtp_use_kda": false,
+  "n_group": 8,
+  "no_kda_lora": true,
+  "norm_topk_prob": true,
+  "num_attention_heads": 16,
+  "num_experts": 128,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 16,
+  "num_kv_heads_for_linear_attn": 0,
+  "num_nextn_predict_layers": 0,
+  "num_shared_experts": 1,
+  "output_dropout": 0.0,
+  "output_router_logits": false,
+  "pad_token_id": 156892,
+  "partial_rotary_factor": 0.5,
+  "pregate_cross_token": true,
+  "pregate_enabled": true,
+  "pregate_hidden": 512,
+  "pregate_inference": false,
+  "pregate_init_router": false,
+  "pregate_opd_ce_weight": 1.0,
+  "pregate_opd_strategy": "none",
+  "pregate_opd_temperature": 1.0,
+  "pregate_opd_topk": 8,
+  "pregate_opd_weight": 1.0,
+  "pregate_opd_weight_mode": "teacher_p",
+  "pregate_shallow_hidden": 1024,
+  "pregate_shallow_layers": 5,
+  "pregate_shallow_loss_weight": 1.5,
+  "pregate_start_layer": 7,
+  "pregate_use_prev_token": true,
+  "pregate_use_prev_topk": true,
+  "q_lora_rank": 256,
+  "qk_head_dim": 192,
+  "qk_nope_head_dim": 128,
+  "qk_rope_head_dim": 64,
+  "rms_norm_eps": 1e-06,
+  "rope_interleave": true,
+  "rope_parameters": {
+    "partial_rotary_factor": 0.5,
+    "rope_theta": 6000000,
+    "rope_type": "default"
+  },
+  "rope_theta": 6000000,
+  "rotary_dim": 64,
+  "routed_scaling_factor": 2.5,
+  "router_dtype": "fp32",
+  "scale_router_input": false,
+  "score_function": "sigmoid",
+  "scoring_func": "sigmoid",
+  "seq_aux": true,
+  "share_expert_swiglu_limit_list": null,
+  "short_conv_kernel_size": 4,
+  "tie_word_embeddings": false,
+  "topk_group": 4,
+  "topk_method": "noaux_tc",
+  "transformers_version": "5.8.1",
+  "up_proj_norm": false,
+  "use_bias": false,
+  "use_cache": false,
+  "use_kda_lora": false,
+  "use_mla_nope": false,
+  "use_nGPT": false,
+  "use_qk_norm": true,
+  "use_qkv_bias": false,
+  "v_head_dim": 128,
+  "value_norm": false,
+  "vocab_size": 157184,
+  "quantization": {
+    "group_size": 64,
+    "bits": 4,
+    "mode": "affine"
+  },
+  "quantization_config": {
+    "group_size": 64,
+    "bits": 4,
+    "mode": "affine"
+  }
+}
diff --git a/ios/Tests/Edge0MLXTests/M1Tests.swift b/ios/Tests/Edge0MLXTests/M1Tests.swift
new file mode 100644
index 0000000..6d03b04
--- /dev/null
+++ b/ios/Tests/Edge0MLXTests/M1Tests.swift
@@ -0,0 +1,161 @@
+import Edge0Core
+@testable import Edge0MLX
+import Foundation
+import MLX
+import Testing
+
+// CPU arithmetic by default; MLX still requires an Xcode-built Metal bundle.
+// EDGE0_TEST_GPU=1 is used by the documented xcodebuild verification.
+private func onDevice(_ body: () throws -> Void) rethrows {
+    try Device.withDefaultDevice(ProcessInfo.processInfo.environment["EDGE0_TEST_GPU"] == "1" ? .gpu : .cpu, body)
+}
+
+@Test func packedAffineLinearMatchesScalarOracle() throws {
+    try onDevice {
+        // Two groups, all nibble positions, nonzero offsets, non-square matrix, two tokens.
+        let width = 128, rows = 64
+        var packed = [UInt32](repeating: 0, count: rows * width / 8)
+        var scales = [Float](), biases = [Float]()
+        for o in 0..<rows {
+            scales += [Float(o % 3 + 1) / 32, Float(o % 5 + 1) / 64]
+            biases += [-0.25, 0.125]
+            for i in 0..<width { packed[o * width / 8 + i / 8] |= UInt32((o * 3 + i) % 16) << (4 * (i % 8)) }
+        }
+        let x = (0..<(width * 2)).map { Float($0 % 19 - 9) / 16 }
+        let linear = try QuantizedExpertLinear(weight: MLXArray(packed, [rows, width/8]),
+            scales: MLXArray(scales, [rows, 2]).asType(.bfloat16),
+            biases: MLXArray(biases, [rows, 2]).asType(.bfloat16))
+        let y = try linear(MLXArray(x, [2, width])).asArray(Float.self)
+        for t in 0..<2 {
+            for o in 0..<rows {
+                var expected: Float = 0
+                for i in 0..<width {
+                    let q = Float((o * 3 + i) % 16)
+                    expected += x[t * width + i] * (q * scales[o*2+i/64] + biases[o*2+i/64])
+                }
+                #expect(abs(y[t*rows+o] - expected) < 1e-5)
+            }
+        }
+        #expect(throws: M1Error.self) { try linear(MLXArray.zeros([1, 64])) }
+    }
+}
+
+@Test func malformedQuantizedShapesRejected() {
+    onDevice {
+        #expect(throws: M1Error.self) {
+            try QuantizedExpertLinear(weight: MLXArray.zeros([64, 8]),
+                scales: MLXArray.zeros([64, 1]), biases: MLXArray.zeros([64, 1]))
+        }
+        #expect(throws: M1Error.self) {
+            try QuantizedExpertLinear(weight: MLXArray.zeros([64, 8], dtype: .uint32),
+                scales: MLXArray.zeros([64, 2]), biases: MLXArray.zeros([64, 2]))
+        }
+    }
+}
+
+@Test func groupedMLXMatchesCPUWithSelectionBias() {
+    onDevice {
+        let logits: [Float] = [-4, -3, 5, 4, 1, 0, 3, 2]
+        let bias: [Float] = [3, 2, 0, 0, 0.2, 0, -1, -1]
+        for keep in [2, 4] {
+            for k in [1, 2] {
+                for normalize in [false, true] {
+                    let cpu = GroupedExpertRouter.select(logits: logits, expertBias: bias, topK: k,
+                        nGroup: 4, topkGroup: keep, normalize: normalize)
+                    let gpu = Edge0GroupedRouter8B.select(logits: MLXArray(logits, [1, 8]),
+                        expertBias: MLXArray(bias), topK: k, nGroup: 4, topkGroup: keep, normalize: normalize)
+                    let ids = gpu.indices.asArray(Int32.self).map(Int.init)
+                    let weights = gpu.weights.asArray(Float.self)
+                    #expect(Set(ids) == Set(cpu.indices))
+                    for (id, w) in zip(ids, weights) {
+                        #expect(abs(w - cpu.weights[cpu.indices.firstIndex(of: id)!]) < 1e-6)
+                    }
+                }
+            }
+        }
+    }
+}
+
+@Test func swigluUsesGateAndDownProjection() throws {
+    try onDevice {
+        // Identity INT4 projections make the expected result exactly x * sigmoid(x) * x.
+        var packed = [UInt32](repeating: 0, count: 64 * 8)
+        for i in 0..<64 { packed[i*8+i/8] = 1 << (4*(i%8)) }
+        let linear = try QuantizedExpertLinear(weight: MLXArray(packed, [64, 8]),
+            scales: MLXArray.ones([64, 1]), biases: MLXArray.zeros([64, 1]))
+        let expert = try QuantizedExpert(up: linear, gate: linear, down: linear)
+        let x = (0..<64).map { Float($0-32)/8 }
+        let actual = try expert(MLXArray(x, [1, 64])).asArray(Float.self)
+        for (v, y) in zip(x, actual) { #expect(abs(y - v*v/(1+exp(-v))) < 2e-6) }
+    }
+}
+
+@Test func streamedMoEIncludesSharedAndOnlySelectedExperts() throws {
+    try onDevice {
+        let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+        try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
+        defer { try? FileManager.default.removeItem(at: directory) }
+        var config = try JSONSerialization.jsonObject(with: Data(contentsOf:
+            Bundle.module.url(forResource: "config", withExtension: "json", subdirectory: "Fixtures")!)) as! [String: Any]
+        for key in ["hidden_size", "moe_intermediate_size", "moe_shared_expert_intermediate_size"] { config[key] = 64 }
+        config["num_experts"] = 8; config["n_group"] = 4
+        config["topk_group"] = 2; config["num_experts_per_tok"] = 2
+        let c = try JSONDecoder().decode(Edge0Configuration8B.self, from: JSONSerialization.data(withJSONObject: config))
+        var header = [String: Any](), payload = Data()
+        func append<T>(_ name: String, _ dtype: String, _ shape: [Int], _ values: [T]) {
+            let bytes = values.withUnsafeBytes { Data($0) }
+            header[name] = ["dtype": dtype, "shape": shape, "data_offsets": [payload.count, payload.count + bytes.count]]
+            payload.append(bytes)
+        }
+        var identity = [UInt32](repeating: 0, count: 64*8)
+        for i in 0..<64 { identity[i*8+i/8] = 1 << (4*(i%8)) }
+        let prefix = "model.layers.1.mlp"
+        let logits: [Float] = [-4, -3, -2, -1, 0, 1, 2, 3]
+        let bias: [Float] = [3, 2, 0, 0, 0, 0, -1, -1]
+        var router = [Float](repeating: 0, count: 8*64)
+        for i in 0..<8 { router[i*64] = logits[i] }
+        append(prefix + ".gate.weight", "F32", [8, 64], router)
+        append(prefix + ".gate.expert_bias", "F32", [8], bias)
+        let route = GroupedExpertRouter.select(logits: logits, expertBias: bias, topK: 2, nGroup: 4, topkGroup: 2)
+        for projection in ExpertProjection.allCases {
+            let name = prefix + ".experts." + projection.rawValue
+            append(name + ".weight", "U32", [8, 64, 8], Array(repeating: identity, count: 8).flatMap { $0 })
+            // Unselected experts deliberately contain NaN scales: their execution contaminates output.
+            let scales: [Float] = (0..<8).flatMap { id in
+                Array(repeating: route.indices.contains(id) ? (projection == .up ? Float(id+1)/8 : 1) : Float.nan, count: 64)
+            }
+            append(name + ".scales", "F32", [8, 64, 1], scales)
+            append(name + ".biases", "F32", [8, 64, 1], [Float](repeating: 0, count: 8*64))
+            let shared = prefix + ".shared_experts." + projection.rawValue
+            append(shared + ".weight", "U32", [64, 8], identity)
+            append(shared + ".scales", "F32", [64, 1], [Float](repeating: projection == .up ? 2 : 1, count: 64))
+            append(shared + ".biases", "F32", [64, 1], [Float](repeating: 0, count: 64))
+        }
+        var json = try JSONSerialization.data(withJSONObject: header, options: .sortedKeys)
+        while json.count % 8 != 0 { json.append(32) }
+        var length = UInt64(json.count).littleEndian
+        var file = withUnsafeBytes(of: &length) { Data($0) }
+        file.append(json); file.append(payload)
+        let url = directory.appendingPathComponent("model.safetensors")
+        try file.write(to: url)
+        let store = try ExpertTensorStore(modelURL: url, expertCount: 8)
+        let block = try StreamingMoE(configuration: c, store: store, layer: 1)
+        var input = (0..<64).map { Float($0-32)/16 }; input[0] = 1
+        let output = try block(MLXArray(input, [1, 64]))
+        #expect(Set(output.expertIndices) == Set(route.indices))
+        let multiplier = zip(route.indices, route.weights).reduce(Float(2)) { $0 + Float($1.0+1)/8 * $1.1 }
+        for (x, y) in zip(input, output.hidden.asArray(Float.self)) {
+            #expect(y.isFinite && abs(y - multiplier*x*x/(1+exp(-x))) < 3e-6)
+        }
+        #expect(throws: M1Error.self) { try block(MLXArray.zeros([2, 64])) }
+        #expect(throws: M1Error.self) { try store.loadExpert(layer: 1, expert: 8) }
+        #expect(throws: M1Error.self) { try StreamingMoE(configuration: c, store: store, layer: 0) }
+        // A copied expert is safe after the originating mapping closes.
+        let detached: QuantizedExpert = try {
+            let temporary = try ExpertTensorStore(modelURL: url, expertCount: 8)
+            return try temporary.loadExpert(layer: 1, expert: route.indices[0])
+        }()
+        let detachedY = try detached(MLXArray(input, [1,64])).asArray(Float.self)
+        #expect(detachedY.allSatisfy { $0.isFinite })
+    }
+}
diff --git a/ios/scripts/test.sh b/ios/scripts/test.sh
new file mode 100755
index 0000000..a405b74
--- /dev/null
+++ b/ios/scripts/test.sh
@@ -0,0 +1,17 @@
+#!/bin/sh
+# Run from any directory. Requires Xcode's Metal Toolchain component.
+set -eu
+cd "$(dirname "$0")/.."
+xcodebuild build -scheme Edge0MLX -destination 'platform=macOS' \
+  -derivedDataPath .build-xcode -skipPackagePluginValidation CODE_SIGNING_ALLOWED=NO
+# SwiftPM cannot build shaders. Stage Xcode's library alongside its test binary.
+swift build --build-tests
+library=.build-xcode/Build/Products/Debug/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib
+if [ ! -f "$library" ]; then
+  echo "Missing Xcode Metal library: $library" >&2
+  exit 1
+fi
+binary_dir=$(swift build --show-bin-path)
+cp "$library" "$binary_dir/mlx.metallib"
+cp "$library" "$binary_dir/Edge0PhonePackageTests.xctest/Contents/MacOS/mlx.metallib"
+swift test --skip-build
diff --git a/ios/tools/convert_pregate.py b/ios/tools/convert_pregate.py
new file mode 100755
index 0000000..20cfaf9
--- /dev/null
+++ b/ios/tools/convert_pregate.py
@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""Convert trained pregate tensors to a memory-mappable safetensors file.
+
+The converter stacks the per-layer heads, stores linear weights transposed for batched
+matrix multiplication, preserves their fp16 values, and records the owning layer IDs.
+It accepts an NPZ or safetensors source.
+
+Usage:
+    convert_pregate.py <input.npz|input.safetensors> <output.safetensors>
+"""
+
+import json
+import struct
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+def load_safetensors(path):
+    """Small NumPy-only reader so deployment does not require another package."""
+    with path.open("rb") as handle:
+        header_size = struct.unpack("<Q", handle.read(8))[0]
+        header = json.loads(handle.read(header_size))
+    base = 8 + header_size
+    dtype = {"F16": np.dtype("<f2"), "F32": np.dtype("<f4"), "I32": np.dtype("<i4")}
+    result = {}
+    for name, entry in header.items():
+        if name == "__metadata__":
+            continue
+        start, end = entry["data_offsets"]
+        result[name] = np.memmap(
+            path, mode="r", dtype=dtype[entry["dtype"]], offset=base + start,
+            shape=tuple(entry["shape"]), order="C")
+    return result
+
+
+def save_safetensors(tensors, path):
+    """Write the subset of safetensors used by these converted fp16/int32 heads."""
+    dtype = {np.dtype("<f2"): "F16", np.dtype("<f4"): "F32", np.dtype("<i4"): "I32"}
+    header = {}
+    offset = 0
+    for name, value in tensors.items():
+        value = np.ascontiguousarray(value)
+        tensors[name] = value
+        size = value.nbytes
+        header[name] = {
+            "dtype": dtype[value.dtype], "shape": list(value.shape),
+            "data_offsets": [offset, offset + size],
+        }
+        offset += size
+    encoded = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    encoded += b" " * ((8 - len(encoded) % 8) % 8)
+    temporary = path.with_suffix(path.suffix + ".partial")
+    with temporary.open("wb", buffering=4 << 20) as handle:
+        handle.write(struct.pack("<Q", len(encoded)))
+        handle.write(encoded)
+        for value in tensors.values():
+            handle.write(memoryview(value).cast("B"))
+    temporary.replace(path)
+
+
+def main():
+    if len(sys.argv) < 3:
+        raise SystemExit(__doc__)
+    source, destination = Path(sys.argv[1]), Path(sys.argv[2])
+
+    # Accept both supported checkpoint containers.
+    if source.suffix == ".safetensors":
+        npz = load_safetensors(source)
+    else:
+        npz = dict(np.load(source))
+    owners = sorted({int(name.split(".")[1]) for name in npz})
+    if not owners:
+        raise SystemExit(f"no `layers.N.*` tensors in {source}")
+
+    # Normalize the checkpoint's `prerouter` names to the runtime's `pregate` names.
+    prefix = "pregate"
+    parts = {}
+    for part in ("fc1", "fc2", "linear_init"):
+        stack = []
+        for owner in owners:
+            key = f"layers.{owner}.{part}.weight"
+            if key not in npz:
+                raise SystemExit(f"{source.name} is missing {key}")
+            # `[out, in]` -> `[in, out]`, contiguous. `ascontiguousarray` is not
+            # decoration: safetensors stores a flat row-major buffer and a transposed
+            # view is not one, so without it the file would either fail to write or, far
+            # worse, write the untransposed bytes under the transposed shape.
+            stack.append(np.ascontiguousarray(npz[key].T))
+        merged = np.stack(stack, axis=0)
+        if merged.dtype != np.float16:
+            raise SystemExit(f"{part} is {merged.dtype}, expected float16")
+        parts[f"{prefix}.{part}"] = merged
+
+    parts[f"{prefix}.owners"] = np.array(owners, dtype=np.int32)
+
+    save_safetensors(parts, destination)
+    total = sum(v.nbytes for v in parts.values())
+    print(f"  {destination}")
+    for name, value in parts.items():
+        print(f"    {name:24s} {str(list(value.shape)):20s} {value.dtype}")
+    print(f"  owners {owners[0]}..{owners[-1]}  n={len(owners)}  "
+          f"{total / 2**20:.1f} MiB")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/ios/tools/convert_tokenizer.py b/ios/tools/convert_tokenizer.py
new file mode 100755
index 0000000..3cb728b
--- /dev/null
+++ b/ios/tools/convert_tokenizer.py
@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Convert `tokenizer.json` into a compact binary the device can load quickly.
+
+The shipped file is 20 MB of JSON holding 248,044 vocabulary entries and 247,587 merge
+rules. Parsing that on a phone at launch is slow and allocates far more than the data
+needs. This rewrites it into the same shape the weights use — a small header plus flat
+arrays that can be read directly:
+
+    vocab_blob     all token strings concatenated, no separators
+    vocab_offsets  uint32[n+1], so token i is blob[offsets[i]..<offsets[i+1]]
+    merge_pairs    uint32[2m], the two token ids of each merge in rank order
+    specials       ids and strings of the 33 added tokens
+
+Merges are stored as **token id pairs rather than strings**, so the device never has to
+look a merge up by text: the BPE loop can key on a packed pair of ids. That is the single
+biggest difference from parsing the JSON directly, and it is why the merge table costs
+about 2 MB instead of 20.
+
+The pre-tokenizer regex is carried across verbatim rather than reimplemented. It is the
+GPT-4 pattern, and the one place where a subtle difference produces a tokenizer that works
+on English and quietly mis-splits CJK.
+
+Usage:
+    convert_tokenizer.py <model-dir> <out-file>
+"""
+
+import json
+import struct
+import sys
+from pathlib import Path
+
+
+def main():
+    if len(sys.argv) < 3:
+        raise SystemExit(__doc__)
+    model_dir, out_path = Path(sys.argv[1]), Path(sys.argv[2])
+
+    spec = json.load(open(model_dir / "tokenizer.json"))
+    model = spec["model"]
+    if model["type"] != "BPE":
+        raise SystemExit(f"expected BPE, found {model['type']}")
+
+    vocab = model["vocab"]                      # token string -> id
+    merges = model["merges"]                    # [[left, right], ...] in rank order
+    added = spec.get("added_tokens", [])
+
+    # Tokens are addressed by id, so build the inverse mapping and check it is dense.
+    by_id = {i: s for s, i in vocab.items()}
+    for entry in added:
+        by_id[entry["id"]] = entry["content"]
+    size = max(by_id) + 1
+    missing = [i for i in range(size) if i not in by_id]
+    if missing:
+        raise SystemExit(f"vocabulary has {len(missing)} holes, first at {missing[0]}")
+
+    blob = bytearray()
+    offsets = [0]
+    for index in range(size):
+        blob.extend(by_id[index].encode("utf-8"))
+        offsets.append(len(blob))
+
+    # Merges as id pairs. A merge naming a token that is not in the vocabulary would
+    # silently never fire, so it is an error rather than a skip.
+    pairs = []
+    for rank, (left, right) in enumerate(merges):
+        if left not in vocab or right not in vocab:
+            raise SystemExit(f"merge {rank} references an unknown token: {left!r} {right!r}")
+        pairs.append((vocab[left], vocab[right]))
+
+    special_ids = [entry["id"] for entry in added]
+    special_blob = bytearray()
+    special_offsets = [0]
+    for entry in added:
+        special_blob.extend(entry["content"].encode("utf-8"))
+        special_offsets.append(len(special_blob))
+
+    regex = spec["pre_tokenizer"]["pretokenizers"][0]["pattern"]["Regex"]
+
+    header = json.dumps({
+        "format": "edge0-tokenizer-v1",
+        "vocab_size": size,
+        "merge_count": len(pairs),
+        "special_count": len(added),
+        "regex": regex,
+        "normalizer": (spec.get("normalizer") or {}).get("type"),
+        "byte_level": True,
+        "sections": ["vocab_offsets", "vocab_blob", "merge_pairs",
+                     "special_ids", "special_offsets", "special_blob"],
+    }, separators=(",", ":")).encode()
+
+    with open(out_path, "wb") as out:
+        out.write(struct.pack("<Q", len(header)))
+        out.write(header)
+        out.write(struct.pack(f"<{len(offsets)}I", *offsets))
+        out.write(bytes(blob))
+        flat = [value for pair in pairs for value in pair]
+        out.write(struct.pack(f"<{len(flat)}I", *flat))
+        out.write(struct.pack(f"<{len(special_ids)}I", *special_ids))
+        out.write(struct.pack(f"<{len(special_offsets)}I", *special_offsets))
+        out.write(bytes(special_blob))
+
+    total = out_path.stat().st_size
+    print(f"  vocab {size:,}   merges {len(pairs):,}   specials {len(added)}")
+    print(f"  ✓ {total / 2**20:.1f} MiB → {out_path.name}  "
+          f"(from {(model_dir / 'tokenizer.json').stat().st_size / 2**20:.1f} MiB of JSON)")
+    print(f"  regex: {regex[:70]}…")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/ios/tools/repack_experts.py b/ios/tools/repack_experts.py
new file mode 100755
index 0000000..6ffed45
--- /dev/null
+++ b/ios/tools/repack_experts.py
@@ -0,0 +1,354 @@
+#!/usr/bin/env python3
+"""Repack routed-expert weights so one expert is one contiguous read.
+
+The shipped safetensors store experts as nine separate stacked tensors per layer —
+`gate_proj`/`up_proj`/`down_proj` × `weight`/`scales`/`biases` — each indexed by expert
+along the leading axis. Reading one expert therefore means nine reads scattered across
+a 19 GB file, and reading the K experts a layer routes to means eighteen. Measured on
+an iPhone 16 Pro that costs 5.73 ms per layer at roughly 430 MB/s, against a device
+ceiling near 2 GB/s.
+
+This rewrites the same bytes grouped by (layer, expert):
+
+    block(layer, expert) = gate.weight ‖ gate.scales ‖ gate.biases
+                         ‖ up.weight   ‖ up.scales   ‖ up.biases
+                         ‖ down.weight ‖ down.scales ‖ down.biases
+
+Every block is 1,769,472 bytes — exactly 108 pages of 16 KB — so with the data starting
+at offset 0 of its own file, **every block is page-aligned**. One expert becomes one
+sequential 1.69 MB read.
+
+Output is **one file per layer**, not one 17 GiB file. Per-layer files make transfers
+retryable one layer at a time while keeping each expert's data contiguous.
+
+Only the experts are rewritten. The resident weights (embeddings, attention, routers,
+shared experts, norms — 2.13 GB) stay in the original shards, which this tool never
+modifies.
+
+The `resident` command extracts the other 2.13 GB — embeddings, attention, routers,
+shared experts, norms — into a single standard safetensors file with its own index, so
+that the original 17.7 GiB of shards need not be kept on the device at all. Nearly all
+of what those shards hold is expert weights that the repacked files already carry.
+
+Usage:
+    repack_experts.py pack     <model-dir> <out-dir>
+    repack_experts.py verify   <model-dir> <out-dir> [samples]
+    repack_experts.py resident <model-dir> <out-dir>
+"""
+
+import json
+import mmap
+import os
+import struct
+import sys
+import time
+
+FORMAT = "edge0-expert-repack-v1"
+PAGE = 16384
+
+# Order within a block. Grouped by projection so that a consumer reading only, say,
+# gate_proj still reads one contiguous run rather than three strided ones.
+PARTS = [
+    ("gate_proj", "weight"), ("gate_proj", "scales"), ("gate_proj", "biases"),
+    ("up_proj", "weight"), ("up_proj", "scales"), ("up_proj", "biases"),
+    ("down_proj", "weight"), ("down_proj", "scales"), ("down_proj", "biases"),
+]
+
+
+def tensor_name(layer, projection, part):
+    return f"language_model.model.layers.{layer}.mlp.switch_mlp.{projection}.{part}"
+
+
+def layer_file(layer):
+    return f"experts-L{layer:02d}.bin"
+
+
+class Shards:
+    """The original shards, mapped read-only, with their headers parsed."""
+
+    def __init__(self, model_dir):
+        self.dir = model_dir
+        index = json.load(open(os.path.join(model_dir, "model.safetensors.index.json")))
+        self.weight_map = index["weight_map"]
+
+        self.maps = {}
+        self.headers = {}
+        self._files = []
+        for filename in sorted(set(self.weight_map.values())):
+            handle = open(os.path.join(model_dir, filename), "rb")
+            self._files.append(handle)
+            header_length = struct.unpack("<Q", handle.read(8))[0]
+            header = json.loads(handle.read(header_length))
+            self.headers[filename] = (header, 8 + header_length)
+            self.maps[filename] = mmap.mmap(
+                handle.fileno(), 0, prot=mmap.PROT_READ)
+
+    def close(self):
+        for mapping in self.maps.values():
+            mapping.close()
+        for handle in self._files:
+            handle.close()
+
+    def entry(self, name):
+        filename = self.weight_map[name]
+        header, base = self.headers[filename]
+        field = header[name]
+        start, end = field["data_offsets"]
+        return filename, base + start, end - start, field["dtype"], field["shape"]
+
+    def slice_of(self, name, expert, experts):
+        """One expert's bytes out of a stacked tensor."""
+        filename, offset, total, _, _ = self.entry(name)
+        per = total // experts
+        start = offset + expert * per
+        return self.maps[filename][start:start + per]
+
+
+def describe(shards, layers, experts):
+    """Block layout, derived from the real headers rather than assumed."""
+    parts = []
+    cursor = 0
+    for projection, part in PARTS:
+        _, _, total, dtype, shape = shards.entry(tensor_name(0, projection, part))
+        per = total // experts
+        parts.append({
+            "name": f"{projection}.{part}",
+            "projection": projection,
+            "part": part,
+            "dtype": dtype,
+            "shape": shape[1:],          # per-expert shape, leading axis dropped
+            "offset": cursor,
+            "bytes": per,
+        })
+        cursor += per
+    return parts, cursor
+
+
+def pack(model_dir, out_dir):
+    shards = Shards(model_dir)
+    try:
+        layers = 1 + max(
+            int(name.split(".layers.")[1].split(".")[0])
+            for name in shards.weight_map
+            if ".layers." in name and ".switch_mlp." in name)
+        _, _, total, _, shape = shards.entry(tensor_name(0, "gate_proj", "weight"))
+        experts = shape[0]
+
+        parts, block_bytes = describe(shards, layers, experts)
+        if block_bytes % PAGE:
+            print(f"  ! block is {block_bytes} B, not a multiple of {PAGE} — "
+                  f"blocks will not be page-aligned")
+
+        os.makedirs(out_dir, exist_ok=True)
+        expected = layers * experts * block_bytes
+        per_layer = experts * block_bytes
+
+        free = os.statvfs(out_dir)
+        if free.f_bavail * free.f_frsize < expected:
+            raise SystemExit(
+                f"need {expected / 2**30:.2f} GiB, "
+                f"{free.f_bavail * free.f_frsize / 2**30:.2f} GiB free — refusing to "
+                f"start rather than fail half-written")
+
+        print(f"  {layers} layers × {experts} experts × {block_bytes:,} B "
+              f"= {expected / 2**30:.2f} GiB")
+        print(f"  block = {block_bytes / PAGE:.2f} pages of {PAGE} B")
+        print(f"  one file per layer, {per_layer / 2**20:.0f} MiB each")
+
+        started = time.time()
+        written = 0
+        for layer in range(layers):
+            sources = {
+                (projection, part): shards.entry(tensor_name(layer, projection, part))
+                for projection, part in PARTS
+            }
+            path = os.path.join(out_dir, layer_file(layer))
+            # Written to a temporary name and renamed only on success, so an interrupted
+            # run cannot leave behind a file that looks complete.
+            temporary = path + ".partial"
+            layer_written = 0
+            with open(temporary, "wb", buffering=1 << 22) as out:
+                for expert in range(experts):
+                    for projection, part in PARTS:
+                        filename, offset, total, _, _ = sources[(projection, part)]
+                        per = total // experts
+                        start = offset + expert * per
+                        out.write(shards.maps[filename][start:start + per])
+                        layer_written += per
+            if layer_written != per_layer:
+                os.unlink(temporary)
+                raise SystemExit(
+                    f"layer {layer}: wrote {layer_written} B, expected {per_layer}")
+            os.rename(temporary, path)
+
+            written += layer_written
+            elapsed = time.time() - started
+            rate = written / elapsed / 2**20 if elapsed else 0
+            print(f"\r  layer {layer + 1}/{layers}  {written / 2**30:.2f} GiB  "
+                  f"{rate:.0f} MB/s", end="", flush=True)
+        print()
+
+        json.dump({
+            "format": FORMAT,
+            "layers": layers,
+            "experts_per_layer": experts,
+            "block_bytes": block_bytes,
+            "layer_bytes": per_layer,
+            "page_size": PAGE,
+            "data_file_pattern": "experts-L%02d.bin",
+            "parts": parts,
+            "note": "expert e of layer L is at offset e * block_bytes in that layer's "
+                    "file; resident weights remain in the original shards",
+        }, open(os.path.join(out_dir, "experts.json"), "w"), indent=2)
+
+        print(f"  ✓ {written / 2**30:.2f} GiB in {time.time() - started:.0f}s")
+    finally:
+        shards.close()
+
+
+def verify(model_dir, out_dir, samples=64):
+    """Compare sampled blocks against the originals, byte for byte."""
+    shards = Shards(model_dir)
+    try:
+        meta = json.load(open(os.path.join(out_dir, "experts.json")))
+        if meta["format"] != FORMAT:
+            raise SystemExit(f"unexpected format {meta['format']}")
+
+        layers = meta["layers"]
+        experts = meta["experts_per_layer"]
+        block_bytes = meta["block_bytes"]
+        per_layer = meta["layer_bytes"]
+
+        # Every layer's file must be present and exactly the right size before any
+        # block is compared. A missing or truncated layer is the failure mode the
+        # per-layer split exists to make recoverable, so it is checked first and
+        # reported in full rather than one file at a time.
+        broken = []
+        for layer in range(layers):
+            path = os.path.join(out_dir, layer_file(layer))
+            if not os.path.exists(path):
+                broken.append(f"layer {layer}: missing")
+            elif os.path.getsize(path) != per_layer:
+                broken.append(
+                    f"layer {layer}: {os.path.getsize(path)} B, expected {per_layer}")
+        if broken:
+            raise SystemExit("✗ " + "\n  ✗ ".join(broken))
+
+        # Corners first, then a spread. A bug in the offset arithmetic shows up at
+        # the last layer's last expert long before it shows up in the middle.
+        picks = [(0, 0), (0, experts - 1), (layers - 1, 0), (layers - 1, experts - 1)]
+        step = max(1, (layers * experts) // max(1, samples - len(picks)))
+        for flat in range(0, layers * experts, step):
+            picks.append((flat // experts, flat % experts))
+
+        checked = 0
+        for layer, expert in sorted(picks):
+            with open(os.path.join(out_dir, layer_file(layer)), "rb") as handle:
+                packed = mmap.mmap(handle.fileno(), 0, prot=mmap.PROT_READ)
+                base = expert * block_bytes
+                for entry in meta["parts"]:
+                    projection, part = entry["projection"], entry["part"]
+                    original = shards.slice_of(
+                        tensor_name(layer, projection, part), expert, experts)
+                    start = base + entry["offset"]
+                    if packed[start:start + entry["bytes"]] != original:
+                        packed.close()
+                        raise SystemExit(
+                            f"✗ mismatch at layer {layer} expert {expert} "
+                            f"{projection}.{part}")
+                packed.close()
+            checked += 1
+            print(f"\r  verified {checked}/{len(picks)} blocks", end="", flush=True)
+        print(f"\n  ✓ {checked} blocks across {layers} layer files match the originals "
+              f"byte for byte")
+    finally:
+        shards.close()
+
+
+def resident(model_dir, out_dir):
+    """Everything that is not a routed expert, in one safetensors file.
+
+    Written as a standard safetensors file plus an index naming it, so the device-side
+    reader opens it exactly as it opens the originals — no format to teach it.
+
+    Verification happens in the same pass rather than as a separate command: every
+    tensor is read back from the finished file and compared against its source, because
+    this file replaces the shards entirely and a silent mismatch here would not surface
+    until inference produced quiet nonsense.
+    """
+    shards = Shards(model_dir)
+    try:
+        names = sorted(n for n in shards.weight_map if ".switch_mlp." not in n)
+        header = {}
+        cursor = 0
+        for name in names:
+            _, _, byte_count, dtype, shape = shards.entry(name)
+            header[name] = {
+                "dtype": dtype,
+                "shape": shape,
+                "data_offsets": [cursor, cursor + byte_count],
+            }
+            cursor += byte_count
+
+        os.makedirs(out_dir, exist_ok=True)
+        blob = json.dumps(header, separators=(",", ":")).encode()
+        path = os.path.join(out_dir, "resident.safetensors")
+        temporary = path + ".partial"
+
+        print(f"  {len(names)} tensors, {cursor / 2**30:.2f} GiB")
+        started = time.time()
+        written = 0
+        with open(temporary, "wb", buffering=1 << 22) as out:
+            out.write(struct.pack("<Q", len(blob)))
+            out.write(blob)
+            for index, name in enumerate(names):
+                filename, offset, byte_count, _, _ = shards.entry(name)
+                out.write(shards.maps[filename][offset:offset + byte_count])
+                written += byte_count
+                if index % 100 == 0:
+                    print(f"\r  {index}/{len(names)}  {written / 2**30:.2f} GiB",
+                          end="", flush=True)
+        print(f"\r  {len(names)}/{len(names)}  {written / 2**30:.2f} GiB")
+
+        if written != cursor:
+            os.unlink(temporary)
+            raise SystemExit(f"wrote {written} B, expected {cursor} — output removed")
+        os.rename(temporary, path)
+
+        json.dump({
+            "metadata": {"note": "resident weights only; experts live in experts-L*.bin"},
+            "weight_map": {name: "resident.safetensors" for name in names},
+        }, open(os.path.join(out_dir, "model.safetensors.index.json"), "w"), indent=2)
+
+        # Read back and compare, all of it.
+        base = 8 + len(blob)
+        with open(path, "rb") as handle:
+            packed = mmap.mmap(handle.fileno(), 0, prot=mmap.PROT_READ)
+            for index, name in enumerate(names):
+                filename, offset, byte_count, _, _ = shards.entry(name)
+                start = base + header[name]["data_offsets"][0]
+                if packed[start:start + byte_count] != \
+                        shards.maps[filename][offset:offset + byte_count]:
+                    packed.close()
+                    raise SystemExit(f"✗ mismatch in {name}")
+                if index % 100 == 0:
+                    print(f"\r  verifying {index}/{len(names)}", end="", flush=True)
+            packed.close()
+        print(f"\r  ✓ all {len(names)} tensors match, "
+              f"{cursor / 2**30:.2f} GiB in {time.time() - started:.0f}s")
+    finally:
+        shards.close()
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 4:
+        raise SystemExit(__doc__)
+    command, model_dir, out_dir = sys.argv[1], sys.argv[2], sys.argv[3]
+    if command == "pack":
+        pack(model_dir, out_dir)
+    elif command == "verify":
+        verify(model_dir, out_dir, int(sys.argv[4]) if len(sys.argv) > 4 else 64)
+    elif command == "resident":
+        resident(model_dir, out_dir)
+    else:
+        raise SystemExit(__doc__)
diff --git a/macos/Cargo.lock b/macos/Cargo.lock
new file mode 100644
index 0000000..97c2db5
--- /dev/null
+++ b/macos/Cargo.lock
@@ -0,0 +1,6190 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "adler2"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
+
+[[package]]
+name = "ahash"
+version = "0.8.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
+dependencies = [
+ "cfg-if",
+ "getrandom 0.3.4",
+ "once_cell",
+ "serde",
+ "version_check",
+ "zerocopy",
+]
+
+[[package]]
+name = "aho-corasick"
+version = "1.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "alloc-no-stdlib"
+version = "2.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3"
+
+[[package]]
+name = "alloc-stdlib"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195"
+dependencies = [
+ "alloc-no-stdlib",
+]
+
+[[package]]
+name = "android_system_properties"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "anstream"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d"
+dependencies = [
+ "anstyle",
+ "anstyle-parse",
+ "anstyle-query",
+ "anstyle-wincon",
+ "colorchoice",
+ "is_terminal_polyfill",
+ "utf8parse",
+]
+
+[[package]]
+name = "anstyle"
+version = "1.0.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
+
+[[package]]
+name = "anstyle-parse"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e"
+dependencies = [
+ "utf8parse",
+]
+
+[[package]]
+name = "anstyle-query"
+version = "1.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc"
+dependencies = [
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "anstyle-wincon"
+version = "3.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d"
+dependencies = [
+ "anstyle",
+ "once_cell_polyfill",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "anyhow"
+version = "1.0.104"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
+
+[[package]]
+name = "arboard"
+version = "3.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0348a1c054491f4bfe6ab86a7b6ab1e44e45d899005de92f58b3df180b36ddaf"
+dependencies = [
+ "clipboard-win",
+ "log",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-foundation",
+ "parking_lot",
+ "percent-encoding",
+ "windows-sys 0.59.0",
+ "x11rb",
+]
+
+[[package]]
+name = "atk"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "241b621213072e993be4f6f3a9e4b45f65b7e6faad43001be957184b7bb1824b"
+dependencies = [
+ "atk-sys",
+ "glib",
+ "libc",
+]
+
+[[package]]
+name = "atk-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c5e48b684b0ca77d2bbadeef17424c2ea3c897d44d566a1617e7e8f30614d086"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "atomic-waker"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
+
+[[package]]
+name = "autocfg"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
+
+[[package]]
+name = "axum"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
+dependencies = [
+ "axum-core",
+ "bytes",
+ "form_urlencoded",
+ "futures-util",
+ "http",
+ "http-body",
+ "http-body-util",
+ "hyper",
+ "hyper-util",
+ "itoa",
+ "matchit",
+ "memchr",
+ "mime",
+ "percent-encoding",
+ "pin-project-lite",
+ "serde_core",
+ "serde_json",
+ "serde_path_to_error",
+ "serde_urlencoded",
+ "sync_wrapper",
+ "tokio",
+ "tower",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
+[[package]]
+name = "axum-core"
+version = "0.5.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
+dependencies = [
+ "bytes",
+ "futures-core",
+ "http",
+ "http-body",
+ "http-body-util",
+ "mime",
+ "pin-project-lite",
+ "sync_wrapper",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
+[[package]]
+name = "base64"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
+
+[[package]]
+name = "base64"
+version = "0.21.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567"
+
+[[package]]
+name = "base64"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
+
+[[package]]
+name = "base64"
+version = "0.23.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5"
+
+[[package]]
+name = "base64ct"
+version = "1.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
+
+[[package]]
+name = "bit-set"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
+dependencies = [
+ "bit-vec",
+]
+
+[[package]]
+name = "bit-vec"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
+
+[[package]]
+name = "bitflags"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
+
+[[package]]
+name = "bitflags"
+version = "2.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "block-buffer"
+version = "0.10.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
+dependencies = [
+ "generic-array",
+]
+
+[[package]]
+name = "block2"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5"
+dependencies = [
+ "objc2",
+]
+
+[[package]]
+name = "brotli"
+version = "8.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3"
+dependencies = [
+ "alloc-no-stdlib",
+ "alloc-stdlib",
+ "brotli-decompressor",
+]
+
+[[package]]
+name = "brotli-decompressor"
+version = "5.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583"
+dependencies = [
+ "alloc-no-stdlib",
+ "alloc-stdlib",
+]
+
+[[package]]
+name = "bs58"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4"
+dependencies = [
+ "tinyvec",
+]
+
+[[package]]
+name = "bumpalo"
+version = "3.20.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649"
+
+[[package]]
+name = "bytemuck"
+version = "1.25.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797"
+
+[[package]]
+name = "byteorder"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
+
+[[package]]
+name = "byteorder-lite"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495"
+
+[[package]]
+name = "bytes"
+version = "1.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "cairo-rs"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ca26ef0159422fb77631dc9d17b102f253b876fe1586b03b803e63a309b4ee2"
+dependencies = [
+ "bitflags 2.13.2",
+ "cairo-sys-rs",
+ "glib",
+ "libc",
+ "once_cell",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "cairo-sys-rs"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "685c9fa8e590b8b3d678873528d83411db17242a73fccaed827770ea0fedda51"
+dependencies = [
+ "glib-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "camino"
+version = "1.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbbad30e4b4c14a39e3cc8aed085a12a327257c316619c93581e017bc52be591"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "cargo-platform"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e35af189006b9c0f00a064685c727031e3ed2d8020f7ba284d78cc2671bd36ea"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "cargo_metadata"
+version = "0.19.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dd5eb614ed4c27c5d706420e4320fbe3216ab31fa1c33cd8246ac36dae4479ba"
+dependencies = [
+ "camino",
+ "cargo-platform",
+ "semver",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.20",
+]
+
+[[package]]
+name = "cargo_toml"
+version = "0.22.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "374b7c592d9c00c1f4972ea58390ac6b18cbb6ab79011f3bdc90a0b82ca06b77"
+dependencies = [
+ "serde",
+ "toml 0.9.12+spec-1.1.0",
+]
+
+[[package]]
+name = "castaway"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a"
+dependencies = [
+ "rustversion",
+]
+
+[[package]]
+name = "cc"
+version = "1.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a3eb0f42d6c360dc3f8a821f6bf2fdea7f72bfd36b3076eb0e6d1e9e0752fff4"
+dependencies = [
+ "find-msvc-tools",
+ "shlex",
+]
+
+[[package]]
+name = "cesu8"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c"
+
+[[package]]
+name = "cfb"
+version = "0.7.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d38f2da7a0a2c4ccf0065be06397cc26a81f4e528be095826eee9d4adbb8c60f"
+dependencies = [
+ "byteorder",
+ "fnv",
+ "uuid",
+]
+
+[[package]]
+name = "cfg-expr"
+version = "0.15.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02"
+dependencies = [
+ "smallvec",
+ "target-lexicon",
+]
+
+[[package]]
+name = "cfg-if"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
+
+[[package]]
+name = "cfg_aliases"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527"
+
+[[package]]
+name = "chacha20"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06"
+dependencies = [
+ "cfg-if",
+ "cpufeatures 0.3.1",
+ "rand_core 0.10.1",
+]
+
+[[package]]
+name = "chrono"
+version = "0.4.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327"
+dependencies = [
+ "iana-time-zone",
+ "js-sys",
+ "num-traits",
+ "serde",
+ "wasm-bindgen",
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "clap"
+version = "4.6.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca"
+dependencies = [
+ "clap_builder",
+ "clap_derive",
+]
+
+[[package]]
+name = "clap_builder"
+version = "4.6.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889"
+dependencies = [
+ "anstream",
+ "anstyle",
+ "clap_lex",
+ "strsim",
+]
+
+[[package]]
+name = "clap_derive"
+version = "4.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061"
+dependencies = [
+ "heck 0.5.0",
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "clap_lex"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
+
+[[package]]
+name = "clipboard-win"
+version = "5.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bde03770d3df201d4fb868f2c9c59e66a3e4e2bd06692a0fe701e7103c7e84d4"
+dependencies = [
+ "error-code",
+]
+
+[[package]]
+name = "colorchoice"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570"
+
+[[package]]
+name = "combine"
+version = "4.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e"
+dependencies = [
+ "bytes",
+ "memchr",
+]
+
+[[package]]
+name = "compact_str"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab"
+dependencies = [
+ "castaway",
+ "cfg-if",
+ "itoa",
+ "rustversion",
+ "ryu",
+ "serde",
+ "static_assertions",
+]
+
+[[package]]
+name = "console"
+version = "0.15.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8"
+dependencies = [
+ "encode_unicode",
+ "libc",
+ "once_cell",
+ "unicode-width",
+ "windows-sys 0.59.0",
+]
+
+[[package]]
+name = "const-oid"
+version = "0.9.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"
+
+[[package]]
+name = "cookie"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87"
+dependencies = [
+ "time",
+ "version_check",
+]
+
+[[package]]
+name = "core-foundation"
+version = "0.9.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f"
+dependencies = [
+ "core-foundation-sys",
+ "libc",
+]
+
+[[package]]
+name = "core-foundation"
+version = "0.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6"
+dependencies = [
+ "core-foundation-sys",
+ "libc",
+]
+
+[[package]]
+name = "core-foundation-sys"
+version = "0.8.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
+
+[[package]]
+name = "core-graphics"
+version = "0.25.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "064badf302c3194842cf2c5d61f56cc88e54a759313879cdf03abdd27d0c3b97"
+dependencies = [
+ "bitflags 2.13.2",
+ "core-foundation 0.10.1",
+ "core-graphics-types",
+ "foreign-types",
+ "libc",
+]
+
+[[package]]
+name = "core-graphics-types"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb"
+dependencies = [
+ "bitflags 2.13.2",
+ "core-foundation 0.10.1",
+ "libc",
+]
+
+[[package]]
+name = "cpufeatures"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "cpufeatures"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "crc32fast"
+version = "1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78"
+dependencies = [
+ "cfg-if",
+]
+
+[[package]]
+name = "crossbeam-channel"
+version = "0.5.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6"
+
+[[package]]
+name = "crypto-common"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a"
+dependencies = [
+ "generic-array",
+ "typenum",
+]
+
+[[package]]
+name = "cssparser"
+version = "0.36.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dae61cf9c0abb83bd659dab65b7e4e38d8236824c85f0f804f173567bda257d2"
+dependencies = [
+ "cssparser-macros",
+ "dtoa-short",
+ "itoa",
+ "phf",
+ "smallvec",
+]
+
+[[package]]
+name = "cssparser-macros"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13b588ba4ac1a99f7f2964d24b3d896ddc6bf847ee3855dbd4366f058cfcd331"
+dependencies = [
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "ctor"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "352d39c2f7bef1d6ad73db6f5160efcaed66d94ef8c6c573a8410c00bf909a98"
+dependencies = [
+ "ctor-proc-macro",
+ "dtor",
+]
+
+[[package]]
+name = "ctor-proc-macro"
+version = "0.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52560adf09603e58c9a7ee1fe1dcb95a16927b17c127f0ac02d6e768a0e25bc1"
+
+[[package]]
+name = "curve25519-dalek"
+version = "4.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be"
+dependencies = [
+ "cfg-if",
+ "cpufeatures 0.2.17",
+ "curve25519-dalek-derive",
+ "digest",
+ "fiat-crypto",
+ "rustc_version",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
+name = "curve25519-dalek-derive"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "darling"
+version = "0.20.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
+dependencies = [
+ "darling_core 0.20.11",
+ "darling_macro 0.20.11",
+]
+
+[[package]]
+name = "darling"
+version = "0.24.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed17f5901b6630b993ca003def43f2f8ef4014fc13b047b57aad617ff32bc2ec"
+dependencies = [
+ "darling_core 0.24.1",
+ "darling_macro 0.24.1",
+]
+
+[[package]]
+name = "darling_core"
+version = "0.20.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e"
+dependencies = [
+ "fnv",
+ "ident_case",
+ "proc-macro2",
+ "quote",
+ "strsim",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "darling_core"
+version = "0.24.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6837e2cf7485aaae18f86181d2f0e9a7ed297a025e220aeabf63fdebd3a2ddff"
+dependencies = [
+ "ident_case",
+ "proc-macro2",
+ "quote",
+ "strsim",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "darling_macro"
+version = "0.20.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
+dependencies = [
+ "darling_core 0.20.11",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "darling_macro"
+version = "0.24.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2ac7135c3ef02b2f7833bbeb1be5ba7f966dcde8a87c6b87f65a778d71a02785"
+dependencies = [
+ "darling_core 0.24.1",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "dary_heap"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "dbus"
+version = "0.9.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3ab69f03cc8c4340c9c8e315114e1658e6775a9b16a04357973aa21cec22b32e"
+dependencies = [
+ "libc",
+ "libdbus-sys",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "defmt"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1"
+dependencies = [
+ "bitflags 1.3.2",
+ "defmt-macros",
+]
+
+[[package]]
+name = "defmt-macros"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8"
+dependencies = [
+ "defmt-parser",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "defmt-parser"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e"
+dependencies = [
+ "thiserror 2.0.20",
+]
+
+[[package]]
+name = "der"
+version = "0.7.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
+dependencies = [
+ "const-oid",
+ "zeroize",
+]
+
+[[package]]
+name = "deranged"
+version = "0.5.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "derive_builder"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947"
+dependencies = [
+ "derive_builder_macro",
+]
+
+[[package]]
+name = "derive_builder_core"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
+dependencies = [
+ "darling 0.20.11",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "derive_builder_macro"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
+dependencies = [
+ "derive_builder_core",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "derive_more"
+version = "2.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134"
+dependencies = [
+ "derive_more-impl",
+]
+
+[[package]]
+name = "derive_more-impl"
+version = "2.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "rustc_version",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "digest"
+version = "0.10.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
+dependencies = [
+ "block-buffer",
+ "crypto-common",
+]
+
+[[package]]
+name = "dirs"
+version = "6.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e"
+dependencies = [
+ "dirs-sys",
+]
+
+[[package]]
+name = "dirs-sys"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab"
+dependencies = [
+ "libc",
+ "option-ext",
+ "redox_users",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "dispatch2"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "libc",
+ "objc2",
+]
+
+[[package]]
+name = "displaydoc"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "dlopen2"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e2c5bd4158e66d1e215c49b837e11d62f3267b30c92f1d171c4d3105e3dc4d4"
+dependencies = [
+ "dlopen2_derive",
+ "libc",
+ "once_cell",
+ "winapi",
+]
+
+[[package]]
+name = "dlopen2_derive"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fbbb781877580993a8707ec48672673ec7b81eeba04cfd2310bd28c08e47c8f"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "dom_query"
+version = "0.27.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "521e380c0c8afb8d9a1e83a1822ee03556fc3e3e7dbc1fd30be14e37f9cb3f89"
+dependencies = [
+ "bit-set",
+ "cssparser",
+ "foldhash",
+ "html5ever",
+ "precomputed-hash",
+ "selectors",
+ "tendril",
+]
+
+[[package]]
+name = "dpi"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8b14ccef22fc6f5a8f4d7d768562a182c04ce9a3b3157b91390b52ddfdf1a76"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "dtoa"
+version = "1.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4c3cf4824e2d5f025c7b531afcb2325364084a16806f6d47fbc1f5fbd9960590"
+
+[[package]]
+name = "dtoa-short"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd1511a7b6a56299bd043a9c167a6d2bfb37bf84a6dfceaba651168adfb43c87"
+dependencies = [
+ "dtoa",
+]
+
+[[package]]
+name = "dtor"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1057d6c64987086ff8ed0fd3fbf377a6b7d205cc7715868cd401705f715cbe4"
+dependencies = [
+ "dtor-proc-macro",
+]
+
+[[package]]
+name = "dtor-proc-macro"
+version = "0.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f678cf4a922c215c63e0de95eb1ff08a958a81d47e485cf9da1e27bf6305cfa5"
+
+[[package]]
+name = "dunce"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
+
+[[package]]
+name = "dyn-clone"
+version = "1.0.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
+
+[[package]]
+name = "ed25519"
+version = "2.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53"
+dependencies = [
+ "pkcs8",
+ "signature",
+]
+
+[[package]]
+name = "ed25519-dalek"
+version = "2.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9"
+dependencies = [
+ "curve25519-dalek",
+ "ed25519",
+ "rand_core 0.6.4",
+ "serde",
+ "sha2",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
+name = "edge0-app"
+version = "0.1.0"
+dependencies = [
+ "arboard",
+ "chrono",
+ "edge0-core",
+ "futures",
+ "libc",
+ "reqwest 0.12.28",
+ "rusqlite",
+ "serde",
+ "serde_json",
+ "tauri",
+ "tauri-build",
+ "thiserror 2.0.20",
+ "tokio",
+ "tracing",
+ "uuid",
+]
+
+[[package]]
+name = "edge0-cli"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "chrono",
+ "clap",
+ "edge0-core",
+ "edge0-pull",
+ "reqwest 0.12.28",
+ "serde",
+ "serde_json",
+ "tokio",
+ "tracing",
+]
+
+[[package]]
+name = "edge0-core"
+version = "0.1.0"
+dependencies = [
+ "chrono",
+ "libc",
+ "libloading 0.8.9",
+ "rand 0.8.8",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.20",
+ "ts-rs",
+ "uuid",
+]
+
+[[package]]
+name = "edge0-engine"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "edge0-core",
+ "libc",
+ "libloading 0.8.9",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.20",
+ "tokio",
+ "tracing",
+ "tracing-subscriber",
+]
+
+[[package]]
+name = "edge0-pull"
+version = "0.1.0"
+dependencies = [
+ "base64 0.22.1",
+ "chrono",
+ "clap",
+ "ed25519-dalek",
+ "edge0-core",
+ "futures",
+ "libc",
+ "rand 0.8.8",
+ "reqwest 0.12.28",
+ "serde",
+ "serde_json",
+ "sha2",
+ "tokio",
+]
+
+[[package]]
+name = "edge0d"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "axum",
+ "chrono",
+ "clap",
+ "edge0-core",
+ "edge0-pull",
+ "futures",
+ "http",
+ "libc",
+ "libloading 0.8.9",
+ "rand 0.8.8",
+ "reqwest 0.12.28",
+ "serde",
+ "serde_json",
+ "sha2",
+ "thiserror 2.0.20",
+ "tokio",
+ "tokio-stream",
+ "tokio-util",
+ "tower",
+ "tracing",
+ "tracing-subscriber",
+ "uuid",
+]
+
+[[package]]
+name = "either"
+version = "1.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34"
+
+[[package]]
+name = "embed-resource"
+version = "3.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbfdaacccebec3b28e4866b8973543c7647797db5ada1bdab552e48fe665fbbd"
+dependencies = [
+ "cc",
+ "memchr",
+ "rustc_version",
+ "toml 1.1.6+spec-1.1.0",
+ "vswhom",
+ "winreg",
+]
+
+[[package]]
+name = "embed_plist"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ef6b89e5b37196644d8796de5268852ff179b44e96276cf4290264843743bb7"
+
+[[package]]
+name = "encode_unicode"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0"
+
+[[package]]
+name = "engine-tok"
+version = "0.1.0"
+dependencies = [
+ "minijinja",
+ "minijinja-contrib",
+ "serde_json",
+ "tokenizers",
+]
+
+[[package]]
+name = "equivalent"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
+
+[[package]]
+name = "erased-serde"
+version = "0.4.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2add8a07dd6a8d93ff627029c51de145e12686fbc36ecb298ac22e74cf02dec"
+dependencies = [
+ "serde",
+ "serde_core",
+ "typeid",
+]
+
+[[package]]
+name = "errno"
+version = "0.3.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
+dependencies = [
+ "libc",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "error-code"
+version = "3.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b5343afd4a8365a643ac588dab4cf234a190c7f6c88c9f6dd6ffe00837661b7"
+
+[[package]]
+name = "esaxx-rs"
+version = "0.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
+dependencies = [
+ "cc",
+]
+
+[[package]]
+name = "fallible-iterator"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649"
+
+[[package]]
+name = "fallible-streaming-iterator"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
+
+[[package]]
+name = "fastrand"
+version = "2.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223"
+
+[[package]]
+name = "fdeflate"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c"
+dependencies = [
+ "simd-adler32",
+]
+
+[[package]]
+name = "fiat-crypto"
+version = "0.2.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d"
+
+[[package]]
+name = "field-offset"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38e2275cc4e4fc009b0669731a1e5ab7ebf11f469eaede2bab9309a5b4d6057f"
+dependencies = [
+ "memoffset",
+ "rustc_version",
+]
+
+[[package]]
+name = "find-msvc-tools"
+version = "0.1.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3e0f1c7c3a72c66fd80abe965175f7523475c0489a87d3ff9d6e8c87d87a9d2d"
+
+[[package]]
+name = "flate2"
+version = "1.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide 0.9.1",
+ "zlib-rs",
+]
+
+[[package]]
+name = "fnv"
+version = "1.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
+
+[[package]]
+name = "foldhash"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
+
+[[package]]
+name = "foreign-types"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965"
+dependencies = [
+ "foreign-types-macros",
+ "foreign-types-shared",
+]
+
+[[package]]
+name = "foreign-types-macros"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ea5190182e6915eb873ddbc16e23b711b6eb1f9c00a0d0a3a91b5f6228475225"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "foreign-types-shared"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b"
+
+[[package]]
+name = "form_urlencoded"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf"
+dependencies = [
+ "percent-encoding",
+]
+
+[[package]]
+name = "futures"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-executor",
+ "futures-io",
+ "futures-sink",
+ "futures-task",
+ "futures-util",
+]
+
+[[package]]
+name = "futures-channel"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4"
+dependencies = [
+ "futures-core",
+ "futures-sink",
+]
+
+[[package]]
+name = "futures-core"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e"
+
+[[package]]
+name = "futures-executor"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432"
+dependencies = [
+ "futures-core",
+ "futures-task",
+ "futures-util",
+]
+
+[[package]]
+name = "futures-io"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed"
+
+[[package]]
+name = "futures-macro"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "futures-sink"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d"
+
+[[package]]
+name = "futures-task"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd"
+
+[[package]]
+name = "futures-util"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-io",
+ "futures-macro",
+ "futures-sink",
+ "futures-task",
+ "memchr",
+ "pin-project-lite",
+ "slab",
+]
+
+[[package]]
+name = "gdk"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d9f245958c627ac99d8e529166f9823fb3b838d1d41fd2b297af3075093c2691"
+dependencies = [
+ "cairo-rs",
+ "gdk-pixbuf",
+ "gdk-sys",
+ "gio",
+ "glib",
+ "libc",
+ "pango",
+]
+
+[[package]]
+name = "gdk-pixbuf"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "50e1f5f1b0bfb830d6ccc8066d18db35c487b1b2b1e8589b5dfe9f07e8defaec"
+dependencies = [
+ "gdk-pixbuf-sys",
+ "gio",
+ "glib",
+ "libc",
+ "once_cell",
+]
+
+[[package]]
+name = "gdk-pixbuf-sys"
+version = "0.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f9839ea644ed9c97a34d129ad56d38a25e6756f99f3a88e15cd39c20629caf7"
+dependencies = [
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "gdk-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c2d13f38594ac1e66619e188c6d5a1adb98d11b2fcf7894fc416ad76aa2f3f7"
+dependencies = [
+ "cairo-sys-rs",
+ "gdk-pixbuf-sys",
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "pango-sys",
+ "pkg-config",
+ "system-deps",
+]
+
+[[package]]
+name = "gdkwayland-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "140071d506d223f7572b9f09b5e155afbd77428cd5cc7af8f2694c41d98dfe69"
+dependencies = [
+ "gdk-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "pkg-config",
+ "system-deps",
+]
+
+[[package]]
+name = "gdkx11"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3caa00e14351bebbc8183b3c36690327eb77c49abc2268dd4bd36b856db3fbfe"
+dependencies = [
+ "gdk",
+ "gdkx11-sys",
+ "gio",
+ "glib",
+ "libc",
+ "x11",
+]
+
+[[package]]
+name = "gdkx11-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e2e7445fe01ac26f11601db260dd8608fe172514eb63b3b5e261ea6b0f4428d"
+dependencies = [
+ "gdk-sys",
+ "glib-sys",
+ "libc",
+ "system-deps",
+ "x11",
+]
+
+[[package]]
+name = "generic-array"
+version = "0.14.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
+dependencies = [
+ "typenum",
+ "version_check",
+]
+
+[[package]]
+name = "gethostname"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1bd49230192a3797a9a4d6abe9b3eed6f7fa4c8a8a4947977c6f80025f92cbd8"
+dependencies = [
+ "rustix",
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
+dependencies = [
+ "cfg-if",
+ "js-sys",
+ "libc",
+ "wasi",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "r-efi 5.3.0",
+ "wasip2",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099"
+dependencies = [
+ "cfg-if",
+ "js-sys",
+ "libc",
+ "r-efi 6.0.0",
+ "rand_core 0.10.1",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "gio"
+version = "0.18.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d4fc8f532f87b79cbc51a79748f16a6828fb784be93145a322fa14d06d354c73"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-io",
+ "futures-util",
+ "gio-sys",
+ "glib",
+ "libc",
+ "once_cell",
+ "pin-project-lite",
+ "smallvec",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "gio-sys"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37566df850baf5e4cb0dfb78af2e4b9898d817ed9263d1090a2df958c64737d2"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+ "winapi",
+]
+
+[[package]]
+name = "glib"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "233daaf6e83ae6a12a52055f568f9d7cf4671dabb78ff9560ab6da230ce00ee5"
+dependencies = [
+ "bitflags 2.13.2",
+ "futures-channel",
+ "futures-core",
+ "futures-executor",
+ "futures-task",
+ "futures-util",
+ "gio-sys",
+ "glib-macros",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "memchr",
+ "once_cell",
+ "smallvec",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "glib-macros"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0bb0228f477c0900c880fd78c8759b95c7636dbd7842707f49e132378aa2acdc"
+dependencies = [
+ "heck 0.4.1",
+ "proc-macro-crate 2.0.2",
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "glib-sys"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "063ce2eb6a8d0ea93d2bf8ba1957e78dbab6be1c2220dd3daca57d5a9d869898"
+dependencies = [
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "glob"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b"
+
+[[package]]
+name = "gobject-sys"
+version = "0.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0850127b514d1c4a4654ead6dedadb18198999985908e6ffe4436f53c785ce44"
+dependencies = [
+ "glib-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "gtk"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd56fb197bfc42bd5d2751f4f017d44ff59fbb58140c6b49f9b3b2bdab08506a"
+dependencies = [
+ "atk",
+ "cairo-rs",
+ "field-offset",
+ "futures-channel",
+ "gdk",
+ "gdk-pixbuf",
+ "gio",
+ "glib",
+ "gtk-sys",
+ "gtk3-macros",
+ "libc",
+ "pango",
+ "pkg-config",
+]
+
+[[package]]
+name = "gtk-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f29a1c21c59553eb7dd40e918be54dccd60c52b049b75119d5d96ce6b624414"
+dependencies = [
+ "atk-sys",
+ "cairo-sys-rs",
+ "gdk-pixbuf-sys",
+ "gdk-sys",
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "pango-sys",
+ "system-deps",
+]
+
+[[package]]
+name = "gtk3-macros"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52ff3c5b21f14f0736fed6dcfc0bfb4225ebf5725f3c0209edeec181e4d73e9d"
+dependencies = [
+ "proc-macro-crate 1.3.1",
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "hashbrown"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+
+[[package]]
+name = "hashbrown"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
+dependencies = [
+ "foldhash",
+]
+
+[[package]]
+name = "hashbrown"
+version = "0.17.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
+dependencies = [
+ "foldhash",
+]
+
+[[package]]
+name = "hashlink"
+version = "0.12.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb"
+dependencies = [
+ "hashbrown 0.17.1",
+]
+
+[[package]]
+name = "heck"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "hex"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
+
+[[package]]
+name = "html5ever"
+version = "0.38.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1054432bae2f14e0061e33d23402fbaa67a921d319d56adc6bcf887ddad1cbc2"
+dependencies = [
+ "log",
+ "markup5ever",
+]
+
+[[package]]
+name = "http"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0"
+dependencies = [
+ "bytes",
+ "itoa",
+]
+
+[[package]]
+name = "http-body"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c"
+dependencies = [
+ "bytes",
+ "http",
+]
+
+[[package]]
+name = "http-body-util"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23169fe34a5fbcdd3f3862e78fb9b6fccd5f02a6dc6f732547005d45631ce71c"
+dependencies = [
+ "bytes",
+ "futures-core",
+ "http",
+ "http-body",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "httparse"
+version = "1.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
+
+[[package]]
+name = "httpdate"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
+
+[[package]]
+name = "hyper"
+version = "1.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "27b501faa50e7a26c3d3560ca625132f4078a17771f4810baf70475ae48cbe43"
+dependencies = [
+ "atomic-waker",
+ "bytes",
+ "futures-channel",
+ "futures-core",
+ "http",
+ "http-body",
+ "httparse",
+ "httpdate",
+ "itoa",
+ "pin-project-lite",
+ "smallvec",
+ "tokio",
+ "want",
+]
+
+[[package]]
+name = "hyper-rustls"
+version = "0.27.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f"
+dependencies = [
+ "http",
+ "hyper",
+ "hyper-util",
+ "rustls",
+ "tokio",
+ "tokio-rustls",
+ "tower-service",
+ "webpki-roots",
+]
+
+[[package]]
+name = "hyper-util"
+version = "0.1.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
+dependencies = [
+ "base64 0.22.1",
+ "bytes",
+ "futures-channel",
+ "futures-util",
+ "http",
+ "http-body",
+ "hyper",
+ "ipnet",
+ "libc",
+ "percent-encoding",
+ "pin-project-lite",
+ "socket2",
+ "system-configuration",
+ "tokio",
+ "tower-service",
+ "tracing",
+ "windows-registry",
+]
+
+[[package]]
+name = "iana-time-zone"
+version = "0.1.65"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
+dependencies = [
+ "android_system_properties",
+ "core-foundation-sys",
+ "iana-time-zone-haiku",
+ "js-sys",
+ "log",
+ "wasm-bindgen",
+ "windows-core 0.62.2",
+]
+
+[[package]]
+name = "iana-time-zone-haiku"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
+dependencies = [
+ "cc",
+]
+
+[[package]]
+name = "ico"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3e795dff5605e0f04bff85ca41b51a96b83e80b281e96231bcaaf1ac35103371"
+dependencies = [
+ "byteorder",
+ "png 0.17.16",
+]
+
+[[package]]
+name = "icu_collections"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa68d21081c4a05d5a901a1c62add574c77048b6a1c67be3b50ce0b60d4ca513"
+dependencies = [
+ "displaydoc",
+ "potential_utf",
+ "utf8_iter",
+ "yoke",
+ "zerofrom",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_locale_core"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb"
+dependencies = [
+ "displaydoc",
+ "litemap",
+ "tinystr",
+ "writeable",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_normalizer"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "12f9cf5f235641ed274641dd81c3f28d870e276763d0797aeeab72317b1c646f"
+dependencies = [
+ "icu_collections",
+ "icu_normalizer_data",
+ "icu_properties",
+ "icu_provider",
+ "smallvec",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_normalizer_data"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1563da1ed3e0b3bf3d74c9b85917ac9c56464d2f57242270c09c9e752f8021a0"
+
+[[package]]
+name = "icu_properties"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e7ca276ad3145661a65914e6daf131ca5120cd3dcee8f8f3214b8875184a148"
+dependencies = [
+ "displaydoc",
+ "icu_collections",
+ "icu_locale_core",
+ "icu_properties_data",
+ "icu_provider",
+ "zerotrie",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_properties_data"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa"
+
+[[package]]
+name = "icu_provider"
+version = "2.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d27bbb9d3abbefac45d55f647c9de1d44aafcd1186eb91879afef17c396c3e73"
+dependencies = [
+ "displaydoc",
+ "icu_locale_core",
+ "writeable",
+ "yoke",
+ "zerofrom",
+ "zerotrie",
+ "zerovec",
+]
+
+[[package]]
+name = "ident_case"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
+
+[[package]]
+name = "idna"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de"
+dependencies = [
+ "idna_adapter",
+ "smallvec",
+ "utf8_iter",
+]
+
+[[package]]
+name = "idna_adapter"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714"
+dependencies = [
+ "icu_normalizer",
+ "icu_properties",
+]
+
+[[package]]
+name = "image"
+version = "0.25.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85ab80394333c02fe689eaf900ab500fbd0c2213da414687ebf995a65d5a6104"
+dependencies = [
+ "bytemuck",
+ "byteorder-lite",
+ "moxcms",
+ "num-traits",
+ "png 0.18.1",
+]
+
+[[package]]
+name = "indexmap"
+version = "1.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
+dependencies = [
+ "autocfg",
+ "hashbrown 0.12.3",
+ "serde",
+]
+
+[[package]]
+name = "indexmap"
+version = "2.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
+dependencies = [
+ "equivalent",
+ "hashbrown 0.17.1",
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "indicatif"
+version = "0.17.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
+dependencies = [
+ "console",
+ "number_prefix",
+ "portable-atomic",
+ "unicode-width",
+ "web-time",
+]
+
+[[package]]
+name = "infer"
+version = "0.19.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a588916bfdfd92e71cacef98a63d9b1f0d74d6599980d11894290e7ddefffcf7"
+dependencies = [
+ "cfb",
+]
+
+[[package]]
+name = "ipnet"
+version = "2.12.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "791930b43c0d5973160d90a8f3894509f2b273430f5c5c73b668636d0287c5c0"
+
+[[package]]
+name = "is_terminal_polyfill"
+version = "1.70.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695"
+
+[[package]]
+name = "itertools"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
+dependencies = [
+ "either",
+]
+
+[[package]]
+name = "itoa"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
+
+[[package]]
+name = "javascriptcore-rs"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca5671e9ffce8ffba57afc24070e906da7fc4b1ba66f2cabebf61bf2ea257fcc"
+dependencies = [
+ "bitflags 1.3.2",
+ "glib",
+ "javascriptcore-rs-sys",
+]
+
+[[package]]
+name = "javascriptcore-rs-sys"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "af1be78d14ffa4b75b66df31840478fef72b51f8c2465d4ca7c194da9f7a5124"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "jiff"
+version = "0.2.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ab1baf72f08796de0260609515130699b890ac25f30e610ad894bc5856cafdb"
+dependencies = [
+ "defmt",
+ "jiff-core",
+ "jiff-static",
+ "jiff-tzdb-platform",
+ "log",
+ "portable-atomic",
+ "portable-atomic-util",
+ "serde_core",
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "jiff-core"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e52fe76043ccecc9005d2305ebaadf7d7fc0cc89ca6baa10a94d6bc68c7128c"
+dependencies = [
+ "defmt",
+ "log",
+]
+
+[[package]]
+name = "jiff-static"
+version = "0.2.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "378268a1116ad67ae6228701118ac9f491d78fda38a40a1f1a9e1348de6f7212"
+dependencies = [
+ "jiff-core",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "jiff-tzdb"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "142bd39932ad231f10513df9ab62661fead8719872150b7ad02a2df79f4e141e"
+
+[[package]]
+name = "jiff-tzdb-platform"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "875a5a69ac2bab1a891711cf5eccbec1ce0341ea805560dcd90b7a2e925132e8"
+dependencies = [
+ "jiff-tzdb",
+]
+
+[[package]]
+name = "jni"
+version = "0.21.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97"
+dependencies = [
+ "cesu8",
+ "cfg-if",
+ "combine",
+ "jni-sys 0.3.1",
+ "log",
+ "thiserror 1.0.69",
+ "walkdir",
+ "windows-sys 0.45.0",
+]
+
+[[package]]
+name = "jni-sys"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258"
+dependencies = [
+ "jni-sys 0.4.1",
+]
+
+[[package]]
+name = "jni-sys"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2"
+dependencies = [
+ "jni-sys-macros",
+]
+
+[[package]]
+name = "jni-sys-macros"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264"
+dependencies = [
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "js-sys"
+version = "0.3.105"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ce57d20d1ea864ce2ac172ab472d409214f4fd359f0b2a2775abdf522e2af99e"
+dependencies = [
+ "cfg-if",
+ "futures-util",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "json-patch"
+version = "3.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "863726d7afb6bc2590eeff7135d923545e5e964f004c2ccf8716c25e70a86f08"
+dependencies = [
+ "jsonptr",
+ "serde",
+ "serde_json",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "jsonptr"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5dea2b27dd239b2556ed7a25ba842fe47fd602e7fc7433c2a8d6106d4d9edd70"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "keyboard-types"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b750dcadc39a09dbadd74e118f6dd6598df77fa01df0cfcdc52c28dece74528a"
+dependencies = [
+ "bitflags 2.13.2",
+ "serde",
+ "unicode-segmentation",
+]
+
+[[package]]
+name = "lazy_static"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+
+[[package]]
+name = "libappindicator"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03589b9607c868cc7ae54c0b2a22c8dc03dd41692d48f2d7df73615c6a95dc0a"
+dependencies = [
+ "glib",
+ "gtk",
+ "gtk-sys",
+ "libappindicator-sys",
+ "log",
+]
+
+[[package]]
+name = "libappindicator-sys"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e9ec52138abedcc58dc17a7c6c0c00a2bdb4f3427c7f63fa97fd0d859155caf"
+dependencies = [
+ "gtk-sys",
+ "libloading 0.7.4",
+ "once_cell",
+]
+
+[[package]]
+name = "libc"
+version = "0.2.189"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
+
+[[package]]
+name = "libdbus-sys"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "328c4789d42200f1eeec05bd86c9c13c7f091d2ba9a6ea35acdf51f31bc0f043"
+dependencies = [
+ "pkg-config",
+]
+
+[[package]]
+name = "libloading"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f"
+dependencies = [
+ "cfg-if",
+ "winapi",
+]
+
+[[package]]
+name = "libloading"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55"
+dependencies = [
+ "cfg-if",
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "libredox"
+version = "0.1.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6480ccc157a1389bb2e4891b24751b0f798ba640d22386f23143fbcc89da195a"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "libsqlite3-sys"
+version = "0.38.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8"
+dependencies = [
+ "cc",
+ "pkg-config",
+ "vcpkg",
+]
+
+[[package]]
+name = "linux-raw-sys"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
+
+[[package]]
+name = "litemap"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae"
+
+[[package]]
+name = "lock_api"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
+dependencies = [
+ "scopeguard",
+]
+
+[[package]]
+name = "log"
+version = "0.4.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6"
+
+[[package]]
+name = "lru-slab"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4050469837a6ff301cd14c1f8f24f88549e6d548f24f64e2148eb0f72cebc51f"
+
+[[package]]
+name = "macro_rules_attribute"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c"
+dependencies = [
+ "macro_rules_attribute-proc_macro",
+ "pastey",
+]
+
+[[package]]
+name = "macro_rules_attribute-proc_macro"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
+
+[[package]]
+name = "markup5ever"
+version = "0.38.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8983d30f2915feeaaab2d6babdd6bc7e9ed1a00b66b5e6d74df19aa9c0e91862"
+dependencies = [
+ "log",
+ "tendril",
+ "web_atoms",
+]
+
+[[package]]
+name = "matchers"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9"
+dependencies = [
+ "regex-automata",
+]
+
+[[package]]
+name = "matchit"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
+
+[[package]]
+name = "memchr"
+version = "2.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
+
+[[package]]
+name = "memo-map"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5449c8c750f1a07ea702bbd212bd999fceece9b3d1508b17023b3e174583124b"
+
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "mime"
+version = "0.3.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
+
+[[package]]
+name = "minijinja"
+version = "2.24.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "86886cf6dbf4e614b19c9a1eec9775f021869d7eadde0fc73921a81b90c9b4c9"
+dependencies = [
+ "memo-map",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "minijinja-contrib"
+version = "2.24.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bd3e5f077bc2379f0f7d911e7cfdd921114ed99fc884533dca502944cb355b11"
+dependencies = [
+ "minijinja",
+ "serde",
+]
+
+[[package]]
+name = "minimal-lexical"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
+
+[[package]]
+name = "miniz_oxide"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316"
+dependencies = [
+ "adler2",
+ "simd-adler32",
+]
+
+[[package]]
+name = "miniz_oxide"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c"
+dependencies = [
+ "adler2",
+ "simd-adler32",
+]
+
+[[package]]
+name = "mio"
+version = "1.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b18443e9c262bfe8fa82f51666e2642c53393f7e5c27b3e1aeab922cff5b9d8"
+dependencies = [
+ "libc",
+ "wasi",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "monostate"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67"
+dependencies = [
+ "monostate-impl",
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "monostate-impl"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "moxcms"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bb85c154ba489f01b25c0d36ae69a87e4a1c73a72631fc6c0eb6dde34a73e44b"
+dependencies = [
+ "num-traits",
+ "pxfm",
+]
+
+[[package]]
+name = "muda"
+version = "0.19.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1dd04e60bc0b07438a6771710ee1698f98f6ebbc7f89b61264af1563b8aeb878"
+dependencies = [
+ "crossbeam-channel",
+ "dpi",
+ "gtk",
+ "keyboard-types",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+ "once_cell",
+ "png 0.18.1",
+ "serde",
+ "thiserror 2.0.20",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "ndk"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3f42e7bbe13d351b6bead8286a43aac9534b82bd3cc43e47037f012ebfd62d4"
+dependencies = [
+ "bitflags 2.13.2",
+ "jni-sys 0.3.1",
+ "log",
+ "ndk-sys",
+ "num_enum",
+ "raw-window-handle",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "ndk-sys"
+version = "0.6.0+11769913"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873"
+dependencies = [
+ "jni-sys 0.3.1",
+]
+
+[[package]]
+name = "new_debug_unreachable"
+version = "1.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086"
+
+[[package]]
+name = "nom"
+version = "7.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
+dependencies = [
+ "memchr",
+ "minimal-lexical",
+]
+
+[[package]]
+name = "nu-ansi-term"
+version = "0.50.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
+dependencies = [
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "num-conv"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441"
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "num_enum"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26"
+dependencies = [
+ "num_enum_derive",
+ "rustversion",
+]
+
+[[package]]
+name = "num_enum_derive"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8"
+dependencies = [
+ "proc-macro-crate 3.5.0",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "number_prefix"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
+
+[[package]]
+name = "objc2"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f"
+dependencies = [
+ "objc2-encode",
+ "objc2-exception-helper",
+]
+
+[[package]]
+name = "objc2-app-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d49e936b501e5c5bf01fda3a9452ff86dc3ea98ad5f283e1455153142d97518c"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-cloud-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73ad74d880bb43877038da939b7427bba67e9dd42004a18b809ba7d87cee241c"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-data"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b402a653efbb5e82ce4df10683b6b28027616a2715e90009947d50b8dd298fa"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-foundation"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536"
+dependencies = [
+ "bitflags 2.13.2",
+ "dispatch2",
+ "objc2",
+]
+
+[[package]]
+name = "objc2-core-graphics"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807"
+dependencies = [
+ "bitflags 2.13.2",
+ "dispatch2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-io-surface",
+]
+
+[[package]]
+name = "objc2-core-image"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5d563b38d2b97209f8e861173de434bd0214cf020e3423a52624cd1d989f006"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-location"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca347214e24bc973fc025fd0d36ebb179ff30536ed1f80252706db19ee452009"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-text"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0cde0dfb48d25d2b4862161a4d5fcc0e3c24367869ad306b0c9ec0073bfed92d"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+]
+
+[[package]]
+name = "objc2-encode"
+version = "4.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33"
+
+[[package]]
+name = "objc2-exception-helper"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7a1c5fbb72d7735b076bb47b578523aedc40f3c439bea6dfd595c089d79d98a"
+dependencies = [
+ "cc",
+]
+
+[[package]]
+name = "objc2-foundation"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-core-foundation",
+]
+
+[[package]]
+name = "objc2-io-surface"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-core-foundation",
+]
+
+[[package]]
+name = "objc2-quartz-core"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-ui-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d87d638e33c06f577498cbcc50491496a3ed4246998a7fbba7ccb98b1e7eab22"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-cloud-kit",
+ "objc2-core-data",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-core-image",
+ "objc2-core-location",
+ "objc2-core-text",
+ "objc2-foundation",
+ "objc2-quartz-core",
+ "objc2-user-notifications",
+]
+
+[[package]]
+name = "objc2-user-notifications"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9df9128cbbfef73cda168416ccf7f837b62737d748333bfe9ab71c245d76613e"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-web-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2e5aaab980c433cf470df9d7af96a7b46a9d892d521a2cbbb2f8a4c16751e7f"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "once_cell_polyfill"
+version = "1.70.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
+
+[[package]]
+name = "onig"
+version = "6.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2"
+dependencies = [
+ "bitflags 2.13.2",
+ "libc",
+ "once_cell",
+ "onig_sys",
+]
+
+[[package]]
+name = "onig_sys"
+version = "69.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7"
+dependencies = [
+ "cc",
+ "pkg-config",
+]
+
+[[package]]
+name = "option-ext"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
+
+[[package]]
+name = "pango"
+version = "0.18.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ca27ec1eb0457ab26f3036ea52229edbdb74dee1edd29063f5b9b010e7ebee4"
+dependencies = [
+ "gio",
+ "glib",
+ "libc",
+ "once_cell",
+ "pango-sys",
+]
+
+[[package]]
+name = "pango-sys"
+version = "0.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "436737e391a843e5933d6d9aa102cb126d501e815b83601365a948a518555dc5"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "parking_lot"
+version = "0.12.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
+dependencies = [
+ "lock_api",
+ "parking_lot_core",
+]
+
+[[package]]
+name = "parking_lot_core"
+version = "0.9.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "redox_syscall",
+ "smallvec",
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "paste"
+version = "1.0.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
+
+[[package]]
+name = "pastey"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
+
+[[package]]
+name = "percent-encoding"
+version = "2.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
+
+[[package]]
+name = "phf"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf"
+dependencies = [
+ "phf_macros",
+ "phf_shared",
+ "serde",
+]
+
+[[package]]
+name = "phf_codegen"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49aa7f9d80421bca176ca8dbfebe668cc7a2684708594ec9f3c0db0805d5d6e1"
+dependencies = [
+ "phf_generator",
+ "phf_shared",
+]
+
+[[package]]
+name = "phf_generator"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "135ace3a761e564ec88c03a77317a7c6b80bb7f7135ef2544dbe054243b89737"
+dependencies = [
+ "fastrand",
+ "phf_shared",
+]
+
+[[package]]
+name = "phf_macros"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "812f032b54b1e759ccd5f8b6677695d5268c588701effba24601f6932f8269ef"
+dependencies = [
+ "phf_generator",
+ "phf_shared",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "phf_shared"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266"
+dependencies = [
+ "siphasher",
+]
+
+[[package]]
+name = "pin-project-lite"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
+
+[[package]]
+name = "pkcs8"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
+dependencies = [
+ "der",
+ "spki",
+]
+
+[[package]]
+name = "pkg-config"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548"
+
+[[package]]
+name = "plist"
+version = "1.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2896bade328c13f7042a297ea5ac5b0951f6cf989dea5f32c2fd98da398195cb"
+dependencies = [
+ "base64 0.23.1",
+ "indexmap 2.14.0",
+ "quick-xml",
+ "serde",
+ "time",
+]
+
+[[package]]
+name = "png"
+version = "0.17.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "82151a2fc869e011c153adc57cf2789ccb8d9906ce52c0b39a6b5697749d7526"
+dependencies = [
+ "bitflags 1.3.2",
+ "crc32fast",
+ "fdeflate",
+ "flate2",
+ "miniz_oxide 0.8.9",
+]
+
+[[package]]
+name = "png"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61"
+dependencies = [
+ "bitflags 2.13.2",
+ "crc32fast",
+ "fdeflate",
+ "flate2",
+ "miniz_oxide 0.8.9",
+]
+
+[[package]]
+name = "portable-atomic"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "potential_utf"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661"
+dependencies = [
+ "zerovec",
+]
+
+[[package]]
+name = "powerfmt"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
+
+[[package]]
+name = "ppv-lite86"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
+dependencies = [
+ "zerocopy",
+]
+
+[[package]]
+name = "precomputed-hash"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c"
+
+[[package]]
+name = "proc-macro-crate"
+version = "1.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f4c021e1093a56626774e81216a4ce732a735e5bad4868a03f3ed65ca0c3919"
+dependencies = [
+ "once_cell",
+ "toml_edit 0.19.15",
+]
+
+[[package]]
+name = "proc-macro-crate"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b00f26d3400549137f92511a46ac1cd8ce37cb5598a96d382381458b992a5d24"
+dependencies = [
+ "toml_datetime 0.6.3",
+ "toml_edit 0.20.2",
+]
+
+[[package]]
+name = "proc-macro-crate"
+version = "3.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f"
+dependencies = [
+ "toml_edit 0.25.15+spec-1.1.0",
+]
+
+[[package]]
+name = "proc-macro-error"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c"
+dependencies = [
+ "proc-macro-error-attr",
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro-error-attr"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.107"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pxfm"
+version = "0.1.30"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d55d956fa96f5ec02be2e13af0e20391a5aa83d6a074e3ad368959d0fab299ea"
+
+[[package]]
+name = "quick-xml"
+version = "0.42.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41b1177fdf999d2321d3fb46ff47159d9c1fb9ad66a4879f8c50a0b504615e9b"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "quinn"
+version = "0.11.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8"
+dependencies = [
+ "bytes",
+ "cfg_aliases",
+ "pin-project-lite",
+ "quinn-proto",
+ "quinn-udp",
+ "rustc-hash",
+ "rustls",
+ "socket2",
+ "thiserror 2.0.20",
+ "tokio",
+ "tracing",
+ "web-time",
+]
+
+[[package]]
+name = "quinn-proto"
+version = "0.11.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04759210543be93709136e28212294a659ef5001836ff4eab4d663e4529bba83"
+dependencies = [
+ "bytes",
+ "getrandom 0.4.3",
+ "lru-slab",
+ "rand 0.10.2",
+ "rand_pcg",
+ "ring",
+ "rustc-hash",
+ "rustls",
+ "rustls-pki-types",
+ "slab",
+ "thiserror 2.0.20",
+ "tinyvec",
+ "tracing",
+ "web-time",
+]
+
+[[package]]
+name = "quinn-udp"
+version = "0.5.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694"
+dependencies = [
+ "cfg_aliases",
+ "libc",
+ "once_cell",
+ "socket2",
+ "tracing",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "r-efi"
+version = "5.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
+
+[[package]]
+name = "r-efi"
+version = "6.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
+
+[[package]]
+name = "rand"
+version = "0.8.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c"
+dependencies = [
+ "libc",
+ "rand_chacha 0.3.1",
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "rand"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41"
+dependencies = [
+ "rand_chacha 0.9.0",
+ "rand_core 0.9.5",
+]
+
+[[package]]
+name = "rand"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80"
+dependencies = [
+ "chacha20",
+ "getrandom 0.4.3",
+ "rand_core 0.10.1",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
+dependencies = [
+ "ppv-lite86",
+ "rand_core 0.9.5",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+dependencies = [
+ "getrandom 0.2.17",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c"
+dependencies = [
+ "getrandom 0.3.4",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69"
+
+[[package]]
+name = "rand_pcg"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a"
+dependencies = [
+ "rand_core 0.10.1",
+]
+
+[[package]]
+name = "raw-window-handle"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"
+
+[[package]]
+name = "rayon"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-cond"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f"
+dependencies = [
+ "either",
+ "itertools",
+ "rayon",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "redox_syscall"
+version = "0.5.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
+dependencies = [
+ "bitflags 2.13.2",
+]
+
+[[package]]
+name = "redox_users"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac"
+dependencies = [
+ "getrandom 0.2.17",
+ "libredox",
+ "thiserror 2.0.20",
+]
+
+[[package]]
+name = "ref-cast"
+version = "1.0.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3"
+dependencies = [
+ "ref-cast-impl",
+]
+
+[[package]]
+name = "ref-cast-impl"
+version = "1.0.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "regex"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
+
+[[package]]
+name = "reqwest"
+version = "0.12.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
+dependencies = [
+ "base64 0.22.1",
+ "bytes",
+ "futures-channel",
+ "futures-core",
+ "futures-util",
+ "http",
+ "http-body",
+ "http-body-util",
+ "hyper",
+ "hyper-rustls",
+ "hyper-util",
+ "js-sys",
+ "log",
+ "percent-encoding",
+ "pin-project-lite",
+ "quinn",
+ "rustls",
+ "rustls-pki-types",
+ "serde",
+ "serde_json",
+ "serde_urlencoded",
+ "sync_wrapper",
+ "tokio",
+ "tokio-rustls",
+ "tokio-util",
+ "tower",
+ "tower-http",
+ "tower-service",
+ "url",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "wasm-streams 0.4.2",
+ "web-sys",
+ "webpki-roots",
+]
+
+[[package]]
+name = "reqwest"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029"
+dependencies = [
+ "base64 0.23.1",
+ "bytes",
+ "futures-core",
+ "futures-util",
+ "http",
+ "http-body",
+ "http-body-util",
+ "hyper",
+ "hyper-util",
+ "js-sys",
+ "log",
+ "percent-encoding",
+ "pin-project-lite",
+ "serde",
+ "serde_json",
+ "sync_wrapper",
+ "tokio",
+ "tokio-util",
+ "tower",
+ "tower-http",
+ "tower-service",
+ "url",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "wasm-streams 0.5.0",
+ "web-sys",
+]
+
+[[package]]
+name = "ring"
+version = "0.17.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7"
+dependencies = [
+ "cc",
+ "cfg-if",
+ "getrandom 0.2.17",
+ "libc",
+ "untrusted",
+ "windows-sys 0.52.0",
+]
+
+[[package]]
+name = "rsqlite-vfs"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c"
+dependencies = [
+ "hashbrown 0.16.1",
+ "thiserror 2.0.20",
+]
+
+[[package]]
+name = "rusqlite"
+version = "0.40.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3"
+dependencies = [
+ "bitflags 2.13.2",
+ "fallible-iterator",
+ "fallible-streaming-iterator",
+ "hashlink",
+ "libsqlite3-sys",
+ "smallvec",
+ "sqlite-wasm-rs",
+]
+
+[[package]]
+name = "rustc-hash"
+version = "2.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d"
+
+[[package]]
+name = "rustc_version"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92"
+dependencies = [
+ "semver",
+]
+
+[[package]]
+name = "rustix"
+version = "1.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190"
+dependencies = [
+ "bitflags 2.13.2",
+ "errno",
+ "libc",
+ "linux-raw-sys",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "rustls"
+version = "0.23.44"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6725596c3f2c3a0aef021139e145d4eafe314a6623e4680ca83852b2c67ab2ba"
+dependencies = [
+ "once_cell",
+ "ring",
+ "rustls-pki-types",
+ "rustls-webpki",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
+name = "rustls-pki-types"
+version = "1.15.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96"
+dependencies = [
+ "web-time",
+ "zeroize",
+]
+
+[[package]]
+name = "rustls-webpki"
+version = "0.103.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2"
+dependencies = [
+ "ring",
+ "rustls-pki-types",
+ "untrusted",
+]
+
+[[package]]
+name = "rustversion"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f"
+
+[[package]]
+name = "ryu"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
+
+[[package]]
+name = "same-file"
+version = "1.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
+dependencies = [
+ "winapi-util",
+]
+
+[[package]]
+name = "schemars"
+version = "0.8.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615"
+dependencies = [
+ "dyn-clone",
+ "indexmap 1.9.3",
+ "schemars_derive",
+ "serde",
+ "serde_json",
+ "url",
+ "uuid",
+]
+
+[[package]]
+name = "schemars"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f"
+dependencies = [
+ "dyn-clone",
+ "ref-cast",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "schemars"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a"
+dependencies = [
+ "dyn-clone",
+ "ref-cast",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "schemars_derive"
+version = "0.8.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "serde_derive_internals",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "scopeguard"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
+
+[[package]]
+name = "selectors"
+version = "0.36.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c5d9c0c92a92d33f08817311cf3f2c29a3538a8240e94a6a3c622ce652d7e00c"
+dependencies = [
+ "bitflags 2.13.2",
+ "cssparser",
+ "derive_more",
+ "log",
+ "new_debug_unreachable",
+ "phf",
+ "phf_codegen",
+ "precomputed-hash",
+ "rustc-hash",
+ "servo_arc",
+ "smallvec",
+]
+
+[[package]]
+name = "semver"
+version = "1.0.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd"
+dependencies = [
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "serde"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
+dependencies = [
+ "serde_core",
+ "serde_derive",
+]
+
+[[package]]
+name = "serde-untagged"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9faf48a4a2d2693be24c6289dbe26552776eb7737074e6722891fadbe6c5058"
+dependencies = [
+ "erased-serde",
+ "serde",
+ "serde_core",
+ "typeid",
+]
+
+[[package]]
+name = "serde_core"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "serde_derive_internals"
+version = "0.29.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "serde_json"
+version = "1.0.151"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14"
+dependencies = [
+ "indexmap 2.14.0",
+ "itoa",
+ "memchr",
+ "serde",
+ "serde_core",
+ "zmij",
+]
+
+[[package]]
+name = "serde_path_to_error"
+version = "0.1.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457"
+dependencies = [
+ "itoa",
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "serde_repr"
+version = "0.1.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8d3b1629de253c70a0508c3899572da79ca359fdab27c7920ff00406df418906"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "serde_spanned"
+version = "0.6.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "serde_spanned"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "serde_urlencoded"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd"
+dependencies = [
+ "form_urlencoded",
+ "itoa",
+ "ryu",
+ "serde",
+]
+
+[[package]]
+name = "serde_with"
+version = "3.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "935177bb8c0cd8ca1a4e6d1a2ac8988bea69cab4f9d3a31311e012ad27868ea4"
+dependencies = [
+ "base64 0.23.1",
+ "bs58",
+ "chrono",
+ "hex",
+ "indexmap 1.9.3",
+ "indexmap 2.14.0",
+ "jiff",
+ "schemars 0.9.0",
+ "schemars 1.2.2",
+ "serde_core",
+ "serde_json",
+ "serde_with_macros",
+ "time",
+]
+
+[[package]]
+name = "serde_with_macros"
+version = "3.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d607aa01a3cb0ad757d6fd216136910db3c97b102fe686585689615a02dbcdc"
+dependencies = [
+ "darling 0.24.1",
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "serialize-to-javascript"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04f3666a07a197cdb77cdf306c32be9b7f598d7060d50cfd4d5aa04bfd92f6c5"
+dependencies = [
+ "serde",
+ "serde_json",
+ "serialize-to-javascript-impl",
+]
+
+[[package]]
+name = "serialize-to-javascript-impl"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "772ee033c0916d670af7860b6e1ef7d658a4629a6d0b4c8c3e67f09b3765b75d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "servo_arc"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "170fb83ab34de17dc69aa7c67482b22218ddb85da56546f9bd6b929e32a05930"
+dependencies = [
+ "stable_deref_trait",
+]
+
+[[package]]
+name = "sha2"
+version = "0.10.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
+dependencies = [
+ "cfg-if",
+ "cpufeatures 0.2.17",
+ "digest",
+]
+
+[[package]]
+name = "sharded-slab"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
+dependencies = [
+ "lazy_static",
+]
+
+[[package]]
+name = "shlex"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba"
+
+[[package]]
+name = "signal-hook-registry"
+version = "1.4.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b"
+dependencies = [
+ "errno",
+ "libc",
+]
+
+[[package]]
+name = "signature"
+version = "2.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de"
+dependencies = [
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "simd-adler32"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea"
+
+[[package]]
+name = "siphasher"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649"
+
+[[package]]
+name = "slab"
+version = "0.4.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5"
+
+[[package]]
+name = "smallvec"
+version = "1.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba467056f1b547ed52077911161fc86985becbc60e8e1857c8a144dab0def891"
+
+[[package]]
+name = "socket2"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4"
+dependencies = [
+ "libc",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "softbuffer"
+version = "0.4.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aac18da81ebbf05109ab275b157c22a653bb3c12cf884450179942f81bcbf6c3"
+dependencies = [
+ "bytemuck",
+ "js-sys",
+ "ndk",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-foundation",
+ "objc2-quartz-core",
+ "raw-window-handle",
+ "redox_syscall",
+ "tracing",
+ "wasm-bindgen",
+ "web-sys",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "soup3"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "471f924a40f31251afc77450e781cb26d55c0b650842efafc9c6cbd2f7cc4f9f"
+dependencies = [
+ "futures-channel",
+ "gio",
+ "glib",
+ "libc",
+ "soup3-sys",
+]
+
+[[package]]
+name = "soup3-sys"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ebe8950a680a12f24f15ebe1bf70db7af98ad242d9db43596ad3108aab86c27"
+dependencies = [
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "spki"
+version = "0.7.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d"
+dependencies = [
+ "base64ct",
+ "der",
+]
+
+[[package]]
+name = "spm_precompiled"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
+dependencies = [
+ "base64 0.13.1",
+ "nom",
+ "serde",
+ "unicode-segmentation",
+]
+
+[[package]]
+name = "sqlite-wasm-rs"
+version = "0.5.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75"
+dependencies = [
+ "cc",
+ "js-sys",
+ "rsqlite-vfs",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "stable_deref_trait"
+version = "1.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596"
+
+[[package]]
+name = "static_assertions"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
+
+[[package]]
+name = "string_cache"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a18596f8c785a729f2819c0f6a7eae6ebeebdfffbfe4214ae6b087f690e31901"
+dependencies = [
+ "new_debug_unreachable",
+ "parking_lot",
+ "phf_shared",
+ "precomputed-hash",
+]
+
+[[package]]
+name = "string_cache_codegen"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "585635e46db231059f76c5849798146164652513eb9e8ab2685939dd90f29b69"
+dependencies = [
+ "phf_generator",
+ "phf_shared",
+ "proc-macro2",
+ "quote",
+]
+
+[[package]]
+name = "strsim"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
+
+[[package]]
+name = "subtle"
+version = "2.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
+
+[[package]]
+name = "swift-rs"
+version = "1.0.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e45c444e496845d3f2a351146bff59aae4975b2280238df1dfaa0c7d1846f38e"
+dependencies = [
+ "base64 0.21.7",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "syn"
+version = "1.0.109"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
+dependencies = [
+ "proc-macro2",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn"
+version = "2.0.119"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn"
+version = "3.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "12df2e0110f65b775f769bb17ef989067a1d931b2eb822bd4346631eeada89f9"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "sync_wrapper"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263"
+dependencies = [
+ "futures-core",
+]
+
+[[package]]
+name = "synstructure"
+version = "0.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "system-configuration"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b"
+dependencies = [
+ "bitflags 2.13.2",
+ "core-foundation 0.9.4",
+ "system-configuration-sys",
+]
+
+[[package]]
+name = "system-configuration-sys"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4"
+dependencies = [
+ "core-foundation-sys",
+ "libc",
+]
+
+[[package]]
+name = "system-deps"
+version = "6.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a3e535eb8dded36d55ec13eddacd30dec501792ff23a0b1682c38601b8cf2349"
+dependencies = [
+ "cfg-expr",
+ "heck 0.5.0",
+ "pkg-config",
+ "toml 0.8.2",
+ "version-compare",
+]
+
+[[package]]
+name = "tao"
+version = "0.35.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d1c93047acf68669466a34690ac58cca7010bd1b201e1ec86f1fd0a75d3dd4a9"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "core-foundation 0.10.1",
+ "core-graphics",
+ "crossbeam-channel",
+ "dbus",
+ "dispatch2",
+ "dlopen2",
+ "dpi",
+ "gdkwayland-sys",
+ "gdkx11-sys",
+ "gtk",
+ "jni",
+ "libc",
+ "log",
+ "ndk",
+ "ndk-sys",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-foundation",
+ "objc2-ui-kit",
+ "once_cell",
+ "parking_lot",
+ "percent-encoding",
+ "raw-window-handle",
+ "tao-macros",
+ "unicode-segmentation",
+ "url",
+ "windows",
+ "windows-core 0.61.2",
+ "windows-version",
+ "x11-dl",
+]
+
+[[package]]
+name = "tao-macros"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5f7eeb6d99155545da6150a1795945f16ac9c178deb2a5f2e74d776107bd5849"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
+[[package]]
+name = "tauri"
+version = "2.11.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "667b20e2726d572dea2de7370da16e188eb06008faf9a92fab7cdc46791190b5"
+dependencies = [
+ "anyhow",
+ "bytes",
+ "cookie",
+ "dirs",
+ "dunce",
+ "embed_plist",
+ "getrandom 0.3.4",
+ "glob",
+ "gtk",
+ "heck 0.5.0",
+ "http",
+ "image",
+ "jni",
+ "libc",
+ "log",
+ "mime",
+ "muda",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-foundation",
+ "objc2-ui-kit",
+ "objc2-web-kit",
+ "percent-encoding",
+ "plist",
+ "raw-window-handle",
+ "reqwest 0.13.5",
+ "serde",
+ "serde_json",
+ "serde_repr",
+ "serialize-to-javascript",
+ "swift-rs",
+ "tauri-build",
+ "tauri-macros",
+ "tauri-runtime",
+ "tauri-runtime-wry",
+ "tauri-utils",
+ "thiserror 2.0.20",
+ "tokio",
+ "tray-icon",
+ "url",
+ "webkit2gtk",
+ "webview2-com",
+ "window-vibrancy",
+ "windows",
+]
+
+[[package]]
+name = "tauri-build"
+version = "2.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc9ce40b16101cb6ea63d3e221567affd1c3a9205f95d7bc574941a10636b632"
+dependencies = [
+ "anyhow",
+ "cargo_toml",
+ "dirs",
+ "glob",
+ "heck 0.5.0",
+ "json-patch",
+ "schemars 0.8.22",
+ "semver",
+ "serde",
+ "serde_json",
+ "tauri-utils",
+ "tauri-winres",
+ "walkdir",
+]
+
+[[package]]
+name = "tauri-codegen"
+version = "2.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08279169ff42f8fc45a1dbc9dcae888893ba95288142e5880c59b93a26d2cfc5"
+dependencies = [
+ "base64 0.22.1",
+ "brotli",
+ "ico",
+ "json-patch",
+ "plist",
+ "png 0.17.16",
+ "proc-macro2",
+ "quote",
+ "semver",
+ "serde",
+ "serde_json",
+ "sha2",
+ "syn 2.0.119",
+ "tauri-utils",
+ "thiserror 2.0.20",
+ "time",
+ "url",
+ "uuid",
+ "walkdir",
+]
+
+[[package]]
+name = "tauri-macros"
+version = "2.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8b394794f399a421811d06966343e7933fcae92d59f5180b9388d1174497a45"
+dependencies = [
+ "heck 0.5.0",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+ "tauri-codegen",
+ "tauri-utils",
+]
+
+[[package]]
+name = "tauri-runtime"
+version = "2.11.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0b4bc95aed361b0019067d189a1174a603d460d0f6c72606512d59fc9c12ec8"
+dependencies = [
+ "cookie",
+ "dpi",
+ "gtk",
+ "http",
+ "jni",
+ "objc2",
+ "objc2-ui-kit",
+ "objc2-web-kit",
+ "raw-window-handle",
+ "serde",
+ "serde_json",
+ "tauri-utils",
+ "thiserror 2.0.20",
+ "url",
+ "webkit2gtk",
+ "webview2-com",
+ "windows",
+]
+
+[[package]]
+name = "tauri-runtime-wry"
+version = "2.11.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4e6fac707727b7a2f48e4ded90976324267371073edbb415ffb73bb0458d203f"
+dependencies = [
+ "gtk",
+ "http",
+ "jni",
+ "log",
+ "objc2",
+ "objc2-app-kit",
+ "once_cell",
+ "percent-encoding",
+ "raw-window-handle",
+ "softbuffer",
+ "tao",
+ "tauri-runtime",
+ "tauri-utils",
+ "url",
+ "webkit2gtk",
+ "webview2-com",
+ "windows",
+ "wry",
+]
+
+[[package]]
+name = "tauri-utils"
+version = "2.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3e176a18e67764923c4f1ce66f25ae4abe5f688384d5eb1a0fa6c77f3d90f887"
+dependencies = [
+ "anyhow",
+ "brotli",
+ "cargo_metadata",
+ "ctor",
+ "dom_query",
+ "dunce",
+ "glob",
+ "http",
+ "infer",
+ "json-patch",
+ "log",
+ "memchr",
+ "phf",
+ "plist",
+ "proc-macro2",
+ "quote",
+ "regex",
+ "schemars 0.8.22",
+ "semver",
+ "serde",
+ "serde-untagged",
+ "serde_json",
+ "serde_with",
+ "swift-rs",
+ "thiserror 2.0.20",
+ "toml 1.1.6+spec-1.1.0",
+ "url",
+ "urlpattern",
+ "uuid",
+ "walkdir",
+]
+
+[[package]]
+name = "tauri-winres"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc65d45c68858bfe420dd29e834b5d15dbecf8a07a8a16cf4d532c7b1f69d4b6"
+dependencies = [
+ "dunce",
+ "embed-resource",
+ "toml 1.1.6+spec-1.1.0",
+]
+
+[[package]]
+name = "tendril"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5fed54709c5b3a53d09bb1c113ea4f5ceafd1e772ddcb0030a82e1d56c087b08"
+dependencies = [
+ "new_debug_unreachable",
+]
+
+[[package]]
+name = "termcolor"
+version = "1.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755"
+dependencies = [
+ "winapi-util",
+]
+
+[[package]]
+name = "thiserror"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
+dependencies = [
+ "thiserror-impl 1.0.69",
+]
+
+[[package]]
+name = "thiserror"
+version = "2.0.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f"
+dependencies = [
+ "thiserror-impl 2.0.20",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "2.0.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "thread_local"
+version = "1.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070"
+dependencies = [
+ "cfg-if",
+]
+
+[[package]]
+name = "time"
+version = "0.3.55"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134"
+dependencies = [
+ "deranged",
+ "num-conv",
+ "powerfmt",
+ "serde_core",
+ "time-core",
+ "time-macros",
+]
+
+[[package]]
+name = "time-core"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109"
+
+[[package]]
+name = "time-macros"
+version = "0.2.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e689342a48d2ea927c87ea50cabf8594854bf940e9310208848d680d668ed85"
+dependencies = [
+ "num-conv",
+ "time-core",
+]
+
+[[package]]
+name = "tinystr"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643"
+dependencies = [
+ "displaydoc",
+ "zerovec",
+]
+
+[[package]]
+name = "tinyvec"
+version = "1.13.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd3ca314f692efd6c868f8408f53fe444634a845f96c028b97d35f6a1f79f0ee"
+
+[[package]]
+name = "tokenizers"
+version = "0.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a620b996116a59e184c2fa2dfd8251ea34a36d0a514758c6f966386bd2e03476"
+dependencies = [
+ "ahash",
+ "aho-corasick",
+ "compact_str",
+ "dary_heap",
+ "derive_builder",
+ "esaxx-rs",
+ "getrandom 0.3.4",
+ "indicatif",
+ "itertools",
+ "log",
+ "macro_rules_attribute",
+ "monostate",
+ "onig",
+ "paste",
+ "rand 0.9.5",
+ "rayon",
+ "rayon-cond",
+ "regex",
+ "regex-syntax",
+ "serde",
+ "serde_json",
+ "spm_precompiled",
+ "thiserror 2.0.20",
+ "unicode-normalization-alignments",
+ "unicode-segmentation",
+ "unicode_categories",
+]
+
+[[package]]
+name = "tokio"
+version = "1.53.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed"
+dependencies = [
+ "bytes",
+ "libc",
+ "mio",
+ "pin-project-lite",
+ "signal-hook-registry",
+ "socket2",
+ "tokio-macros",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "tokio-macros"
+version = "2.7.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "tokio-rustls"
+version = "0.26.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0c85f2c3ef0b1cd58b36682f4b17aaa995f0e5db534d85692b4903abce21f67"
+dependencies = [
+ "rustls",
+ "tokio",
+]
+
+[[package]]
+name = "tokio-stream"
+version = "0.1.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b"
+dependencies = [
+ "futures-core",
+ "pin-project-lite",
+ "tokio",
+]
+
+[[package]]
+name = "tokio-util"
+version = "0.7.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52"
+dependencies = [
+ "bytes",
+ "futures-core",
+ "futures-sink",
+ "pin-project-lite",
+ "tokio",
+]
+
+[[package]]
+name = "toml"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "185d8ab0dfbb35cf1399a6344d8484209c088f75f8f68230da55d48d95d43e3d"
+dependencies = [
+ "serde",
+ "serde_spanned 0.6.9",
+ "toml_datetime 0.6.3",
+ "toml_edit 0.20.2",
+]
+
+[[package]]
+name = "toml"
+version = "0.9.12+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863"
+dependencies = [
+ "indexmap 2.14.0",
+ "serde_core",
+ "serde_spanned 1.1.1",
+ "toml_datetime 0.7.5+spec-1.1.0",
+ "toml_parser",
+ "toml_writer",
+ "winnow 0.7.15",
+]
+
+[[package]]
+name = "toml"
+version = "1.1.6+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "920602543f0911ab71da12c50d59701da54c196d1a2bf5cb4b75667f137a406a"
+dependencies = [
+ "indexmap 2.14.0",
+ "serde_core",
+ "serde_spanned 1.1.1",
+ "toml_datetime 1.1.1+spec-1.1.0",
+ "toml_parser",
+ "toml_writer",
+ "winnow 1.0.4",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7cda73e2f1397b1262d6dfdcef8aafae14d1de7748d66822d3bfeeb6d03e5e4b"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "0.7.5+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "1.1.1+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.19.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421"
+dependencies = [
+ "indexmap 2.14.0",
+ "toml_datetime 0.6.3",
+ "winnow 0.5.40",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "396e4d48bbb2b7554c944bde63101b5ae446cff6ec4a24227428f15eb72ef338"
+dependencies = [
+ "indexmap 2.14.0",
+ "serde",
+ "serde_spanned 0.6.9",
+ "toml_datetime 0.6.3",
+ "winnow 0.5.40",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.25.15+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1340ea94a5856333492c9064b02c778b191dd2c853778d9609debdcdfea3a614"
+dependencies = [
+ "indexmap 2.14.0",
+ "toml_datetime 1.1.1+spec-1.1.0",
+ "toml_parser",
+ "winnow 1.0.4",
+]
+
+[[package]]
+name = "toml_parser"
+version = "1.1.3+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56"
+dependencies = [
+ "winnow 1.0.4",
+]
+
+[[package]]
+name = "toml_writer"
+version = "1.1.2+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2"
+
+[[package]]
+name = "tower"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4"
+dependencies = [
+ "futures-core",
+ "futures-util",
+ "pin-project-lite",
+ "sync_wrapper",
+ "tokio",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
+[[package]]
+name = "tower-http"
+version = "0.6.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840"
+dependencies = [
+ "bitflags 2.13.2",
+ "bytes",
+ "futures-util",
+ "http",
+ "http-body",
+ "pin-project-lite",
+ "tower",
+ "tower-layer",
+ "tower-service",
+ "url",
+]
+
+[[package]]
+name = "tower-layer"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e"
+
+[[package]]
+name = "tower-service"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3"
+
+[[package]]
+name = "tracing"
+version = "0.1.44"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
+dependencies = [
+ "log",
+ "pin-project-lite",
+ "tracing-attributes",
+ "tracing-core",
+]
+
+[[package]]
+name = "tracing-attributes"
+version = "0.1.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "tracing-core"
+version = "0.1.36"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a"
+dependencies = [
+ "once_cell",
+ "valuable",
+]
+
+[[package]]
+name = "tracing-log"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3"
+dependencies = [
+ "log",
+ "once_cell",
+ "tracing-core",
+]
+
+[[package]]
+name = "tracing-subscriber"
+version = "0.3.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
+dependencies = [
+ "matchers",
+ "nu-ansi-term",
+ "once_cell",
+ "regex-automata",
+ "sharded-slab",
+ "smallvec",
+ "thread_local",
+ "tracing",
+ "tracing-core",
+ "tracing-log",
+]
+
+[[package]]
+name = "tray-icon"
+version = "0.24.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "045979e3f037cd18ad1cb2a419dfda133c5c29c9f3453370079f2255d46c257e"
+dependencies = [
+ "crossbeam-channel",
+ "dirs",
+ "libappindicator",
+ "muda",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-foundation",
+ "once_cell",
+ "png 0.18.1",
+ "serde",
+ "thiserror 2.0.20",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "try-lock"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
+
+[[package]]
+name = "ts-rs"
+version = "12.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "756050066659291d47a554a9f558125db17428b073c5ffce1daf5dcb0f7231d8"
+dependencies = [
+ "chrono",
+ "serde_json",
+ "thiserror 2.0.20",
+ "ts-rs-macros",
+]
+
+[[package]]
+name = "ts-rs-macros"
+version = "12.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38d90eea51bc7988ef9e674bf80a85ba6804739e535e9cab48e4bb34a8b652aa"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+ "termcolor",
+]
+
+[[package]]
+name = "typeid"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c"
+
+[[package]]
+name = "typenum"
+version = "1.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
+
+[[package]]
+name = "unic-char-property"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221"
+dependencies = [
+ "unic-char-range",
+]
+
+[[package]]
+name = "unic-char-range"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc"
+
+[[package]]
+name = "unic-common"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc"
+
+[[package]]
+name = "unic-ucd-ident"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987"
+dependencies = [
+ "unic-char-property",
+ "unic-char-range",
+ "unic-ucd-version",
+]
+
+[[package]]
+name = "unic-ucd-version"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4"
+dependencies = [
+ "unic-common",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
+
+[[package]]
+name = "unicode-normalization-alignments"
+version = "0.1.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
+dependencies = [
+ "smallvec",
+]
+
+[[package]]
+name = "unicode-segmentation"
+version = "1.13.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
+
+[[package]]
+name = "unicode-width"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
+
+[[package]]
+name = "unicode_categories"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
+
+[[package]]
+name = "untrusted"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
+
+[[package]]
+name = "url"
+version = "2.5.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed"
+dependencies = [
+ "form_urlencoded",
+ "idna",
+ "percent-encoding",
+ "serde",
+ "serde_derive",
+]
+
+[[package]]
+name = "urlpattern"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "70acd30e3aa1450bc2eece896ce2ad0d178e9c079493819301573dae3c37ba6d"
+dependencies = [
+ "regex",
+ "serde",
+ "unic-ucd-ident",
+ "url",
+]
+
+[[package]]
+name = "utf8_iter"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
+
+[[package]]
+name = "utf8parse"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
+
+[[package]]
+name = "uuid"
+version = "1.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2ef6dac1e96601b4fb3acccccff2139741fcb757cb9a36089bf5be91cfb285ce"
+dependencies = [
+ "getrandom 0.4.3",
+ "js-sys",
+ "serde_core",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "valuable"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
+
+[[package]]
+name = "vcpkg"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
+
+[[package]]
+name = "version-compare"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03c2856837ef78f57382f06b2b8563a2f512f7185d732608fd9176cb3b8edf0e"
+
+[[package]]
+name = "version_check"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
+
+[[package]]
+name = "vswhom"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "be979b7f07507105799e854203b470ff7c78a1639e330a58f183b5fea574608b"
+dependencies = [
+ "libc",
+ "vswhom-sys",
+]
+
+[[package]]
+name = "vswhom-sys"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb067e4cbd1ff067d1df46c9194b5de0e98efd2810bbc95c5d5e5f25a3231150"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "walkdir"
+version = "2.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
+dependencies = [
+ "same-file",
+ "winapi-util",
+]
+
+[[package]]
+name = "want"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e"
+dependencies = [
+ "try-lock",
+]
+
+[[package]]
+name = "wasi"
+version = "0.11.1+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
+
+[[package]]
+name = "wasip2"
+version = "1.0.4+wasi-0.2.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487"
+dependencies = [
+ "wit-bindgen",
+]
+
+[[package]]
+name = "wasm-bindgen"
+version = "0.2.128"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aecb87a33d3b0c5e3b7aa46336eaf486cffafbd281b195e4c8b80d50df2351bf"
+dependencies = [
+ "cfg-if",
+ "once_cell",
+ "rustversion",
+ "wasm-bindgen-macro",
+ "wasm-bindgen-shared",
+]
+
+[[package]]
+name = "wasm-bindgen-futures"
+version = "0.4.78"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6ef4c5d3d2cdf5c54f4231181768f5510842e350db025faf1f7163b1030ed928"
+dependencies = [
+ "js-sys",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "wasm-bindgen-macro"
+version = "0.2.128"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a690d511e3c1a8b3a55e33511e3c2c00c78415cd23650f32b808627f5696b9ed"
+dependencies = [
+ "quote",
+ "wasm-bindgen-macro-support",
+]
+
+[[package]]
+name = "wasm-bindgen-macro-support"
+version = "0.2.128"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "411e4887f0071ef2d2164a9d5fdf2d20efbef78fccd3a78b0c10a1dc5295e48a"
+dependencies = [
+ "bumpalo",
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+ "wasm-bindgen-shared",
+]
+
+[[package]]
+name = "wasm-bindgen-shared"
+version = "0.2.128"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "81941cd78d0c92026c33e5e01312845a4cb1e9af3407f9134b100dd03144103e"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "wasm-streams"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
+dependencies = [
+ "futures-util",
+ "js-sys",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "web-sys",
+]
+
+[[package]]
+name = "wasm-streams"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb"
+dependencies = [
+ "futures-util",
+ "js-sys",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "web-sys",
+]
+
+[[package]]
+name = "web-sys"
+version = "0.3.105"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fbddc4a036f00ec4f18c83445bd3115cb306a91da554919a099d9222fe4a7f8"
+dependencies = [
+ "js-sys",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "web-time"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb"
+dependencies = [
+ "js-sys",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "web_atoms"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba8b815c1b593dc0baf78dd0f4fc8fdb2de53198fb1163738093e9a311c33fb3"
+dependencies = [
+ "phf",
+ "phf_codegen",
+ "string_cache",
+ "string_cache_codegen",
+]
+
+[[package]]
+name = "webkit2gtk"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1027150013530fb2eaf806408df88461ae4815a45c541c8975e61d6f2fc4793"
+dependencies = [
+ "bitflags 1.3.2",
+ "cairo-rs",
+ "gdk",
+ "gdk-sys",
+ "gio",
+ "gio-sys",
+ "glib",
+ "glib-sys",
+ "gobject-sys",
+ "gtk",
+ "gtk-sys",
+ "javascriptcore-rs",
+ "libc",
+ "once_cell",
+ "soup3",
+ "webkit2gtk-sys",
+]
+
+[[package]]
+name = "webkit2gtk-sys"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "916a5f65c2ef0dfe12fff695960a2ec3d4565359fdbb2e9943c974e06c734ea5"
+dependencies = [
+ "bitflags 1.3.2",
+ "cairo-sys-rs",
+ "gdk-sys",
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "gtk-sys",
+ "javascriptcore-rs-sys",
+ "libc",
+ "pkg-config",
+ "soup3-sys",
+ "system-deps",
+]
+
+[[package]]
+name = "webpki-roots"
+version = "1.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a"
+dependencies = [
+ "rustls-pki-types",
+]
+
+[[package]]
+name = "webview2-com"
+version = "0.38.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7130243a7a5b33c54a444e54842e6a9e133de08b5ad7b5861cd8ed9a6a5bc96a"
+dependencies = [
+ "webview2-com-macros",
+ "webview2-com-sys",
+ "windows",
+ "windows-core 0.61.2",
+ "windows-implement",
+ "windows-interface",
+]
+
+[[package]]
+name = "webview2-com-macros"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67a921c1b6914c367b2b823cd4cde6f96beec77d30a939c8199bb377cf9b9b54"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "webview2-com-sys"
+version = "0.38.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "381336cfffd772377d291702245447a5251a2ffa5bad679c99e61bc48bacbf9c"
+dependencies = [
+ "thiserror 2.0.20",
+ "windows",
+ "windows-core 0.61.2",
+]
+
+[[package]]
+name = "winapi"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
+dependencies = [
+ "winapi-i686-pc-windows-gnu",
+ "winapi-x86_64-pc-windows-gnu",
+]
+
+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
+
+[[package]]
+name = "winapi-util"
+version = "0.1.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
+dependencies = [
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
+
+[[package]]
+name = "window-vibrancy"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d9bec5a31f3f9362f2258fd0e9c9dd61a9ca432e7306cc78c444258f0dce9a9c"
+dependencies = [
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+ "raw-window-handle",
+ "windows-sys 0.59.0",
+ "windows-version",
+]
+
+[[package]]
+name = "windows"
+version = "0.61.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893"
+dependencies = [
+ "windows-collections",
+ "windows-core 0.61.2",
+ "windows-future",
+ "windows-link 0.1.3",
+ "windows-numerics",
+]
+
+[[package]]
+name = "windows-collections"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8"
+dependencies = [
+ "windows-core 0.61.2",
+]
+
+[[package]]
+name = "windows-core"
+version = "0.61.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3"
+dependencies = [
+ "windows-implement",
+ "windows-interface",
+ "windows-link 0.1.3",
+ "windows-result 0.3.4",
+ "windows-strings 0.4.2",
+]
+
+[[package]]
+name = "windows-core"
+version = "0.62.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
+dependencies = [
+ "windows-implement",
+ "windows-interface",
+ "windows-link 0.2.1",
+ "windows-result 0.4.1",
+ "windows-strings 0.5.1",
+]
+
+[[package]]
+name = "windows-future"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e"
+dependencies = [
+ "windows-core 0.61.2",
+ "windows-link 0.1.3",
+ "windows-threading",
+]
+
+[[package]]
+name = "windows-implement"
+version = "0.60.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "windows-interface"
+version = "0.59.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "windows-link"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a"
+
+[[package]]
+name = "windows-link"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
+
+[[package]]
+name = "windows-numerics"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1"
+dependencies = [
+ "windows-core 0.61.2",
+ "windows-link 0.1.3",
+]
+
+[[package]]
+name = "windows-registry"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720"
+dependencies = [
+ "windows-link 0.2.1",
+ "windows-result 0.4.1",
+ "windows-strings 0.5.1",
+]
+
+[[package]]
+name = "windows-result"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6"
+dependencies = [
+ "windows-link 0.1.3",
+]
+
+[[package]]
+name = "windows-result"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5"
+dependencies = [
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "windows-strings"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57"
+dependencies = [
+ "windows-link 0.1.3",
+]
+
+[[package]]
+name = "windows-strings"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091"
+dependencies = [
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.45.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0"
+dependencies = [
+ "windows-targets 0.42.2",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.52.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
+dependencies = [
+ "windows-targets 0.52.6",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.59.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
+dependencies = [
+ "windows-targets 0.52.6",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.61.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
+dependencies = [
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071"
+dependencies = [
+ "windows_aarch64_gnullvm 0.42.2",
+ "windows_aarch64_msvc 0.42.2",
+ "windows_i686_gnu 0.42.2",
+ "windows_i686_msvc 0.42.2",
+ "windows_x86_64_gnu 0.42.2",
+ "windows_x86_64_gnullvm 0.42.2",
+ "windows_x86_64_msvc 0.42.2",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
+dependencies = [
+ "windows_aarch64_gnullvm 0.52.6",
+ "windows_aarch64_msvc 0.52.6",
+ "windows_i686_gnu 0.52.6",
+ "windows_i686_gnullvm",
+ "windows_i686_msvc 0.52.6",
+ "windows_x86_64_gnu 0.52.6",
+ "windows_x86_64_gnullvm 0.52.6",
+ "windows_x86_64_msvc 0.52.6",
+]
+
+[[package]]
+name = "windows-threading"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b66463ad2e0ea3bbf808b7f1d371311c80e115c0b71d60efc142cafbcfb057a6"
+dependencies = [
+ "windows-link 0.1.3",
+]
+
+[[package]]
+name = "windows-version"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4060a1da109b9d0326b7262c8e12c84df67cc0dbc9e33cf49e01ccc2eb63631"
+dependencies = [
+ "windows-link 0.2.1",
+]
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8"
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
+
+[[package]]
+name = "windows_i686_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
+
+[[package]]
+name = "winnow"
+version = "0.5.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f593a95398737aeed53e489c785df13f3618e41dbcd6718c6addbf1395aa6876"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "winnow"
+version = "0.7.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945"
+
+[[package]]
+name = "winnow"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "winreg"
+version = "0.55.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb5a765337c50e9ec252c2069be9bf91c7df47afb103b642ba3a53bf8101be97"
+dependencies = [
+ "cfg-if",
+ "windows-sys 0.59.0",
+]
+
+[[package]]
+name = "wit-bindgen"
+version = "0.57.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
+
+[[package]]
+name = "writeable"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc"
+
+[[package]]
+name = "wry"
+version = "0.55.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "186f9871daa55fd9c016578b810d149de58367113db7fb72b462d2323ce19514"
+dependencies = [
+ "base64 0.22.1",
+ "block2",
+ "cookie",
+ "crossbeam-channel",
+ "dirs",
+ "dom_query",
+ "dpi",
+ "dunce",
+ "gdkx11",
+ "gtk",
+ "http",
+ "javascriptcore-rs",
+ "jni",
+ "libc",
+ "ndk",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+ "objc2-ui-kit",
+ "objc2-web-kit",
+ "once_cell",
+ "percent-encoding",
+ "raw-window-handle",
+ "sha2",
+ "soup3",
+ "tao-macros",
+ "thiserror 2.0.20",
+ "url",
+ "webkit2gtk",
+ "webkit2gtk-sys",
+ "webview2-com",
+ "windows",
+ "windows-core 0.61.2",
+ "windows-version",
+ "x11-dl",
+]
+
+[[package]]
+name = "x11"
+version = "2.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "502da5464ccd04011667b11c435cb992822c2c0dbde1770c988480d312a0db2e"
+dependencies = [
+ "libc",
+ "pkg-config",
+]
+
+[[package]]
+name = "x11-dl"
+version = "2.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38735924fedd5314a6e548792904ed8c6de6636285cb9fec04d5b1db85c1516f"
+dependencies = [
+ "libc",
+ "once_cell",
+ "pkg-config",
+]
+
+[[package]]
+name = "x11rb"
+version = "0.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9993aa5be5a26815fe2c3eacfc1fde061fc1a1f094bf1ad2a18bf9c495dd7414"
+dependencies = [
+ "gethostname",
+ "rustix",
+ "x11rb-protocol",
+]
+
+[[package]]
+name = "x11rb-protocol"
+version = "0.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ea6fc2961e4ef194dcbfe56bb845534d0dc8098940c7e5c012a258bfec6701bd"
+
+[[package]]
+name = "yoke"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5"
+dependencies = [
+ "stable_deref_trait",
+ "yoke-derive",
+ "zerofrom",
+]
+
+[[package]]
+name = "yoke-derive"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+ "synstructure",
+]
+
+[[package]]
+name = "zerocopy"
+version = "0.8.57"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d35102a9f36d089ccae9e4c6802bc118be4487b80aaffc0ab4e0cf5ce92d2873"
+dependencies = [
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.8.57"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "146c01f5ab44258da43cf276c74a2763db2ff3969c9c652c3f2de07041d0b2bc"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "zerofrom"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272"
+dependencies = [
+ "zerofrom-derive",
+]
+
+[[package]]
+name = "zerofrom-derive"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+ "synstructure",
+]
+
+[[package]]
+name = "zeroize"
+version = "1.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e"
+
+[[package]]
+name = "zerotrie"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ea269c3bd32f0a32c321907a2ae912ba6f4649bb0fc764a15627e99a7095a3f"
+dependencies = [
+ "displaydoc",
+ "yoke",
+ "zerofrom",
+]
+
+[[package]]
+name = "zerovec"
+version = "0.11.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bb0464e17806c1d976d5cba29399c7f08e516e279e2ba493f63123b5fca67dd8"
+dependencies = [
+ "yoke",
+ "zerofrom",
+ "zerovec-derive",
+]
+
+[[package]]
+name = "zerovec-derive"
+version = "0.11.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34df6fc39dbd26ddc9c10e6a2984476e13acce22e64e4487636ef494369225da"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.5",
+]
+
+[[package]]
+name = "zlib-rs"
+version = "0.6.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12"
+
+[[package]]
+name = "zmij"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
diff --git a/macos/Cargo.toml b/macos/Cargo.toml
new file mode 100644
index 0000000..d3bdb39
--- /dev/null
+++ b/macos/Cargo.toml
@@ -0,0 +1,60 @@
+[workspace]
+resolver = "3"
+members = [
+    "crates/edge0-core",
+    "crates/edge0-pull",
+    "crates/edge0d",
+    "crates/edge0-cli",
+    "crates/edge0-engine",
+    "crates/engine-tok",
+    "app/src-tauri",
+]
+
+[workspace.package]
+version = "0.1.0"
+edition = "2021"
+rust-version = "1.88.0"
+license = "Apache-2.0"
+
+[workspace.dependencies]
+edge0-core = { path = "crates/edge0-core" }
+engine-tok = { path = "crates/engine-tok" }
+edge0-pull = { path = "crates/edge0-pull" }
+
+tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time", "sync", "net", "io-util", "process", "fs"] }
+axum = { version = "0.8", features = ["json"] }
+tower = "0.5"
+http = "1"
+serde = { version = "1", features = ["derive"] }
+serde_json = { version = "1", features = ["preserve_order"] }
+chrono = { version = "0.4", features = ["serde"] }
+uuid = { version = "1", features = ["v4"] }
+clap = { version = "4", features = ["derive"] }
+# system-proxy: outbound requests honor HTTPS_PROXY and friends.
+reqwest = { version = "0.12", default-features = false, features = [
+  "blocking",
+  "json",
+  "rustls-tls",
+  "stream",
+  "system-proxy",
+] }
+thiserror = "2"
+anyhow = "1"
+libc = "0.2"
+rand = "0.8"
+sha2 = "0.10"
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+futures = "0.3"
+tokio-util = { version = "0.7", features = ["io"] }
+tokio-stream = "0.1"
+libloading = "0.8"
+tokenizers = "0.21"
+minijinja = { version = "2", features = ["json", "serde"] }
+minijinja-contrib = { version = "2", features = ["pycompat"] }
+ed25519-dalek = { version = "2", features = ["rand_core"] }
+base64 = "0.22"
+tauri = { version = "2", features = ["tray-icon", "image-png"] }
+tauri-build = { version = "2" }
+ts-rs = { version = "12", features = ["chrono-impl", "serde-json-impl"] }
+arboard = { version = "3", default-features = false }
diff --git a/macos/LICENSE b/macos/LICENSE
new file mode 100644
index 0000000..f3c4155
--- /dev/null
+++ b/macos/LICENSE
@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Edge0
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
diff --git a/macos/Makefile b/macos/Makefile
new file mode 100644
index 0000000..6d3861f
--- /dev/null
+++ b/macos/Makefile
@@ -0,0 +1,38 @@
+export PATH := $(HOME)/.local/bin:$(PATH)
+export CARGO_TARGET_DIR := $(abspath target)
+
+ENGINE_BUILD := work/engine-build
+ENGINE_OUT   := target/engine
+ENGINE_VERSION ?= 0.1.0
+
+.PHONY: help setup engine app check-prereqs
+
+help:
+	@echo "  make setup           install JS deps and check the toolchain"
+	@echo "  make engine          build the native Metal engine"
+	@echo "  make app             build, bundle, and write dist/edge0-*-arm64.app"
+
+check-prereqs:
+	bash scripts/check-prereqs.sh
+
+setup: check-prereqs
+	cd app && yarn install --frozen-lockfile
+
+engine:
+	cmake -S crates/engine-native -B $(ENGINE_BUILD) -DCMAKE_BUILD_TYPE=Release -DE0N_VERSION=$(ENGINE_VERSION)
+	cmake --build $(ENGINE_BUILD) --parallel 8
+	@mkdir -p $(ENGINE_OUT)
+	rm -f $(ENGINE_OUT)/libedge0_engine_native.dylib $(ENGINE_OUT)/libmlx.dylib
+	cp -f $(ENGINE_BUILD)/out/libedge0_engine_native.dylib $(ENGINE_OUT)/
+	cp -f $(ENGINE_BUILD)/out/libmlx.dylib $(ENGINE_OUT)/ 2>/dev/null || \
+	  cp -f $(ENGINE_BUILD)/mlx/libmlx.dylib $(ENGINE_OUT)/
+	cp -f $(ENGINE_BUILD)/out/mlx.metallib $(ENGINE_OUT)/ 2>/dev/null || \
+	  cp -f $(ENGINE_BUILD)/mlx/mlx/backend/metal/kernels/mlx.metallib $(ENGINE_OUT)/
+	codesign --force --sign - $(ENGINE_OUT)/libmlx.dylib $(ENGINE_OUT)/mlx.metallib \
+	  $(ENGINE_OUT)/libedge0_engine_native.dylib
+	@echo "engine: $(ENGINE_OUT)/libedge0_engine_native.dylib"
+
+app: setup
+	cargo build --release -p edge0-cli -p edge0d -p edge0-engine
+	$(MAKE) ENGINE_VERSION=$$(awk '$$0 == "[workspace.package]" { in_pkg=1; next } /^\[/ { in_pkg=0 } in_pkg && $$1 == "version" { gsub(/"/, "", $$3); print $$3; exit }' Cargo.toml) engine
+	bash scripts/package-macos.sh
diff --git a/macos/NOTICE b/macos/NOTICE
new file mode 100644
index 0000000..4e92d82
--- /dev/null
+++ b/macos/NOTICE
@@ -0,0 +1,23 @@
+edge0
+Copyright 2026 Edge0
+
+This product is licensed under the Apache License, Version 2.0.
+See the LICENSE file.
+
+Third-party notices
+-------------------
+
+mlx (third_party/mlx)
+  Copyright 2023 Apple Inc.
+  Licensed under the MIT License. See third_party/mlx/LICENSE.
+
+Models
+------
+
+Edge0-8B-A1B-preview and Edge0-35B-A3B-preview are published on Hugging Face
+under Apache-2.0:
+
+  https://huggingface.co/Edge0/Edge0-8B-A1B-preview
+  https://huggingface.co/Edge0/Edge0-35B-A3B-preview
+
+The upstream inference framework is https://github.com/Edge0-AI/edge0 (Apache-2.0).
diff --git a/macos/README.md b/macos/README.md
new file mode 100644
index 0000000..e2857e0
--- /dev/null
+++ b/macos/README.md
@@ -0,0 +1,99 @@
+# edge0
+
+English | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+On-device inference for Apple Silicon. The app talks to a local OpenAI-compatible server that runs Edge0 MoE models (8B and 35B) with streaming expert weights, SSD offload, and Metal.
+
+Requires **macOS 14+** on **Apple Silicon (M3 or later)**. All commands below are from this directory's root. MLX v0.30.6 is already under `third_party/mlx`.
+
+## Quick Start
+
+### Prerequisites
+
+- Xcode Command Line Tools (Metal compiler)
+- CMake 3.24+
+- Rust 1.88 (`rustup` will pick the version from `rust-toolchain.toml`)
+- Node 22 or later (yarn 1.22 is used if already installed; otherwise `check-prereqs` enables it via Corepack)
+
+`make app` runs this check. To verify first:
+
+```bash
+bash scripts/check-prereqs.sh
+```
+
+### Build and package
+
+```bash
+make app
+open dist/edge0-0.1.0-arm64.app
+```
+
+`make app` installs JavaScript dependencies, builds `edge0`, `edge0d`, and `edge0-engine`, compiles the native engine, then writes `dist/edge0-<version>-arm64.app`. It does not start the app or the local service. The `.app` is ad-hoc signed.
+
+### First launch
+
+1. Open the `.app`. The local service (`edge0d` on `127.0.0.1:8000`) starts with it.
+2. On the Models page, download a model from Hugging Face (stored under `~/.edge0`), load it, then chat.
+3. Chat: Enter sends. Replies stream as Markdown, including `$inline$` and `$$block$$` math.
+4. HTTP API (loopback needs no token):
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "edge0-8b",
+    "messages": [{"role": "user", "content": "Say hello in one sentence."}],
+    "stream": false
+  }'
+```
+
+Non-loopback access requires `Authorization: Bearer <token>` from the Service page.
+
+### Development
+
+To iterate without packaging:
+
+```bash
+make setup
+cargo build -p edge0d -p edge0-engine
+make engine
+cd app && yarn tauri dev
+```
+
+This starts a development window that can launch the local service. `make engine` builds the native MLX/Metal library.
+
+### App usage
+
+| Screen | What it does |
+| --- | --- |
+| Chat | Streaming chat, preset prompts, generation style (Precise / Balanced / Creative), Markdown + math |
+| Models | Download, load/unload, keep-alive |
+| Service | Start/stop the daemon, LAN bind, LaunchAgent, API token, environment check, request log |
+| Settings | Theme, font size, optional system prompt |
+| Menu bar | Status, Open window, Restart service, Quit (stops the daemon this app started) |
+
+Data lives in `~/.edge0` (models, session database, logs). The CLI is `Contents/Resources/bin/edge0` inside the `.app`. `edge0 uninstall` stops the service and any LaunchAgent; add `--purge` to delete user data after confirmation.
+
+## Performance
+
+Measured on **MacBook Air, Apple M3**.
+
+| Model | Decode | Prefill | Resident memory |
+| --- | --- | --- | --- |
+| Edge0-35B-A3B | 10–12 tok/s | ~70–130 tok/s | 4 GB |
+| Edge0-8B-A1B | 18–20 tok/s | ~300–570 tok/s | 1.7 GB |
+
+## Technical Details
+
+```
+edge0.app (Tauri 2 + React 19)
+    HTTP 127.0.0.1:8000  →  edge0d (OpenAI-compatible + /v1/edge0/*)
+    Unix socket frames   →  edge0-engine
+    dlopen               →  libedge0_engine_native.dylib + libmlx.dylib
+```
+
+- **UI:** Tauri 2, React 19, TanStack Router, Tailwind, Vercel AI SDK, streamdown + `@streamdown/math` (KaTeX).
+- **Daemon:** Rust. `/v1/chat/completions` (SSE). Downloads from Hugging Face by default.
+- **Engine:** C++/Metal on mlx 0.30.6. Int4 streaming experts (mmap slices + LRU hot stack + SSD offload), prerouter prefetch, unmerged LoRA, KV prefix cache. NAX GEMM is off by default (`MLX_METAL_NO_NAX`) because mlx 0.30.6 is numerically wrong for edge0-8b on M5.
+
+`EDGE0_HOME` defaults to `~/.edge0`.
diff --git a/macos/README_es.md b/macos/README_es.md
new file mode 100644
index 0000000..c72895e
--- /dev/null
+++ b/macos/README_es.md
@@ -0,0 +1,99 @@
+# edge0
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | Español | [Français](README_fr.md)
+
+Inferencia en el dispositivo para Apple Silicon. La aplicación se comunica con un servidor local compatible con OpenAI que ejecuta modelos MoE de Edge0 (8B y 35B) con pesos de expertos en streaming, offload a SSD y Metal.
+
+Requiere **macOS 14+** en **Apple Silicon (M3 o posterior)**. Todos los comandos siguientes se ejecutan desde la raíz de este directorio. MLX v0.30.6 ya está en `third_party/mlx`.
+
+## Inicio rápido
+
+### Requisitos previos
+
+- Xcode Command Line Tools (compilador de Metal)
+- CMake 3.24+
+- Rust 1.88 (`rustup` elegirá la versión desde `rust-toolchain.toml`)
+- Node 22 o posterior (se usa yarn 1.22 si ya está instalado; en caso contrario `check-prereqs` lo habilita mediante Corepack)
+
+`make app` ejecuta esta verificación. Para comprobarlo antes:
+
+```bash
+bash scripts/check-prereqs.sh
+```
+
+### Compilación y empaquetado
+
+```bash
+make app
+open dist/edge0-0.1.0-arm64.app
+```
+
+`make app` instala las dependencias de JavaScript, compila `edge0`, `edge0d` y `edge0-engine`, compila el motor nativo y después genera `dist/edge0-<version>-arm64.app`. No inicia la aplicación ni el servicio local. El `.app` está firmado de forma ad-hoc.
+
+### Primer inicio
+
+1. Abre el `.app`. El servicio local (`edge0d` en `127.0.0.1:8000`) se inicia con él.
+2. En la página Modelos, descarga un modelo desde Hugging Face (se guarda en `~/.edge0`), cárgalo y chatea.
+3. Chat: Enter envía. Las respuestas se transmiten en Markdown, incluidas las fórmulas matemáticas `$inline$` y `$$block$$`.
+4. API HTTP (loopback no necesita token):
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "edge0-8b",
+    "messages": [{"role": "user", "content": "Say hello in one sentence."}],
+    "stream": false
+  }'
+```
+
+El acceso fuera de loopback requiere `Authorization: Bearer <token>` desde la página Servicio.
+
+### Desarrollo
+
+Para iterar sin empaquetar:
+
+```bash
+make setup
+cargo build -p edge0d -p edge0-engine
+make engine
+cd app && yarn tauri dev
+```
+
+Esto inicia una ventana de desarrollo que puede lanzar el servicio local. `make engine` compila la biblioteca nativa de MLX/Metal.
+
+### Uso de la aplicación
+
+| Pantalla | Qué hace |
+| --- | --- |
+| Chat | Chat en streaming, prompts predefinidos, estilo de generación (Preciso / Equilibrado / Creativo), Markdown + fórmulas matemáticas |
+| Modelos | Descargar, cargar/liberar, keep-alive |
+| Servicio | Iniciar/detener el demonio, bind en LAN, LaunchAgent, token de API, verificación del entorno, registro de peticiones |
+| Configuración | Tema, tamaño de fuente, system prompt opcional |
+| Barra de menús | Estado, Abrir ventana, Reiniciar servicio, Salir (detiene el demonio que inició esta aplicación) |
+
+Los datos viven en `~/.edge0` (modelos, base de datos de sesiones, logs). La CLI es `Contents/Resources/bin/edge0` dentro del `.app`. `edge0 uninstall` detiene el servicio y cualquier LaunchAgent; añade `--purge` para eliminar los datos de usuario tras confirmación.
+
+## Rendimiento
+
+Medido en **MacBook Air, Apple M3**.
+
+| Modelo | Decode | Prefill | Memoria residente |
+| --- | --- | --- | --- |
+| Edge0-35B-A3B | 10–12 tok/s | ~70–130 tok/s | 4 GB |
+| Edge0-8B-A1B | 18–20 tok/s | ~300–570 tok/s | 1.7 GB |
+
+## Detalles técnicos
+
+```
+edge0.app (Tauri 2 + React 19)
+    HTTP 127.0.0.1:8000  →  edge0d (OpenAI-compatible + /v1/edge0/*)
+    Unix socket frames   →  edge0-engine
+    dlopen               →  libedge0_engine_native.dylib + libmlx.dylib
+```
+
+- **UI:** Tauri 2, React 19, TanStack Router, Tailwind, Vercel AI SDK, streamdown + `@streamdown/math` (KaTeX).
+- **Demonio:** Rust. `/v1/chat/completions` (SSE). Descarga desde Hugging Face por defecto.
+- **Motor:** C++/Metal sobre mlx 0.30.6. Expertos en streaming int4 (slices mmap + pila caliente LRU + SSD offload), prefetch del prerouter, LoRA sin fusionar, caché de prefijos KV. NAX GEMM está desactivado por defecto (`MLX_METAL_NO_NAX`) porque mlx 0.30.6 es numéricamente incorrecto para edge0-8b en M5.
+
+`EDGE0_HOME` tiene por defecto `~/.edge0`.
diff --git a/macos/README_fr.md b/macos/README_fr.md
new file mode 100644
index 0000000..d752915
--- /dev/null
+++ b/macos/README_fr.md
@@ -0,0 +1,99 @@
+# edge0
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | Français
+
+Inférence embarquée pour Apple Silicon. L'application communique avec un serveur local compatible OpenAI qui exécute les modèles MoE Edge0 (8B et 35B) avec des poids d'experts en streaming, SSD offload et Metal.
+
+Nécessite **macOS 14+** sur **Apple Silicon (M3 ou ultérieur)**. Toutes les commandes ci-dessous s'exécutent depuis la racine de ce répertoire. MLX v0.30.6 est déjà présent sous `third_party/mlx`.
+
+## Démarrage rapide
+
+### Prérequis
+
+- Xcode Command Line Tools (compilateur Metal)
+- CMake 3.24+
+- Rust 1.88 (`rustup` choisira la version depuis `rust-toolchain.toml`)
+- Node 22 ou ultérieur (yarn 1.22 est utilisé s'il est déjà installé ; sinon `check-prereqs` l'active via Corepack)
+
+`make app` exécute cette vérification. Pour vérifier d'abord :
+
+```bash
+bash scripts/check-prereqs.sh
+```
+
+### Compilation et packaging
+
+```bash
+make app
+open dist/edge0-0.1.0-arm64.app
+```
+
+`make app` installe les dépendances JavaScript, compile `edge0`, `edge0d` et `edge0-engine`, compile le moteur natif, puis génère `dist/edge0-<version>-arm64.app`. Il ne lance ni l'application ni le service local. Le `.app` est signé ad hoc.
+
+### Premier lancement
+
+1. Ouvrez le `.app`. Le service local (`edge0d` sur `127.0.0.1:8000`) démarre avec lui.
+2. Sur la page Models, téléchargez un modèle depuis Hugging Face (stocké sous `~/.edge0`), chargez-le, puis discutez.
+3. Chat : Entrée envoie. Les réponses arrivent en streaming au format Markdown, y compris les formules mathématiques `$inline$` et `$$block$$`.
+4. API HTTP (le loopback ne nécessite pas de token) :
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "edge0-8b",
+    "messages": [{"role": "user", "content": "Say hello in one sentence."}],
+    "stream": false
+  }'
+```
+
+L'accès hors loopback nécessite `Authorization: Bearer <token>` depuis la page Service.
+
+### Développement
+
+Pour itérer sans packaging :
+
+```bash
+make setup
+cargo build -p edge0d -p edge0-engine
+make engine
+cd app && yarn tauri dev
+```
+
+Cela lance une fenêtre de développement qui peut démarrer le service local. `make engine` compile la bibliothèque native MLX/Metal.
+
+### Utilisation de l'application
+
+| Écran | Rôle |
+| --- | --- |
+| Chat | Chat en streaming, prompts prédéfinis, style de génération (Precise / Balanced / Creative), Markdown + formules mathématiques |
+| Models | Téléchargement, chargement/déchargement, maintien actif (keep-alive) |
+| Service | Démarrer/arrêter le démon, liaison LAN, LaunchAgent, token API, vérification de l'environnement, journal des requêtes |
+| Settings | Thème, taille de police, prompt système optionnel |
+| Menu bar | Status, Open window, Restart service, Quit (arrête le démon que cette application a démarré) |
+
+Les données résident dans `~/.edge0` (modèles, base de données des sessions, journaux). La CLI se trouve dans `Contents/Resources/bin/edge0` à l'intérieur du `.app`. `edge0 uninstall` arrête le service et tout LaunchAgent ; ajoutez `--purge` pour supprimer les données utilisateur après confirmation.
+
+## Performances
+
+Mesurées sur **MacBook Air, Apple M3**.
+
+| Modèle | Decode | Prefill | Mémoire résidente |
+| --- | --- | --- | --- |
+| Edge0-35B-A3B | 10–12 tok/s | ~70–130 tok/s | 4 GB |
+| Edge0-8B-A1B | 18–20 tok/s | ~300–570 tok/s | 1.7 GB |
+
+## Détails techniques
+
+```
+edge0.app (Tauri 2 + React 19)
+    HTTP 127.0.0.1:8000  →  edge0d (OpenAI-compatible + /v1/edge0/*)
+    Unix socket frames   →  edge0-engine
+    dlopen               →  libedge0_engine_native.dylib + libmlx.dylib
+```
+
+- **UI :** Tauri 2, React 19, TanStack Router, Tailwind, Vercel AI SDK, streamdown + `@streamdown/math` (KaTeX).
+- **Démon :** Rust. `/v1/chat/completions` (SSE). Téléchargements depuis Hugging Face par défaut.
+- **Moteur :** C++/Metal sur mlx 0.30.6. Experts int4 en streaming (tranches mmap + pile LRU à chaud + SSD offload), préchargement prerouter, LoRA non fusionné, KV cache de préfixe. NAX GEMM est désactivé par défaut (`MLX_METAL_NO_NAX`) car mlx 0.30.6 est numériquement incorrect pour edge0-8b sur M5.
+
+`EDGE0_HOME` a pour valeur par défaut `~/.edge0`.
diff --git a/macos/README_ja.md b/macos/README_ja.md
new file mode 100644
index 0000000..839157c
--- /dev/null
+++ b/macos/README_ja.md
@@ -0,0 +1,99 @@
+# edge0
+
+[English](README.md) | [中文](README_zh.md) | 日本語 | [Español](README_es.md) | [Français](README_fr.md)
+
+Apple Silicon 向けオンデバイス推論。アプリはローカルの OpenAI 互換サーバーと通信し、そのサーバーが Edge0 MoE モデル(8B および 35B)を、ストリーミングされるエキスパート重み、SSD オフロード、Metal で実行します。
+
+**Apple Silicon(M3 以降)**搭載の **macOS 14+** が必要です。以下のすべてのコマンドはこのディレクトリのルートから実行します。MLX v0.30.6 はすでに `third_party/mlx` に含まれています。
+
+## クイックスタート
+
+### 前提条件
+
+- Xcode Command Line Tools(Metal コンパイラ)
+- CMake 3.24+
+- Rust 1.88(`rustup` が `rust-toolchain.toml` からバージョンを選択します)
+- Node 22 以降(すでにインストールされている場合は yarn 1.22 が使われます;そうでない場合は `check-prereqs` が Corepack 経由で有効化します)
+
+`make app` はこのチェックを実行します。事前に確認するには:
+
+```bash
+bash scripts/check-prereqs.sh
+```
+
+### ビルドとパッケージング
+
+```bash
+make app
+open dist/edge0-0.1.0-arm64.app
+```
+
+`make app` は JavaScript の依存関係をインストールし、`edge0`、`edge0d`、`edge0-engine` をビルドし、ネイティブエンジンをコンパイルしてから、`dist/edge0-<version>-arm64.app` を書き出します。アプリやローカルサービスの起動は行いません。`.app` は ad-hoc 署名されています。
+
+### 初回起動
+
+1. `.app` を開きます。ローカルサービス(`127.0.0.1:8000` の `edge0d`)も同時に起動します。
+2. Models ページで、Hugging Face からモデルをダウンロードし(`~/.edge0` に保存されます)、ロードしてからチャットします。
+3. チャット: Enter で送信します。返信は Markdown としてストリーミング表示され、`$inline$` および `$$block$$` 数式を含みます。
+4. HTTP API(ループバックにはトークン不要):
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "edge0-8b",
+    "messages": [{"role": "user", "content": "Say hello in one sentence."}],
+    "stream": false
+  }'
+```
+
+ループバック以外からのアクセスには、Service ページで取得できる `Authorization: Bearer <token>` が必要です。
+
+### 開発
+
+パッケージングせずに反復作業するには:
+
+```bash
+make setup
+cargo build -p edge0d -p edge0-engine
+make engine
+cd app && yarn tauri dev
+```
+
+ローカルサービスを起動できる開発ウィンドウが開きます。`make engine` はネイティブの MLX/Metal ライブラリをビルドします。
+
+### アプリの使い方
+
+| 画面 | 機能 |
+| --- | --- |
+| Chat | ストリーミングチャット、プリセットプロンプト、生成スタイル(Precise / Balanced / Creative)、Markdown + 数式 |
+| Models | ダウンロード、ロード/アンロード、keep-alive |
+| Service | デーモンの起動/停止、LAN バインド、LaunchAgent、API トークン、環境チェック、リクエストログ |
+| Settings | テーマ、フォントサイズ、オプションのシステムプロンプト |
+| メニューバー | ステータス、ウィンドウを開く、サービスの再起動、終了(このアプリが起動したデーモンを停止します)|
+
+データは `~/.edge0` に保存されます(モデル、セッションデータベース、ログ)。CLI は `.app` 内の `Contents/Resources/bin/edge0` です。`edge0 uninstall` はサービスと LaunchAgent を停止します。確認後にユーザーデータを削除するには `--purge` を追加します。
+
+## パフォーマンス
+
+**MacBook Air、Apple M3** で測定。
+
+| モデル | デコード | プリフィル | 常駐メモリ |
+| --- | --- | --- | --- |
+| Edge0-35B-A3B | 10–12 tok/s | ~70–130 tok/s | 4 GB |
+| Edge0-8B-A1B | 18–20 tok/s | ~300–570 tok/s | 1.7 GB |
+
+## 技術詳細
+
+```
+edge0.app (Tauri 2 + React 19)
+    HTTP 127.0.0.1:8000  →  edge0d (OpenAI-compatible + /v1/edge0/*)
+    Unix socket frames   →  edge0-engine
+    dlopen               →  libedge0_engine_native.dylib + libmlx.dylib
+```
+
+- **UI:** Tauri 2、React 19、TanStack Router、Tailwind、Vercel AI SDK、streamdown + `@streamdown/math`(KaTeX)。
+- **デーモン:** Rust。`/v1/chat/completions`(SSE)。デフォルトでは Hugging Face からダウンロードします。
+- **エンジン:** mlx 0.30.6 上の C++/Metal。Int4 ストリーミングエキスパート(mmap スライス + LRU ホットスタック + SSD オフロード)、prerouter プリフェッチ、非マージ LoRA、KV プレフィックスキャッシュ。NAX GEMM はデフォルトでオフです(`MLX_METAL_NO_NAX`)。mlx 0.30.6 が M5 上で edge0-8b に対して数値的に誤った結果を出すためです。
+
+`EDGE0_HOME` のデフォルトは `~/.edge0` です。
diff --git a/macos/README_zh.md b/macos/README_zh.md
new file mode 100644
index 0000000..c8dbb08
--- /dev/null
+++ b/macos/README_zh.md
@@ -0,0 +1,99 @@
+# edge0
+
+[English](README.md) | 中文 | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+面向 Apple Silicon 的端侧推理。App 与本地 OpenAI 兼容 server 通信，由其以流式专家权重、SSD offload 与 Metal 运行 Edge0 MoE 模型（8B 与 35B）。
+
+需要 **Apple Silicon（M3 或更新机型）**上的 **macOS 14+**。下文所有命令均在本目录根下运行。MLX v0.30.6 已置于 `third_party/mlx`。
+
+## 快速开始
+
+### 环境要求
+
+- Xcode Command Line Tools（Metal 编译器）
+- CMake 3.24+
+- Rust 1.88（`rustup` 会按 `rust-toolchain.toml` 选择版本）
+- Node 22 或更新版本（若已安装则使用 yarn 1.22；否则 `check-prereqs` 会通过 Corepack 启用）
+
+`make app` 会执行该检查。如需先行验证：
+
+```bash
+bash scripts/check-prereqs.sh
+```
+
+### 构建与打包
+
+```bash
+make app
+open dist/edge0-0.1.0-arm64.app
+```
+
+`make app` 会安装 JavaScript 依赖，构建 `edge0`、`edge0d` 与 `edge0-engine`，编译原生引擎，然后生成 `dist/edge0-<version>-arm64.app`。它不会启动 App 或本地服务。`.app` 为 ad-hoc 签名。
+
+### 首次启动
+
+1. 打开 `.app`。本地服务（监听 `127.0.0.1:8000` 的 `edge0d`）随之启动。
+2. 在 Models 页面从 Hugging Face 下载模型（存放于 `~/.edge0`），加载后即可对话。
+3. 对话：Enter 发送。回复以 Markdown 流式呈现，支持 `$inline$` 与 `$$block$$` 数学公式。
+4. HTTP API（回环访问无需 token）：
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "edge0-8b",
+    "messages": [{"role": "user", "content": "Say hello in one sentence."}],
+    "stream": false
+  }'
+```
+
+非回环访问需要携带 Service 页面提供的 `Authorization: Bearer <token>`。
+
+### 开发
+
+无需打包即可迭代开发：
+
+```bash
+make setup
+cargo build -p edge0d -p edge0-engine
+make engine
+cd app && yarn tauri dev
+```
+
+这会启动一个可拉起本地服务的开发窗口。`make engine` 构建原生 MLX/Metal 库。
+
+### App 使用
+
+| 页面 | 功能 |
+| --- | --- |
+| Chat | 流式对话、预置提示词、生成风格（Precise / Balanced / Creative）、Markdown + 数学公式 |
+| Models | 下载、加载/卸载、keep-alive |
+| Service | 启停 daemon、局域网绑定、LaunchAgent、API token、环境检查、请求日志 |
+| Settings | 主题、字号、可选系统提示词 |
+| 菜单栏 | 状态、打开窗口、重启服务、退出（会停止本 App 启动的 daemon） |
+
+数据存放于 `~/.edge0`（模型、会话数据库、日志）。CLI 位于 `.app` 内的 `Contents/Resources/bin/edge0`。`edge0 uninstall` 会停止服务与所有 LaunchAgent；加 `--purge` 可在确认后删除用户数据。
+
+## 性能实测
+
+在 **MacBook Air（Apple M3）**上实测。
+
+| 模型 | 解码 | Prefill | 常驻内存 |
+| --- | --- | --- | --- |
+| Edge0-35B-A3B | 10–12 tok/s | ~70–130 tok/s | 4 GB |
+| Edge0-8B-A1B | 18–20 tok/s | ~300–570 tok/s | 1.7 GB |
+
+## 技术细节
+
+```
+edge0.app (Tauri 2 + React 19)
+    HTTP 127.0.0.1:8000  →  edge0d (OpenAI-compatible + /v1/edge0/*)
+    Unix socket frames   →  edge0-engine
+    dlopen               →  libedge0_engine_native.dylib + libmlx.dylib
+```
+
+- **UI：**Tauri 2、React 19、TanStack Router、Tailwind、Vercel AI SDK、streamdown + `@streamdown/math`（KaTeX）。
+- **Daemon：**Rust。`/v1/chat/completions`（SSE）。默认从 Hugging Face 下载。
+- **引擎：**基于 mlx 0.30.6 的 C++/Metal。Int4 流式专家（mmap 切片 + LRU 热栈 + SSD offload）、prerouter 预取、不合并的 LoRA、KV 前缀缓存。NAX GEMM 默认关闭（`MLX_METAL_NO_NAX`），因为 mlx 0.30.6 在 M5 上对 edge0-8b 的数值结果有误。
+
+`EDGE0_HOME` 默认为 `~/.edge0`。
diff --git a/macos/app/index.html b/macos/app/index.html
new file mode 100644
index 0000000..b144af3
--- /dev/null
+++ b/macos/app/index.html
@@ -0,0 +1,12 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>edge0</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/main.tsx"></script>
+  </body>
+</html>
diff --git a/macos/app/package.json b/macos/app/package.json
new file mode 100644
index 0000000..a0d64a8
--- /dev/null
+++ b/macos/app/package.json
@@ -0,0 +1,44 @@
+{
+  "name": "edge0-app-ui",
+  "private": true,
+  "version": "0.1.0",
+  "type": "module",
+  "packageManager": "yarn@1.22.22",
+  "engines": {
+    "node": ">=22"
+  },
+  "scripts": {
+    "dev": "vite",
+    "build": "tsc --noEmit && vite build",
+    "typecheck": "tsc --noEmit"
+  },
+  "dependencies": {
+    "@ai-sdk/openai-compatible": "1.0.54",
+    "@ai-sdk/react": "2.0.261",
+    "@radix-ui/react-progress": "^1.1.16",
+    "@radix-ui/react-slot": "^1.3.3",
+    "@radix-ui/react-toggle-group": "^1.1.19",
+    "@streamdown/math": "^1.0.2",
+    "@tanstack/react-router": "^1.170.36",
+    "@tauri-apps/api": "^2.11.1",
+    "ai": "5.0.258",
+    "i18next": "^26.4.2",
+    "lucide-react": "^1.46.0",
+    "react": "^19.3.0",
+    "react-dom": "^19.3.0",
+    "react-i18next": "^17.0.14",
+    "streamdown": "^2.6.0",
+    "zod": "^4.1.8"
+  },
+  "devDependencies": {
+    "@tailwindcss/vite": "^4.3.3",
+    "@tauri-apps/cli": "^2.11.4",
+    "@types/node": "22",
+    "@types/react": "^19.3.0",
+    "@types/react-dom": "^19.3.0",
+    "@vitejs/plugin-react": "^6.1.1",
+    "tailwindcss": "^4.3.3",
+    "typescript": "^7.0.2",
+    "vite": "^8.3.0"
+  }
+}
diff --git a/macos/app/src-tauri/Cargo.toml b/macos/app/src-tauri/Cargo.toml
new file mode 100644
index 0000000..c9ac1dc
--- /dev/null
+++ b/macos/app/src-tauri/Cargo.toml
@@ -0,0 +1,29 @@
+[package]
+name = "edge0-app"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+[lib]
+name = "edge0_app"
+crate-type = ["rlib"]
+
+[build-dependencies]
+tauri-build.workspace = true
+
+[dependencies]
+tauri.workspace = true
+edge0-core.workspace = true
+tokio.workspace = true
+reqwest.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+chrono.workspace = true
+futures.workspace = true
+thiserror.workspace = true
+tracing.workspace = true
+libc.workspace = true
+uuid.workspace = true
+arboard.workspace = true
+rusqlite = { version = "0.40.2", features = ["bundled"] }
diff --git a/macos/app/src-tauri/Entitlements.plist b/macos/app/src-tauri/Entitlements.plist
new file mode 100644
index 0000000..9ab52e1
--- /dev/null
+++ b/macos/app/src-tauri/Entitlements.plist
@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+  <key>com.apple.security.cs.allow-jit</key>
+  <true/>
+  <key>com.apple.security.cs.allow-unsigned-executable-memory</key>
+  <true/>
+</dict>
+</plist>
diff --git a/macos/app/src-tauri/build.rs b/macos/app/src-tauri/build.rs
new file mode 100644
index 0000000..d860e1e
--- /dev/null
+++ b/macos/app/src-tauri/build.rs
@@ -0,0 +1,3 @@
+fn main() {
+    tauri_build::build()
+}
diff --git a/macos/app/src-tauri/capabilities/default.json b/macos/app/src-tauri/capabilities/default.json
new file mode 100644
index 0000000..8bc248b
--- /dev/null
+++ b/macos/app/src-tauri/capabilities/default.json
@@ -0,0 +1,7 @@
+{
+  "$schema": "../gen/schemas/desktop-schema.json",
+  "identifier": "default",
+  "description": "edge0.app main webview: custom commands and core events only; no network/HTTP plugin permissions",
+  "windows": ["main"],
+  "permissions": ["core:default"]
+}
diff --git a/macos/app/src-tauri/icons/edge0_logo.png b/macos/app/src-tauri/icons/edge0_logo.png
new file mode 100644
index 0000000..712a001
Binary files /dev/null and b/macos/app/src-tauri/icons/edge0_logo.png differ
diff --git a/macos/app/src-tauri/icons/icon.icns b/macos/app/src-tauri/icons/icon.icns
new file mode 100644
index 0000000..2b1fc68
Binary files /dev/null and b/macos/app/src-tauri/icons/icon.icns differ
diff --git a/macos/app/src-tauri/icons/icon.png b/macos/app/src-tauri/icons/icon.png
new file mode 100644
index 0000000..1c757ee
Binary files /dev/null and b/macos/app/src-tauri/icons/icon.png differ
diff --git a/macos/app/src-tauri/icons/tray-idle.png b/macos/app/src-tauri/icons/tray-idle.png
new file mode 100644
index 0000000..5f7cb03
Binary files /dev/null and b/macos/app/src-tauri/icons/tray-idle.png differ
diff --git a/macos/app/src-tauri/icons/tray-loading.png b/macos/app/src-tauri/icons/tray-loading.png
new file mode 100644
index 0000000..efd2679
Binary files /dev/null and b/macos/app/src-tauri/icons/tray-loading.png differ
diff --git a/macos/app/src-tauri/icons/tray-resident.png b/macos/app/src-tauri/icons/tray-resident.png
new file mode 100644
index 0000000..5888448
Binary files /dev/null and b/macos/app/src-tauri/icons/tray-resident.png differ
diff --git a/macos/app/src-tauri/icons/tray-throttled.png b/macos/app/src-tauri/icons/tray-throttled.png
new file mode 100644
index 0000000..8405431
Binary files /dev/null and b/macos/app/src-tauri/icons/tray-throttled.png differ
diff --git a/macos/app/src-tauri/src/bridge/client.rs b/macos/app/src-tauri/src/bridge/client.rs
new file mode 100644
index 0000000..bff1ea7
--- /dev/null
+++ b/macos/app/src-tauri/src/bridge/client.rs
@@ -0,0 +1,299 @@
+//! The shell holds every daemon connection. The webview never sees tokens or host addresses.
+
+use std::sync::RwLock;
+
+use edge0_core::paths::Home;
+use edge0_core::protocol::SystemSnapshot;
+use serde_json::Value;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct BridgeTarget {
+    pub base: String,
+    pub loopback: bool,
+}
+
+impl BridgeTarget {
+    pub fn from_base(base: &str) -> Option<Self> {
+        let u = url_like(base)?;
+        Some(Self {
+            base: base.trim_end_matches('/').to_string(),
+            loopback: {
+                let lo = std::net::Ipv4Addr::LOCALHOST.to_string();
+                u.0 == lo || u.0 == "localhost" || u.0 == "::1"
+            },
+        })
+    }
+}
+
+fn url_like(s: &str) -> Option<(String, Option<u16>)> {
+    let rest = s
+        .strip_prefix("http://")
+        .or_else(|| s.strip_prefix("https://"))?;
+    let hostport = rest.trim_end_matches('/');
+    match hostport.rsplit_once(':') {
+        Some((h, p)) => Some((h.to_string(), p.parse().ok())),
+        None => Some((hostport.to_string(), None)),
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum BridgeError {
+    #[error("Local service not found: {0}. Open edge0.app or run `edge0 serve`")]
+    NoService(String),
+    #[error("Authentication required: non-loopback targets need EDGE0_TOKEN")]
+    TokenRequired,
+    #[error("{message}")]
+    Api {
+        status: u16,
+        code: Option<String>,
+        message: String,
+    },
+    #[error("Protocol decode failed: {0}")]
+    Decode(String),
+}
+
+impl BridgeError {
+    pub fn code(&self) -> Option<&str> {
+        match self {
+            BridgeError::Api { code, .. } => code.as_deref(),
+            _ => None,
+        }
+    }
+    pub fn status(&self) -> Option<u16> {
+        match self {
+            BridgeError::Api { status, .. } => Some(*status),
+            _ => None,
+        }
+    }
+
+    pub fn to_json(&self) -> String {
+        serde_json::json!({
+            "status": self.status(),
+            "code": self.code(),
+            "message": self.to_string(),
+        })
+        .to_string()
+    }
+}
+
+pub struct Bridge {
+    http: reqwest::Client,
+    /// No overall timeout: reqwest `.timeout()` covers the whole body and would cut long chat streams.
+    stream_http: reqwest::Client,
+    probe_http: reqwest::Client,
+    target: RwLock<Option<BridgeTarget>>,
+    pinned: bool,
+    home: Home,
+    env_token: Option<String>,
+}
+
+impl Bridge {
+    pub fn new(home: Home) -> Self {
+        Self::with_override(
+            home,
+            std::env::var("EDGE0_HOST").ok(),
+            std::env::var("EDGE0_TOKEN").ok(),
+        )
+    }
+
+    pub fn with_override(home: Home, host: Option<String>, token: Option<String>) -> Self {
+        let http = reqwest::Client::builder()
+            .timeout(std::time::Duration::from_secs(30))
+            .no_proxy() // user outbound proxies must not intercept loopback
+            .build()
+            .expect("reqwest client");
+        let stream_http = reqwest::Client::builder()
+            .connect_timeout(std::time::Duration::from_secs(5))
+            .no_proxy()
+            .build()
+            .expect("reqwest stream client");
+        let (pinned, target) = match host
+            .filter(|s| !s.trim().is_empty())
+            .and_then(|s| BridgeTarget::from_base(&s))
+        {
+            Some(t) => (true, Some(t)),
+            None => (false, None),
+        };
+        let probe_http = reqwest::Client::builder()
+            .timeout(std::time::Duration::from_millis(2500))
+            .connect_timeout(std::time::Duration::from_millis(1000))
+            .no_proxy()
+            .build()
+            .expect("probe client");
+        Self {
+            http,
+            stream_http,
+            probe_http,
+            target: RwLock::new(target),
+            pinned,
+            home,
+            env_token: token.filter(|s| !s.trim().is_empty()),
+        }
+    }
+
+    pub fn set_target(&self, t: BridgeTarget) {
+        if !self.pinned {
+            *self.target.write().unwrap() = Some(t);
+        }
+    }
+
+    pub fn clear_target(&self) {
+        if !self.pinned {
+            *self.target.write().unwrap() = None;
+        }
+    }
+
+    pub fn target(&self) -> Option<BridgeTarget> {
+        self.target.read().unwrap().clone()
+    }
+
+    fn token(&self, t: &BridgeTarget) -> Option<String> {
+        // Loopback is unauthenticated by contract; attaching a token would mis-label GUI traffic as cli.
+        if !t.loopback {
+            return self.env_token.clone();
+        }
+        None
+    }
+
+    fn req(&self, method: &str, path: &str) -> Result<reqwest::RequestBuilder, BridgeError> {
+        self.req_with(&self.http, method, path)
+    }
+
+    fn req_with(
+        &self,
+        client: &reqwest::Client,
+        method: &str,
+        path: &str,
+    ) -> Result<reqwest::RequestBuilder, BridgeError> {
+        let t = self
+            .target()
+            .ok_or_else(|| BridgeError::NoService(path.to_string()))?;
+        if !t.loopback && self.env_token.is_none() {
+            return Err(BridgeError::TokenRequired);
+        }
+        let url = format!("{}{}", t.base, path);
+        let mut r = match method {
+            "GET" => client.get(&url),
+            "DELETE" => client.delete(&url),
+            _ => client.post(&url),
+        };
+        if let Some(tok) = self.token(&t) {
+            r = r.bearer_auth(tok);
+        }
+        r = r.header(
+            "x-edge0-client-version",
+            edge0_core::version::daemon_version(),
+        );
+        Ok(r)
+    }
+
+    pub async fn get_json(&self, path: &str) -> Result<Value, BridgeError> {
+        let resp = self
+            .req("GET", path)?
+            .send()
+            .await
+            .map_err(|_| BridgeError::NoService(path.to_string()))?;
+        decode(resp).await
+    }
+
+    pub async fn post_json(&self, path: &str, body: Value) -> Result<Value, BridgeError> {
+        let resp = self
+            .req("POST", path)?
+            .json(&body)
+            .send()
+            .await
+            .map_err(|_| BridgeError::NoService(path.to_string()))?;
+        decode(resp).await
+    }
+
+    pub async fn req_delete(&self, path: &str) -> Result<Value, BridgeError> {
+        let resp = self
+            .req("DELETE", path)?
+            .send()
+            .await
+            .map_err(|_| BridgeError::NoService(path.to_string()))?;
+        decode(resp).await
+    }
+
+    pub async fn system(&self) -> Result<SystemSnapshot, BridgeError> {
+        let v = self.get_json("/v1/edge0/system").await?;
+        serde_json::from_value(v).map_err(|e| BridgeError::Decode(e.to_string()))
+    }
+
+    pub async fn open_stream(
+        &self,
+        path: &str,
+        body: Value,
+    ) -> Result<reqwest::Response, BridgeError> {
+        self.req_with(&self.stream_http, "POST", path)?
+            .header("accept", "text/event-stream")
+            .json(&body)
+            .send()
+            .await
+            .map_err(|_| BridgeError::NoService(path.to_string()))
+    }
+
+    pub async fn open_events(&self, since: Option<u64>) -> Result<reqwest::Response, BridgeError> {
+        let mut r = self.req_with(&self.stream_http, "GET", "/v1/edge0/events")?;
+        if let Some(seq) = since {
+            r = r.header("last-event-id", seq.to_string());
+        }
+        r.header("accept", "text/event-stream")
+            .send()
+            .await
+            .map_err(|_| BridgeError::NoService("/v1/edge0/events".to_string()))
+    }
+
+    pub async fn health(&self, base: &str) -> Option<serde_json::Value> {
+        let mut r = self.probe_http.get(format!("{base}/health"));
+        if let Some(t) = BridgeTarget::from_base(base).filter(|t| !t.loopback) {
+            if let Some(tok) = self.token(&t) {
+                r = r.bearer_auth(tok);
+            }
+        }
+        let resp = r.send().await.ok()?;
+        if !resp.status().is_success() {
+            return None;
+        }
+        resp.json().await.ok()
+    }
+
+    pub async fn health_ok(&self, base: &str) -> bool {
+        self.health(base).await.is_some()
+    }
+
+    pub fn http(&self) -> &reqwest::Client {
+        &self.http
+    }
+
+    pub fn home(&self) -> &Home {
+        &self.home
+    }
+}
+
+async fn decode(resp: reqwest::Response) -> Result<Value, BridgeError> {
+    let status = resp.status();
+    let v: Value = resp.json().await.unwrap_or(Value::Null);
+    if status.is_success() {
+        return Ok(v);
+    }
+    let err = v.get("error");
+    let code = err.and_then(|e| e.get("code")).and_then(Value::as_str);
+    let message = err
+        .and_then(|e| e.get("message"))
+        .and_then(Value::as_str)
+        .map(str::to_string)
+        .unwrap_or_else(|| format!("HTTP {status}"));
+    if let Some(id) = code {
+        if edge0_core::errorcodes::lookup(id).is_none() {
+            return Err(BridgeError::Decode(format!(
+                "Unknown error code {id} (not in error-codes.md)"
+            )));
+        }
+    }
+    Err(BridgeError::Api {
+        status: status.as_u16(),
+        code: code.map(str::to_string),
+        message,
+    })
+}
diff --git a/macos/app/src-tauri/src/bridge/events.rs b/macos/app/src-tauri/src/bridge/events.rs
new file mode 100644
index 0000000..1ef13d3
--- /dev/null
+++ b/macos/app/src-tauri/src/bridge/events.rs
@@ -0,0 +1,281 @@
+//! Status-event driver. UI rebuilds from one snapshot plus later deltas; foldable types coalesce by subject.
+
+use std::collections::{HashMap, VecDeque};
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+use futures::StreamExt;
+
+use edge0_core::protocol::{foldable, known_type, Event, HelloPayload, SystemSnapshot};
+use serde::Serialize;
+use tokio::sync::watch;
+
+use super::client::Bridge;
+
+#[derive(Debug, Clone, Serialize)]
+#[serde(tag = "kind", rename_all = "camelCase")]
+pub enum EventFrame {
+    Snapshot { snapshot: Box<SystemSnapshot> },
+    Event { event: Event },
+    Status { state: LinkState },
+}
+
+#[derive(Debug, Clone, Copy, Serialize)]
+#[serde(rename_all = "lowercase")]
+pub enum LinkState {
+    Live,
+    Reconnecting,
+    Resyncing,
+}
+
+pub struct EventDriver {
+    bridge: Arc<Bridge>,
+    stop: watch::Receiver<bool>,
+    /// Pause before stop/restart: daemon graceful shutdown waits on open SSE, so this connection must drop first.
+    pause_tx: watch::Sender<bool>,
+    pause: watch::Receiver<bool>,
+    kick_tx: watch::Sender<()>,
+    /// Keep a Receiver; a watch channel with no subscriber silently drops send().
+    kick_rx: watch::Receiver<()>,
+    cursor: AtomicU64,
+    session: Mutex<Option<String>>,
+}
+
+impl EventDriver {
+    pub fn new(bridge: Arc<Bridge>) -> (Arc<Self>, watch::Sender<bool>) {
+        let (stop_tx, stop_rx) = watch::channel(false);
+        let (pause_tx, pause_rx) = watch::channel(false);
+        let (kick_tx, kick_rx) = watch::channel(());
+        let drv = Arc::new(Self {
+            bridge,
+            stop: stop_rx,
+            pause_tx,
+            pause: pause_rx,
+            kick_tx,
+            kick_rx,
+            cursor: AtomicU64::new(0),
+            session: Mutex::new(None),
+        });
+        (drv, stop_tx)
+    }
+
+    pub fn notify(&self) {
+        let _ = self.kick_tx.send(());
+    }
+
+    pub fn set_paused(&self, paused: bool) {
+        let _ = self.pause_tx.send(paused);
+        if !paused {
+            self.notify();
+        }
+    }
+
+    pub async fn run(self: Arc<Self>, mut sink: impl FnMut(EventFrame) -> bool + Send + 'static) {
+        let mut backoff = Duration::from_millis(250);
+        let cap = Duration::from_secs(5);
+        loop {
+            let mut stop_now = self.stop.clone();
+            if *stop_now.borrow_and_update() {
+                return;
+            }
+            if *self.pause.borrow() && !self.wait_resume().await {
+                return;
+            }
+            match self.bridge.system().await {
+                Ok(snap) => {
+                    self.cursor.store(snap.events_seq, Ordering::SeqCst);
+                    if !sink(EventFrame::Snapshot {
+                        snapshot: Box::new(snap),
+                    }) {
+                        return;
+                    }
+                }
+                Err(_) => {
+                    if !sink(EventFrame::Status {
+                        state: LinkState::Reconnecting,
+                    }) {
+                        return;
+                    }
+                    if !self.wait_tick(backoff).await {
+                        return;
+                    }
+                    backoff = (backoff * 2).min(cap);
+                    continue;
+                }
+            }
+            backoff = Duration::from_millis(250);
+            match self.pump_stream(&mut sink).await {
+                PumpOut::Resynced | PumpOut::Paused => continue,
+                PumpOut::Gone => return,
+                PumpOut::Retry => {
+                    if !sink(EventFrame::Status {
+                        state: LinkState::Reconnecting,
+                    }) {
+                        return;
+                    }
+                    if !self.wait_tick(backoff).await {
+                        return;
+                    }
+                    backoff = (backoff * 2).min(cap);
+                }
+            }
+        }
+    }
+
+    async fn wait_resume(&self) -> bool {
+        let mut pause = self.pause.clone();
+        let _ = pause.wait_for(|p| !*p).await;
+        !*self.stop.clone().borrow_and_update()
+    }
+
+    async fn wait_tick(&self, d: Duration) -> bool {
+        let mut kick = self.kick_rx.clone();
+        let _ = kick.borrow_and_update();
+        let mut stop = self.stop.clone();
+        let _ = stop.borrow_and_update();
+        tokio::select! {
+            _ = tokio::time::sleep(d) => true,
+            r = kick.changed() => r.is_ok(),
+            r = stop.wait_for(|s| *s) => r.is_err(),
+        }
+    }
+
+    async fn pump_stream(
+        self: &Arc<Self>,
+        sink: &mut (impl FnMut(EventFrame) -> bool + Send),
+    ) -> PumpOut {
+        let since = self.cursor.load(Ordering::SeqCst);
+        let resp = match self.bridge.open_events(Some(since)).await {
+            Ok(r) => r,
+            Err(_) => return PumpOut::Retry,
+        };
+        if !resp.status().is_success() {
+            return PumpOut::Retry;
+        }
+        if !sink(EventFrame::Status {
+            state: LinkState::Live,
+        }) {
+            return PumpOut::Gone;
+        }
+        let mut stream = resp.bytes_stream();
+        let mut parser = super::sse::SseParser::new();
+        let mut pending: VecDeque<Event> = VecDeque::new();
+        let mut fold_idx: HashMap<(String, String), usize> = HashMap::new();
+        let mut kicked = self.kick_rx.clone();
+        let _ = kicked.borrow_and_update();
+        let mut stop = self.stop.clone();
+        let _ = stop.borrow_and_update();
+        let mut pause = self.pause.clone();
+        let _ = pause.borrow_and_update();
+
+        loop {
+            tokio::select! {
+                chunk = stream.next() => {
+                    match chunk {
+                        None => return PumpOut::Retry,
+                        Some(Err(_)) => return PumpOut::Retry,
+                        Some(Ok(bytes)) => {
+                            for block in parser.push_str(&String::from_utf8_lossy(&bytes)) {
+                                let Some(ev) = decode_event(&block.data) else {
+                                    continue;
+                                };
+                                self.cursor.fetch_max(ev.seq, Ordering::SeqCst);
+                                match ev.etype.as_str() {
+                                    "daemon.hello" => {
+                                        let session = hello_session(&ev);
+                                        let changed = {
+                                            let mut cur = self.session.lock().unwrap();
+                                            let changed = matches!(
+                                                (cur.as_ref(), session.as_ref()),
+                                                (Some(prev), Some(s)) if prev != s
+                                            );
+                                            if session.is_some() {
+                                                *cur = session;
+                                            }
+                                            changed
+                                        };
+                                        if changed {
+                                            return PumpOut::Resynced;
+                                        }
+                                        fold(&mut pending, &mut fold_idx, ev);
+                                    }
+                                    "resync.required" => {
+                                        if !sink(EventFrame::Status { state: LinkState::Resyncing }) {
+                                            return PumpOut::Gone;
+                                        }
+                                        return PumpOut::Resynced;
+                                    }
+                                    _ => fold(&mut pending, &mut fold_idx, ev),
+                                }
+                            }
+                            if !flush(sink, &mut pending, &mut fold_idx) {
+                                return PumpOut::Gone;
+                            }
+                        }
+                    }
+                }
+                _ = tokio::time::sleep(Duration::from_millis(50)) => {
+                    if !flush(sink, &mut pending, &mut fold_idx) {
+                        return PumpOut::Gone;
+                    }
+                }
+                Ok(()) = kicked.changed() => {
+                    return PumpOut::Resynced;
+                }
+                r = stop.wait_for(|s| *s) => { if r.is_ok() { return PumpOut::Gone; } },
+                Ok(_) = pause.wait_for(|p| *p) => { return PumpOut::Paused; }
+            }
+        }
+    }
+}
+
+enum PumpOut {
+    Resynced,
+    Retry,
+    Gone,
+    Paused,
+}
+
+fn decode_event(data: &str) -> Option<Event> {
+    let ev: Event = serde_json::from_str(data).ok()?;
+    known_type(&ev.etype).then_some(ev)
+}
+
+fn hello_session(ev: &Event) -> Option<String> {
+    serde_json::from_value::<HelloPayload>(ev.payload.clone())
+        .ok()
+        .map(|p| p.session_id)
+}
+
+fn fold(pending: &mut VecDeque<Event>, idx: &mut HashMap<(String, String), usize>, ev: Event) {
+    if foldable(&ev.etype) {
+        let key = (ev.etype.clone(), ev.subject.clone());
+        if let Some(&slot) = idx.get(&key) {
+            pending[slot] = ev;
+            return;
+        }
+        idx.insert(key, pending.len());
+        pending.push_back(ev);
+    } else {
+        pending.push_back(ev);
+    }
+}
+
+fn flush(
+    sink: &mut (impl FnMut(EventFrame) -> bool + Send),
+    pending: &mut VecDeque<Event>,
+    idx: &mut HashMap<(String, String), usize>,
+) -> bool {
+    while let Some(ev) = pending.pop_front() {
+        idx.retain(|_, slot| *slot > 0);
+        for v in idx.values_mut() {
+            *v -= 1;
+        }
+        if !sink(EventFrame::Event { event: ev }) {
+            return false;
+        }
+    }
+    idx.clear();
+    true
+}
diff --git a/macos/app/src-tauri/src/bridge/mod.rs b/macos/app/src-tauri/src/bridge/mod.rs
new file mode 100644
index 0000000..1ae5269
--- /dev/null
+++ b/macos/app/src-tauri/src/bridge/mod.rs
@@ -0,0 +1,8 @@
+//! Daemon connection layer: token injection, target resolution, decode, streaming, and event folding.
+
+pub mod client;
+pub mod events;
+pub mod sse;
+pub mod stream;
+
+pub use client::{Bridge, BridgeError, BridgeTarget};
diff --git a/macos/app/src-tauri/src/bridge/sse.rs b/macos/app/src-tauri/src/bridge/sse.rs
new file mode 100644
index 0000000..d8af9ba
--- /dev/null
+++ b/macos/app/src-tauri/src/bridge/sse.rs
@@ -0,0 +1,71 @@
+//! Minimal SSE splitter: frame boundaries only; callers interpret the payload.
+
+#[derive(Debug, Clone)]
+pub struct SseBlock {
+    pub id: Option<String>,
+    pub event: Option<String>,
+    pub data: String,
+    pub raw: String,
+}
+
+#[derive(Default)]
+pub struct SseParser {
+    buf: String,
+}
+
+impl SseParser {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn push_str(&mut self, s: &str) -> Vec<SseBlock> {
+        // Normalize line endings; this daemon emits `\n` only.
+        if s.contains('\r') {
+            self.buf
+                .push_str(&s.replace("\r\n", "\n").replace('\r', "\n"));
+        } else {
+            self.buf.push_str(s);
+        }
+        let mut out = Vec::new();
+        while let Some(end) = self.buf.find("\n\n") {
+            let block = self.buf[..end].to_string();
+            self.buf.drain(..end + 2);
+            if let Some(b) = parse_block(&block) {
+                out.push(b);
+            }
+        }
+        out
+    }
+
+    pub fn finish(&mut self) -> Option<SseBlock> {
+        if self.buf.trim().is_empty() {
+            return None;
+        }
+        let block = std::mem::take(&mut self.buf);
+        parse_block(&block)
+    }
+}
+
+fn parse_block(block: &str) -> Option<SseBlock> {
+    let mut id = None;
+    let mut event = None;
+    let mut data: Vec<String> = Vec::new();
+    for line in block.lines() {
+        if let Some(v) = line.strip_prefix("id:") {
+            id = Some(v.trim().to_string());
+        } else if let Some(v) = line.strip_prefix("event:") {
+            event = Some(v.trim().to_string());
+        } else if let Some(v) = line.strip_prefix("data:") {
+            data.push(v.strip_prefix(' ').unwrap_or(v).to_string());
+        }
+    }
+    if data.is_empty() && id.is_none() && event.is_none() {
+        return None;
+    }
+    Some(SseBlock {
+        id,
+        event,
+        data: data.join("\n"),
+        raw: format!("{block}\n\n"),
+    })
+}
diff --git a/macos/app/src-tauri/src/bridge/stream.rs b/macos/app/src-tauri/src/bridge/stream.rs
new file mode 100644
index 0000000..b295910
--- /dev/null
+++ b/macos/app/src-tauri/src/bridge/stream.rs
@@ -0,0 +1,106 @@
+//! Chat SSE forwarder. Cancel drops the TCP stream so the daemon observes `e0_cancel`.
+
+use futures::StreamExt;
+use serde::Serialize;
+use serde_json::Value;
+
+use super::client::{Bridge, BridgeError};
+use super::sse::SseParser;
+
+#[derive(Debug, Clone, Serialize)]
+#[serde(tag = "kind", rename_all = "camelCase")]
+pub enum StreamFrame {
+    Chunk {
+        data: String,
+    },
+    Done,
+    Error {
+        code: Option<String>,
+        message: String,
+    },
+}
+
+pub async fn run(
+    bridge: &Bridge,
+    body: &Value,
+    mut emit: impl FnMut(StreamFrame),
+    should_stop: impl FnMut() -> bool,
+) {
+    match bridge
+        .open_stream("/v1/chat/completions", body.clone())
+        .await
+    {
+        Ok(resp) => pump(resp, emit, should_stop).await,
+        Err(e) => {
+            emit(error_frame(&e));
+            emit(StreamFrame::Done);
+        }
+    }
+}
+
+pub fn error_frame(e: &BridgeError) -> StreamFrame {
+    StreamFrame::Error {
+        code: e.code().map(str::to_string),
+        message: e.to_string(),
+    }
+}
+
+pub async fn pump(
+    resp: reqwest::Response,
+    mut emit: impl FnMut(StreamFrame),
+    mut should_stop: impl FnMut() -> bool,
+) {
+    let status = resp.status();
+    if !status.is_success() {
+        let v: Value = resp.json().await.unwrap_or(Value::Null);
+        let err = v.get("error");
+        emit(StreamFrame::Error {
+            code: err
+                .and_then(|e| e.get("code"))
+                .and_then(Value::as_str)
+                .map(str::to_string),
+            message: err
+                .and_then(|e| e.get("message"))
+                .and_then(Value::as_str)
+                .unwrap_or("daemon returned an error")
+                .to_string(),
+        });
+        emit(StreamFrame::Done);
+        return;
+    }
+    let mut stream = resp.bytes_stream();
+    let mut parser = SseParser::new();
+    while let Some(chunk) = stream.next().await {
+        if should_stop() {
+            return;
+        }
+        match chunk {
+            Ok(bytes) => {
+                for block in parser.push_str(&String::from_utf8_lossy(&bytes)) {
+                    if should_stop() {
+                        return;
+                    }
+                    if block.data == "[DONE]" {
+                        emit(StreamFrame::Done);
+                        return;
+                    }
+                    emit(StreamFrame::Chunk { data: block.data });
+                }
+            }
+            Err(_) => {
+                emit(StreamFrame::Error {
+                    code: None,
+                    message: "stream interrupted (network or daemon error)".to_string(),
+                });
+                emit(StreamFrame::Done);
+                return;
+            }
+        }
+    }
+    if let Some(block) = parser.finish() {
+        if !block.data.is_empty() && block.data != "[DONE]" {
+            emit(StreamFrame::Chunk { data: block.data });
+        }
+    }
+    emit(StreamFrame::Done);
+}
diff --git a/macos/app/src-tauri/src/lib.rs b/macos/app/src-tauri/src/lib.rs
new file mode 100644
index 0000000..97e7d94
--- /dev/null
+++ b/macos/app/src-tauri/src/lib.rs
@@ -0,0 +1,530 @@
+//! Tauri 2 shell: no listen port, no engine worker, no blob access. All daemon I/O goes through `bridge`.
+
+pub mod bridge;
+pub mod lifecycle;
+pub mod service;
+pub mod store;
+mod tray;
+
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+use serde_json::Value;
+use tauri::ipc::Channel;
+use tauri::{Manager, State};
+
+use bridge::events::EventDriver;
+use bridge::stream::{self, StreamFrame};
+use bridge::Bridge;
+use lifecycle::{Effect, Lifecycle, Phase};
+use store::{MessageRow, Store, ThreadRow};
+use tokio::sync::watch;
+
+type TrayStatusItem = tauri::menu::MenuItem<tauri::Wry>;
+
+type Streams = Arc<Mutex<HashMap<String, tauri::async_runtime::JoinHandle<()>>>>;
+
+#[derive(Default)]
+pub struct EventSink {
+    chan: Mutex<Option<Channel<String>>>,
+}
+
+pub struct AppState {
+    bridge: Arc<Bridge>,
+    streams: Streams,
+    sink: Arc<EventSink>,
+    driver: Arc<EventDriver>,
+    driver_stop: watch::Sender<bool>,
+    lifecycle: Arc<tokio::sync::Mutex<Lifecycle>>,
+    /// Acquire svc_gate before lifecycle; never nest the other way.
+    svc_gate: Arc<tokio::sync::Mutex<()>>,
+    view: Arc<std::sync::Mutex<tray::TrayView>>,
+    tray_status: std::sync::Mutex<Option<TrayStatusItem>>,
+    store: std::sync::Mutex<Store>,
+    reconciled: usize,
+}
+
+#[tauri::command]
+fn app_version() -> &'static str {
+    env!("CARGO_PKG_VERSION")
+}
+
+fn checked_path(path: &str) -> Result<&str, String> {
+    if !path.starts_with("/v1/") || path.contains("//") || path.contains("://") {
+        return Err("path must be an absolute daemon path (/v1/ prefix; no scheme/host)".into());
+    }
+    Ok(path)
+}
+
+#[tauri::command]
+async fn api_get(state: State<'_, AppState>, path: String) -> Result<Value, String> {
+    let p = checked_path(&path)?;
+    state.bridge.get_json(p).await.map_err(|e| e.to_json())
+}
+
+#[tauri::command]
+async fn api_post(state: State<'_, AppState>, path: String, body: Value) -> Result<Value, String> {
+    let p = checked_path(&path)?;
+    state
+        .bridge
+        .post_json(p, body)
+        .await
+        .map_err(|e| e.to_json())
+}
+
+#[tauri::command]
+async fn api_delete(state: State<'_, AppState>, path: String) -> Result<Value, String> {
+    let p = checked_path(&path)?;
+    state.bridge.req_delete(p).await.map_err(|e| e.to_json())
+}
+
+#[derive(serde::Serialize)]
+struct StreamInit {
+    status: u16,
+    id: String,
+}
+
+#[tauri::command]
+async fn chat_stream(
+    app: tauri::AppHandle,
+    state: State<'_, AppState>,
+    body: Value,
+    on_frame: Channel<String>,
+) -> Result<StreamInit, String> {
+    let resp = state
+        .bridge
+        .open_stream("/v1/chat/completions", body)
+        .await
+        .map_err(|e| e.to_json())?;
+    let status = resp.status().as_u16();
+    let id = format!("st-{}", uuid::Uuid::new_v4());
+    let streams = app.state::<AppState>().streams.clone();
+    let handle = tauri::async_runtime::spawn({
+        let id = id.clone();
+        let streams = streams.clone();
+        async move {
+            let mut emit = move |f: StreamFrame| {
+                if let Ok(json) = serde_json::to_string(&f) {
+                    let _ = on_frame.send(json);
+                }
+            };
+            stream::pump(resp, &mut emit, || false).await;
+            streams.lock().unwrap().remove(&id);
+        }
+    });
+    streams.lock().unwrap().insert(id.clone(), handle);
+    Ok(StreamInit { status, id })
+}
+
+#[tauri::command]
+fn events_subscribe(state: State<'_, AppState>, on_frame: Channel<String>) {
+    *state.sink.chan.lock().unwrap() = Some(on_frame);
+    state.driver.notify();
+}
+
+#[tauri::command]
+fn events_unsubscribe(state: State<'_, AppState>) {
+    *state.sink.chan.lock().unwrap() = None;
+}
+
+#[tauri::command]
+async fn service_state(state: State<'_, AppState>) -> Result<Value, String> {
+    let l = state.lifecycle.lock().await;
+    let mut out = serde_json::json!({
+        "phase": l.phase,
+        "spawnWasDev": l.spawn_was_dev,
+        "pinned": l.pinned,
+    });
+    out["phase"]["spawnedByUs"] = Value::Bool(l.child_pid().is_some());
+    out["appVersion"] = Value::String(env!("CARGO_PKG_VERSION").into());
+    out["versionMismatch"] = Value::Bool(matches!(
+        &l.phase,
+        Phase::Adopted { health, .. } | Phase::Spawned { health } if health.version_mismatch()
+    ));
+    Ok(out)
+}
+
+#[tauri::command]
+async fn service_start(state: State<'_, AppState>) -> Result<(), String> {
+    let mut l = state.lifecycle.lock().await;
+    l.start_requested();
+    l.spawn_child()?;
+    drop(l);
+    state.driver.notify();
+    Ok(())
+}
+
+pub(crate) async fn hush_for_lifecycle(
+    app: &tauri::AppHandle,
+) -> Box<dyn FnOnce() + Send + 'static> {
+    let st = app.state::<AppState>();
+    st.driver.set_paused(true);
+    for h in st.streams.lock().unwrap().values() {
+        h.abort();
+    }
+    st.streams.lock().unwrap().clear();
+    // One scheduler tick so the pause gate can drop the SSE connection; not a retry loop.
+    tokio::time::sleep(std::time::Duration::from_millis(150)).await;
+    let app2 = app.clone();
+    Box::new(move || {
+        app2.state::<AppState>().driver.set_paused(false);
+        app2.state::<AppState>().driver.notify();
+    })
+}
+
+pub async fn do_service_restart(app: &tauri::AppHandle) -> Result<(), String> {
+    let st = app.state::<AppState>();
+    let bridge = st.bridge.clone();
+    let resume = hush_for_lifecycle(app).await;
+    let r = {
+        let mut l = st.lifecycle.lock().await;
+        l.restart_service(&bridge).await
+    };
+    resume();
+    r
+}
+
+#[tauri::command]
+async fn service_restart(app: tauri::AppHandle) -> Result<(), String> {
+    do_service_restart(&app).await
+}
+
+#[tauri::command]
+fn chat_cancel(state: State<'_, AppState>, id: String) {
+    if let Some(h) = state.streams.lock().unwrap().remove(&id) {
+        h.abort();
+    }
+}
+
+
+fn parse_enum<T: serde::de::DeserializeOwned>(label: &str, raw: &str) -> Result<T, String> {
+    serde_json::from_value::<T>(serde_json::Value::String(raw.into()))
+        .map_err(|_| format!("invalid {label}: {raw}"))
+}
+
+fn page_cursor(
+    before_created_at: Option<i64>,
+    before_id: Option<&str>,
+) -> Result<Option<(i64, &str)>, String> {
+    match (before_created_at, before_id) {
+        (Some(ca), Some(id)) => Ok(Some((ca, id))),
+        (None, None) => Ok(None),
+        _ => Err("page cursor fields must be paired (before_created_at + before_id)".into()),
+    }
+}
+
+#[tauri::command]
+async fn thread_list(state: State<'_, AppState>) -> Result<Vec<ThreadRow>, String> {
+    state.store.lock().unwrap().threads()
+}
+
+#[tauri::command]
+async fn thread_create(
+    state: State<'_, AppState>,
+    title: String,
+    model: Option<String>,
+) -> Result<ThreadRow, String> {
+    state.store.lock().unwrap().create_thread(&title, model)
+}
+
+#[tauri::command]
+async fn thread_rename(
+    state: State<'_, AppState>,
+    id: String,
+    title: String,
+) -> Result<(), String> {
+    state.store.lock().unwrap().rename_thread(&id, &title)
+}
+
+#[tauri::command]
+async fn thread_set_model(
+    state: State<'_, AppState>,
+    id: String,
+    model: Option<String>,
+) -> Result<(), String> {
+    state.store.lock().unwrap().set_thread_model(&id, model)
+}
+
+#[tauri::command]
+async fn thread_delete(state: State<'_, AppState>, id: String) -> Result<(), String> {
+    state.store.lock().unwrap().delete_thread(&id)
+}
+
+#[derive(serde::Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct StoreInfo {
+    pub reconciled: usize,
+}
+
+#[tauri::command]
+async fn store_info(state: State<'_, AppState>) -> Result<StoreInfo, String> {
+    Ok(StoreInfo {
+        reconciled: state.reconciled,
+    })
+}
+
+#[tauri::command]
+async fn messages_page(
+    state: State<'_, AppState>,
+    thread_id: String,
+    before_created_at: Option<i64>,
+    before_id: Option<String>,
+    limit: Option<u32>,
+) -> Result<Vec<MessageRow>, String> {
+    let before = page_cursor(before_created_at, before_id.as_deref())?;
+    state
+        .store
+        .lock()
+        .unwrap()
+        .messages_page(&thread_id, before, limit.unwrap_or(50).min(200))
+}
+
+#[tauri::command]
+async fn message_start(
+    state: State<'_, AppState>,
+    thread_id: String,
+    role: String,
+    content: String,
+    model: String,
+    params: Option<Value>,
+    id: Option<String>,
+) -> Result<MessageRow, String> {
+    let role = parse_enum::<store::Role>("role", &role)?;
+    state
+        .store
+        .lock()
+        .unwrap()
+        .message_start(&thread_id, role, &content, &model, params, id)
+}
+
+#[tauri::command]
+async fn messages_delete_from(
+    state: State<'_, AppState>,
+    thread_id: String,
+    from_id: String,
+) -> Result<usize, String> {
+    state
+        .store
+        .lock()
+        .unwrap()
+        .messages_delete_from(&thread_id, &from_id)
+}
+
+#[tauri::command]
+async fn message_append(
+    state: State<'_, AppState>,
+    id: String,
+    content_delta: Option<String>,
+    reasoning_delta: Option<String>,
+) -> Result<(), String> {
+    if content_delta.is_none() && reasoning_delta.is_none() {
+        return Err("empty delta is not persisted".into());
+    }
+    state.store.lock().unwrap().message_append(
+        &id,
+        content_delta.as_deref(),
+        reasoning_delta.as_deref(),
+    )
+}
+
+#[tauri::command]
+async fn message_finish(
+    state: State<'_, AppState>,
+    id: String,
+    status: String,
+    code: Option<String>,
+    usage: Option<Value>,
+) -> Result<MessageRow, String> {
+    let status = parse_enum::<store::MessageStatus>("message status", &status)?;
+    let usage: Option<store::Usage> = match usage {
+        Some(v) => Some(serde_json::from_value(v).map_err(|e| format!("invalid usage shape: {e}"))?),
+        None => None,
+    };
+    state
+        .store
+        .lock()
+        .unwrap()
+        .message_finish(&id, status, code.as_deref(), usage)
+}
+
+#[tauri::command]
+async fn message_set_content(
+    state: State<'_, AppState>,
+    id: String,
+    content: String,
+) -> Result<(), String> {
+    state
+        .store
+        .lock()
+        .unwrap()
+        .message_set_content(&id, &content)
+}
+
+#[tauri::command]
+async fn thread_truncate_and_start(
+    state: State<'_, AppState>,
+    thread_id: String,
+    from_created_at: i64,
+    role: String,
+    content: String,
+    model: String,
+    params: Option<Value>,
+) -> Result<MessageRow, String> {
+    let role = parse_enum::<store::Role>("role", &role)?;
+    state.store.lock().unwrap().thread_truncate_and_start(
+        &thread_id,
+        from_created_at,
+        role,
+        &content,
+        &model,
+        params,
+    )
+}
+
+pub fn run() {
+    let builder = tauri::Builder::default()
+        .setup(|app| {
+            let handle = app.handle().clone();
+
+            let home = edge0_core::paths::Home::from_env()
+                .map_err(|e| format!("failed to resolve EDGE0_HOME: {e}"))?;
+            let bridge = Arc::new(Bridge::new(home.clone()));
+            let (store, reconciled) = store::Store::open(&home)?;
+            if reconciled > 0 {
+                tracing::warn!(
+                    "reconciled {reconciled} leftover streaming row(s) to aborted (previous process did not exit cleanly)"
+                );
+            }
+            let (driver, driver_stop) = EventDriver::new(bridge.clone());
+            let sink = Arc::new(EventSink::default());
+            let lifecycle = Arc::new(tokio::sync::Mutex::new(Lifecycle::new(
+                edge0_core::paths::Home::from_env().map_err(|e| e.to_string())?,
+                std::env::var("EDGE0_HOST").is_ok(),
+            )));
+            app.manage(AppState {
+                bridge: bridge.clone(),
+                streams: Arc::new(Mutex::new(HashMap::new())),
+                sink: sink.clone(),
+                driver: driver.clone(),
+                driver_stop: driver_stop.clone(),
+                lifecycle: lifecycle.clone(),
+                svc_gate: Arc::new(tokio::sync::Mutex::new(())),
+                view: Arc::new(std::sync::Mutex::default()),
+                tray_status: std::sync::Mutex::new(None),
+                store: std::sync::Mutex::new(store),
+                reconciled,
+            });
+
+            let st = app.state::<AppState>();
+            tray::build(&handle, &st.tray_status).map_err(|e| format!("failed to build tray: {e}"))?;
+            if !std::env::var("EDGE0_WINDOWLESS").is_ok_and(|v| v == "1") {
+                tray::show_main_window_public(&handle);
+            }
+
+            let h_sink = handle.clone();
+            tauri::async_runtime::spawn(driver.run(move |f| {
+                {
+                    let st = h_sink.state::<AppState>();
+                    tray::apply_frame(&st.view, &f);
+                }
+                tray::refresh(&h_sink);
+                if let Ok(json) = serde_json::to_string(&f) {
+                    if let Some(ch) = sink.chan.lock().unwrap().as_ref() {
+                        let _ = ch.send(json);
+                    }
+                }
+                true
+            }));
+
+            {
+                let h2 = handle.clone();
+                tauri::async_runtime::spawn(async move {
+                    loop {
+                        let (effect, phase) = {
+                            let st = h2.state::<AppState>();
+                            let mut l = st.lifecycle.lock().await;
+                            let e = l.tick(&st.bridge).await;
+                            (e, l.phase.clone())
+                        };
+                        match effect {
+                            Effect::Target(t) => {
+                                let st = h2.state::<AppState>();
+                                st.bridge.set_target(t);
+                                st.driver.notify();
+                            }
+                            Effect::NeedSpawn => {
+                                let st = h2.state::<AppState>();
+                                let mut l = st.lifecycle.lock().await;
+                                if let Err(e) = l.spawn_child() {
+                                    l.set_phase_stopped(format!("failed to launch daemon: {e}"));
+                                }
+                                drop(l);
+                                st.driver.notify();
+                            }
+                            Effect::NoTarget | Effect::Idle => {}
+                        }
+                        {
+                            let st = h2.state::<AppState>();
+                            tray::set_phase(&st.view, &phase);
+                        }
+                        tray::refresh(&h2);
+                        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+                    }
+                });
+            }
+            Ok(())
+        })
+        .invoke_handler(tauri::generate_handler![
+            app_version,
+            api_get,
+            api_post,
+            api_delete,
+            chat_stream,
+            chat_cancel,
+            events_subscribe,
+            events_unsubscribe,
+            service_state,
+            service_start,
+            service::service_stop,
+            service_restart,
+            service::lan_set,
+            service::agent_set,
+            service::token_copy,
+            thread_list,
+            thread_create,
+            thread_rename,
+            thread_set_model,
+            thread_delete,
+            store_info,
+            messages_page,
+            message_start,
+            messages_delete_from,
+            message_append,
+            message_finish,
+            message_set_content,
+            thread_truncate_and_start
+        ]);
+
+    let app = builder
+        .build(tauri::generate_context!())
+        .expect("failed to start edge0.app");
+    app.run(|app, event| {
+        match event {
+            // Closing the last window is not Quit; Quit is the tray cascade (explicit app.exit(0)).
+            tauri::RunEvent::ExitRequested { ref api, code, .. } => {
+                if code.is_none() {
+                    api.prevent_exit();
+                }
+            }
+            #[cfg(target_os = "macos")]
+            tauri::RunEvent::Reopen {
+                has_visible_windows, ..
+            } => {
+                if !has_visible_windows {
+                    tray::show_main_window_public(app);
+                }
+            }
+            _ => {}
+        }
+    });
+}
diff --git a/macos/app/src-tauri/src/lifecycle.rs b/macos/app/src-tauri/src/lifecycle.rs
new file mode 100644
index 0000000..2822a0d
--- /dev/null
+++ b/macos/app/src-tauri/src/lifecycle.rs
@@ -0,0 +1,398 @@
+//! `daemon.json` is a hint; `GET /health` is the source of truth. No path silently restarts or swaps the binary.
+
+use std::path::PathBuf;
+use std::process::{Child, Command, Stdio};
+use std::time::{Duration, Instant};
+
+use edge0_core::paths::Home;
+use edge0_core::state::{DaemonJson, Scope};
+
+use crate::bridge::client::{Bridge, BridgeTarget};
+
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct Health {
+    pub base: String,
+    pub daemon_version: String,
+    pub host_app_version: Option<String>,
+    pub scope: Scope,
+    pub started_by: String,
+    pub session_id: String,
+    pub pid: u32,
+    pub uptime_s: u64,
+    pub throttled: bool,
+}
+
+impl Health {
+    pub fn version_mismatch(&self) -> bool {
+        match &self.host_app_version {
+            Some(v) => v != env!("CARGO_PKG_VERSION"),
+            None => false,
+        }
+    }
+}
+
+#[derive(Debug, Clone, serde::Serialize)]
+#[serde(tag = "phase", rename_all = "camelCase")]
+pub enum Phase {
+    Probing,
+    Adopted {
+        health: Health,
+        #[serde(skip)]
+        spawned_by_us: bool,
+    },
+    Spawned { health: Health },
+    Exited { code: Option<i32> },
+    Stopped { detail: String },
+    Remote { base: String },
+}
+
+#[derive(Debug, PartialEq)]
+pub enum Effect {
+    Target(BridgeTarget),
+    NoTarget,
+    NeedSpawn,
+    Idle,
+}
+
+pub struct Lifecycle {
+    pub phase: Phase,
+    child: Option<Child>,
+    first_spawn_used: bool,
+    pub home: Home,
+    pub pinned: bool,
+    pub spawn_was_dev: bool,
+    pub port_override: Option<u16>,
+    saw_launchagent_owner: bool,
+}
+
+impl Lifecycle {
+    pub fn new(home: Home, pinned: bool) -> Self {
+        Self {
+            phase: Phase::Probing,
+            child: None,
+            first_spawn_used: false,
+            home,
+            pinned,
+            spawn_was_dev: false,
+            port_override: None,
+            saw_launchagent_owner: false,
+        }
+    }
+
+    pub fn child_pid(&self) -> Option<u32> {
+        self.child.as_ref().map(|c| c.id())
+    }
+
+    pub async fn tick(&mut self, bridge: &Bridge) -> Effect {
+        if self.pinned {
+            self.phase = Phase::Remote {
+                base: bridge.target().map(|t| t.base).unwrap_or_default(),
+            };
+            return Effect::Idle;
+        }
+        if let Some(ch) = &mut self.child {
+            if let Ok(Some(st)) = ch.try_wait() {
+                self.child = None;
+                self.phase = Phase::Exited { code: st.code() };
+                bridge.clear_target();
+                return Effect::NoTarget;
+            }
+        }
+        match probe(bridge, &self.home).await {
+            Some(health) => {
+                if health.scope == Scope::LaunchAgent {
+                    self.saw_launchagent_owner = true;
+                }
+                let spawned_by_us = self.child.as_ref().is_some_and(|c| c.id() == health.pid);
+                let target = health.base.clone();
+                self.phase = if spawned_by_us {
+                    Phase::Spawned {
+                        health: health.clone(),
+                    }
+                } else {
+                    if !spawned_by_us {
+                        self.child = None;
+                    }
+                    Phase::Adopted {
+                        health,
+                        spawned_by_us,
+                    }
+                };
+                match BridgeTarget::from_base(&target) {
+                    Some(t) => Effect::Target(t),
+                    None => Effect::NoTarget,
+                }
+            }
+            None => {
+                bridge.clear_target();
+                if self.child.is_some() {
+                    return Effect::NoTarget;
+                }
+                if matches!(self.phase, Phase::Exited { .. }) {
+                    return Effect::NoTarget;
+                }
+                if matches!(self.phase, Phase::Adopted { health: ref h, .. } if h.scope == Scope::LaunchAgent)
+                {
+                    self.phase = Phase::Stopped {
+                        detail: "LaunchAgent owns the service; waiting for launchd (the app only adopts it)"
+                            .into(),
+                    };
+                    return Effect::Idle;
+                }
+                if !self.first_spawn_used {
+                    self.first_spawn_used = true;
+                    return Effect::NeedSpawn;
+                }
+                self.phase = Phase::Stopped {
+                    detail: "No local service (this session will not auto-start again; use Start service)".into(),
+                };
+                Effect::NoTarget
+            }
+        }
+    }
+
+    pub fn spawn_child(&mut self) -> Result<(), String> {
+        if let Some(ch) = &mut self.child {
+            match ch.try_wait() {
+                Ok(None) => {
+                    return Err(
+                        "Service is already running (the app-owned daemon is alive; use Restart service to apply new configuration)"
+                            .into(),
+                    )
+                }
+                _ => {
+                    let _ = ch.wait();
+                    self.child = None;
+                }
+            }
+        }
+        let (bin, dev) = daemon_binary_path()?;
+        self.spawn_was_dev = dev;
+        let owner = std::process::id().to_string();
+        let home_s = self.home.root.display().to_string();
+        let mut cmd = Command::new(&bin);
+        cmd.args([
+            "--scope",
+            "menu-agent",
+            "--started-by",
+            "gui",
+            "--owner-pid",
+            &owner,
+            "--home",
+            &home_s,
+        ]);
+        if let Some(p) = self.port_override {
+            cmd.args([
+                "--port",
+                &p.to_string(),
+                "--bind",
+                &std::net::Ipv4Addr::LOCALHOST.to_string(),
+            ]);
+        }
+        cmd.env("EDGE0_HOST_APP_VERSION", env!("CARGO_PKG_VERSION"));
+        if !dev {
+            let resources = bin
+                .parent()
+                .and_then(|p| p.parent())
+                .ok_or("Cannot resolve the app bundle Resources path")?;
+            let lib = resources.join("lib").join("libedge0_engine_native.dylib");
+            let worker = resources.join("bin").join("edge0-engine");
+            cmd.env("EDGE0_ENGINE_LIB", lib)
+                .env("EDGE0_WORKER_BIN", worker);
+        }
+        let child = cmd
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .spawn()
+            .map_err(|e| format!("Failed to launch {}: {e}", bin.display()))?;
+        self.child = Some(child);
+        Ok(())
+    }
+
+    pub fn start_requested(&mut self) {
+        self.first_spawn_used = false;
+    }
+
+    pub async fn stop(&mut self, bridge: &Bridge) -> Result<(), String> {
+        let Some(mut ch) = self.child.take() else {
+            return Err(if self.pinned {
+                "This app does not own the service: remote EDGE0_HOST services cannot be stopped here".into()
+            } else {
+                "This app does not own the service: stop the external service or LaunchAgent from its owner".into()
+            });
+        };
+        let pid = ch.id().to_string();
+        let _ = Command::new("/bin/kill").args(["-TERM", &pid]).status();
+        let t0 = Instant::now();
+        loop {
+            match ch.try_wait() {
+                Ok(Some(_)) => break,
+                Ok(None) if t0.elapsed() < Duration::from_secs(10) => {
+                    tokio::time::sleep(Duration::from_millis(50)).await
+                }
+                _ => {
+                    let _ = ch.kill();
+                    let _ = ch.wait();
+                    break;
+                }
+            }
+        }
+        bridge.clear_target();
+        self.first_spawn_used = true;
+        self.phase = Phase::Stopped {
+            detail: "Service stopped (the app remains in the menu bar; use Start service to launch it again)".into(),
+        };
+        Ok(())
+    }
+
+    pub fn agent_released(&mut self) {
+        self.saw_launchagent_owner = false;
+    }
+
+    pub async fn restart_service(&mut self, bridge: &Bridge) -> Result<(), String> {
+        let pid = self.child_pid().or_else(|| daemon_pid_of(&self.home));
+        if let Some(pid) = pid {
+            if self.child.is_none() {
+                let _ = Command::new("/bin/kill")
+                    .args(["-TERM", &pid.to_string()])
+                    .status();
+            }
+            let t0 = Instant::now();
+            loop {
+                let gone = match &mut self.child {
+                    Some(ch) => matches!(ch.try_wait(), Ok(Some(_))),
+                    None => {
+                        !self.home.daemon_json().exists() && !probe_now(bridge, &self.home).await
+                    }
+                };
+                if gone || t0.elapsed() > Duration::from_secs(10) {
+                    break;
+                }
+                tokio::time::sleep(Duration::from_millis(100)).await;
+            }
+            if let Some(mut ch) = self.child.take() {
+                let _ = ch.kill();
+                let _ = ch.wait();
+            }
+        }
+        self.first_spawn_used = false;
+        self.spawn_child()
+    }
+
+    pub fn quit_cascade(&mut self) {
+        if let Some(ch) = &mut self.child {
+            let pid = ch.id().to_string();
+            let _ = Command::new("/bin/kill").args(["-TERM", &pid]).status();
+            let t0 = Instant::now();
+            loop {
+                match ch.try_wait() {
+                    Ok(Some(_)) => break,
+                    Ok(None) if t0.elapsed() < Duration::from_secs(10) => {
+                        std::thread::sleep(Duration::from_millis(50))
+                    }
+                    _ => {
+                        let _ = ch.kill();
+                        let _ = ch.wait();
+                        break;
+                    }
+                }
+            }
+            self.child = None;
+        }
+    }
+
+    pub fn set_phase_stopped(&mut self, detail: String) {
+        self.phase = Phase::Stopped { detail };
+    }
+
+    pub fn suppress_first_spawn(&mut self) {
+        self.first_spawn_used = true;
+    }
+}
+
+pub async fn probe(bridge: &Bridge, home: &Home) -> Option<Health> {
+    let dj = home.daemon_json();
+    let candidate = format!(
+        "http://{}:{}",
+        std::net::Ipv4Addr::LOCALHOST,
+        daemon_json_hint(home)?
+    );
+    match probe_health(bridge, &candidate).await {
+        Some(mut h) => {
+            h.base = candidate;
+            Some(h)
+        }
+        None => {
+            // /health failed: treat the stale daemon.json as absent.
+            if dj.exists() {
+                let _ = std::fs::remove_file(&dj);
+            }
+            None
+        }
+    }
+}
+
+async fn probe_now(bridge: &Bridge, home: &Home) -> bool {
+    probe(bridge, home).await.is_some()
+}
+
+pub async fn probe_health(bridge: &Bridge, base: &str) -> Option<Health> {
+    let v = bridge.health(base).await?;
+    Some(Health {
+        base: base.to_string(),
+        daemon_version: v["daemon_version"].as_str()?.to_string(),
+        host_app_version: v["host_app_version"].as_str().map(str::to_string),
+        scope: Scope::parse(v["scope"].as_str()?).ok()?,
+        started_by: v["started_by"].as_str()?.to_string(),
+        session_id: v["session_id"].as_str()?.to_string(),
+        pid: v["pid"].as_u64()? as u32,
+        uptime_s: v["uptime_s"].as_u64().unwrap_or(0),
+        throttled: v["throttled"].as_bool().unwrap_or(false),
+    })
+}
+
+pub fn daemon_binary_path() -> Result<(PathBuf, bool), String> {
+    let exe = std::env::current_exe().map_err(|e| e.to_string())?;
+    let dir = exe.parent().ok_or("current_exe has no parent directory")?;
+    let bundle = dir.ancestors().nth(2).map(|root| {
+        root.join("Contents")
+            .join("Resources")
+            .join("bin")
+            .join("edge0d")
+    });
+    if let Some(b) = &bundle {
+        if b.exists() {
+            return Ok((b.clone(), false));
+        }
+    }
+    for cand in [
+        dir.to_path_buf(),
+        dir.parent().map(PathBuf::from).unwrap_or_default(),
+    ] {
+        let dev = cand.join("edge0d");
+        if dev.exists() {
+            return Ok((dev, true));
+        }
+    }
+    Err(format!(
+        "Bundled edge0d not found (tried {}/Contents/Resources/bin/edge0d and {})",
+        bundle
+            .and_then(|b| b.parent().map(|p| p.display().to_string()))
+            .unwrap_or_default(),
+        dir.display()
+    ))
+}
+
+pub fn daemon_json_hint(home: &Home) -> Option<u16> {
+    daemon_json(home).map(|d| d.port)
+}
+
+pub fn daemon_json(home: &Home) -> Option<DaemonJson> {
+    std::fs::read_to_string(home.daemon_json())
+        .ok()
+        .and_then(|t| serde_json::from_str(&t).ok())
+}
+
+fn daemon_pid_of(home: &Home) -> Option<u32> {
+    daemon_json(home).map(|d| d.pid)
+}
diff --git a/macos/app/src-tauri/src/main.rs b/macos/app/src-tauri/src/main.rs
new file mode 100644
index 0000000..95c8079
--- /dev/null
+++ b/macos/app/src-tauri/src/main.rs
@@ -0,0 +1,5 @@
+#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
+
+fn main() {
+    edge0_app::run();
+}
diff --git a/macos/app/src-tauri/src/service.rs b/macos/app/src-tauri/src/service.rs
new file mode 100644
index 0000000..22b7118
--- /dev/null
+++ b/macos/app/src-tauri/src/service.rs
@@ -0,0 +1,204 @@
+//! Service panel commands. Lock order is always svc_gate, then lifecycle.
+
+use std::process::Command;
+use std::time::{Duration, Instant};
+
+use edge0_core::paths::Home;
+use edge0_core::state::Scope;
+use tauri::{AppHandle, Manager, State};
+
+use crate::lifecycle::{Effect, Lifecycle, Phase};
+use crate::AppState;
+
+// ---------------------------------------------------------------- token_copy
+
+pub fn copy_token(home: &Home) -> Result<(), String> {
+    let raw = std::fs::read_to_string(home.token())
+        .ok()
+        .map(|s| s.trim().to_string())
+        .filter(|s| !s.is_empty());
+    let Some(token) = raw else {
+        return Err("no-credential".into());
+    };
+    let mut clipboard = arboard::Clipboard::new().map_err(|_| "io".to_string())?;
+    clipboard.set_text(token).map_err(|_| "io".to_string())
+}
+
+#[tauri::command]
+pub(crate) fn token_copy(state: State<'_, AppState>) -> Result<(), String> {
+    copy_token(state.bridge.home())
+}
+
+// ---------------------------------------------------------------- service_stop
+
+#[tauri::command]
+pub(crate) async fn service_stop(app: AppHandle, state: State<'_, AppState>) -> Result<(), String> {
+    let (gate, bridge, lifecycle) = (
+        state.svc_gate.clone(),
+        state.bridge.clone(),
+        state.lifecycle.clone(),
+    );
+    let _g = gate.lock().await;
+    let resume = crate::hush_for_lifecycle(&app).await;
+    let r = {
+        let mut l = lifecycle.lock().await;
+        l.stop(&bridge).await
+    };
+    resume();
+    r
+}
+
+// ---------------------------------------------------------------- lan_set
+
+pub fn apply_lan_config(home: &Home, on: bool) -> Result<(), String> {
+    let lan = std::net::Ipv4Addr::UNSPECIFIED.to_string();
+    let loopback = std::net::Ipv4Addr::LOCALHOST.to_string();
+    let (bind, enable) = if on {
+        (lan.as_str(), "1")
+    } else {
+        (loopback.as_str(), "false")
+    };
+    edge0_core::config::upsert_edge0_keys(
+        &home.config_toml(),
+        &[("bind", bind), ("enable_lan", enable)],
+    )
+}
+
+#[tauri::command]
+pub(crate) async fn lan_set(
+    app: AppHandle,
+    state: State<'_, AppState>,
+    on: bool,
+) -> Result<(), String> {
+    let (gate, home, lifecycle) = (
+        state.svc_gate.clone(),
+        state.bridge.home().clone(),
+        state.lifecycle.clone(),
+    );
+    let _g = gate.lock().await;
+    if lifecycle.lock().await.pinned {
+        return Err("Remote mode (EDGE0_HOST pinned) does not manage the local service".into());
+    }
+    apply_lan_config(&home, on)?;
+    crate::do_service_restart(&app).await
+}
+
+// ---------------------------------------------------------------- agent_set
+
+pub fn agent_args(home: &Home, enable: bool) -> Vec<String> {
+    vec![
+        if enable {
+            "--enable-agent"
+        } else {
+            "--disable-agent"
+        }
+        .to_string(),
+        "--home".to_string(),
+        home.root.display().to_string(),
+    ]
+}
+
+async fn yield_ownership(l: &mut Lifecycle, bridge: &crate::bridge::Bridge) {
+    let _ = l.stop(bridge).await;
+}
+
+async fn launchagent_adopted(app: &AppHandle) -> bool {
+    let st = app.state::<AppState>();
+    let mut l = st.lifecycle.lock().await;
+    let effect = l.tick(&st.bridge).await;
+    if let Effect::Target(t) = effect {
+        st.bridge.set_target(t);
+    }
+    matches!(
+        &l.phase,
+        Phase::Adopted { health, .. } if health.scope == Scope::LaunchAgent
+    )
+}
+
+#[tauri::command]
+pub(crate) async fn agent_set(
+    app: AppHandle,
+    state: State<'_, AppState>,
+    enable: bool,
+) -> Result<(), String> {
+    let (gate, bridge, lifecycle, driver) = (
+        state.svc_gate.clone(),
+        state.bridge.clone(),
+        state.lifecycle.clone(),
+        state.driver.clone(),
+    );
+    let _g = gate.lock().await;
+    if lifecycle.lock().await.pinned {
+        return Err("Remote mode (EDGE0_HOST pinned) does not manage the local service".into());
+    }
+    let home = bridge.home().clone();
+    let (bin, _dev) = crate::lifecycle::daemon_binary_path()?;
+    let bin_for_msg = bin.clone();
+
+    let resume = crate::hush_for_lifecycle(&app).await;
+    {
+        let mut l = lifecycle.lock().await;
+        yield_ownership(&mut l, &bridge).await;
+    }
+
+    let args = agent_args(&home, enable);
+    let spawned = tokio::task::spawn_blocking(move || {
+        Command::new(&bin)
+            .args(&args)
+            .stdout(std::process::Stdio::null())
+            .stderr(std::process::Stdio::null())
+            .status()
+    })
+    .await;
+    let failed: Option<String> = match &spawned {
+        Err(e) => Some(format!("Failed to schedule registration: {e}")),
+        Ok(Err(e)) => Some(format!("Failed to launch {}: {e}", bin_for_msg.display())),
+        Ok(Ok(st)) if !st.success() => Some(if enable {
+            "LaunchAgent registration failed (edge0d --enable-agent returned a non-zero status)"
+                .into()
+        } else {
+            "LaunchAgent removal failed (edge0d --disable-agent returned a non-zero status)".into()
+        }),
+        Ok(Ok(_)) => None,
+    };
+    if let Some(msg) = failed {
+        resume();
+        return Err(msg);
+    }
+
+    if enable {
+        let deadline = Instant::now() + Duration::from_secs(20);
+        let mut adopted = false;
+        while Instant::now() < deadline {
+            if launchagent_adopted(&app).await {
+                adopted = true;
+                break;
+            }
+            tokio::time::sleep(Duration::from_millis(500)).await;
+        }
+        resume();
+        if !adopted {
+            return Err(
+                "LaunchAgent registered, but it did not take ownership within 20 seconds (the port may be occupied or the login session may restrict launchd)".into(),
+            );
+        }
+    } else {
+        {
+            let mut l = lifecycle.lock().await;
+            l.agent_released();
+            let _ = l.tick(&bridge).await;
+            let occupied = matches!(l.phase, Phase::Spawned { .. } | Phase::Adopted { .. });
+            if !occupied {
+                l.start_requested();
+                if let Err(e) = l.spawn_child() {
+                    l.set_phase_stopped(format!(
+                        "Failed to return ownership to the app after removing LaunchAgent: {e}"
+                    ));
+                }
+            }
+        }
+        resume();
+    }
+    driver.notify();
+    Ok(())
+}
diff --git a/macos/app/src-tauri/src/store.rs b/macos/app/src-tauri/src/store.rs
new file mode 100644
index 0000000..d4fdad8
--- /dev/null
+++ b/macos/app/src-tauri/src/store.rs
@@ -0,0 +1,552 @@
+//! Conversation store. This process is the only writer; the webview never sees SQL or file paths.
+
+use std::path::Path;
+use std::sync::atomic::{AtomicI64, Ordering};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use rusqlite::{params, Connection};
+use serde::Serialize;
+
+use edge0_core::conversation::{MessageRecord, ThreadRecord};
+
+pub use edge0_core::conversation::{MessageStatus, Role, Usage};
+use edge0_core::paths::Home;
+
+const SCHEMA_VERSION: u32 = 1;
+
+const MIGRATIONS: &[(u32, &[&str])] = &[(
+    1,
+    &[
+        "CREATE TABLE IF NOT EXISTS threads(
+            id         TEXT PRIMARY KEY,
+            data       TEXT NOT NULL,
+            created_at INTEGER NOT NULL,
+            updated_at INTEGER NOT NULL
+        )",
+        "CREATE TABLE IF NOT EXISTS messages(
+            id         TEXT PRIMARY KEY,
+            thread_id  TEXT NOT NULL,
+            data       TEXT NOT NULL,
+            created_at INTEGER NOT NULL
+        )",
+        "CREATE INDEX IF NOT EXISTS idx_messages_thread_created ON messages(thread_id, created_at)",
+        "CREATE INDEX IF NOT EXISTS idx_threads_updated ON threads(updated_at)",
+    ],
+)];
+
+fn wall_ms() -> i64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_millis() as i64)
+        .unwrap_or(0)
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ThreadRow {
+    pub id: String,
+    pub created_at: i64,
+    pub updated_at: i64,
+    pub data: ThreadRecord,
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct MessageRow {
+    pub id: String,
+    pub created_at: i64,
+    pub data: MessageRecord,
+}
+
+#[derive(Debug)]
+pub struct Store {
+    conn: Connection,
+    /// Write timestamps are `max(wall, last+1)` so `(created_at, id)` pagination is deterministic when user/assistant land in the same millisecond (ids are random UUIDs).
+    clock: AtomicI64,
+}
+
+impl Store {
+    pub fn open(home: &Home) -> Result<(Store, usize), String> {
+        Self::open_path(&home.edge0_db())
+    }
+
+    pub fn open_path(path: &Path) -> Result<(Store, usize), String> {
+        if let Some(dir) = path.parent() {
+            std::fs::create_dir_all(dir)
+                .map_err(|e| format!("failed to create {}: {e}", dir.display()))?;
+        }
+        let conn = Connection::open(path).map_err(|e| format!("failed to open conversation store: {e}"))?;
+        let mut s = Store {
+            conn,
+            clock: AtomicI64::new(0),
+        };
+        s.migrate()?;
+        let hi = s
+            .conn
+            .query_row(
+                "SELECT COALESCE(MAX(x), 0) FROM (
+                    SELECT created_at AS x FROM messages
+                    UNION ALL SELECT created_at FROM threads
+                    UNION ALL SELECT updated_at FROM threads
+                 )",
+                [],
+                |r| r.get::<_, i64>(0),
+            )
+            .unwrap_or(0);
+        s.clock.store(hi.max(wall_ms()), Ordering::Relaxed);
+        let reconciled = s.reconcile_streaming()?;
+        Ok((s, reconciled))
+    }
+
+    fn next_ts(&self) -> i64 {
+        loop {
+            let last = self.clock.load(Ordering::Relaxed);
+            let next = wall_ms().max(last + 1);
+            if self
+                .clock
+                .compare_exchange_weak(last, next, Ordering::Relaxed, Ordering::Relaxed)
+                .is_ok()
+            {
+                return next;
+            }
+        }
+    }
+
+    fn user_version(&self) -> Result<u32, String> {
+        self.conn
+            .query_row("PRAGMA user_version", [], |r| r.get::<_, u32>(0))
+            .map_err(|e| e.to_string())
+    }
+
+    fn migrate(&mut self) -> Result<(), String> {
+        let v = self.user_version()?;
+        if v > SCHEMA_VERSION {
+            return Err(format!(
+                "conversation store generation {v} is newer than this app supports ({SCHEMA_VERSION}) — upgrade the app; do not guess downward"
+            ));
+        }
+        for (ver, ddls) in MIGRATIONS {
+            if *ver <= v {
+                continue;
+            }
+            let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+            for ddl in *ddls {
+                tx.execute_batch(ddl)
+                    .map_err(|e| format!("migration {ver} failed: {e}"))?;
+            }
+            tx.pragma_update(None, "user_version", ver)
+                .map_err(|e| e.to_string())?;
+            tx.commit().map_err(|e| e.to_string())?;
+        }
+        Ok(())
+    }
+
+    pub fn reconcile_streaming(&mut self) -> Result<usize, String> {
+        let mut rows: Vec<(String, String)> = self
+            .conn
+            .prepare("SELECT id, data FROM messages")
+            .map_err(|e| e.to_string())?
+            .query_map([], |r| Ok((r.get(0)?, r.get(1)?)))
+            .map_err(|e| e.to_string())?
+            .collect::<Result<_, _>>()
+            .map_err(|e| e.to_string())?;
+        let mut n = 0;
+        let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+        for (id, data) in rows.drain(..) {
+            let mut rec: MessageRecord = match serde_json::from_str(&data) {
+                Ok(r) => r,
+                Err(_) => continue,
+            };
+            if rec.status != MessageStatus::Streaming {
+                continue;
+            }
+            rec.status = MessageStatus::Aborted;
+            rec.validate()
+                .map_err(|e| format!("reconciled record is invalid: {e}"))?;
+            tx.execute(
+                "UPDATE messages SET data = ?1 WHERE id = ?2",
+                params![serde_json::to_string(&rec).map_err(|e| e.to_string())?, id],
+            )
+            .map_err(|e| e.to_string())?;
+            n += 1;
+        }
+        tx.commit().map_err(|e| e.to_string())?;
+        Ok(n)
+    }
+
+    // ------------------------------------------------------------------ threads
+
+    pub fn threads(&self) -> Result<Vec<ThreadRow>, String> {
+        let mut st = self
+            .conn
+            .prepare(
+                "SELECT id, created_at, updated_at, data FROM threads ORDER BY updated_at DESC",
+            )
+            .map_err(|e| e.to_string())?;
+        let rows = st
+            .query_map([], |r| {
+                Ok((
+                    r.get::<_, String>(0)?,
+                    r.get::<_, i64>(1)?,
+                    r.get::<_, i64>(2)?,
+                    r.get::<_, String>(3)?,
+                ))
+            })
+            .map_err(|e| e.to_string())?;
+        let mut out = Vec::new();
+        for r in rows {
+            let (id, created_at, updated_at, data) = r.map_err(|e| e.to_string())?;
+            out.push(ThreadRow {
+                id,
+                created_at,
+                updated_at,
+                data: serde_json::from_str(&data).map_err(|e| format!("corrupt thread record: {e}"))?,
+            });
+        }
+        Ok(out)
+    }
+
+    pub fn create_thread(&self, title: &str, model: Option<String>) -> Result<ThreadRow, String> {
+        let id = format!("thr-{}", uuid::Uuid::new_v4());
+        let now = self.next_ts();
+        let data = ThreadRecord::new(title, model);
+        self.conn
+            .execute(
+                "INSERT INTO threads(id, data, created_at, updated_at) VALUES (?1, ?2, ?3, ?4)",
+                params![id, serde_json::to_string(&data).unwrap(), now, now],
+            )
+            .map_err(|e| e.to_string())?;
+        Ok(ThreadRow {
+            id,
+            created_at: now,
+            updated_at: now,
+            data,
+        })
+    }
+
+    fn thread_data(&self, id: &str) -> Result<(String, i64), String> {
+        self.conn
+            .query_row(
+                "SELECT data, updated_at FROM threads WHERE id = ?1",
+                params![id],
+                |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?)),
+            )
+            .map_err(|_| format!("thread not found: {id}"))
+    }
+
+    pub fn rename_thread(&self, id: &str, title: &str) -> Result<(), String> {
+        let (data, _) = self.thread_data(id)?;
+        let mut rec: ThreadRecord = serde_json::from_str(&data).map_err(|e| e.to_string())?;
+        rec.title = title.into();
+        self.conn
+            .execute(
+                "UPDATE threads SET data = ?1, updated_at = ?2 WHERE id = ?3",
+                params![serde_json::to_string(&rec).unwrap(), self.next_ts(), id],
+            )
+            .map_err(|e| e.to_string())?;
+        Ok(())
+    }
+
+    pub fn set_thread_model(&self, id: &str, model: Option<String>) -> Result<(), String> {
+        let (data, _) = self.thread_data(id)?;
+        let mut rec: ThreadRecord = serde_json::from_str(&data).map_err(|e| e.to_string())?;
+        rec.model = model;
+        self.conn
+            .execute(
+                "UPDATE threads SET data = ?1 WHERE id = ?2",
+                params![serde_json::to_string(&rec).unwrap(), id],
+            )
+            .map_err(|e| e.to_string())?;
+        Ok(())
+    }
+
+    pub fn delete_thread(&mut self, id: &str) -> Result<(), String> {
+        let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+        tx.execute("DELETE FROM messages WHERE thread_id = ?1", params![id])
+            .map_err(|e| e.to_string())?;
+        tx.execute("DELETE FROM threads WHERE id = ?1", params![id])
+            .map_err(|e| e.to_string())?;
+        tx.commit().map_err(|e| e.to_string())
+    }
+}
+
+fn touch_thread(tx: &rusqlite::Transaction<'_>, thread_id: &str, ts: i64) -> Result<(), String> {
+    tx.execute(
+        "UPDATE threads SET updated_at = ?1 WHERE id = ?2",
+        params![ts, thread_id],
+    )
+    .map_err(|e| e.to_string())?;
+    Ok(())
+}
+
+fn write_message_in(
+    tx: &rusqlite::Transaction<'_>,
+    id: &str,
+    rec: &MessageRecord,
+) -> Result<(), String> {
+    rec.validate()?;
+    tx.execute(
+        "UPDATE messages SET data = ?1 WHERE id = ?2",
+        params![serde_json::to_string(rec).map_err(|e| e.to_string())?, id],
+    )
+    .map_err(|e| e.to_string())?;
+    Ok(())
+}
+
+impl Store {
+    // ------------------------------------------------------------------ messages
+
+    pub fn messages_page(
+        &self,
+        thread_id: &str,
+        before: Option<(i64, &str)>,
+        limit: u32,
+    ) -> Result<Vec<MessageRow>, String> {
+        let (sql, pv): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = match before {
+            None => (
+                "SELECT id, created_at, data FROM messages
+                 WHERE thread_id = ?1 ORDER BY created_at DESC, id DESC LIMIT ?2",
+                vec![Box::new(thread_id.to_string()), Box::new(limit)],
+            ),
+            Some((ca, id)) => (
+                "SELECT id, created_at, data FROM messages
+                 WHERE thread_id = ?1 AND (created_at < ?2 OR (created_at = ?2 AND id < ?3))
+                 ORDER BY created_at DESC, id DESC LIMIT ?4",
+                vec![
+                    Box::new(thread_id.to_string()),
+                    Box::new(ca),
+                    Box::new(id.to_string()),
+                    Box::new(limit),
+                ],
+            ),
+        };
+        let mut st = self.conn.prepare(sql).map_err(|e| e.to_string())?;
+        let refs: Vec<&dyn rusqlite::types::ToSql> = pv.iter().map(|p| p.as_ref()).collect();
+        let rows = st
+            .query_map(refs.as_slice(), |r| {
+                Ok((
+                    r.get::<_, String>(0)?,
+                    r.get::<_, i64>(1)?,
+                    r.get::<_, String>(2)?,
+                ))
+            })
+            .map_err(|e| e.to_string())?;
+        let mut out = Vec::new();
+        for r in rows {
+            let (id, created_at, data) = r.map_err(|e| e.to_string())?;
+            out.push(MessageRow {
+                id,
+                created_at,
+                data: serde_json::from_str(&data).map_err(|e| format!("corrupt message record: {e}"))?,
+            });
+        }
+        out.reverse();
+        Ok(out)
+    }
+
+    fn message_data(&self, id: &str) -> Result<(String, String), String> {
+        self.conn
+            .query_row(
+                "SELECT thread_id, data FROM messages WHERE id = ?1",
+                params![id],
+                |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)),
+            )
+            .map_err(|_| format!("message not found: {id}"))
+    }
+
+    fn write_message(&self, id: &str, rec: &MessageRecord) -> Result<(), String> {
+        rec.validate()?;
+        self.conn
+            .execute(
+                "UPDATE messages SET data = ?1 WHERE id = ?2",
+                params![serde_json::to_string(rec).map_err(|e| e.to_string())?, id],
+            )
+            .map_err(|e| e.to_string())?;
+        Ok(())
+    }
+
+    pub fn message_start(
+        &mut self,
+        thread_id: &str,
+        role: Role,
+        content: &str,
+        model: &str,
+        params_snapshot: Option<serde_json::Value>,
+        id: Option<String>,
+    ) -> Result<MessageRow, String> {
+        let id = match id {
+            Some(given) => {
+                if !given.starts_with("msg-")
+                    || given.len() > 80
+                    || !given
+                        .chars()
+                        .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
+                {
+                    return Err("explicit message id must use the msg- prefix and a safe character set".into());
+                }
+                let exists: i64 = self
+                    .conn
+                    .query_row(
+                        "SELECT COUNT(*) FROM messages WHERE id = ?1",
+                        params![given],
+                        |r| r.get(0),
+                    )
+                    .map_err(|e| e.to_string())?;
+                if exists > 0 {
+                    return Err(format!("message id already exists: {given}"));
+                }
+                given
+            }
+            None => format!("msg-{}", uuid::Uuid::new_v4()),
+        };
+        let now = self.next_ts();
+        let mut rec = MessageRecord::new(role, content, model, now);
+        rec.thread_id = Some(thread_id.into());
+        rec.params = params_snapshot;
+        rec.validate()?;
+        let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+        tx.execute(
+            "INSERT INTO messages(id, thread_id, data, created_at) VALUES (?1, ?2, ?3, ?4)",
+            params![id, thread_id, serde_json::to_string(&rec).unwrap(), now],
+        )
+        .map_err(|e| e.to_string())?;
+        touch_thread(&tx, thread_id, now)?;
+        tx.commit().map_err(|e| e.to_string())?;
+        Ok(MessageRow {
+            id,
+            created_at: now,
+            data: rec,
+        })
+    }
+
+    pub fn message_append(
+        &mut self,
+        id: &str,
+        content_delta: Option<&str>,
+        reasoning_delta: Option<&str>,
+    ) -> Result<(), String> {
+        let (thread_id, data) = self.message_data(id)?;
+        let mut rec: MessageRecord = serde_json::from_str(&data).map_err(|e| e.to_string())?;
+        if let Some(d) = content_delta {
+            rec.content.push_str(d);
+        }
+        if let Some(d) = reasoning_delta {
+            rec.reasoning_content
+                .get_or_insert_with(String::new)
+                .push_str(d);
+        }
+        let ts = self.next_ts();
+        let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+        write_message_in(&tx, id, &rec)?;
+        touch_thread(&tx, &thread_id, ts)?;
+        tx.commit().map_err(|e| e.to_string())
+    }
+
+    pub fn message_finish(
+        &mut self,
+        id: &str,
+        status: MessageStatus,
+        code: Option<&str>,
+        usage: Option<Usage>,
+    ) -> Result<MessageRow, String> {
+        let (thread_id, data) = self.message_data(id)?;
+        let mut rec: MessageRecord = serde_json::from_str(&data).map_err(|e| e.to_string())?;
+        if rec.status == MessageStatus::Streaming && status == MessageStatus::Streaming {
+            return Err("streaming is not a terminal status".into());
+        }
+        rec.status = status;
+        rec.code = code.map(str::to_string);
+        if usage.is_some() {
+            rec.usage = usage;
+        }
+        rec.validate()?;
+        let ts = self.next_ts();
+        let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+        write_message_in(&tx, id, &rec)?;
+        touch_thread(&tx, &thread_id, ts)?;
+        tx.commit().map_err(|e| e.to_string())?;
+        Ok(MessageRow {
+            id: id.into(),
+            created_at: rec.created_at,
+            data: rec,
+        })
+    }
+
+    pub fn message_set_content(&self, id: &str, content: &str) -> Result<(), String> {
+        let (_, data) = self.message_data(id)?;
+        let mut rec: MessageRecord = serde_json::from_str(&data).map_err(|e| e.to_string())?;
+        rec.content = content.into();
+        self.write_message(id, &rec)
+    }
+
+    pub fn thread_truncate_and_start(
+        &mut self,
+        thread_id: &str,
+        from_created_at: i64,
+        role: Role,
+        content: &str,
+        model: &str,
+        params_snapshot: Option<serde_json::Value>,
+    ) -> Result<MessageRow, String> {
+        let id = format!("msg-{}", uuid::Uuid::new_v4());
+        let now = self.next_ts();
+        let mut rec = MessageRecord::new(role, content, model, now);
+        rec.thread_id = Some(thread_id.into());
+        rec.params = params_snapshot;
+        rec.validate()?;
+        let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+        tx.execute(
+            "DELETE FROM messages WHERE thread_id = ?1 AND created_at >= ?2",
+            params![thread_id, from_created_at],
+        )
+        .map_err(|e| e.to_string())?;
+        tx.execute(
+            "INSERT INTO messages(id, thread_id, data, created_at) VALUES (?1, ?2, ?3, ?4)",
+            params![id, thread_id, serde_json::to_string(&rec).unwrap(), now],
+        )
+        .map_err(|e| e.to_string())?;
+        touch_thread(&tx, thread_id, now)?;
+        tx.commit().map_err(|e| e.to_string())?;
+        Ok(MessageRow {
+            id,
+            created_at: now,
+            data: rec,
+        })
+    }
+
+    pub fn messages_delete_from(
+        &mut self,
+        thread_id: &str,
+        from_id: &str,
+    ) -> Result<usize, String> {
+        let from_created_at: i64 = self
+            .conn
+            .query_row(
+                "SELECT created_at FROM messages WHERE id = ?1 AND thread_id = ?2",
+                params![from_id, thread_id],
+                |r| r.get(0),
+            )
+            .map_err(|_| format!("message is not in this thread: {from_id}"))?;
+        let ts = self.next_ts();
+        let tx = self.conn.transaction().map_err(|e| e.to_string())?;
+        let n = tx
+            .execute(
+                "DELETE FROM messages WHERE thread_id = ?1 AND created_at >= ?2",
+                params![thread_id, from_created_at],
+            )
+            .map_err(|e| e.to_string())?;
+        touch_thread(&tx, thread_id, ts)?;
+        tx.commit().map_err(|e| e.to_string())?;
+        Ok(n)
+    }
+
+    pub fn with_conn<T>(&self, f: impl FnOnce(&Connection) -> T) -> T {
+        f(&self.conn)
+    }
+}
+
+impl Drop for Store {
+    fn drop(&mut self) {
+        let _ = self.conn.execute_batch("PRAGMA wal_checkpoint(TRUNCATE)");
+    }
+}
diff --git a/macos/app/src-tauri/src/tray.rs b/macos/app/src-tauri/src/tray.rs
new file mode 100644
index 0000000..740c07b
--- /dev/null
+++ b/macos/app/src-tauri/src/tray.rs
@@ -0,0 +1,250 @@
+//! Menu-bar status. Numbers come from the snapshot; the tray does not estimate.
+
+use std::sync::{Arc, Mutex};
+
+use edge0_core::protocol::SystemSnapshot;
+use tauri::image::Image;
+use tauri::menu::{Menu, MenuItem, PredefinedMenuItem};
+use tauri::tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent};
+use tauri::{AppHandle, Manager, WebviewUrl, WebviewWindowBuilder, Wry};
+
+use crate::bridge::events::{EventFrame, LinkState};
+use crate::lifecycle::Phase;
+
+#[derive(Default, Clone)]
+pub struct TrayView {
+    pub snapshot: Option<SystemSnapshot>,
+    #[allow(dead_code)]
+    pub link: Option<LinkState>,
+    pub loading: bool,
+    pub phase: Option<String>,
+}
+
+pub const TRAY_ID: &str = "edge0-tray";
+const MENU_SHOW: &str = "menu-show";
+const MENU_RESTART: &str = "menu-restart";
+const MENU_QUIT: &str = "menu-quit";
+
+pub fn build(app: &AppHandle, status_slot: &Mutex<Option<MenuItem<Wry>>>) -> Result<(), String> {
+    let status = MenuItem::with_id(
+        app,
+        "menu-status",
+        "edge0: probing service...",
+        false,
+        None::<&str>,
+    )
+    .map_err(|e| e.to_string())?;
+    let show = MenuItem::with_id(app, MENU_SHOW, "Open window", true, None::<&str>)
+        .map_err(|e| e.to_string())?;
+    let restart = MenuItem::with_id(app, MENU_RESTART, "Restart service", true, None::<&str>)
+        .map_err(|e| e.to_string())?;
+    let quit = MenuItem::with_id(app, MENU_QUIT, "Quit edge0", true, None::<&str>)
+        .map_err(|e| e.to_string())?;
+    let menu2 = Menu::with_items(
+        app,
+        &[
+            &status,
+            &PredefinedMenuItem::separator(app).map_err(|e| e.to_string())?,
+            &show,
+            &restart,
+            &PredefinedMenuItem::separator(app).map_err(|e| e.to_string())?,
+            &quit,
+        ],
+    )
+    .map_err(|e| e.to_string())?;
+
+    *status_slot.lock().unwrap() = Some(status);
+
+    TrayIconBuilder::with_id(TRAY_ID)
+        .menu(&menu2)
+        .show_menu_on_left_click(false)
+        .icon(Image::from_bytes(include_bytes!("../icons/tray-idle.png")).expect("tray-idle"))
+        .on_tray_icon_event({
+            let app = app.clone();
+            move |tray, ev| {
+                if let TrayIconEvent::Click {
+                    button: MouseButton::Left,
+                    button_state: MouseButtonState::Up,
+                    ..
+                } = ev
+                {
+                    let _ = tray;
+                    show_main_window(&app);
+                }
+            }
+        })
+        .on_menu_event({
+            let app = app.clone();
+            move |_tray, ev| match ev.id().as_ref() {
+                MENU_SHOW => show_main_window(&app),
+                MENU_RESTART => {
+                    let a = app.clone();
+                    tauri::async_runtime::spawn(async move {
+                        if let Err(e) = crate::do_service_restart(&a).await {
+                            tracing::warn!("Tray restart action failed: {e}");
+                        }
+                    });
+                }
+                MENU_QUIT => quit_cascade(&app),
+                _ => {}
+            }
+        })
+        .build(app)
+        .map_err(|e| e.to_string())?;
+    Ok(())
+}
+
+pub fn show_main_window_public(app: &AppHandle) {
+    show_main_window(app);
+}
+
+fn show_main_window(app: &AppHandle) {
+    if let Some(w) = app.get_webview_window("main") {
+        let _ = w.unminimize();
+        let _ = w.show();
+        let _ = w.set_focus();
+        return;
+    }
+    let headless = std::env::var("EDGE0_WINDOWLESS").is_ok_and(|v| v == "1");
+    match WebviewWindowBuilder::new(app, "main", WebviewUrl::App("index.html".into()))
+        .title("edge0")
+        .inner_size(1100.0, 720.0)
+        .visible(!headless)
+        .build()
+    {
+        Ok(w) => {
+            let _ = w.set_focus();
+        }
+        Err(e) => tracing::warn!("Failed to recreate the main window: {e}"),
+    }
+}
+
+pub fn quit_cascade(app: &AppHandle) {
+    {
+        let st = app.state::<crate::AppState>();
+        let _ = st.driver_stop.clone().send(true);
+        for h in st.streams.lock().unwrap().values() {
+            h.abort();
+        }
+        let lc = st.lifecycle.clone();
+        let mut l = tauri::async_runtime::block_on(lc.lock());
+        l.quit_cascade();
+    }
+    app.exit(0);
+}
+
+pub fn refresh(app: &AppHandle) {
+    let st = app.state::<crate::AppState>();
+    let view = st.view.lock().unwrap().clone();
+    let Some(tray) = app.tray_by_id(TRAY_ID) else {
+        return;
+    };
+    let icon: &'static [u8] = if view.snapshot.as_ref().is_some_and(|s| s.daemon.throttled) {
+        include_bytes!("../icons/tray-throttled.png")
+    } else if view.loading {
+        include_bytes!("../icons/tray-loading.png")
+    } else if view
+        .snapshot
+        .as_ref()
+        .is_some_and(|s| s.models.iter().any(|m| m.state == "resident"))
+    {
+        include_bytes!("../icons/tray-resident.png")
+    } else {
+        include_bytes!("../icons/tray-idle.png")
+    };
+    if let Ok(img) = Image::from_bytes(icon) {
+        let _ = tray.set_icon(Some(img));
+    }
+    let text = match &view.snapshot {
+        Some(s) => {
+            let resident: Vec<String> = s
+                .models
+                .iter()
+                .filter(|m| m.state == "resident")
+                .map(|m| {
+                    format!(
+                        "{} until {}",
+                        m.id,
+                        m.unload_at
+                            .as_deref()
+                            .map(hhmm)
+                            .unwrap_or_else(|| "inf".to_string())
+                    )
+                })
+                .collect();
+            let mut line = format!(
+                "Service :{} - daemon v{} - {:.1} GiB free",
+                s.daemon.port, s.daemon.version, s.hardware.disk_free_gib
+            );
+            if !resident.is_empty() {
+                line.push_str(&format!(" - resident {}", resident.join(", ")));
+            }
+            if s.daemon.throttled {
+                line.push_str(" - throttled");
+            }
+            if matches!(view.phase.as_deref(), Some(p) if p.contains("version mismatch") || p.contains("mismatch"))
+            {
+                line.push_str(" - version mismatch");
+            }
+            line
+        }
+        None => format!(
+            "edge0: {}",
+            view.phase
+                .clone()
+                .unwrap_or_else(|| "probing service...".into())
+        ),
+    };
+    let item = st.tray_status.lock().unwrap().clone();
+    if let Some(item) = item {
+        let _ = item.set_text(text);
+    }
+}
+
+pub fn apply_frame(view: &Arc<Mutex<TrayView>>, f: &EventFrame) {
+    let mut g = view.lock().unwrap();
+    match f {
+        EventFrame::Snapshot { snapshot } => {
+            g.snapshot = Some((**snapshot).clone());
+            if !snapshot.models.iter().any(|m| m.state == "resident") {
+                g.loading = false;
+            }
+        }
+        EventFrame::Event { event } => match event.etype.as_str() {
+            "model.load.started" => g.loading = true,
+            "model.load.ready" | "model.load.failed" | "model.unloaded" => g.loading = false,
+            _ => {}
+        },
+        EventFrame::Status { state } => g.link = Some(*state),
+    }
+}
+
+pub fn set_phase(view: &Arc<Mutex<TrayView>>, phase: &Phase) {
+    let label = match phase {
+        Phase::Probing => "probing service...".to_string(),
+        Phase::Adopted { health, .. } => format!(
+            "Connected to local service v{}, scope={} (adopt)",
+            health.daemon_version,
+            health.scope.as_str()
+        ),
+        Phase::Spawned { health } => {
+            format!(
+                "Connected to local service v{} (menu-agent, app-owned)",
+                health.daemon_version
+            )
+        }
+        Phase::Exited { code } => format!("daemon exited (code {code:?})"),
+        Phase::Stopped { detail } => detail.clone(),
+        Phase::Remote { base } => format!("remote mode: {base}"),
+    };
+    view.lock().unwrap().phase = Some(label);
+}
+
+fn hhmm(iso: &str) -> String {
+    chrono::DateTime::parse_from_rfc3339(iso)
+        .map(|d| {
+            let local: chrono::DateTime<chrono::Local> = d.into();
+            local.format("%H:%M").to_string()
+        })
+        .unwrap_or_else(|_| iso.to_string())
+}
diff --git a/macos/app/src-tauri/tauri.conf.json b/macos/app/src-tauri/tauri.conf.json
new file mode 100644
index 0000000..4471620
--- /dev/null
+++ b/macos/app/src-tauri/tauri.conf.json
@@ -0,0 +1,32 @@
+{
+  "$schema": "https://schema.tauri.app/config/2",
+  "productName": "edge0",
+  "version": "0.1.0",
+  "identifier": "app.edge0.client",
+  "build": {
+    "beforeDevCommand": "yarn dev",
+    "devUrl": "http://localhost:1420",
+    "beforeBuildCommand": "yarn build",
+    "frontendDist": "../dist"
+  },
+  "app": {
+    "windows": [],
+    "security": {
+      "csp": "default-src 'self'; connect-src 'self' ipc: http://ipc.localhost; img-src 'self' data:; style-src 'self' 'unsafe-inline'; font-src 'self' data:"
+    }
+  },
+  "bundle": {
+    "active": true,
+    "category": "DeveloperTool",
+    "icon": ["icons/icon.icns", "icons/edge0_logo.png"],
+    "macOS": {
+      "minimumSystemVersion": "14.0"
+    }
+  },
+  "plugins": {
+    "updater": {
+      "pubkey": "",
+      "endpoints": []
+    }
+  }
+}
diff --git a/macos/app/src/bridge/edge0Fetch.ts b/macos/app/src/bridge/edge0Fetch.ts
new file mode 100644
index 0000000..640a468
--- /dev/null
+++ b/macos/app/src/bridge/edge0Fetch.ts
@@ -0,0 +1,135 @@
+import { Channel, invoke } from "@tauri-apps/api/core";
+
+import type { StreamFrame } from "./frames";
+import { toApiError } from "./frames";
+
+export const FAKE_ORIGIN = "daemon://edge0";
+
+export interface Edge0FetchInit {
+  method?: "GET" | "POST" | "DELETE";
+  body?: unknown;
+  stream?: boolean;
+  signal?: AbortSignal;
+  headers?: Record<string, string>;
+}
+
+function resolvePath(input: string | URL): string {
+  const s = typeof input === "string" ? input : input.toString();
+  if (s.startsWith(FAKE_ORIGIN)) {
+    const p = s.slice(FAKE_ORIGIN.length);
+    if (!p.startsWith("/")) throw new TypeError("missing absolute path");
+    return assertPath(p);
+  }
+  if (/^[a-z]+:\/\//i.test(s)) {
+    throw new TypeError("edge0Fetch does not accept a scheme/host target (the webview must not hold the daemon address)");
+  }
+  return assertPath(s);
+}
+
+function assertPath(p: string): string {
+  if (!p.startsWith("/v1/") || p.includes("//") || p.includes("://")) {
+    throw new TypeError("path must be an absolute /v1/ path on the daemon");
+  }
+  return p;
+}
+
+function checkHeaders(headers?: Record<string, string>): void {
+  for (const k of Object.keys(headers ?? {})) {
+    if (k.toLowerCase() === "authorization") {
+      throw new TypeError("Authorization is injected by the shell; the UI must not handle the token");
+    }
+  }
+}
+
+export async function edge0Fetch(
+  input: string | URL,
+  init: Edge0FetchInit = {},
+): Promise<Response> {
+  const path = resolvePath(input);
+  const method = init.method ?? "GET";
+  checkHeaders(init.headers);
+
+  if (init.stream) {
+    if (method !== "POST" || path !== "/v1/chat/completions") {
+      throw new TypeError("stream is only valid for POST /v1/chat/completions");
+    }
+    return streamChat(init);
+  }
+
+  const json =
+    method === "GET"
+      ? await invoke<unknown>("api_get", { path })
+      : method === "DELETE"
+        ? await invoke<unknown>("api_delete", { path })
+        : await invoke<unknown>("api_post", { path, body: init.body ?? null });
+  return Response.json(json);
+}
+
+async function streamChat(init: Edge0FetchInit): Promise<Response> {
+  const channel = new Channel<string>();
+  let controllerRef: ReadableStreamDefaultController<Uint8Array> | null = null;
+  let streamId: string | null = null;
+  let closed = false;
+  const enc = new TextEncoder();
+
+  const body = new ReadableStream<Uint8Array>({
+    start(controller) {
+      controllerRef = controller;
+    },
+    cancel() {
+      // Consumer abort (AI SDK) must close the TCP stream so the daemon observes e0_cancel.
+      if (streamId) void invoke("chat_cancel", { id: streamId });
+    },
+  });
+
+  const onFrame = (raw: string) => {
+    let frame: StreamFrame;
+    try {
+      frame = JSON.parse(raw) as StreamFrame;
+    } catch {
+      return;
+    }
+    const c = controllerRef;
+    if (!c || closed) return;
+    switch (frame.kind) {
+      case "chunk":
+        for (const line of frame.data.split("\n")) {
+          c.enqueue(enc.encode(`data: ${line}\n`));
+        }
+        c.enqueue(enc.encode("\n"));
+        break;
+      case "done":
+        closed = true;
+        c.close();
+        break;
+      case "error":
+        closed = true;
+        c.error(toApiError(JSON.stringify(frame)));
+        break;
+    }
+  };
+  channel.onmessage = onFrame;
+
+  let status: number;
+  try {
+    const init_ = await invoke<{ status: number; id: string }>("chat_stream", {
+      body: init.body ?? {},
+      onFrame: channel,
+    });
+    status = init_.status;
+    streamId = init_.id;
+  } catch (e) {
+    throw toApiError(String(e));
+  }
+
+  const abort = () => {
+    if (streamId) void invoke("chat_cancel", { id: streamId });
+  };
+  init.signal?.addEventListener("abort", abort, { once: true });
+
+  const res = new Response(body, {
+    status,
+    headers: { "content-type": "text/event-stream" },
+  });
+  return res;
+}
diff --git a/macos/app/src/bridge/frames.ts b/macos/app/src/bridge/frames.ts
new file mode 100644
index 0000000..ad50ee3
--- /dev/null
+++ b/macos/app/src/bridge/frames.ts
@@ -0,0 +1,41 @@
+import type { Event } from "../gen/Event";
+import type { SystemSnapshot } from "../gen/SystemSnapshot";
+
+export type StreamFrame =
+  | { kind: "chunk"; data: string }
+  | { kind: "done" }
+  | { kind: "error"; code: string | null; message: string };
+
+export type EventFrame =
+  | { kind: "snapshot"; snapshot: SystemSnapshot }
+  | { kind: "event"; event: Event }
+  | { kind: "status"; state: "live" | "reconnecting" | "resyncing" };
+
+export interface ApiErrorShape {
+  status: number | null;
+  code: string | null;
+  message: string;
+}
+
+export class ApiError extends Error implements ApiErrorShape {
+  status: number | null;
+  code: string | null;
+  constructor(status: number | null, code: string | null, message: string) {
+    super(message);
+    this.name = "ApiError";
+    this.status = status;
+    this.code = code;
+  }
+}
+
+export function toApiError(raw: string): ApiError {
+  try {
+    const o = JSON.parse(raw) as { status?: number | null; code?: string | null; message?: string };
+    if (o && typeof o === "object" && "message" in o) {
+      return new ApiError(o.status ?? null, o.code ?? null, o.message ?? raw);
+    }
+  } catch {
+    /* non-JSON errors are passed through as-is */
+  }
+  return new ApiError(null, null, raw);
+}
diff --git a/macos/app/src/bridge/ipc.ts b/macos/app/src/bridge/ipc.ts
new file mode 100644
index 0000000..2e5f447
--- /dev/null
+++ b/macos/app/src/bridge/ipc.ts
@@ -0,0 +1,43 @@
+import { Channel, invoke } from "@tauri-apps/api/core";
+
+import { toApiError } from "./frames";
+import type { EventFrame, StreamFrame } from "./frames";
+
+export async function apiGet(path: string): Promise<unknown> {
+  try {
+    return await invoke<unknown>("api_get", { path });
+  } catch (e) {
+    throw toApiError(String(e));
+  }
+}
+
+export async function apiPost(path: string, body: unknown): Promise<unknown> {
+  try {
+    return await invoke<unknown>("api_post", { path, body });
+  } catch (e) {
+    throw toApiError(String(e));
+  }
+}
+
+export async function apiDelete(path: string): Promise<unknown> {
+  try {
+    return await invoke<unknown>("api_delete", { path });
+  } catch (e) {
+    throw toApiError(String(e));
+  }
+}
+
+export function subscribeEventDriver(onFrame: (f: EventFrame) => void): () => void {
+  const channel = new Channel<string>();
+  channel.onmessage = (raw: string) => {
+    try {
+      onFrame(JSON.parse(raw) as EventFrame);
+    } catch {
+      /* ignore malformed event frames */
+    }
+  };
+  void invoke("events_subscribe", { onFrame: channel }).catch(() => undefined);
+  return () => void invoke("events_unsubscribe").catch(() => undefined);
+}
+
+export type { StreamFrame };
diff --git a/macos/app/src/bridge/state.ts b/macos/app/src/bridge/state.ts
new file mode 100644
index 0000000..37a248d
--- /dev/null
+++ b/macos/app/src/bridge/state.ts
@@ -0,0 +1,89 @@
+import type { Event, ModelEntry, SystemSnapshot } from "../gen";
+import type { EventFrame } from "./frames";
+import { subscribeEventDriver } from "./ipc";
+
+export type LinkState = "idle" | "live" | "reconnecting" | "resyncing";
+
+type Listener = () => void;
+
+class DaemonStore {
+  snapshot: SystemSnapshot | null = null;
+  link: LinkState = "idle";
+  lastSeq = 0;
+  private folded = new Map<string, Event>();
+  private listeners = new Set<Listener>();
+  private unsubscribe: (() => void) | null = null;
+
+  ingest(snapshot: SystemSnapshot): void {
+    this.snapshot = snapshot;
+    this.emit();
+  }
+
+  subscribe(fn: Listener): () => void {
+    this.listeners.add(fn);
+    if (!this.unsubscribe) {
+      this.unsubscribe = subscribeEventDriver((f) => this.applyFrame(f));
+    }
+    return () => {
+      this.listeners.delete(fn);
+    };
+  }
+
+  private applyFrame(f: EventFrame): void {
+    switch (f.kind) {
+      case "snapshot":
+        this.snapshot = f.snapshot;
+        for (const key of [...this.folded.keys()]) {
+          if (key.startsWith("api.request.finished:")) this.folded.delete(key);
+        }
+        break;
+      case "event": {
+        this.lastSeq = f.event.seq;
+        this.folded.set(`${f.event.type}:${f.event.subject}`, f.event);
+        if (this.snapshot && f.event.type === "model.load.ready") {
+          for (const m of this.snapshot.models) {
+            if (m.id === f.event.subject) m.state = "resident" satisfies ModelEntry["state"];
+          }
+        }
+        break;
+      }
+      case "status":
+        this.link = f.state;
+        break;
+    }
+    this.emit();
+  }
+
+  latestFolded(type: string, subject: string): Event | undefined {
+    return this.folded.get(`${type}:${subject}`);
+  }
+
+  requestsWindow(): Event[] {
+    const out: Event[] = [];
+    for (const [key, ev] of this.folded) {
+      if (key.startsWith("api.request.finished:")) out.push(ev);
+    }
+    return out.sort((a, b) => b.seq - a.seq);
+  }
+
+  clear(): void {
+    this.snapshot = null;
+    this.folded.clear();
+    this.lastSeq = 0;
+  }
+
+  maxSeqAmong(types: readonly string[]): number {
+    let best = 0;
+    for (const [key, ev] of this.folded) {
+      const t = key.slice(0, key.lastIndexOf(":"));
+      if (types.includes(t) && ev.seq > best) best = ev.seq;
+    }
+    return best;
+  }
+
+  private emit(): void {
+    for (const l of this.listeners) l();
+  }
+}
+
+export const daemonStore = new DaemonStore();
diff --git a/macos/app/src/chat/buckets.ts b/macos/app/src/chat/buckets.ts
new file mode 100644
index 0000000..88444d5
--- /dev/null
+++ b/macos/app/src/chat/buckets.ts
@@ -0,0 +1,30 @@
+export interface Bucket<T> {
+  key: "today" | "yesterday" | "week" | "earlier";
+  items: T[];
+}
+
+const DAY_MS = 24 * 60 * 60 * 1000;
+
+function localDayStart(ms: number): number {
+  const d = new Date(ms);
+  d.setHours(0, 0, 0, 0);
+  return d.getTime();
+}
+
+export function bucketThreads<T extends { updatedAt: number }>(
+  threads: T[],
+  nowMs: number,
+): Bucket<T>[] {
+  const today = localDayStart(nowMs);
+  const buckets: Record<Bucket<T>["key"], T[]> = { today: [], yesterday: [], week: [], earlier: [] };
+  for (const th of threads) {
+    const day = localDayStart(th.updatedAt);
+    if (day >= today) buckets.today.push(th);
+    else if (day >= today - DAY_MS) buckets.yesterday.push(th);
+    else if (day >= today - 6 * DAY_MS) buckets.week.push(th);
+    else buckets.earlier.push(th);
+  }
+  return (["today", "yesterday", "week", "earlier"] as const)
+    .filter((k) => buckets[k].length > 0)
+    .map((k) => ({ key: k, items: buckets[k] }));
+}
diff --git a/macos/app/src/chat/errorMap.ts b/macos/app/src/chat/errorMap.ts
new file mode 100644
index 0000000..2e92698
--- /dev/null
+++ b/macos/app/src/chat/errorMap.ts
@@ -0,0 +1,19 @@
+const CHAT_ERROR_KEYS: Record<string, string> = {
+  "E-SRV-NOTLOAD": "chat.err.notload",
+  "E-SRV-BUSY": "chat.err.busy",
+  "E-SRV-CANCEL": "chat.err.cancel",
+  "E-SRV-WORKER": "chat.err.worker",
+  "E-SRV-PARAM": "chat.err.param",
+  "E-MODEL-MISSING": "chat.err.missing",
+  "E-MEM-LOAD": "chat.err.memload",
+  "E-GPU-DEVICE": "chat.err.gpu",
+  "E-SRV-VERSION": "chat.err.version",
+};
+
+export function errorNextKey(code: string | null): string {
+  return (code && CHAT_ERROR_KEYS[code]) || "chat.err.generic";
+}
+
+export function hasKnownError(code: string | null): boolean {
+  return code !== null && code in CHAT_ERROR_KEYS;
+}
diff --git a/macos/app/src/chat/ipc.ts b/macos/app/src/chat/ipc.ts
new file mode 100644
index 0000000..ef94593
--- /dev/null
+++ b/macos/app/src/chat/ipc.ts
@@ -0,0 +1,108 @@
+import { invoke } from "@tauri-apps/api/core";
+
+import type { MessageRecord, ThreadRecord, Usage } from "../gen";
+
+export interface ThreadRow {
+  id: string;
+  createdAt: number;
+  updatedAt: number;
+  data: ThreadRecord;
+}
+
+export interface MessageRow {
+  id: string;
+  createdAt: number;
+  data: MessageRecord;
+}
+
+export function threadList(): Promise<ThreadRow[]> {
+  return invoke<ThreadRow[]>("thread_list");
+}
+
+export function threadCreate(title: string, model: string | null): Promise<ThreadRow> {
+  return invoke<ThreadRow>("thread_create", { title, model });
+}
+
+export function threadRename(id: string, title: string): Promise<void> {
+  return invoke<void>("thread_rename", { id, title });
+}
+
+export function threadSetModel(id: string, model: string | null): Promise<void> {
+  return invoke<void>("thread_set_model", { id, model });
+}
+
+export function threadDelete(id: string): Promise<void> {
+  return invoke<void>("thread_delete", { id });
+}
+
+export function storeInfo(): Promise<{ reconciled: number }> {
+  return invoke<{ reconciled: number }>("store_info");
+}
+
+export interface MessagePage {
+  threadId: string;
+  beforeCreatedAt?: number;
+  beforeId?: string;
+  limit?: number;
+}
+
+export function messagesPage(p: MessagePage): Promise<MessageRow[]> {
+  if ((p.beforeCreatedAt === undefined) !== (p.beforeId === undefined)) {
+    return Promise.reject(new Error("page cursor fields must be provided together"));
+  }
+  return invoke<MessageRow[]>("messages_page", {
+    threadId: p.threadId,
+    beforeCreatedAt: p.beforeCreatedAt ?? null,
+    beforeId: p.beforeId ?? null,
+    limit: p.limit ?? 50,
+  });
+}
+
+export interface StartInput {
+  threadId: string;
+  role: "user" | "assistant" | "tool";
+  content: string;
+  model: string;
+  params?: Record<string, unknown> | null;
+  id?: string;
+}
+
+export function messageStart(in_: StartInput): Promise<MessageRow> {
+  return invoke<MessageRow>("message_start", {
+    threadId: in_.threadId,
+    role: in_.role,
+    content: in_.content,
+    model: in_.model,
+    params: in_.params ?? null,
+    id: in_.id ?? null,
+  });
+}
+
+export function messageAppend(
+  id: string,
+  contentDelta: string | null,
+  reasoningDelta: string | null,
+): Promise<void> {
+  return invoke<void>("message_append", { id, contentDelta, reasoningDelta });
+}
+
+export function messageFinish(
+  id: string,
+  status: "streaming" | "done" | "aborted" | "error",
+  code: string | null,
+  usage: Usage | null,
+): Promise<MessageRow> {
+  return invoke<MessageRow>("message_finish", { id, status, code, usage });
+}
+
+export function messageSetContent(id: string, content: string): Promise<void> {
+  return invoke<void>("message_set_content", { id, content });
+}
+
+export function messagesDeleteFrom(threadId: string, fromId: string): Promise<number> {
+  return invoke<number>("messages_delete_from", { threadId, fromId });
+}
+
+export function newMessageId(): string {
+  return `msg-${crypto.randomUUID()}`;
+}
diff --git a/macos/app/src/chat/modelStatus.ts b/macos/app/src/chat/modelStatus.ts
new file mode 100644
index 0000000..3d55fdd
--- /dev/null
+++ b/macos/app/src/chat/modelStatus.ts
@@ -0,0 +1,40 @@
+import type { Event, SystemSnapshot } from "../gen";
+
+export type ModelLight = "resident" | "loading" | "installed" | "absent";
+
+const MODEL_EVENT_TYPES = [
+  "model.load.started",
+  "model.load.ready",
+  "model.load.failed",
+  "model.unloaded",
+] as const;
+
+export function latestModelEvent(
+  store: { latestFolded(type: string, subject: string): Event | undefined },
+  modelId: string,
+): Event | null {
+  let best: Event | null = null;
+  for (const t of MODEL_EVENT_TYPES) {
+    const e = store.latestFolded(t, modelId);
+    if (e && (!best || e.seq > best.seq)) best = e;
+  }
+  return best;
+}
+
+export function deriveLight(
+  snapshot: SystemSnapshot | null,
+  modelId: string,
+  latest: Event | null,
+): ModelLight {
+  if (!snapshot) return "absent";
+  const entry = snapshot.models.find((m) => m.id === modelId);
+  if (!entry) return "absent";
+  if (latest && latest.subject === modelId && latest.type === "model.load.started") {
+    if (entry.state !== "resident") return "loading";
+  }
+  return entry.state === "resident" ? "resident" : "installed";
+}
+
+export function candidateModels(snapshot: SystemSnapshot | null): string[] {
+  return (snapshot?.models ?? []).map((m) => m.id);
+}
diff --git a/macos/app/src/chat/records.ts b/macos/app/src/chat/records.ts
new file mode 100644
index 0000000..e5242d4
--- /dev/null
+++ b/macos/app/src/chat/records.ts
@@ -0,0 +1,50 @@
+import type { UIMessage } from "ai";
+
+import type { MessageRow } from "./ipc";
+import type { PerfEdge0 } from "./sse";
+
+export interface ChatMessageMeta {
+  status?: "streaming" | "done" | "aborted" | "error";
+  code?: string | null;
+  usage?: MessageRow["data"]["usage"] | null;
+  perf?: PerfEdge0 | null;
+  model?: string;
+  saveFailed?: boolean;
+}
+
+export type ChatUIMessage = UIMessage<ChatMessageMeta>;
+
+export function rowToUiMessage(row: MessageRow): ChatUIMessage {
+  const parts: ChatUIMessage["parts"] = [];
+  if (row.data.reasoning_content) {
+    parts.push({ type: "reasoning", text: row.data.reasoning_content });
+  }
+  if (row.data.content) {
+    parts.push({ type: "text", text: row.data.content });
+  }
+  return {
+    id: row.id,
+    role: row.data.role === "user" ? "user" : "assistant",
+    parts,
+    metadata: {
+      status: row.data.status,
+      code: row.data.code ?? null,
+      usage: row.data.usage ?? null,
+      model: row.data.model,
+    },
+  };
+}
+
+export function messageText(m: ChatUIMessage): string {
+  return m.parts
+    .filter((p): p is { type: "text"; text: string } => p.type === "text")
+    .map((p) => p.text)
+    .join("");
+}
+
+export function messageReasoning(m: ChatUIMessage): string {
+  return m.parts
+    .filter((p): p is { type: "reasoning"; text: string } => p.type === "reasoning")
+    .map((p) => p.text)
+    .join("");
+}
diff --git a/macos/app/src/chat/scroll.ts b/macos/app/src/chat/scroll.ts
new file mode 100644
index 0000000..b5c28aa
--- /dev/null
+++ b/macos/app/src/chat/scroll.ts
@@ -0,0 +1,11 @@
+export const STICK_BOTTOM_THRESHOLD = 40;
+
+export interface ScrollGeometry {
+  scrollHeight: number;
+  scrollTop: number;
+  clientHeight: number;
+}
+
+export function isAtBottom(el: ScrollGeometry, threshold = STICK_BOTTOM_THRESHOLD): boolean {
+  return el.scrollHeight - el.scrollTop - el.clientHeight < threshold;
+}
diff --git a/macos/app/src/chat/sse.ts b/macos/app/src/chat/sse.ts
new file mode 100644
index 0000000..f12f3f4
--- /dev/null
+++ b/macos/app/src/chat/sse.ts
@@ -0,0 +1,127 @@
+import type { Usage } from "../gen";
+import { ApiError } from "../bridge/frames";
+
+export interface SseEvent {
+  data: string;
+}
+
+export function splitSseEvents(buf: string): { events: SseEvent[]; rest: string } {
+  const events: SseEvent[] = [];
+  let rest = buf;
+  for (;;) {
+    const m = rest.match(/\r?\n\r?\n/);
+    if (!m || m.index === undefined) break;
+    const block = rest.slice(0, m.index);
+    rest = rest.slice(m.index + m[0].length);
+    const dataLines: string[] = [];
+    for (const line of block.split(/\r?\n/)) {
+      if (line.startsWith("data:")) {
+        dataLines.push(line.slice(5).trimStart());
+      }
+    }
+    if (dataLines.length > 0) events.push({ data: dataLines.join("\n") });
+  }
+  return { events, rest };
+}
+
+/** Keep `delta.content` and `delta.reasoning_content` on separate sinks so reasoning never enters the visible body. */
+export interface ChatDeltaSink {
+  onText(text: string): void;
+  onReasoning(text: string): void;
+  onUsage(usage: Usage, perf?: PerfEdge0): void;
+}
+
+export interface PerfEdge0 {
+  ttft_ms?: number;
+  prefill_tps?: number;
+  decode_tps?: number;
+  footprint_mb?: number;
+}
+
+export async function consumeChatSse(
+  res: Response,
+  sink: ChatDeltaSink,
+  signal?: AbortSignal,
+): Promise<void> {
+  const reader = res.body?.getReader();
+  if (!reader) throw new ApiError(res.status, null, "streaming response has no body");
+  let aborted = false;
+  const onAbort = () => {
+    aborted = true;
+    void reader.cancel();
+  };
+  signal?.addEventListener("abort", onAbort, { once: true });
+  const decoder = new TextDecoder();
+  let buf = "";
+  let doneFlagged = false;
+  for (;;) {
+    const { done, value } = await reader.read();
+    if (done) break;
+    buf += decoder.decode(value, { stream: true });
+    const { events, rest } = splitSseEvents(buf);
+    buf = rest;
+    for (const ev of events) {
+      if (ev.data === "[DONE]") {
+        doneFlagged = true;
+        continue;
+      }
+      let obj: {
+        choices?: { delta?: { content?: string; reasoning_content?: string }; finish_reason?: string | null }[];
+        usage?: {
+          prompt_tokens?: number;
+          completion_tokens?: number;
+          total_tokens?: number;
+          edge0?: Record<string, unknown> | null;
+        } | null;
+        error?: { code?: string; message?: string };
+      };
+      try {
+        obj = JSON.parse(ev.data);
+      } catch {
+        continue;
+      }
+      if (obj.error) {
+        throw new ApiError(res.status, obj.error.code ?? "E-SRV-WORKER", obj.error.message ?? "in-stream error");
+      }
+      const delta = obj.choices?.[0]?.delta;
+      if (delta?.content) sink.onText(delta.content);
+      if (delta?.reasoning_content) sink.onReasoning(delta.reasoning_content);
+      if (obj.usage && typeof obj.usage.total_tokens === "number") {
+        const e = obj.usage.edge0 ?? {};
+        const num = (k: string): number | undefined =>
+          typeof e[k] === "number" ? (e[k] as number) : undefined;
+        const perf: PerfEdge0 = {
+          ttft_ms: num("ttft_ms"),
+          prefill_tps: num("prefill_tps"),
+          decode_tps: num("decode_tps"),
+          footprint_mb: num("footprint_mb"),
+        };
+        sink.onUsage(
+          {
+            prompt_tokens: obj.usage.prompt_tokens ?? 0,
+            completion_tokens: obj.usage.completion_tokens ?? 0,
+            total_tokens: obj.usage.total_tokens,
+          },
+          perf,
+        );
+      }
+    }
+  }
+  signal?.removeEventListener("abort", onAbort);
+  if (aborted) return;
+  if (!doneFlagged && buf.trim().length > 0) {
+    throw new ApiError(res.status, "E-SRV-WORKER", "stream ended without [DONE] (connection truncated)");
+  }
+}
+
+export async function parseErrorEnvelope(res: Response): Promise<ApiError> {
+  let code: string | null = null;
+  let message = `HTTP ${res.status}`;
+  try {
+    const j = (await res.json()) as { error?: { code?: string; message?: string } };
+    code = j?.error?.code ?? null;
+    message = j?.error?.message ?? message;
+  } catch {
+  }
+  return new ApiError(res.status, code, message);
+}
diff --git a/macos/app/src/chat/transport.ts b/macos/app/src/chat/transport.ts
new file mode 100644
index 0000000..f866fdc
--- /dev/null
+++ b/macos/app/src/chat/transport.ts
@@ -0,0 +1,239 @@
+import type { ChatTransport, UIMessage, UIMessageChunk } from "ai";
+
+import { ApiError } from "../bridge/frames";
+import type { Usage } from "../gen";
+import { edge0Fetch, FAKE_ORIGIN } from "../bridge/edge0Fetch";
+import { currentTemperature } from "../lib/generation";
+import { consumeChatSse, parseErrorEnvelope, type PerfEdge0 } from "./sse";
+
+export interface AssistantPersistence {
+  beginAssistant(): Promise<string>;
+  append(id: string, contentDelta: string | null, reasoningDelta: string | null): Promise<void>;
+  finish(
+    id: string,
+    status: "done" | "aborted" | "error",
+    code: string | null,
+    usage: Usage | null,
+  ): Promise<void>;
+  dropAssistant(fromId: string): Promise<void>;
+}
+
+export interface TransportModelCtx {
+  model: string;
+  params?: Record<string, unknown>;
+  systemPrompt?: string | null;
+}
+
+export interface TransportDeps {
+  getModel: () => TransportModelCtx | null;
+  getPersistence: () => AssistantPersistence;
+  flushMs?: number;
+  flushEvery?: number;
+}
+
+type Fetcher = typeof edge0Fetch;
+
+const TEXT_ID = "text-main";
+const REASONING_ID = "reasoning-main";
+
+export function toWireMessages(
+  messages: UIMessage[],
+  systemPrompt?: string | null,
+): { role: "system" | "user" | "assistant"; content: string }[] {
+  const out: { role: "system" | "user" | "assistant"; content: string }[] = [];
+  if (systemPrompt?.trim()) {
+    out.push({ role: "system", content: systemPrompt });
+  }
+  for (const m of messages) {
+    if (m.role !== "user" && m.role !== "assistant") continue;
+    let text = "";
+    for (const part of m.parts) {
+      if (part.type === "text") text += part.text;
+    }
+    if (text.length > 0 || m.role === "user") out.push({ role: m.role, content: text });
+  }
+  return out;
+}
+
+class DeltaSink {
+  pendingContent = "";
+  pendingReasoning = "";
+  sinceFlush = 0;
+  constructor(
+    private readonly onFlush: () => void,
+    private readonly flushEvery: number,
+  ) {}
+  collect(content: string | null, reasoning: string | null) {
+    if (content) this.pendingContent += content;
+    if (reasoning) this.pendingReasoning += reasoning;
+    this.sinceFlush += content ? 1 : 0;
+    if (this.sinceFlush >= this.flushEvery) this.onFlush();
+  }
+  take(): [string | null, string | null] {
+    const pair: [string | null, string | null] = [
+      this.pendingContent || null,
+      this.pendingReasoning || null,
+    ];
+    this.pendingContent = "";
+    this.pendingReasoning = "";
+    this.sinceFlush = 0;
+    return pair;
+  }
+}
+
+export function createEdge0Transport(
+  deps: TransportDeps,
+  fetcher: Fetcher = edge0Fetch,
+): ChatTransport<UIMessage> {
+  const flushMs = deps.flushMs ?? 250;
+  const flushEvery = deps.flushEvery ?? 32;
+
+  return {
+    async sendMessages(options): Promise<ReadableStream<UIMessageChunk>> {
+      const modelCtx = deps.getModel();
+      const persist = deps.getPersistence();
+
+      const stream = new ReadableStream<UIMessageChunk>({
+        start(controller) {
+          void run(controller);
+        },
+      });
+
+      async function run(controller: ReadableStreamDefaultController<UIMessageChunk>) {
+        const push = (c: UIMessageChunk) => controller.enqueue(c);
+
+        if (options.trigger === "regenerate-message" && options.messageId) {
+          try {
+            await persist.dropAssistant(options.messageId);
+          } catch {
+          }
+        }
+
+        let assistantId: string;
+        try {
+          assistantId = await persist.beginAssistant();
+        } catch (e) {
+          push({ type: "start" });
+          push({ type: "error", errorText: JSON.stringify({ local: true, message: String(e) }) });
+          push({ type: "finish" });
+          controller.close();
+          return;
+        }
+        push({ type: "start", messageId: assistantId });
+
+        if (!modelCtx?.model) {
+          await persist
+            .finish(assistantId, "error", "E-MODEL-MISSING", null)
+            .catch(() => undefined);
+          push({ type: "error", errorText: JSON.stringify({ code: "E-MODEL-MISSING", message: "No model selected" }) });
+          push({ type: "finish" });
+          controller.close();
+          return;
+        }
+
+        let saveFailed = false;
+        const sink = new DeltaSink(() => void flush(), flushEvery);
+        let openChannel: "reasoning" | "text" | null = null;
+        let usage: Usage | null = null;
+        let perf: PerfEdge0 | null = null;
+
+        async function flush() {
+          const [c, r] = sink.take();
+          if (!c && !r) return;
+          try {
+            await persist.append(assistantId, c, r);
+          } catch {
+            saveFailed = true;
+          }
+        }
+        const timer = setInterval(() => void flush(), flushMs);
+
+        const switchTo = (ch: "text" | "reasoning") => {
+          if (openChannel === ch) return;
+          if (openChannel === "text") push({ type: "text-end", id: TEXT_ID });
+          if (openChannel === "reasoning") push({ type: "reasoning-end", id: REASONING_ID });
+          if (ch === "text") push({ type: "text-start", id: TEXT_ID });
+          else push({ type: "reasoning-start", id: REASONING_ID });
+          openChannel = ch;
+        };
+
+        let terminal: { status: "done" | "aborted" | "error"; code: string | null };
+        try {
+          const res = await fetcher(`${FAKE_ORIGIN}/v1/chat/completions`, {
+            method: "POST",
+            stream: true,
+            body: {
+              model: modelCtx.model,
+              messages: toWireMessages(options.messages, modelCtx.systemPrompt),
+              stream: true,
+              temperature: currentTemperature(),
+              // Chat needs a full reply; daemon default 64 is a fixture cap, not a chat cap.
+              max_tokens: 1024,
+              ...(modelCtx.params ?? {}),
+            },
+            signal: options.abortSignal ?? undefined,
+          });
+          if (!res.ok) {
+            throw await parseErrorEnvelope(res); // 404 NOTLOAD / 429 BUSY / 507 …
+          }
+          await consumeChatSse(res, {
+            onText: (s) => {
+              switchTo("text");
+              push({ type: "text-delta", id: TEXT_ID, delta: s });
+              sink.collect(s, null);
+            },
+            onReasoning: (s) => {
+              switchTo("reasoning");
+              push({ type: "reasoning-delta", id: REASONING_ID, delta: s });
+              sink.collect(null, s);
+            },
+            onUsage: (u, p) => {
+              usage = u;
+              perf = p ?? null;
+            },
+          }, options.abortSignal ?? undefined);
+          terminal = options.abortSignal?.aborted ? { status: "aborted", code: null } : { status: "done", code: null };
+        } catch (e) {
+          const apiErr = e instanceof ApiError ? e : null;
+          if (options.abortSignal?.aborted) {
+            terminal = { status: "aborted", code: null };
+          } else {
+            const code = apiErr?.code ?? "E-SRV-WORKER";
+            const message = apiErr?.message ?? String(e);
+            push({ type: "error", errorText: JSON.stringify({ code, message, status: apiErr?.status ?? null }) });
+            terminal = { status: "error", code };
+          }
+        } finally {
+          clearInterval(timer);
+        }
+
+        if (openChannel === "text") push({ type: "text-end", id: TEXT_ID });
+        if (openChannel === "reasoning") push({ type: "reasoning-end", id: REASONING_ID });
+
+        await flush();
+        try {
+          await persist.finish(assistantId, terminal.status, terminal.code, usage);
+        } catch {
+          saveFailed = true;
+        }
+        push({ type: "finish", messageMetadata: { usage, perf, saveFailed, status: terminal.status } });
+        controller.close();
+      }
+
+      return stream;
+    },
+
+    async reconnectToStream() {
+      return null;
+    },
+  };
+}
+
+export function decodeErrorChunk(errorText: string): { code: string | null; message: string; status: number | null } {
+  try {
+    const o = JSON.parse(errorText) as { code?: string | null; message?: string; status?: number | null };
+    return { code: o.code ?? null, message: o.message ?? errorText, status: o.status ?? null };
+  } catch {
+    return { code: null, message: errorText, status: null };
+  }
+}
diff --git a/macos/app/src/chat/useThreadChat.ts b/macos/app/src/chat/useThreadChat.ts
new file mode 100644
index 0000000..e3edd1f
--- /dev/null
+++ b/macos/app/src/chat/useThreadChat.ts
@@ -0,0 +1,161 @@
+import { useChat } from "@ai-sdk/react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+
+import {
+  messageFinish,
+  messageAppend,
+  messagesDeleteFrom,
+  messagesPage,
+  messageStart,
+  newMessageId,
+  threadRename,
+  type ThreadRow,
+} from "./ipc";
+import { rowToUiMessage, type ChatUIMessage } from "./records";
+import { createEdge0Transport } from "./transport";
+
+export interface ThreadChatApi {
+  messages: ChatUIMessage[];
+  status: "submitted" | "streaming" | "ready" | "error";
+  send: (text: string) => Promise<void>;
+  stop: () => void;
+  regenerate: () => Promise<void>;
+  editResend: (messageId: string, newText: string) => Promise<void>;
+  loadEarlier: () => Promise<boolean>;
+  hasMoreEarlier: boolean;
+  historyLoaded: boolean;
+  busy: boolean;
+}
+
+export function useThreadChat(
+  thread: ThreadRow,
+  model: string,
+  onThreadChanged: () => void,
+  onRoundSettled?: (assistantText: string) => void,
+  systemPrompt?: string | null,
+): ThreadChatApi {
+  const [hasMoreEarlier, setHasMoreEarlier] = useState(false);
+  const oldestRef = useRef<{ createdAt: number; id: string } | null>(null);
+  const modelRef = useRef(model);
+  modelRef.current = model;
+
+  const transport = useMemo(
+    () =>
+      createEdge0Transport({
+        getModel: () =>
+          modelRef.current
+            ? { model: modelRef.current, params: {}, systemPrompt }
+            : null,
+        getPersistence: () => ({
+          beginAssistant: async () => {
+            const row = await messageStart({
+              threadId: thread.id,
+              role: "assistant",
+              content: "",
+              model: modelRef.current,
+            });
+            return row.id;
+          },
+          append: (id, c, r) => messageAppend(id, c, r),
+          finish: (id, s, code, usage) => messageFinish(id, s, code, usage).then(() => undefined),
+          dropAssistant: (fromId) => messagesDeleteFrom(thread.id, fromId).then(() => undefined),
+        }),
+      }),
+    [thread.id, systemPrompt],
+  );
+
+  const { messages, sendMessage, regenerate: chatRegenerate, stop, setMessages, status } = useChat<ChatUIMessage>({
+    id: thread.id,
+    messages: [],
+    transport,
+    onFinish: () => {
+      onThreadChanged();
+      onRoundSettledRef.current?.("");
+    },
+  });
+
+  const [historyLoaded, setHistoryLoaded] = useState(false);
+  useEffect(() => {
+    setHistoryLoaded(false);
+    void (async () => {
+      const rows = await messagesPage({ threadId: thread.id, limit: 50 });
+      setHasMoreEarlier(rows.length === 50);
+      oldestRef.current = rows.length > 0 ? { createdAt: rows[0].createdAt, id: rows[0].id } : null;
+      const list = rows.map(rowToUiMessage);
+      if (list.length > 0) {
+        setMessages((cur) => (cur.length === 0 ? list : cur));
+      }
+      setHistoryLoaded(true);
+    })();
+  }, [thread.id, setMessages]);
+
+  const onRoundSettledRef = useRef(onRoundSettled);
+  onRoundSettledRef.current = onRoundSettled;
+
+  const send = useCallback(
+    async (text: string) => {
+      const trimmed = text.trim();
+      if (!trimmed) return;
+      const id = newMessageId();
+      await messageStart({
+        threadId: thread.id,
+        role: "user",
+        content: trimmed,
+        model: modelRef.current,
+        id,
+      });
+      await messageFinish(id, "done", null, null);
+      if (messages.length === 0) {
+        const t = trimmed.slice(0, 30);
+        await threadRename(thread.id, trimmed.length > 30 ? t + "…" : t).catch(() => undefined);
+      }
+      await sendMessage({ id, role: "user", parts: [{ type: "text", text: trimmed }] });
+      onThreadChanged();
+    },
+    [sendMessage, thread.id, onThreadChanged, messages.length],
+  );
+
+  const regenerate = useCallback(async () => {
+    await chatRegenerate();
+    onThreadChanged();
+  }, [chatRegenerate, onThreadChanged]);
+
+  const editResend = useCallback(
+    async (messageId: string, newText: string) => {
+      const idx = messages.findIndex((m) => m.id === messageId);
+      if (idx < 0) return;
+      await messagesDeleteFrom(thread.id, messageId);
+      setMessages(messages.slice(0, idx));
+      await send(newText);
+    },
+    [messages, setMessages, send, thread.id],
+  );
+
+  const loadEarlier = useCallback(async (): Promise<boolean> => {
+    const before = oldestRef.current;
+    if (!before) return false;
+    const rows = await messagesPage({ threadId: thread.id, beforeCreatedAt: before.createdAt, beforeId: before.id, limit: 50 });
+    if (rows.length === 0) {
+      setHasMoreEarlier(false);
+      return false;
+    }
+    oldestRef.current = { createdAt: rows[0].createdAt, id: rows[0].id };
+    const older = rows.map(rowToUiMessage);
+    setMessages((cur) => [...older, ...cur]);
+    setHasMoreEarlier(rows.length === 50);
+    return true;
+  }, [setMessages, thread.id]);
+
+  return {
+    messages,
+    status,
+    send,
+    stop,
+    regenerate,
+    editResend,
+    loadEarlier,
+    hasMoreEarlier,
+    historyLoaded,
+    busy: status === "streaming" || status === "submitted",
+  };
+}
diff --git a/macos/app/src/components/AppShell.tsx b/macos/app/src/components/AppShell.tsx
new file mode 100644
index 0000000..fb9a02d
--- /dev/null
+++ b/macos/app/src/components/AppShell.tsx
@@ -0,0 +1,268 @@
+import { Link, Outlet, useRouter } from "@tanstack/react-router";
+import { BookOpen, Home, Menu, Plus, Server, Settings, Sparkles, X } from "lucide-react";
+import { useEffect, useState, useSyncExternalStore } from "react";
+import { useTranslation } from "react-i18next";
+
+import { daemonStore } from "../bridge/state";
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { threadCreate, threadDelete, threadRename } from "../chat/ipc";
+import { bucketThreads } from "../chat/buckets";
+import { threadStore } from "../state/threads";
+import { autoLoad } from "../models/autoLoad";
+import { refreshSystem } from "../models/system";
+import { SidebarThreadItem } from "./SidebarThreadItem";
+
+export function AppShell() {
+  const { t } = useTranslation();
+  const router = useRouter();
+  const [sidebarOpen, setSidebarOpen] = useState(false);
+  const threadsState = useSyncExternalStore(
+    (cb) => threadStore.subscribe(cb),
+    () => threadStore.threads,
+  );
+  const activeId = useSyncExternalStore(
+    (cb) => threadStore.subscribe(cb),
+    () => threadStore.activeId,
+  );
+  const snapTick = useSyncExternalStore(
+    (cb) => daemonStore.subscribe(cb),
+    () => daemonStore.snapshot,
+  );
+
+  useEffect(() => {
+    void threadStore.refresh();
+  }, []);
+
+  useEffect(() => {
+    const check = () => {
+      const cur = autoLoad.current;
+      if (!cur) return;
+      if (cur.phase === "download") {
+        const comp = daemonStore.latestFolded("download.completed", cur.taskId);
+        const fail = daemonStore.latestFolded("download.failed", cur.taskId);
+        if (fail && (!comp || fail.seq >= comp.seq)) {
+          autoLoad.cancel();
+          return;
+        }
+        if (comp) {
+          const tier = autoLoad.claimCompleted(comp.subject);
+          if (tier) {
+            void (async () => {
+              await refreshSystem();
+              try {
+                const res = await edge0Fetch(`/v1/edge0/models/${tier}/load`, {
+                  method: "POST",
+                  body: {},
+                });
+                if (!res.ok) autoLoad.cancel();
+              } catch {
+                autoLoad.cancel();
+              }
+            })();
+          }
+        }
+      } else {
+        const ready = daemonStore.latestFolded("model.load.ready", cur.tier);
+        const fail = daemonStore.latestFolded("model.load.failed", cur.tier);
+        if (fail) {
+          autoLoad.cancel();
+          return;
+        }
+        if (ready && autoLoad.consume(ready.subject)) {
+          void (async () => {
+            await refreshSystem();
+            const row = await threadCreate(t("chat.newThread"), cur.tier);
+            await threadStore.refresh();
+            threadStore.setActive(row.id);
+            void router.navigate({ to: "/" });
+          })();
+        }
+      }
+    };
+    const off1 = daemonStore.subscribe(check);
+    const off2 = autoLoad.subscribe(check);
+    check();
+    return () => {
+      off1();
+      off2();
+    };
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
+  const snap = snapTick;
+  const buckets = bucketThreads(threadsState, Date.now());
+  const resident = snap?.models.find((m) => m.state === "resident");
+  const serviceLine = snap
+    ? `:${snap.daemon.port} - ${resident ? `${t("sidebar.resident")} ${resident.id}` : t("sidebar.idle")}`
+    : t("sidebar.noService");
+
+  async function newChat() {
+    const defaultModel = (resident ?? snap?.models[0])?.id ?? null;
+    const row = await threadCreate(t("chat.newThread"), defaultModel);
+    await threadStore.refresh();
+    threadStore.setActive(row.id);
+    void router.navigate({ to: "/" });
+  }
+
+  function goHome() {
+    threadStore.setActive(null);
+    setSidebarOpen(false);
+    void router.navigate({ to: "/" });
+  }
+
+  return (
+    <div className="e0-app-shell text-fg">
+      {sidebarOpen && (
+        <button
+          type="button"
+          aria-label={t("nav.closeSidebar")}
+          className="e0-sidebar-backdrop"
+          onClick={() => setSidebarOpen(false)}
+        />
+      )}
+      <aside
+        data-testid="app-sidebar"
+        className={"e0-sidebar " + (sidebarOpen ? "e0-sidebar-open" : "")}
+        data-open={sidebarOpen ? "true" : "false"}
+      >
+        <div className="flex items-center justify-between px-4 pb-2 pt-4">
+          <div className="flex items-center gap-2">
+            <span className="grid size-8 place-items-center rounded-xl bg-fg text-surface shadow-sm">
+              <Sparkles size={16} aria-hidden />
+            </span>
+            <div>
+              <span className="block text-sm font-bold tracking-tight">{t("app.name")}</span>
+              <span className="block text-[10px] font-medium uppercase tracking-[0.16em] text-muted">{t("app.localLabel")}</span>
+            </div>
+          </div>
+          <button
+            type="button"
+            aria-label={t("nav.closeSidebar")}
+            className="e0-icon-btn md:hidden"
+            onClick={() => setSidebarOpen(false)}
+          >
+            <X size={16} aria-hidden />
+          </button>
+        </div>
+        <div className="px-3 pb-3">
+          <button
+            type="button"
+            onClick={() => void newChat()}
+            className="e0-btn e0-btn-primary w-full"
+          >
+            <Plus size={16} aria-hidden />
+            {t("sidebar.newChat")}
+          </button>
+        </div>
+        <nav aria-label={t("sidebar.navigation")} className="space-y-1 px-2 pb-2">
+          <button
+            type="button"
+            aria-label={t("nav.home")}
+            data-testid="nav-home"
+            onClick={goHome}
+            className={
+              "flex w-full items-center gap-2 rounded-lg px-3 py-2 text-sm hover:bg-surface2 hover:text-fg " +
+              (activeId === null && router.state.location.pathname === "/" ? "bg-surface2 text-fg" : "text-muted")
+            }
+          >
+            <Home size={16} aria-hidden />
+            {t("nav.home")}
+          </button>
+          <Link
+            to="/models"
+            aria-label={t("nav.models")}
+            data-testid="nav-models"
+            activeProps={{ className: "bg-surface2 text-fg" }}
+            className="flex items-center gap-2 rounded-lg px-3 py-2 text-sm text-muted hover:bg-surface2 hover:text-fg"
+            onClick={() => setSidebarOpen(false)}
+          >
+            <BookOpen size={16} aria-hidden />
+            {t("nav.models")}
+          </Link>
+          <Link
+            to="/service"
+            aria-label={t("nav.service")}
+            data-testid="nav-service"
+            activeProps={{ className: "bg-surface2 text-fg" }}
+            className="flex items-center gap-2 rounded-lg px-3 py-2 text-sm text-muted hover:bg-surface2 hover:text-fg"
+            onClick={() => setSidebarOpen(false)}
+          >
+            <Server size={16} aria-hidden />
+            {t("nav.service")}
+          </Link>
+        </nav>
+        <nav aria-label={t("sidebar.history")} className="min-h-0 flex-1 overflow-y-auto px-2">
+          <h2 className="px-3 pb-2 pt-2 text-[10px] font-bold uppercase tracking-[0.16em] text-muted">
+            {t("sidebar.history")}
+          </h2>
+          {buckets.map((b) => (
+            <section key={b.key} className="mb-2">
+              <h2 className="px-2 pb-1 pt-2 text-[11px] font-semibold uppercase tracking-wide text-muted">
+                {t(`sidebar.bucket.${b.key}`)}
+              </h2>
+              <ul className="space-y-0.5">
+                {b.items.map((th) => (
+                  <SidebarThreadItem
+                    key={th.id}
+                    thread={th}
+                    active={th.id === activeId}
+                    onSelect={() => {
+                      threadStore.setActive(th.id);
+                      setSidebarOpen(false);
+                      void router.navigate({ to: "/" });
+                    }}
+                    onRename={(title) => void threadRename(th.id, title).then(() => threadStore.refresh())}
+                    onDelete={() =>
+                      void threadDelete(th.id).then(() => {
+                        if (activeId === th.id) threadStore.setActive(null);
+                        void threadStore.refresh();
+                      })
+                    }
+                  />
+                ))}
+              </ul>
+            </section>
+          ))}
+          {threadsState.length === 0 && <p className="px-3 py-3 text-xs leading-relaxed text-muted">{t("chat.noThreads")}</p>}
+        </nav>
+        <div className="e0-divider mx-3" />
+        <div className="flex items-center gap-2 px-3 py-3 text-xs text-muted">
+          <span
+            aria-label={serviceLine}
+            className={
+              "inline-block size-2 shrink-0 rounded-full " +
+              (snap ? "bg-ok" : "bg-muted")
+            }
+          />
+          <span className="min-w-0 flex-1 truncate" data-testid="service-summary">
+            {serviceLine}
+          </span>
+          <Link
+            to="/settings"
+            aria-label={t("nav.settings")}
+            className="e0-icon-btn ml-auto"
+            onClick={() => setSidebarOpen(false)}
+          >
+            <Settings size={16} aria-hidden />
+          </Link>
+        </div>
+      </aside>
+      <main data-testid="app-main" className="e0-main-frame">
+        <div className="e0-mobile-toggle absolute left-4 top-4 z-10">
+          <button
+            type="button"
+            aria-label={t("nav.openSidebar")}
+            title={t("nav.openSidebar")}
+            className="e0-icon-btn border border-line bg-surface shadow-sm"
+            onClick={() => setSidebarOpen(true)}
+          >
+            <Menu size={17} aria-hidden />
+          </button>
+        </div>
+        <div className="h-full min-h-0 overflow-y-auto">
+          <Outlet />
+        </div>
+      </main>
+    </div>
+  );
+}
diff --git a/macos/app/src/components/ShellControls.tsx b/macos/app/src/components/ShellControls.tsx
new file mode 100644
index 0000000..7a4dd30
--- /dev/null
+++ b/macos/app/src/components/ShellControls.tsx
@@ -0,0 +1,67 @@
+import * as RadixToggleGroup from "@radix-ui/react-toggle-group";
+import { useEffect, useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import {
+  applyFont,
+  applyTheme,
+  storedFont,
+  storedTheme,
+  type FontScale,
+  type Theme,
+} from "../lib/theme";
+
+const ToggleGroup = RadixToggleGroup.Root;
+const Toggle = RadixToggleGroup.Item;
+
+export function ShellControls() {
+  const { t } = useTranslation();
+  const [theme, setTheme] = useState<Theme>(storedTheme());
+  const [font, setFont] = useState<FontScale>(storedFont());
+
+  useEffect(() => {
+    applyTheme(theme);
+  }, [theme]);
+  useEffect(() => {
+    applyFont(font);
+  }, [font]);
+  useEffect(() => {
+    // Re-resolve when the OS theme changes while the app is set to "system".
+    const mq = window.matchMedia("(prefers-color-scheme: dark)");
+    const onChange = () => theme === "system" && applyTheme("system");
+    mq.addEventListener("change", onChange);
+    return () => mq.removeEventListener("change", onChange);
+  }, [theme]);
+
+  const seg =
+    "e0-toggle";
+  return (
+    <div className="flex flex-wrap items-center gap-3 border-b border-line pb-3" data-testid="shell-controls">
+      <ToggleGroup
+        type="single"
+        value={theme}
+        onValueChange={(v) => v && setTheme(v as Theme)}
+        className="e0-toggle-group"
+        aria-label={t("status.title")}
+      >
+        {(["light", "dark", "system"] as const).map((v) => (
+          <Toggle key={v} value={v} className={seg}>
+            {t(`theme.${v}`)}
+          </Toggle>
+        ))}
+      </ToggleGroup>
+      <ToggleGroup
+        type="single"
+        value={font}
+        onValueChange={(v) => v && setFont(v as FontScale)}
+        className="e0-toggle-group"
+      >
+        {(["sm", "md", "lg"] as const).map((v) => (
+          <Toggle key={v} value={v} className={seg}>
+            {t(`font.${v}`)}
+          </Toggle>
+        ))}
+      </ToggleGroup>
+    </div>
+  );
+}
diff --git a/macos/app/src/components/SidebarThreadItem.tsx b/macos/app/src/components/SidebarThreadItem.tsx
new file mode 100644
index 0000000..b7945ad
--- /dev/null
+++ b/macos/app/src/components/SidebarThreadItem.tsx
@@ -0,0 +1,85 @@
+import { Pencil, Trash2 } from "lucide-react";
+import { useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import type { ThreadRow } from "../chat/ipc";
+
+export function SidebarThreadItem({
+  thread,
+  active,
+  onSelect,
+  onRename,
+  onDelete,
+}: {
+  thread: ThreadRow;
+  active: boolean;
+  onSelect: () => void;
+  onRename: (title: string) => void;
+  onDelete: () => void;
+}) {
+  const { t } = useTranslation();
+  const [editing, setEditing] = useState(false);
+  const [draft, setDraft] = useState("");
+
+  if (editing) {
+    return (
+      <li>
+        <form
+          onSubmit={(e) => {
+            e.preventDefault();
+            onRename(draft.trim() || thread.data.title);
+            setEditing(false);
+          }}
+        >
+          <input
+            autoFocus
+            value={draft}
+            onChange={(e) => setDraft(e.target.value)}
+            onBlur={() => setEditing(false)}
+            className="w-full rounded-lg border border-accent bg-surface px-3 py-2 text-sm outline-none"
+          />
+        </form>
+      </li>
+    );
+  }
+
+  return (
+    <li
+      className="group flex items-center rounded-lg transition-colors hover:bg-surface2/80"
+      data-active={active || undefined}
+    >
+      <button
+        type="button"
+        onClick={onSelect}
+        className={
+          "min-w-0 flex-1 truncate rounded-md px-2.5 py-1.5 text-left text-sm " +
+          (active ? "bg-surface2 font-semibold text-fg shadow-sm" : "text-fg/85")
+        }
+        title={thread.data.title}
+      >
+        {thread.data.title}
+      </button>
+      <span className="mr-1 hidden shrink-0 items-center gap-0.5 group-hover:flex group-focus-within:flex">
+        <button
+          type="button"
+          aria-label={t("chat.rename")}
+          onClick={() => {
+            setEditing(true);
+            setDraft(thread.data.title);
+          }}
+          className="e0-icon-btn !size-7 !rounded-md"
+        >
+          <Pencil size={13} />
+        </button>
+        <button
+          type="button"
+          aria-label={t("chat.deleteThread")}
+          onClick={onDelete}
+          className="e0-icon-btn !size-7 !rounded-md hover:!text-err"
+        >
+          <Trash2 size={13} />
+        </button>
+      </span>
+    </li>
+  );
+}
diff --git a/macos/app/src/components/chat/ComposerBar.tsx b/macos/app/src/components/chat/ComposerBar.tsx
new file mode 100644
index 0000000..59c50eb
--- /dev/null
+++ b/macos/app/src/components/chat/ComposerBar.tsx
@@ -0,0 +1,88 @@
+import { ArrowUp, Square } from "lucide-react";
+import { useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import type { ThreadChatApi } from "../../chat/useThreadChat";
+import { GenPill } from "./GenPill";
+import { ModelPill } from "./ModelPill";
+
+export function ComposerBar({
+  chat,
+  model,
+  onModelChange,
+  wireJson,
+  meterNonce,
+}: {
+  chat: ThreadChatApi;
+  model: string;
+  onModelChange: (id: string) => void;
+  wireJson: () => { role: string; content: string }[];
+  meterNonce: number;
+}) {
+  const { t } = useTranslation();
+  const [text, setText] = useState("");
+  const [sending, setSending] = useState(false);
+
+  async function submit() {
+    const trimmed = text.trim();
+    if (!trimmed || sending) return;
+    setSending(true);
+    setText("");
+    try {
+      await chat.send(trimmed);
+    } finally {
+      setSending(false);
+    }
+  }
+
+  return (
+    <div className="e0-composer" data-testid="composer">
+      <div className="e0-composer-shell">
+        <textarea
+          value={text}
+          onChange={(e) => setText(e.target.value)}
+          onKeyDown={(e) => {
+            if (e.key === "Enter" && !e.shiftKey) {
+              e.preventDefault();
+              void submit();
+            }
+          }}
+          rows={2}
+          placeholder={t("chat.placeholder")}
+          className="e0-composer-textarea w-full resize-none bg-transparent text-sm outline-none placeholder:text-muted"
+        />
+        <div className="e0-composer-toolbar">
+          <ModelPill model={model} onModelChange={onModelChange} wireJson={wireJson} meterNonce={meterNonce} />
+          <GenPill />
+          <span className="min-w-0 flex-1 truncate text-xs text-muted">{!model ? t("chat.pickModelFirst") : ""}</span>
+          {chat.busy ? (
+            <button
+              type="button"
+              aria-label={t("chat.stop")}
+              title={t("chat.stop")}
+              onClick={chat.stop}
+              className="e0-send-btn bg-fg text-surface"
+            >
+              <Square size={14} fill="currentColor" />
+            </button>
+          ) : (
+            <button
+              type="button"
+              aria-label={t("chat.send")}
+              title={t("chat.send")}
+              onClick={() => void submit()}
+              disabled={disabledNow()}
+              className="e0-send-btn disabled:opacity-40"
+            >
+              <ArrowUp size={16} />
+            </button>
+          )}
+        </div>
+      </div>
+    </div>
+  );
+
+  function disabledNow(): boolean {
+    return !model || !text.trim();
+  }
+}
diff --git a/macos/app/src/components/chat/GenPill.tsx b/macos/app/src/components/chat/GenPill.tsx
new file mode 100644
index 0000000..4de0d61
--- /dev/null
+++ b/macos/app/src/components/chat/GenPill.tsx
@@ -0,0 +1,36 @@
+import { useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import { GEN_TEMPERATURE, setGen, storedGen, type GenPreset } from "../../lib/generation";
+
+const PRESETS: GenPreset[] = ["precise", "balanced", "creative"];
+
+export function GenPill() {
+  const { t } = useTranslation();
+  const [gen, setGenState] = useState<GenPreset>(() => storedGen());
+  return (
+    <div
+      className="flex shrink-0 items-center gap-0.5 rounded-full border border-line bg-surface2 px-1 py-0.5"
+      data-testid="gen-pill"
+      title={t("gen.hint")}
+    >
+      {PRESETS.map((p) => (
+        <button
+          key={p}
+          type="button"
+          data-testid={`gen-${p}`}
+          data-state={gen === p ? "on" : "off"}
+          aria-pressed={gen === p}
+          onClick={() => {
+            setGen(p);
+            setGenState(p);
+          }}
+          className="rounded-full px-2 py-0.5 text-[11px] text-muted transition-colors hover:text-fg data-[state=on]:bg-accent data-[state=on]:text-accentink"
+          title={`${t(`gen.${p}`)} - ${GEN_TEMPERATURE[p]}`}
+        >
+          {t(`gen.${p}`)}
+        </button>
+      ))}
+    </div>
+  );
+}
diff --git a/macos/app/src/components/chat/MessageItem.tsx b/macos/app/src/components/chat/MessageItem.tsx
new file mode 100644
index 0000000..566fc68
--- /dev/null
+++ b/macos/app/src/components/chat/MessageItem.tsx
@@ -0,0 +1,213 @@
+import { Streamdown } from "streamdown";
+import { createMathPlugin } from "@streamdown/math";
+// Models commonly write `$d$`; leaving single-dollar math off skips those. Currency like "$5 and $10" can false-positive.
+const mathPlugin = createMathPlugin({ singleDollarTextMath: true });
+import { Brain, Check, Copy, Pencil, RotateCcw } from "lucide-react";
+import { useTranslation } from "react-i18next";
+import { useState } from "react";
+
+import {
+  messageReasoning,
+  messageText,
+  type ChatMessageMeta,
+  type ChatUIMessage,
+} from "../../chat/records";
+import { errorNextKey } from "../../chat/errorMap";
+
+export function MessageItem({
+  message,
+  streaming,
+  isLastAssistant,
+  busy,
+  onRegenerate,
+  onEditResend,
+}: {
+  message: ChatUIMessage;
+  streaming: boolean;
+  isLastAssistant: boolean;
+  busy: boolean;
+  onRegenerate: () => void;
+  onEditResend: (newText: string) => void;
+}) {
+  const { t } = useTranslation();
+  const [editing, setEditing] = useState(false);
+  const [draft, setDraft] = useState("");
+  const [copied, setCopied] = useState<"text" | "reasoning" | null>(null);
+  const isUser = message.role === "user";
+  const text = messageText(message);
+  const reasoning = messageReasoning(message);
+  const meta = message.metadata ?? {};
+  const live = streaming && isLastAssistant && !isUser;
+  const isAborted = meta.status === "aborted" || (live && false);
+  const isError = meta.status === "error";
+
+  async function copy(value: string, which: "text" | "reasoning") {
+    try {
+      await navigator.clipboard.writeText(value);
+      setCopied(which);
+      setTimeout(() => setCopied(null), 1200);
+    } catch {
+      /* clipboard may be unavailable */
+    }
+  }
+
+  if (isUser) {
+    return (
+    <article data-msg-role="user" className="group/msg flex flex-col items-end py-2.5">
+        {editing ? (
+          <form
+            onSubmit={(e) => {
+              e.preventDefault();
+              setEditing(false);
+              if (draft.trim()) onEditResend(draft);
+            }}
+            className="w-full max-w-[88%]"
+          >
+            <textarea
+              autoFocus
+              value={draft}
+              onChange={(e) => setDraft(e.target.value)}
+              rows={3}
+              className="w-full resize-y rounded-2xl border border-accent/50 bg-surface2 p-3 text-sm outline-none"
+            />
+            <div className="flex justify-end gap-2 pt-1.5">
+              <button type="submit" className="e0-btn e0-btn-primary">
+                {t("chat.resend")}
+              </button>
+              <button type="button" onClick={() => setEditing(false)} className="e0-btn e0-btn-ghost">
+                {t("common.cancel")}
+              </button>
+            </div>
+          </form>
+        ) : (
+          <>
+            <div className="max-w-[88%] whitespace-pre-wrap break-words rounded-2xl rounded-br-md bg-accent px-4 py-3 text-sm leading-relaxed text-accentink shadow-sm">
+              {text}
+            </div>
+            <div className="mt-0.5 flex h-6 gap-1 opacity-0 transition-opacity group-hover/msg:opacity-100 focus-within:opacity-100 [@media(hover:none)]:opacity-100">
+              {text && (
+                <IconAction label={t("chat.copy")} onClick={() => void copy(text, "text")}>
+                  {copied === "text" ? <Check size={13} /> : <Copy size={13} />}
+                </IconAction>
+              )}
+              {!busy && (
+                <IconAction
+                  label={t("chat.editResend")}
+                  onClick={() => {
+                    setEditing(true);
+                    setDraft(text);
+                  }}
+                >
+                  <Pencil size={13} />
+                </IconAction>
+              )}
+            </div>
+          </>
+        )}
+      </article>
+    );
+  }
+
+  return (
+    <article data-msg-role="assistant" className="group/msg flex gap-3 border-t border-line/40 py-5 first:border-t-0">
+      <Avatar />
+      <div className="min-w-0 flex-1">
+        {reasoning && (
+          <details className="mb-3 max-w-full rounded-xl border border-line bg-surface2/70 px-3 py-2" data-testid="reasoning-block">
+            <summary className="flex cursor-pointer select-none items-center gap-1.5 text-xs font-medium text-muted">
+              <Brain size={13} aria-hidden />
+              {t("chat.reasoning")}
+            </summary>
+            <div className="mt-1.5 grid gap-1.5">
+              <div className="whitespace-pre-wrap text-xs leading-relaxed text-muted">{reasoning}</div>
+              <button type="button" onClick={() => void copy(reasoning, "reasoning")} className="flex w-fit items-center gap-1 text-xs text-muted hover:text-fg">
+                {copied === "reasoning" ? <Check size={12} /> : <Copy size={12} />}
+                {t("chat.copyReasoning")}
+              </button>
+            </div>
+          </details>
+        )}
+
+        <div
+          className="prose prose-sm dark:prose-invert max-w-none break-words text-sm [&>*:first-child]:mt-0 [&>*:last-child]:mb-0"
+          data-testid="message-body"
+        >
+          <Streamdown mode={live ? "streaming" : "static"} plugins={{ math: mathPlugin }}>{text}</Streamdown>
+          {live && !text && <span className="inline-block h-4 w-1 animate-pulse bg-fg align-middle" />}
+        </div>
+
+        {isError && meta.code && (
+          <p className="mt-1.5 flex items-center gap-2 rounded-lg bg-err/10 px-2.5 py-1.5 text-xs text-err" role="alert">
+            <span>{t(errorNextKey(meta.code))}</span>
+            <code className="rounded bg-err/10 px-1">{meta.code}</code>
+          </p>
+        )}
+        {isAborted && <p className="mt-1 text-xs text-warn">{t("chat.aborted")}</p>}
+        {meta.saveFailed && <p className="mt-1 text-xs text-warn">{t("chat.saveFailed")}</p>}
+
+        
+        {!live && <PerfLine perf={meta.perf} />}
+
+        <div className="mt-0.5 flex h-6 items-center gap-1 opacity-0 transition-opacity group-hover/msg:opacity-100 focus-within:opacity-100 [@media(hover:none)]:opacity-100">
+          {text && (
+            <IconAction label={t("chat.copy")} onClick={() => void copy(text, "text")}>
+              {copied === "text" ? <Check size={13} /> : <Copy size={13} />}
+            </IconAction>
+          )}
+          {isLastAssistant && !busy && (
+            <IconAction label={t("chat.regenerate")} onClick={onRegenerate}>
+              <RotateCcw size={13} />
+            </IconAction>
+          )}
+          {meta.usage && (
+            <span className="ml-auto text-[11px] tabular-nums text-muted" data-testid="usage-line">
+              {t("chat.usage", { p: meta.usage.prompt_tokens, c: meta.usage.completion_tokens, t: meta.usage.total_tokens })}
+            </span>
+          )}
+        </div>
+      </div>
+    </article>
+  );
+}
+
+function IconAction({ label, onClick, children }: { label: string; onClick: () => void; children: React.ReactNode }) {
+  return (
+    <button
+      type="button"
+      aria-label={label}
+      title={label}
+      onClick={onClick}
+      className="e0-icon-btn !size-7 !rounded-md"
+    >
+      {children}
+    </button>
+  );
+}
+
+function PerfLine({ perf }: { perf: ChatMessageMeta["perf"] }) {
+  const { t } = useTranslation();
+  if (!perf || !Object.values(perf).some((v) => typeof v === "number")) return null;
+  return (
+    <p className="mt-1 flex flex-wrap gap-x-2 text-[11px] tabular-nums text-muted" data-testid="perf-line">
+      {typeof perf.decode_tps === "number" && <span>{t("chat.perf.decode", { v: perf.decode_tps })}</span>}
+      {typeof perf.prefill_tps === "number" && <span>{t("chat.perf.prefill", { v: perf.prefill_tps })}</span>}
+      {typeof perf.ttft_ms === "number" && <span>{t("chat.perf.ttft", { ms: perf.ttft_ms })}</span>}
+      {typeof perf.footprint_mb === "number" && (
+        <span>{t("chat.perf.mem", { gi: (perf.footprint_mb / 1024).toFixed(2) })}</span>
+      )}
+    </p>
+  );
+}
+
+function Avatar() {
+  return (
+    <span
+      aria-hidden
+      className="mt-0.5 grid h-7 w-7 shrink-0 place-content-center rounded-full border border-line bg-surface2 text-fg"
+    >
+      <svg width="14" height="14" viewBox="0 0 14 14" fill="none">
+        <path d="M7 1l1.6 4.4L13 7l-4.4 1.6L7 13 5.4 8.6 1 7l4.4-1.6L7 1z" fill="currentColor" />
+      </svg>
+    </span>
+  );
+}
diff --git a/macos/app/src/components/chat/MessageList.tsx b/macos/app/src/components/chat/MessageList.tsx
new file mode 100644
index 0000000..561312a
--- /dev/null
+++ b/macos/app/src/components/chat/MessageList.tsx
@@ -0,0 +1,33 @@
+import { useTranslation } from "react-i18next";
+import { useMemo } from "react";
+
+import type { ThreadChatApi } from "../../chat/useThreadChat";
+import { MessageItem } from "./MessageItem";
+
+export function MessageList({ chat }: { chat: ThreadChatApi }) {
+  const { t } = useTranslation();
+  const lastAssistantIdx = useMemo(
+    () => chat.messages.map((m) => m.role).lastIndexOf("assistant"),
+    [chat.messages],
+  );
+  return (
+    <div data-testid="message-list">
+      {chat.messages.map((m, i) => (
+        <MessageItem
+          key={m.id ?? i}
+          message={m}
+          streaming={chat.busy}
+          isLastAssistant={i === lastAssistantIdx}
+          busy={chat.busy}
+          onRegenerate={() => void chat.regenerate()}
+          onEditResend={(next) => void chat.editResend(m.id, next)}
+        />
+      ))}
+      {chat.busy && chat.messages[chat.messages.length - 1]?.role !== "assistant" && (
+        <p className="px-3 text-xs text-muted" data-testid="waiting-first-chunk">
+          {t("chat.waitingFirstChunk")}
+        </p>
+      )}
+    </div>
+  );
+}
diff --git a/macos/app/src/components/chat/ModelPill.tsx b/macos/app/src/components/chat/ModelPill.tsx
new file mode 100644
index 0000000..3a107db
--- /dev/null
+++ b/macos/app/src/components/chat/ModelPill.tsx
@@ -0,0 +1,164 @@
+import { useEffect, useRef, useState } from "react";
+import { ChevronDown, Loader2 } from "lucide-react";
+import { useTranslation } from "react-i18next";
+
+import { daemonStore } from "../../bridge/state";
+import { edge0Fetch } from "../../bridge/edge0Fetch";
+import { ApiError } from "../../bridge/frames";
+import { candidateModels, deriveLight, latestModelEvent, type ModelLight } from "../../chat/modelStatus";
+import { errorNextKey } from "../../chat/errorMap";
+
+export function ModelPill({
+  model,
+  onModelChange,
+  wireJson,
+  meterNonce,
+}: {
+  model: string;
+  onModelChange: (id: string) => void;
+  wireJson: () => { role: string; content: string }[];
+  meterNonce: number;
+}) {
+  const { t } = useTranslation();
+  const [, force] = useState(0);
+  const [open, setOpen] = useState(false);
+  const [loadError, setLoadError] = useState<{ code: string | null; message: string } | null>(null);
+  const [contextTokens, setContextTokens] = useState<number | null>(null);
+  const debounceRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const wireRef = useRef(wireJson);
+  wireRef.current = wireJson;
+  const rootRef = useRef<HTMLDivElement>(null);
+
+  useEffect(() => daemonStore.subscribe(() => force((n) => n + 1)), []);
+
+  useEffect(() => {
+    if (!open) return;
+    const onDoc = (e: MouseEvent) => {
+      if (!rootRef.current?.contains(e.target as Node)) setOpen(false);
+    };
+    document.addEventListener("mousedown", onDoc);
+    return () => document.removeEventListener("mousedown", onDoc);
+  }, [open]);
+
+  const snap = daemonStore.snapshot;
+  const candidates = candidateModels(snap);
+  const light = model ? deriveLight(snap, model, latestModelEvent(daemonStore, model)) : "absent";
+
+  useEffect(() => {
+    if (debounceRef.current) clearTimeout(debounceRef.current);
+    debounceRef.current = setTimeout(() => void measure(), 400);
+    return () => {
+      if (debounceRef.current) clearTimeout(debounceRef.current);
+    };
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [meterNonce, model, snap?.daemon.port]);
+
+  async function measure() {
+    if (!model || snap === null) {
+      setContextTokens(null);
+      return;
+    }
+    try {
+      const res = await edge0Fetch("/v1/edge0/apply-template", {
+        method: "POST",
+        body: { model, messages: wireRef.current(), add_generation_prompt: true },
+      });
+      if (!res.ok) throw new ApiError(res.status, null, `HTTP ${res.status}`);
+      const j = (await res.json()) as { token_count?: number };
+      setContextTokens(typeof j.token_count === "number" ? j.token_count : null);
+    } catch {
+      setContextTokens(null);
+    }
+  }
+
+  async function load(id: string) {
+    setLoadError(null);
+    try {
+      const res = await edge0Fetch(`/v1/edge0/models/${id}/load`, { method: "POST", body: {} });
+      if (!res.ok) throw await apiErrFrom(res);
+    } catch (e) {
+      const ae = e instanceof ApiError ? e : null;
+      setLoadError({ code: ae?.code ?? null, message: ae?.message ?? String(e) });
+    }
+  }
+
+  return (
+    <div ref={rootRef} className="relative" data-testid="model-pill">
+      <button
+        type="button"
+        aria-haspopup="listbox"
+        aria-expanded={open}
+        onClick={() => setOpen((o) => !o)}
+        className="flex min-w-0 max-w-[18rem] items-center gap-1.5 rounded-full border border-line bg-surface2 px-2.5 py-1.5 text-xs font-medium text-fg hover:border-accent/50 hover:bg-surface"
+      >
+        <LightDot light={light} />
+        <span className="max-w-[10rem] truncate">{model || t("chat.noModel")}</span>
+        <span className="text-muted" data-testid="context-meter">
+          {contextTokens === null ? t("chat.contextUnknown") : t("chat.contextUsed", { n: contextTokens })}
+        </span>
+        <ChevronDown size={12} aria-hidden />
+      </button>
+
+      {open && (
+        <div
+          role="listbox"
+          className="absolute bottom-full left-0 z-20 mb-2 w-72 rounded-2xl border border-line bg-surface p-2 shadow-lg"
+        >
+          {candidates.length === 0 && <p className="px-2 py-1.5 text-sm text-muted">{t("chat.noModel")}</p>}
+          {candidates.map((id) => {
+            const l = deriveLight(snap, id, latestModelEvent(daemonStore, id));
+            return (
+              <div key={id} className="flex items-center gap-2 rounded-lg px-2 py-1.5 text-sm hover:bg-surface2">
+                <button
+                  type="button"
+                  role="option"
+                  aria-selected={id === model}
+                  onClick={() => {
+                    onModelChange(id);
+                    setOpen(false);
+                  }}
+                  className="flex min-w-0 flex-1 items-center gap-2 text-left"
+                >
+                  <LightDot light={l} />
+                  <span className="min-w-0 flex-1 truncate">{id}</span>
+                  <span className="text-xs text-muted">{t(`chat.light.${l}`)}</span>
+                </button>
+                {(l === "installed" || l === "loading") && id === model && (
+                  <button
+                    type="button"
+                    onClick={() => void load(id)}
+                    disabled={l === "loading"}
+                    className="flex items-center gap-1 rounded-full border border-line px-2 py-0.5 text-xs text-muted hover:text-fg disabled:opacity-60"
+                  >
+                    {l === "loading" && <Loader2 size={11} className="animate-spin" aria-hidden />}
+                    {t("chat.loadModel")}
+                  </button>
+                )}
+              </div>
+            );
+          })}
+          {loadError && (
+            <p className="mt-1 rounded-lg bg-err/10 px-2 py-1.5 text-xs text-err" role="alert">
+              {t(errorNextKey(loadError.code))} {loadError.code ? ` -  ${loadError.code}` : ""}
+            </p>
+          )}
+        </div>
+      )}
+    </div>
+  );
+}
+
+function LightDot({ light }: { light: ModelLight }) {
+  const cls =
+    light === "resident" ? "bg-ok" : light === "loading" ? "bg-warn animate-pulse" : light === "installed" ? "bg-muted" : "bg-line";
+  return <span data-testid="model-light" data-light={light} className={"inline-block h-2 w-2 shrink-0 rounded-full " + cls} aria-hidden />;
+}
+
+async function apiErrFrom(res: Response): Promise<ApiError> {
+  try {
+    const j = (await res.json()) as { error?: { code?: string; message?: string } };
+    return new ApiError(res.status, j?.error?.code ?? null, j?.error?.message ?? `HTTP ${res.status}`);
+  } catch {
+    return new ApiError(res.status, null, `HTTP ${res.status}`);
+  }
+}
diff --git a/macos/app/src/gen/Attachment.ts b/macos/app/src/gen/Attachment.ts
new file mode 100644
index 0000000..23b21e5
--- /dev/null
+++ b/macos/app/src/gen/Attachment.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type Attachment = { "kind": "text", text: string, } | { "kind": "local_path", path: string, };
diff --git a/macos/app/src/gen/CatalogResponse.ts b/macos/app/src/gen/CatalogResponse.ts
new file mode 100644
index 0000000..71356b8
--- /dev/null
+++ b/macos/app/src/gen/CatalogResponse.ts
@@ -0,0 +1,9 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { CatalogTier } from "./CatalogTier";
+
+export type CatalogResponse = { 
+generated_at: string, 
+expires_at: string, 
+fetched_at: string, 
+stale: boolean, 
+signer_key_id: string, tiers: Array<CatalogTier>, };
diff --git a/macos/app/src/gen/CatalogTier.ts b/macos/app/src/gen/CatalogTier.ts
new file mode 100644
index 0000000..d432902
--- /dev/null
+++ b/macos/app/src/gen/CatalogTier.ts
@@ -0,0 +1,8 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type CatalogTier = { tier: string, rev: string, min_cli_version: string, 
+bytes_total: number, files_count: number, license: string | null, notes_url: string | null, 
+primary_source: string, 
+has_modelscope: boolean, has_huggingface: boolean, 
+installed: string | null, 
+update_available: boolean, };
diff --git a/macos/app/src/gen/DaemonInfo.ts b/macos/app/src/gen/DaemonInfo.ts
new file mode 100644
index 0000000..e193f37
--- /dev/null
+++ b/macos/app/src/gen/DaemonInfo.ts
@@ -0,0 +1,7 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type DaemonInfo = { version: string, abi_version: number, 
+host_app_version: string | null, scope: string, started_by: string, 
+owner_pid: number | null, session_id: string, started_at: string, uptime_s: number, throttled: boolean, port: number, bound: string, 
+home: string, 
+engine: string, };
diff --git a/macos/app/src/gen/DownloadEntry.ts b/macos/app/src/gen/DownloadEntry.ts
new file mode 100644
index 0000000..da298bf
--- /dev/null
+++ b/macos/app/src/gen/DownloadEntry.ts
@@ -0,0 +1,7 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { DownloadFailure } from "./DownloadFailure";
+
+export type DownloadEntry = { id: string, tier: string, rev: string, phase: string, source: string | null, source_reason: string | null, bytes_done: number, bytes_total: number, paused: boolean, failed: DownloadFailure | null, 
+rate_bps?: number | null, 
+eta_s?: number | null, 
+source_switched?: boolean, };
diff --git a/macos/app/src/gen/DownloadFailure.ts b/macos/app/src/gen/DownloadFailure.ts
new file mode 100644
index 0000000..f110d79
--- /dev/null
+++ b/macos/app/src/gen/DownloadFailure.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type DownloadFailure = { code: string, message: string, };
diff --git a/macos/app/src/gen/Event.ts b/macos/app/src/gen/Event.ts
new file mode 100644
index 0000000..fbadfd0
--- /dev/null
+++ b/macos/app/src/gen/Event.ts
@@ -0,0 +1,4 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { JsonValue } from "./serde_json/JsonValue";
+
+export type Event = { seq: number, ts: string, type: string, subject: string, payload: JsonValue, };
diff --git a/macos/app/src/gen/EventType.ts b/macos/app/src/gen/EventType.ts
new file mode 100644
index 0000000..812781e
--- /dev/null
+++ b/macos/app/src/gen/EventType.ts
@@ -0,0 +1,16 @@
+export type EventType = "daemon.hello" |
+  "model.load.started" |
+  "model.load.ready" |
+  "model.load.failed" |
+  "model.unloaded" |
+  "download.progress" |
+  "download.paused" |
+  "download.completed" |
+  "download.failed" |
+  "worker.crashed" |
+  "request.rejected" |
+  "api.request.finished" |
+  "config.changed" |
+  "service.restart" |
+  "resync.required"; 
+export const FOLDED_EVENT_TYPES: EventType[] = ["download.progress", "request.rejected", "config.changed"];
diff --git a/macos/app/src/gen/HardwareInfo.ts b/macos/app/src/gen/HardwareInfo.ts
new file mode 100644
index 0000000..252adf6
--- /dev/null
+++ b/macos/app/src/gen/HardwareInfo.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type HardwareInfo = { chip: string, macos: string, mem_total_gib: number | null, metal_device: boolean, disk_free_gib: number, };
diff --git a/macos/app/src/gen/HelloPayload.ts b/macos/app/src/gen/HelloPayload.ts
new file mode 100644
index 0000000..e7557ce
--- /dev/null
+++ b/macos/app/src/gen/HelloPayload.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type HelloPayload = { session_id: string, seq: number, abi_version: number, };
diff --git a/macos/app/src/gen/MessageRecord.ts b/macos/app/src/gen/MessageRecord.ts
new file mode 100644
index 0000000..32efe6d
--- /dev/null
+++ b/macos/app/src/gen/MessageRecord.ts
@@ -0,0 +1,15 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { Attachment } from "./Attachment";
+import type { MessageStatus } from "./MessageStatus";
+import type { Role } from "./Role";
+import type { Usage } from "./Usage";
+import type { JsonValue } from "./serde_json/JsonValue";
+
+export type MessageRecord = { v: number, role: Role, content: string, 
+reasoning_content?: string | null, attachments?: Array<Attachment>, tool_calls?: Array<JsonValue> | null, usage?: Usage | null, status: MessageStatus, 
+code?: string | null, model: string, 
+params?: JsonValue | null, 
+/**
+ * epoch ms.
+ */
+created_at: number, thread_id?: string | null, } & ({ [key in string]: number | string | boolean | Array<JsonValue> | { [key in string]: JsonValue } | null });
diff --git a/macos/app/src/gen/MessageStatus.ts b/macos/app/src/gen/MessageStatus.ts
new file mode 100644
index 0000000..13f945a
--- /dev/null
+++ b/macos/app/src/gen/MessageStatus.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type MessageStatus = "streaming" | "done" | "aborted" | "error";
diff --git a/macos/app/src/gen/ModelEntry.ts b/macos/app/src/gen/ModelEntry.ts
new file mode 100644
index 0000000..1b3ca67
--- /dev/null
+++ b/macos/app/src/gen/ModelEntry.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type ModelEntry = { id: string, rev: string, bytes_total: number, state: string, unload_at: string | null, active_requests: number, context_used_tokens: number, bundled: boolean, };
diff --git a/macos/app/src/gen/ResyncPayload.ts b/macos/app/src/gen/ResyncPayload.ts
new file mode 100644
index 0000000..25cb8f9
--- /dev/null
+++ b/macos/app/src/gen/ResyncPayload.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type ResyncPayload = { gap_from: number, buffer_min: number, session_id: string, };
diff --git a/macos/app/src/gen/Role.ts b/macos/app/src/gen/Role.ts
new file mode 100644
index 0000000..49d909a
--- /dev/null
+++ b/macos/app/src/gen/Role.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type Role = "user" | "assistant" | "tool";
diff --git a/macos/app/src/gen/SystemSnapshot.ts b/macos/app/src/gen/SystemSnapshot.ts
new file mode 100644
index 0000000..2fc7c08
--- /dev/null
+++ b/macos/app/src/gen/SystemSnapshot.ts
@@ -0,0 +1,7 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { DaemonInfo } from "./DaemonInfo";
+import type { DownloadEntry } from "./DownloadEntry";
+import type { HardwareInfo } from "./HardwareInfo";
+import type { ModelEntry } from "./ModelEntry";
+
+export type SystemSnapshot = { daemon: DaemonInfo, hardware: HardwareInfo, models: Array<ModelEntry>, downloads: Array<DownloadEntry>, events_seq: number, };
diff --git a/macos/app/src/gen/ThreadRecord.ts b/macos/app/src/gen/ThreadRecord.ts
new file mode 100644
index 0000000..45de178
--- /dev/null
+++ b/macos/app/src/gen/ThreadRecord.ts
@@ -0,0 +1,6 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { JsonValue } from "./serde_json/JsonValue";
+
+export type ThreadRecord = { v: number, title: string, 
+model?: string | null, 
+params?: JsonValue | null, } & ({ [key in string]: number | string | boolean | Array<JsonValue> | { [key in string]: JsonValue } | null });
diff --git a/macos/app/src/gen/Usage.ts b/macos/app/src/gen/Usage.ts
new file mode 100644
index 0000000..06bc1b1
--- /dev/null
+++ b/macos/app/src/gen/Usage.ts
@@ -0,0 +1,6 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+/**
+ * schema `usage`.
+ */
+export type Usage = { prompt_tokens: number, completion_tokens: number, total_tokens: number, };
diff --git a/macos/app/src/gen/index.ts b/macos/app/src/gen/index.ts
new file mode 100644
index 0000000..bd222a2
--- /dev/null
+++ b/macos/app/src/gen/index.ts
@@ -0,0 +1,19 @@
+export * from "./Event";
+export * from "./HelloPayload";
+export * from "./ResyncPayload";
+export * from "./SystemSnapshot";
+export * from "./DaemonInfo";
+export * from "./HardwareInfo";
+export * from "./ModelEntry";
+export * from "./DownloadEntry";
+export * from "./DownloadFailure";
+export * from "./CatalogTier";
+export * from "./CatalogResponse";
+export * from "./EventType";
+export * from "./Role";
+export * from "./MessageStatus";
+export * from "./Usage";
+export * from "./Attachment";
+export * from "./MessageRecord";
+export * from "./ThreadRecord";
+
diff --git a/macos/app/src/gen/serde_json/JsonValue.ts b/macos/app/src/gen/serde_json/JsonValue.ts
new file mode 100644
index 0000000..320a852
--- /dev/null
+++ b/macos/app/src/gen/serde_json/JsonValue.ts
@@ -0,0 +1,3 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+export type JsonValue = number | string | boolean | Array<JsonValue> | { [key in string]: JsonValue } | null;
diff --git a/macos/app/src/i18n/index.ts b/macos/app/src/i18n/index.ts
new file mode 100644
index 0000000..957d107
--- /dev/null
+++ b/macos/app/src/i18n/index.ts
@@ -0,0 +1,16 @@
+import i18n from "i18next";
+import { initReactI18next } from "react-i18next";
+
+import en from "./locales/en.json";
+
+export const LANGS = ["en"] as const;
+export type Lang = (typeof LANGS)[number];
+
+void i18n.use(initReactI18next).init({
+  resources: { en: { translation: en } },
+  lng: "en",
+  fallbackLng: "en",
+  interpolation: { escapeValue: false },
+});
+
+export default i18n;
diff --git a/macos/app/src/i18n/locales/en.json b/macos/app/src/i18n/locales/en.json
new file mode 100644
index 0000000..49d9e90
--- /dev/null
+++ b/macos/app/src/i18n/locales/en.json
@@ -0,0 +1,331 @@
+{
+  "app": {
+    "name": "edge0",
+    "localLabel": "LOCAL INFERENCE"
+  },
+  "platform": {
+    "note": "Requires macOS (Apple M3 or later)"
+  },
+  "status": {
+    "title": "Service status",
+    "shellVersion": "Shell version",
+    "daemon": "Service process",
+    "scope": "Owner",
+    "uptime": "Uptime",
+    "waiting": "Waiting for local service…",
+    "unavailable": "The local service is unavailable.",
+    "hostApp": "App version",
+    "port": "Endpoint",
+    "diskFree": "Disk free",
+    "throttled": "Throttled by macOS",
+    "yes": "yes",
+    "no": "no",
+    "seq": "Last event number",
+    "models": "Loaded models",
+    "noResident": "No model is currently loaded",
+    "residentUntil": "Loaded until",
+    "hours": "h",
+    "minutes": "min",
+    "stream_live": "Events: live",
+    "stream_idle": "Events: idle",
+    "stream_reconnecting": "Events: reconnecting",
+    "stream_resyncing": "Events: resyncing",
+    "mismatch": "Service was started by an older app ({{daemon}}); this app is {{app}}. Use the button to restart (no silent binary swap).",
+    "restartService": "Restart service",
+    "exited": "Service stopped unexpectedly (exit code {{code}})",
+    "startService": "Start service",
+    "unknown": "unknown",
+    "remote": "Connected to the remote service at {{base}}. Local service controls are unavailable.",
+    "devDaemon": "dev: daemon from local target dir (release ships inside .app)"
+  },
+  "theme": {
+    "light": "Light",
+    "dark": "Dark",
+    "system": "System"
+  },
+  "font": {
+    "sm": "Small",
+    "md": "Default",
+    "lg": "Large"
+  },
+  "nav": {
+    "status": "Status",
+    "chat": "Chat",
+    "settings": "Settings",
+    "models": "Models",
+    "service": "Service",
+    "home": "Home",
+    "openSidebar": "Open sidebar",
+    "closeSidebar": "Close sidebar"
+  },
+  "common": {
+    "cancel": "Cancel"
+  },
+  "chat": {
+    "newThread": "New chat",
+    "rename": "Rename",
+    "deleteThread": "Delete chat",
+    "noThreads": "No chats yet",
+    "placeholder": "Type a message, Enter to send",
+    "send": "Send",
+    "stop": "Stop",
+    "regenerate": "Regenerate",
+    "editResend": "Edit & resend",
+    "resend": "Resend",
+    "copy": "Copy",
+    "copyReasoning": "Copy thinking",
+    "reasoning": "Thinking",
+    "aborted": "Response interrupted",
+    "saveFailed": "Not saved (conversation store write failed)",
+    "usage": "tokens in {{p}}/out {{c}}/total {{t}}",
+    "perf": {
+      "decode": "{{v}} tok/s decode",
+      "prefill": "{{v}} tok/s prefill",
+      "ttft": "first token {{ms}} ms",
+      "mem": "memory {{gi}} GiB"
+    },
+    "model": "Model",
+    "noModel": "(no installed model)",
+    "light": {
+      "resident": "loaded",
+      "loading": "loading",
+      "installed": "not loaded",
+      "absent": "unknown"
+    },
+    "loadModel": "Load",
+    "contextUsed": "context {{n}} tok",
+    "contextUnknown": "context: unknown",
+    "waitingService": "Waiting for local service…",
+    "greeting": "What would you like to explore?",
+    "suggestion": {
+      "water": "Will the water level rise when floating ice melts? Explain in two sentences.",
+      "story": "Write a sci-fi story with a twist in exactly six words.",
+      "shirts": "Five shirts take an hour to dry. How long will 30 take if they all dry at once under the same conditions?",
+      "averageSpeed": "I drive somewhere at 30 km/h. How fast must I drive back along the same route to average 60 km/h for the whole trip? Explain briefly."
+    },
+    "serviceAgent": "LaunchAgent service",
+    "serviceMenu": "Menu bar service",
+    "pickOrCreate": "Select or create a chat",
+    "noInstalledModel": "No models to chat with yet — download one first",
+    "goDownload": "Go to downloads",
+    "loadEarlier": "Load earlier messages",
+    "waitingFirstChunk": "Waiting for the model…",
+    "pickModelFirst": "Pick a model first",
+    "err": {
+      "notload": "Model not loaded — press Load, or download it",
+      "busy": "Service busy, retry shortly",
+      "cancel": "Generation cancelled",
+      "worker": "Engine worker crashed and was recycled — retry",
+      "param": "Request parameters rejected — check the model selection",
+      "missing": "Model not installed — download it first",
+      "memload": "Not enough memory to load this model — try a smaller model or free memory",
+      "gpu": "No GPU device — run edge0 diagnostics to check the environment",
+      "version": "Service version mismatch — restart it from the status page",
+      "generic": "Request failed"
+    }
+  },
+  "sidebar": {
+    "newChat": "New chat",
+    "history": "Chat history",
+    "navigation": "Navigation",
+    "resident": "loaded",
+    "idle": "idle",
+    "noService": "No service",
+    "bucket": {
+      "today": "Today",
+      "yesterday": "Yesterday",
+      "week": "Previous 7 days",
+      "earlier": "Earlier"
+    }
+  },
+  "settings": {
+    "title": "Settings",
+    "kicker": "Tune edge0 appearance and local service access",
+    "appearance": "Appearance",
+    "appearanceNote": "Theme and font size apply across the app immediately.",
+    "systemPrompt": "System prompt",
+    "systemPromptNote": "Optional instructions sent before each chat. Leave blank to disable.",
+    "systemPromptEnabled": "System prompt enabled",
+    "systemPromptDisabled": "No system prompt will be sent",
+    "systemPromptReset": "Use default",
+    "service": "Service",
+    "serviceNote": "View a compact summary; full controls, logs, and diagnostics live in the service console.",
+    "serviceUnavailable": "Service not connected",
+    "advancedNote": "Advanced service actions, request logs, and diagnostics live in the service console.",
+    "openConsole": "Open service console"
+  },
+  "gen": {
+    "precise": "Precise",
+    "balanced": "Balanced",
+    "creative": "Creative",
+    "hint": "Generation style: sampling temperature (Precise 0.2 / Balanced 0.7 / Creative 1.0), applied from the next message."
+  },
+  "models": {
+    "title": "Models & Downloads",
+    "kicker": "Manage models, keep them loaded, and track downloads",
+    "refresh": "Refresh",
+    "loading": "Reading catalog…",
+    "retry": "Retry",
+    "stale": "Couldn't refresh the model list. You can browse installed models, but downloads are paused.",
+    "rev": "Version",
+    "bytes": "Size",
+    "files": "Files",
+    "source": "Source",
+    "singleSource": "only one source",
+    "license": "License",
+    "notes": "Model card",
+    "state": {
+      "absent": "Not installed",
+      "installed": "Installed",
+      "update": "Update available"
+    },
+    "updateFrom": "Updatable from {{rev}}",
+    "download": "Download",
+    "updateDownload": "Download update",
+    "sourceAuto": "Auto source",
+    "sourceMs": "ModelScope",
+    "sourceHf": "Hugging Face",
+    "keepDefault": "Default keep time",
+    "keep30": "Keep loaded for 30 min",
+    "keepInf": "Keep loaded",
+    "load": "Load model",
+    "unload": "Unload model",
+    "loadingModel": "Loading…",
+    "delete": "Delete",
+    "bundled": "Bundled",
+    "bundledHint": "Shipped with this app build; delete is disabled.",
+    "deleteConfirm": "Delete this model? Run `edge0 gc` in Terminal to free space used by files no longer needed by any model.",
+    "residentUntil": "Loaded until",
+    "remaining": "~{{n}} min",
+    "unknown": "unknown",
+    "pause": "Pause",
+    "resume": "Resume",
+    "deleteTask": "Delete task",
+    "sourceSwitched": "Switched source",
+    "eta": "~{{n}} min left",
+    "estimating": "estimating speed…",
+    "verifying": "verifying…",
+    "phase": {
+      "probe": "Checking download sources",
+      "select-source": "Choosing a source",
+      "preflight": "Checking disk space",
+      "fetch": "Downloading",
+      "verify": "Verifying",
+      "stage": "Preparing files",
+      "commit": "Installing",
+      "register": "Finishing installation"
+    },
+    "err": {
+      "net": "Couldn't connect to the download provider. Check your network and try again.",
+      "src": "Couldn't reach either download provider. Try again later or switch networks.",
+      "disk": "Not enough disk space — free space and try again, or choose a smaller model.",
+      "hash": "Downloaded files failed integrity checks from both providers. Delete this task and try again.",
+      "sig": "The model catalog could not be verified. Do not bypass this check; try again later.",
+      "expired": "The model catalog is out of date. You can browse installed models, but try downloads again later.",
+      "minver": "This model needs a newer version of edge0. Update the app and try again.",
+      "conflict": "This model is being used by another request. Wait or cancel the request, then try again.",
+      "generic": "Request failed",
+      "invalid": "Model files do not match the installation record. They may be damaged or changed; delete the model and download it again."
+    },
+    "sourceDirect": "HuggingFace (direct)"
+  },
+  "welcome": {
+    "title": "Get started with edge0 on this Mac",
+    "body": "Recommended: {{model}} ({{bytes}}, stored in {{path}}) — inference runs entirely on-device.",
+    "startupTitle": "Service startup",
+    "headlessLabel": "Keep edge0 in the menu bar after login",
+    "headlessNote": "Registers a LaunchAgent; when unchecked, this app session owns the service.",
+    "errors": {
+      "disk": "Not enough disk space",
+      "memory": "There is not enough memory to load the recommended model",
+      "engine": "The local engine is unavailable",
+      "generic": "First launch did not complete"
+    },
+    "bytesUnknown": "size unknown",
+    "pathUnknown": "local app directory",
+    "downloadEnable": "Download & enable",
+    "skip": "Skip for now"
+  },
+  "service": {
+    "title": "Local API Service",
+    "kicker": "Manage lifecycle, access, and runtime diagnostics",
+    "sections": {
+      "lifecycle": "Lifecycle & ownership",
+      "lifecycleNote": "Start or stop the service and inspect ownership, version, and uptime.",
+      "logs": "Request log",
+      "logsNote": "Shows the request window since the current service start.",
+      "doctor": "Environment checks",
+      "doctorNote": "Shows daemon-provided findings and next actions."
+    },
+    "panel": {
+      "title": "Service Panel",
+      "start": "Start service",
+      "stop": "Stop service",
+      "busy": "Service operation in progress… (restart ends in-flight requests)",
+      "actionFailed": "The service action could not be completed.",
+      "port": "Listening",
+      "bind": "Bind address",
+      "scope": "Ownership",
+      "uptime": "Uptime",
+      "version": "Service version",
+      "versionMismatch": "Service was started by an older app ({{daemon}}); current app is {{app}}",
+      "home": "Storage location",
+      "token": "Access token",
+      "mask": "••••••••",
+      "copy": "Copy",
+      "copiedWarn": "Copied — clipboard now contains a sensitive credential",
+      "noCredential": "No credential needed (loopback is exempt)",
+      "notOwned": "This service was not started by this app; stop it via its owner",
+      "lan": "Allow LAN",
+      "lanOn": "On",
+      "lanOff": "Off",
+      "lanGoOn": "Confirm",
+      "lanConfirm": "The service will restart on a non-loopback address; in-flight requests will end. Continue?",
+      "lanGuide": "LAN allowed: other devices connect with `edge0 token` or the copied credential (/health stays token-free).",
+      "agent": "Startup owner",
+      "agentToMenu": "Stay in menu bar",
+      "agentToLaunchd": "Switch to LaunchAgent"
+    },
+    "log": {
+      "title": "Request Log",
+      "count": "{{n}} requests since start",
+      "note": "Window = events since this service started (ring buffer, 4096 max); cleared on restart",
+      "empty": "No requests since this service started",
+      "noCode": "(no error code)",
+      "th": {
+        "status": "Status",
+        "model": "Model",
+        "latency": "Latency / TTFT",
+        "tokens": "tokens",
+        "source": "Source"
+      },
+      "status": {
+        "done": "done",
+        "cancelled": "cancelled",
+        "error": "error"
+      },
+      "source": {
+        "loopback": "local",
+        "lan": "LAN",
+        "cli": "credential"
+      }
+    },
+    "doctor": {
+      "title": "Environment Check",
+      "loading": "Reading doctor results…",
+      "failed": "Cannot fetch doctor results (service not running)",
+      "retry": "Retry",
+      "next": "Next",
+      "goModels": "Go to Models",
+      "overall": {
+        "pass": "All passing",
+        "warn": "Warnings"
+      },
+      "verdict": {
+        "pass": "pass",
+        "warn": "warn",
+        "fail": "fail"
+      }
+    }
+  }
+}
diff --git a/macos/app/src/lib/generation.ts b/macos/app/src/lib/generation.ts
new file mode 100644
index 0000000..7e65375
--- /dev/null
+++ b/macos/app/src/lib/generation.ts
@@ -0,0 +1,23 @@
+// App-layer sampling presets. Engine default 0.7 is left alone for parity; precise 0.2 is the product default.
+export type GenPreset = "precise" | "balanced" | "creative";
+
+const GEN_KEY = "edge0.gen";
+
+export const GEN_TEMPERATURE: Record<GenPreset, number> = {
+  precise: 0.2,
+  balanced: 0.7,
+  creative: 1.0,
+};
+
+export function storedGen(): GenPreset {
+  const v = localStorage.getItem(GEN_KEY);
+  return v === "balanced" || v === "creative" ? v : "precise";
+}
+
+export function setGen(p: GenPreset): void {
+  localStorage.setItem(GEN_KEY, p);
+}
+
+export function currentTemperature(): number {
+  return GEN_TEMPERATURE[storedGen()];
+}
diff --git a/macos/app/src/lib/systemPrompt.ts b/macos/app/src/lib/systemPrompt.ts
new file mode 100644
index 0000000..7081701
--- /dev/null
+++ b/macos/app/src/lib/systemPrompt.ts
@@ -0,0 +1,21 @@
+export const DEFAULT_SYSTEM_PROMPT = "You're edge0, an on-device AI assistant.";
+
+const SYSTEM_PROMPT_KEY = "edge0.system-prompt";
+const SYSTEM_PROMPT_EVENT = "edge0:system-prompt";
+
+export function storedSystemPrompt(): string | null {
+  const stored = localStorage.getItem(SYSTEM_PROMPT_KEY);
+  if (stored === null) return DEFAULT_SYSTEM_PROMPT;
+  const trimmed = stored.trim();
+  return trimmed.length > 0 ? stored : null;
+}
+
+export function setSystemPrompt(value: string): void {
+  localStorage.setItem(SYSTEM_PROMPT_KEY, value);
+  window.dispatchEvent(new Event(SYSTEM_PROMPT_EVENT));
+}
+
+export function subscribeSystemPrompt(listener: () => void): () => void {
+  window.addEventListener(SYSTEM_PROMPT_EVENT, listener);
+  return () => window.removeEventListener(SYSTEM_PROMPT_EVENT, listener);
+}
diff --git a/macos/app/src/lib/theme.ts b/macos/app/src/lib/theme.ts
new file mode 100644
index 0000000..b76908a
--- /dev/null
+++ b/macos/app/src/lib/theme.ts
@@ -0,0 +1,31 @@
+// Theme is applied via data-theme on the document; "system" is resolved in JS, not prefers-color-scheme CSS.
+export type Theme = "light" | "dark" | "system";
+export type FontScale = "sm" | "md" | "lg";
+
+const THEME_KEY = "edge0.t…heme";
+const FONT_KEY = "edge0.token…font";
+
+export function resolveTheme(t: Theme): "light" | "dark" {
+  if (t !== "system") return t;
+  return window.matchMedia("(prefers-color-scheme: dark)").matches ? "dark" : "light";
+}
+
+export function storedTheme(): Theme {
+  const v = localStorage.getItem(THEME_KEY);
+  return v === "light" || v === "dark" || v === "system" ? v : "system";
+}
+
+export function storedFont(): FontScale {
+  const v = localStorage.getItem(FONT_KEY);
+  return v === "sm" || v === "md" || v === "lg" ? v : "md";
+}
+
+export function applyTheme(t: Theme): void {
+  localStorage.setItem(THEME_KEY, t);
+  document.documentElement.dataset.theme = resolveTheme(t);
+}
+
+export function applyFont(f: FontScale): void {
+  localStorage.setItem(FONT_KEY, f);
+  document.documentElement.dataset.font = f;
+}
diff --git a/macos/app/src/main.tsx b/macos/app/src/main.tsx
new file mode 100644
index 0000000..d56459b
--- /dev/null
+++ b/macos/app/src/main.tsx
@@ -0,0 +1,18 @@
+import React from "react";
+import ReactDOM from "react-dom/client";
+import { RouterProvider } from "@tanstack/react-router";
+
+import "./i18n";
+import { applyFont, applyTheme, storedFont, storedTheme } from "./lib/theme";
+import { router } from "./router";
+import "./styles.css";
+import "katex/dist/katex.min.css";
+
+applyTheme(storedTheme());
+applyFont(storedFont());
+
+ReactDOM.createRoot(document.getElementById("root")!).render(
+  <React.StrictMode>
+    <RouterProvider router={router} />
+  </React.StrictMode>,
+);
diff --git a/macos/app/src/models/autoLoad.ts b/macos/app/src/models/autoLoad.ts
new file mode 100644
index 0000000..40108d6
--- /dev/null
+++ b/macos/app/src/models/autoLoad.ts
@@ -0,0 +1,58 @@
+type Phase = "download" | "load";
+
+interface Token {
+  tier: string;
+  taskId: string;
+  phase: Phase;
+}
+
+type Listener = () => void;
+
+class AutoLoadStore {
+  private token: Token | null = null;
+  private listeners = new Set<Listener>();
+
+  get current(): Token | null {
+    return this.token;
+  }
+
+  subscribe(fn: Listener): () => void {
+    this.listeners.add(fn);
+    return () => {
+      this.listeners.delete(fn);
+    };
+  }
+
+  private emit(): void {
+    for (const l of this.listeners) l();
+  }
+
+  begin(tier: string, taskId: string): void {
+    this.token = { tier, taskId, phase: "download" };
+    this.emit();
+  }
+
+  claimCompleted(taskId: string): string | null {
+    if (!this.token || this.token.phase !== "download" || this.token.taskId !== taskId) return null;
+    const tier = this.token.tier;
+    this.token = { ...this.token, phase: "load" };
+    this.emit();
+    return tier;
+  }
+
+  consume(tier: string): boolean {
+    if (!this.token || this.token.phase !== "load" || this.token.tier !== tier) return false;
+    this.token = null;
+    this.emit();
+    return true;
+  }
+
+  cancel(): void {
+    if (this.token) {
+      this.token = null;
+      this.emit();
+    }
+  }
+}
+
+export const autoLoad = new AutoLoadStore();
diff --git a/macos/app/src/models/byteFormat.ts b/macos/app/src/models/byteFormat.ts
new file mode 100644
index 0000000..bbcafc3
--- /dev/null
+++ b/macos/app/src/models/byteFormat.ts
@@ -0,0 +1,29 @@
+export function humanBytes(n: number): string {
+  if (n >= 1e12) return `${(n / 1e12).toFixed(2)} TB`;
+  if (n >= 1e9) return `${(n / 1e9).toFixed(2)} GB`;
+  if (n >= 1e6) return `${(n / 1e6).toFixed(2)} MB`;
+  if (n >= 1e3) return `${(n / 1e3).toFixed(1)} KB`;
+  return `${n} B`;
+}
+
+export function humanRate(bytesPerSec: number): string {
+  return `${humanBytes(bytesPerSec)}/s`;
+}
+
+export function etaMinutes(etaS: number): number {
+  return Math.max(0.5, Math.round(etaS / 30) * 30 / 60);
+}
+
+export function formatClock(iso: string | null | undefined): string | null {
+  if (!iso) return null;
+  const d = new Date(iso);
+  if (Number.isNaN(d.getTime())) return null;
+  return d.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" });
+}
+
+export function remainingSeconds(unloadAtIso: string | null | undefined, nowMs: number): number | null {
+  if (!unloadAtIso) return null;
+  const t = new Date(unloadAtIso).getTime();
+  if (Number.isNaN(t)) return null;
+  return Math.max(0, Math.floor((t - nowMs) / 1000));
+}
diff --git a/macos/app/src/models/cards.tsx b/macos/app/src/models/cards.tsx
new file mode 100644
index 0000000..716ff90
--- /dev/null
+++ b/macos/app/src/models/cards.tsx
@@ -0,0 +1,131 @@
+import { useEffect, useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import { daemonStore } from "../bridge/state";
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { refreshSystem, taskView, type TaskView } from "./system";
+import { downloadErrKey } from "./errMap";
+import { etaMinutes, formatClock, humanBytes, humanRate, remainingSeconds } from "./byteFormat";
+
+export function DownloadCardForTier({ tier }: { tier: string }) {
+  const [, force] = useState(0);
+  useEffect(() => daemonStore.subscribe(() => force((n) => n + 1)), []);
+  const task = taskView(daemonStore.snapshot, daemonStore, tier);
+  if (!task) return null;
+  return <DownloadCardBody task={task} />;
+}
+
+export function DownloadCardBody({ task }: { task: TaskView }) {
+  const { t } = useTranslation();
+  const pct = task.bytesTotal > 0 ? Math.min(100, (task.bytesDone / task.bytesTotal) * 100) : 0;
+  const showEta = task.phase === "fetch" && task.etaS !== null;
+  return (
+    <div data-testid={`dl-${task.id}`} className="space-y-1 rounded-lg bg-surface2/50 p-3 text-left text-xs">
+      <div className="flex items-center justify-between">
+        <span data-testid={`dl-phase-${task.id}`}>{t(`models.phase.${task.phase}`)}</span>
+        <span className="text-muted">
+          {humanBytes(task.bytesDone)} / {humanBytes(task.bytesTotal)}
+          
+          {task.phase === "fetch" && task.rateBps !== null && ` - ${humanRate(task.rateBps)}`}
+        </span>
+      </div>
+      <div
+        className="h-1.5 w-full overflow-hidden rounded bg-line"
+        role="progressbar"
+        aria-valuenow={Math.round(pct)}
+        data-testid={`dl-bar-${task.id}`}
+      >
+        <div className="h-full bg-accent" style={{ width: `${pct}%` }} />
+      </div>
+      <div className="flex items-center justify-between gap-3">
+        <span
+          data-testid={`dl-eta-${task.id}`}
+          className="min-w-0 flex-1 truncate text-muted"
+        >
+          {task.failedCode
+            ? `${t(downloadErrKey(task.failedCode))} - ${task.failedCode}`
+            : showEta
+              ? t("models.eta", { n: etaMinutes(task.etaS as number) })
+              : task.phase === "fetch"
+                ? t("models.estimating")
+                : t("models.verifying")}
+        </span>
+        <div className="flex shrink-0 items-center gap-2">
+          {task.sourceSwitched && (
+            <span data-testid={`dl-switched-${task.id}`} className="text-warn">
+              {t("models.sourceSwitched")}
+            </span>
+          )}
+          {task.failedCode ? (
+            <button
+              type="button"
+              data-testid={`dl-retry-${task.id}`}
+              onClick={() => void post(`/v1/edge0/downloads/${task.id}/resume`)}
+              className="e0-btn e0-btn-secondary"
+            >
+              {t("models.retry")}
+            </button>
+          ) : task.paused ? (
+            <button
+              type="button"
+              data-testid={`dl-resume-${task.id}`}
+              onClick={() => void post(`/v1/edge0/downloads/${task.id}/resume`)}
+              className="e0-btn e0-btn-secondary"
+            >
+              {t("models.resume")}
+            </button>
+          ) : task.phase === "fetch" ? (
+            <button
+              type="button"
+              data-testid={`dl-pause-${task.id}`}
+              onClick={() => void post(`/v1/edge0/downloads/${task.id}/pause`)}
+              className="e0-btn e0-btn-secondary"
+            >
+              {t("models.pause")}
+            </button>
+          ) : null}
+          <button
+            type="button"
+            data-testid={`dl-cancel-${task.id}`}
+            onClick={() =>
+              void edge0Fetch(`/v1/edge0/downloads/${task.id}`, { method: "DELETE" })
+                .then(refreshSystem)
+                .catch(() => undefined)
+            }
+            className="rounded-full border border-line px-2 py-0.5 text-muted hover:text-err"
+          >
+            {t("models.deleteTask")}
+          </button>
+        </div>
+      </div>
+    </div>
+  );
+}
+
+export function ResidentCountdown({ unloadAt, resident }: { unloadAt: string | null; resident: boolean }) {
+  const { t } = useTranslation();
+  const [now, setNow] = useState(() => Date.now());
+  useEffect(() => {
+    if (!resident || !unloadAt) return;
+    const id = setInterval(() => setNow(Date.now()), 1000);
+    return () => clearInterval(id);
+  }, [resident, unloadAt]);
+  if (!resident) return null;
+  const clock = formatClock(unloadAt);
+  const rem = remainingSeconds(unloadAt, now);
+  return (
+    <span data-testid="resident-countdown" className="text-xs text-muted">
+      {t("models.residentUntil")} {clock ?? t("models.unknown")}
+      {rem !== null && ` - ${t("models.remaining", { n: Math.ceil(rem / 60) })}`}
+    </span>
+  );
+}
+
+async function post(path: string): Promise<void> {
+  try {
+    await edge0Fetch(path, { method: "POST", body: {} });
+  } catch {
+    /* ignore */
+  }
+  await refreshSystem();
+}
diff --git a/macos/app/src/models/catalog.ts b/macos/app/src/models/catalog.ts
new file mode 100644
index 0000000..aba5865
--- /dev/null
+++ b/macos/app/src/models/catalog.ts
@@ -0,0 +1,66 @@
+import type { CatalogResponse } from "../gen";
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { ApiError } from "../bridge/frames";
+
+export type CatalogStatus = "idle" | "loading" | "ready" | "error";
+
+interface Listener {
+  (): void;
+}
+
+class CatalogStore {
+  status: CatalogStatus = "idle";
+  data: CatalogResponse | null = null;
+  error: { code: string | null; message: string } | null = null;
+  lastError: Error | null = null;
+  private listeners = new Set<Listener>();
+  private inflight: Promise<void> | null = null;
+
+  subscribe(fn: Listener): () => void {
+    this.listeners.add(fn);
+    return () => {
+      this.listeners.delete(fn);
+    };
+  }
+
+  private emit(): void {
+    for (const l of this.listeners) l();
+  }
+
+  async refresh(): Promise<void> {
+    if (this.inflight) return this.inflight;
+    this.status = "loading";
+    this.emit();
+    const p = (async () => {
+      try {
+        const res = await edge0Fetch("/v1/edge0/catalog");
+        if (!res.ok) throw await apiErrFrom(res);
+        const j = (await res.json()) as CatalogResponse;
+        this.data = j;
+        this.error = null;
+        this.status = "ready";
+      } catch (e) {
+        const ae = e instanceof ApiError ? e : null;
+        this.error = { code: ae?.code ?? null, message: ae?.message ?? String(e) };
+        this.lastError = e instanceof Error ? e : new Error(String(e));
+        this.status = "error";
+      } finally {
+        this.inflight = null;
+        this.emit();
+      }
+    })();
+    this.inflight = p;
+    return p;
+  }
+}
+
+export const catalogStore = new CatalogStore();
+
+async function apiErrFrom(res: Response): Promise<ApiError> {
+  try {
+    const j = (await res.json()) as { error?: { code?: string; message?: string } };
+    return new ApiError(res.status, j?.error?.code ?? null, j?.error?.message ?? `HTTP ${res.status}`);
+  } catch {
+    return new ApiError(res.status, null, `HTTP ${res.status}`);
+  }
+}
diff --git a/macos/app/src/models/errMap.ts b/macos/app/src/models/errMap.ts
new file mode 100644
index 0000000..3a296e5
--- /dev/null
+++ b/macos/app/src/models/errMap.ts
@@ -0,0 +1,19 @@
+const DL_ERROR_KEYS: Record<string, string> = {
+  "E-DL-NET": "models.err.net",
+  "E-DL-SRC": "models.err.src",
+  "E-DL-DISK": "models.err.disk",
+  "E-MODEL-HASH": "models.err.hash",
+  "E-MODEL-INVALID": "models.err.invalid",
+  "E-MANIFEST-SIG": "models.err.sig",
+  "E-MANIFEST-EXPIRED": "models.err.expired",
+  "E-MANIFEST-MINVER": "models.err.minver",
+  "E-SRV-CONFLICT": "models.err.conflict",
+};
+
+export function downloadErrKey(code: string | null): string {
+  return (code && DL_ERROR_KEYS[code]) || "models.err.generic";
+}
+
+export function hasDownloadErr(code: string | null): boolean {
+  return code !== null && code in DL_ERROR_KEYS;
+}
diff --git a/macos/app/src/models/system.ts b/macos/app/src/models/system.ts
new file mode 100644
index 0000000..f15a301
--- /dev/null
+++ b/macos/app/src/models/system.ts
@@ -0,0 +1,100 @@
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { daemonStore } from "../bridge/state";
+import type { Event, SystemSnapshot } from "../gen";
+
+export const TERMINAL_EVENT_TYPES = [
+  "download.completed",
+  "download.failed",
+  "download.paused",
+  "model.load.started",
+  "model.load.ready",
+  "model.load.failed",
+  "model.unloaded",
+  "worker.crashed",
+  "service.restart",
+] as const;
+
+export async function refreshSystem(): Promise<void> {
+  try {
+    const res = await edge0Fetch("/v1/edge0/system");
+    if (!res.ok) return;
+    daemonStore.ingest((await res.json()) as SystemSnapshot);
+  } catch {
+    /* ignore */
+  }
+}
+
+export interface TaskView {
+  id: string;
+  tier: string;
+  rev: string;
+  phase: string;
+  source: string | null;
+  bytesDone: number;
+  bytesTotal: number;
+  paused: boolean;
+  rateBps: number | null;
+  etaS: number | null;
+  sourceSwitched: boolean;
+  failedCode: string | null;
+}
+
+const DL_TYPES = [
+  "download.progress",
+  "download.paused",
+  "download.completed",
+  "download.failed",
+] as const;
+
+export function latestDownloadEvent(
+  store: { latestFolded(type: string, subject: string): Event | undefined },
+  taskId: string,
+): Event | null {
+  let best: Event | null = null;
+  for (const t of DL_TYPES) {
+    const e = store.latestFolded(t, taskId);
+    if (e && (!best || e.seq > best.seq)) best = e;
+  }
+  return best;
+}
+
+export function taskView(
+  snap: SystemSnapshot | null,
+  store: { latestFolded(type: string, subject: string): Event | undefined },
+  tier: string,
+): TaskView | null {
+  const entry = (snap?.downloads ?? []).find((d) => d.tier === tier);
+  if (!entry) return null;
+  const v: TaskView = {
+    id: entry.id,
+    tier: entry.tier,
+    rev: entry.rev,
+    phase: entry.phase,
+    source: entry.source,
+    bytesDone: entry.bytes_done,
+    bytesTotal: entry.bytes_total,
+    paused: entry.paused,
+    rateBps: entry.rate_bps ?? null,
+    etaS: entry.eta_s ?? null,
+    sourceSwitched: entry.source_switched ?? false,
+    failedCode: entry.failed?.code ?? null,
+  };
+  const ev = latestDownloadEvent(store, entry.id);
+  if (ev && snap && ev.seq > snap.events_seq && ev.subject === entry.id) {
+    const p = ev.payload as Record<string, unknown>;
+    const num = (k: string): number | undefined =>
+      typeof p[k] === "number" ? (p[k] as number) : undefined;
+    if (typeof p.phase === "string") v.phase = p.phase;
+    if (typeof p.source === "string") v.source = p.source;
+    const bd = num("bytes_done");
+    if (bd !== undefined && bd >= v.bytesDone) v.bytesDone = bd;
+    if (typeof p.paused === "boolean") v.paused = p.paused;
+    const r = num("rate_bps");
+    if (r !== undefined) v.rateBps = r;
+    const eta = num("eta_s");
+    if (eta !== undefined) v.etaS = eta;
+    if (p.source_switched === true) v.sourceSwitched = true;
+    if (ev.type === "download.failed" && typeof p.code === "string") v.failedCode = p.code;
+  }
+  return v;
+}
diff --git a/macos/app/src/router.tsx b/macos/app/src/router.tsx
new file mode 100644
index 0000000..5aaa38f
--- /dev/null
+++ b/macos/app/src/router.tsx
@@ -0,0 +1,59 @@
+import { createRootRoute, createRoute, createRouter, redirect } from "@tanstack/react-router";
+
+import { AppShell } from "./components/AppShell";
+import { ChatPage } from "./routes/chat";
+import { ModelsPage } from "./routes/models";
+import { ServicePage } from "./routes/service";
+import { SettingsPage } from "./routes/settings";
+
+const rootRoute = createRootRoute({
+  component: AppShell,
+});
+
+const chatRoute = createRoute({
+  getParentRoute: () => rootRoute,
+  path: "/",
+  component: ChatPage,
+});
+
+const modelsRoute = createRoute({
+  getParentRoute: () => rootRoute,
+  path: "/models",
+  component: ModelsPage,
+});
+
+const settingsRoute = createRoute({
+  getParentRoute: () => rootRoute,
+  path: "/settings",
+  component: SettingsPage,
+});
+
+const serviceRoute = createRoute({
+  getParentRoute: () => rootRoute,
+  path: "/service",
+  component: ServicePage,
+});
+
+const legacyChatRedirect = createRoute({
+  getParentRoute: () => rootRoute,
+  path: "/chat",
+  beforeLoad: () => {
+    throw redirect({ to: "/" });
+  },
+});
+
+const routeTree = rootRoute.addChildren([
+  chatRoute,
+  modelsRoute,
+  legacyChatRedirect,
+  settingsRoute,
+  serviceRoute,
+]);
+
+export const router = createRouter({ routeTree });
+
+declare module "@tanstack/react-router" {
+  interface Register {
+    router: typeof router;
+  }
+}
diff --git a/macos/app/src/routes/chat.tsx b/macos/app/src/routes/chat.tsx
new file mode 100644
index 0000000..2d11183
--- /dev/null
+++ b/macos/app/src/routes/chat.tsx
@@ -0,0 +1,198 @@
+import { useSyncExternalStore } from "react";
+import { useCallback, useEffect, useRef, useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import { daemonStore } from "../bridge/state";
+import type { SystemSnapshot } from "../gen";
+import { useThreadChat } from "../chat/useThreadChat";
+import { threadCreate, threadSetModel, type ThreadRow } from "../chat/ipc";
+import { type ChatUIMessage } from "../chat/records";
+import { toWireMessages } from "../chat/transport";
+import { threadStore } from "../state/threads";
+import { isAtBottom } from "../chat/scroll";
+import { MessageList } from "../components/chat/MessageList";
+import { ComposerBar } from "../components/chat/ComposerBar";
+import { WelcomeCard } from "./welcome";
+import { storedSystemPrompt, subscribeSystemPrompt } from "../lib/systemPrompt";
+
+export function ChatPage() {
+  const { t } = useTranslation();
+  const threads = useSyncExternalStore((cb) => threadStore.subscribe(cb), () => threadStore.threads);
+  const activeId = useSyncExternalStore((cb) => threadStore.subscribe(cb), () => threadStore.activeId);
+  const snap = useSyncExternalStore((cb) => daemonStore.subscribe(cb), () => daemonStore.snapshot);
+  const [seed, setSeed] = useState<{ threadId: string; text: string } | null>(null);
+
+  const active = threads.find((x) => x.id === activeId) ?? null;
+  const hasInstalled = (snap?.models ?? []).length > 0;
+
+  if (!snap) {
+    return (
+      <div className="e0-chat-empty text-sm text-muted">
+        <span className="mx-auto grid size-10 place-items-center rounded-2xl bg-surface2 text-fg">
+          <span className="size-2 rounded-full bg-warn animate-pulse" aria-hidden />
+        </span>
+        <p className="font-medium text-fg">{t("chat.waitingService")}</p>
+        <p className="text-xs">{t("platform.note")}</p>
+      </div>
+    );
+  }
+  if (!hasInstalled) {
+    return <WelcomeCard />;
+  }
+  if (!active) {
+    const promptSuggestions = [
+      t("chat.suggestion.water"),
+      t("chat.suggestion.story"),
+      t("chat.suggestion.shirts"),
+      t("chat.suggestion.averageSpeed"),
+    ];
+    const sendSuggestion = async (text: string) => {
+      try {
+        const row = await threadCreate(t("chat.newThread"), defaultModel(snap));
+        await threadStore.refresh();
+        setSeed({ threadId: row.id, text });
+        threadStore.setActive(row.id);
+      } catch {
+        setSeed(null);
+      }
+    };
+    return (
+      <div className="e0-chat-empty">
+        <span className="mx-auto grid size-12 place-items-center rounded-2xl bg-fg text-surface shadow-sm">
+          <span className="text-lg font-semibold" aria-hidden>+</span>
+        </span>
+        <h2 className="mt-1 text-3xl font-semibold tracking-tight">{t("chat.greeting")}</h2>
+        <p className="text-sm text-muted">{t("chat.pickOrCreate")}</p>
+        <div className="mx-auto mt-4 grid w-full max-w-[760px] gap-2 sm:grid-cols-2" data-testid="chat-suggestions">
+          {promptSuggestions.map((text) => (
+            <button
+              key={text}
+              type="button"
+              onClick={() => void sendSuggestion(text)}
+              className="rounded-xl border border-line bg-surface2/60 px-4 py-3 text-left text-sm leading-snug text-fg/80 transition-colors hover:bg-surface2 hover:text-fg"
+            >
+              {text}
+            </button>
+          ))}
+        </div>
+      </div>
+    );
+  }
+  return (
+    <ChatPane
+      key={active.id}
+      thread={active}
+      snapshot={snap}
+      initialPrompt={seed?.threadId === active.id ? seed.text : null}
+      onInitialConsumed={() => setSeed((s) => (s?.threadId === active.id ? null : s))}
+      onChanged={() => void threadStore.refresh()}
+    />
+  );
+}
+
+function ChatPane({
+  thread,
+  snapshot,
+  initialPrompt,
+  onInitialConsumed,
+  onChanged,
+}: {
+  thread: ThreadRow;
+  snapshot: SystemSnapshot;
+  initialPrompt: string | null;
+  onInitialConsumed: () => void;
+  onChanged: () => void;
+}) {
+  const { t } = useTranslation();
+  const [model, setModel] = useState(thread.data.model ?? defaultModel(snapshot) ?? "");
+  const systemPrompt = useSyncExternalStore(
+    subscribeSystemPrompt,
+    storedSystemPrompt,
+    storedSystemPrompt,
+  );
+  const scrollerRef = useRef<HTMLDivElement>(null);
+  const stickBottom = useRef(true);
+  const [meterNonce, setMeterNonce] = useState(0);
+
+  const chat = useThreadChat(thread, model, () => {
+    onChanged();
+    setMeterNonce((n) => n + 1);
+  }, undefined, systemPrompt);
+
+  const onModelChange = useCallback(
+    (id: string) => {
+      setModel(id);
+      void threadSetModel(thread.id, id);
+    },
+    [thread.id],
+  );
+
+  const seedSentRef = useRef(false);
+  useEffect(() => {
+    // History must finish filling first, or the refill treats the in-flight seed as empty UI and duplicates it.
+    if (!initialPrompt || !chat.historyLoaded || seedSentRef.current) return;
+    seedSentRef.current = true;
+    onInitialConsumed();
+    void chat.send(initialPrompt);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [initialPrompt, chat.historyLoaded]);
+
+  const wireJson = useCallback(
+    () => toWireMessages(chat.messages as ChatUIMessage[] as never, systemPrompt),
+    [chat.messages, systemPrompt],
+  );
+
+  useEffect(() => {
+    // Stick to the bottom only when the user was already near it (or just sent).
+    const last = chat.messages[chat.messages.length - 1];
+    if (last?.role === "user") stickBottom.current = true;
+    const el = scrollerRef.current;
+    if (el && stickBottom.current) el.scrollTo({ top: el.scrollHeight });
+  }, [chat.messages]);
+
+  const onScrollerScroll = () => {
+    const el = scrollerRef.current;
+    if (el) stickBottom.current = isAtBottom(el);
+  };
+
+  return (
+    <div className="flex h-full min-h-0 flex-col">
+      <header className="e0-chat-header">
+        <div className="min-w-0 overflow-hidden">{/* min-w-0 is required for title ellipsis */}
+          <div className="e0-chat-header-title">{thread.data.title}</div>
+          <div className="e0-chat-header-meta">
+            <span className={"size-2 rounded-full " + (snapshot ? "bg-ok" : "bg-muted")} aria-hidden />
+            {snapshot.daemon.scope === "launchagent" ? t("chat.serviceAgent") : t("chat.serviceMenu")}
+          </div>
+        </div>
+        <span className="hidden text-xs text-muted sm:inline">{model || t("chat.noModel")}</span>
+      </header>
+      <div ref={scrollerRef} onScroll={onScrollerScroll} className="min-h-0 flex-1 overflow-y-auto">
+        <div className="e0-chat-column space-y-1 px-5 pb-4 pt-7 sm:px-6">
+          {chat.hasMoreEarlier && (
+            <button
+              type="button"
+              onClick={() => void chat.loadEarlier()}
+              className="mx-auto block rounded-full px-3 py-1 text-xs text-muted hover:bg-surface2 hover:text-fg"
+            >
+              {t("chat.loadEarlier")}
+            </button>
+          )}
+          <MessageList chat={chat} />
+        </div>
+      </div>
+      <ComposerBar
+        chat={chat}
+        model={model}
+        onModelChange={onModelChange}
+        wireJson={wireJson}
+        meterNonce={meterNonce}
+      />
+    </div>
+  );
+}
+
+function defaultModel(snap: SystemSnapshot | null): string | null {
+  const resident = snap?.models.find((m) => m.state === "resident");
+  return (resident ?? snap?.models[0])?.id ?? null;
+}
diff --git a/macos/app/src/routes/models.tsx b/macos/app/src/routes/models.tsx
new file mode 100644
index 0000000..da31f02
--- /dev/null
+++ b/macos/app/src/routes/models.tsx
@@ -0,0 +1,367 @@
+import { useEffect, useState } from "react";
+import { useTranslation } from "react-i18next";
+import { BookOpen, Download, RefreshCw, Trash2 } from "lucide-react";
+
+import { daemonStore } from "../bridge/state";
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { ApiError } from "../bridge/frames";
+import type { CatalogTier } from "../gen";
+import { catalogStore } from "../models/catalog";
+import {
+  refreshSystem,
+  taskView,
+  TERMINAL_EVENT_TYPES,
+} from "../models/system";
+import { DownloadCardBody, ResidentCountdown } from "../models/cards";
+import { downloadErrKey } from "../models/errMap";
+import { humanBytes } from "../models/byteFormat";
+import { deriveLight, latestModelEvent, type ModelLight } from "../chat/modelStatus";
+
+export function ModelsPage() {
+  const { t } = useTranslation();
+  const [, force] = useState(0);
+  useEffect(() => {
+    const off1 = daemonStore.subscribe(() => force((n) => n + 1));
+    const off2 = catalogStore.subscribe(() => force((n) => n + 1));
+    return () => {
+      off1();
+      off2();
+    };
+  }, []);
+
+  useEffect(() => {
+    let lastTerm = 0;
+    return daemonStore.subscribe(() => {
+      const m = daemonStore.maxSeqAmong(TERMINAL_EVENT_TYPES);
+      if (m > lastTerm) {
+        lastTerm = m;
+        void refreshSystem().then(() => void catalogStore.refresh());
+      }
+    });
+  }, []);
+
+  useEffect(() => {
+    void catalogStore.refresh();
+    void refreshSystem();
+  }, []);
+
+  const catalog = catalogStore;
+  const stale = catalog.data?.stale ?? false;
+
+  return (
+    <section className="e0-page space-y-4" data-testid="models-page">
+      <div className="e0-page-header">
+        <div>
+          <h2 className="e0-page-title flex items-center gap-2">
+          <BookOpen size={18} aria-hidden />
+          {t("models.title")}
+          </h2>
+          <p className="e0-page-kicker">{t("models.kicker")}</p>
+        </div>
+        <button
+          type="button"
+          data-testid="models-refresh"
+          onClick={() => {
+            void catalogStore.refresh();
+            void refreshSystem();
+          }}
+          className="e0-btn e0-btn-secondary"
+        >
+          <RefreshCw
+            size={12}
+            aria-hidden
+            className={catalog.status === "loading" ? "animate-spin" : ""}
+          />
+          {t("models.refresh")}
+        </button>
+      </div>
+
+      {stale && (
+        <p data-testid="models-stale" className="e0-status e0-status-warn w-full justify-start rounded-xl px-3 py-2">
+          {t("models.stale")}
+        </p>
+      )}
+
+      {catalog.status === "error" && !catalog.data && (
+        <p
+          data-testid="models-error"
+          className="e0-status e0-status-err w-full justify-start rounded-xl px-3 py-2"
+        >
+          {t(downloadErrKey(catalog.error?.code ?? null))}
+          {catalog.error?.code ? ` - ${catalog.error.code}` : ""}
+          <button
+            type="button"
+            data-testid="models-retry"
+            onClick={() => void catalogStore.refresh()}
+            className="ml-2 underline"
+          >
+            {t("models.retry")}
+          </button>
+        </p>
+      )}
+
+      {!catalog.data && catalog.status !== "error" && (
+        <p className="text-sm text-muted">{t("models.loading")}</p>
+      )}
+
+      {catalog.data?.tiers.map((tier) => (
+        <TierCard
+          key={tier.tier}
+          tier={tier}
+          stale={stale}
+          direct={catalog.data?.signer_key_id === "direct"}
+        />
+      ))}
+    </section>
+  );
+}
+
+function TierCard({ tier, stale, direct }: { tier: CatalogTier; stale: boolean; direct?: boolean }) {
+  const { t } = useTranslation();
+  const [, force] = useState(0);
+  useEffect(() => daemonStore.subscribe(() => force((n) => n + 1)), []);
+  const [source, setSource] = useState<string>("");
+  const [keepAlive, setKeepAlive] = useState<string>("");
+  const [opError, setOpError] = useState<string | null>(null);
+  const [confirmDelete, setConfirmDelete] = useState(false);
+
+  const snap = daemonStore.snapshot;
+  const task = taskView(snap, daemonStore, tier.tier);
+  const installed = snap?.models.find((m) => m.id === tier.tier);
+  const light: ModelLight = deriveLight(snap, tier.tier, latestModelEvent(daemonStore, tier.tier));
+  const stateLabel: "absent" | "installed" | "update" = tier.update_available
+    ? "update"
+    : installed || tier.installed
+      ? "installed"
+      : "absent";
+
+  async function run(fn: () => Promise<void>) {
+    setOpError(null);
+    setConfirmDelete(false);
+    try {
+      await fn();
+      void refreshSystem();
+    } catch (e) {
+      const ae = e instanceof ApiError ? e : null;
+      setOpError(ae ? `${t(downloadErrKey(ae.code))}${ae.code ? ` - ${ae.code}` : ""}` : t("models.err.generic"));
+    }
+  }
+
+  return (
+    <div
+      data-testid={`tier-${tier.tier}`}
+      className="e0-section flex flex-col gap-3 p-5"
+    >
+      <div className="flex items-baseline justify-between gap-2">
+        <div>
+          <h3 className="text-base font-semibold">{tier.tier}</h3>
+          <p className="text-xs text-muted" data-testid={`tier-state-${tier.tier}`}>
+            {t(`models.state.${stateLabel}`)}
+            {installed?.bundled ? ` - ${t("models.bundled")}` : ""}
+            {tier.update_available && tier.installed
+              ? ` - ${t("models.updateFrom", { rev: tier.installed })}`
+              : ""}
+          </p>
+        </div>
+        <span
+          data-testid={`light-${tier.tier}`}
+          data-light={light}
+          className={
+            "e0-status " +
+            (light === "resident" ? "e0-status-ok" : light === "loading" ? "e0-status-warn" : "")
+          }
+        >
+          <span className={"size-1.5 rounded-full " + (light === "resident" ? "bg-ok" : light === "loading" ? "bg-warn" : "bg-muted")} aria-hidden />
+          {t(`chat.light.${light}`)}
+        </span>
+      </div>
+
+      <dl className="order-3 grid grid-cols-2 gap-x-4 gap-y-1 border-t border-line/60 pt-3 text-xs">
+        <Field k={t("models.rev")} v={tier.rev} />
+        <Field k={t("models.bytes")} v={humanBytes(tier.bytes_total)} />
+        <Field k={t("models.files")} v={String(tier.files_count)} />
+        <Field
+          k={t("models.source")}
+          v={
+            tier.primary_source +
+            (tier.has_modelscope && tier.has_huggingface ? "" : ` - ${t("models.singleSource")}`)
+          }
+        />
+        {tier.license && <Field k={t("models.license")} v={tier.license} />}
+        {tier.notes_url && (
+          <div className="col-span-2">
+            <a
+              href={tier.notes_url}
+              target="_blank"
+              rel="noreferrer"
+              data-testid={`tier-notes-${tier.tier}`}
+              className="text-accent underline"
+            >
+              {t("models.notes")}
+            </a>
+          </div>
+        )}
+      </dl>
+
+      {task && <div className="order-2"><DownloadCardBody task={task} /></div>}
+
+      {opError && (
+        <p data-testid={`tier-operr-${tier.tier}`} role="alert" className="e0-status e0-status-err order-2 w-full justify-start rounded-lg">
+          {opError}
+        </p>
+      )}
+
+      {!task && (stateLabel === "absent" || stateLabel === "update") && (
+        <div className="order-1 flex flex-wrap items-center gap-2">
+          {direct ? (
+            <span data-testid={`source-direct-${tier.tier}`} className="rounded-full border border-line bg-surface2 px-3 py-1.5 text-xs text-muted">
+              {t("models.sourceDirect")}
+            </span>
+          ) : (
+            <select
+              data-testid={`source-select-${tier.tier}`}
+              value={source}
+              onChange={(e) => setSource(e.target.value)}
+              className="rounded-lg border border-line bg-surface px-2.5 py-1.5 text-xs"
+            >
+              <option value="">{t("models.sourceAuto")}</option>
+              <option value="modelscope">{t("models.sourceMs")}</option>
+              <option value="huggingface">{t("models.sourceHf")}</option>
+            </select>
+          )}
+          <button
+            type="button"
+            data-testid={`download-${tier.tier}`}
+            disabled={stale}
+            title={stale ? t("models.stale") : undefined}
+            onClick={() =>
+              void run(async () => {
+                const res = await edge0Fetch("/v1/edge0/downloads", {
+                  method: "POST",
+                  body: direct ? { tier: tier.tier } : { tier: tier.tier, source: source || null },
+                });
+                if (!res.ok) throw await apiErr(res);
+              })
+            }
+            className="e0-btn e0-btn-primary"
+          >
+            <Download size={12} aria-hidden />
+            {stateLabel === "update" ? t("models.updateDownload") : t("models.download")}
+          </button>
+        </div>
+      )}
+
+      {installed && !task && (
+        <div className="order-1 flex flex-wrap items-center gap-2">
+          <select
+            data-testid={`keepalive-${tier.tier}`}
+            value={keepAlive}
+            onChange={(e) => setKeepAlive(e.target.value)}
+            className="rounded-lg border border-line bg-surface px-2.5 py-1.5 text-xs"
+          >
+            <option value="">{t("models.keepDefault")}</option>
+            <option value="30m">{t("models.keep30")}</option>
+            <option value="inf">{t("models.keepInf")}</option>
+          </select>
+          {light !== "resident" && light !== "loading" && (
+            <button
+              type="button"
+              data-testid={`load-${tier.tier}`}
+              onClick={() =>
+                void run(async () => {
+                  const res = await edge0Fetch(`/v1/edge0/models/${tier.tier}/load`, {
+                    method: "POST",
+                    body: keepAlive ? { keep_alive: keepAlive } : {},
+                  });
+                  if (!res.ok) throw await apiErr(res);
+                })
+              }
+              className="e0-btn e0-btn-secondary"
+            >
+              {t("models.load")}
+            </button>
+          )}
+          {light === "loading" && <span className="text-xs text-muted">{t("models.loadingModel")}</span>}
+          {light === "resident" && (
+            <button
+              type="button"
+              data-testid={`unload-${tier.tier}`}
+              onClick={() =>
+                void run(async () => {
+                  const res = await edge0Fetch(`/v1/edge0/models/${tier.tier}/unload`, {
+                    method: "POST",
+                    body: {},
+                  });
+                  if (!res.ok) throw await apiErr(res);
+                })
+              }
+              className="e0-btn e0-btn-secondary"
+            >
+              {t("models.unload")}
+            </button>
+          )}
+          <ResidentCountdown unloadAt={installed.unload_at ?? null} resident={light === "resident"} />
+          {installed.bundled ? (
+            <span className="text-xs text-muted">{t("models.bundledHint")}</span>
+          ) : confirmDelete ? (
+            <span className="flex items-center gap-1 text-xs">
+              <span>{t("models.deleteConfirm")}</span>
+              <button
+                type="button"
+                data-testid={`delete-yes-${tier.tier}`}
+                onClick={() =>
+                  void run(async () => {
+                    const res = await edge0Fetch(`/v1/edge0/models/${tier.tier}`, {
+                      method: "DELETE",
+                    });
+                    if (!res.ok) throw await apiErr(res);
+                    void catalogStore.refresh();
+                  })
+                }
+                className="e0-btn bg-err text-white"
+              >
+                {t("models.delete")}
+              </button>
+              <button
+                type="button"
+                data-testid={`delete-no-${tier.tier}`}
+                onClick={() => setConfirmDelete(false)}
+                className="e0-btn e0-btn-secondary"
+              >
+                {t("common.cancel")}
+              </button>
+            </span>
+          ) : (
+            <button
+              type="button"
+              data-testid={`delete-${tier.tier}`}
+              onClick={() => setConfirmDelete(true)}
+              className="e0-btn e0-btn-danger"
+            >
+              <Trash2 size={12} aria-hidden />
+              {t("models.delete")}
+            </button>
+          )}
+        </div>
+      )}
+    </div>
+  );
+}
+
+function Field({ k, v }: { k: string; v: string }) {
+  return (
+    <>
+      <dt className="text-muted">{k}</dt>
+      <dd className="text-fg">{v}</dd>
+    </>
+  );
+}
+
+async function apiErr(res: Response): Promise<ApiError> {
+  try {
+    const j = (await res.json()) as { error?: { code?: string; message?: string } };
+    return new ApiError(res.status, j?.error?.code ?? null, j?.error?.message ?? `HTTP ${res.status}`);
+  } catch {
+    return new ApiError(res.status, null, `HTTP ${res.status}`);
+  }
+}
diff --git a/macos/app/src/routes/service.tsx b/macos/app/src/routes/service.tsx
new file mode 100644
index 0000000..165bf02
--- /dev/null
+++ b/macos/app/src/routes/service.tsx
@@ -0,0 +1,44 @@
+import { useTranslation } from "react-i18next";
+import { Server } from "lucide-react";
+
+import { DoctorSection } from "../service/doctor";
+import { ServicePanel } from "../service/panel";
+import { RequestLogSection } from "../service/reqlog";
+
+export function ServicePage() {
+  const { t } = useTranslation();
+  return (
+    <div className="e0-page space-y-5" data-testid="service-page">
+      <header className="e0-page-header">
+        <div>
+          <h2 className="e0-page-title flex items-center gap-2">
+          <Server size={18} aria-hidden />
+          {t("service.title")}
+          </h2>
+          <p className="e0-page-kicker">{t("service.kicker")}</p>
+        </div>
+      </header>
+      <section aria-labelledby="service-lifecycle-title" data-testid="service-lifecycle-section">
+        <div className="mb-2 px-1">
+          <h3 id="service-lifecycle-title" className="e0-section-title">{t("service.sections.lifecycle")}</h3>
+          <p className="e0-section-note">{t("service.sections.lifecycleNote")}</p>
+        </div>
+        <ServicePanel />
+      </section>
+      <section aria-labelledby="service-log-title" data-testid="service-log-section">
+        <div className="mb-2 px-1">
+          <h3 id="service-log-title" className="e0-section-title">{t("service.sections.logs")}</h3>
+          <p className="e0-section-note">{t("service.sections.logsNote")}</p>
+        </div>
+        <RequestLogSection />
+      </section>
+      <section aria-labelledby="service-doctor-title" data-testid="service-doctor-section">
+        <div className="mb-2 px-1">
+          <h3 id="service-doctor-title" className="e0-section-title">{t("service.sections.doctor")}</h3>
+          <p className="e0-section-note">{t("service.sections.doctorNote")}</p>
+        </div>
+        <DoctorSection />
+      </section>
+    </div>
+  );
+}
diff --git a/macos/app/src/routes/settings.tsx b/macos/app/src/routes/settings.tsx
new file mode 100644
index 0000000..5ee9afa
--- /dev/null
+++ b/macos/app/src/routes/settings.tsx
@@ -0,0 +1,102 @@
+import { Link } from "@tanstack/react-router";
+import { useState, useSyncExternalStore } from "react";
+import { Settings2, SlidersHorizontal, Server, RotateCcw } from "lucide-react";
+import { useTranslation } from "react-i18next";
+
+import { ShellControls } from "../components/ShellControls";
+import { daemonStore } from "../bridge/state";
+import { DEFAULT_SYSTEM_PROMPT, setSystemPrompt, storedSystemPrompt } from "../lib/systemPrompt";
+
+export function SettingsPage() {
+  const { t } = useTranslation();
+  const [systemPrompt, setPrompt] = useState(() => storedSystemPrompt() ?? "");
+  const snap = useSyncExternalStore(
+    (cb) => daemonStore.subscribe(cb),
+    () => daemonStore.snapshot,
+  );
+  return (
+    <div className="e0-page space-y-5">
+      <header className="e0-page-header">
+        <div>
+          <h2 className="e0-page-title flex items-center gap-2">
+            <Settings2 size={19} aria-hidden />
+            {t("settings.title")}
+          </h2>
+          <p className="e0-page-kicker">{t("settings.kicker")}</p>
+        </div>
+      </header>
+      <section className="e0-section overflow-hidden" aria-labelledby="settings-appearance-title">
+        <div className="e0-section-header">
+          <div>
+            <h3 id="settings-appearance-title" className="e0-section-title">{t("settings.appearance")}</h3>
+            <p className="e0-section-note">{t("settings.appearanceNote")}</p>
+          </div>
+          <SlidersHorizontal size={17} className="text-muted" aria-hidden />
+        </div>
+        <div className="px-5 pb-5">
+          <ShellControls />
+        </div>
+      </section>
+      <section className="e0-section overflow-hidden" aria-labelledby="settings-prompt-title">
+        <div className="e0-section-header">
+          <div>
+            <h3 id="settings-prompt-title" className="e0-section-title">{t("settings.systemPrompt")}</h3>
+            <p className="e0-section-note">{t("settings.systemPromptNote")}</p>
+          </div>
+        </div>
+        <div className="border-t border-line/60 px-5 pb-5 pt-4">
+          <textarea
+            value={systemPrompt}
+            onChange={(e) => {
+              setPrompt(e.target.value);
+              setSystemPrompt(e.target.value);
+            }}
+            placeholder={DEFAULT_SYSTEM_PROMPT}
+            rows={4}
+            aria-label={t("settings.systemPrompt")}
+            className="w-full resize-y rounded-xl border border-line bg-surface2 p-3 text-sm leading-relaxed outline-none focus:border-accent"
+          />
+          <div className="mt-3 flex items-center justify-between gap-3">
+            <span className="text-xs text-muted">
+              {systemPrompt.trim() ? t("settings.systemPromptEnabled") : t("settings.systemPromptDisabled")}
+            </span>
+            <button
+              type="button"
+              className="e0-btn e0-btn-secondary"
+              onClick={() => {
+                setPrompt(DEFAULT_SYSTEM_PROMPT);
+                setSystemPrompt(DEFAULT_SYSTEM_PROMPT);
+              }}
+            >
+              <RotateCcw size={13} aria-hidden />
+              {t("settings.systemPromptReset")}
+            </button>
+          </div>
+        </div>
+      </section>
+      <section className="e0-section overflow-hidden" aria-labelledby="settings-service-title">
+        <div className="e0-section-header">
+          <div>
+            <h3 id="settings-service-title" className="e0-section-title">{t("settings.service")}</h3>
+            <p className="e0-section-note">{t("settings.serviceNote")}</p>
+          </div>
+          <Server size={17} className="text-muted" aria-hidden />
+        </div>
+        <div className="flex items-center gap-3 border-t border-line/60 px-5 py-4">
+          <span className={"size-2 rounded-full " + (snap ? "bg-ok" : "bg-muted")} aria-hidden />
+          <span className="min-w-0 flex-1 truncate text-sm text-muted">
+            {snap ? `${snap.daemon.bound}:${snap.daemon.port} - v${snap.daemon.version}` : t("settings.serviceUnavailable")}
+          </span>
+          <Link
+            to="/service"
+            data-testid="service-console-link"
+            className="e0-btn e0-btn-secondary shrink-0"
+          >
+            {t("settings.openConsole")}
+          </Link>
+        </div>
+      </section>
+      <p className="px-1 text-xs text-muted">{t("settings.advancedNote")}</p>
+    </div>
+  );
+}
diff --git a/macos/app/src/routes/status.tsx b/macos/app/src/routes/status.tsx
new file mode 100644
index 0000000..66a9763
--- /dev/null
+++ b/macos/app/src/routes/status.tsx
@@ -0,0 +1,257 @@
+import { invoke } from "@tauri-apps/api/core";
+import { useCallback, useEffect, useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { daemonStore } from "../bridge/state";
+import type { SystemSnapshot } from "../gen";
+
+interface ServiceState {
+  phase: {
+    phase: string;
+    health?: {
+      daemon_version: string;
+      host_app_version: string | null;
+      scope: string;
+      started_by: string;
+      uptime_s: number;
+      pid: number;
+    };
+    code?: number | null;
+    detail?: string;
+    base?: string;
+    spawnedByUs?: boolean;
+  };
+  appVersion: string;
+  versionMismatch: boolean;
+  spawnWasDev: boolean;
+  pinned: boolean;
+}
+
+export function StatusPage() {
+  const { t } = useTranslation();
+  const [, force] = useState(0);
+  const [error, setError] = useState<string | null>(null);
+  const [svc, setSvc] = useState<ServiceState | null>(null);
+
+  useEffect(() => {
+    const off = daemonStore.subscribe(() => force((n) => n + 1));
+    void refresh();
+    const poll = setInterval(() => {
+      invoke<ServiceState>("service_state").then(setSvc).catch(() => undefined);
+    }, 2000);
+    return () => {
+      off();
+      clearInterval(poll);
+    };
+  }, []);
+
+  const act = useCallback(async (cmd: "service_start" | "service_restart") => {
+    try {
+      await invoke(cmd);
+    } catch (e) {
+      setError(String(e));
+    }
+  }, []);
+
+  async function refresh() {
+    try {
+      const res = await edge0Fetch("/v1/edge0/system");
+      if (!res.ok) throw new Error(`HTTP ${res.status}`);
+      daemonStore.ingest((await res.json()) as SystemSnapshot);
+      setError(null);
+    } catch (e) {
+      setError(t("status.unavailable"));
+    }
+  }
+
+  const snap = daemonStore.snapshot;
+  const resident = snap?.models.filter((m) => m.state === "resident") ?? [];
+  return (
+    <section className="space-y-3" data-testid="status-page">
+      <div className="flex items-baseline justify-between">
+        <h2 className="text-lg font-semibold">{t("status.title")}</h2>
+        <span
+          data-testid="stream-link"
+          className={
+            "rounded px-2 py-0.5 text-xs " +
+            (daemonStore.link === "live" || daemonStore.lastSeq > 0
+              ? "bg-accent text-white"
+              : "border border-line text-muted")
+          }
+        >
+          {t(`status.${linkKey(daemonStore.link, daemonStore.lastSeq)}`)}
+        </span>
+      </div>
+
+      {svc && <LifecycleBanner svc={svc} onStart={() => void act("service_start")} onRestart={() => void act("service_restart")} />}
+
+      {error && !snap && (
+        <p data-testid="status-error" className="rounded border border-line p-3 text-sm text-muted">
+          {error}
+        </p>
+      )}
+
+      {snap ? (
+        <dl
+          className="grid grid-cols-[auto_1fr] gap-x-4 gap-y-1 text-sm"
+          data-testid="snapshot-fields"
+        >
+          <Field k={t("status.daemon")} v={`v${snap.daemon.version}`} testid="daemon-version" />
+          <Field
+            k={t("status.hostApp")}
+            v={snap.daemon.host_app_version ?? "—"}
+            testid="host-app"
+          />
+          <Field
+            k={t("status.scope")}
+            v={snap.daemon.scope}
+            extra={snap.daemon.started_by}
+            testid="scope"
+          />
+          <Field
+            k={t("status.port")}
+            v={`${snap.daemon.bound}:${snap.daemon.port}`}
+            testid="port"
+          />
+          <Field k={t("status.uptime")} v={fmtUptime(snap.daemon.uptime_s, t)} testid="uptime" />
+          <Field
+            k={t("status.diskFree")}
+            v={`${snap.hardware.disk_free_gib.toFixed(1)} GiB`}
+            testid="disk"
+          />
+          <Field
+            k={t("status.throttled")}
+            v={snap.daemon.throttled ? t("status.yes") : t("status.no")}
+            testid="throttled"
+          />
+          <Field k={t("status.seq")} v={String(snap.events_seq)} testid="events-seq" />
+        </dl>
+      ) : (
+        !error && <p className="text-sm text-muted">{t("status.waiting")}</p>
+      )}
+
+      <div data-testid="resident-models" className="space-y-1">
+        <h3 className="text-sm font-medium">{t("status.models")}</h3>
+        {resident.length === 0 && (
+          <p className="text-sm text-muted" data-testid="no-resident">
+            {t("status.noResident")}
+          </p>
+        )}
+        {resident.map((m) => (
+          <p key={m.id} className="text-sm">
+            <span data-testid={`resident-${m.id}`}>{m.id}</span>
+            {m.unload_at && (
+              <span className="text-muted"> - {t("status.residentUntil")} {fmtHHMM(m.unload_at)}</span>
+            )}
+          </p>
+        ))}
+      </div>
+
+      
+      <p className="text-xs text-muted">{t("platform.note")}</p>
+    </section>
+  );
+}
+
+function LifecycleBanner({
+  svc,
+  onStart,
+  onRestart,
+}: {
+  svc: ServiceState;
+  onStart: () => void;
+  onRestart: () => void;
+}) {
+  const { t } = useTranslation();
+  const p = svc.phase;
+  if (svc.versionMismatch) {
+    return (
+      <div
+        data-testid="mismatch-banner"
+        className="flex items-center justify-between gap-3 rounded border border-yellow-500/50 p-3 text-sm"
+      >
+        <span>
+          {t("status.mismatch", {
+            daemon: p.health?.host_app_version ?? "?",
+            app: svc.appVersion,
+          })}
+        </span>
+        <button
+          data-testid="restart-service"
+          className="rounded bg-accent px-3 py-1 text-white"
+          onClick={onRestart}
+        >
+          {t("status.restartService")}
+        </button>
+      </div>
+    );
+  }
+  if (p.phase === "exited") {
+    return (
+      <div
+        data-testid="exited-banner"
+        className="flex items-center justify-between gap-3 rounded border border-line p-3 text-sm"
+      >
+        <span>{t("status.exited", { code: p.code ?? t("status.unknown") })}</span>
+        <button data-testid="start-service" className="rounded bg-accent px-3 py-1 text-white" onClick={onStart}>
+          {t("status.startService")}
+        </button>
+      </div>
+    );
+  }
+  if (p.phase === "stopped") {
+    return (
+      <div data-testid="stopped-banner" className="flex items-center justify-between gap-3 rounded border border-line p-3 text-sm">
+        <span className="text-muted">{p.detail}</span>
+        <button data-testid="start-service" className="rounded bg-accent px-3 py-1 text-white" onClick={onStart}>
+          {t("status.startService")}
+        </button>
+      </div>
+    );
+  }
+  if (p.phase === "remote") {
+    return (
+      <p data-testid="remote-note" className="text-sm text-muted">
+        {t("status.remote", { base: p.base ?? "?" })}
+      </p>
+    );
+  }
+  if (svc.spawnWasDev && (p.phase === "spawned" || p.spawnedByUs)) {
+    return (
+      <p data-testid="dev-badge" className="text-xs text-muted">
+        {t("status.devDaemon")}
+      </p>
+    );
+  }
+  return null;
+}
+
+function Field({ k, v, extra, testid }: { k: string; v: string; extra?: string; testid: string }) {
+  return (
+    <>
+      <dt className="text-muted">{k}</dt>
+      <dd data-testid={testid}>
+        {v}
+        {extra ? <span className="text-muted"> - {extra}</span> : null}
+      </dd>
+    </>
+  );
+}
+
+function linkKey(link: string, lastSeq: number): string {
+  if (link !== "live") return `stream_${link}`;
+  return lastSeq > 0 ? "stream_live" : "stream_idle";
+}
+
+function fmtUptime(s: number, t: (k: string) => string): string {
+  if (s < 60) return `${s}s`;
+  const h = Math.floor(s / 3600);
+  const m = Math.floor((s % 3600) / 60);
+  return h > 0 ? `${h}${t("status.hours")}${m}${t("status.minutes")}` : `${m}${t("status.minutes")}`;
+}
+
+function fmtHHMM(iso: string): string {
+  const d = new Date(iso);
+  return `${String(d.getHours()).padStart(2, "0")}:${String(d.getMinutes()).padStart(2, "0")}`;
+}
diff --git a/macos/app/src/routes/welcome.tsx b/macos/app/src/routes/welcome.tsx
new file mode 100644
index 0000000..b3eb02f
--- /dev/null
+++ b/macos/app/src/routes/welcome.tsx
@@ -0,0 +1,175 @@
+import { useEffect, useState } from "react";
+import { Link } from "@tanstack/react-router";
+import { invoke } from "@tauri-apps/api/core";
+import { Download, Sparkles } from "lucide-react";
+import { useTranslation } from "react-i18next";
+
+import { daemonStore } from "../bridge/state";
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { ApiError } from "../bridge/frames";
+import { catalogStore } from "../models/catalog";
+import { autoLoad } from "../models/autoLoad";
+import { refreshSystem, taskView } from "../models/system";
+import { downloadErrKey } from "../models/errMap";
+import { humanBytes } from "../models/byteFormat";
+import { DownloadCardForTier } from "../models/cards";
+
+export function WelcomeCard() {
+  const { t } = useTranslation();
+  const [, force] = useState(0);
+  useEffect(() => {
+    const off1 = daemonStore.subscribe(() => force((n) => n + 1));
+    const off2 = catalogStore.subscribe(() => force((n) => n + 1));
+    const off3 = autoLoad.subscribe(() => force((n) => n + 1));
+    return () => {
+      off1();
+      off2();
+      off3();
+    };
+  }, []);
+  useEffect(() => {
+    void catalogStore.refresh();
+  }, []);
+  const [dismissed, setDismissed] = useState(false);
+  const [err, setErr] = useState<{ code: string | null; message: string } | null>(null);
+  const [busy, setBusy] = useState(false);
+  const [headless, setHeadless] = useState(false);
+
+  const snap = daemonStore.snapshot;
+  const rec = catalogStore.data?.tiers.find((x) => x.tier === "edge0-8b") ?? null;
+  const task = taskView(snap, daemonStore, "edge0-8b");
+
+  if (dismissed) {
+    return (
+      <div className="e0-chat-empty">
+        <p className="text-lg font-medium">{t("chat.noInstalledModel")}</p>
+        <Link
+          to="/models"
+          className="justify-self-center rounded-full border border-line px-4 py-1.5 text-sm text-muted hover:bg-surface2 hover:text-fg"
+        >
+          {t("chat.goDownload")}
+        </Link>
+      </div>
+    );
+  }
+
+  async function downloadAndEnable() {
+    setErr(null);
+    setBusy(true);
+    try {
+      if (headless) {
+        await invoke("agent_set", { enable: true });
+      }
+      const res = await edge0Fetch("/v1/edge0/downloads", {
+        method: "POST",
+        body: { tier: "edge0-8b", source: null },
+      });
+      if (!res.ok) throw await errFrom(res);
+      const j = (await res.json()) as { id: string };
+      autoLoad.begin("edge0-8b", j.id);
+      void refreshSystem();
+    } catch (e) {
+      const ae = e instanceof ApiError ? e : null;
+      setErr({ code: ae?.code ?? null, message: ae?.message ?? String(e) });
+    } finally {
+      setBusy(false);
+    }
+  }
+
+  return (
+    <div className="e0-chat-empty px-6" data-testid="welcome-card">
+      <div className="e0-chat-column mx-auto space-y-5 text-center">
+        <span className="mx-auto grid size-12 place-items-center rounded-2xl bg-fg text-surface shadow-sm">
+          <Sparkles size={20} aria-hidden />
+        </span>
+        <div>
+          <h2 className="text-3xl font-semibold tracking-tight">{t("welcome.title")}</h2>
+          <p className="mx-auto mt-2 max-w-lg text-sm leading-relaxed text-muted">
+          {t("welcome.body", {
+            model: "edge0-8b",
+            bytes: rec ? humanBytes(rec.bytes_total) : t("welcome.bytesUnknown"),
+            path: snap?.daemon.home ?? t("welcome.pathUnknown"),
+          })}
+          </p>
+        </div>
+        <fieldset className="mx-auto w-full max-w-lg rounded-2xl border border-line bg-surface2/60 p-4 text-left">
+          <legend className="px-1 text-xs font-semibold text-fg">{t("welcome.startupTitle")}</legend>
+          <label className="mt-1 flex cursor-pointer items-start gap-3 text-sm">
+            <input
+              type="checkbox"
+              checked={headless}
+              onChange={(e) => setHeadless(e.target.checked)}
+              className="mt-0.5 size-4 accent-accent"
+            />
+            <span>
+              <span className="block font-medium">{t("welcome.headlessLabel")}</span>
+              <span className="mt-0.5 block text-xs leading-relaxed text-muted">{t("welcome.headlessNote")}</span>
+            </span>
+          </label>
+        </fieldset>
+        {task ? (
+          <DownloadCardForTier tier="edge0-8b" />
+        ) : (
+          <button
+            type="button"
+            data-testid="welcome-go"
+            disabled={busy || !catalogStore.data || catalogStore.data.stale}
+            title={catalogStore.data?.stale ? t("models.stale") : undefined}
+            onClick={() => void downloadAndEnable()}
+            className="e0-btn e0-btn-primary mx-auto min-h-10 px-5 text-sm"
+          >
+            <Download size={15} aria-hidden />
+            {t("welcome.downloadEnable")}
+          </button>
+        )}
+        {err && (
+            <p data-testid="welcome-err" role="alert" className="e0-status e0-status-err mx-auto max-w-lg justify-start rounded-xl px-3 py-2 text-left">
+            <span>
+              <span className="block font-semibold">{t(welcomeErrorTitle(err.code))}</span>
+              <span className="mt-0.5 block font-normal">{t(downloadErrKey(err.code))} - {err.code ?? ""}</span>
+            </span>
+          </p>
+        )}
+        {catalogStore.status === "error" && !catalogStore.data && (
+          <p data-testid="welcome-catalog-err" className="text-xs text-muted">
+            {t(downloadErrKey(catalogStore.error?.code ?? null))}
+            <button
+              type="button"
+              data-testid="welcome-retry"
+              className="ml-2 underline"
+              onClick={() => void catalogStore.refresh()}
+            >
+              {t("models.retry")}
+            </button>
+          </p>
+        )}
+        <div>
+          <button
+            type="button"
+            data-testid="welcome-skip"
+            onClick={() => setDismissed(true)}
+            className="text-xs text-muted underline hover:text-fg"
+          >
+            {t("welcome.skip")}
+          </button>
+        </div>
+      </div>
+    </div>
+  );
+}
+
+async function errFrom(res: Response): Promise<ApiError> {
+  try {
+    const j = (await res.json()) as { error?: { code?: string; message?: string } };
+    return new ApiError(res.status, j?.error?.code ?? null, j?.error?.message ?? `HTTP ${res.status}`);
+  } catch {
+    return new ApiError(res.status, null, `HTTP ${res.status}`);
+  }
+}
+
+function welcomeErrorTitle(code: string | null): string {
+  if (code === "E-DL-DISK") return "welcome.errors.disk";
+  if (code === "E-MEM-LOAD") return "welcome.errors.memory";
+  if (code === "E-GPU-DEVICE" || code === "E-SRV-WORKER") return "welcome.errors.engine";
+  return "welcome.errors.generic";
+}
diff --git a/macos/app/src/service/codeMap.ts b/macos/app/src/service/codeMap.ts
new file mode 100644
index 0000000..c343301
--- /dev/null
+++ b/macos/app/src/service/codeMap.ts
@@ -0,0 +1,9 @@
+import { errorNextKey, hasKnownError } from "../chat/errorMap";
+import { downloadErrKey, hasDownloadErr } from "../models/errMap";
+
+export function requestCodeKey(code: string | null | undefined): string {
+  if (!code) return "service.log.noCode";
+  if (hasKnownError(code)) return errorNextKey(code);
+  if (hasDownloadErr(code)) return downloadErrKey(code);
+  return "chat.err.generic";
+}
diff --git a/macos/app/src/service/doctor.tsx b/macos/app/src/service/doctor.tsx
new file mode 100644
index 0000000..6fd8a93
--- /dev/null
+++ b/macos/app/src/service/doctor.tsx
@@ -0,0 +1,128 @@
+import { useCallback, useEffect, useState } from "react";
+import { useTranslation } from "react-i18next";
+import { Link } from "@tanstack/react-router";
+import { Stethoscope } from "lucide-react";
+
+import { edge0Fetch } from "../bridge/edge0Fetch";
+import { requestCodeKey } from "./codeMap";
+
+interface DoctorCheck {
+  id: string;
+  verdict: string;
+  detail: string;
+  code?: string | null;
+  next?: string | null;
+}
+interface DoctorReport {
+  overall: string;
+  checks: DoctorCheck[];
+}
+
+const ACTION_ROUTES: Record<string, { to: "/models"; labelKey: string }> = {
+  "E-DL-DISK": { to: "/models", labelKey: "service.doctor.goModels" },
+  "E-MODEL-HASH": { to: "/models", labelKey: "service.doctor.goModels" },
+};
+
+export function DoctorSection() {
+  const { t } = useTranslation();
+  const [report, setReport] = useState<DoctorReport | null>(null);
+  const [failed, setFailed] = useState(false);
+  const [loading, setLoading] = useState(true);
+
+  const fetchDoctor = useCallback(async () => {
+    setLoading(true);
+    setFailed(false);
+    try {
+      const res = await edge0Fetch("/v1/edge0/doctor");
+      if (!res.ok) throw new Error(`HTTP ${res.status}`);
+      setReport((await res.json()) as DoctorReport);
+    } catch {
+      setReport(null);
+      setFailed(true);
+    } finally {
+      setLoading(false);
+    }
+  }, []);
+
+  useEffect(() => {
+    void fetchDoctor();
+  }, [fetchDoctor]);
+
+  const verdictTone = (v: string) => (v === "pass" ? "ok" : v === "warn" ? "warn" : "err");
+  const cls = (tone: string) =>
+    tone === "ok" ? "bg-accent/15 text-accent" : tone === "warn" ? "bg-warn/15 text-warn" : "bg-err/15 text-err";
+
+  return (
+    <section className="e0-section overflow-hidden p-5" data-testid="svc-doctor">
+      <div className="mb-2 flex items-center justify-between">
+        <h3 className="flex items-center gap-2 text-sm font-semibold">
+          <Stethoscope size={14} aria-hidden />
+          {t("service.doctor.title")}
+        </h3>
+        {report && (
+          <span data-testid="doctor-overall" className={"rounded-full px-2 py-0.5 text-xs " + cls(verdictTone(report.overall))}>
+            {t(`service.doctor.overall.${report.overall}`)}
+          </span>
+        )}
+      </div>
+
+      {loading && !report && <p className="text-xs text-muted">{t("service.doctor.loading")}</p>}
+
+      {failed && (
+        <p data-testid="doctor-failed" className="text-xs text-muted">
+          {t("service.doctor.failed")}
+          <button
+            type="button"
+            data-testid="doctor-retry"
+            onClick={() => void fetchDoctor()}
+            className="ml-2 underline"
+          >
+            {t("service.doctor.retry")}
+          </button>
+        </p>
+      )}
+
+      {report && (
+        <ul className="space-y-1">
+          {report.checks.map((c) => {
+            const action = c.code ? ACTION_ROUTES[c.code] : undefined;
+            return (
+              <li
+                key={c.id}
+                data-testid={`doctor-item-${c.id}`}
+                className="flex items-baseline gap-2 border-t border-line/60 py-1.5 text-xs first:border-t-0"
+              >
+                <span className={"w-14 shrink-0 rounded-full px-1.5 py-0.5 text-center " + cls(verdictTone(c.verdict))}>
+                  {t(`service.doctor.verdict.${c.verdict}`)}
+                </span>
+                <span data-testid={`doctor-detail-${c.id}`} className="min-w-0 flex-1">
+                  <span className="text-muted">{c.id}</span> - {c.detail}
+                  {c.code && (
+                    <span data-testid={`doctor-code-${c.id}`} className="text-muted">
+                      {" "}
+                      - {c.code} {t(requestCodeKey(c.code))}
+                    </span>
+                  )}
+                  {c.next && (
+                    <span data-testid={`doctor-next-${c.id}`} className="block text-muted">
+                      {t("service.doctor.next")}: {c.next}
+                    </span>
+                  )}
+                </span>
+                {action && (
+                  <Link
+                    to={action.to}
+                    data-testid={`doctor-action-${c.id}`}
+                    className="e0-btn e0-btn-secondary shrink-0"
+                  >
+                    {t(action.labelKey)}
+                  </Link>
+                )}
+              </li>
+            );
+          })}
+        </ul>
+      )}
+    </section>
+  );
+}
diff --git a/macos/app/src/service/panel.tsx b/macos/app/src/service/panel.tsx
new file mode 100644
index 0000000..8350232
--- /dev/null
+++ b/macos/app/src/service/panel.tsx
@@ -0,0 +1,248 @@
+import { invoke } from "@tauri-apps/api/core";
+import { useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import { daemonStore } from "../bridge/state";
+import type { SystemSnapshot } from "../gen";
+import { isLoopback, useServiceState } from "./state";
+
+export function ServicePanel() {
+  const { t } = useTranslation();
+  const [svc] = useServiceState();
+  const [, force] = useState(0);
+  const snap = daemonStore.snapshot;
+  const phase = svc?.phase.phase ?? "probing";
+  const running = phase === "spawned" || phase === "adopted";
+
+  const [busy, setBusy] = useState<string | null>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [copyState, setCopyState] = useState<"masked" | "copied" | "no-credential">("masked");
+
+  async function act(name: string, fn: () => Promise<unknown>) {
+    setBusy(name);
+    setError(null);
+    try {
+      await fn();
+    } catch (e) {
+      const raw = typeof e === "string" ? e : e instanceof Error ? e.message : String(e);
+      const msg = raw === "no-credential" ? t("service.panel.noCredential") : t("service.panel.actionFailed");
+      setError(msg);
+      if (raw.includes("no-credential")) setCopyState("no-credential");
+    } finally {
+      setBusy(null);
+      force((n) => n + 1);
+    }
+  }
+
+  const inFlight = busy !== null;
+  const versionLine = snap
+    ? svc?.versionMismatch
+      ? t("service.panel.versionMismatch", {
+          daemon: snap.daemon.host_app_version ?? "?",
+          app: svc?.appVersion ?? "?",
+        })
+      : `v${snap.daemon.version}`
+    : "—";
+
+  return (
+    <section className="e0-section overflow-hidden p-5" data-testid="svc-panel">
+      <div className="mb-2 flex items-center justify-between">
+        <h3 className="text-sm font-semibold">{t("service.panel.title")}</h3>
+        <div className="flex items-center gap-2">
+          {running ? (
+            <button
+              type="button"
+              data-testid="svc-stop"
+              disabled={inFlight || phase === "adopted"}
+              title={phase === "adopted" ? t("service.panel.notOwned") : undefined}
+              onClick={() => void act("stop", () => invoke("service_stop"))}
+              className="e0-btn e0-btn-secondary"
+            >
+              {t("service.panel.stop")}
+            </button>
+          ) : (
+            (phase === "stopped" || phase === "exited") && (
+              <button
+                type="button"
+                data-testid="svc-start"
+                disabled={inFlight}
+                onClick={() => void act("start", () => invoke("service_start"))}
+                className="rounded-full bg-accent px-3 py-1 text-xs font-medium text-accentink hover:opacity-90 disabled:opacity-50"
+              >
+                {t("service.panel.start")}
+              </button>
+            )
+          )}
+        </div>
+      </div>
+
+      {inFlight && (
+        <p data-testid="svc-busy" className="mb-2 text-xs text-muted">
+          {t("service.panel.busy")}
+        </p>
+      )}
+      {error && (
+        <p data-testid="svc-error" role="alert" className="mb-2 rounded-lg bg-err/10 px-3 py-2 text-xs text-err">
+          {error}
+        </p>
+      )}
+
+      <dl
+        className="grid grid-cols-[auto_1fr] gap-x-4 gap-y-1 text-xs"
+        data-testid="svc-fields"
+      >
+        <dt className="text-muted">{t("service.panel.port")}</dt>
+        <dd data-testid="svc-port">{snap ? `${snap.daemon.bound}:${snap.daemon.port}` : "—"}</dd>
+        <dt className="text-muted">{t("service.panel.bind")}</dt>
+        <dd data-testid="svc-bind">{snap?.daemon.bound ?? "—"}</dd>
+        <dt className="text-muted">{t("service.panel.scope")}</dt>
+        <dd data-testid="svc-scope">
+          {snap ? `${snap.daemon.scope} - ${snap.daemon.started_by}` : "—"}
+        </dd>
+        <dt className="text-muted">{t("service.panel.uptime")}</dt>
+        <dd data-testid="svc-uptime">{snap ? `${snap.daemon.uptime_s}s` : "—"}</dd>
+        <dt className="text-muted">{t("service.panel.version")}</dt>
+        <dd data-testid="svc-version" className={svc?.versionMismatch ? "text-warn" : ""}>
+          {versionLine}
+        </dd>
+        <dt className="text-muted">{t("service.panel.home")}</dt>
+        <dd data-testid="svc-home" className="break-all">
+          {snap?.daemon.home ?? "—"}
+        </dd>
+      </dl>
+
+      
+      <div className="mt-3 flex items-center gap-2 text-xs" data-testid="svc-token-row">
+        <span className="text-muted">{t("service.panel.token")}</span>
+        {copyState === "no-credential" ? (
+          <span data-testid="svc-token-none">{t("service.panel.noCredential")}</span>
+        ) : (
+          <>
+            <span data-testid="svc-token-mask" className="tracking-widest text-muted">
+              {t("service.panel.mask")}
+            </span>
+            <button
+              type="button"
+              data-testid="svc-token-copy"
+              disabled={inFlight}
+              onClick={() =>
+                void act("copy", async () => {
+                  await invoke("token_copy");
+                  setCopyState("copied");
+                })
+              }
+              className="e0-btn e0-btn-secondary disabled:opacity-50"
+            >
+              {t("service.panel.copy")}
+            </button>
+            {copyState === "copied" && (
+              <span data-testid="svc-token-copied" className="text-warn">
+                {t("service.panel.copiedWarn")}
+              </span>
+            )}
+          </>
+        )}
+      </div>
+
+      <LanSwitch snap={snap} busy={inFlight} onAct={act} />
+      <AgentSwitch snap={snap} phase={phase} busy={inFlight} onAct={act} />
+    </section>
+  );
+}
+
+function LanSwitch({
+  snap,
+  busy,
+  onAct,
+}: {
+  snap: SystemSnapshot | null;
+  busy: boolean;
+  onAct: (name: string, fn: () => Promise<unknown>) => Promise<void>;
+}) {
+  const { t } = useTranslation();
+  const [confirming, setConfirming] = useState(false);
+  const lanOn = snap ? !isLoopback(snap.daemon.bound) : false;
+  const target = !lanOn;
+
+  return (
+    <div className="mt-3 space-y-1 border-t border-line pt-3 text-xs" data-testid="svc-lan-row">
+      <div className="flex items-center justify-between">
+        <span className="text-muted">{t("service.panel.lan")}</span>
+        {confirming ? (
+          <span className="flex items-center gap-2">
+            <span data-testid="svc-lan-confirm-text">{t("service.panel.lanConfirm")}</span>
+            <button
+              type="button"
+              data-testid="svc-lan-yes"
+              disabled={busy}
+              onClick={() => {
+                setConfirming(false);
+                void onAct("lan", () => invoke("lan_set", { on: target }));
+              }}
+              className="rounded bg-accent px-2 py-0.5 text-accentink disabled:opacity-50"
+            >
+              {t("service.panel.lanGoOn")}
+            </button>
+            <button
+              type="button"
+              data-testid="svc-lan-no"
+              onClick={() => setConfirming(false)}
+              className="rounded border border-line px-2 py-0.5"
+            >
+              {t("common.cancel")}
+            </button>
+          </span>
+        ) : (
+          <button
+            type="button"
+            role="switch"
+            aria-checked={lanOn}
+            data-testid="svc-lan-switch"
+            data-on={lanOn}
+            disabled={busy || !snap}
+            onClick={() => setConfirming(true)}
+            className="e0-btn e0-btn-secondary disabled:opacity-50"
+          >
+            {lanOn ? t("service.panel.lanOn") : t("service.panel.lanOff")}
+          </button>
+        )}
+      </div>
+      {lanOn && (
+        <p data-testid="svc-lan-guide" className="text-muted">
+          {t("service.panel.lanGuide")}
+        </p>
+      )}
+    </div>
+  );
+}
+
+function AgentSwitch({
+  snap,
+  phase,
+  busy,
+  onAct,
+}: {
+  snap: SystemSnapshot | null;
+  phase: string;
+  busy: boolean;
+  onAct: (name: string, fn: () => Promise<unknown>) => Promise<void>;
+}) {
+  const { t } = useTranslation();
+  if (!snap || (phase !== "spawned" && phase !== "adopted")) return null;
+  const onAgent = snap.daemon.scope === "launchagent";
+  return (
+    <div className="mt-3 flex items-center justify-between border-t border-line pt-3 text-xs" data-testid="svc-agent-row">
+      <span className="text-muted">{t("service.panel.agent")}</span>
+      <button
+        type="button"
+        data-testid="svc-agent-toggle"
+        data-to={onAgent ? "menu" : "agent"}
+        disabled={busy}
+        onClick={() => void onAct("agent", () => invoke("agent_set", { enable: !onAgent }))}
+        className="e0-btn e0-btn-secondary disabled:opacity-50"
+      >
+        {onAgent ? t("service.panel.agentToMenu") : t("service.panel.agentToLaunchd")}
+      </button>
+    </div>
+  );
+}
diff --git a/macos/app/src/service/reqlog.tsx b/macos/app/src/service/reqlog.tsx
new file mode 100644
index 0000000..a1a939b
--- /dev/null
+++ b/macos/app/src/service/reqlog.tsx
@@ -0,0 +1,118 @@
+import { useEffect, useState } from "react";
+import { useTranslation } from "react-i18next";
+
+import { daemonStore } from "../bridge/state";
+import type { Event } from "../gen";
+import { requestCodeKey } from "./codeMap";
+
+export interface RequestRow {
+  seq: number;
+  subject: string;
+  ts: string;
+  status: string;
+  code: string | null;
+  model: string;
+  latencyMs: number | null;
+  ttftMs: number | null;
+  tokens: number | null;
+  source: string;
+}
+
+const num = (v: unknown): number | null => (typeof v === "number" ? v : null);
+const str = (v: unknown): string => (typeof v === "string" ? v : "");
+
+export function rowFromEvent(e: Event): RequestRow {
+  const p = (e.payload ?? {}) as Record<string, unknown>;
+  return {
+    seq: e.seq,
+    subject: e.subject,
+    ts: e.ts,
+    status: str(p.status),
+    code: p.code == null ? null : str(p.code),
+    model: str(p.model),
+    latencyMs: num(p.latency_ms),
+    ttftMs: num(p.ttft_ms),
+    tokens: num(p.tokens),
+    source: str(p.source),
+  };
+}
+
+export function statusTone(status: string): "ok" | "warn" | "err" {
+  if (status === "done") return "ok";
+  if (status === "error") return "err";
+  return "warn";
+}
+
+export function RequestLogSection() {
+  const { t } = useTranslation();
+  const [, force] = useState(0);
+  useEffect(() => daemonStore.subscribe(() => force((n) => n + 1)), []);
+
+  const rows = daemonStore.requestsWindow().map(rowFromEvent);
+  return (
+    <section className="e0-section overflow-hidden p-5" data-testid="svc-log">
+      <div className="mb-1 flex items-baseline justify-between">
+        <h3 className="text-sm font-semibold">{t("service.log.title")}</h3>
+        <span data-testid="svc-log-count" className="text-xs text-muted">
+          {t("service.log.count", { n: rows.length })}
+        </span>
+      </div>
+      <p data-testid="svc-log-note" className="mb-2 text-xs text-muted">
+        {t("service.log.note")}
+      </p>
+      {rows.length === 0 ? (
+        <p data-testid="svc-log-empty" className="text-xs text-muted">
+          {t("service.log.empty")}
+        </p>
+      ) : (
+        <table className="w-full text-xs">
+          <thead>
+            <tr className="text-left text-muted">
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.status")}</th>
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.model")}</th>
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.latency")}</th>
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.tokens")}</th>
+              <th className="py-1 font-normal">{t("service.log.th.source")}</th>
+            </tr>
+          </thead>
+          <tbody>
+            {rows.map((r) => {
+              const tone = statusTone(r.status);
+              return (
+                <tr key={r.seq} data-testid={`svc-log-row-${r.subject}`} className="border-t border-line/60">
+                  <td className="py-1 pr-2">
+                    <span className="flex items-center gap-1">
+                      <span
+                        data-testid={`svc-log-dot-${r.subject}`}
+                        className={
+                          "inline-block size-2 rounded-full " +
+                          (tone === "ok" ? "bg-accent" : tone === "err" ? "bg-err" : "bg-warn")
+                        }
+                        aria-hidden
+                      />
+                      {t(`service.log.status.${r.status}`)}
+                      {r.code ? (
+                        <span data-testid={`svc-log-code-${r.subject}`} className="text-muted">
+                          {r.code} - {t(requestCodeKey(r.code))}
+                        </span>
+                      ) : null}
+                    </span>
+                  </td>
+                  <td className="py-1 pr-2">{r.model || "—"}</td>
+                  <td className="py-1 pr-2">
+                    {r.latencyMs ?? "—"}ms
+                    {r.ttftMs !== null && (
+                      <span className="text-muted"> - TTFT {r.ttftMs}ms</span>
+                    )}
+                  </td>
+                  <td className="py-1 pr-2">{r.tokens ?? "—"}</td>
+                  <td className="py-1">{t(`service.log.source.${r.source}`)}</td>
+                </tr>
+              );
+            })}
+          </tbody>
+        </table>
+      )}
+    </section>
+  );
+}
diff --git a/macos/app/src/service/state.ts b/macos/app/src/service/state.ts
new file mode 100644
index 0000000..530f1b6
--- /dev/null
+++ b/macos/app/src/service/state.ts
@@ -0,0 +1,47 @@
+import { invoke } from "@tauri-apps/api/core";
+import { useCallback, useEffect, useState } from "react";
+
+export interface ServiceHealth {
+  daemon_version: string;
+  host_app_version: string | null;
+  scope: string;
+  started_by: string;
+  uptime_s: number;
+  pid: number;
+}
+
+export interface ServicePhase {
+  phase: string;
+  health?: ServiceHealth;
+  code?: number | null;
+  detail?: string;
+  base?: string;
+  spawnedByUs?: boolean;
+}
+
+export interface ServiceState {
+  phase: ServicePhase;
+  appVersion: string;
+  versionMismatch: boolean;
+  spawnWasDev: boolean;
+  pinned: boolean;
+}
+
+export function isLoopback(bound: string): boolean {
+  return bound === "::1" || bound.startsWith("127.");
+}
+
+export function useServiceState(): [ServiceState | null, () => void] {
+  const [svc, setSvc] = useState<ServiceState | null>(null);
+  const refresh = useCallback(() => {
+    invoke<ServiceState>("service_state")
+      .then(setSvc)
+      .catch(() => undefined);
+  }, []);
+  useEffect(() => {
+    refresh();
+    const poll = setInterval(refresh, 2000);
+    return () => clearInterval(poll);
+  }, [refresh]);
+  return [svc, refresh];
+}
diff --git a/macos/app/src/state/threads.ts b/macos/app/src/state/threads.ts
new file mode 100644
index 0000000..d0dd1bc
--- /dev/null
+++ b/macos/app/src/state/threads.ts
@@ -0,0 +1,40 @@
+import type { ThreadRow } from "../chat/ipc";
+
+type Listener = () => void;
+
+class ThreadStore {
+  threads: ThreadRow[] = [];
+  activeId: string | null = null;
+  private listeners = new Set<Listener>();
+  private loaded = false;
+
+  subscribe(fn: Listener): () => void {
+    this.listeners.add(fn);
+    return () => this.listeners.delete(fn);
+  }
+
+  setActive(id: string | null): void {
+    this.activeId = id;
+    this.emit();
+  }
+
+  async refresh(): Promise<void> {
+    const { threadList } = await import("../chat/ipc");
+    try {
+      this.threads = await threadList();
+      this.loaded = true;
+    } finally {
+      this.emit();
+    }
+  }
+
+  get isLoaded(): boolean {
+    return this.loaded;
+  }
+
+  private emit(): void {
+    for (const l of this.listeners) l();
+  }
+}
+
+export const threadStore = new ThreadStore();
diff --git a/macos/app/src/styles.css b/macos/app/src/styles.css
new file mode 100644
index 0000000..1fefbe2
--- /dev/null
+++ b/macos/app/src/styles.css
@@ -0,0 +1,464 @@
+@import "tailwindcss";
+
+@custom-variant dark (&:where([data-theme="dark"], [data-theme="dark"] *));
+
+:root {
+  --e0-app: #f3f4f6;
+  --e0-surface: #ffffff;
+  --e0-surface-2: #f7f7f8;
+  --e0-sidebar: #eceef1;
+  --e0-text: #181a1f;
+  --e0-muted: #70747d;
+  --e0-border: #dfe2e7;
+  --e0-accent: #2f6feb;
+  --e0-accent-ink: #ffffff;
+  --e0-ok: #168052;
+  --e0-warn: #a66700;
+  --e0-err: #c53636;
+  --e0-focus: #2f6feb;
+  --e0-shadow: 0 12px 36px rgba(29, 34, 43, 0.08);
+  --e0-radius-lg: 18px;
+  --e0-radius-md: 12px;
+}
+[data-theme="dark"] {
+  --e0-app: #111317;
+  --e0-surface: #181b20;
+  --e0-surface-2: #20242a;
+  --e0-sidebar: #1b1e24;
+  --e0-text: #f1f3f5;
+  --e0-muted: #a2a8b2;
+  --e0-border: #30353d;
+  --e0-accent: #6ea3ff;
+  --e0-accent-ink: #111317;
+  --e0-ok: #55d69a;
+  --e0-warn: #efb456;
+  --e0-err: #ff8585;
+  --e0-focus: #8bb5ff;
+  --e0-shadow: 0 16px 42px rgba(0, 0, 0, 0.28);
+}
+
+html[data-font="sm"] { font-size: 14px; }
+html[data-font="md"] { font-size: 16px; }
+html[data-font="lg"] { font-size: 18px; }
+
+@theme inline {
+  --color-app: var(--e0-app);
+  --color-surface: var(--e0-surface);
+  --color-surface2: var(--e0-surface-2);
+  --color-sidebar: var(--e0-sidebar);
+  --color-fg: var(--e0-text);
+  --color-muted: var(--e0-muted);
+  --color-line: var(--e0-border);
+  --color-accent: var(--e0-accent);
+  --color-accentink: var(--e0-accent-ink);
+  --color-ok: var(--e0-ok);
+  --color-warn: var(--e0-warn);
+  --color-err: var(--e0-err);
+}
+
+body {
+  @apply bg-app text-fg antialiased;
+  margin: 0;
+  min-width: 320px;
+}
+
+button,
+input,
+textarea,
+select {
+  font: inherit;
+}
+
+button,
+a,
+input,
+textarea,
+select {
+  -webkit-tap-highlight-color: transparent;
+}
+
+button:focus-visible,
+a:focus-visible,
+input:focus-visible,
+textarea:focus-visible,
+select:focus-visible {
+  outline: 2px solid var(--e0-focus);
+  outline-offset: 2px;
+}
+
+/* Composer focus is the rounded shell's focus-within. This rule must stay unlayered so it wins over the global :focus-visible outline (component outline-none cannot). */
+.e0-composer-shell textarea:focus-visible {
+  outline: none;
+}
+
+@layer components {
+  .e0-app-shell {
+    display: flex;
+    gap: 12px;
+    height: 100vh;
+    min-height: 0;
+    max-height: 100vh;
+    padding: 12px;
+    background: var(--e0-app);
+    box-sizing: border-box;
+    overflow: hidden;
+  }
+
+  .e0-sidebar {
+    position: relative;
+    z-index: 30;
+    display: flex;
+    width: 258px;
+    height: 100%;
+    flex: 0 0 258px;
+    flex-direction: column;
+    overflow: hidden;
+    border: 1px solid var(--e0-border);
+    border-radius: var(--e0-radius-lg);
+    background: var(--e0-sidebar);
+    box-shadow: var(--e0-shadow);
+  }
+
+  .e0-main-frame {
+    position: relative;
+    height: 100%;
+    min-width: 0;
+    min-height: 0;
+    flex: 1;
+    overflow: hidden;
+    border: 1px solid var(--e0-border);
+    border-radius: var(--e0-radius-lg);
+    background: var(--e0-surface);
+    box-shadow: var(--e0-shadow);
+  }
+
+  .e0-page {
+    width: min(100%, 980px);
+    margin: 0 auto;
+    padding: 34px clamp(20px, 4vw, 56px) 44px;
+  }
+
+  .e0-page-header {
+    display: flex;
+    align-items: flex-start;
+    justify-content: space-between;
+    gap: 16px;
+    margin-bottom: 28px;
+  }
+
+  .e0-page-title {
+    margin: 0;
+    font-size: 1.35rem;
+    font-weight: 650;
+    letter-spacing: -0.01em;
+  }
+
+  .e0-page-kicker {
+    margin: 5px 0 0;
+    color: var(--e0-muted);
+    font-size: 0.8rem;
+    line-height: 1.5;
+  }
+
+  .e0-section {
+    border: 1px solid var(--e0-border);
+    border-radius: var(--e0-radius-md);
+    background: var(--e0-surface);
+    box-shadow: 0 4px 16px rgba(29, 34, 43, 0.04);
+  }
+
+  .e0-section-header {
+    display: flex;
+    align-items: flex-start;
+    justify-content: space-between;
+    gap: 16px;
+    padding: 18px 20px 12px;
+  }
+
+  .e0-section-title {
+    margin: 0;
+    font-size: 0.9rem;
+    font-weight: 650;
+  }
+
+  .e0-section-note {
+    margin: 4px 0 0;
+    color: var(--e0-muted);
+    font-size: 0.75rem;
+    line-height: 1.5;
+  }
+
+  .e0-btn {
+    display: inline-flex;
+    min-height: 32px;
+    align-items: center;
+    justify-content: center;
+    gap: 7px;
+    border: 1px solid transparent;
+    border-radius: 999px;
+    padding: 0 13px;
+    font-size: 0.75rem;
+    font-weight: 600;
+    transition: background-color 140ms ease, border-color 140ms ease, color 140ms ease, opacity 140ms ease;
+  }
+
+  .e0-btn:disabled {
+    cursor: not-allowed;
+    opacity: 0.48;
+  }
+
+  .e0-btn-primary {
+    background: var(--e0-accent);
+    color: var(--e0-accent-ink);
+  }
+
+  .e0-btn-primary:hover:not(:disabled) {
+    filter: brightness(0.96);
+  }
+
+  .e0-btn-secondary {
+    border-color: var(--e0-border);
+    background: var(--e0-surface);
+    color: var(--e0-text);
+  }
+
+  .e0-btn-secondary:hover:not(:disabled),
+  .e0-btn-ghost:hover:not(:disabled) {
+    background: var(--e0-surface-2);
+  }
+
+  .e0-btn-ghost {
+    color: var(--e0-muted);
+  }
+
+  .e0-btn-danger {
+    border-color: color-mix(in srgb, var(--e0-err) 35%, var(--e0-border));
+    color: var(--e0-err);
+  }
+
+  .e0-icon-btn {
+    display: inline-grid;
+    width: 32px;
+    height: 32px;
+    place-items: center;
+    border-radius: 10px;
+    color: var(--e0-muted);
+    transition: background-color 140ms ease, color 140ms ease;
+  }
+
+  .e0-icon-btn:hover {
+    background: color-mix(in srgb, var(--e0-text) 8%, transparent);
+    color: var(--e0-text);
+  }
+
+  .e0-status {
+    display: inline-flex;
+    min-height: 24px;
+    align-items: center;
+    gap: 6px;
+    border-radius: 999px;
+    padding: 0 9px;
+    font-size: 0.7rem;
+    font-weight: 600;
+  }
+
+  .e0-status-ok { background: color-mix(in srgb, var(--e0-ok) 13%, transparent); color: var(--e0-ok); }
+  .e0-status-warn { background: color-mix(in srgb, var(--e0-warn) 14%, transparent); color: var(--e0-warn); }
+  .e0-status-err { background: color-mix(in srgb, var(--e0-err) 13%, transparent); color: var(--e0-err); }
+
+  .e0-toggle-group {
+    display: inline-flex;
+    gap: 2px;
+    border: 1px solid var(--e0-border);
+    border-radius: 999px;
+    padding: 3px;
+    background: var(--e0-surface-2);
+  }
+
+  .e0-toggle {
+    border-radius: 999px;
+    padding: 5px 9px;
+    color: var(--e0-muted);
+    font-size: 0.7rem;
+    font-weight: 600;
+  }
+
+  .e0-toggle[data-state="on"] {
+    background: var(--e0-accent);
+    color: var(--e0-accent-ink);
+  }
+
+  .e0-divider {
+    border-top: 1px solid color-mix(in srgb, var(--e0-border) 72%, transparent);
+  }
+
+  .e0-mobile-toggle {
+    display: none;
+  }
+
+  .e0-chat-header {
+    display: flex;
+    min-height: 58px;
+    align-items: center;
+    justify-content: space-between;
+    gap: 16px;
+    border-bottom: 1px solid color-mix(in srgb, var(--e0-border) 72%, transparent);
+    padding: 0 28px;
+  }
+
+  .e0-chat-header-title {
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+    font-size: 0.8rem;
+    font-weight: 650;
+  }
+
+  .e0-chat-header-meta {
+    display: inline-flex;
+    align-items: center;
+    gap: 7px;
+    color: var(--e0-muted);
+    font-size: 0.7rem;
+  }
+
+  .e0-chat-empty {
+    display: grid;
+    min-height: 100%;
+    place-content: center;
+    gap: 8px;
+    padding: 32px;
+    text-align: center;
+  }
+
+  .e0-chat-column {
+    width: min(100%, 760px);
+    margin: 0 auto;
+  }
+
+  .e0-composer {
+    padding: 12px 22px 20px;
+  }
+
+  .e0-composer-shell {
+    width: min(100%, 760px);
+    margin: 0 auto;
+    overflow: visible;
+    border: 1px solid var(--e0-border);
+    border-radius: 20px;
+    background: var(--e0-surface);
+    box-shadow: 0 10px 28px rgba(29, 34, 43, 0.09);
+    transition: border-color 140ms ease, box-shadow 140ms ease;
+  }
+
+  .e0-composer-shell:focus-within {
+    border-color: color-mix(in srgb, var(--e0-accent) 55%, var(--e0-border));
+    box-shadow: 0 10px 30px color-mix(in srgb, var(--e0-accent) 13%, transparent);
+  }
+
+  .e0-composer-textarea {
+    min-height: 82px;
+    max-height: 220px;
+    padding: 17px 18px 8px;
+    line-height: 1.5;
+  }
+
+  .e0-composer-toolbar {
+    display: flex;
+    min-width: 0;
+    align-items: center;
+    gap: 8px;
+    padding: 6px 10px 10px;
+  }
+
+  .e0-send-btn {
+    display: inline-grid;
+    width: 34px;
+    height: 34px;
+    flex: 0 0 auto;
+    place-items: center;
+    border-radius: 999px;
+    background: var(--e0-accent);
+    color: var(--e0-accent-ink);
+    transition: transform 140ms ease, opacity 140ms ease;
+  }
+
+  .e0-send-btn:hover:not(:disabled) {
+    transform: translateY(-1px);
+  }
+
+  .e0-send-btn:disabled {
+    cursor: not-allowed;
+    opacity: 0.38;
+  }
+}
+
+@media (max-width: 760px) {
+  .e0-app-shell {
+    display: block;
+    padding: 0;
+  }
+
+  .e0-sidebar {
+    position: fixed;
+    inset: 10px auto 10px 10px;
+    width: min(280px, calc(100vw - 28px));
+    transform: translateX(calc(-100% - 20px));
+    transition: transform 180ms ease;
+  }
+
+  .e0-sidebar.e0-sidebar-open {
+    transform: translateX(0);
+  }
+
+  .e0-sidebar-backdrop {
+    position: fixed;
+    z-index: 20;
+    inset: 0;
+    background: rgba(13, 17, 23, 0.34);
+  }
+
+  .e0-main-frame {
+    height: 100vh;
+    min-height: 100vh;
+    border: 0;
+    border-radius: 0;
+  }
+
+  .e0-mobile-toggle {
+    display: inline-grid;
+  }
+
+  .e0-page {
+    padding: 26px 16px 32px;
+  }
+
+  .e0-page-header {
+    align-items: center;
+    margin-bottom: 22px;
+  }
+
+  .e0-chat-header {
+    min-height: 54px;
+    padding: 0 16px 0 58px;
+  }
+
+  .e0-chat-column {
+    width: 100%;
+  }
+
+  .e0-composer {
+    padding: 10px 12px 14px;
+  }
+}
+
+@media (prefers-reduced-motion: reduce) {
+  *,
+  *::before,
+  *::after {
+    scroll-behavior: auto !important;
+    transition-duration: 0.01ms !important;
+    animation-duration: 0.01ms !important;
+  }
+}
diff --git a/macos/app/tsconfig.json b/macos/app/tsconfig.json
new file mode 100644
index 0000000..ade52b7
--- /dev/null
+++ b/macos/app/tsconfig.json
@@ -0,0 +1,19 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "lib": ["ES2022", "DOM", "DOM.Iterable"],
+    "module": "ESNext",
+    "moduleResolution": "bundler",
+    "jsx": "react-jsx",
+    "strict": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "noFallthroughCasesInSwitch": true,
+    "skipLibCheck": true,
+    "resolveJsonModule": true,
+    "isolatedModules": true,
+    "noEmit": true,
+    "types": ["vite/client"]
+  },
+  "include": ["src"]
+}
diff --git a/macos/app/vite.config.ts b/macos/app/vite.config.ts
new file mode 100644
index 0000000..faa32f5
--- /dev/null
+++ b/macos/app/vite.config.ts
@@ -0,0 +1,20 @@
+import tailwindcss from "@tailwindcss/vite";
+import react from "@vitejs/plugin-react";
+import { defineConfig } from "vite";
+
+// Tauri expects a fixed port, no port fallback, and only TAURI_/VITE_ env passthrough.
+export default defineConfig({
+  plugins: [react(), tailwindcss()],
+  clearScreen: false,
+  server: {
+    port: 1420,
+    strictPort: true,
+    watch: { ignored: ["**/src-tauri/**"] },
+  },
+  envPrefix: ["VITE_", "TAURI_"],
+  build: {
+    target: "es2022",
+    outDir: "dist",
+    sourcemap: !!process.env.TAURI_ENV_DEBUG,
+  },
+});
diff --git a/macos/app/yarn.lock b/macos/app/yarn.lock
new file mode 100644
index 0000000..b32f51a
--- /dev/null
+++ b/macos/app/yarn.lock
@@ -0,0 +1,2284 @@
+# THIS IS AN AUTOGENERATED FILE. DO NOT EDIT THIS FILE DIRECTLY.
+# yarn lockfile v1
+
+
+"@ai-sdk/gateway@2.0.152":
+  version "2.0.152"
+  resolved "https://registry.yarnpkg.com/@ai-sdk/gateway/-/gateway-2.0.152.tgz#9300c19883734ed4c0c0be5c19c7a8144fbcfeea"
+  integrity sha512-ypuos9pHRawPlimfTC8wf0hCvp8Mx571zZfuDZvmECDYfaFze2qOo81yKIxacnQxnOekd6uV9ooMvFS+fB51dg==
+  dependencies:
+    "@ai-sdk/provider" "2.0.4"
+    "@ai-sdk/provider-utils" "3.0.37"
+    "@vercel/oidc" "3.1.0"
+
+"@ai-sdk/openai-compatible@1.0.54":
+  version "1.0.54"
+  resolved "https://registry.yarnpkg.com/@ai-sdk/openai-compatible/-/openai-compatible-1.0.54.tgz#c443e3b27960fdc01c6222be05fb4f60b364b84c"
+  integrity sha512-3wY1L21FLhfKN6I1nQeWgFzDCeBiRxXwU2guE9BCixSLfOPtkJ0vyF7zo7lbnAg1y7tz6RDPq/o+AJf6QSL55g==
+  dependencies:
+    "@ai-sdk/provider" "2.0.4"
+    "@ai-sdk/provider-utils" "3.0.37"
+
+"@ai-sdk/provider-utils@3.0.37":
+  version "3.0.37"
+  resolved "https://registry.yarnpkg.com/@ai-sdk/provider-utils/-/provider-utils-3.0.37.tgz#cb634d69fd8a844c93d903c2019e2762f99da487"
+  integrity sha512-mLx1SgE20xKQ87xME+9uG6pFlYPTqGDrMFrZ6XcytDnmbWGwmCIh2Xb9zFrfXUrKycSFRDet4Zovj388L9b1dg==
+  dependencies:
+    "@ai-sdk/provider" "2.0.4"
+    "@standard-schema/spec" "^1.0.0"
+    eventsource-parser "^3.0.6"
+    undici "^5.29.0"
+
+"@ai-sdk/provider@2.0.4":
+  version "2.0.4"
+  resolved "https://registry.yarnpkg.com/@ai-sdk/provider/-/provider-2.0.4.tgz#d450b7e6bf080e732108b0e79cb361590434669c"
+  integrity sha512-B0M50w0W43jTzZGyZrNwwaK1M79aKPw6PD/dQhkBR7vkR68yCINWBo2CbBLJrbrqYc46TaiSfsTiT+S9EmVl0w==
+  dependencies:
+    json-schema "^0.4.0"
+
+"@ai-sdk/react@2.0.261":
+  version "2.0.261"
+  resolved "https://registry.yarnpkg.com/@ai-sdk/react/-/react-2.0.261.tgz#1080005b2189e1fd15886959862480660c11d63d"
+  integrity sha512-5F4qIKUiDf2x8eVBAJvmNiMIhLRQLEqyaB1HL73WJUrF+ZPUVvcixXDmItB5R0AUx+7i+33SWJBWTO+4XVCaIg==
+  dependencies:
+    "@ai-sdk/provider-utils" "3.0.37"
+    ai "5.0.258"
+    swr "^2.2.5"
+    throttleit "2.1.0"
+
+"@babel/runtime@^7.29.7":
+  version "7.29.7"
+  resolved "https://registry.yarnpkg.com/@babel/runtime/-/runtime-7.29.7.tgz#12022450c45a4da6d8d8287b18a4ff2ddb23f768"
+  integrity sha512-Nq8OhGWiZIZGV6hLHoyAKLLcJihP/xFeBMGJoUrxTX2psI8dCifzLhZISFb+VWS3wFMRDmCGw5R+dOySCqPLhw==
+
+"@emnapi/core@^1.11.1":
+  version "1.11.3"
+  resolved "https://registry.yarnpkg.com/@emnapi/core/-/core-1.11.3.tgz#5e95348a42cd1e06f0b9aa380cd74091daa4d520"
+  integrity sha512-zLpS5asjEb7lq8jYLq37N6XKaE41DIexlY1rF/z4/tIl3wo13Sqm28fRyfIsKZD+NZ8mM5RoKkpW/rBcuoSZSg==
+  dependencies:
+    "@emnapi/wasi-threads" "1.2.3"
+    tslib "^2.4.0"
+
+"@emnapi/runtime@^1.11.1":
+  version "1.11.3"
+  resolved "https://registry.yarnpkg.com/@emnapi/runtime/-/runtime-1.11.3.tgz#84257ae3b0531eb2aec1ffa23d70700da007ba95"
+  integrity sha512-Xz4Tpyki7XyrpbUK1jR1AhdAdaXyhhY4lZ3neLodmhpuWfy2PAQN5B46sAiU4liOXGLkHypn/qU+jvfWSCYYLA==
+  dependencies:
+    tslib "^2.4.0"
+
+"@emnapi/wasi-threads@1.2.3", "@emnapi/wasi-threads@^1.2.2":
+  version "1.2.3"
+  resolved "https://registry.yarnpkg.com/@emnapi/wasi-threads/-/wasi-threads-1.2.3.tgz#c9bf72fd4be5b928aee894820e8d814ed73916e9"
+  integrity sha512-ELEBe8PsLvvJ6QMr0zLt8ffvOHW/dc1m3CEzNMg7aJUv3bMaoDtw2TXyDAwkYBuroxxuHEwhRTLJSe5sya547g==
+  dependencies:
+    tslib "^2.4.0"
+
+"@fastify/busboy@^2.0.0":
+  version "2.1.1"
+  resolved "https://registry.yarnpkg.com/@fastify/busboy/-/busboy-2.1.1.tgz#b9da6a878a371829a0502c9b6c1c143ef6663f4d"
+  integrity sha512-vBZP4NlzfOlerQTnba4aqZoMhE/a9HY7HRqoOPaETQcSQuWEIyZMHGfVu6w9wGtGK5fED5qRs2DteVCjOH60sA==
+
+"@jridgewell/gen-mapping@^0.3.5":
+  version "0.3.13"
+  resolved "https://registry.yarnpkg.com/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz#6342a19f44347518c93e43b1ac69deb3c4656a1f"
+  integrity sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==
+  dependencies:
+    "@jridgewell/sourcemap-codec" "^1.5.0"
+    "@jridgewell/trace-mapping" "^0.3.24"
+
+"@jridgewell/remapping@^2.3.5":
+  version "2.3.5"
+  resolved "https://registry.yarnpkg.com/@jridgewell/remapping/-/remapping-2.3.5.tgz#375c476d1972947851ba1e15ae8f123047445aa1"
+  integrity sha512-LI9u/+laYG4Ds1TDKSJW2YPrIlcVYOwi2fUC6xB43lueCjgxV4lffOCZCtYFiH6TNOX+tQKXx97T4IKHbhyHEQ==
+  dependencies:
+    "@jridgewell/gen-mapping" "^0.3.5"
+    "@jridgewell/trace-mapping" "^0.3.24"
+
+"@jridgewell/resolve-uri@^3.1.0":
+  version "3.1.2"
+  resolved "https://registry.yarnpkg.com/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz#7a0ee601f60f99a20c7c7c5ff0c80388c1189bd6"
+  integrity sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==
+
+"@jridgewell/sourcemap-codec@^1.4.14", "@jridgewell/sourcemap-codec@^1.5.0", "@jridgewell/sourcemap-codec@^1.5.5":
+  version "1.6.0"
+  resolved "https://registry.yarnpkg.com/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.6.0.tgz#f4c663e862f06dc98ca4d453862c46902789a18d"
+  integrity sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==
+
+"@jridgewell/trace-mapping@^0.3.24":
+  version "0.3.31"
+  resolved "https://registry.yarnpkg.com/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz#db15d6781c931f3a251a3dac39501c98a6082fd0"
+  integrity sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==
+  dependencies:
+    "@jridgewell/resolve-uri" "^3.1.0"
+    "@jridgewell/sourcemap-codec" "^1.4.14"
+
+"@napi-rs/wasm-runtime@^1.1.4":
+  version "1.2.4"
+  resolved "https://registry.yarnpkg.com/@napi-rs/wasm-runtime/-/wasm-runtime-1.2.4.tgz#ae38d9ea8483942054c5d82f00c0163374cd1e5f"
+  integrity sha512-AJxoUD2/15ESHbvpcyjU274nsAPLuOtPHCk0vKJM5pj//Fg/B1FXNWjPnXTT9PymCYYiHo4zPj0ZomXBKhoy7g==
+  dependencies:
+    "@tybys/wasm-util" "^0.10.3"
+
+"@opentelemetry/api@1.9.0":
+  version "1.9.0"
+  resolved "https://registry.yarnpkg.com/@opentelemetry/api/-/api-1.9.0.tgz#d03eba68273dc0f7509e2a3d5cba21eae10379fe"
+  integrity sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==
+
+"@oxc-project/types@=0.149.0":
+  version "0.149.0"
+  resolved "https://registry.yarnpkg.com/@oxc-project/types/-/types-0.149.0.tgz#328c1d19403980199c869676871db9ed46932ff8"
+  integrity sha512-Efcc+iF0j3Bf67YjEqIqWXbX5XddXoK/Mw4K1/JuXwRCZ8N16VR7iT23nlCc9XrveFVh/E5Rqs2StT0V8v9LdA==
+
+"@radix-ui/primitive@1.1.7":
+  version "1.1.7"
+  resolved "https://registry.yarnpkg.com/@radix-ui/primitive/-/primitive-1.1.7.tgz#0d0929d20299f60b0cd8a0deafad79f8048f9306"
+  integrity sha512-rqWnm76nYT8HoNNqEjpgJ7Pw/DrBj5iBTrmEPo6HTX5+VJyBNOqTdv4g89G63HuR5g0AaENoAcH7Is5fF2kZ8Q==
+
+"@radix-ui/react-collection@1.1.15":
+  version "1.1.15"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-collection/-/react-collection-1.1.15.tgz#6f045630d6f9778ab1b0c456ae97b474aec8a20b"
+  integrity sha512-9W+B9NPF0NaaPh/1NJd3+KqsnlLqU9H7T2rvww+fp+T/evVXdNAyYcnfRQZFOjkR1ajQp3yORlqnI8soawLvNA==
+  dependencies:
+    "@radix-ui/react-compose-refs" "1.1.5"
+    "@radix-ui/react-context" "1.2.2"
+    "@radix-ui/react-primitive" "2.1.10"
+    "@radix-ui/react-slot" "1.3.3"
+
+"@radix-ui/react-compose-refs@1.1.5":
+  version "1.1.5"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-compose-refs/-/react-compose-refs-1.1.5.tgz#883e642c5ec0ba393dd1bf586e7e772985770d87"
+  integrity sha512-+48PbAAbq3didjJxa+OaWY2ZwgAKsNiRGyeHKszblZMQ+kcpd9pAaT11cMkGEie0vsOi3QdeTE6d5Fe3Gn61kA==
+
+"@radix-ui/react-context@1.2.2":
+  version "1.2.2"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-context/-/react-context-1.2.2.tgz#fe2548e98465f5f9b776c5953de35d04a8919b6b"
+  integrity sha512-RHCUGwKHDr0hDGg4X7ma4JG4/+12qxw8rkh5QKdDldlCvtja6nUx1Ef/8HVrJze81lEsgLQlqjzjGNHantgnQA==
+
+"@radix-ui/react-direction@1.1.4":
+  version "1.1.4"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-direction/-/react-direction-1.1.4.tgz#e7b7049f3c3ab8e6e014c4600f9f1974a52c05f1"
+  integrity sha512-5pzg4FGQNpExhnhT2zlrP1wZFaYCd1K0nYWoFAdcYoYK868IEigqMX3B3f8yIoRlAhAeDWciLI6ZdCKHF9P4Vg==
+
+"@radix-ui/react-id@1.1.4":
+  version "1.1.4"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-id/-/react-id-1.1.4.tgz#e642714bdaf551c1bfacbe51508ed843a3a27980"
+  integrity sha512-TMQp2llA+RYn7JcjnrMnz7wN4pcVttPZnRZo52PLQsoLVKzNlVwUeHmfePgTgRluXFvlD3GD5g5MOVVTJCO0qA==
+  dependencies:
+    "@radix-ui/react-use-layout-effect" "1.1.4"
+
+"@radix-ui/react-primitive@2.1.10":
+  version "2.1.10"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-primitive/-/react-primitive-2.1.10.tgz#292b7476499d17227f337c8d1c59f270c683842a"
+  integrity sha512-MucOnzh6hR5mid6VpkbglRAMYMjKLqRnGBbjXkzjK52fuQDd1qbkx78a5P40mkcnVXJdEVxm26E9OPAiUq7nBg==
+  dependencies:
+    "@radix-ui/react-slot" "1.3.3"
+
+"@radix-ui/react-progress@^1.1.16":
+  version "1.1.16"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-progress/-/react-progress-1.1.16.tgz#250b03532573ba603f074b7702388a9427a5bcc4"
+  integrity sha512-5XnomAsoZZCY+KNTxbIghpGqPruZvKFNlvcAljVAOdDRDsH4/OZQxhtwo5wdtoDM5R6MhJBb2sPnDuRFep3lzg==
+  dependencies:
+    "@radix-ui/react-context" "1.2.2"
+    "@radix-ui/react-primitive" "2.1.10"
+
+"@radix-ui/react-roving-focus@1.1.19":
+  version "1.1.19"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-roving-focus/-/react-roving-focus-1.1.19.tgz#2abc0676448f4ebae8c43a61f9531aebe0f5c823"
+  integrity sha512-V9jI6hDjT7l3jsCQD9bLNvDLM3tH/gdbOTp7Tefp3hbbgCGQoK7tUvrWiRlcoBHIZ809ElXwNQwVo0B98LuTXQ==
+  dependencies:
+    "@radix-ui/primitive" "1.1.7"
+    "@radix-ui/react-collection" "1.1.15"
+    "@radix-ui/react-compose-refs" "1.1.5"
+    "@radix-ui/react-context" "1.2.2"
+    "@radix-ui/react-direction" "1.1.4"
+    "@radix-ui/react-id" "1.1.4"
+    "@radix-ui/react-primitive" "2.1.10"
+    "@radix-ui/react-use-callback-ref" "1.1.4"
+    "@radix-ui/react-use-controllable-state" "1.2.6"
+    "@radix-ui/react-use-is-hydrated" "0.1.3"
+    "@radix-ui/react-use-layout-effect" "1.1.4"
+
+"@radix-ui/react-slot@1.3.3", "@radix-ui/react-slot@^1.3.3":
+  version "1.3.3"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-slot/-/react-slot-1.3.3.tgz#ccc8c8936fff12f7e324d6917fe5167357097676"
+  integrity sha512-qx7oqnYbxnK9kYI9m317qmFmEgo6ywqWvbTogdj7cL9p3/yx4M48p7Rnw5z3H890cL/ow/EeWJsuTykeZVXP5Q==
+  dependencies:
+    "@radix-ui/react-compose-refs" "1.1.5"
+
+"@radix-ui/react-toggle-group@^1.1.19":
+  version "1.1.19"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-toggle-group/-/react-toggle-group-1.1.19.tgz#cefc32db4e117e5a15812bb944676a4d14d9e59b"
+  integrity sha512-OtnwuSVjd1Ofi+AdnvhsjQdyuhCDwYs1w9RyB5BN/OavXOVQo42SYqQjwUnbPnaiPFBpQ9aX70dWeee+v2oBLA==
+  dependencies:
+    "@radix-ui/primitive" "1.1.7"
+    "@radix-ui/react-context" "1.2.2"
+    "@radix-ui/react-direction" "1.1.4"
+    "@radix-ui/react-primitive" "2.1.10"
+    "@radix-ui/react-roving-focus" "1.1.19"
+    "@radix-ui/react-toggle" "1.1.18"
+    "@radix-ui/react-use-controllable-state" "1.2.6"
+
+"@radix-ui/react-toggle@1.1.18":
+  version "1.1.18"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-toggle/-/react-toggle-1.1.18.tgz#3b464b235e76193142d9ea276611a9c8ef3babc7"
+  integrity sha512-7lonPlKfSacd20GlOBx2ltuVKz9oqWYZz+oMQyOltw6t1y2nyftj2ZmwwUHYn49kqfDWcp8dNZm5NgV+5Z+mug==
+  dependencies:
+    "@radix-ui/primitive" "1.1.7"
+    "@radix-ui/react-primitive" "2.1.10"
+    "@radix-ui/react-use-controllable-state" "1.2.6"
+
+"@radix-ui/react-use-callback-ref@1.1.4":
+  version "1.1.4"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-use-callback-ref/-/react-use-callback-ref-1.1.4.tgz#1e55d37148e60a3f47de4bc7ea7cff6fa3f425ae"
+  integrity sha512-R6OUY2e2fA6Yn6s+VSx5KBV6Nx8LQEhu+cz7LCej18rQ1HLyg9PSC9jP/ZNx0o6FAIK9c0F1kHylzSxKsdlkrQ==
+
+"@radix-ui/react-use-controllable-state@1.2.6":
+  version "1.2.6"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-use-controllable-state/-/react-use-controllable-state-1.2.6.tgz#0f109eff004c3c1ff5f9709570f45dee9955513e"
+  integrity sha512-uEQJGT97ZA/TgP/Hydw47lHu+/vQj6z/0jA+WeTbK1o9Rx45GImjpD0tc3W5ad3D6XTSR6e1yEO0FvGq6WQfVQ==
+  dependencies:
+    "@radix-ui/primitive" "1.1.7"
+    "@radix-ui/react-use-effect-event" "0.0.5"
+    "@radix-ui/react-use-layout-effect" "1.1.4"
+
+"@radix-ui/react-use-effect-event@0.0.5":
+  version "0.0.5"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-use-effect-event/-/react-use-effect-event-0.0.5.tgz#7732a4ae6bc032c7f654e658203df92d3520f8d1"
+  integrity sha512-7cshFL8HGS/7HEiHH+9kL9HBwp2sa9yX18Knwek6KYWmXwM7pegMgta2AXMQKI+rq3JnfSj9x8wYqFMTdG1Jgg==
+  dependencies:
+    "@radix-ui/react-use-layout-effect" "1.1.4"
+
+"@radix-ui/react-use-is-hydrated@0.1.3":
+  version "0.1.3"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-use-is-hydrated/-/react-use-is-hydrated-0.1.3.tgz#f0ffa47394726046ba6c6d0b7c97239f48c8df45"
+  integrity sha512-umO/aJ+82CpOnhDZUTbILCQf7kU/g0iv+oGs/Q8jw7IkhWBzaEP4sA268PhFAJTFetbwp3ICc6ktpI4TqtxcIw==
+
+"@radix-ui/react-use-layout-effect@1.1.4":
+  version "1.1.4"
+  resolved "https://registry.yarnpkg.com/@radix-ui/react-use-layout-effect/-/react-use-layout-effect-1.1.4.tgz#51a0bbc342315070983383b3f8f00061083782a8"
+  integrity sha512-K20DkRkUwDnxEYMBPcg3Y6voLkEy5p5QQmszZgLngKKiC7dzBR/aEuK3w1qlx2JWDUNH6FluahYdgR3BP+QbYw==
+
+"@rolldown/binding-android-arm-eabi@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-android-arm-eabi/-/binding-android-arm-eabi-1.2.8.tgz#55f0a8e97eb87a0873ea1466d2b4af9510739589"
+  integrity sha512-tN5aztYkKCte4i5SIrrz5yK/HMjEuCqCSCJa418jOV8tZ1cBY3YF2otxB1ktPxzsLA1BeTqwapK0bfjxNvHJVw==
+
+"@rolldown/binding-android-arm64@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-android-arm64/-/binding-android-arm64-1.2.8.tgz#3b783ee41120b2fefac41bd8e6d711f4ce7f82a7"
+  integrity sha512-dIYTWl9XprMUiQFoc55KUyk/oS8SKYH3zFl0LTR7RT0Xj4hgSVyuJcroH8JUu8RcpF8fTB6E0aOwCkZoYPcDSQ==
+
+"@rolldown/binding-darwin-arm64@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.2.8.tgz#56c27646e8faae70aa56065a3b686b1d7de9d8b4"
+  integrity sha512-PCSDQGXD2IyTEFrcgPyBM8jJuGmrbCMuoIOXdbEGVemruKACXoLQJrb+A45Z0L5t1RQkdfJprAYPkikbh7dzdA==
+
+"@rolldown/binding-darwin-x64@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.2.8.tgz#3b9da1032c897ce191cb23601dbf9887ce6895ae"
+  integrity sha512-Uk7lRsGhPFHVX/sAUC6D5H9Ol30dFHd6iquokll2th3LpdJ3F5CzQB+7DHn0Ri2mG+U7k2zXiPHDrwZenXhwSA==
+
+"@rolldown/binding-freebsd-x64@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.2.8.tgz#43d4257704b15204aa42059e36a638b63c0e402d"
+  integrity sha512-DjszaTEVogPqA5bYzsEeqDCQxbcp2fexQwKcRspYji2yzR68fCf+e4fx6kBSRDwX5/brZaHw/hWS9+A/+/w9sQ==
+
+"@rolldown/binding-linux-arm-gnueabihf@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.2.8.tgz#d95b9fe9f04f2278cb2eef98e2cf20199c3fe5d0"
+  integrity sha512-zmwa7FTmdzB6aaEEuuls18H6Ap5JmJPSoPTuXixeJZV6tG40SyLkApQtz1g8ptZtiEKqj9OM0oNLPh1AgvE31Q==
+
+"@rolldown/binding-linux-arm64-gnu@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.2.8.tgz#07f33840b650b6a3f5954ccfc43b13629374182a"
+  integrity sha512-KdYQDPHwJVnbFwdTGMgxsI9SqblBlz6STGM+w1We/d5B8OWWidYH0MwkU/uA1wM5fIpO2MkOVxXrNzzuZhw9ew==
+
+"@rolldown/binding-linux-arm64-musl@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.2.8.tgz#357922929ec19c05f6d8eabd0899a053bab1a08b"
+  integrity sha512-jFJTifHnNPY+yzOoNZQfSIysrVyXzEQPhPnOUjmD1bcQGHH6s7c8cViKWar8YplQImE5N9JRqMCLrM2CdxOrZA==
+
+"@rolldown/binding-linux-ppc64-gnu@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.2.8.tgz#0bc7b29262d3022d8cf57876d70d819387bd3dc9"
+  integrity sha512-FhiOziBDWPBjbcmRzfLyIJnaP7AVMFXT7YCXPjXxj7wKU3vx24RjrCNN/zjvVa+N2vVoHJwCoUBvsrN/DG3zIA==
+
+"@rolldown/binding-linux-s390x-gnu@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.2.8.tgz#3a3d653becc48975025b8922056822662b7d4414"
+  integrity sha512-WnHfADMzOV2Y55wlx1hzzQnar/wDt/VdvWSD99r18Mz9ylNieIGOkRx3UV21h7m/eJvjySYJkO26VvGNFkwsIQ==
+
+"@rolldown/binding-linux-x64-gnu@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.2.8.tgz#c14ec9af7dbc2b1e459b36e292d57b27565f6737"
+  integrity sha512-H9tRr5ibfXFVLxbPOseVewewFpl28zcEdjRDt2FTUZU7odxP0gEv1ki4/kGmcGOh78oRwZuuQllGLZ9zTJp84g==
+
+"@rolldown/binding-linux-x64-musl@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.2.8.tgz#40c2888d1ac87d9e39c6fc6799c060510e7a4d69"
+  integrity sha512-UefiqfM3D6IVNlZ8tSGs9+Ejjud2T+oxO0IHADU45Y+lyEjD2dVFyZHbkfX0LUb5Zugo/oIv1eCO/KVYhgYJYA==
+
+"@rolldown/binding-openharmony-arm64@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.2.8.tgz#3e021f4e77283ca111c709d3ed6775419f471056"
+  integrity sha512-637Ke4kWSy6rp9cxQ9gMOXlxPgIw/c1beASV4M//3+9I4uwBVOOl74G+e3zyU3u19U7RkRl/HuewixZ/Z6+Rjg==
+
+"@rolldown/binding-win32-arm64-msvc@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.2.8.tgz#32b6c6a335ad14f97025bb2bc90135222405bb76"
+  integrity sha512-xWBkPOF1Q9k/Gv1nQXnVdLxKu74jXppuOM4Z3mnypVUJJJwLsMl7hNJGRAUJoG8A5MgOI1ACKM+wBFxSJzKy4A==
+
+"@rolldown/binding-win32-x64-msvc@1.2.8":
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.2.8.tgz#df4e965896c61411b96566abf345f46141f498c4"
+  integrity sha512-uz2ZvfgXbxqNwijjjbxrnvALwpyODDcgc1T1N8N3rf/DXKQmaFwmB4LX4yyjggpwN2obdQLb2rgirX5ffCWYng==
+
+"@rolldown/pluginutils@^1.0.0", "@rolldown/pluginutils@^1.0.1":
+  version "1.0.1"
+  resolved "https://registry.yarnpkg.com/@rolldown/pluginutils/-/pluginutils-1.0.1.tgz#e3fcee093fbb5ce765e1ad088ff4de2889f6f9be"
+  integrity sha512-2j9bGt5Jh8hj+vPtgzPtl72j0yRxHAyumoo6TNfAjsLB04UtpSvPbPcDcBMxz7n+9CYB0c1GxQFxYRg2jimqGw==
+
+"@standard-schema/spec@^1.0.0":
+  version "1.1.0"
+  resolved "https://registry.yarnpkg.com/@standard-schema/spec/-/spec-1.1.0.tgz#a79b55dbaf8604812f52d140b2c9ab41bc150bb8"
+  integrity sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w==
+
+"@streamdown/math@^1.0.2":
+  version "1.0.2"
+  resolved "https://registry.yarnpkg.com/@streamdown/math/-/math-1.0.2.tgz#b1bce87f47a045ef5d8f8b421ca456cddab0ce00"
+  integrity sha512-r8Ur9/lBuFnzZAFdEWrLUF2s/gRwRRRwruqltdZibyjbCBnuW7SJbFm26nXqvpJPW/gzpBUMrBVBzd88z05D5g==
+  dependencies:
+    katex "^0.16.27"
+    rehype-katex "^7.0.1"
+    remark-math "^6.0.0"
+
+"@tailwindcss/node@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/node/-/node-4.3.3.tgz#38ff04309ff036ea3589a7bad9069c44ec9d3883"
+  integrity sha512-/T8IKEsf9VTU6tLjgC7+sv2mOPtQxzE2jMw7u4Tt40Tx+QSZxpzh95/H6cMKoja9XuW7iMdLJYBB0o9G1CaAgg==
+  dependencies:
+    "@jridgewell/remapping" "^2.3.5"
+    enhanced-resolve "^5.24.1"
+    jiti "^2.7.0"
+    lightningcss "1.32.0"
+    magic-string "^0.30.21"
+    source-map-js "^1.2.1"
+    tailwindcss "4.3.3"
+
+"@tailwindcss/oxide-android-arm64@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-android-arm64/-/oxide-android-arm64-4.3.3.tgz#848f93034155daf7892185028dfb18143bfc7d07"
+  integrity sha512-Y85A2gmPSkl5Ve5qR86GL4HT509cFqQh1aes9p3sSkyTPwt0Pppf3GkwGe4JPACcRYjgJIEhQgM6dBClnr0NYw==
+
+"@tailwindcss/oxide-darwin-arm64@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-darwin-arm64/-/oxide-darwin-arm64-4.3.3.tgz#5c779e32c1c361beb75136fd8e2c627fcb9a5b28"
+  integrity sha512-BiaWatpBcERQFDlOjRDpIVXuFK5PJez5SA4JMg6VYZdBYU+qKfV/vqjcIs+IYmtitf1xYQZTwXvU/8y4lfZUGw==
+
+"@tailwindcss/oxide-darwin-x64@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-darwin-x64/-/oxide-darwin-x64-4.3.3.tgz#ba292ff52fa3264139f7fe7874719ce0a4666875"
+  integrity sha512-fAeUqfV5ndhxRwai8cXGzdLvul9utWOmeTkv69unv4ZXixjn61Z+p9lCWdwOwA3TYboG3BwdVuN/RDjhBRl0mw==
+
+"@tailwindcss/oxide-freebsd-x64@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-freebsd-x64/-/oxide-freebsd-x64-4.3.3.tgz#07e589487b9a636235a4ade48cefc1f9eeeec57f"
+  integrity sha512-iyf5bV6+wnAlflVeEy7R25dupxTNECZN5QMI0qNT6eT+EgaGdZcKhGkr5SdoaWiLJ3spLqIY9VCeSGrwmtg4kw==
+
+"@tailwindcss/oxide-linux-arm-gnueabihf@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-linux-arm-gnueabihf/-/oxide-linux-arm-gnueabihf-4.3.3.tgz#0c8abc228d9e19e0065ddb707eba52e21855b3ff"
+  integrity sha512-aAYUprJAJQWWbRrPvtjdroZ56Md+JM8pMiopS6xGEwDfLhqj+2ver2p4nU4Mb3CRqcMmNBjo8KkUgcxhkzVQGQ==
+
+"@tailwindcss/oxide-linux-arm64-gnu@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-linux-arm64-gnu/-/oxide-linux-arm64-gnu-4.3.3.tgz#5067dc7afd15d2b97adc77a91177dc4bec8d1b8b"
+  integrity sha512-nDxldcEENOxZRzC2uu9jrutZdAAQtb+8WWDCSnWL1zvBk1+FN+x6MtDViPB5AJMfttVCUhehGWus3XBPgatM/w==
+
+"@tailwindcss/oxide-linux-arm64-musl@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-linux-arm64-musl/-/oxide-linux-arm64-musl-4.3.3.tgz#29ae5f279be2ce64db368711985b6ad8e4562e57"
+  integrity sha512-Md44bD6veX/PC5iyF8cDVnw4HBIANZepRZZ7a8DQOvkfo5WUBwcp6iAuCUz23u+4SUkhJlD3eL7hNdW8ezd/kA==
+
+"@tailwindcss/oxide-linux-x64-gnu@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-linux-x64-gnu/-/oxide-linux-x64-gnu-4.3.3.tgz#7d693b40f69875744b499481359b227fd3ac3a3c"
+  integrity sha512-tx7us1muwOKAKWao2v/GaafFeQboE6aj88vC6ziN2NCGcRm8gWUhwjzg+YdVB1e4boAtdtma4L43onunI6NS4w==
+
+"@tailwindcss/oxide-linux-x64-musl@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-linux-x64-musl/-/oxide-linux-x64-musl-4.3.3.tgz#6a55d664f47c5ff9ad3f46bf05fe07d410b8cfd8"
+  integrity sha512-SJxX60smvHgasZoBy11dX6YRjXJFovwWBoedhbQPOBzgFWBHGB+TVPWB9BxzR7TTxU8FQZAI2AyiNCMzFm8Img==
+
+"@tailwindcss/oxide-wasm32-wasi@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-wasm32-wasi/-/oxide-wasm32-wasi-4.3.3.tgz#408a9bc620e68f1aaf0dfea33da7bd3b65f9e2ef"
+  integrity sha512-jx1+rPhY/5Ympkktd656HBWEBLxP7dH06losBLjjf5vgCODXvi9KhtftWcMIwTFIDqBr7cRnQkdLnAG+IOlGvQ==
+  dependencies:
+    "@emnapi/core" "^1.11.1"
+    "@emnapi/runtime" "^1.11.1"
+    "@emnapi/wasi-threads" "^1.2.2"
+    "@napi-rs/wasm-runtime" "^1.1.4"
+    "@tybys/wasm-util" "^0.10.2"
+    tslib "^2.8.1"
+
+"@tailwindcss/oxide-win32-arm64-msvc@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.3.3.tgz#3d582b00180fa4680e0b6c90be69fa5f4a209092"
+  integrity sha512-3rc292Ca2ceK6Ulcc/bAVnTs/3nDtoPhyEKlgPv+yQJQi/JS/AMJlqzxvlDacL1nekbrcf6bTqp/jV4qgnPxNQ==
+
+"@tailwindcss/oxide-win32-x64-msvc@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide-win32-x64-msvc/-/oxide-win32-x64-msvc-4.3.3.tgz#ce4c663fef3b4fde67611a4c1391866f8c0aa79b"
+  integrity sha512-yJ0pwIVc/nYeGoV02WtsN8KYyLQv7kyI2wDnkezyJlGGjkd4QLwDGAwl47YpPJeuI0M0ObaXGSPjvWDPeTPggw==
+
+"@tailwindcss/oxide@4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/oxide/-/oxide-4.3.3.tgz#6266109d025cfcb04f8e4c58954a6f630a415b7f"
+  integrity sha512-krXjAikiaFSPaK/FkAQT5UTx3VormQaiZ5hBFlJZ9UFQGB/rwg1MZIhHAG9smMQRTdyJxP6Qt5MwMtdyU5FWrA==
+  optionalDependencies:
+    "@tailwindcss/oxide-android-arm64" "4.3.3"
+    "@tailwindcss/oxide-darwin-arm64" "4.3.3"
+    "@tailwindcss/oxide-darwin-x64" "4.3.3"
+    "@tailwindcss/oxide-freebsd-x64" "4.3.3"
+    "@tailwindcss/oxide-linux-arm-gnueabihf" "4.3.3"
+    "@tailwindcss/oxide-linux-arm64-gnu" "4.3.3"
+    "@tailwindcss/oxide-linux-arm64-musl" "4.3.3"
+    "@tailwindcss/oxide-linux-x64-gnu" "4.3.3"
+    "@tailwindcss/oxide-linux-x64-musl" "4.3.3"
+    "@tailwindcss/oxide-wasm32-wasi" "4.3.3"
+    "@tailwindcss/oxide-win32-arm64-msvc" "4.3.3"
+    "@tailwindcss/oxide-win32-x64-msvc" "4.3.3"
+
+"@tailwindcss/vite@^4.3.3":
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/@tailwindcss/vite/-/vite-4.3.3.tgz#1f4037badcb5e6c6bb574d3d8a0172c83f26b97f"
+  integrity sha512-yYU8cogLeSh/ms2jh8Fj7jaba/EWa7Ja6GoUqYZaraEuCI5YS6ms6ObZgjjedm+jm6XZjdNRWBpPP6Z86oOxcw==
+  dependencies:
+    "@tailwindcss/node" "4.3.3"
+    "@tailwindcss/oxide" "4.3.3"
+    tailwindcss "4.3.3"
+
+"@tanstack/history@1.162.4":
+  version "1.162.4"
+  resolved "https://registry.yarnpkg.com/@tanstack/history/-/history-1.162.4.tgz#bbb72e9d2ad8885fc7beafcbd97ed34b2cd38d1b"
+  integrity sha512-utTS5L2OkeYUzXGohL1Z8sefu1GLNOJcxe8Hd6iIdc/Xo1K1nDB2JEp4iSFhvYh33xKC9V91TxrS8qfrpoKobQ==
+
+"@tanstack/react-router@^1.170.36":
+  version "1.170.36"
+  resolved "https://registry.yarnpkg.com/@tanstack/react-router/-/react-router-1.170.36.tgz#5db32d9d61a1d9a75360127961630d58f9362b31"
+  integrity sha512-AbgcI/xO8AJY1r6y5tF7TZyHH0VKm0XcfY3QVIVCrHq8k04Hkj1N4nfLL+1xCGC3VnMI3AYGAVWkatnqE2dxAQ==
+  dependencies:
+    "@tanstack/history" "1.162.4"
+    "@tanstack/react-store" "^0.11.0"
+    "@tanstack/router-core" "1.171.30"
+    isbot "^5.1.22"
+
+"@tanstack/react-store@^0.11.0":
+  version "0.11.1"
+  resolved "https://registry.yarnpkg.com/@tanstack/react-store/-/react-store-0.11.1.tgz#e1e58036ffffa58f5b5d8826c2389f97502fac2c"
+  integrity sha512-HaIGKI3YLmjBYIvy5DFDY23oNaYZIsTZfngey07Uh5iLVJgM3bIGCnZeOFOqzjFld9JHWcaHJnasD/bKoGKwJQ==
+  dependencies:
+    "@tanstack/store" "0.11.1"
+    use-sync-external-store "^1.6.0"
+
+"@tanstack/router-core@1.171.30":
+  version "1.171.30"
+  resolved "https://registry.yarnpkg.com/@tanstack/router-core/-/router-core-1.171.30.tgz#315d91701abb4385fa5b4e85f9c1b47138e36dad"
+  integrity sha512-wtS1yCAb/k5CJBauvrCyqf2n4/kRc3aiz81tFBJ1QXH4Utsu81xFTMq15AdmNWgWmt2UstUSMfWi4Ua7ph2F+w==
+  dependencies:
+    "@tanstack/history" "1.162.4"
+    cookie-es "^3.0.0"
+    seroval "^1.6.2"
+    seroval-plugins "^1.6.2"
+
+"@tanstack/store@0.11.1":
+  version "0.11.1"
+  resolved "https://registry.yarnpkg.com/@tanstack/store/-/store-0.11.1.tgz#063f634fe3cc925ad83704581a7be75301bec612"
+  integrity sha512-mzTOBhypOuDJAy/D8n2MfUZ1HFkXnmSETviRyhqEC8LUE7/IZQExOTxMANj3KjTofYTkFNpBY67qaVrT41YccA==
+
+"@tauri-apps/api@^2.11.1":
+  version "2.11.1"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/api/-/api-2.11.1.tgz#cd6b13fc26403ca095a02e39ecdbec8048d2872d"
+  integrity sha512-M2FPuYND2m+wh5hfW9ZpSdxMPdEJovPBWwoHJmwUpysTYNHaOkVFN419m/K0LIgjb/7KU2vBgsUepJWugQCvAA==
+
+"@tauri-apps/cli-darwin-arm64@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.11.4.tgz#2884e6f2e62318e4ca1db0d6d3be8380dcb6c542"
+  integrity sha512-1ryOF3ZhpZ/nemHV5zVwBQBz9jDGKmKPvWPADOhc83ig0P4bMc2iER4NbC6r9sjeIZ6RVQ4g3RZIYvezhcl4TQ==
+
+"@tauri-apps/cli-darwin-x64@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-darwin-x64/-/cli-darwin-x64-2.11.4.tgz#d9e2d92b15ee2ff0b7819ea5c6df398f2fc754ba"
+  integrity sha512-uFsGQAAfuyz1k/yGLmkWfkBlgKAqZfxqlHmLWx81QU27RJWfmbNHCIq8T8w1e+VClleIuZUjpHWfoE4E3DLo3A==
+
+"@tauri-apps/cli-linux-arm-gnueabihf@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-linux-arm-gnueabihf/-/cli-linux-arm-gnueabihf-2.11.4.tgz#af14ae237d37a403acdb56c7421476b65057db27"
+  integrity sha512-IaHZn5CdBL21oUmjiVOS1ctw6Ip1O0pjp70FwOWmYz1myWe0SY96ZIj2FYf7pT0m8bI2h/hrs5ZbEXXh44/MkQ==
+
+"@tauri-apps/cli-linux-arm64-gnu@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-linux-arm64-gnu/-/cli-linux-arm64-gnu-2.11.4.tgz#fe521657fb011eb9a5236c5c9d0fa63ebce39dd4"
+  integrity sha512-N41/ukTRVe6XSuUTESuFdGeOW2i7k62tK+6gHK5Kd5/q5RPvvi19GaWAVPPb9u95HSGmTChSolBfzynUsssFaA==
+
+"@tauri-apps/cli-linux-arm64-musl@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-linux-arm64-musl/-/cli-linux-arm64-musl-2.11.4.tgz#4ac7dcc22c25b1738b1d390432714aa9d9905a10"
+  integrity sha512-v277UnT/fB64xAfSroL5N3Km3tLmvATWqJJw/wRI+g6o+HkeD0slyE7gOhNs1MbjE41R7bQOTxMVoL3aomUJmw==
+
+"@tauri-apps/cli-linux-riscv64-gnu@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-linux-riscv64-gnu/-/cli-linux-riscv64-gnu-2.11.4.tgz#f0cd8de249c8ede745cf69b57d96ad9f1dc55743"
+  integrity sha512-qqgNkQ2u1yZHxjhxsZaxUtRDW8dIqIYm33rx/mzwQv0SfY9x1B+iraj8vWeFiXjjSVVhEMepXSOts1TqPzvXNQ==
+
+"@tauri-apps/cli-linux-x64-gnu@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.11.4.tgz#944636a54656c95bf0eb6ca4af3fe317f4d9a82e"
+  integrity sha512-2VRNWl84FOH0m2giiDkO2h0QXlcMJeX+zJDpI5kDIQAx6s+geF3v48F4DXfJez4GS/FdoDGnPnw1C2iYGbQ7bQ==
+
+"@tauri-apps/cli-linux-x64-musl@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.11.4.tgz#050f98b7800eeb7812d6d583423d37a5b629f571"
+  integrity sha512-o9GyhYor/nc7xarmwDE3ka2szuW3uuZzXjHWh64Q8YX5AtSgxdQkFWzrY4O8KiGtVNvFBI14H3Q49Qj5TOIP/A==
+
+"@tauri-apps/cli-win32-arm64-msvc@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-win32-arm64-msvc/-/cli-win32-arm64-msvc-2.11.4.tgz#24be91dc58f89cd8c9142b52797544991d8cc16c"
+  integrity sha512-ld5Ehb598m0VkYyylRPNeCFsBe/km0jxis6KgMpl3IGY6I/i1RwQXO05I1AsXUXO2WC6AvB/Lw4qTf/asiuEiQ==
+
+"@tauri-apps/cli-win32-ia32-msvc@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-win32-ia32-msvc/-/cli-win32-ia32-msvc-2.11.4.tgz#0fb80df916a4bd0576ca3c30d2308e786833f433"
+  integrity sha512-12Hxi0XX/H5VFxO/bGgHkFWhml9VMgEOu9CidjeCeTNQ1l6fpUlbiGgSP7CLI3PFtW9/FfbeHieZ+kyWK5H7CA==
+
+"@tauri-apps/cli-win32-x64-msvc@2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.11.4.tgz#31917e255214731b31a0fb847279be71fe3e05d8"
+  integrity sha512-+vDiqBIU5dMISg/wNvX3sF+ZHfgJGJ5T0AcO+EHNXV9GGAG+P5fzodlDXD3QdKCRgZxMoCm5PPvj3BqLNjBthw==
+
+"@tauri-apps/cli@^2.11.4":
+  version "2.11.4"
+  resolved "https://registry.yarnpkg.com/@tauri-apps/cli/-/cli-2.11.4.tgz#eafeb0d3a7bd364d27014843fb0a84d9c972e626"
+  integrity sha512-R8xGtMpwyetawSqm9kYOuMmEqkhUbvcUy8n0aNXIxollKBLESUu5f4Fx+64hgASYm1H+jSWq6jCW6zqTnH6hqQ==
+  optionalDependencies:
+    "@tauri-apps/cli-darwin-arm64" "2.11.4"
+    "@tauri-apps/cli-darwin-x64" "2.11.4"
+    "@tauri-apps/cli-linux-arm-gnueabihf" "2.11.4"
+    "@tauri-apps/cli-linux-arm64-gnu" "2.11.4"
+    "@tauri-apps/cli-linux-arm64-musl" "2.11.4"
+    "@tauri-apps/cli-linux-riscv64-gnu" "2.11.4"
+    "@tauri-apps/cli-linux-x64-gnu" "2.11.4"
+    "@tauri-apps/cli-linux-x64-musl" "2.11.4"
+    "@tauri-apps/cli-win32-arm64-msvc" "2.11.4"
+    "@tauri-apps/cli-win32-ia32-msvc" "2.11.4"
+    "@tauri-apps/cli-win32-x64-msvc" "2.11.4"
+
+"@tybys/wasm-util@^0.10.2", "@tybys/wasm-util@^0.10.3":
+  version "0.10.4"
+  resolved "https://registry.yarnpkg.com/@tybys/wasm-util/-/wasm-util-0.10.4.tgz#3e85ecb266f9d1722250e89bc69528615720b154"
+  integrity sha512-W3c4gRigFS0T/Ma4qIYF3GDAc5AQdHb1yL5znJT1Zv1YaD9Kitx656wBjvr19qbiosmZT8lWDM5BEMynUqX65A==
+  dependencies:
+    tslib "^2.4.0"
+
+"@types/debug@^4.0.0":
+  version "4.1.13"
+  resolved "https://registry.yarnpkg.com/@types/debug/-/debug-4.1.13.tgz#22d1cc9d542d3593caea764f974306ab36286ee7"
+  integrity sha512-KSVgmQmzMwPlmtljOomayoR89W4FynCAi3E8PPs7vmDVPe84hT+vGPKkJfThkmXs0x0jAaa9U8uW8bbfyS2fWw==
+  dependencies:
+    "@types/ms" "*"
+
+"@types/estree-jsx@^1.0.0":
+  version "1.0.5"
+  resolved "https://registry.yarnpkg.com/@types/estree-jsx/-/estree-jsx-1.0.5.tgz#858a88ea20f34fe65111f005a689fa1ebf70dc18"
+  integrity sha512-52CcUVNFyfb1A2ALocQw/Dd1BQFNmSdkuC3BkZ6iqhdMfQz7JWOFRuJFloOzjk+6WijU56m9oKXFAXc7o3Towg==
+  dependencies:
+    "@types/estree" "*"
+
+"@types/estree@*", "@types/estree@^1.0.0":
+  version "1.0.9"
+  resolved "https://registry.yarnpkg.com/@types/estree/-/estree-1.0.9.tgz#cf3f0e876d7bee15a93ab925b82bf570a3904a24"
+  integrity sha512-GhdPgy1el4/ImP05X05Uw4cw2/M93BCUmnEvWZNStlCzEKME4Fkk+YpoA5OiHNQmoS7Cafb8Xa3Pya8m1Qrzeg==
+
+"@types/hast@^3.0.0":
+  version "3.0.5"
+  resolved "https://registry.yarnpkg.com/@types/hast/-/hast-3.0.5.tgz#48020de4c0e63492f4ca9db42068c108f68b7f8f"
+  integrity sha512-rp/ezSWaD1m44dPKICGhiskI13nVr7qTloFwDa/IYkhhf5nzwP+zIQcIJh3WIFSBOy/H1PzB40jPjMDksN4F+g==
+  dependencies:
+    "@types/unist" "*"
+
+"@types/katex@^0.16.0":
+  version "0.16.8"
+  resolved "https://registry.yarnpkg.com/@types/katex/-/katex-0.16.8.tgz#80bf3e0814d09a846412a0b0f140946b79c36c3e"
+  integrity sha512-trgaNyfU+Xh2Tc+ABIb44a5AYUpicB3uwirOioeOkNPPbmgRNtcWyDeeFRzjPZENO9Vq8gvVqfhaaXWLlevVwg==
+
+"@types/mdast@^4.0.0":
+  version "4.0.4"
+  resolved "https://registry.yarnpkg.com/@types/mdast/-/mdast-4.0.4.tgz#7ccf72edd2f1aa7dd3437e180c64373585804dd6"
+  integrity sha512-kGaNbPh1k7AFzgpud/gMdvIm5xuECykRR+JnWKQno9TAXVa6WIVCGTPvYGekIDL4uwCZQSYbUxNBSb1aUo79oA==
+  dependencies:
+    "@types/unist" "*"
+
+"@types/ms@*":
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/@types/ms/-/ms-2.1.0.tgz#052aa67a48eccc4309d7f0191b7e41434b90bb78"
+  integrity sha512-GsCCIZDE/p3i96vtEqx+7dBUGXrc7zeSK3wwPHIaRThS+9OhWIXRqzs4d6k1SVU8g91DrNRWxWUGhp5KXQb2VA==
+
+"@types/node@22":
+  version "22.20.2"
+  resolved "https://registry.yarnpkg.com/@types/node/-/node-22.20.2.tgz#daae777b5f5965a587f50cc80d5364c2ddf27e36"
+  integrity sha512-xlvWf4Vs9n1PEVYwP1n4vvG07M6y8WgvJ2t0vbrWTmijsIHp1cS+uJ2kMIRdY3nHZK0nCYKrPeD171+SzF4/zw==
+  dependencies:
+    undici-types "~6.21.0"
+
+"@types/react-dom@^19.3.0":
+  version "19.3.0"
+  resolved "https://registry.yarnpkg.com/@types/react-dom/-/react-dom-19.3.0.tgz#90294a9266efd928a5266f2110d60d5dba5015d2"
+  integrity sha512-ZI7bU42mZXXKHn/qNLEw2IrbiINU7X5+vfgdixBHkCNpYWXjKgfQ/P+uyGb5CjOLB9UcnTeg3rylQtV2hym44Q==
+
+"@types/react@^19.3.0":
+  version "19.3.0"
+  resolved "https://registry.yarnpkg.com/@types/react/-/react-19.3.0.tgz#3b3ce8d8c1627fb5ccd4dadd57ef445c5a153b77"
+  integrity sha512-N0rFCuH9YoxG9/m61l9MfpJKfmLOVU0em7ipIz6TRgSSkvReLB9vL85GB+yr8Bs5leqpvg96JSwF4ZS1s4viQg==
+  dependencies:
+    csstype "^3.2.2"
+
+"@types/unist@*", "@types/unist@^3.0.0":
+  version "3.0.3"
+  resolved "https://registry.yarnpkg.com/@types/unist/-/unist-3.0.3.tgz#acaab0f919ce69cce629c2d4ed2eb4adc1b6c20c"
+  integrity sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q==
+
+"@types/unist@^2.0.0":
+  version "2.0.11"
+  resolved "https://registry.yarnpkg.com/@types/unist/-/unist-2.0.11.tgz#11af57b127e32487774841f7a4e54eab166d03c4"
+  integrity sha512-CmBKiL6NNo/OqgmMn95Fk9Whlp2mtvIv+KNpQKN2F4SjvrEesubTRWGYSg+BnWZOnlCaSTU1sMpsBOzgbYhnsA==
+
+"@typescript/typescript-aix-ppc64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-aix-ppc64/-/typescript-aix-ppc64-7.0.2.tgz#cdc7ce81d60f1e09034960ddfb1fb880d7a776b6"
+  integrity sha512-MTKKkWB7p/0E9xi1d1tHtZ5PiLkGEMIq88pK2CubZjOsLtYTLqhgIgi6zepFa+9GHZ6h05NMCkQxGKiPXMxXtQ==
+
+"@typescript/typescript-darwin-arm64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-darwin-arm64/-/typescript-darwin-arm64-7.0.2.tgz#a55fdfcfa58df58d27db2237cde6a5c1e35a7235"
+  integrity sha512-gowzar9MwS/aRWp6f3a4KUqzRjAZjOsmGNCM6LcTgXum+dBfgsBVMN+AgvOCCbguXyick6LJhpBszxMebJ8syA==
+
+"@typescript/typescript-darwin-x64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-darwin-x64/-/typescript-darwin-x64-7.0.2.tgz#38d1c9172800a91d707bec64d2a370a016634db4"
+  integrity sha512-SZ9xZInqApNlNGc9s0W1VSsktYSOe9cFqNOIqmN1Gs8SmkjKZYFt017G4VwPxASInODuAdbTW7sXiFUf893RgA==
+
+"@typescript/typescript-freebsd-arm64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-freebsd-arm64/-/typescript-freebsd-arm64-7.0.2.tgz#f1ff8810030b35d2b5be0db6a2dc650460ea94fa"
+  integrity sha512-W5NH4y/J0plIIS5b2xvTEkU7JFxyqdMAOgf+Ilhl0vHQXKO5dZoxd+C/jEtq56c4F3wk71RB4BMRQ2XdI+bwYQ==
+
+"@typescript/typescript-freebsd-x64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-freebsd-x64/-/typescript-freebsd-x64-7.0.2.tgz#3d86b03f353c5b1ba95162eb6ce35533bfc294bd"
+  integrity sha512-UMGDx5sTpzNw3WiPebH7l90IWfJggEd+egHt/q6p7/Cm3zqoV7VxkGXt+3DxPIw8CcmvAB0j3sVVfbhX+M4Tpw==
+
+"@typescript/typescript-linux-arm64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-arm64/-/typescript-linux-arm64-7.0.2.tgz#d9334d96d6dac6ff85da9c865588948de939e91f"
+  integrity sha512-Qh4eU4/y3yDjnfjjyPYihMj5/ODIlmt+Bzu17OI+fiSRDW57QmU5SiN63exPRNJPKUzcc1INa1NXdrJ+MqHjUQ==
+
+"@typescript/typescript-linux-arm@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-arm/-/typescript-linux-arm-7.0.2.tgz#ad94b41e1aee2a4dcc6a298c7b67c43345fde32e"
+  integrity sha512-gffT3xPz9sR7j/YJExkyPntrI0P2EP9XbOyWzth2/Gs0RstK+90RBcO0ncXoXy/beYll1SXw846Nf2zdnEz0QQ==
+
+"@typescript/typescript-linux-loong64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-loong64/-/typescript-linux-loong64-7.0.2.tgz#2965aee4fc873360139d893daafe6397a29138ad"
+  integrity sha512-uEHck9i8hoAzXPiYRib1O7miOnz23SxIeVl6F4LXox+qov1K35jHcEW6VHKvZI+pyvl7fZEP4MCU5LYvIq1GuQ==
+
+"@typescript/typescript-linux-mips64el@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-mips64el/-/typescript-linux-mips64el-7.0.2.tgz#1a887a311bed3a833f80bfd4a9ed37c271936cf0"
+  integrity sha512-R4KvAMnE43W5Qeqb0Ly56O3mWMWIAgsMyz36DCaycd5nbg/9kzm0liw3JocfRqyJY0KPmzFjbswozXyW0DnIYA==
+
+"@typescript/typescript-linux-ppc64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-ppc64/-/typescript-linux-ppc64-7.0.2.tgz#8b63c9b2f445b393eb4e43ec21da225dade3577d"
+  integrity sha512-DORx5b3sd/4S7eayxm4FQv+A7CrkUIGRaHiwI8oiHTAI1fAPWhF4J0vAlkC8biAlHSVVwxMQ3tjZ2/DVbnQiiA==
+
+"@typescript/typescript-linux-riscv64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-riscv64/-/typescript-linux-riscv64-7.0.2.tgz#b6e8a35c289b3ea97a92a41d461aaeed0d3b36e1"
+  integrity sha512-wf0jqEDOjrPRnKwYRyyJDRo11KMbvMFrU+q4zqKyChODBzvlkbhNQfKvLxQCcwTpdDaXSHZTVuh0JoCrKCUMHQ==
+
+"@typescript/typescript-linux-s390x@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-s390x/-/typescript-linux-s390x-7.0.2.tgz#2ef96693be4861f6d17965427e5b009cbbed1a3e"
+  integrity sha512-IkwJc3L7yhytWd/ewjyxNDfOmswCm9GWMJT/ue/dU4aZNbwZeYAetq42VyLmsmSjvoX7z74X6ZaYCtzAr0EuGw==
+
+"@typescript/typescript-linux-x64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-linux-x64/-/typescript-linux-x64-7.0.2.tgz#73269cb0baba50aea0ca060445a6b88e583f1ce2"
+  integrity sha512-EYdf2cNg7rgCWJnxCdJ+F3V39O8ihb37eHAu1LK8oAFizgTQbPOK7zHHXbPt8rX24COqODXeI3sIf0fCXG7H/A==
+
+"@typescript/typescript-netbsd-arm64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-netbsd-arm64/-/typescript-netbsd-arm64-7.0.2.tgz#3a3649f97fafa210b4e6e3798c15e06605c8a901"
+  integrity sha512-+polYF4MF04aPpO5FTkHran9yUQDSXqy5GiSDKpsll5jy3l3+g9QLhpf39T+ePtefhXLOGrLl0QIjkQP6VnelA==
+
+"@typescript/typescript-netbsd-x64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-netbsd-x64/-/typescript-netbsd-x64-7.0.2.tgz#47ec59491a40c470d2807dc4d2b825528fd979ab"
+  integrity sha512-8YIT0EHM/3dq10ZOVF/A7pc/YSMtbcecct4rWtexrnSCHOPcpC2KTLXfTCR6vDpnSiY12heNb1GiN/wu+T/FyA==
+
+"@typescript/typescript-openbsd-arm64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-openbsd-arm64/-/typescript-openbsd-arm64-7.0.2.tgz#796be8da0bd989d8a3fb96f2801e38a8365b4baf"
+  integrity sha512-APT8+ClYnuYm1u9+kgGXoMj2VzWzcymwh2gNSQVySHfkRDGOTVkoWLjCmOQSaO+PoqQ57B0flRp9SA+7GnnkzQ==
+
+"@typescript/typescript-openbsd-x64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-openbsd-x64/-/typescript-openbsd-x64-7.0.2.tgz#d37fe2a729eb942c076c454ee7f1815faf7d560f"
+  integrity sha512-yX7s+Q0Dln0Dt9tEzZsAjXXR/+ytBM7AlglaqyeMPxQszJ1JhlJdZ6jLA+IzldHtflX81em7lDao1xXu+aRRkg==
+
+"@typescript/typescript-sunos-x64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-sunos-x64/-/typescript-sunos-x64-7.0.2.tgz#aba8d3464c3565a7044789baba96916bd4ab2c88"
+  integrity sha512-dLJDGaLZ1D4HPQn62u1n8mBDkJREwMsAkCdkwd4Ieqw+x3TUyTsqY0YiBCtE6H6OzzgGk3iuZ3vFWRS+E8/d1g==
+
+"@typescript/typescript-win32-arm64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-win32-arm64/-/typescript-win32-arm64-7.0.2.tgz#b9de50a17196383f62620b5f9d0a2f34ad3b60d7"
+  integrity sha512-Gyl1Vy6OsWesLzmq+EP0Fb7b4Nid5232AvcA2SFcdYreldpNtYFFofPjnt62y9hQy7VTaZp65ICJjuAQRaVcIQ==
+
+"@typescript/typescript-win32-x64@7.0.2":
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/@typescript/typescript-win32-x64/-/typescript-win32-x64-7.0.2.tgz#cf3b7b0d6ce5635daca4c8e01c189cdcde47ec3c"
+  integrity sha512-0BQ3HkAHHlKLSp1qRvf3SUhGpGsDuhB/jgFw75guyqbxJqEaS0Cw/VFO8i2nHglJUzQCRtMMR/IBAKE3ETMC4g==
+
+"@ungap/structured-clone@^1.0.0":
+  version "1.4.0"
+  resolved "https://registry.yarnpkg.com/@ungap/structured-clone/-/structured-clone-1.4.0.tgz#5e2e1374c0a30b5a42e8b083523a225c6945f88a"
+  integrity sha512-1mEZtMKPM09vDmQt5y7YvmN2+DFTP7Tg0EWXdic8/C6VRnpb33e4ghisCIE3WZjsE2N8mf+QV1Zqh7ZFYLWInQ==
+
+"@vercel/oidc@3.1.0":
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/@vercel/oidc/-/oidc-3.1.0.tgz#066caee449b84079f33c7445fc862464fe10ec32"
+  integrity sha512-Fw28YZpRnA3cAHHDlkt7xQHiJ0fcL+NRcIqsocZQUSmbzeIKRpwttJjik5ZGanXP+vlA4SbTg+AbA3bP363l+w==
+
+"@vitejs/plugin-react@^6.1.1":
+  version "6.1.1"
+  resolved "https://registry.yarnpkg.com/@vitejs/plugin-react/-/plugin-react-6.1.1.tgz#47fa1a2f956e88eaf778ebf65ddce8e2b36597e7"
+  integrity sha512-yxLaQV9gkhS8ezJqCM6+ndU7mDY6gqAg75NQ+0IjwEI8IYOmQCgkRwHKVSfWXW076DsqMo0Dk+0FK1U+M5RgFw==
+  dependencies:
+    "@rolldown/pluginutils" "^1.0.1"
+
+ai@5.0.258:
+  version "5.0.258"
+  resolved "https://registry.yarnpkg.com/ai/-/ai-5.0.258.tgz#4469dad5467971322c0578a97733867bbb90777e"
+  integrity sha512-oCy6WZBVNo6N4oJ594qZtCE6L/sXYBu01YoerCru9ixD6ki2S2UVD/MSpFv7pq3BU7EjJnuEhlnn2boTHra/Uw==
+  dependencies:
+    "@ai-sdk/gateway" "2.0.152"
+    "@ai-sdk/provider" "2.0.4"
+    "@ai-sdk/provider-utils" "3.0.37"
+    "@opentelemetry/api" "1.9.0"
+
+bail@^2.0.0:
+  version "2.0.2"
+  resolved "https://registry.yarnpkg.com/bail/-/bail-2.0.2.tgz#d26f5cd8fe5d6f832a31517b9f7c356040ba6d5d"
+  integrity sha512-0xO6mYd7JB2YesxDKplafRpsiOzPt9V02ddPCLbY1xYGPOX24NTyN50qnUxgCPcSoYMhKpAuBTjQoRZCAkUDRw==
+
+ccount@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/ccount/-/ccount-2.0.1.tgz#17a3bf82302e0870d6da43a01311a8bc02a3ecf5"
+  integrity sha512-eyrF0jiFpY+3drT6383f1qhkbGsLSifNAjA61IUjZjmLCWjItY6LB9ft9YhoDgwfmclB2zhu51Lc7+95b8NRAg==
+
+character-entities-html4@^2.0.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/character-entities-html4/-/character-entities-html4-2.1.0.tgz#1f1adb940c971a4b22ba39ddca6b618dc6e56b2b"
+  integrity sha512-1v7fgQRj6hnSwFpq1Eu0ynr/CDEw0rXo2B61qXrLNdHZmPKgb7fqS1a2JwF0rISo9q77jDI8VMEHoApn8qDoZA==
+
+character-entities-legacy@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.yarnpkg.com/character-entities-legacy/-/character-entities-legacy-3.0.0.tgz#76bc83a90738901d7bc223a9e93759fdd560125b"
+  integrity sha512-RpPp0asT/6ufRm//AJVwpViZbGM/MkjQFxJccQRHmISF/22NBtsHqAWmL+/pmkPWoIUJdWyeVleTl1wydHATVQ==
+
+character-entities@^2.0.0:
+  version "2.0.2"
+  resolved "https://registry.yarnpkg.com/character-entities/-/character-entities-2.0.2.tgz#2d09c2e72cd9523076ccb21157dff66ad43fcc22"
+  integrity sha512-shx7oQ0Awen/BRIdkjkvz54PnEEI/EjwXDSIZp86/KKdbafHh1Df/RYGBhn4hbe2+uKC9FnT5UCEdyPz3ai9hQ==
+
+character-reference-invalid@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/character-reference-invalid/-/character-reference-invalid-2.0.1.tgz#85c66b041e43b47210faf401278abf808ac45cb9"
+  integrity sha512-iBZ4F4wRbyORVsu0jPV7gXkOsGYjGHPmAyv+HiHG8gi5PtC9KI2j1+v8/tlibRvjoWX027ypmG/n0HtO5t7unw==
+
+clsx@^2.1.1:
+  version "2.1.1"
+  resolved "https://registry.yarnpkg.com/clsx/-/clsx-2.1.1.tgz#eed397c9fd8bd882bfb18deab7102049a2f32999"
+  integrity sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==
+
+comma-separated-tokens@^2.0.0:
+  version "2.0.3"
+  resolved "https://registry.yarnpkg.com/comma-separated-tokens/-/comma-separated-tokens-2.0.3.tgz#4e89c9458acb61bc8fef19f4529973b2392839ee"
+  integrity sha512-Fu4hJdvzeylCfQPp9SGWidpzrMs7tTrlu6Vb8XGaRGck8QSNZJJp538Wrb60Lax4fPwR64ViY468OIUTbRlGZg==
+
+commander@^8.3.0:
+  version "8.3.0"
+  resolved "https://registry.yarnpkg.com/commander/-/commander-8.3.0.tgz#4837ea1b2da67b9c616a67afbb0fafee567bca66"
+  integrity sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==
+
+cookie-es@^3.0.0:
+  version "3.1.1"
+  resolved "https://registry.yarnpkg.com/cookie-es/-/cookie-es-3.1.1.tgz#c4a8a16cf88cb5a185b23f4a61d6e9a85eb53287"
+  integrity sha512-UaXxwISYJPTr9hwQxMFYZ7kNhSXboMXP+Z3TRX6f1/NyaGPfuNUZOWP1pUEb75B2HjfklIYLVRfWiFZJyC6Npg==
+
+csstype@^3.2.2:
+  version "3.2.3"
+  resolved "https://registry.yarnpkg.com/csstype/-/csstype-3.2.3.tgz#ec48c0f3e993e50648c86da559e2610995cf989a"
+  integrity sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==
+
+debug@^4.0.0:
+  version "4.4.3"
+  resolved "https://registry.yarnpkg.com/debug/-/debug-4.4.3.tgz#c6ae432d9bd9662582fce08709b038c58e9e3d6a"
+  integrity sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==
+  dependencies:
+    ms "^2.1.3"
+
+decode-named-character-reference@^1.0.0:
+  version "1.3.0"
+  resolved "https://registry.yarnpkg.com/decode-named-character-reference/-/decode-named-character-reference-1.3.0.tgz#3e40603760874c2e5867691b599d73a7da25b53f"
+  integrity sha512-GtpQYB283KrPp6nRw50q3U9/VfOutZOe103qlN7BPP6Ad27xYnOIWv4lPzo8HCAL+mMZofJ9KEy30fq6MfaK6Q==
+  dependencies:
+    character-entities "^2.0.0"
+
+dequal@^2.0.0, dequal@^2.0.3:
+  version "2.0.3"
+  resolved "https://registry.yarnpkg.com/dequal/-/dequal-2.0.3.tgz#2644214f1997d39ed0ee0ece72335490a7ac67be"
+  integrity sha512-0je+qPKHEMohvfRTCEo3CrPG6cAzAYgmzKyxRiYSSDkS6eGJdyVJm7WaYA5ECaAD9wLB2T4EEeymA5aFVcYXCA==
+
+detect-libc@^2.0.3:
+  version "2.1.2"
+  resolved "https://registry.yarnpkg.com/detect-libc/-/detect-libc-2.1.2.tgz#689c5dcdc1900ef5583a4cb9f6d7b473742074ad"
+  integrity sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==
+
+devlop@^1.0.0, devlop@^1.1.0:
+  version "1.1.0"
+  resolved "https://registry.yarnpkg.com/devlop/-/devlop-1.1.0.tgz#4db7c2ca4dc6e0e834c30be70c94bbc976dc7018"
+  integrity sha512-RWmIqhcFf1lRYBvNmr7qTNuyCt/7/ns2jbpp1+PalgE/rDQcBT0fioSMUpJ93irlUhC5hrg4cYqe6U+0ImW0rA==
+  dependencies:
+    dequal "^2.0.0"
+
+enhanced-resolve@^5.24.1:
+  version "5.25.1"
+  resolved "https://registry.yarnpkg.com/enhanced-resolve/-/enhanced-resolve-5.25.1.tgz#9d919bfa898b116e58f20e95127c047693158771"
+  integrity sha512-nGXts5znJzmWPu+mIE9izCOzdg63oJca2mDzGWWTth7sr4aCToKcoyFVBQwN75Ij5Pf6p510EwkTqViTRzDV+w==
+  dependencies:
+    graceful-fs "^4.2.4"
+    tapable "^2.3.3"
+
+entities@^6.0.0:
+  version "6.0.1"
+  resolved "https://registry.yarnpkg.com/entities/-/entities-6.0.1.tgz#c28c34a43379ca7f61d074130b2f5f7020a30694"
+  integrity sha512-aN97NXWF6AWBTahfVOIrB/NShkzi5H7F9r1s9mD3cDj4Ko5f2qhhVoYMibXF7GlLveb/D2ioWay8lxI97Ven3g==
+
+escape-string-regexp@^5.0.0:
+  version "5.0.0"
+  resolved "https://registry.yarnpkg.com/escape-string-regexp/-/escape-string-regexp-5.0.0.tgz#4683126b500b61762f2dbebace1806e8be31b1c8"
+  integrity sha512-/veY75JbMK4j1yjvuUxuVsiS/hr/4iHs9FTT6cgTexxdE0Ly/glccBAkloH/DofkjRbZU3bnoj38mOmhkZ0lHw==
+
+estree-util-is-identifier-name@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.yarnpkg.com/estree-util-is-identifier-name/-/estree-util-is-identifier-name-3.0.0.tgz#0b5ef4c4ff13508b34dcd01ecfa945f61fce5dbd"
+  integrity sha512-hFtqIDZTIUZ9BXLb8y4pYGyk6+wekIivNVTcmvk8NoOh+VeRn5y6cEHzbURrWbfp1fIqdVipilzj+lfaadNZmg==
+
+eventsource-parser@^3.0.6:
+  version "3.1.1"
+  resolved "https://registry.yarnpkg.com/eventsource-parser/-/eventsource-parser-3.1.1.tgz#b96cbb7dace4f3774f58a9e3b1ae9a4a524872e2"
+  integrity sha512-EKN1vKAMcZ8MlYMpaNuxN6R9yakzH6uajHcHVTqWJzvu5pWw9DyhbP35HH8MVBQ+dZjAfDxk+A8NiR9KWaXiyQ==
+
+extend@^3.0.0:
+  version "3.0.2"
+  resolved "https://registry.yarnpkg.com/extend/-/extend-3.0.2.tgz#f8b1136b4071fbd8eb140aff858b1019ec2915fa"
+  integrity sha512-fjquC59cD7CyW6urNXK0FBufkZcoiGG80wTuPujX590cB5Ttln20E2UB4S/WARVqhXffZl2LNgS+gQdPIIim/g==
+
+fdir@^6.5.0:
+  version "6.5.0"
+  resolved "https://registry.yarnpkg.com/fdir/-/fdir-6.5.0.tgz#ed2ab967a331ade62f18d077dae192684d50d350"
+  integrity sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==
+
+fsevents@~2.3.3:
+  version "2.3.3"
+  resolved "https://registry.yarnpkg.com/fsevents/-/fsevents-2.3.3.tgz#cac6407785d03675a2a5e1a5305c697b347d90d6"
+  integrity sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==
+
+graceful-fs@^4.2.4:
+  version "4.2.11"
+  resolved "https://registry.yarnpkg.com/graceful-fs/-/graceful-fs-4.2.11.tgz#4183e4e8bf08bb6e05bbb2f7d2e0c8f712ca40e3"
+  integrity sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==
+
+hast-util-from-dom@^5.0.0:
+  version "5.0.1"
+  resolved "https://registry.yarnpkg.com/hast-util-from-dom/-/hast-util-from-dom-5.0.1.tgz#c3c92fbd8d4e1c1625edeb3a773952b9e4ad64a8"
+  integrity sha512-N+LqofjR2zuzTjCPzyDUdSshy4Ma6li7p/c3pA78uTwzFgENbgbUrm2ugwsOdcjI1muO+o6Dgzp9p8WHtn/39Q==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    hastscript "^9.0.0"
+    web-namespaces "^2.0.0"
+
+hast-util-from-html-isomorphic@^2.0.0:
+  version "2.0.0"
+  resolved "https://registry.yarnpkg.com/hast-util-from-html-isomorphic/-/hast-util-from-html-isomorphic-2.0.0.tgz#b31baee386a899a2472326a3c5692f29f86d1d3c"
+  integrity sha512-zJfpXq44yff2hmE0XmwEOzdWin5xwH+QIhMLOScpX91e/NSGPsAzNCvLQDIEPyO2TXi+lBmU6hjLIhV8MwP2kw==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    hast-util-from-dom "^5.0.0"
+    hast-util-from-html "^2.0.0"
+    unist-util-remove-position "^5.0.0"
+
+hast-util-from-html@^2.0.0:
+  version "2.0.3"
+  resolved "https://registry.yarnpkg.com/hast-util-from-html/-/hast-util-from-html-2.0.3.tgz#485c74785358beb80c4ba6346299311ac4c49c82"
+  integrity sha512-CUSRHXyKjzHov8yKsQjGOElXy/3EKpyX56ELnkHH34vDVw1N1XSQ1ZcAvTyAPtGqLTuKP/uxM+aLkSPqF/EtMw==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    devlop "^1.1.0"
+    hast-util-from-parse5 "^8.0.0"
+    parse5 "^7.0.0"
+    vfile "^6.0.0"
+    vfile-message "^4.0.0"
+
+hast-util-from-parse5@^8.0.0:
+  version "8.0.3"
+  resolved "https://registry.yarnpkg.com/hast-util-from-parse5/-/hast-util-from-parse5-8.0.3.tgz#830a35022fff28c3fea3697a98c2f4cc6b835a2e"
+  integrity sha512-3kxEVkEKt0zvcZ3hCRYI8rqrgwtlIOFMWkbclACvjlDw8Li9S2hk/d51OI0nr/gIpdMHNepwgOKqZ/sy0Clpyg==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@types/unist" "^3.0.0"
+    devlop "^1.0.0"
+    hastscript "^9.0.0"
+    property-information "^7.0.0"
+    vfile "^6.0.0"
+    vfile-location "^5.0.0"
+    web-namespaces "^2.0.0"
+
+hast-util-is-element@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.yarnpkg.com/hast-util-is-element/-/hast-util-is-element-3.0.0.tgz#6e31a6532c217e5b533848c7e52c9d9369ca0932"
+  integrity sha512-Val9mnv2IWpLbNPqc/pUem+a7Ipj2aHacCwgNfTiK0vJKl0LF+4Ba4+v1oPHFpf3bLYmreq0/l3Gud9S5OH42g==
+  dependencies:
+    "@types/hast" "^3.0.0"
+
+hast-util-parse-selector@^4.0.0:
+  version "4.0.0"
+  resolved "https://registry.yarnpkg.com/hast-util-parse-selector/-/hast-util-parse-selector-4.0.0.tgz#352879fa86e25616036037dd8931fb5f34cb4a27"
+  integrity sha512-wkQCkSYoOGCRKERFWcxMVMOcYE2K1AaNLU8DXS9arxnLOUEWbOXKXiJUNzEpqZ3JOKpnha3jkFrumEjVliDe7A==
+  dependencies:
+    "@types/hast" "^3.0.0"
+
+hast-util-raw@^9.0.0:
+  version "9.1.0"
+  resolved "https://registry.yarnpkg.com/hast-util-raw/-/hast-util-raw-9.1.0.tgz#79b66b26f6f68fb50dfb4716b2cdca90d92adf2e"
+  integrity sha512-Y8/SBAHkZGoNkpzqqfCldijcuUKh7/su31kEBp67cFY09Wy0mTRgtsLYsiIxMJxlu0f6AA5SUTbDR8K0rxnbUw==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@types/unist" "^3.0.0"
+    "@ungap/structured-clone" "^1.0.0"
+    hast-util-from-parse5 "^8.0.0"
+    hast-util-to-parse5 "^8.0.0"
+    html-void-elements "^3.0.0"
+    mdast-util-to-hast "^13.0.0"
+    parse5 "^7.0.0"
+    unist-util-position "^5.0.0"
+    unist-util-visit "^5.0.0"
+    vfile "^6.0.0"
+    web-namespaces "^2.0.0"
+    zwitch "^2.0.0"
+
+hast-util-sanitize@^5.0.0:
+  version "5.0.2"
+  resolved "https://registry.yarnpkg.com/hast-util-sanitize/-/hast-util-sanitize-5.0.2.tgz#edb260d94e5bba2030eb9375790a8753e5bf391f"
+  integrity sha512-3yTWghByc50aGS7JlGhk61SPenfE/p1oaFeNwkOOyrscaOkMGrcW9+Cy/QAIOBpZxP1yqDIzFMR0+Np0i0+usg==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@ungap/structured-clone" "^1.0.0"
+    unist-util-position "^5.0.0"
+
+hast-util-to-jsx-runtime@^2.3.6:
+  version "2.3.6"
+  resolved "https://registry.yarnpkg.com/hast-util-to-jsx-runtime/-/hast-util-to-jsx-runtime-2.3.6.tgz#ff31897aae59f62232e21594eac7ef6b63333e98"
+  integrity sha512-zl6s8LwNyo1P9uw+XJGvZtdFF1GdAkOg8ujOw+4Pyb76874fLps4ueHXDhXWdk6YHQ6OgUtinliG7RsYvCbbBg==
+  dependencies:
+    "@types/estree" "^1.0.0"
+    "@types/hast" "^3.0.0"
+    "@types/unist" "^3.0.0"
+    comma-separated-tokens "^2.0.0"
+    devlop "^1.0.0"
+    estree-util-is-identifier-name "^3.0.0"
+    hast-util-whitespace "^3.0.0"
+    mdast-util-mdx-expression "^2.0.0"
+    mdast-util-mdx-jsx "^3.0.0"
+    mdast-util-mdxjs-esm "^2.0.0"
+    property-information "^7.0.0"
+    space-separated-tokens "^2.0.0"
+    style-to-js "^1.0.0"
+    unist-util-position "^5.0.0"
+    vfile-message "^4.0.0"
+
+hast-util-to-parse5@^8.0.0:
+  version "8.0.1"
+  resolved "https://registry.yarnpkg.com/hast-util-to-parse5/-/hast-util-to-parse5-8.0.1.tgz#95aa391cc0514b4951418d01c883d1038af42f5d"
+  integrity sha512-MlWT6Pjt4CG9lFCjiz4BH7l9wmrMkfkJYCxFwKQic8+RTZgWPuWxwAfjJElsXkex7DJjfSJsQIt931ilUgmwdA==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    comma-separated-tokens "^2.0.0"
+    devlop "^1.0.0"
+    property-information "^7.0.0"
+    space-separated-tokens "^2.0.0"
+    web-namespaces "^2.0.0"
+    zwitch "^2.0.0"
+
+hast-util-to-text@^4.0.0:
+  version "4.0.2"
+  resolved "https://registry.yarnpkg.com/hast-util-to-text/-/hast-util-to-text-4.0.2.tgz#57b676931e71bf9cb852453678495b3080bfae3e"
+  integrity sha512-KK6y/BN8lbaq654j7JgBydev7wuNMcID54lkRav1P0CaE1e47P72AWWPiGKXTJU271ooYzcvTAn/Zt0REnvc7A==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@types/unist" "^3.0.0"
+    hast-util-is-element "^3.0.0"
+    unist-util-find-after "^5.0.0"
+
+hast-util-whitespace@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.yarnpkg.com/hast-util-whitespace/-/hast-util-whitespace-3.0.0.tgz#7778ed9d3c92dd9e8c5c8f648a49c21fc51cb621"
+  integrity sha512-88JUN06ipLwsnv+dVn+OIYOvAuvBMy/Qoi6O7mQHxdPXpjy+Cd6xRkWwux7DKO+4sYILtLBRIKgsdpS2gQc7qw==
+  dependencies:
+    "@types/hast" "^3.0.0"
+
+hastscript@^9.0.0:
+  version "9.0.1"
+  resolved "https://registry.yarnpkg.com/hastscript/-/hastscript-9.0.1.tgz#dbc84bef6051d40084342c229c451cd9dc567dff"
+  integrity sha512-g7df9rMFX/SPi34tyGCyUBREQoKkapwdY/T04Qn9TDWfHhAYt4/I0gMVirzK5wEzeUqIjEB+LXC/ypb7Aqno5w==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    comma-separated-tokens "^2.0.0"
+    hast-util-parse-selector "^4.0.0"
+    property-information "^7.0.0"
+    space-separated-tokens "^2.0.0"
+
+html-parse-stringify@^4.0.1:
+  version "4.0.1"
+  resolved "https://registry.yarnpkg.com/html-parse-stringify/-/html-parse-stringify-4.0.1.tgz#7264be901df0fef8829eb04632191be0aa94f55e"
+  integrity sha512-0zHsZJrK7S3K2aucXWL6ycoYJ/iNtIcFHC/nYQgFklPtrv5LpJctIiSCroWZWeuoXvuyFdzp6KzjJQ+OT5MfFw==
+
+html-url-attributes@^3.0.1:
+  version "3.0.1"
+  resolved "https://registry.yarnpkg.com/html-url-attributes/-/html-url-attributes-3.0.1.tgz#83b052cd5e437071b756cd74ae70f708870c2d87"
+  integrity sha512-ol6UPyBWqsrO6EJySPz2O7ZSr856WDrEzM5zMqp+FJJLGMW35cLYmmZnl0vztAZxRUoNZJFTCohfjuIJ8I4QBQ==
+
+html-void-elements@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.yarnpkg.com/html-void-elements/-/html-void-elements-3.0.0.tgz#fc9dbd84af9e747249034d4d62602def6517f1d7"
+  integrity sha512-bEqo66MRXsUGxWHV5IP0PUiAWwoEjba4VCzg0LjFJBpchPaTfyfCKTG6bc5F8ucKec3q5y6qOdGyYTSBEvhCrg==
+
+i18next@^26.4.2:
+  version "26.4.2"
+  resolved "https://registry.yarnpkg.com/i18next/-/i18next-26.4.2.tgz#824233630f147359296bcaddf1404945d2b211bf"
+  integrity sha512-RX+R0VLg13IbvRuJSxnqykUFS9vQZTl8wYpWPCIUDWVrSGjsQywB5Y+pjzrkboxGAuYfJZVH1InFTdgBdxq6ug==
+
+inline-style-parser@0.2.7:
+  version "0.2.7"
+  resolved "https://registry.yarnpkg.com/inline-style-parser/-/inline-style-parser-0.2.7.tgz#b1fc68bfc0313b8685745e4464e37f9376b9c909"
+  integrity sha512-Nb2ctOyNR8DqQoR0OwRG95uNWIC0C1lCgf5Naz5H6Ji72KZ8OcFZLz2P5sNgwlyoJ8Yif11oMuYs5pBQa86csA==
+
+is-alphabetical@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/is-alphabetical/-/is-alphabetical-2.0.1.tgz#01072053ea7c1036df3c7d19a6daaec7f19e789b"
+  integrity sha512-FWyyY60MeTNyeSRpkM2Iry0G9hpr7/9kD40mD/cGQEuilcZYS4okz8SN2Q6rLCJ8gbCt6fN+rC+6tMGS99LaxQ==
+
+is-alphanumerical@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/is-alphanumerical/-/is-alphanumerical-2.0.1.tgz#7c03fbe96e3e931113e57f964b0a368cc2dfd875"
+  integrity sha512-hmbYhX/9MUMF5uh7tOXyK/n0ZvWpad5caBA17GsC6vyuCqaWliRG5K1qS9inmUhEMaOBIW7/whAnSwveW/LtZw==
+  dependencies:
+    is-alphabetical "^2.0.0"
+    is-decimal "^2.0.0"
+
+is-decimal@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/is-decimal/-/is-decimal-2.0.1.tgz#9469d2dc190d0214fd87d78b78caecc0cc14eef7"
+  integrity sha512-AAB9hiomQs5DXWcRB1rqsxGUstbRroFOPPVAomNk/3XHR5JyEZChOyTWe2oayKnsSsr/kcGqF+z6yuH6HHpN0A==
+
+is-hexadecimal@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/is-hexadecimal/-/is-hexadecimal-2.0.1.tgz#86b5bf668fca307498d319dfc03289d781a90027"
+  integrity sha512-DgZQp241c8oO6cA1SbTEWiXeoxV42vlcJxgH+B3hi1AiqqKruZR3ZGF8In3fj4+/y/7rHvlOZLZtgJ/4ttYGZg==
+
+is-plain-obj@^4.0.0:
+  version "4.1.0"
+  resolved "https://registry.yarnpkg.com/is-plain-obj/-/is-plain-obj-4.1.0.tgz#d65025edec3657ce032fd7db63c97883eaed71f0"
+  integrity sha512-+Pgi+vMuUNkJyExiMBt5IlFoMyKnr5zhJ4Uspz58WOhBF5QoIZkFyNHIbBAtHwzVAgk5RtndVNsDRN61/mmDqg==
+
+isbot@^5.1.22:
+  version "5.2.2"
+  resolved "https://registry.yarnpkg.com/isbot/-/isbot-5.2.2.tgz#4c6c4c9b1391a91e490077f9323418693140c416"
+  integrity sha512-iQcBXcd+Rv/pkubRyGh2utW2j1oPG5hZY6TUhVPpqK4G+o3IbxpJNx04hgksjc/N7GK5pEorUxDeg31cFgEk/w==
+
+jiti@^2.7.0:
+  version "2.7.0"
+  resolved "https://registry.yarnpkg.com/jiti/-/jiti-2.7.0.tgz#974228f2f4ca2bc21885a1797b45fea68e950c64"
+  integrity sha512-AC/7JofJvZGrrneWNaEnJeOLUx+JlGt7tNa0wZiRPT4MY1wmfKjt2+6O2p2uz2+skll8OZZmJMNqeke7kKbNgQ==
+
+json-schema@^0.4.0:
+  version "0.4.0"
+  resolved "https://registry.yarnpkg.com/json-schema/-/json-schema-0.4.0.tgz#f7de4cf6efab838ebaeb3236474cbba5a1930ab5"
+  integrity sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA==
+
+katex@^0.16.0, katex@^0.16.27:
+  version "0.16.47"
+  resolved "https://registry.yarnpkg.com/katex/-/katex-0.16.47.tgz#0a13a42c2deb4f74e61f162d440b9165a548030f"
+  integrity sha512-Eeo8Ys1doU1z+x8AZsPpQu+p/QcZBI5PeOo7QGQdy2x2m0MU/hYagBbGOmXwr5KVbEfVuWv9LpnQWeehogurjg==
+  dependencies:
+    commander "^8.3.0"
+
+lightningcss-android-arm64@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-android-arm64/-/lightningcss-android-arm64-1.32.0.tgz#f033885116dfefd9c6f54787523e3514b61e1968"
+  integrity sha512-YK7/ClTt4kAK0vo6w3X+Pnm0D2cf2vPHbhOXdoNti1Ga0al1P4TBZhwjATvjNwLEBCnKvjJc2jQgHXH0NEwlAg==
+
+lightningcss-android-arm64@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-android-arm64/-/lightningcss-android-arm64-1.33.0.tgz#9a6841f88ae50fc83502903892b41af41bc2b907"
+  integrity sha512-gEpRTalKdosp4Bb8qWtc2iOgE5SeIHlpS1up9bFq2wAyYhl1UdTObYiHe98zEM9SQvSoqQZ1IQD0JNpg3Ml5pg==
+
+lightningcss-darwin-arm64@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-darwin-arm64/-/lightningcss-darwin-arm64-1.32.0.tgz#50b71871b01c8199584b649e292547faea7af9b5"
+  integrity sha512-RzeG9Ju5bag2Bv1/lwlVJvBE3q6TtXskdZLLCyfg5pt+HLz9BqlICO7LZM7VHNTTn/5PRhHFBSjk5lc4cmscPQ==
+
+lightningcss-darwin-arm64@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-darwin-arm64/-/lightningcss-darwin-arm64-1.33.0.tgz#c0f2c31c0bfd19fa4dd3f18e957a1f1a152097d6"
+  integrity sha512-Sciaz8eenNTKn9b3t7+xr0ipTp9YxKQY4npwQ3mrRuL0BAVHBLyZxofhaKBAVtzmtRZ/zTyo0/to4B1uWG/Djg==
+
+lightningcss-darwin-x64@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-darwin-x64/-/lightningcss-darwin-x64-1.32.0.tgz#35f3e97332d130b9ca181e11b568ded6aebc6d5e"
+  integrity sha512-U+QsBp2m/s2wqpUYT/6wnlagdZbtZdndSmut/NJqlCcMLTWp5muCrID+K5UJ6jqD2BFshejCYXniPDbNh73V8w==
+
+lightningcss-darwin-x64@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-darwin-x64/-/lightningcss-darwin-x64-1.33.0.tgz#cb0705965acb538c6683949ce6925fb3cdf7c361"
+  integrity sha512-Z5UPAxzrjlWNNyGy6i65cJzzvgJ5D3T6wMvs+gWpY9d7qRhANrxqAp6LhxIgZhWEw18RfJTGcRxjuLIBr+m8XQ==
+
+lightningcss-freebsd-x64@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-freebsd-x64/-/lightningcss-freebsd-x64-1.32.0.tgz#9777a76472b64ed6ff94342ad64c7bafd794a575"
+  integrity sha512-JCTigedEksZk3tHTTthnMdVfGf61Fky8Ji2E4YjUTEQX14xiy/lTzXnu1vwiZe3bYe0q+SpsSH/CTeDXK6WHig==
+
+lightningcss-freebsd-x64@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-freebsd-x64/-/lightningcss-freebsd-x64-1.33.0.tgz#763538828b26bab2680dadafcc84ee78b0eb502b"
+  integrity sha512-QQM/Ti/hQajJwCY+RiWuCZ9sdtI/XQk7nDK5vC8kkdwixezOlDgvDx7+RT+QjK6FcFT4MpsuoBnHIo/O3StRRg==
+
+lightningcss-linux-arm-gnueabihf@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-arm-gnueabihf/-/lightningcss-linux-arm-gnueabihf-1.32.0.tgz#13ae652e1ab73b9135d7b7da172f666c410ad53d"
+  integrity sha512-x6rnnpRa2GL0zQOkt6rts3YDPzduLpWvwAF6EMhXFVZXD4tPrBkEFqzGowzCsIWsPjqSK+tyNEODUBXeeVHSkw==
+
+lightningcss-linux-arm-gnueabihf@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-arm-gnueabihf/-/lightningcss-linux-arm-gnueabihf-1.33.0.tgz#6862e3176a331aedbdec1ed352b4d7d0dd0784de"
+  integrity sha512-N7FVBe6iS24MlM6R/4RBTxGhQheZGs7tiQ9U32UtF75NzP5Q7xWPRqLBCKxlRQRk3rY1jCIPLzx7WzOhuUIRLQ==
+
+lightningcss-linux-arm64-gnu@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-arm64-gnu/-/lightningcss-linux-arm64-gnu-1.32.0.tgz#417858795a94592f680123a1b1f9da8a0e1ef335"
+  integrity sha512-0nnMyoyOLRJXfbMOilaSRcLH3Jw5z9HDNGfT/gwCPgaDjnx0i8w7vBzFLFR1f6CMLKF8gVbebmkUN3fa/kQJpQ==
+
+lightningcss-linux-arm64-gnu@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-arm64-gnu/-/lightningcss-linux-arm64-gnu-1.33.0.tgz#c6a3a2ed15141daf6bdc2628930f8e39bdf473aa"
+  integrity sha512-j2v/itmy4HlNxlc6voKXYgBqNi0Ng2LShg4z7GufpEgs05P+2suBVyi9I6YHq5uoVFx9ETin3eCEhLVyXGQnKg==
+
+lightningcss-linux-arm64-musl@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-arm64-musl/-/lightningcss-linux-arm64-musl-1.32.0.tgz#6be36692e810b718040802fd809623cffe732133"
+  integrity sha512-UpQkoenr4UJEzgVIYpI80lDFvRmPVg6oqboNHfoH4CQIfNA+HOrZ7Mo7KZP02dC6LjghPQJeBsvXhJod/wnIBg==
+
+lightningcss-linux-arm64-musl@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-arm64-musl/-/lightningcss-linux-arm64-musl-1.33.0.tgz#7fa1334971fc82845f9827df6ef8a0b20914bac6"
+  integrity sha512-yiO5ROMuYQgXbC60yjZU5CYSFZGKXL0HFATXt9mHJn1+zW55oCtMI9NfcVhYLMFDL7gV7oBPon/EmMMGg2OvtQ==
+
+lightningcss-linux-x64-gnu@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-x64-gnu/-/lightningcss-linux-x64-gnu-1.32.0.tgz#0b7803af4eb21cfd38dd39fe2abbb53c7dd091f6"
+  integrity sha512-V7Qr52IhZmdKPVr+Vtw8o+WLsQJYCTd8loIfpDaMRWGUZfBOYEJeyJIkqGIDMZPwPx24pUMfwSxxI8phr/MbOA==
+
+lightningcss-linux-x64-gnu@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-x64-gnu/-/lightningcss-linux-x64-gnu-1.33.0.tgz#8b927862ea8c2bbc6831a46509244b50d9936e55"
+  integrity sha512-ar+Ju7LmcN0Jo4FpL4hpFybwNG9/3A/Br5KW2n2jyODg3MEZXaDYADdemoNS+BDNfMgKvylJLj4S5tyRActuAg==
+
+lightningcss-linux-x64-musl@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-x64-musl/-/lightningcss-linux-x64-musl-1.32.0.tgz#88dc8ba865ddddb1ac5ef04b0f161804418c163b"
+  integrity sha512-bYcLp+Vb0awsiXg/80uCRezCYHNg1/l3mt0gzHnWV9XP1W5sKa5/TCdGWaR/zBM2PeF/HbsQv/j2URNOiVuxWg==
+
+lightningcss-linux-x64-musl@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-linux-x64-musl/-/lightningcss-linux-x64-musl-1.33.0.tgz#0c525bb077dfd94404c059cfe42dad797e96aeaf"
+  integrity sha512-RYiYbkokw0trfKqqzfF55lginwEPrD3OJDfTuJzFs1MK6iFnDenaz1fqLLtX4ITG3OktJQXOeTaw1awrBAlZPw==
+
+lightningcss-win32-arm64-msvc@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-win32-arm64-msvc/-/lightningcss-win32-arm64-msvc-1.32.0.tgz#4f30ba3fa5e925f5b79f945e8cc0d176c3b1ab38"
+  integrity sha512-8SbC8BR40pS6baCM8sbtYDSwEVQd4JlFTOlaD3gWGHfThTcABnNDBda6eTZeqbofalIJhFx0qKzgHJmcPTnGdw==
+
+lightningcss-win32-arm64-msvc@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-win32-arm64-msvc/-/lightningcss-win32-arm64-msvc-1.33.0.tgz#850ee1103dac989cfab50e3ac22d1a69e394e63d"
+  integrity sha512-1K+MPfLSFVpphzpdbfkhlWk6wBrTObBzS2T6db10PNOZgR9GoVsAWzwNyuhUYYbTp23j+4RrncfujZ4uAzXvwA==
+
+lightningcss-win32-x64-msvc@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-win32-x64-msvc/-/lightningcss-win32-x64-msvc-1.32.0.tgz#141aa5605645064928902bb4af045fa7d9f4220a"
+  integrity sha512-Amq9B/SoZYdDi1kFrojnoqPLxYhQ4Wo5XiL8EVJrVsB8ARoC1PWW6VGtT0WKCemjy8aC+louJnjS7U18x3b06Q==
+
+lightningcss-win32-x64-msvc@1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss-win32-x64-msvc/-/lightningcss-win32-x64-msvc-1.33.0.tgz#e343ae152eed3609dc6e11949d1a3bf39a1c946f"
+  integrity sha512-OlEICDx/Xl0FqSp4bry8zFnCvGpig3Gl4gCquvYwHuqJKEC1+n9NgDniFvqHGmMv1ZkqDJrDqKKSykTDX+ehuA==
+
+lightningcss@1.32.0:
+  version "1.32.0"
+  resolved "https://registry.yarnpkg.com/lightningcss/-/lightningcss-1.32.0.tgz#b85aae96486dcb1bf49a7c8571221273f4f1e4a9"
+  integrity sha512-NXYBzinNrblfraPGyrbPoD19C1h9lfI/1mzgWYvXUTe414Gz/X1FD2XBZSZM7rRTrMA8JL3OtAaGifrIKhQ5yQ==
+  dependencies:
+    detect-libc "^2.0.3"
+  optionalDependencies:
+    lightningcss-android-arm64 "1.32.0"
+    lightningcss-darwin-arm64 "1.32.0"
+    lightningcss-darwin-x64 "1.32.0"
+    lightningcss-freebsd-x64 "1.32.0"
+    lightningcss-linux-arm-gnueabihf "1.32.0"
+    lightningcss-linux-arm64-gnu "1.32.0"
+    lightningcss-linux-arm64-musl "1.32.0"
+    lightningcss-linux-x64-gnu "1.32.0"
+    lightningcss-linux-x64-musl "1.32.0"
+    lightningcss-win32-arm64-msvc "1.32.0"
+    lightningcss-win32-x64-msvc "1.32.0"
+
+lightningcss@^1.33.0:
+  version "1.33.0"
+  resolved "https://registry.yarnpkg.com/lightningcss/-/lightningcss-1.33.0.tgz#c08867d71a79385c6e190214fd72fef3e5f95f0b"
+  integrity sha512-WkUDrojuJs0xkgGf2udWxa3yGBRxPtxUkB79i6aCZLRgc7PM8fZe9TosfPDcvEpQZbuFASnHYmRLBLUbmLOIIA==
+  dependencies:
+    detect-libc "^2.0.3"
+  optionalDependencies:
+    lightningcss-android-arm64 "1.33.0"
+    lightningcss-darwin-arm64 "1.33.0"
+    lightningcss-darwin-x64 "1.33.0"
+    lightningcss-freebsd-x64 "1.33.0"
+    lightningcss-linux-arm-gnueabihf "1.33.0"
+    lightningcss-linux-arm64-gnu "1.33.0"
+    lightningcss-linux-arm64-musl "1.33.0"
+    lightningcss-linux-x64-gnu "1.33.0"
+    lightningcss-linux-x64-musl "1.33.0"
+    lightningcss-win32-arm64-msvc "1.33.0"
+    lightningcss-win32-x64-msvc "1.33.0"
+
+longest-streak@^3.0.0:
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/longest-streak/-/longest-streak-3.1.0.tgz#62fa67cd958742a1574af9f39866364102d90cd4"
+  integrity sha512-9Ri+o0JYgehTaVBBDoMqIl8GXtbWg711O3srftcHhZ0dqnETqLaoIK0x17fUw9rFSlK/0NlsKe0Ahhyl5pXE2g==
+
+lucide-react@^1.46.0:
+  version "1.46.0"
+  resolved "https://registry.yarnpkg.com/lucide-react/-/lucide-react-1.46.0.tgz#f485e29a7b448250ca96977450e529c473b5b496"
+  integrity sha512-Bv+FZXgZPrxc/NCl1e7JJVQFLdiCxYgxNVhqoV7X0p6I8ADJo8DxBnK1auH0fZz4AmqOJ3jgneL4f1i8LJQRAA==
+
+magic-string@^0.30.21:
+  version "0.30.21"
+  resolved "https://registry.yarnpkg.com/magic-string/-/magic-string-0.30.21.tgz#56763ec09a0fa8091df27879fd94d19078c00d91"
+  integrity sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ==
+  dependencies:
+    "@jridgewell/sourcemap-codec" "^1.5.5"
+
+markdown-table@^3.0.0:
+  version "3.0.4"
+  resolved "https://registry.yarnpkg.com/markdown-table/-/markdown-table-3.0.4.tgz#fe44d6d410ff9d6f2ea1797a3f60aa4d2b631c2a"
+  integrity sha512-wiYz4+JrLyb/DqW2hkFJxP7Vd7JuTDm77fvbM8VfEQdmSMqcImWeeRbHwZjBjIFki/VaMK2BhFi7oUUZeM5bqw==
+
+marked@^17.0.1:
+  version "17.0.6"
+  resolved "https://registry.yarnpkg.com/marked/-/marked-17.0.6.tgz#2a97586a272d3be5880f198e020b74ad27cf86ba"
+  integrity sha512-gB0gkNafnonOw0obSTEGZTT86IuhILt2Wfx0mWH/1Au83kybTayroZ/V6nS25mN7u8ASy+5fMhgB3XPNrOZdmA==
+
+mdast-util-find-and-replace@^3.0.0:
+  version "3.0.2"
+  resolved "https://registry.yarnpkg.com/mdast-util-find-and-replace/-/mdast-util-find-and-replace-3.0.2.tgz#70a3174c894e14df722abf43bc250cbae44b11df"
+  integrity sha512-Tmd1Vg/m3Xz43afeNxDIhWRtFZgM2VLyaf4vSTYwudTyeuTneoL3qtWMA5jeLyz/O1vDJmmV4QuScFCA2tBPwg==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    escape-string-regexp "^5.0.0"
+    unist-util-is "^6.0.0"
+    unist-util-visit-parents "^6.0.0"
+
+mdast-util-from-markdown@^2.0.0:
+  version "2.0.3"
+  resolved "https://registry.yarnpkg.com/mdast-util-from-markdown/-/mdast-util-from-markdown-2.0.3.tgz#c95822b91aab75f18a4cbe8b2f51b873ed2cf0c7"
+  integrity sha512-W4mAWTvSlKvf8L6J+VN9yLSqQ9AOAAvHuoDAmPkz4dHf553m5gVj2ejadHJhoJmcmxEnOv6Pa8XJhpxE93kb8Q==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    "@types/unist" "^3.0.0"
+    decode-named-character-reference "^1.0.0"
+    devlop "^1.0.0"
+    mdast-util-to-string "^4.0.0"
+    micromark "^4.0.0"
+    micromark-util-decode-numeric-character-reference "^2.0.0"
+    micromark-util-decode-string "^2.0.0"
+    micromark-util-normalize-identifier "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+    unist-util-stringify-position "^4.0.0"
+
+mdast-util-gfm-autolink-literal@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/mdast-util-gfm-autolink-literal/-/mdast-util-gfm-autolink-literal-2.0.1.tgz#abd557630337bd30a6d5a4bd8252e1c2dc0875d5"
+  integrity sha512-5HVP2MKaP6L+G6YaxPNjuL0BPrq9orG3TsrZ9YXbA3vDw/ACI4MEsnoDpn6ZNm7GnZgtAcONJyPhOP8tNJQavQ==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    ccount "^2.0.0"
+    devlop "^1.0.0"
+    mdast-util-find-and-replace "^3.0.0"
+    micromark-util-character "^2.0.0"
+
+mdast-util-gfm-footnote@^2.0.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-gfm-footnote/-/mdast-util-gfm-footnote-2.1.0.tgz#7778e9d9ca3df7238cc2bd3fa2b1bf6a65b19403"
+  integrity sha512-sqpDWlsHn7Ac9GNZQMeUzPQSMzR6Wv0WKRNvQRg0KqHh02fpTz69Qc1QSseNX29bhz1ROIyNyxExfawVKTm1GQ==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    devlop "^1.1.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+    micromark-util-normalize-identifier "^2.0.0"
+
+mdast-util-gfm-strikethrough@^2.0.0:
+  version "2.0.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-gfm-strikethrough/-/mdast-util-gfm-strikethrough-2.0.0.tgz#d44ef9e8ed283ac8c1165ab0d0dfd058c2764c16"
+  integrity sha512-mKKb915TF+OC5ptj5bJ7WFRPdYtuHv0yTRxK2tJvi+BDqbkiG7h7u/9SI89nRAYcmap2xHQL9D+QG/6wSrTtXg==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+
+mdast-util-gfm-table@^2.0.0:
+  version "2.0.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-gfm-table/-/mdast-util-gfm-table-2.0.0.tgz#7a435fb6223a72b0862b33afbd712b6dae878d38"
+  integrity sha512-78UEvebzz/rJIxLvE7ZtDd/vIQ0RHv+3Mh5DR96p7cS7HsBhYIICDBCu8csTNWNO6tBWfqXPWekRuj2FNOGOZg==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    devlop "^1.0.0"
+    markdown-table "^3.0.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+
+mdast-util-gfm-task-list-item@^2.0.0:
+  version "2.0.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-gfm-task-list-item/-/mdast-util-gfm-task-list-item-2.0.0.tgz#e68095d2f8a4303ef24094ab642e1047b991a936"
+  integrity sha512-IrtvNvjxC1o06taBAVJznEnkiHxLFTzgonUdy8hzFVeDun0uTjxxrRGVaNFqkU1wJR3RBPEfsxmU6jDWPofrTQ==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    devlop "^1.0.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+
+mdast-util-gfm@^3.0.0:
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-gfm/-/mdast-util-gfm-3.1.0.tgz#2cdf63b92c2a331406b0fb0db4c077c1b0331751"
+  integrity sha512-0ulfdQOM3ysHhCJ1p06l0b0VKlhU0wuQs3thxZQagjcjPrlFRqY215uZGHHJan9GEAXd9MbfPjFJz+qMkVR6zQ==
+  dependencies:
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-gfm-autolink-literal "^2.0.0"
+    mdast-util-gfm-footnote "^2.0.0"
+    mdast-util-gfm-strikethrough "^2.0.0"
+    mdast-util-gfm-table "^2.0.0"
+    mdast-util-gfm-task-list-item "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+
+mdast-util-math@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-math/-/mdast-util-math-3.0.0.tgz#8d79dd3baf8ab8ac781f62b8853768190b9a00b0"
+  integrity sha512-Tl9GBNeG/AhJnQM221bJR2HPvLOSnLE/T9cJI9tlc6zwQk2nPk/4f0cHkOdEixQPC/j8UtKDdITswvLAy1OZ1w==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@types/mdast" "^4.0.0"
+    devlop "^1.0.0"
+    longest-streak "^3.0.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.1.0"
+    unist-util-remove-position "^5.0.0"
+
+mdast-util-mdx-expression@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/mdast-util-mdx-expression/-/mdast-util-mdx-expression-2.0.1.tgz#43f0abac9adc756e2086f63822a38c8d3c3a5096"
+  integrity sha512-J6f+9hUp+ldTZqKRSg7Vw5V6MqjATc+3E4gf3CFNcuZNWD8XdyI6zQ8GqH7f8169MM6P7hMBRDVGnn7oHB9kXQ==
+  dependencies:
+    "@types/estree-jsx" "^1.0.0"
+    "@types/hast" "^3.0.0"
+    "@types/mdast" "^4.0.0"
+    devlop "^1.0.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+
+mdast-util-mdx-jsx@^3.0.0:
+  version "3.2.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-mdx-jsx/-/mdast-util-mdx-jsx-3.2.0.tgz#fd04c67a2a7499efb905a8a5c578dddc9fdada0d"
+  integrity sha512-lj/z8v0r6ZtsN/cGNNtemmmfoLAFZnjMbNyLzBafjzikOM+glrjNHPlf6lQDOTccj9n5b0PPihEBbhneMyGs1Q==
+  dependencies:
+    "@types/estree-jsx" "^1.0.0"
+    "@types/hast" "^3.0.0"
+    "@types/mdast" "^4.0.0"
+    "@types/unist" "^3.0.0"
+    ccount "^2.0.0"
+    devlop "^1.1.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+    parse-entities "^4.0.0"
+    stringify-entities "^4.0.0"
+    unist-util-stringify-position "^4.0.0"
+    vfile-message "^4.0.0"
+
+mdast-util-mdxjs-esm@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/mdast-util-mdxjs-esm/-/mdast-util-mdxjs-esm-2.0.1.tgz#019cfbe757ad62dd557db35a695e7314bcc9fa97"
+  integrity sha512-EcmOpxsZ96CvlP03NghtH1EsLtr0n9Tm4lPUJUBccV9RwUOneqSycg19n5HGzCf+10LozMRSObtVr3ee1WoHtg==
+  dependencies:
+    "@types/estree-jsx" "^1.0.0"
+    "@types/hast" "^3.0.0"
+    "@types/mdast" "^4.0.0"
+    devlop "^1.0.0"
+    mdast-util-from-markdown "^2.0.0"
+    mdast-util-to-markdown "^2.0.0"
+
+mdast-util-phrasing@^4.0.0:
+  version "4.1.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-phrasing/-/mdast-util-phrasing-4.1.0.tgz#7cc0a8dec30eaf04b7b1a9661a92adb3382aa6e3"
+  integrity sha512-TqICwyvJJpBwvGAMZjj4J2n0X8QWp21b9l0o7eXyVJ25YNWYbJDVIyD1bZXE6WtV6RmKJVYmQAKWa0zWOABz2w==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    unist-util-is "^6.0.0"
+
+mdast-util-to-hast@^13.0.0:
+  version "13.2.1"
+  resolved "https://registry.yarnpkg.com/mdast-util-to-hast/-/mdast-util-to-hast-13.2.1.tgz#d7ff84ca499a57e2c060ae67548ad950e689a053"
+  integrity sha512-cctsq2wp5vTsLIcaymblUriiTcZd0CwWtCbLvrOzYCDZoWyMNV8sZ7krj09FSnsiJi3WVsHLM4k6Dq/yaPyCXA==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@types/mdast" "^4.0.0"
+    "@ungap/structured-clone" "^1.0.0"
+    devlop "^1.0.0"
+    micromark-util-sanitize-uri "^2.0.0"
+    trim-lines "^3.0.0"
+    unist-util-position "^5.0.0"
+    unist-util-visit "^5.0.0"
+    vfile "^6.0.0"
+
+mdast-util-to-markdown@^2.0.0:
+  version "2.1.2"
+  resolved "https://registry.yarnpkg.com/mdast-util-to-markdown/-/mdast-util-to-markdown-2.1.2.tgz#f910ffe60897f04bb4b7e7ee434486f76288361b"
+  integrity sha512-xj68wMTvGXVOKonmog6LwyJKrYXZPvlwabaryTjLh9LuvovB/KAH+kvi8Gjj+7rJjsFi23nkUxRQv1KqSroMqA==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    "@types/unist" "^3.0.0"
+    longest-streak "^3.0.0"
+    mdast-util-phrasing "^4.0.0"
+    mdast-util-to-string "^4.0.0"
+    micromark-util-classify-character "^2.0.0"
+    micromark-util-decode-string "^2.0.0"
+    unist-util-visit "^5.0.0"
+    zwitch "^2.0.0"
+
+mdast-util-to-markdown@^2.1.0:
+  version "2.1.3"
+  resolved "https://registry.yarnpkg.com/mdast-util-to-markdown/-/mdast-util-to-markdown-2.1.3.tgz#098abd393c0fe047c44bd3396849255c55d0b55f"
+  integrity sha512-wgyJtgUkUcdU7zci7uuwc/tzoAhm0TswWhaXuvnWolqug+8jY7bgNkBBL0dYBnHgc/tHL7lCzRdPsdUxUKKBqw==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    "@types/unist" "^3.0.0"
+    longest-streak "^3.0.0"
+    mdast-util-phrasing "^4.0.0"
+    mdast-util-to-string "^4.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-classify-character "^2.0.0"
+    micromark-util-decode-string "^2.0.0"
+    micromark-util-html-tag-name "^2.0.0"
+    unist-util-visit "^5.0.0"
+    zwitch "^2.0.0"
+
+mdast-util-to-string@^4.0.0:
+  version "4.0.0"
+  resolved "https://registry.yarnpkg.com/mdast-util-to-string/-/mdast-util-to-string-4.0.0.tgz#7a5121475556a04e7eddeb67b264aae79d312814"
+  integrity sha512-0H44vDimn51F0YwvxSJSm0eCDOJTRlmN0R1yBh4HLj9wiV1Dn0QoXGbvFAWj2hSItVTlCmBF1hqKlIyUBVFLPg==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+
+micromark-core-commonmark@^2.0.0:
+  version "2.0.3"
+  resolved "https://registry.yarnpkg.com/micromark-core-commonmark/-/micromark-core-commonmark-2.0.3.tgz#c691630e485021a68cf28dbc2b2ca27ebf678cd4"
+  integrity sha512-RDBrHEMSxVFLg6xvnXmb1Ayr2WzLAWjeSATAoxwKYJV94TeNavgoIdA0a9ytzDSVzBy2YKFK+emCPOEibLeCrg==
+  dependencies:
+    decode-named-character-reference "^1.0.0"
+    devlop "^1.0.0"
+    micromark-factory-destination "^2.0.0"
+    micromark-factory-label "^2.0.0"
+    micromark-factory-space "^2.0.0"
+    micromark-factory-title "^2.0.0"
+    micromark-factory-whitespace "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-chunked "^2.0.0"
+    micromark-util-classify-character "^2.0.0"
+    micromark-util-html-tag-name "^2.0.0"
+    micromark-util-normalize-identifier "^2.0.0"
+    micromark-util-resolve-all "^2.0.0"
+    micromark-util-subtokenize "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-extension-gfm-autolink-literal@^2.0.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/micromark-extension-gfm-autolink-literal/-/micromark-extension-gfm-autolink-literal-2.1.0.tgz#6286aee9686c4462c1e3552a9d505feddceeb935"
+  integrity sha512-oOg7knzhicgQ3t4QCjCWgTmfNhvQbDDnJeVu9v81r7NltNCVmhPy1fJRX27pISafdjL+SVc4d3l48Gb6pbRypw==
+  dependencies:
+    micromark-util-character "^2.0.0"
+    micromark-util-sanitize-uri "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-extension-gfm-footnote@^2.0.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/micromark-extension-gfm-footnote/-/micromark-extension-gfm-footnote-2.1.0.tgz#4dab56d4e398b9853f6fe4efac4fc9361f3e0750"
+  integrity sha512-/yPhxI1ntnDNsiHtzLKYnE3vf9JZ6cAisqVDauhp4CEHxlb4uoOTxOCJ+9s51bIB8U1N1FJ1RXOKTIlD5B/gqw==
+  dependencies:
+    devlop "^1.0.0"
+    micromark-core-commonmark "^2.0.0"
+    micromark-factory-space "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-normalize-identifier "^2.0.0"
+    micromark-util-sanitize-uri "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-extension-gfm-strikethrough@^2.0.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/micromark-extension-gfm-strikethrough/-/micromark-extension-gfm-strikethrough-2.1.0.tgz#86106df8b3a692b5f6a92280d3879be6be46d923"
+  integrity sha512-ADVjpOOkjz1hhkZLlBiYA9cR2Anf8F4HqZUO6e5eDcPQd0Txw5fxLzzxnEkSkfnD0wziSGiv7sYhk/ktvbf1uw==
+  dependencies:
+    devlop "^1.0.0"
+    micromark-util-chunked "^2.0.0"
+    micromark-util-classify-character "^2.0.0"
+    micromark-util-resolve-all "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-extension-gfm-table@^2.0.0:
+  version "2.1.2"
+  resolved "https://registry.yarnpkg.com/micromark-extension-gfm-table/-/micromark-extension-gfm-table-2.1.2.tgz#54dc246873c0c0bebd84600e578aeb8a5f5de78d"
+  integrity sha512-pRzm4kDTu0MjlmBkxmS9yYhw60nncfcEwu9NNdPFSQEFXS95ZKyIIyTSHu/o3ReBUrLKYEq+7YaXCRn/bPB4MA==
+  dependencies:
+    devlop "^1.0.0"
+    micromark-factory-space "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-extension-gfm-tagfilter@^2.0.0:
+  version "2.0.0"
+  resolved "https://registry.yarnpkg.com/micromark-extension-gfm-tagfilter/-/micromark-extension-gfm-tagfilter-2.0.0.tgz#f26d8a7807b5985fba13cf61465b58ca5ff7dc57"
+  integrity sha512-xHlTOmuCSotIA8TW1mDIM6X2O1SiX5P9IuDtqGonFhEK0qgRI4yeC6vMxEV2dgyr2TiD+2PQ10o+cOhdVAcwfg==
+  dependencies:
+    micromark-util-types "^2.0.0"
+
+micromark-extension-gfm-task-list-item@^2.0.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/micromark-extension-gfm-task-list-item/-/micromark-extension-gfm-task-list-item-2.1.0.tgz#bcc34d805639829990ec175c3eea12bb5b781f2c"
+  integrity sha512-qIBZhqxqI6fjLDYFTBIa4eivDMnP+OZqsNwmQ3xNLE4Cxwc+zfQEfbs6tzAo2Hjq+bh6q5F+Z8/cksrLFYWQQw==
+  dependencies:
+    devlop "^1.0.0"
+    micromark-factory-space "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-extension-gfm@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.yarnpkg.com/micromark-extension-gfm/-/micromark-extension-gfm-3.0.0.tgz#3e13376ab95dd7a5cfd0e29560dfe999657b3c5b"
+  integrity sha512-vsKArQsicm7t0z2GugkCKtZehqUm31oeGBV/KVSorWSy8ZlNAv7ytjFhvaryUiCUJYqs+NoE6AFhpQvBTM6Q4w==
+  dependencies:
+    micromark-extension-gfm-autolink-literal "^2.0.0"
+    micromark-extension-gfm-footnote "^2.0.0"
+    micromark-extension-gfm-strikethrough "^2.0.0"
+    micromark-extension-gfm-table "^2.0.0"
+    micromark-extension-gfm-tagfilter "^2.0.0"
+    micromark-extension-gfm-task-list-item "^2.0.0"
+    micromark-util-combine-extensions "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-extension-math@^3.0.0:
+  version "3.1.0"
+  resolved "https://registry.yarnpkg.com/micromark-extension-math/-/micromark-extension-math-3.1.0.tgz#c42ee3b1dd5a9a03584e83dd8f08e3de510212c1"
+  integrity sha512-lvEqd+fHjATVs+2v/8kg9i5Q0AP2k85H0WUOwpIVvUML8BapsMvh1XAogmQjOCsLpoKRCVQqEkQBB3NhVBcsOg==
+  dependencies:
+    "@types/katex" "^0.16.0"
+    devlop "^1.0.0"
+    katex "^0.16.0"
+    micromark-factory-space "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-factory-destination@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-factory-destination/-/micromark-factory-destination-2.0.1.tgz#8fef8e0f7081f0474fbdd92deb50c990a0264639"
+  integrity sha512-Xe6rDdJlkmbFRExpTOmRj9N3MaWmbAgdpSrBQvCFqhezUn4AHqJHbaEnfbVYYiexVSs//tqOdY/DxhjdCiJnIA==
+  dependencies:
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-factory-label@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-factory-label/-/micromark-factory-label-2.0.1.tgz#5267efa97f1e5254efc7f20b459a38cb21058ba1"
+  integrity sha512-VFMekyQExqIW7xIChcXn4ok29YE3rnuyveW3wZQWWqF4Nv9Wk5rgJ99KzPvHjkmPXF93FXIbBp6YdW3t71/7Vg==
+  dependencies:
+    devlop "^1.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-factory-space@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-factory-space/-/micromark-factory-space-2.0.1.tgz#36d0212e962b2b3121f8525fc7a3c7c029f334fc"
+  integrity sha512-zRkxjtBxxLd2Sc0d+fbnEunsTj46SWXgXciZmHq0kDYGnck/ZSGj9/wULTV95uoeYiK5hRXP2mJ98Uo4cq/LQg==
+  dependencies:
+    micromark-util-character "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-factory-title@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-factory-title/-/micromark-factory-title-2.0.1.tgz#237e4aa5d58a95863f01032d9ee9b090f1de6e94"
+  integrity sha512-5bZ+3CjhAd9eChYTHsjy6TGxpOFSKgKKJPJxr293jTbfry2KDoWkhBb6TcPVB4NmzaPhMs1Frm9AZH7OD4Cjzw==
+  dependencies:
+    micromark-factory-space "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-factory-whitespace@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-factory-whitespace/-/micromark-factory-whitespace-2.0.1.tgz#06b26b2983c4d27bfcc657b33e25134d4868b0b1"
+  integrity sha512-Ob0nuZ3PKt/n0hORHyvoD9uZhr+Za8sFoP+OnMcnWK5lngSzALgQYKMr9RJVOWLqQYuyn6ulqGWSXdwf6F80lQ==
+  dependencies:
+    micromark-factory-space "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-util-character@^2.0.0:
+  version "2.1.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-character/-/micromark-util-character-2.1.1.tgz#2f987831a40d4c510ac261e89852c4e9703ccda6"
+  integrity sha512-wv8tdUTJ3thSFFFJKtpYKOYiGP2+v96Hvk4Tu8KpCAsTMs6yi+nVmGh1syvSCsaxz45J6Jbw+9DD6g97+NV67Q==
+  dependencies:
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-util-chunked@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-chunked/-/micromark-util-chunked-2.0.1.tgz#47fbcd93471a3fccab86cff03847fc3552db1051"
+  integrity sha512-QUNFEOPELfmvv+4xiNg2sRYeS/P84pTW0TCgP5zc9FpXetHY0ab7SxKyAQCNCc1eK0459uoLI1y5oO5Vc1dbhA==
+  dependencies:
+    micromark-util-symbol "^2.0.0"
+
+micromark-util-classify-character@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-classify-character/-/micromark-util-classify-character-2.0.1.tgz#d399faf9c45ca14c8b4be98b1ea481bced87b629"
+  integrity sha512-K0kHzM6afW/MbeWYWLjoHQv1sgg2Q9EccHEDzSkxiP/EaagNzCm7T/WMKZ3rjMbvIpvBiZgwR3dKMygtA4mG1Q==
+  dependencies:
+    micromark-util-character "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-util-combine-extensions@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-combine-extensions/-/micromark-util-combine-extensions-2.0.1.tgz#2a0f490ab08bff5cc2fd5eec6dd0ca04f89b30a9"
+  integrity sha512-OnAnH8Ujmy59JcyZw8JSbK9cGpdVY44NKgSM7E9Eh7DiLS2E9RNQf0dONaGDzEG9yjEl5hcqeIsj4hfRkLH/Bg==
+  dependencies:
+    micromark-util-chunked "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-util-decode-numeric-character-reference@^2.0.0:
+  version "2.0.2"
+  resolved "https://registry.yarnpkg.com/micromark-util-decode-numeric-character-reference/-/micromark-util-decode-numeric-character-reference-2.0.2.tgz#fcf15b660979388e6f118cdb6bf7d79d73d26fe5"
+  integrity sha512-ccUbYk6CwVdkmCQMyr64dXz42EfHGkPQlBj5p7YVGzq8I7CtjXZJrubAYezf7Rp+bjPseiROqe7G6foFd+lEuw==
+  dependencies:
+    micromark-util-symbol "^2.0.0"
+
+micromark-util-decode-string@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-decode-string/-/micromark-util-decode-string-2.0.1.tgz#6cb99582e5d271e84efca8e61a807994d7161eb2"
+  integrity sha512-nDV/77Fj6eH1ynwscYTOsbK7rR//Uj0bZXBwJZRfaLEJ1iGBR6kIfNmlNqaqJf649EP0F3NWNdeJi03elllNUQ==
+  dependencies:
+    decode-named-character-reference "^1.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-decode-numeric-character-reference "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+
+micromark-util-encode@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-encode/-/micromark-util-encode-2.0.1.tgz#0d51d1c095551cfaac368326963cf55f15f540b8"
+  integrity sha512-c3cVx2y4KqUnwopcO9b/SCdo2O67LwJJ/UyqGfbigahfegL9myoEFoDYZgkT7f36T0bLrM9hZTAaAyH+PCAXjw==
+
+micromark-util-html-tag-name@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-html-tag-name/-/micromark-util-html-tag-name-2.0.1.tgz#e40403096481986b41c106627f98f72d4d10b825"
+  integrity sha512-2cNEiYDhCWKI+Gs9T0Tiysk136SnR13hhO8yW6BGNyhOC4qYFnwF1nKfD3HFAIXA5c45RrIG1ub11GiXeYd1xA==
+
+micromark-util-normalize-identifier@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-normalize-identifier/-/micromark-util-normalize-identifier-2.0.1.tgz#c30d77b2e832acf6526f8bf1aa47bc9c9438c16d"
+  integrity sha512-sxPqmo70LyARJs0w2UclACPUUEqltCkJ6PhKdMIDuJ3gSf/Q+/GIe3WKl0Ijb/GyH9lOpUkRAO2wp0GVkLvS9Q==
+  dependencies:
+    micromark-util-symbol "^2.0.0"
+
+micromark-util-resolve-all@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-resolve-all/-/micromark-util-resolve-all-2.0.1.tgz#e1a2d62cdd237230a2ae11839027b19381e31e8b"
+  integrity sha512-VdQyxFWFT2/FGJgwQnJYbe1jjQoNTS4RjglmSjTUlpUMa95Htx9NHeYW4rGDJzbjvCsl9eLjMQwGeElsqmzcHg==
+  dependencies:
+    micromark-util-types "^2.0.0"
+
+micromark-util-sanitize-uri@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-sanitize-uri/-/micromark-util-sanitize-uri-2.0.1.tgz#ab89789b818a58752b73d6b55238621b7faa8fd7"
+  integrity sha512-9N9IomZ/YuGGZZmQec1MbgxtlgougxTodVwDzzEouPKo3qFWvymFHWcnDi2vzV1ff6kas9ucW+o3yzJK9YB1AQ==
+  dependencies:
+    micromark-util-character "^2.0.0"
+    micromark-util-encode "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+
+micromark-util-subtokenize@^2.0.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/micromark-util-subtokenize/-/micromark-util-subtokenize-2.1.0.tgz#d8ade5ba0f3197a1cf6a2999fbbfe6357a1a19ee"
+  integrity sha512-XQLu552iSctvnEcgXw6+Sx75GflAPNED1qx7eBJ+wydBb2KCbRZe+NwvIEEMM83uml1+2WSXpBAcp9IUCgCYWA==
+  dependencies:
+    devlop "^1.0.0"
+    micromark-util-chunked "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+micromark-util-symbol@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/micromark-util-symbol/-/micromark-util-symbol-2.0.1.tgz#e5da494e8eb2b071a0d08fb34f6cefec6c0a19b8"
+  integrity sha512-vs5t8Apaud9N28kgCrRUdEed4UJ+wWNvicHLPxCa9ENlYuAY31M0ETy5y1vA33YoNPDFTghEbnh6efaE8h4x0Q==
+
+micromark-util-types@^2.0.0:
+  version "2.0.2"
+  resolved "https://registry.yarnpkg.com/micromark-util-types/-/micromark-util-types-2.0.2.tgz#f00225f5f5a0ebc3254f96c36b6605c4b393908e"
+  integrity sha512-Yw0ECSpJoViF1qTU4DC6NwtC4aWGt1EkzaQB8KPPyCRR8z9TWeV0HbEFGTO+ZY1wB22zmxnJqhPyTpOVCpeHTA==
+
+micromark@^4.0.0:
+  version "4.0.2"
+  resolved "https://registry.yarnpkg.com/micromark/-/micromark-4.0.2.tgz#91395a3e1884a198e62116e33c9c568e39936fdb"
+  integrity sha512-zpe98Q6kvavpCr1NPVSCMebCKfD7CA2NqZ+rykeNhONIJBpc1tFKt9hucLGwha3jNTNI8lHpctWJWoimVF4PfA==
+  dependencies:
+    "@types/debug" "^4.0.0"
+    debug "^4.0.0"
+    decode-named-character-reference "^1.0.0"
+    devlop "^1.0.0"
+    micromark-core-commonmark "^2.0.0"
+    micromark-factory-space "^2.0.0"
+    micromark-util-character "^2.0.0"
+    micromark-util-chunked "^2.0.0"
+    micromark-util-combine-extensions "^2.0.0"
+    micromark-util-decode-numeric-character-reference "^2.0.0"
+    micromark-util-encode "^2.0.0"
+    micromark-util-normalize-identifier "^2.0.0"
+    micromark-util-resolve-all "^2.0.0"
+    micromark-util-sanitize-uri "^2.0.0"
+    micromark-util-subtokenize "^2.0.0"
+    micromark-util-symbol "^2.0.0"
+    micromark-util-types "^2.0.0"
+
+ms@^2.1.3:
+  version "2.1.3"
+  resolved "https://registry.yarnpkg.com/ms/-/ms-2.1.3.tgz#574c8138ce1d2b5861f0b44579dbadd60c6615b2"
+  integrity sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==
+
+nanoid@^3.3.18:
+  version "3.3.19"
+  resolved "https://registry.yarnpkg.com/nanoid/-/nanoid-3.3.19.tgz#336d4aa4bcd4fb24d2cddede7ffeae40bec03f0a"
+  integrity sha512-Y2tUNy4ouw6tq5oDSKeQYGOyhkUBhNOcGV/02KC+6kd9eDGqdZd++mjMiIDilrBYvjEnCYvVtsuHCuP+okSfug==
+
+parse-entities@^4.0.0:
+  version "4.0.2"
+  resolved "https://registry.yarnpkg.com/parse-entities/-/parse-entities-4.0.2.tgz#61d46f5ed28e4ee62e9ddc43d6b010188443f159"
+  integrity sha512-GG2AQYWoLgL877gQIKeRPGO1xF9+eG1ujIb5soS5gPvLQ1y2o8FL90w2QWNdf9I361Mpp7726c+lj3U0qK1uGw==
+  dependencies:
+    "@types/unist" "^2.0.0"
+    character-entities-legacy "^3.0.0"
+    character-reference-invalid "^2.0.0"
+    decode-named-character-reference "^1.0.0"
+    is-alphanumerical "^2.0.0"
+    is-decimal "^2.0.0"
+    is-hexadecimal "^2.0.0"
+
+parse5@^7.0.0:
+  version "7.3.0"
+  resolved "https://registry.yarnpkg.com/parse5/-/parse5-7.3.0.tgz#d7e224fa72399c7a175099f45fc2ad024b05ec05"
+  integrity sha512-IInvU7fabl34qmi9gY8XOVxhYyMyuH2xUNpb2q8/Y+7552KlejkRvqvD19nMoUW/uQGGbqNpA6Tufu5FL5BZgw==
+  dependencies:
+    entities "^6.0.0"
+
+picocolors@^1.1.1:
+  version "1.1.1"
+  resolved "https://registry.yarnpkg.com/picocolors/-/picocolors-1.1.1.tgz#3d321af3eab939b083c8f929a1d12cda81c26b6b"
+  integrity sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==
+
+picomatch@^4.0.4, picomatch@^4.0.7:
+  version "4.0.7"
+  resolved "https://registry.yarnpkg.com/picomatch/-/picomatch-4.0.7.tgz#6313360034ccb36b3dc61ecbdff78121f90fe21f"
+  integrity sha512-qcJu88Q2IWqJsDD529JKMdwGm/dvInW4HvQnRwiH9JtihJvzGOscDtHE3x1pBKeUOTysQ8kVmLnJ2kJu7yhcGA==
+
+postcss@^8.5.28:
+  version "8.5.28"
+  resolved "https://registry.yarnpkg.com/postcss/-/postcss-8.5.28.tgz#da4563a99a06e62d6c1cd1acae363224bcaed6e9"
+  integrity sha512-RRuzqDtt5Y9h3quz5hWhK+TPnsmVs6WwSU6LkJMeY4HstUEDuYTG8UJSdawMRzmzAtV+KEoG8N3Qg2qLy5vM/A==
+  dependencies:
+    nanoid "^3.3.18"
+    picocolors "^1.1.1"
+    source-map-js "^1.2.1"
+
+property-information@^7.0.0:
+  version "7.2.0"
+  resolved "https://registry.yarnpkg.com/property-information/-/property-information-7.2.0.tgz#0809b34264e995c0bfcd3227028a1e35210af80a"
+  integrity sha512-IAtzIB6sUiWaJYrX9smp3V46pBGbBeLFRGdh25kg1334VcBlD8HzhPeNIWQH9zhGmo2itIe25EHt9dQP7G5hmg==
+
+react-dom@^19.3.0:
+  version "19.3.0"
+  resolved "https://registry.yarnpkg.com/react-dom/-/react-dom-19.3.0.tgz#a0ff19161d7f4e9be5469a07278b0c06e00f9001"
+  integrity sha512-JDk8dgif51OjFoDE70+OT9ICyYr+69HlmihNwp1+Nsfbna3t5sIiCa9ZJktDmQ4/1b/rn26hIAR2uYXDMr5r0Q==
+  dependencies:
+    scheduler "^0.28.0"
+
+react-i18next@^17.0.14:
+  version "17.0.14"
+  resolved "https://registry.yarnpkg.com/react-i18next/-/react-i18next-17.0.14.tgz#a895ff593e461877953732c194a242fa3895ed61"
+  integrity sha512-ZpMBfJL3BiXPuYHj5QMY1GwvKJNhE1vCxOJQ8BoMBdHP5XONTfi/Qss1nuRlnFdIl72PKo2jPmw+fxuLjRFgaQ==
+  dependencies:
+    "@babel/runtime" "^7.29.7"
+    html-parse-stringify "^4.0.1"
+    use-sync-external-store "^1.6.0"
+
+react@^19.3.0:
+  version "19.3.0"
+  resolved "https://registry.yarnpkg.com/react/-/react-19.3.0.tgz#e20aa3a14687075ce863351b5878b0608c3fbfb2"
+  integrity sha512-E8LUcbtBWt20bbl2YoHfx4ZDBdxVTfOKtCZn9cDSJ4l6/nuoApcpIBcj47t2wZoVX8g2ZHuMHbiShgCR1T5Sog==
+
+rehype-harden@^1.1.8:
+  version "1.1.8"
+  resolved "https://registry.yarnpkg.com/rehype-harden/-/rehype-harden-1.1.8.tgz#c5326fb691f64daa40711e44f0b7ce1a75e620d3"
+  integrity sha512-Qn7vR1xrf6fZCrkm9TDWi/AB4ylrHy+jqsNm1EHOAmbARYA6gsnVJBq/sdBh6kmT4NEZxH5vgIjrscefJAOXcw==
+  dependencies:
+    unist-util-visit "^5.0.0"
+
+rehype-katex@^7.0.1:
+  version "7.0.1"
+  resolved "https://registry.yarnpkg.com/rehype-katex/-/rehype-katex-7.0.1.tgz#832e6d7af2744a228981d1b0fe89483a9e7c93a1"
+  integrity sha512-OiM2wrZ/wuhKkigASodFoo8wimG3H12LWQaH8qSPVJn9apWKFSH3YOCtbKpBorTVw/eI7cuT21XBbvwEswbIOA==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@types/katex" "^0.16.0"
+    hast-util-from-html-isomorphic "^2.0.0"
+    hast-util-to-text "^4.0.0"
+    katex "^0.16.0"
+    unist-util-visit-parents "^6.0.0"
+    vfile "^6.0.0"
+
+rehype-raw@^7.0.0:
+  version "7.0.0"
+  resolved "https://registry.yarnpkg.com/rehype-raw/-/rehype-raw-7.0.0.tgz#59d7348fd5dbef3807bbaa1d443efd2dd85ecee4"
+  integrity sha512-/aE8hCfKlQeA8LmyeyQvQF3eBiLRGNlfBJEvWH7ivp9sBqs7TNqBL5X3v157rM4IFETqDnIOO+z5M/biZbo9Ww==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    hast-util-raw "^9.0.0"
+    vfile "^6.0.0"
+
+rehype-sanitize@^6.0.0:
+  version "6.0.0"
+  resolved "https://registry.yarnpkg.com/rehype-sanitize/-/rehype-sanitize-6.0.0.tgz#16e95f4a67a69cbf0f79e113c8e0df48203db73c"
+  integrity sha512-CsnhKNsyI8Tub6L4sm5ZFsme4puGfc6pYylvXo1AeqaGbjOYyzNv3qZPwvs0oMJ39eryyeOdmxwUIo94IpEhqg==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    hast-util-sanitize "^5.0.0"
+
+remark-gfm@^4.0.1:
+  version "4.0.1"
+  resolved "https://registry.yarnpkg.com/remark-gfm/-/remark-gfm-4.0.1.tgz#33227b2a74397670d357bf05c098eaf8513f0d6b"
+  integrity sha512-1quofZ2RQ9EWdeN34S79+KExV1764+wCUGop5CPL1WGdD0ocPpu91lzPGbwWMECpEpd42kJGQwzRfyov9j4yNg==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    mdast-util-gfm "^3.0.0"
+    micromark-extension-gfm "^3.0.0"
+    remark-parse "^11.0.0"
+    remark-stringify "^11.0.0"
+    unified "^11.0.0"
+
+remark-math@^6.0.0:
+  version "6.0.0"
+  resolved "https://registry.yarnpkg.com/remark-math/-/remark-math-6.0.0.tgz#0acdf74675f1c195fea6efffa78582f7ed7fc0d7"
+  integrity sha512-MMqgnP74Igy+S3WwnhQ7kqGlEerTETXMvJhrUzDikVZ2/uogJCb+WHUg97hK9/jcfc0dkD73s3LN8zU49cTEtA==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    mdast-util-math "^3.0.0"
+    micromark-extension-math "^3.0.0"
+    unified "^11.0.0"
+
+remark-parse@^11.0.0:
+  version "11.0.0"
+  resolved "https://registry.yarnpkg.com/remark-parse/-/remark-parse-11.0.0.tgz#aa60743fcb37ebf6b069204eb4da304e40db45a1"
+  integrity sha512-FCxlKLNGknS5ba/1lmpYijMUzX2esxW5xQqjWxw2eHFfS2MSdaHVINFmhjo+qN1WhZhNimq0dZATN9pH0IDrpA==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    mdast-util-from-markdown "^2.0.0"
+    micromark-util-types "^2.0.0"
+    unified "^11.0.0"
+
+remark-rehype@^11.1.2:
+  version "11.1.2"
+  resolved "https://registry.yarnpkg.com/remark-rehype/-/remark-rehype-11.1.2.tgz#2addaadda80ca9bd9aa0da763e74d16327683b37"
+  integrity sha512-Dh7l57ianaEoIpzbp0PC9UKAdCSVklD8E5Rpw7ETfbTl3FqcOOgq5q2LVDhgGCkaBv7p24JXikPdvhhmHvKMsw==
+  dependencies:
+    "@types/hast" "^3.0.0"
+    "@types/mdast" "^4.0.0"
+    mdast-util-to-hast "^13.0.0"
+    unified "^11.0.0"
+    vfile "^6.0.0"
+
+remark-stringify@^11.0.0:
+  version "11.0.0"
+  resolved "https://registry.yarnpkg.com/remark-stringify/-/remark-stringify-11.0.0.tgz#4c5b01dd711c269df1aaae11743eb7e2e7636fd3"
+  integrity sha512-1OSmLd3awB/t8qdoEOMazZkNsfVTeY4fTsgzcQFdXNq8ToTN4ZGwrMnlda4K6smTFKD+GRV6O48i6Z4iKgPPpw==
+  dependencies:
+    "@types/mdast" "^4.0.0"
+    mdast-util-to-markdown "^2.0.0"
+    unified "^11.0.0"
+
+remend@1.3.1:
+  version "1.3.1"
+  resolved "https://registry.yarnpkg.com/remend/-/remend-1.3.1.tgz#474575f709d39d01f5d3b04333d4fa40d4491133"
+  integrity sha512-N3DiY5qbRPoa5vkxn1oDLMyXOVTeo6Hp+XOj6SIqJAYUgLS0Q587gILPMom/qm86AQ/ZrcOdwEIzCz8V3J0nxQ==
+
+rolldown@~1.2.6:
+  version "1.2.8"
+  resolved "https://registry.yarnpkg.com/rolldown/-/rolldown-1.2.8.tgz#3a18ad3c74809f05abe9130229b2b8b4249e9e21"
+  integrity sha512-Z67nTmhZe7anqnM/EjI392w5i/ANUinjip7QYsOyN37oayduxt3ksdX0hf5OOamkAd53BiIHfbfSzfUmzKFQqQ==
+  dependencies:
+    "@oxc-project/types" "=0.149.0"
+    "@rolldown/pluginutils" "^1.0.0"
+  optionalDependencies:
+    "@rolldown/binding-android-arm-eabi" "1.2.8"
+    "@rolldown/binding-android-arm64" "1.2.8"
+    "@rolldown/binding-darwin-arm64" "1.2.8"
+    "@rolldown/binding-darwin-x64" "1.2.8"
+    "@rolldown/binding-freebsd-x64" "1.2.8"
+    "@rolldown/binding-linux-arm-gnueabihf" "1.2.8"
+    "@rolldown/binding-linux-arm64-gnu" "1.2.8"
+    "@rolldown/binding-linux-arm64-musl" "1.2.8"
+    "@rolldown/binding-linux-ppc64-gnu" "1.2.8"
+    "@rolldown/binding-linux-s390x-gnu" "1.2.8"
+    "@rolldown/binding-linux-x64-gnu" "1.2.8"
+    "@rolldown/binding-linux-x64-musl" "1.2.8"
+    "@rolldown/binding-openharmony-arm64" "1.2.8"
+    "@rolldown/binding-win32-arm64-msvc" "1.2.8"
+    "@rolldown/binding-win32-x64-msvc" "1.2.8"
+
+scheduler@^0.28.0:
+  version "0.28.0"
+  resolved "https://registry.yarnpkg.com/scheduler/-/scheduler-0.28.0.tgz#03d3ca348aeaf8b7502f64468f217b27523f9695"
+  integrity sha512-juorfCmIkIw8tT+p5BXSm6PJjQF/ycEYmKyzURCIt/RaZIhL+PulbQ9Yu2z1HdOJDdqDTlxA1+xKBmHXJsczAw==
+
+seroval-plugins@^1.6.2:
+  version "1.6.7"
+  resolved "https://registry.yarnpkg.com/seroval-plugins/-/seroval-plugins-1.6.7.tgz#4aba839a2fdafa9c115d965a71861d4b288b90c1"
+  integrity sha512-4Nk35ttD3DTDJW4hgw5StsVAPeU6qnDFnULAouw6tQ7oLTV/ICXrWpsXo2EE52eSP2joUMazbVf52mFEcADqRw==
+
+seroval@^1.6.2:
+  version "1.6.7"
+  resolved "https://registry.yarnpkg.com/seroval/-/seroval-1.6.7.tgz#b33953ac59aec5bde14f99e010bdaa7a507ae2db"
+  integrity sha512-AeDcLh0yO2SFm9W71essgnSzLV9DI8ZH0x0knXn2DMnUZj728mpLbxjlbB6IqKCmqh8JA3cEqRyGoNkt584JcQ==
+
+source-map-js@^1.2.1:
+  version "1.2.1"
+  resolved "https://registry.yarnpkg.com/source-map-js/-/source-map-js-1.2.1.tgz#1ce5650fddd87abc099eda37dcff024c2667ae46"
+  integrity sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==
+
+space-separated-tokens@^2.0.0:
+  version "2.0.2"
+  resolved "https://registry.yarnpkg.com/space-separated-tokens/-/space-separated-tokens-2.0.2.tgz#1ecd9d2350a3844572c3f4a312bceb018348859f"
+  integrity sha512-PEGlAwrG8yXGXRjW32fGbg66JAlOAwbObuqVoJpv/mRgoWDQfgH1wDPvtzWyUSNAXBGSk8h755YDbbcEy3SH2Q==
+
+streamdown@^2.6.0:
+  version "2.6.0"
+  resolved "https://registry.yarnpkg.com/streamdown/-/streamdown-2.6.0.tgz#0d1793135b047a374ba8330882319849ecc048d9"
+  integrity sha512-nQZVUn4GvB2R5SAlDNph20iKZeW+RM4gv6G1H3ACypEnmQf8PgTHZ/1Ta2OACRwQ2coK7aW5AeIAa7Ls5zk+RA==
+  dependencies:
+    clsx "^2.1.1"
+    hast-util-to-jsx-runtime "^2.3.6"
+    html-url-attributes "^3.0.1"
+    marked "^17.0.1"
+    rehype-harden "^1.1.8"
+    rehype-raw "^7.0.0"
+    rehype-sanitize "^6.0.0"
+    remark-gfm "^4.0.1"
+    remark-parse "^11.0.0"
+    remark-rehype "^11.1.2"
+    remend "1.3.1"
+    tailwind-merge "^3.6.0"
+    unified "^11.0.5"
+    unist-util-visit "^5.1.0"
+    unist-util-visit-parents "^6.0.2"
+
+stringify-entities@^4.0.0:
+  version "4.0.4"
+  resolved "https://registry.yarnpkg.com/stringify-entities/-/stringify-entities-4.0.4.tgz#b3b79ef5f277cc4ac73caeb0236c5ba939b3a4f3"
+  integrity sha512-IwfBptatlO+QCJUo19AqvrPNqlVMpW9YEL2LIVY+Rpv2qsjCGxaDLNRgeGsQWJhfItebuJhsGSLjaBbNSQ+ieg==
+  dependencies:
+    character-entities-html4 "^2.0.0"
+    character-entities-legacy "^3.0.0"
+
+style-to-js@^1.0.0:
+  version "1.1.21"
+  resolved "https://registry.yarnpkg.com/style-to-js/-/style-to-js-1.1.21.tgz#2908941187f857e79e28e9cd78008b9a0b3e0e8d"
+  integrity sha512-RjQetxJrrUJLQPHbLku6U/ocGtzyjbJMP9lCNK7Ag0CNh690nSH8woqWH9u16nMjYBAok+i7JO1NP2pOy8IsPQ==
+  dependencies:
+    style-to-object "1.0.14"
+
+style-to-object@1.0.14:
+  version "1.0.14"
+  resolved "https://registry.yarnpkg.com/style-to-object/-/style-to-object-1.0.14.tgz#1d22f0e7266bb8c6d8cae5caf4ec4f005e08f611"
+  integrity sha512-LIN7rULI0jBscWQYaSswptyderlarFkjQ+t79nzty8tcIAceVomEVlLzH5VP4Cmsv6MtKhs7qaAiwlcp+Mgaxw==
+  dependencies:
+    inline-style-parser "0.2.7"
+
+swr@^2.2.5:
+  version "2.5.1"
+  resolved "https://registry.yarnpkg.com/swr/-/swr-2.5.1.tgz#3d1801fafe06d13b149bc52c62323dcd43ca570c"
+  integrity sha512-BRw55e8r0B7SpDN20CAzoQAHl7y1yP7/Zt7oqUjMv0vSt2u2Xnkm88Ws+VypbV9BXHQVuSuyVq7zMjO16wSExw==
+  dependencies:
+    dequal "^2.0.3"
+    use-sync-external-store "^1.6.0"
+
+tailwind-merge@^3.6.0:
+  version "3.7.0"
+  resolved "https://registry.yarnpkg.com/tailwind-merge/-/tailwind-merge-3.7.0.tgz#61e9a67ab11b64138dea7d4fa5b74a33e17f2c87"
+  integrity sha512-XPPUyAc+cvspz3lHTcR/QgPfW2A0lv/xQNIjX3HGhLR+Nq2lHaLq5MtTesHn8GUr3W3DguT2KT5x3NVgRtYwmA==
+
+tailwindcss@4.3.3, tailwindcss@^4.3.3:
+  version "4.3.3"
+  resolved "https://registry.yarnpkg.com/tailwindcss/-/tailwindcss-4.3.3.tgz#c006861611c213c1877893ab5b23daa16be2bb55"
+  integrity sha512-gOhV3P7ufE62QDGg1zVaTgCR+EtPv92k2nIhVcVKcLmxT1sUBsQGhnZj175j+MqRt4zLF7ic+sCYjfhxMxj7YQ==
+
+tapable@^2.3.3:
+  version "2.3.3"
+  resolved "https://registry.yarnpkg.com/tapable/-/tapable-2.3.3.tgz#5da7c9992c46038221267985ab28421a8879f160"
+  integrity sha512-uxc/zpqFg6x7C8vOE7lh6Lbda8eEL9zmVm/PLeTPBRhh1xCgdWaQ+J1CUieGpIfm2HdtsUpRv+HshiasBMcc6A==
+
+throttleit@2.1.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/throttleit/-/throttleit-2.1.0.tgz#a7e4aa0bf4845a5bd10daa39ea0c783f631a07b4"
+  integrity sha512-nt6AMGKW1p/70DF/hGBdJB57B8Tspmbp5gfJ8ilhLnt7kkr2ye7hzD6NVG8GGErk2HWF34igrL2CXmNIkzKqKw==
+
+tinyglobby@^0.2.17:
+  version "0.2.17"
+  resolved "https://registry.yarnpkg.com/tinyglobby/-/tinyglobby-0.2.17.tgz#562a9a6c9eb2b3b123d39719f9af5bb44fcd7631"
+  integrity sha512-wXR/dYpcqKmfWpEdZjiKJOwCNFndD0DMnrW/cYjVGttEkBfVgcLFHoNrlj47mjOVic9yyNu65alsgF4NQyTa2g==
+  dependencies:
+    fdir "^6.5.0"
+    picomatch "^4.0.4"
+
+trim-lines@^3.0.0:
+  version "3.0.1"
+  resolved "https://registry.yarnpkg.com/trim-lines/-/trim-lines-3.0.1.tgz#d802e332a07df861c48802c04321017b1bd87338"
+  integrity sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg==
+
+trough@^2.0.0:
+  version "2.2.0"
+  resolved "https://registry.yarnpkg.com/trough/-/trough-2.2.0.tgz#94a60bd6bd375c152c1df911a4b11d5b0256f50f"
+  integrity sha512-tmMpK00BjZiUyVyvrBK7knerNgmgvcV/KLVyuma/SC+TQN167GrMRciANTz09+k3zW8L8t60jWO1GpfkZdjTaw==
+
+tslib@^2.4.0, tslib@^2.8.1:
+  version "2.8.1"
+  resolved "https://registry.yarnpkg.com/tslib/-/tslib-2.8.1.tgz#612efe4ed235d567e8aba5f2a5fab70280ade83f"
+  integrity sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==
+
+typescript@^7.0.2:
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/typescript/-/typescript-7.0.2.tgz#9ec773d7954a8c182c17cc5bbd575aa28bc51582"
+  integrity sha512-8FYau96o3NKOhbjKi/qNvG/W5jhzxkbdm5sj9AbZ/5T5sWqn3hJgLfGx27sRKZWTvyzCP8dLRBTf5tBTSRVUNA==
+  optionalDependencies:
+    "@typescript/typescript-aix-ppc64" "7.0.2"
+    "@typescript/typescript-darwin-arm64" "7.0.2"
+    "@typescript/typescript-darwin-x64" "7.0.2"
+    "@typescript/typescript-freebsd-arm64" "7.0.2"
+    "@typescript/typescript-freebsd-x64" "7.0.2"
+    "@typescript/typescript-linux-arm" "7.0.2"
+    "@typescript/typescript-linux-arm64" "7.0.2"
+    "@typescript/typescript-linux-loong64" "7.0.2"
+    "@typescript/typescript-linux-mips64el" "7.0.2"
+    "@typescript/typescript-linux-ppc64" "7.0.2"
+    "@typescript/typescript-linux-riscv64" "7.0.2"
+    "@typescript/typescript-linux-s390x" "7.0.2"
+    "@typescript/typescript-linux-x64" "7.0.2"
+    "@typescript/typescript-netbsd-arm64" "7.0.2"
+    "@typescript/typescript-netbsd-x64" "7.0.2"
+    "@typescript/typescript-openbsd-arm64" "7.0.2"
+    "@typescript/typescript-openbsd-x64" "7.0.2"
+    "@typescript/typescript-sunos-x64" "7.0.2"
+    "@typescript/typescript-win32-arm64" "7.0.2"
+    "@typescript/typescript-win32-x64" "7.0.2"
+
+undici-types@~6.21.0:
+  version "6.21.0"
+  resolved "https://registry.yarnpkg.com/undici-types/-/undici-types-6.21.0.tgz#691d00af3909be93a7faa13be61b3a5b50ef12cb"
+  integrity sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==
+
+undici@^5.29.0:
+  version "5.29.0"
+  resolved "https://registry.yarnpkg.com/undici/-/undici-5.29.0.tgz#419595449ae3f2cdcba3580a2e8903399bd1f5a3"
+  integrity sha512-raqeBD6NQK4SkWhQzeYKd1KmIG6dllBOTt55Rmkt4HtI9mwdWtJljnrXjAFUBLTSN67HWrOIZ3EPF4kjUw80Bg==
+  dependencies:
+    "@fastify/busboy" "^2.0.0"
+
+unified@^11.0.0, unified@^11.0.5:
+  version "11.0.5"
+  resolved "https://registry.yarnpkg.com/unified/-/unified-11.0.5.tgz#f66677610a5c0a9ee90cab2b8d4d66037026d9e1"
+  integrity sha512-xKvGhPWw3k84Qjh8bI3ZeJjqnyadK+GEFtazSfZv/rKeTkTjOJho6mFqh2SM96iIcZokxiOpg78GazTSg8+KHA==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    bail "^2.0.0"
+    devlop "^1.0.0"
+    extend "^3.0.0"
+    is-plain-obj "^4.0.0"
+    trough "^2.0.0"
+    vfile "^6.0.0"
+
+unist-util-find-after@^5.0.0:
+  version "5.0.0"
+  resolved "https://registry.yarnpkg.com/unist-util-find-after/-/unist-util-find-after-5.0.0.tgz#3fccc1b086b56f34c8b798e1ff90b5c54468e896"
+  integrity sha512-amQa0Ep2m6hE2g72AugUItjbuM8X8cGQnFoHk0pGfrFeT9GZhzN5SW8nRsiGKK7Aif4CrACPENkA6P/Lw6fHGQ==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    unist-util-is "^6.0.0"
+
+unist-util-is@^6.0.0:
+  version "6.0.1"
+  resolved "https://registry.yarnpkg.com/unist-util-is/-/unist-util-is-6.0.1.tgz#d0a3f86f2dd0db7acd7d8c2478080b5c67f9c6a9"
+  integrity sha512-LsiILbtBETkDz8I9p1dQ0uyRUWuaQzd/cuEeS1hoRSyW5E5XGmTzlwY1OrNzzakGowI9Dr/I8HVaw4hTtnxy8g==
+  dependencies:
+    "@types/unist" "^3.0.0"
+
+unist-util-position@^5.0.0:
+  version "5.0.0"
+  resolved "https://registry.yarnpkg.com/unist-util-position/-/unist-util-position-5.0.0.tgz#678f20ab5ca1207a97d7ea8a388373c9cf896be4"
+  integrity sha512-fucsC7HjXvkB5R3kTCO7kUjRdrS0BJt3M/FPxmHMBOm8JQi2BsHAHFsy27E0EolP8rp0NzXsJ+jNPyDWvOJZPA==
+  dependencies:
+    "@types/unist" "^3.0.0"
+
+unist-util-remove-position@^5.0.0:
+  version "5.0.0"
+  resolved "https://registry.yarnpkg.com/unist-util-remove-position/-/unist-util-remove-position-5.0.0.tgz#fea68a25658409c9460408bc6b4991b965b52163"
+  integrity sha512-Hp5Kh3wLxv0PHj9m2yZhhLt58KzPtEYKQQ4yxfYFEO7EvHwzyDYnduhHnY1mDxoqr7VUwVuHXk9RXKIiYS1N8Q==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    unist-util-visit "^5.0.0"
+
+unist-util-stringify-position@^4.0.0:
+  version "4.0.0"
+  resolved "https://registry.yarnpkg.com/unist-util-stringify-position/-/unist-util-stringify-position-4.0.0.tgz#449c6e21a880e0855bf5aabadeb3a740314abac2"
+  integrity sha512-0ASV06AAoKCDkS2+xw5RXJywruurpbC4JZSm7nr7MOt1ojAzvyyaO+UxZf18j8FCF6kmzCZKcAgN/yu2gm2XgQ==
+  dependencies:
+    "@types/unist" "^3.0.0"
+
+unist-util-visit-parents@^6.0.0, unist-util-visit-parents@^6.0.2:
+  version "6.0.2"
+  resolved "https://registry.yarnpkg.com/unist-util-visit-parents/-/unist-util-visit-parents-6.0.2.tgz#777df7fb98652ce16b4b7cd999d0a1a40efa3a02"
+  integrity sha512-goh1s1TBrqSqukSc8wrjwWhL0hiJxgA8m4kFxGlQ+8FYQ3C/m11FcTs4YYem7V664AhHVvgoQLk890Ssdsr2IQ==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    unist-util-is "^6.0.0"
+
+unist-util-visit@^5.0.0, unist-util-visit@^5.1.0:
+  version "5.1.0"
+  resolved "https://registry.yarnpkg.com/unist-util-visit/-/unist-util-visit-5.1.0.tgz#9a2a28b0aa76a15e0da70a08a5863a2f060e2468"
+  integrity sha512-m+vIdyeCOpdr/QeQCu2EzxX/ohgS8KbnPDgFni4dQsfSCtpz8UqDyY5GjRru8PDKuYn7Fq19j1CQ+nJSsGKOzg==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    unist-util-is "^6.0.0"
+    unist-util-visit-parents "^6.0.0"
+
+use-sync-external-store@^1.6.0:
+  version "1.7.0"
+  resolved "https://registry.yarnpkg.com/use-sync-external-store/-/use-sync-external-store-1.7.0.tgz#6dcb66ef569e02f186af6b3d575f414ce746e18f"
+  integrity sha512-6L+EeigHMQhdaIPNIFUKwfWJSwWFQ8gJbJ2DLOs5sDIegTwR9fRxvnM3uciHKjIZhFz+KAv2emhWMRvDmMcY8A==
+
+vfile-location@^5.0.0:
+  version "5.0.3"
+  resolved "https://registry.yarnpkg.com/vfile-location/-/vfile-location-5.0.3.tgz#cb9eacd20f2b6426d19451e0eafa3d0a846225c3"
+  integrity sha512-5yXvWDEgqeiYiBe1lbxYF7UMAIm/IcopxMHrMQDq3nvKcjPKIhZklUKL+AE7J7uApI4kwe2snsK+eI6UTj9EHg==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    vfile "^6.0.0"
+
+vfile-message@^4.0.0:
+  version "4.0.3"
+  resolved "https://registry.yarnpkg.com/vfile-message/-/vfile-message-4.0.3.tgz#87b44dddd7b70f0641c2e3ed0864ba73e2ea8df4"
+  integrity sha512-QTHzsGd1EhbZs4AsQ20JX1rC3cOlt/IWJruk893DfLRr57lcnOeMaWG4K0JrRta4mIJZKth2Au3mM3u03/JWKw==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    unist-util-stringify-position "^4.0.0"
+
+vfile@^6.0.0:
+  version "6.0.3"
+  resolved "https://registry.yarnpkg.com/vfile/-/vfile-6.0.3.tgz#3652ab1c496531852bf55a6bac57af981ebc38ab"
+  integrity sha512-KzIbH/9tXat2u30jf+smMwFCsno4wHVdNmzFyL+T/L3UGqqk6JKfVqOFOZEpZSHADH1k40ab6NUIXZq422ov3Q==
+  dependencies:
+    "@types/unist" "^3.0.0"
+    vfile-message "^4.0.0"
+
+vite@^8.3.0:
+  version "8.3.0"
+  resolved "https://registry.yarnpkg.com/vite/-/vite-8.3.0.tgz#f9565cfd4879d58d28fa64d3c7aff3b372d2ccc8"
+  integrity sha512-lhZBVvEHefgE+HQZC9O7EBJgCU/nVzFNl7vkS4RE0APtWLP02/8QVIkQtzBxPquh7lq5/78NHipTj7ODQ6XuyQ==
+  dependencies:
+    lightningcss "^1.33.0"
+    picomatch "^4.0.7"
+    postcss "^8.5.28"
+    rolldown "~1.2.6"
+    tinyglobby "^0.2.17"
+  optionalDependencies:
+    fsevents "~2.3.3"
+
+web-namespaces@^2.0.0:
+  version "2.0.1"
+  resolved "https://registry.yarnpkg.com/web-namespaces/-/web-namespaces-2.0.1.tgz#1010ff7c650eccb2592cebeeaf9a1b253fd40692"
+  integrity sha512-bKr1DkiNa2krS7qxNtdrtHAmzuYGFQLiQ13TsorsdT6ULTkPLKuu5+GsFpDlg6JFjUTwX2DyhMPG2be8uPrqsQ==
+
+zod@^4.1.8:
+  version "4.6.5"
+  resolved "https://registry.yarnpkg.com/zod/-/zod-4.6.5.tgz#e3af3edbdc260239c52f73caad05dbbffe6e82cd"
+  integrity sha512-v5l/aFXZQeai4awLbOpSoHecE9UiMrnfx75tEXLjNonXVARxQ5mOeipTjROUchszUNCqnE+hqAMujRsRHsut2Q==
+
+zwitch@^2.0.0:
+  version "2.0.4"
+  resolved "https://registry.yarnpkg.com/zwitch/-/zwitch-2.0.4.tgz#c827d4b0acb76fc3e685a4c6ec2902d51070e9d7"
+  integrity sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A==
diff --git a/macos/crates/edge0-cli/Cargo.toml b/macos/crates/edge0-cli/Cargo.toml
new file mode 100644
index 0000000..970194e
--- /dev/null
+++ b/macos/crates/edge0-cli/Cargo.toml
@@ -0,0 +1,22 @@
+[package]
+name = "edge0-cli"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+[[bin]]
+name = "edge0"
+path = "src/main.rs"
+
+[dependencies]
+edge0-core.workspace = true
+edge0-pull.workspace = true
+clap.workspace = true
+reqwest.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+chrono.workspace = true
+anyhow.workspace = true
+tracing = { workspace = true }
+tokio.workspace = true
diff --git a/macos/crates/edge0-cli/src/main.rs b/macos/crates/edge0-cli/src/main.rs
new file mode 100644
index 0000000..3601052
--- /dev/null
+++ b/macos/crates/edge0-cli/src/main.rs
@@ -0,0 +1,961 @@
+//! Thin `edge0` CLI: `serve` execs edge0d; other commands are HTTP clients or local registry ops.
+
+use std::process::Command;
+
+use clap::{Parser, Subcommand};
+
+use edge0_core::config::{Config, Overrides};
+use edge0_core::paths::Home;
+
+#[derive(Parser)]
+#[command(
+    name = "edge0",
+    version,
+    about = "edge0 local runtime CLI (thin client; the service is edge0d)"
+)]
+struct Cli {
+    /// Connection target (EDGE0_HOST / config.toml / http://127.0.0.1:8000)
+    #[arg(long, global = true)]
+    host: Option<String>,
+    #[arg(long, global = true)]
+    home: Option<String>,
+    #[command(subcommand)]
+    cmd: Cmd,
+}
+
+#[derive(Subcommand)]
+enum Cmd {
+    /// Start the daemon in the foreground (the only command that launches the service)
+    Serve {
+        #[arg(long, default_value = "foreground")]
+        scope: String,
+        #[arg(long)]
+        port: Option<u16>,
+        #[arg(long)]
+        bind: Option<String>,
+        #[arg(long)]
+        origins: Option<String>,
+        #[arg(long)]
+        concurrency: Option<u32>,
+        #[arg(long)]
+        keep_alive: Option<String>,
+        /// Print the token once (for copying to a phone; only set on non-loopback)
+        #[arg(long)]
+        print_token: bool,
+        /// Idempotent headless LaunchAgent install/uninstall
+        #[arg(long)]
+        enable_agent: bool,
+        #[arg(long)]
+        disable_agent: bool,
+        /// Prefill experts on demand (default); explicit production restores whole-layer/hot-window
+        #[arg(long)]
+        prefill_ondemand: bool,
+    },
+    /// Who is serving: scope/pid/version/uptime, resident models, disk
+    Status,
+    /// Stop the service (menu-agent → app Quit; LaunchAgent → --disable-agent)
+    Stop,
+    /// Uninstall runtime bits; keeps models, conversations, and logs unless --purge
+    Uninstall {
+        #[arg(long)]
+        purge: bool,
+    },
+    /// Print today's daemon log path and tail
+    Logs {
+        #[arg(long, default_value = "40")]
+        tail: usize,
+    },
+    /// Environment self-check (from daemon `/v1/edge0/doctor`)
+    Doctor,
+    /// Print the access token (also prints an existing file on loopback)
+    Token,
+    /// Load a tier
+    Load {
+        tier: String,
+        #[arg(long)]
+        keep_alive: Option<String>,
+    },
+    /// Unload a tier
+    Unload { tier: String },
+    /// Download a tier (HTTP if the daemon is up; otherwise local lock + same runner)
+    Pull {
+        tier: String,
+        /// Pin source (modelscope|huggingface)
+        #[arg(long)]
+        source: Option<String>,
+    },
+    /// List installed tiers
+    List,
+    /// Remove a tier (decrements blob refs; `edge0 gc` reclaims)
+    Rm { tier: String },
+    /// Reclaim unreferenced blobs (skips in-flight tasks)
+    Gc,
+}
+
+fn main() {
+    let cli = Cli::parse();
+    match run(cli) {
+        Ok(()) => {}
+        Err(msg) => {
+            eprintln!("error: {msg}");
+            std::process::exit(1);
+        }
+    }
+}
+
+fn home_path(cli_home: &Option<String>) -> Home {
+    match cli_home {
+        Some(h) => Home::new(h),
+        None => Home::from_env().unwrap_or_else(|_| Home::new(dirs_home().join(".edge0"))),
+    }
+}
+
+fn dirs_home() -> std::path::PathBuf {
+    std::env::var_os("HOME")
+        .map(std::path::PathBuf::from)
+        .unwrap_or_else(|| ".".into())
+}
+
+fn config(
+    cli: &Cli,
+    port: Option<u16>,
+    bind: Option<String>,
+    origins: Option<String>,
+    concurrency: Option<u32>,
+    keep_alive: Option<String>,
+) -> Result<Config, String> {
+    let ov = Overrides {
+        host: cli.host.clone(),
+        port,
+        bind,
+        origins,
+        source: None,
+        concurrency,
+        keep_alive,
+        home: cli.home.clone(),
+    };
+    let home = home_path(&cli.home);
+    Config::load(&ov, Some(&home.config_toml()))
+}
+
+fn client_for(url: &str) -> reqwest::blocking::Client {
+    let mut b = reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(15));
+    if url.contains("://127.") || url.contains("://localhost") {
+        b = b.no_proxy();
+    }
+    b.build().expect("client")
+}
+
+fn no_service_msg(host: &str) -> String {
+    format!(
+        "No service at {host}. Open edge0.app or run `edge0 serve`; for a remote host, enable LAN and pass a token (`edge0 token`)."
+    )
+}
+
+fn token_of(home: &Home) -> Option<String> {
+    std::fs::read_to_string(home.token())
+        .ok()
+        .map(|s| s.trim().to_string())
+        .filter(|s| !s.is_empty())
+}
+
+fn api(
+    cfg: &Config,
+    method: &str,
+    path: &str,
+    body: Option<serde_json::Value>,
+    home: &Home,
+) -> Result<(u16, serde_json::Value), String> {
+    let url = format!("{}{path}", cfg.host_url);
+    let c = client_for(&url);
+    let mut req = match method {
+        "GET" => c.get(&url),
+        "DELETE" => c.delete(&url),
+        _ => c.post(&url),
+    };
+    if let Some(t) = token_of(home) {
+        req = req.bearer_auth(t);
+    }
+    req = req.header(
+        "x-edge0-client-version",
+        edge0_core::version::daemon_version(),
+    );
+    if let Some(b) = body {
+        req = req.json(&b);
+    }
+    let resp = req.send().map_err(|_| no_service_msg(&cfg.host_url))?;
+    let status = resp.status().as_u16();
+    let v: serde_json::Value = resp.json().unwrap_or_else(|_| serde_json::Value::Null);
+    Ok((status, v))
+}
+
+fn run(cli: Cli) -> Result<(), String> {
+    let home = home_path(&cli.home);
+    match &cli.cmd {
+        Cmd::Serve {
+            scope,
+            port,
+            bind,
+            origins,
+            concurrency,
+            keep_alive,
+            print_token,
+            enable_agent,
+            disable_agent,
+            prefill_ondemand,
+        } => {
+            let cfg = config(
+                &cli,
+                *port,
+                bind.clone(),
+                origins.clone(),
+                *concurrency,
+                keep_alive.clone(),
+            )?;
+            if cfg.port == 0 {
+                return Err("invalid port".into());
+            }
+            let exe_dir = std::env::current_exe()
+                .map_err(|e| e.to_string())?
+                .parent()
+                .unwrap()
+                .to_path_buf();
+            let d = exe_dir.join("edge0d");
+            if !d.exists() {
+                return Err(format!("edge0d not found (should ship beside edge0): {d:?}"));
+            }
+            let mut cmd = Command::new(&d);
+            cmd.arg("--scope")
+                .arg(scope)
+                .arg("--port")
+                .arg(cfg.port.to_string())
+                .arg("--bind")
+                .arg(cfg.bind.to_string())
+                .arg("--home")
+                .arg(home.root.display().to_string())
+                .arg("--started-by")
+                .arg("cli");
+            if *print_token {
+                cmd.arg("--print-token");
+            }
+            if *enable_agent {
+                cmd.arg("--enable-agent");
+            }
+            if *disable_agent {
+                cmd.arg("--disable-agent");
+            }
+            if *prefill_ondemand {
+                cmd.env("EDGE0_PREFILL", "ondemand");
+            }
+            if let Some(o) = origins {
+                cmd.arg("--origins").arg(o);
+            }
+            if let Some(c) = concurrency {
+                cmd.arg("--concurrency").arg(c.to_string());
+            }
+            if let Some(k) = keep_alive {
+                cmd.arg("--keep-alive").arg(k);
+            }
+            if let Some(t) = &cli.host {
+                cmd.env("EDGE0_HOST", t);
+            }
+            use std::os::unix::process::CommandExt;
+            let err = cmd.exec();
+            Err(format!("exec edge0d failed: {err}"))
+        }
+        Cmd::Status => {
+            let cfg = config(&cli, None, None, None, None, None)?;
+            let (_, h) = api(&cfg, "GET", "/health", None, &home)?;
+            if h["session_id"].is_null() {
+                return Err("unexpected service response (missing session_id)".into());
+            }
+            let (_, s) = api(&cfg, "GET", "/v1/edge0/system", None, &home)?;
+            println!(
+                "service: v{} scope={} started_by={} uptime={}s session={} port {} ({}) throttled={}",
+                h["daemon_version"].as_str().unwrap_or("?"),
+                h["scope"].as_str().unwrap_or("?"),
+                h["started_by"].as_str().unwrap_or("?"),
+                h["uptime_s"].as_u64().unwrap_or(0),
+                h["session_id"].as_str().unwrap_or("?"),
+                s["daemon"]["port"].as_u64().unwrap_or(0),
+                s["daemon"]["bound"].as_str().unwrap_or("?"),
+                h["throttled"].as_bool().unwrap_or(false),
+            );
+            for m in s["models"].as_array().unwrap_or(&vec![]) {
+                println!(
+                    "  model {}@{} state={} unload_at={} in-flight={}",
+                    m["id"].as_str().unwrap_or("?"),
+                    m["rev"].as_str().unwrap_or("?"),
+                    m["state"].as_str().unwrap_or("?"),
+                    m["unload_at"].as_str().unwrap_or("never"),
+                    m["active_requests"].as_u64().unwrap_or(0),
+                );
+            }
+            let df = s["hardware"]["disk_free_gib"].as_f64();
+            if let Some(d) = df {
+                println!("  disk free {d:.1} GiB");
+            }
+            Ok(())
+        }
+        Cmd::Stop => {
+            let cfg = config(&cli, None, None, None, None, None)?;
+            let (_, h) = api(&cfg, "GET", "/health", None, &home)?;
+            match h["scope"].as_str() {
+                Some("menu-agent") => Err("Under menu-agent, stop from the app menu Quit (cascade-stops the service and workers); the CLI does not do this, to avoid a second entry point.".into()),
+                Some("launchagent") => Err("To unload headless persistence: `edge0 serve --disable-agent`".into()),
+                Some("foreground") => {
+                    if let Some(pid) = h["pid"].as_u64() {
+                        let _ = Command::new("/bin/kill").arg("-TERM").arg(pid.to_string()).status();
+                        println!("sent SIGTERM to pid {pid}");
+                        Ok(())
+                    } else {
+                        Err("health did not report pid; cannot stop".into())
+                    }
+                }
+                _ => Err(no_service_msg(&cfg.host_url)),
+            }
+        }
+        Cmd::Uninstall { purge } => cmd_uninstall(&cli, &home, *purge),
+        Cmd::Logs { tail } => {
+            let name = format!("edge0d.{}.log", chrono::Local::now().format("%Y-%m-%d"));
+            let p = home.logs().join(&name);
+            if !p.exists() {
+                return Err(format!(
+                    "no log for today {} (logs are named by date; see the logs/ directory)",
+                    p.display()
+                ));
+            }
+            println!("# {}", p.display());
+            let text = std::fs::read_to_string(&p).map_err(|e| e.to_string())?;
+            let lines: Vec<&str> = text.lines().collect();
+            for l in lines[lines.len().saturating_sub(*tail)..].iter() {
+                println!("{l}");
+            }
+            Ok(())
+        }
+        Cmd::Doctor => {
+            let cfg = config(&cli, None, None, None, None, None)?;
+            let (st, v) = api(&cfg, "GET", "/v1/edge0/doctor", None, &home)?;
+            if st != 200 {
+                return Err(format!(
+                    "doctor call failed (HTTP {st}): {}",
+                    v["error"]["message"].as_str().unwrap_or("?")
+                ));
+            }
+            println!("overall: {}", v["overall"].as_str().unwrap_or("?"));
+            for c in v["checks"].as_array().unwrap_or(&vec![]) {
+                let code = c["code"]
+                    .as_str()
+                    .map(|s| format!(" [{s}]"))
+                    .unwrap_or_default();
+                let next = c["next"]
+                    .as_str()
+                    .map(|s| format!(" → {s}"))
+                    .unwrap_or_default();
+                println!(
+                    "  {:<9} {:<6} {}{}{}",
+                    c["id"].as_str().unwrap_or("?"),
+                    c["verdict"].as_str().unwrap_or("?"),
+                    c["detail"].as_str().unwrap_or(""),
+                    code,
+                    next
+                );
+            }
+            Ok(())
+        }
+        Cmd::Token => match token_of(&home) {
+            Some(t) => {
+                println!("{t}");
+                Ok(())
+            }
+            None => Err("no token yet (only a non-loopback bind mints one; loopback does not need it)".into()),
+        },
+        Cmd::Load { tier, keep_alive } => {
+            let cfg = config(&cli, None, None, None, None, None)?;
+            let body = keep_alive
+                .as_ref()
+                .map(|k| serde_json::json!({"keep_alive": k}))
+                .unwrap_or_else(|| serde_json::json!({}));
+            let (st, v) = api(
+                &cfg,
+                "POST",
+                &format!("/v1/edge0/models/{tier}/load"),
+                Some(body),
+                &home,
+            )?;
+            match st {
+                202 => {
+                    println!("load started ({tier}); watch `edge0 status` or the event stream");
+                    Ok(())
+                }
+                _ => Err(format!(
+                    "HTTP {st}: {}",
+                    v["error"]["message"].as_str().unwrap_or("?")
+                )),
+            }
+        }
+        Cmd::Unload { tier } => {
+            let cfg = config(&cli, None, None, None, None, None)?;
+            let (st, v) = api(
+                &cfg,
+                "POST",
+                &format!("/v1/edge0/models/{tier}/unload"),
+                None,
+                &home,
+            )?;
+            if st == 200 {
+                println!("unload requested for {tier}");
+                Ok(())
+            } else {
+                Err(format!(
+                    "HTTP {st}: {}",
+                    v["error"]["message"].as_str().unwrap_or("?")
+                ))
+            }
+        }
+        Cmd::Pull { tier, source } => {
+            let cfg = config(&cli, None, None, None, None, None)?;
+            cmd_pull(&cfg, &home, tier, source.as_deref())
+        }
+        Cmd::List => cmd_list(&home),
+        Cmd::Rm { tier } => {
+            let cfg = config(&cli, None, None, None, None, None)?;
+            cmd_rm(&cfg, &home, tier)
+        }
+        Cmd::Gc => cmd_gc(&home),
+    }
+}
+
+fn cmd_uninstall(cli: &Cli, home: &Home, purge: bool) -> Result<(), String> {
+    let mut failures = 0usize;
+    let cfg = config(cli, None, None, None, None, None).ok();
+    let daemon = cfg
+        .as_ref()
+        .and_then(|c| api(c, "GET", "/health", None, home).ok());
+
+    if let Some((_, health)) = daemon.as_ref() {
+        if let Some(pid) = health["pid"].as_u64() {
+            println!("OK  stop daemon pid={pid}");
+            let status = Command::new("/bin/kill")
+                .args(["-TERM", &pid.to_string()])
+                .status()
+                .map_err(|e| e.to_string())?;
+            if !status.success() {
+                failures += 1;
+                println!("FAIL stop daemon pid={pid}");
+            }
+        } else {
+            println!("SKIP stop daemon (health did not report pid)");
+        }
+    } else {
+        println!("SKIP stop daemon (service is not running)");
+    }
+
+    let daemon_bin = std::env::var_os("EDGE0_DAEMON_BIN")
+        .map(std::path::PathBuf::from)
+        .or_else(|| {
+            std::env::current_exe()
+                .ok()
+                .and_then(|p| p.parent().map(|d| d.join("edge0d")))
+        });
+    let scope = daemon
+        .as_ref()
+        .and_then(|(_, h)| h["scope"].as_str())
+        .unwrap_or("");
+    if scope == "launchagent" {
+        match (&daemon_bin, daemon_bin.as_ref()) {
+            (Some(bin), Some(_)) if bin.exists() => {
+                let status = Command::new(bin)
+                    .args([
+                        "--disable-agent",
+                        "--home",
+                        &home.root.display().to_string(),
+                    ])
+                    .status()
+                    .map_err(|e| e.to_string())?;
+                if status.success() {
+                    println!("OK  disable LaunchAgent");
+                } else {
+                    failures += 1;
+                    println!("FAIL disable LaunchAgent");
+                }
+            }
+            _ => {
+                failures += 1;
+                println!("FAIL disable LaunchAgent (edge0d not found)");
+            }
+        }
+    } else {
+        println!(
+            "SKIP disable LaunchAgent (current scope={})",
+            if scope.is_empty() { "none" } else { scope }
+        );
+    }
+
+    let link = std::env::var_os("EDGE0_BIN_LINK")
+        .map(std::path::PathBuf::from)
+        .unwrap_or_else(|| "/usr/local/bin/edge0".into());
+    if link.is_symlink() {
+        match std::fs::remove_file(&link) {
+            Ok(()) => println!("OK  remove PATH symlink {}", link.display()),
+            Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
+                failures += 1;
+                println!(
+                    "FAIL remove PATH symlink {}; run: sudo unlink {}",
+                    link.display(),
+                    link.display()
+                );
+            }
+            Err(e) => {
+                failures += 1;
+                println!("FAIL remove PATH symlink {}: {e}", link.display());
+            }
+        }
+    } else {
+        println!("SKIP remove PATH symlink {} (not present)", link.display());
+    }
+
+    let bytes = directory_bytes(&home.root);
+    println!(
+        "KEEP user data {} ({})",
+        home.root.display(),
+        human_bytes(bytes)
+    );
+    if purge {
+        println!("PURGE requested; this removes models, conversations, logs, config, and state.");
+        println!("Type PURGE to continue:");
+        let mut answer = String::new();
+        std::io::stdin()
+            .read_line(&mut answer)
+            .map_err(|e| e.to_string())?;
+        if answer.trim() != "PURGE" {
+            println!("SKIP purge (confirmation did not match)");
+        } else if home.root.exists() {
+            std::fs::remove_dir_all(&home.root).map_err(|e| e.to_string())?;
+            println!("OK  purge {}", home.root.display());
+        } else {
+            println!("SKIP purge (data directory not present)");
+        }
+    } else {
+        println!("SKIP purge (use --purge to remove user data)");
+    }
+
+    if failures == 0 {
+        Ok(())
+    } else {
+        Err(format!("uninstall completed with {failures} failure(s)"))
+    }
+}
+
+fn directory_bytes(path: &std::path::Path) -> u64 {
+    if path.is_file() {
+        return std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
+    }
+    std::fs::read_dir(path)
+        .ok()
+        .into_iter()
+        .flatten()
+        .filter_map(Result::ok)
+        .map(|entry| directory_bytes(&entry.path()))
+        .sum()
+}
+
+fn human_bytes(bytes: u64) -> String {
+    const UNITS: [&str; 4] = ["B", "KiB", "MiB", "GiB"];
+    let mut value = bytes as f64;
+    let mut idx = 0usize;
+    while value >= 1024.0 && idx < UNITS.len() - 1 {
+        value /= 1024.0;
+        idx += 1;
+    }
+    if idx == 0 {
+        format!("{bytes} {}", UNITS[idx])
+    } else {
+        format!("{value:.1} {}", UNITS[idx])
+    }
+}
+
+fn host_is_loopback(url: &str) -> bool {
+    let rest = url
+        .strip_prefix("http://")
+        .or_else(|| url.strip_prefix("https://"))
+        .unwrap_or(url);
+    let host = rest.rsplit_once(':').map_or(rest, |(h, _)| h);
+    host == "localhost" || host.starts_with("127.")
+}
+
+fn daemon_alive(cfg: &Config) -> bool {
+    let url = format!("{}/health", cfg.host_url);
+    let mut b =
+        reqwest::blocking::Client::builder().timeout(std::time::Duration::from_millis(1200));
+    if host_is_loopback(&cfg.host_url) {
+        b = b.no_proxy();
+    }
+    match b.build() {
+        Ok(c) => c
+            .get(&url)
+            .send()
+            .map(|r| r.status().is_success())
+            .unwrap_or(false),
+        Err(_) => false,
+    }
+}
+
+fn cmd_pull(cfg: &Config, home: &Home, tier: &str, source: Option<&str>) -> Result<(), String> {
+    if daemon_alive(cfg) {
+        return pull_via_daemon(cfg, home, tier, source);
+    }
+    if !host_is_loopback(&cfg.host_url) {
+        return Err(format!(
+            "{} has no daemon: remote pull is HTTP-only (no local self-run). On the target machine, open edge0.app or run `edge0 serve`.",
+            cfg.host_url
+        ));
+    }
+    pull_self_run(cfg, home, tier, source)
+}
+
+fn pull_via_daemon(
+    cfg: &Config,
+    home: &Home,
+    tier: &str,
+    source: Option<&str>,
+) -> Result<(), String> {
+    let body = serde_json::json!({ "tier": tier, "source": source });
+    let (st, v) = api(cfg, "POST", "/v1/edge0/downloads", Some(body), home)?;
+    if st != 201 {
+        return Err(format!(
+            "HTTP {st}: {}",
+            v["error"]["message"].as_str().unwrap_or("?")
+        ));
+    }
+    let id = v["id"].as_str().ok_or("daemon response missing id")?.to_string();
+    let total = v["bytes_total"].as_u64().unwrap_or(0);
+    println!(
+        "task {id}: {}@{} {total} bytes, daemon is running it (GUI/CLI share progress; interrupt stops at the last verified range)",
+        v["tier"].as_str().unwrap_or(tier),
+        v["rev"].as_str().unwrap_or("?"),
+    );
+    match sse_terminal(cfg, home, &id, total) {
+        Some(r) => r,
+        None => poll_watch(cfg, home, &id),
+    }
+}
+
+fn sse_terminal(cfg: &Config, home: &Home, id: &str, total: u64) -> Option<Result<(), String>> {
+    use std::io::Read as _;
+    let url = format!("{}/v1/edge0/events", cfg.host_url);
+    let mut b = reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(3600));
+    if host_is_loopback(&cfg.host_url) {
+        b = b.no_proxy();
+    }
+    let mut req = b.build().ok()?.get(&url);
+    if let Some(t) = token_of(home) {
+        req = req.bearer_auth(t);
+    }
+    let mut resp = req.send().ok()?;
+    let mut buf = String::new();
+    let mut chunk = [0u8; 4096];
+    let mut last_pct: i64 = -1;
+    loop {
+        match resp.read(&mut chunk) {
+            Ok(0) | Err(_) => return None,
+            Ok(n) => {
+                buf.push_str(&String::from_utf8_lossy(&chunk[..n]));
+                while let Some(pos) = buf.find("\n\n") {
+                    let block: String = buf.drain(..pos + 2).collect();
+                    for line in block.lines() {
+                        let Some(data) = line
+                            .strip_prefix("data: ")
+                            .or_else(|| line.strip_prefix("data:"))
+                        else {
+                            continue;
+                        };
+                        let Ok(e) = serde_json::from_str::<serde_json::Value>(data) else {
+                            continue;
+                        };
+                        if e["subject"].as_str() != Some(id) {
+                            continue;
+                        }
+                        let p = &e["payload"];
+                        match e["type"].as_str().unwrap_or("") {
+                            "download.progress" => {
+                                let d = p["bytes_done"].as_u64().unwrap_or(0);
+                                let pct = if total > 0 {
+                                    (d * 100 / total) as i64
+                                } else {
+                                    -1
+                                };
+                                if pct != last_pct {
+                                    last_pct = pct;
+                                    use std::io::Write as _;
+                                    if pct >= 0 {
+                                        print!("\r  {pct:>3}%  {d}/{total}");
+                                    } else {
+                                        print!(
+                                            "\r  {d} bytes  phase={}",
+                                            p["phase"].as_str().unwrap_or("?")
+                                        );
+                                    }
+                                    let _ = std::io::stdout().flush();
+                                }
+                            }
+                            "download.completed" => {
+                                println!("\ncomplete: task {id} is registered in the local model directory");
+                                return Some(Ok(()));
+                            }
+                            "download.failed" => {
+                                println!();
+                                let code = p["code"].as_str().unwrap_or("?");
+                                return Some(Err(format!(
+                                    "[{code}] download stopped at {}/{}; the task is kept and verified ranges are intact (rerun the same command to resume, or `edge0 pull --source <other>` to switch source)",
+                                    p["bytes_done"].as_u64().unwrap_or(0),
+                                    p["bytes_total"].as_u64().unwrap_or(total)
+                                )));
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+fn poll_watch(cfg: &Config, home: &Home, id: &str) -> Result<(), String> {
+    let t0 = std::time::Instant::now();
+    loop {
+        let (_, v) = api(cfg, "GET", "/v1/edge0/downloads", None, home)?;
+        let list = v["data"].as_array().cloned().unwrap_or_default();
+        match list.iter().find(|t| t["id"].as_str() == Some(id)) {
+            None => return Ok(()),
+            Some(t) if t["failed"].is_object() => {
+                return Err(format!(
+                    "[{}] {}",
+                    t["failed"]["code"].as_str().unwrap_or("?"),
+                    t["failed"]["message"].as_str().unwrap_or("download failed")
+                ))
+            }
+            Some(t) => {
+                println!(
+                    "  phase={} {}/{}",
+                    t["phase"].as_str().unwrap_or("?"),
+                    t["bytes_done"].as_u64().unwrap_or(0),
+                    t["bytes_total"].as_u64().unwrap_or(0)
+                );
+            }
+        }
+        if t0.elapsed() > std::time::Duration::from_secs(3600) {
+            return Err("waited more than 1h; giving up tracking (the task is still running)".into());
+        }
+        std::thread::sleep(std::time::Duration::from_secs(1));
+    }
+}
+
+fn pull_self_run(
+    cfg: &Config,
+    home: &Home,
+    tier: &str,
+    source: Option<&str>,
+) -> Result<(), String> {
+    home.ensure_dirs()?;
+    let lock = edge0_core::state::RegistryLock::acquire(&home.registry_lock())?;
+    println!("local run (no daemon): holding the registry lock, same state machine as the daemon.");
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .map_err(|e| e.to_string())?;
+    let client = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(30))
+        .build()
+        .map_err(|e| e.to_string())?;
+    let plan = rt
+        .block_on(edge0_pull::catalog::plan_for_tier(
+            cfg.pull_mode,
+            &client,
+            &cfg.registry_urls,
+            &cfg.hf_base,
+            tier,
+            home,
+            chrono::Utc::now(),
+            edge0_core::version::daemon_version(),
+        ))
+        .map_err(|e| format!("{e}"))?;
+    let total = plan.bytes_total();
+    if matches!(plan, edge0_pull::catalog::CreationPlan::Registry(..)) {
+        std::fs::create_dir_all(home.blobs()).map_err(|e| e.to_string())?;
+    }
+    println!(
+        "{tier}@{} {total} bytes (~{:.1} GiB, mode={}): disk preflight…",
+        plan.rev(),
+        total as f64 / (1u64 << 30) as f64,
+        cfg.pull_mode.as_str()
+    );
+    edge0_pull::registry_ops::preflight_check(home, total).map_err(|e| format!("{e}"))?;
+    let now = chrono::Utc::now();
+    let (snapshot, task_plan, source_reason) = match plan {
+        edge0_pull::catalog::CreationPlan::Registry(list, _) => (Some(list.manifest), None, None),
+        edge0_pull::catalog::CreationPlan::Direct(snap) => (
+            None,
+            Some(snap),
+            Some(edge0_pull::catalog::CreationPlan::source_reason_direct(
+                &cfg.hf_base,
+                source.is_some() || cfg.source.is_some(),
+            )),
+        ),
+    };
+    let task = edge0_core::state::DownloadTask {
+        id: format!("dl-{}", now.timestamp_micros()),
+        tier: tier.into(),
+        rev: snapshot
+            .as_ref()
+            .and_then(|m| m.find_tier(tier).map(|t| t.rev.clone()))
+            .or_else(|| task_plan.as_ref().map(|p| p.commit.clone()))
+            .unwrap_or_default(),
+        phase: "probe".into(),
+        source: None,
+        bytes_total: total,
+        bytes_done: 0,
+        paused: false,
+        files: vec![],
+        created_at: now,
+        updated_at: now,
+        snapshot,
+        source_reason,
+        failed: None,
+        rate_bps: None,
+        eta_s: None,
+        source_switched: false,
+        plan: task_plan,
+    };
+    let runner = edge0_pull::runner::TaskRunner {
+        home: home.clone(),
+        fetcher: std::sync::Arc::new(edge0_pull::fetch::HttpFetcher::with_env_proxy()),
+        store: std::sync::Arc::new(edge0_pull::runner::FileStore { home: home.clone() }),
+        events: std::sync::Arc::new(CliPrint),
+        source_pin: source.or(cfg.source.as_deref()).map(str::to_string),
+        parallel: 4,
+    };
+    let done = rt
+        .block_on(runner.run(task))
+        .map_err(|e| format!("{e} (task progress kept; rerun the same command to resume)"))?;
+    println!(
+        "\ninstalled: models/{tier}/{} (`edge0 list` to inspect, `edge0 load {tier}` to load)",
+        done.rev
+    );
+    drop(lock);
+    Ok(())
+}
+
+struct CliPrint;
+impl edge0_pull::runner::DownloadEvents for CliPrint {
+    fn event(
+        &self,
+        etype: &'static str,
+        task: &edge0_core::state::DownloadTask,
+        code: Option<&'static str>,
+    ) {
+        use std::io::Write as _;
+        match etype {
+            "download.progress" => print!(
+                "\r  {}/{} bytes  phase={}",
+                task.bytes_done, task.bytes_total, task.phase
+            ),
+            "download.completed" => {}
+            "download.failed" => {
+                println!();
+                if let Some(c) = code {
+                    eprintln!("  stopped at [{c}] {}", task.bytes_done);
+                }
+            }
+            _ => {}
+        }
+        let _ = std::io::stdout().flush();
+    }
+}
+
+fn cmd_list(home: &Home) -> Result<(), String> {
+    let tiers = edge0_core::manifest::scan_installed(home);
+    if tiers.is_empty() {
+        println!(
+            "No tiers installed. Start with `edge0 pull edge0-8b` (4.6 GB) or `edge0 pull edge0-35b` (19.7 GB)"
+        );
+        return Ok(());
+    }
+    println!("{:<11} {:<22} {:>16}  {:<8}", "tier", "rev", "bytes", "mode");
+    for t in tiers {
+        println!(
+            "{:<11} {:<22} {:>16}  {:<8}",
+            t.tier, t.rev, t.bytes_total, t.mode
+        );
+    }
+    Ok(())
+}
+
+fn cmd_rm(cfg: &Config, home: &Home, tier: &str) -> Result<(), String> {
+    if daemon_alive(cfg) {
+        let (st, v) = api(
+            cfg,
+            "DELETE",
+            &format!("/v1/edge0/models/{tier}"),
+            None,
+            home,
+        )?;
+        return match st {
+            200 => {
+                println!("removed {tier} (direct installs delete the directory immediately; otherwise run `edge0 gc`)");
+                Ok(())
+            }
+            _ => Err(format!(
+                "HTTP {st}: {}",
+                v["error"]["message"].as_str().unwrap_or("?")
+            )),
+        };
+    }
+    if !host_is_loopback(&cfg.host_url) {
+        return Err(format!(
+            "No daemon on {}: remote deletes must go over HTTP (the peer owns the registry)",
+            cfg.host_url
+        ));
+    }
+    home.ensure_dirs()?;
+    let _lock = edge0_core::state::RegistryLock::acquire(&home.registry_lock())?;
+    let found = edge0_core::manifest::scan_installed(home)
+        .into_iter()
+        .find(|t| t.tier == tier);
+    if found.as_ref().is_some_and(|t| t.bundled) {
+        return Err(format!("{tier} is bundled with the app and cannot be deleted"));
+    }
+    let mode = found.map(|t| t.mode).unwrap_or_default();
+    let shas = edge0_pull::registry_ops::remove_tier(home, tier).map_err(|e| format!("{e}"))?;
+    if mode == "direct" {
+        println!("removed {tier} (direct single-rev install: directory deleted)");
+    } else {
+        println!(
+            "removed {tier} ({} blob refs decremented; run `edge0 gc` to free disk)",
+            shas.len()
+        );
+    }
+    Ok(())
+}
+
+fn cmd_gc(home: &Home) -> Result<(), String> {
+    home.ensure_dirs()?;
+    let _lock = edge0_core::state::RegistryLock::acquire(&home.registry_lock()).map_err(|e| {
+        format!("cannot reclaim: {e} — registry changes (including gc) require the lock holder; gc via CLI is unavailable while the daemon is running")
+    })?;
+    let orphans = edge0_pull::registry_ops::gc_orphans(home);
+    let rep = edge0_pull::registry_ops::gc(home).map_err(|e| format!("{e}"))?;
+    if !orphans.is_empty() {
+        println!("cleared {} orphan staging paths: {}", orphans.len(), orphans.join(", "));
+    }
+    println!(
+        "gc: reclaimed {} unreferenced blobs; kept {} still referenced, {} in-flight",
+        rep.deleted.len(),
+        rep.kept_referenced,
+        rep.kept_inflight.len()
+    );
+    if rep.deleted.is_empty() && rep.kept_referenced == 0 && rep.kept_inflight.is_empty() {
+        println!("(direct installs have no blob refs to reclaim; orphan staging is marked on rescan and cleaned by rm/reinstall)");
+    }
+    Ok(())
+}
diff --git a/macos/crates/edge0-core/Cargo.toml b/macos/crates/edge0-core/Cargo.toml
new file mode 100644
index 0000000..a123117
--- /dev/null
+++ b/macos/crates/edge0-core/Cargo.toml
@@ -0,0 +1,22 @@
+[package]
+name = "edge0-core"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+[dependencies]
+serde.workspace = true
+serde_json.workspace = true
+chrono.workspace = true
+uuid.workspace = true
+rand.workspace = true
+thiserror.workspace = true
+libc.workspace = true
+libloading.workspace = true
+ts-rs.workspace = true
+
+[[bin]]
+# TypeScript binding generator: cargo run -p edge0-core --bin gen-ts -- app/src/gen
+name = "gen-ts"
+path = "src/bin/gen_ts.rs"
diff --git a/macos/crates/edge0-core/src/abi_load.rs b/macos/crates/edge0-core/src/abi_load.rs
new file mode 100644
index 0000000..53056ab
--- /dev/null
+++ b/macos/crates/edge0-core/src/abi_load.rs
@@ -0,0 +1,189 @@
+//! Shared `dlopen` + symbol resolve + ABI handshake for daemon and worker.
+
+use std::path::Path;
+
+use libloading::{Library, Symbol};
+
+use crate::version::E0_ABI_VERSION;
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct E0TokenOutC {
+    pub size: u32,
+    pub token: i32,
+    pub flags: u32,
+}
+
+pub type FnAbiVersion = extern "C" fn() -> u32;
+pub type FnEngineCreate = unsafe extern "C" fn(*const i8, *const i8, *const i8) -> *mut c_void;
+pub type FnGenerateBegin =
+    unsafe extern "C" fn(*mut c_void, *const i32, u32, *const i8) -> *mut c_void;
+pub type FnNext = unsafe extern "C" fn(*mut c_void, *mut c_void, *mut E0TokenOutC) -> i32;
+pub type FnCancel = unsafe extern "C" fn(*mut c_void, *mut c_void) -> i32;
+pub type FnStats = unsafe extern "C" fn(*mut c_void, *mut i8, usize, *mut usize) -> i32;
+pub type FnTurnReset = unsafe extern "C" fn(*mut c_void) -> i32;
+pub type FnLastError = unsafe extern "C" fn(*mut c_void, *mut i8, usize, *mut usize) -> i32;
+pub type FnEngineDestroy = unsafe extern "C" fn(*mut c_void);
+pub type FnTokCreate = unsafe extern "C" fn(*const i8) -> *mut c_void;
+pub type FnTokCall =
+    unsafe extern "C" fn(*mut c_void, *const i8, *mut i8, usize, *mut usize) -> i32;
+pub type FnTokDestroy = unsafe extern "C" fn(*mut c_void);
+
+pub use std::ffi::c_void;
+
+pub const E0_OK: i32 = 0;
+pub const E0_NOTIFIED: i32 = 1;
+pub const E0_WOULD_BLOCK: i32 = 2;
+
+/// Report what was actually loaded; unknown basenames are passed through, not classified.
+pub fn engine_backend(path: &Path) -> String {
+    let base = path
+        .file_name()
+        .map(|s| s.to_string_lossy().into_owned())
+        .unwrap_or_else(|| path.display().to_string());
+    if base.contains("native") {
+        "native".into()
+    } else if base.contains("replay") {
+        "replay".into()
+    } else {
+        base
+    }
+}
+
+/// Version mismatch refuses to load; missing symbols (including `e0_tok_decode`) fail load.
+pub struct EngineLib {
+    _lib: Library,
+    pub path: std::path::PathBuf,
+    pub abi_version: FnAbiVersion,
+    pub engine_create: FnEngineCreate,
+    pub generate_begin: FnGenerateBegin,
+    pub next: FnNext,
+    pub cancel: FnCancel,
+    pub stats: FnStats,
+    pub turn_reset: FnTurnReset,
+    pub last_error: FnLastError,
+    pub engine_destroy: FnEngineDestroy,
+    pub tok_create: FnTokCreate,
+    pub tok_tokenize: FnTokCall,
+    pub tok_apply_template: FnTokCall,
+    pub tok_decode: FnTokCall,
+    pub tok_destroy: FnTokDestroy,
+}
+
+// dyld is thread-safe; function pointers are read-only addresses.
+unsafe impl Send for EngineLib {}
+unsafe impl Sync for EngineLib {}
+
+macro_rules! sym {
+    ($lib:expr, $ty:ty, $name:literal) => {{
+        let s: Symbol<$ty> = (unsafe { $lib.get($name.as_bytes()) })
+            .map_err(|e| format!("symbol {} missing or wrong type: {e}", $name))?;
+        *s
+    }};
+}
+
+impl EngineLib {
+    pub fn backend(&self) -> String {
+        engine_backend(&self.path)
+    }
+
+    pub fn load(path: &Path) -> Result<Self, String> {
+        let lib = unsafe { Library::new(path) }
+            .map_err(|e| format!("dlopen {} failed: {e}", path.display()))?;
+        let abi_version: FnAbiVersion = sym!(lib, FnAbiVersion, "e0_abi_version");
+        let got = abi_version();
+        if got != E0_ABI_VERSION {
+            return Err(format!(
+                "ABI handshake failed: library reports {got}, daemon compiled with E0_ABI_VERSION={E0_ABI_VERSION} — refusing to load"
+            ));
+        }
+        let e = Self {
+            path: path.to_path_buf(),
+            abi_version,
+            engine_create: sym!(lib, FnEngineCreate, "e0_engine_create"),
+            generate_begin: sym!(lib, FnGenerateBegin, "e0_generate_begin"),
+            next: sym!(lib, FnNext, "e0_next"),
+            cancel: sym!(lib, FnCancel, "e0_cancel"),
+            stats: sym!(lib, FnStats, "e0_stats"),
+            turn_reset: sym!(lib, FnTurnReset, "e0_turn_reset"),
+            last_error: sym!(lib, FnLastError, "e0_last_error"),
+            engine_destroy: sym!(lib, FnEngineDestroy, "e0_engine_destroy"),
+            tok_create: sym!(lib, FnTokCreate, "e0_tok_create"),
+            tok_tokenize: sym!(lib, FnTokCall, "e0_tok_tokenize"),
+            tok_apply_template: sym!(lib, FnTokCall, "e0_tok_apply_template"),
+            tok_decode: sym!(lib, FnTokCall, "e0_tok_decode"),
+            tok_destroy: sym!(lib, FnTokDestroy, "e0_tok_destroy"),
+            _lib: lib,
+        };
+        Ok(e)
+    }
+
+    /// Two-probe JSON out-param wrapper (returns an owned String).
+    ///
+    /// # Safety
+    /// `f` must report the required length when `buf==NULL` and write exactly `*written` bytes
+    /// into a sufficient buffer, capturing only state valid for this call.
+    pub unsafe fn call_json(
+        &self,
+        f: impl Fn(*mut i8, usize, *mut usize) -> i32,
+    ) -> Result<String, i32> {
+        let mut written = 0usize;
+        let st = f(std::ptr::null_mut(), 0, &mut written);
+        if st != 0 || written == 0 {
+            return Err(st);
+        }
+        let mut buf = vec![0u8; written];
+        let st = f(buf.as_mut_ptr() as *mut i8, written, &mut written);
+        if st != 0 {
+            return Err(st);
+        }
+        String::from_utf8(buf).map_err(|_| -1)
+    }
+}
+
+/// Per-tier tokenizer handle on the daemon; independent of worker keep_alive.
+/// Pointers are only touched inside this wrapper, under a Mutex at the call site.
+pub struct TokHandle {
+    pub lib: std::sync::Arc<EngineLib>,
+    ptr: *mut c_void,
+}
+unsafe impl Send for TokHandle {}
+
+impl TokHandle {
+    pub fn open(lib: std::sync::Arc<EngineLib>, model_dir: &Path) -> Result<Self, String> {
+        let c =
+            std::ffi::CString::new(model_dir.display().to_string()).map_err(|e| e.to_string())?;
+        let ptr = unsafe { (lib.tok_create)(c.as_ptr()) };
+        if ptr.is_null() {
+            return Err(format!(
+                "e0_tok_create({}) returned NULL (tokenizer files missing or unreadable)",
+                model_dir.display()
+            ));
+        }
+        Ok(Self { lib, ptr })
+    }
+    fn call(&self, f: FnTokCall, input: &str) -> Result<String, i32> {
+        let c = match std::ffi::CString::new(input) {
+            Ok(c) => c,
+            Err(_) => return Err(-1),
+        };
+        let lib = &self.lib;
+        let p = self.ptr;
+        unsafe { lib.call_json(|b, l, w| f(p, c.as_ptr(), b, l, w)) }
+    }
+    pub fn tokenize(&self, text_json: &str) -> Result<String, i32> {
+        self.call(self.lib.tok_tokenize, text_json)
+    }
+    pub fn apply_template(&self, messages_json: &str) -> Result<String, i32> {
+        self.call(self.lib.tok_apply_template, messages_json)
+    }
+    pub fn decode(&self, tokens_json: &str) -> Result<String, i32> {
+        self.call(self.lib.tok_decode, tokens_json)
+    }
+}
+
+impl Drop for TokHandle {
+    fn drop(&mut self) {
+        unsafe { (self.lib.tok_destroy)(self.ptr) };
+    }
+}
diff --git a/macos/crates/edge0-core/src/bin/gen_ts.rs b/macos/crates/edge0-core/src/bin/gen_ts.rs
new file mode 100644
index 0000000..4217a30
--- /dev/null
+++ b/macos/crates/edge0-core/src/bin/gen_ts.rs
@@ -0,0 +1,136 @@
+//! TypeScript binding generator: `cargo run -p edge0-core --bin gen-ts -- <out-dir>`.
+
+use std::path::Path;
+
+use ts_rs::{Config, TS};
+
+fn cfg() -> Config {
+    // On-wire u64 values stay below 2^53, so JS number is safe; bigint would need conversion.
+    Config::new().with_large_int("number")
+}
+
+fn write_ts(dir: &Path, ty: &str, content: &str) {
+    std::fs::write(dir.join(format!("{ty}.ts")), content).unwrap_or_else(|e| {
+        panic!("failed to write {ty}.ts: {e}");
+    });
+}
+
+fn ts<T: TS + 'static>(name: &str) -> String {
+    T::export_to_string(&cfg()).unwrap_or_else(|e| panic!("failed to generate {name}: {e}"))
+}
+
+fn main() {
+    let out = std::env::args()
+        .nth(1)
+        .unwrap_or_else(|| panic!("usage: gen-ts <out-dir>"));
+    let dir = Path::new(&out);
+    std::fs::create_dir_all(dir).expect("out-dir");
+
+    use edge0_core::protocol as p;
+    let dep = dir.join("serde_json");
+    std::fs::create_dir_all(&dep).expect("serde_json dir");
+    std::fs::write(
+        dep.join("JsonValue.ts"),
+        ts::<serde_json::Value>("JsonValue"),
+    )
+    .unwrap();
+    write_ts(dir, "Event", &ts::<p::Event>("Event"));
+    write_ts(dir, "HelloPayload", &ts::<p::HelloPayload>("HelloPayload"));
+    write_ts(
+        dir,
+        "ResyncPayload",
+        &ts::<p::ResyncPayload>("ResyncPayload"),
+    );
+    write_ts(
+        dir,
+        "SystemSnapshot",
+        &ts::<p::SystemSnapshot>("SystemSnapshot"),
+    );
+    write_ts(dir, "DaemonInfo", &ts::<p::DaemonInfo>("DaemonInfo"));
+    write_ts(dir, "HardwareInfo", &ts::<p::HardwareInfo>("HardwareInfo"));
+    write_ts(dir, "ModelEntry", &ts::<p::ModelEntry>("ModelEntry"));
+    write_ts(
+        dir,
+        "DownloadEntry",
+        &ts::<p::DownloadEntry>("DownloadEntry"),
+    );
+    write_ts(
+        dir,
+        "DownloadFailure",
+        &ts::<p::DownloadFailure>("DownloadFailure"),
+    );
+    write_ts(dir, "CatalogTier", &ts::<p::CatalogTier>("CatalogTier"));
+    write_ts(
+        dir,
+        "CatalogResponse",
+        &ts::<p::CatalogResponse>("CatalogResponse"),
+    );
+
+    use edge0_core::conversation as c;
+    write_ts(dir, "Role", &ts::<c::Role>("Role"));
+    write_ts(
+        dir,
+        "MessageStatus",
+        &ts::<c::MessageStatus>("MessageStatus"),
+    );
+    write_ts(dir, "Usage", &ts::<c::Usage>("Usage"));
+    write_ts(dir, "Attachment", &ts::<c::Attachment>("Attachment"));
+    write_ts(
+        dir,
+        "MessageRecord",
+        &ts::<c::MessageRecord>("MessageRecord"),
+    );
+    write_ts(dir, "ThreadRecord", &ts::<c::ThreadRecord>("ThreadRecord"));
+
+    let union: Vec<String> = p::EVENT_TYPES.iter().map(|t| format!("\"{t}\"")).collect();
+    let folded: Vec<String> = p::EVENT_TYPES
+        .iter()
+        .filter(|t| p::foldable(t))
+        .map(|t| format!("\"{t}\""))
+        .collect();
+    std::fs::write(
+        dir.join("EventType.ts"),
+        format!(
+            "// Generated by gen-ts; do not edit.\nexport type EventType = {}; \
+             \nexport const FOLDED_EVENT_TYPES: EventType[] = [{}];\n",
+            union.join(" |\n  "),
+            folded.join(", ")
+        ),
+    )
+    .unwrap();
+
+    let names = [
+        "Event",
+        "HelloPayload",
+        "ResyncPayload",
+        "SystemSnapshot",
+        "DaemonInfo",
+        "HardwareInfo",
+        "ModelEntry",
+        "DownloadEntry",
+        "DownloadFailure",
+        "CatalogTier",
+        "CatalogResponse",
+        "EventType",
+        "Role",
+        "MessageStatus",
+        "Usage",
+        "Attachment",
+        "MessageRecord",
+        "ThreadRecord",
+    ];
+    std::fs::write(
+        dir.join("index.ts"),
+        format!(
+            "// Generated by gen-ts; do not edit.\n{}\n",
+            names
+                .iter()
+                .map(|n| format!("export * from \"./{n}\";\n"))
+                .collect::<Vec<_>>()
+                .join(""),
+        ),
+    )
+    .unwrap();
+
+    println!("gen-ts: wrote {dir:?}");
+}
diff --git a/macos/crates/edge0-core/src/config.rs b/macos/crates/edge0-core/src/config.rs
new file mode 100644
index 0000000..71fd0f2
--- /dev/null
+++ b/macos/crates/edge0-core/src/config.rs
@@ -0,0 +1,362 @@
+//! Config precedence shared by daemon and CLI: flags > `EDGE0_*` env > `~/.edge0/config.toml` > defaults.
+
+use std::collections::BTreeMap;
+use std::net::{IpAddr, SocketAddr};
+
+use serde::Deserialize;
+
+/// `direct` = Hugging Face; `registry` = signed-manifest world. Default is `direct`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum PullMode {
+    Direct,
+    Registry,
+}
+
+impl PullMode {
+    pub fn parse(v: &str) -> Result<Self, String> {
+        match v.trim().to_ascii_lowercase().as_str() {
+            "direct" => Ok(PullMode::Direct),
+            "registry" => Ok(PullMode::Registry),
+            other => Err(format!(
+                "invalid EDGE0_PULL_MODE `{other}`: expected direct | registry"
+            )),
+        }
+    }
+    pub fn as_str(self) -> &'static str {
+        match self {
+            PullMode::Direct => "direct",
+            PullMode::Registry => "registry",
+        }
+    }
+}
+
+/// Flag-layer input (`None` = that source did not set the key).
+#[derive(Debug, Default, Clone)]
+pub struct Overrides {
+    pub host: Option<String>,
+    pub port: Option<u16>,
+    pub bind: Option<String>,
+    pub origins: Option<String>,
+    pub source: Option<String>,
+    pub concurrency: Option<u32>,
+    pub keep_alive: Option<String>,
+    pub home: Option<String>,
+}
+
+impl Overrides {
+    fn to_section(&self) -> Edge0Section {
+        Edge0Section {
+            host: self.host.clone(),
+            port: self.port,
+            bind: self.bind.clone(),
+            origins: self.origins.clone(),
+            source: self.source.clone(),
+            concurrency: self.concurrency,
+            keep_alive: self.keep_alive.clone(),
+            registry_url: None,
+            pull_mode: None,
+            hf_base: None,
+            // LAN admission is env/config only; the GUI writes the config layer.
+            enable_lan: None,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default, Deserialize)]
+struct TomlFile {
+    edge0: Option<Edge0Section>,
+}
+
+#[derive(Debug, Clone, Default, Deserialize)]
+struct Edge0Section {
+    host: Option<String>,
+    port: Option<u16>,
+    bind: Option<String>,
+    origins: Option<String>,
+    source: Option<String>,
+    concurrency: Option<u32>,
+    keep_alive: Option<String>,
+    registry_url: Option<String>,
+    pull_mode: Option<String>,
+    hf_base: Option<String>,
+    /// Non-loopback bind with no existing token may mint the first token.
+    enable_lan: Option<String>,
+}
+
+#[derive(Debug, Clone)]
+pub struct Config {
+    pub host_url: String,
+    pub bind: IpAddr,
+    pub port: u16,
+    pub origins: Vec<String>,
+    pub source: Option<String>,
+    pub concurrency: u32,
+    pub keep_alive_default: String,
+    pub registry_urls: Vec<String>,
+    pub pull_mode: PullMode,
+    pub hf_base: String,
+    /// Explicit permission to mint a token on first non-loopback bind.
+    pub enable_lan: bool,
+    /// Unknown `config.toml` keys are kept; meaning of known keys is never changed.
+    pub extra: BTreeMap<String, String>,
+}
+
+pub const DEFAULT_ORIGINS: &str =
+    "http://127.0.0.1:8000,http://localhost:8000,tauri://localhost,http://tauri.localhost";
+
+pub const DEFAULT_REGISTRY_URL: &str = "https://registry.edge0.app/manifest.json";
+
+pub const DEFAULT_HF_BASE: &str = "https://huggingface.co";
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            host_url: "http://127.0.0.1:8000".into(),
+            bind: IpAddr::V4(std::net::Ipv4Addr::LOCALHOST),
+            port: 8000,
+            origins: DEFAULT_ORIGINS
+                .split(',')
+                .map(|s| s.trim().to_string())
+                .collect(),
+            source: None,
+            concurrency: 4,
+            keep_alive_default: "10m".into(),
+            registry_urls: vec![DEFAULT_REGISTRY_URL.into()],
+            pull_mode: PullMode::Direct,
+            hf_base: DEFAULT_HF_BASE.into(),
+            enable_lan: false,
+            extra: BTreeMap::new(),
+        }
+    }
+}
+
+impl Config {
+    /// Merge order: file < env < flag.
+    pub fn load(
+        overrides: &Overrides,
+        config_toml_path: Option<&std::path::Path>,
+    ) -> Result<Self, String> {
+        let mut cfg = Config::default();
+        if let Some(p) = config_toml_path.filter(|p| p.exists()) {
+            let text = std::fs::read_to_string(p)
+                .map_err(|e| format!("failed to read {}: {e}", p.display()))?;
+            let file: TomlFile = toml_parse(&text)?;
+            if let Some(s) = file.edge0 {
+                apply(&mut cfg, s);
+            }
+        }
+        let env = Edge0Section {
+            host: std::env::var("EDGE0_HOST").ok(),
+            port: std::env::var("EDGE0_PORT")
+                .ok()
+                .map(|v| v.trim().to_string())
+                .and_then(|v| v.parse().ok()),
+            bind: std::env::var("EDGE0_BIND").ok(),
+            origins: std::env::var("EDGE0_ORIGINS").ok(),
+            source: std::env::var("EDGE0_SOURCE").ok(),
+            concurrency: std::env::var("EDGE0_CONCURRENCY")
+                .ok()
+                .and_then(|v| v.parse().ok()),
+            keep_alive: std::env::var("EDGE0_KEEP_ALIVE").ok(),
+            registry_url: std::env::var("EDGE0_REGISTRY_URL").ok(),
+            pull_mode: std::env::var("EDGE0_PULL_MODE").ok(),
+            hf_base: std::env::var("EDGE0_HF_BASE").ok(),
+            enable_lan: std::env::var("EDGE0_ENABLE_LAN").ok(),
+        };
+        apply(&mut cfg, env);
+        apply(&mut cfg, overrides.to_section());
+        if let Some(h) = &overrides.home {
+            cfg.extra.insert("EDGE0_HOME".into(), h.clone());
+        }
+        cfg.validate()?;
+        Ok(cfg)
+    }
+
+    pub fn socket_addr(&self) -> SocketAddr {
+        SocketAddr::new(self.bind, self.port)
+    }
+
+    pub fn is_loopback_bind(&self) -> bool {
+        self.bind.is_loopback()
+    }
+
+    fn validate(&mut self) -> Result<(), String> {
+        let u = url_origin(&self.host_url)
+            .ok_or_else(|| format!("EDGE0_HOST is not a valid HTTP origin: {}", self.host_url))?;
+        self.host_url = u;
+        if self.port == 0 {
+            return Err("EDGE0_PORT must be in 1–65535".into());
+        }
+        for u in &self.registry_urls {
+            if !(u.starts_with("http://") || u.starts_with("https://")) {
+                return Err(format!(
+                    "EDGE0_REGISTRY_URL has invalid entry `{u}`: mirror URLs must start with http(s):// (comma-separated)"
+                ));
+            }
+        }
+        if let Some(v) = self.extra.get("EDGE0_PULL_MODE").cloned() {
+            self.pull_mode = PullMode::parse(&v)?;
+            self.extra.remove("EDGE0_PULL_MODE");
+        }
+        if !(self.hf_base.starts_with("http://") || self.hf_base.starts_with("https://")) {
+            return Err(format!(
+                "invalid EDGE0_HF_BASE `{}`: must start with http(s):// (e.g. {DEFAULT_HF_BASE})",
+                self.hf_base
+            ));
+        }
+        self.hf_base = self.hf_base.trim_end_matches('/').to_string();
+        Ok(())
+    }
+}
+
+fn apply(cfg: &mut Config, s: Edge0Section) {
+    if let Some(v) = nonempty(s.host) {
+        cfg.host_url = v;
+    }
+    if let Some(v) = s.port {
+        cfg.port = v;
+    }
+    if let Some(v) = nonempty(s.bind) {
+        if let Ok(ip) = v.parse::<IpAddr>() {
+            cfg.bind = ip;
+        }
+    }
+    if let Some(v) = nonempty(s.origins) {
+        cfg.origins = v
+            .split(',')
+            .map(|s| s.trim().to_string())
+            .filter(|s| !s.is_empty())
+            .collect();
+    }
+    if let Some(v) = nonempty(s.source) {
+        cfg.source = Some(v);
+    }
+    if let Some(v) = s.concurrency {
+        cfg.concurrency = v;
+    }
+    if let Some(v) = nonempty(s.keep_alive) {
+        cfg.keep_alive_default = v;
+    }
+    if let Some(v) = nonempty(s.registry_url) {
+        let urls: Vec<String> = v
+            .split(',')
+            .map(|u| u.trim().to_string())
+            .filter(|u| !u.is_empty())
+            .collect();
+        if !urls.is_empty() {
+            cfg.registry_urls = urls;
+        }
+    }
+    if let Some(v) = nonempty(s.pull_mode) {
+        // Invalid values fail in validate so the three-layer merge stays uniform.
+        cfg.extra.insert("EDGE0_PULL_MODE".into(), v);
+    }
+    if let Some(v) = nonempty(s.hf_base) {
+        cfg.hf_base = v;
+    }
+    if let Some(v) = s.enable_lan {
+        // Set and not an explicit negative (`0`/`false`/`off`/`no`) means admitted.
+        cfg.enable_lan = !matches!(v.trim(), "" | "0" | "false" | "off" | "no");
+    }
+}
+
+fn nonempty(s: Option<String>) -> Option<String> {
+    s.map(|v| v.trim().to_string()).filter(|v| !v.is_empty())
+}
+
+fn url_origin(s: &str) -> Option<String> {
+    let s = s.trim().trim_end_matches('/').to_string();
+    let rest = s
+        .strip_prefix("http://")
+        .or_else(|| s.strip_prefix("https://"))?;
+    if rest.is_empty() || rest.contains('/') || rest.contains(char::is_whitespace) {
+        return None;
+    }
+    Some(s)
+}
+
+/// Minimal `[edge0]` parser: `key = "value"` / integers. Unknown keys are ignored.
+fn toml_parse(text: &str) -> Result<TomlFile, String> {
+    let mut sec: Edge0Section = Default::default();
+    let mut cur_section = String::new();
+    for line in text.lines() {
+        let t = line.trim();
+        if t.is_empty() || t.starts_with('#') {
+            continue;
+        }
+        if let Some(s) = t.strip_prefix('[').and_then(|r| r.strip_suffix(']')) {
+            cur_section = s.trim().to_string();
+            continue;
+        }
+        let Some((k, v)) = t.split_once('=') else {
+            continue;
+        };
+        if cur_section != "edge0" {
+            continue;
+        }
+        let k = k.trim().to_string();
+        let v = v.trim().trim_matches('"').to_string();
+        match k.as_str() {
+            "host" => sec.host = Some(v),
+            "port" => sec.port = v.parse().ok(),
+            "bind" => sec.bind = Some(v),
+            "origins" => sec.origins = Some(v),
+            "source" => sec.source = Some(v),
+            "concurrency" => sec.concurrency = v.parse().ok(),
+            "keep_alive" => sec.keep_alive = Some(v),
+            "registry_url" => sec.registry_url = Some(v),
+            "pull_mode" => sec.pull_mode = Some(v),
+            "hf_base" => sec.hf_base = Some(v),
+            "enable_lan" => sec.enable_lan = Some(v),
+            _ => {}
+        }
+    }
+    Ok(TomlFile { edge0: Some(sec) })
+}
+
+/// Upsert keys in the `[edge0]` section; create the file if missing; leave other sections intact.
+pub fn upsert_edge0_keys(path: &std::path::Path, updates: &[(&str, &str)]) -> Result<(), String> {
+    let mut lines: Vec<String> = match std::fs::read_to_string(path) {
+        Ok(t) => t.lines().map(str::to_string).collect(),
+        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Vec::new(),
+        Err(e) => return Err(format!("failed to read {}: {e}", path.display())),
+    };
+    let start = match lines.iter().position(|l| l.trim() == "[edge0]") {
+        Some(i) => i,
+        None => {
+            if !lines.is_empty() && !lines.last().unwrap().trim().is_empty() {
+                lines.push(String::new());
+            }
+            lines.push("[edge0]".to_string());
+            lines.len() - 1
+        }
+    };
+    let mut end = start
+        + 1
+        + lines[start + 1..]
+            .iter()
+            .position(|l| l.trim().starts_with('['))
+            .unwrap_or(lines.len() - start - 1);
+    for (k, v) in updates {
+        let hit = lines[start + 1..end].iter().position(|l| {
+            l.trim()
+                .split_once('=')
+                .is_some_and(|(kk, _)| kk.trim() == *k)
+        });
+        match hit {
+            Some(rel) => lines[start + 1 + rel] = format!("{k} = \"{v}\""),
+            None => {
+                lines.insert(end, format!("{k} = \"{v}\""));
+                end += 1;
+            }
+        }
+    }
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
+    }
+    let mut out = lines.join("\n");
+    if !out.ends_with('\n') {
+        out.push('\n');
+    }
+    std::fs::write(path, out).map_err(|e| format!("failed to write {}: {e}", path.display()))
+}
diff --git a/macos/crates/edge0-core/src/conversation.rs b/macos/crates/edge0-core/src/conversation.rs
new file mode 100644
index 0000000..21bd144
--- /dev/null
+++ b/macos/crates/edge0-core/src/conversation.rs
@@ -0,0 +1,143 @@
+//! On-disk conversation record types shared by the app store and generated TS bindings.
+
+use serde::{Deserialize, Serialize};
+use ts_rs::TS;
+
+pub const RECORD_V: u32 = 1;
+
+/// System prompts are thread-level settings, so there is no `system` role.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
+#[serde(rename_all = "lowercase")]
+pub enum Role {
+    User,
+    Assistant,
+    Tool,
+}
+
+/// Terminal transitions are only `streaming → done | aborted | error`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
+#[serde(rename_all = "lowercase")]
+pub enum MessageStatus {
+    Streaming,
+    Done,
+    Aborted,
+    Error,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
+pub struct Usage {
+    pub prompt_tokens: u64,
+    pub completion_tokens: u64,
+    pub total_tokens: u64,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TS)]
+#[serde(tag = "kind", rename_all = "lowercase")]
+pub enum Attachment {
+    Text {
+        text: String,
+    },
+    #[serde(rename = "local_path")]
+    LocalPath {
+        path: String,
+    },
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
+pub struct MessageRecord {
+    pub v: u32,
+    pub role: Role,
+    pub content: String,
+    /// Kept separate from `content` (maps to `delta.reasoning_content`).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub reasoning_content: Option<String>,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub attachments: Vec<Attachment>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub tool_calls: Option<Vec<serde_json::Value>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub usage: Option<Usage>,
+    pub status: MessageStatus,
+    /// Required when `status=error`; must exist in the error-code table.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub code: Option<String>,
+    pub model: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub params: Option<serde_json::Value>,
+    /// Epoch milliseconds.
+    pub created_at: i64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub thread_id: Option<String>,
+    /// Unknown fields are preserved on read and written back unchanged.
+    #[serde(flatten, default)]
+    pub extra: serde_json::Map<String, serde_json::Value>,
+}
+
+impl MessageRecord {
+    pub fn validate(&self) -> Result<(), String> {
+        if self.v != RECORD_V && self.v > RECORD_V {
+            return Err(format!(
+                "unknown record version v={} (this build supports up to {RECORD_V})",
+                self.v
+            ));
+        }
+        if self.v == 0 {
+            return Err("missing record version v (schema requires v ≥ 1)".into());
+        }
+        match (self.status, self.code.as_deref()) {
+            (MessageStatus::Error, None) => Err("status=error must include code".into()),
+            (MessageStatus::Error, Some(c)) if crate::errorcodes::lookup(c).is_none() => {
+                Err(format!("error code {c} is not in the error-code table"))
+            }
+            _ => Ok(()),
+        }
+    }
+
+    pub fn new(
+        role: Role,
+        content: impl Into<String>,
+        model: impl Into<String>,
+        created_at: i64,
+    ) -> Self {
+        Self {
+            v: RECORD_V,
+            role,
+            content: content.into(),
+            reasoning_content: None,
+            attachments: Vec::new(),
+            tool_calls: None,
+            usage: None,
+            status: MessageStatus::Streaming,
+            code: None,
+            model: model.into(),
+            params: None,
+            created_at,
+            thread_id: None,
+            extra: serde_json::Map::new(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
+pub struct ThreadRecord {
+    pub v: u32,
+    pub title: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub model: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub params: Option<serde_json::Value>,
+    #[serde(flatten, default)]
+    pub extra: serde_json::Map<String, serde_json::Value>,
+}
+
+impl ThreadRecord {
+    pub fn new(title: impl Into<String>, model: Option<String>) -> Self {
+        Self {
+            v: RECORD_V,
+            title: title.into(),
+            model,
+            params: None,
+            extra: serde_json::Map::new(),
+        }
+    }
+}
diff --git a/macos/crates/edge0-core/src/errorcodes.rs b/macos/crates/edge0-core/src/errorcodes.rs
new file mode 100644
index 0000000..0235b2b
--- /dev/null
+++ b/macos/crates/edge0-core/src/errorcodes.rs
@@ -0,0 +1,309 @@
+//! Error-code table and the unified HTTP error envelope. User-visible errors must not include a raw stack trace.
+
+use serde::Serialize;
+
+/// `plain` / `next` / `retry` are the user-facing strings for each code.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct Code {
+    pub id: &'static str,
+    pub http: u16, // 0 = not an HTTP surface (startup / in-stream)
+    pub kind: &'static str,
+    pub plain: &'static str,
+    pub next: &'static str,
+    pub retry: &'static str,
+}
+
+pub const ALL: &[Code] = &[
+    Code {
+        id: "E-DL-NET",
+        http: 502,
+        kind: "network",
+        plain: "Network interrupted or timed out; download stopped at the last verified byte range",
+        next: "Check the network, then `edge0 pull` the same tier (resumes automatically); or switch source with --source",
+        retry: "yes",
+    },
+    Code {
+        id: "E-DL-SRC",
+        http: 502,
+        kind: "network",
+        plain: "Source refused service (403/429/5xx) or is unreachable",
+        next: "Use --source to pick another source; retry later",
+        retry: "yes",
+    },
+    Code {
+        id: "E-DL-DISK",
+        http: 507,
+        kind: "storage",
+        plain: "Not enough disk space (preflight or while writing)",
+        next: "Free space as prompted or move EDGE0_HOME to another volume; run edge0 gc",
+        retry: "no (needs manual action)",
+    },
+    Code {
+        id: "E-MANIFEST-SIG",
+        http: 409,
+        kind: "manifest",
+        plain: "Manifest signature check failed: untrusted source or tampered content",
+        next: "Retry from another mirror; if it still fails, stop and report (do not bypass)",
+        retry: "yes, with caution (switch source once)",
+    },
+    Code {
+        id: "E-MANIFEST-EXPIRED",
+        http: 409,
+        kind: "manifest",
+        plain: "Manifest expired and no mirror returned a fresh copy",
+        next: "Contact the publisher / retry later; in-progress downloads are unaffected (snapshot resume)",
+        retry: "yes (retry later)",
+    },
+    Code {
+        id: "E-MANIFEST-MINVER",
+        http: 409,
+        kind: "manifest",
+        plain: "This tier requires a newer CLI (min_cli_version not met)",
+        next: "edge0 upgrade (macOS: update the app; headless: upgrade the CLI)",
+        retry: "no (upgrade first)",
+    },
+    Code {
+        id: "E-MEM-LOAD",
+        http: 507,
+        kind: "memory",
+        plain: "Not enough memory to load the model (weights/cache exceed available RAM)",
+        next: "Close memory-heavy apps; or stay on edge0-8b; edge0 doctor shows a minimum-RAM hint",
+        retry: "no (free memory first)",
+    },
+    Code {
+        id: "E-MEM-WORKSET",
+        http: 500,
+        kind: "memory",
+        plain: "Generation aborted because the working set exceeded budget (protects the rest of the machine)",
+        next: "Shorten the prompt or lower n_ctx/concurrency, then retry",
+        retry: "yes (change parameters)",
+    },
+    Code {
+        id: "E-GPU-DEVICE",
+        http: 503,
+        kind: "gpu",
+        plain: "No Metal device (no GUI session / chip older than M3 / device lost)",
+        next: "Run under a logged-in graphical session; edge0 doctor checks chip and macOS version",
+        retry: "no (change environment)",
+    },
+    Code {
+        id: "E-GPU-COMMAND",
+        http: 500,
+        kind: "gpu",
+        plain: "Metal command failed (driver or GPU-memory error)",
+        next: "Retry once; if it keeps failing, run edge0 status and share logs",
+        retry: "yes (once)",
+    },
+    Code {
+        id: "E-MODEL-MISSING",
+        http: 404,
+        kind: "model",
+        plain: "Tier is not installed or was deleted (files/ path missing)",
+        next: "Run edge0 pull <tier> as prompted (or download in the UI)",
+        retry: "yes (works after install)",
+    },
+    Code {
+        id: "E-MODEL-HASH",
+        http: 500,
+        kind: "model",
+        plain: "File hash does not match the manifest (local corruption or two sources diverged without a manifest)",
+        next: "edge0 pull this rev again (retries once from the other source); or edge0 rm and reinstall",
+        retry: "yes",
+    },
+    Code {
+        id: "E-MODEL-LAYOUT",
+        http: 409,
+        kind: "model",
+        plain: "Weight layout not recognized (new tier needs a newer engine)",
+        next: "Upgrade the app/CLI; if it still fails, check engine_requirements in the manifest",
+        retry: "no (upgrade first)",
+    },
+    Code {
+        id: "E-SRV-PARAM",
+        http: 400,
+        kind: "service",
+        plain: "Invalid request parameters (unknown field, wrong type, or a value this release does not support)",
+        next: "Fix the field named in message; the compatibility API only accepts existing OpenAI parameters (except keep_alive)",
+        retry: "no (fix parameters first)",
+    },
+    Code {
+        id: "E-SRV-NOTLOAD",
+        http: 404,
+        kind: "service",
+        plain: "Model is not loaded",
+        next: "edge0 load <tier> / load in the UI; pull first if not installed",
+        retry: "yes",
+    },
+    Code {
+        id: "E-SRV-BUSY",
+        http: 429,
+        kind: "service",
+        plain: "Concurrency limit exceeded",
+        next: "Wait for Retry-After, or lower client concurrency",
+        retry: "yes",
+    },
+    Code {
+        id: "E-SRV-CANCEL",
+        http: 499,
+        kind: "service",
+        plain: "Generation cancelled (client disconnect / e0_cancel / user stopped)",
+        next: "Expected; resend if you want to continue",
+        retry: "yes (resend)",
+    },
+    Code {
+        id: "E-SRV-PORT",
+        http: 0,
+        kind: "service",
+        plain: "Port is in use by another process; daemon refuses to start (does not auto-pick another port)",
+        next: "Free port 8000 or set EDGE0_PORT; edge0 doctor shows who holds it",
+        retry: "no (change config)",
+    },
+    Code {
+        id: "E-SRV-VERSION",
+        http: 409,
+        kind: "service",
+        plain: "Service version does not match client version",
+        next: "Restart the service or upgrade the app as prompted; no silent compatibility",
+        retry: "no (align versions first)",
+    },
+    Code {
+        id: "E-SRV-CONFLICT",
+        http: 409,
+        kind: "service",
+        plain: "Registry operation refused (for example, rm of a resident tier with in-flight requests)",
+        next: "Wait or cancel active requests, then retry; the UI shows why it is blocked",
+        retry: "yes",
+    },
+    Code {
+        id: "E-SRV-AUTH",
+        http: 401,
+        kind: "service",
+        plain: "Non-loopback access is missing a Bearer token",
+        next: "Copy the token from state/token or the app service panel into the client",
+        retry: "no (credentials required)",
+    },
+    Code {
+        id: "E-SRV-WORKER",
+        http: 500,
+        kind: "service",
+        plain: "Worker crashed; the streaming request ends with this code (SSE error event)",
+        next: "The daemon has reaped the worker and kept diagnostics; retry reloads it",
+        retry: "yes",
+    },
+];
+
+pub fn lookup(id: &str) -> Option<&'static Code> {
+    ALL.iter().find(|c| c.id == id)
+}
+
+/// OpenAI-style `type` for the HTTP envelope. `code` remains the authoritative id.
+pub fn openai_type(code: &Code) -> &'static str {
+    match code.id {
+        "E-SRV-AUTH" => "authentication_error",
+        "E-SRV-CONFLICT" | "E-MANIFEST-SIG" | "E-MANIFEST-EXPIRED" | "E-MANIFEST-MINVER" => {
+            "conflict_error"
+        }
+        "E-SRV-VERSION" => "version_error",
+        "E-SRV-BUSY" => "rate_limit_error",
+        "E-SRV-NOTLOAD" | "E-SRV-PARAM" | "E-MODEL-MISSING" => "invalid_request_error",
+        "E-SRV-WORKER" | "E-GPU-COMMAND" | "E-MODEL-HASH" => "server_error",
+        _ => "api_error",
+    }
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct ErrEnvelope {
+    pub error: ErrBody,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct ErrBody {
+    pub message: String,
+    pub r#type: &'static str,
+    pub code: &'static str,
+    pub param: Option<String>,
+}
+
+impl ErrEnvelope {
+    pub fn new(code: &Code, message: impl Into<String>, param: Option<String>) -> Self {
+        Self {
+            error: ErrBody {
+                message: message.into(),
+                r#type: openai_type(code),
+                code: code.id,
+                param,
+            },
+        }
+    }
+}
+
+/// Internal API error. Rendered as an envelope; never includes a stack trace.
+#[derive(Debug, Clone)]
+pub struct ApiError {
+    pub code: &'static Code,
+    pub message: String,
+    pub param: Option<String>,
+    pub retry_after: Option<u32>,
+    /// Origin rejection uses 403 + `type=forbidden_error`.
+    pub status_override: Option<u16>,
+    pub kind_override: Option<&'static str>,
+}
+
+impl ApiError {
+    pub fn new(code: &'static Code, message: impl Into<String>) -> Self {
+        Self {
+            code,
+            message: message.into(),
+            param: None,
+            retry_after: None,
+            status_override: None,
+            kind_override: None,
+        }
+    }
+    pub fn with_param(mut self, p: impl Into<String>) -> Self {
+        self.param = Some(p.into());
+        self
+    }
+    pub fn cors(message: impl Into<String>) -> Self {
+        let mut e = Self::new(lookup("E-SRV-AUTH").unwrap(), message);
+        e.status_override = Some(403);
+        e.kind_override = Some("forbidden_error");
+        e
+    }
+    pub fn with_retry_after(mut self, secs: u32) -> Self {
+        self.retry_after = Some(secs);
+        self
+    }
+    pub fn with_status(mut self, code: u16) -> Self {
+        self.status_override = Some(code);
+        self
+    }
+    pub fn envelope(&self) -> ErrEnvelope {
+        match self.kind_override {
+            None => ErrEnvelope::new(self.code, self.message.clone(), self.param.clone()),
+            Some(kind) => ErrEnvelope {
+                error: ErrBody {
+                    message: self.message.clone(),
+                    r#type: kind,
+                    code: self.code.id,
+                    param: self.param.clone(),
+                },
+            },
+        }
+    }
+}
+
+/// Map ABI `E0_ERR_*` (values match `engine-abi.h`) onto a table code.
+pub fn from_e0_err(e0_status: i32, details: Option<&str>) -> ApiError {
+    let code = match e0_status {
+        -3 => lookup("E-MODEL-MISSING").unwrap(), // E0_ERR_MODEL
+        -4 => lookup("E-MEM-LOAD").unwrap(),      // E0_ERR_MEM
+        -5 => lookup("E-GPU-DEVICE").unwrap(),    // E0_ERR_GPU
+        -7 => lookup("E-SRV-CANCEL").unwrap(),    // E0_ERR_CANCELLED
+        _ => lookup("E-SRV-WORKER").unwrap(),
+    };
+    let msg = details.unwrap_or(code.plain).to_string();
+    let mut e = ApiError::new(code, msg);
+    e.status_override = None;
+    e
+}
diff --git a/macos/crates/edge0-core/src/frames.rs b/macos/crates/edge0-core/src/frames.rs
new file mode 100644
index 0000000..8cc4b62
--- /dev/null
+++ b/macos/crates/edge0-core/src/frames.rs
@@ -0,0 +1,111 @@
+//! Length-prefixed JSON frames between daemon and worker.
+
+use serde::{Deserialize, Serialize};
+
+use crate::version::FRAME_PROTO_VERSION;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum Channel {
+    Content,
+    Reasoning,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum EndReason {
+    /// EOS or max_new_tokens.
+    Natural,
+    Cancelled,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "t", rename_all = "snake_case")]
+pub enum Frame {
+    Hello {
+        v: u32,
+        tier: String,
+        engine_lib: String,
+        abi: u32,
+        pid: u32,
+    },
+    Begin {
+        id: u64,
+        prompt_tokens: Vec<i32>,
+        params_json: String,
+    },
+    /// One token per frame; never coalesce.
+    Token {
+        id: u64,
+        token: i32,
+        channel: Channel,
+    },
+    Finished { id: u64, reason: EndReason },
+    /// Cancel the engine immediately; do not drain the queue.
+    Cancel { id: u64 },
+    Stats,
+    StatsResp { stats_json: String },
+    Ping,
+    Pong,
+    /// `details_json` is for logs, not the user-facing error body.
+    Error {
+        id: Option<u64>,
+        e0_status: i32,
+        details_json: String,
+    },
+    Bye { reason: String },
+}
+
+pub const MAX_FRAME_BYTES: u32 = 8 * 1024 * 1024;
+
+pub fn encode(frame: &Frame) -> Result<Vec<u8>, String> {
+    let mut body = serde_json::to_vec(frame).map_err(|e| e.to_string())?;
+    let len = body.len() as u32;
+    let mut out = Vec::with_capacity(4 + body.len());
+    out.extend_from_slice(&len.to_be_bytes());
+    out.append(&mut body);
+    Ok(out)
+}
+
+/// Peer close returns `Ok(None)`. A bad length is a protocol error.
+pub fn decode<R: std::io::BufRead>(r: &mut R) -> Result<Option<Frame>, String> {
+    let mut head = [0u8; 4];
+    if !read_exact_or_eof(r, &mut head)? {
+        return Ok(None);
+    }
+    let len = u32::from_be_bytes(head);
+    if len == 0 || len > MAX_FRAME_BYTES {
+        return Err(format!("invalid frame length: {len}"));
+    }
+    let mut buf = vec![0u8; len as usize];
+    if !read_exact_or_eof(r, &mut buf)? {
+        return Err("truncated frame body".into());
+    }
+    let frame: Frame = serde_json::from_slice(&buf).map_err(|e| format!("invalid frame JSON: {e}"))?;
+    Ok(Some(frame))
+}
+
+fn read_exact_or_eof<R: std::io::BufRead>(r: &mut R, buf: &mut [u8]) -> Result<bool, String> {
+    let mut filled = 0;
+    while filled < buf.len() {
+        match r.read(&mut buf[filled..]) {
+            Ok(0) if filled == 0 => return Ok(false),
+            Ok(0) => return Err("unexpected EOF mid-frame".into()),
+            Ok(n) => filled += n,
+            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
+            Err(e) => return Err(e.to_string()),
+        }
+    }
+    Ok(true)
+}
+
+pub fn write_to<W: std::io::Write>(w: &mut W, frame: &Frame) -> Result<(), String> {
+    w.write_all(&encode(frame)?).map_err(|e| e.to_string())?;
+    w.flush().map_err(|e| e.to_string())
+}
+
+impl Frame {
+    pub fn proto_version(&self) -> u32 {
+        FRAME_PROTO_VERSION
+    }
+}
diff --git a/macos/crates/edge0-core/src/keepalive.rs b/macos/crates/edge0-core/src/keepalive.rs
new file mode 100644
index 0000000..b71100c
--- /dev/null
+++ b/macos/crates/edge0-core/src/keepalive.rs
@@ -0,0 +1,38 @@
+//! `keep_alive` duration syntax: `5m` / `120s` / `1h` / `inf` (server default `10m`).
+//! `None` means `inf` (never auto-unload).
+
+use chrono::{DateTime, Utc};
+use std::time::Duration;
+
+pub fn parse(s: &str) -> Result<Option<Duration>, String> {
+    let t = s.trim().to_ascii_lowercase();
+    if t.is_empty() {
+        return Err(format!("invalid duration: {s}"));
+    }
+    if t == "inf" {
+        return Ok(None);
+    }
+    let (num, unit) = t.split_at(t.len() - 1);
+    let n: u64 = num.trim().parse().map_err(|_| format!("invalid duration: {s}"))?;
+    if n == 0 {
+        return Err(format!("invalid duration: {s}"));
+    }
+    let secs = match unit {
+        "s" => n,
+        "m" => n * 60,
+        "h" => n * 3600,
+        _ => return Err(format!("invalid duration: {s}")),
+    };
+    Ok(Some(Duration::from_secs(secs)))
+}
+
+pub fn default() -> Option<Duration> {
+    parse("10m").unwrap()
+}
+
+/// `inf` → None; otherwise now + duration.
+pub fn unload_at(dur: Option<Duration>, now: DateTime<Utc>) -> Option<DateTime<Utc>> {
+    let d = dur?;
+    let chrono_dur = chrono::Duration::from_std(d).ok()?;
+    Some(now + chrono_dur)
+}
diff --git a/macos/crates/edge0-core/src/lib.rs b/macos/crates/edge0-core/src/lib.rs
new file mode 100644
index 0000000..c94c315
--- /dev/null
+++ b/macos/crates/edge0-core/src/lib.rs
@@ -0,0 +1,16 @@
+//! Shared foundation for the edge0 runtime: paths, config precedence, error codes,
+//! on-disk state schema, worker frame protocol and keep_alive parsing.
+
+pub mod abi_load;
+pub mod config;
+pub mod conversation;
+pub mod errorcodes;
+pub mod frames;
+pub mod keepalive;
+pub mod manifest;
+pub mod paths;
+pub mod protocol;
+pub mod state;
+pub mod version;
+
+pub use errorcodes::{ApiError, Code, ErrEnvelope};
diff --git a/macos/crates/edge0-core/src/manifest.rs b/macos/crates/edge0-core/src/manifest.rs
new file mode 100644
index 0000000..820ebbf
--- /dev/null
+++ b/macos/crates/edge0-core/src/manifest.rs
@@ -0,0 +1,324 @@
+//! Signed-manifest types and the single parse/validate entry used by both rescan and install.
+
+use std::path::Path;
+
+use serde::{Deserialize, Serialize};
+
+use crate::paths::Home;
+use crate::state::InstalledTier;
+
+pub const TIER_ENUM: &[&str] = &["edge0-8b", "edge0-35b"];
+pub const SOURCE_ENUM: &[&str] = &["huggingface", "modelscope"];
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct Manifest {
+    pub schema_version: u32,
+    pub generated_at: String,
+    pub expires_at: String,
+    pub tiers: Vec<TierEntry>,
+    pub signature: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub signer_key_id: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct TierEntry {
+    pub tier: String,
+    pub rev: String,
+    pub min_cli_version: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub engine_requirements: Option<serde_json::Value>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub notes_url: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub license: Option<String>,
+    pub files: Vec<FileEntry>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct FileEntry {
+    pub path: String,
+    pub size: u64,
+    pub sha256: String,
+    /// Optional on a seeded copy; required on the install path (`edge0-pull::check`).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub sources: Option<Sources>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub cross_check: Option<std::collections::BTreeMap<String, String>>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct Sources {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub huggingface: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub modelscope: Option<String>,
+    pub primary: String,
+}
+
+impl Manifest {
+    pub fn from_json_str(text: &str) -> Result<Self, String> {
+        let m: Manifest = serde_json::from_str(text).map_err(|e| format!("manifest parse failed: {e}"))?;
+        m.validate_structure()?;
+        Ok(m)
+    }
+
+    pub fn validate_structure(&self) -> Result<(), String> {
+        if self.schema_version != 1 {
+            return Err(format!("unsupported schema_version: {}", self.schema_version));
+        }
+        parse_ts(&self.generated_at, "generated_at")?;
+        parse_ts(&self.expires_at, "expires_at")?;
+        if self.tiers.is_empty() {
+            return Err("tiers must not be empty".into());
+        }
+        for t in &self.tiers {
+            if !TIER_ENUM.contains(&t.tier.as_str()) {
+                return Err(format!("unknown tier: {}", t.tier));
+            }
+            validate_rev(&t.rev)?;
+            validate_semver(&t.min_cli_version)?;
+            if t.files.is_empty() {
+                return Err(format!("{}: files must not be empty", t.tier));
+            }
+            for f in &t.files {
+                validate_path(&f.path)?;
+                validate_sha256(&f.sha256)?;
+                if let Some(s) = &f.sources {
+                    if !SOURCE_ENUM.contains(&s.primary.as_str()) {
+                        return Err(format!("{}: invalid primary {}", f.path, s.primary));
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    /// Install path: every file must have a URL and `primary` (the two sources are not byte-equivalent).
+    pub fn validate_installable(&self) -> Result<(), String> {
+        self.validate_structure()?;
+        validate_signature_format(&self.signature, self.signer_key_id.as_deref())?;
+        for t in &self.tiers {
+            for f in &t.files {
+                let s = f
+                    .sources
+                    .as_ref()
+                    .ok_or_else(|| format!("{}: missing sources (required to install)", f.path))?;
+                s.huggingface
+                    .as_ref()
+                    .or(s.modelscope.as_ref())
+                    .ok_or_else(|| format!("{}: sources has no URL", f.path))?;
+            }
+        }
+        Ok(())
+    }
+
+    pub fn find_tier(&self, tier: &str) -> Option<&TierEntry> {
+        self.tiers.iter().find(|t| t.tier == tier)
+    }
+}
+
+/// Schema pattern `^[A-Za-z0-9+/]{86}==$` (base64 of 64 bytes = 88 chars).
+pub fn validate_signature_format(sig: &str, key_id: Option<&str>) -> Result<(), String> {
+    fn b64_body(s: &str, body_len: usize, pad: &str) -> bool {
+        s.len() == body_len + pad.len()
+            && s.ends_with(pad)
+            && s[..body_len]
+                .bytes()
+                .all(|c| c.is_ascii_alphanumeric() || matches!(c, b'+' | b'/'))
+    }
+    if !b64_body(sig, 86, "==") {
+        return Err("signature is not 88-char base64(Ed25519 64B)".into());
+    }
+    if let Some(k) = key_id {
+        if !b64_body(k, 11, "=") {
+            return Err("signer_key_id is not 12-char base64(8B)".into());
+        }
+    }
+    Ok(())
+}
+
+/// rev: `^rYYYY-MM-DD(-[a-z]+)?$`.
+fn validate_rev(rev: &str) -> Result<(), String> {
+    let rest = rev
+        .strip_prefix('r')
+        .ok_or_else(|| format!("rev must start with r: {rev}"))?;
+    let b = rest.as_bytes();
+    if b.len() < 10
+        || b[4] != b'-'
+        || b[7] != b'-'
+        || !b[..4].iter().all(|c| c.is_ascii_digit())
+        || !b[5..7].iter().all(|c| c.is_ascii_digit())
+        || !b[8..10].iter().all(|c| c.is_ascii_digit())
+    {
+        return Err(format!("rev does not match rYYYY-MM-DD: {rev}"));
+    }
+    let mut i = 10;
+    if i < b.len() {
+        if b[i] != b'-' {
+            return Err(format!("rev suffix must start with -: {rev}"));
+        }
+        i += 1;
+        if i == b.len() || !b[i..].iter().all(|c| c.is_ascii_lowercase()) {
+            return Err(format!("rev suffix must be lowercase letters: {rev}"));
+        }
+    }
+    Ok(())
+}
+
+fn validate_semver(s: &str) -> Result<(), String> {
+    parse_semver(s)
+        .map(|_| ())
+        .ok_or_else(|| format!("invalid version: {s}"))
+}
+
+pub fn parse_semver(s: &str) -> Option<(u64, u64, u64)> {
+    let mut it = s.split('.');
+    let a = it.next()?.parse().ok()?;
+    let b = it.next()?.parse().ok()?;
+    let c = it.next()?.parse().ok()?;
+    if it.next().is_some() {
+        return None;
+    }
+    Some((a, b, c))
+}
+
+/// Relative path only; no `..`; charset `[A-Za-z0-9._/ -]`, first char alphanumeric.
+pub fn validate_path(p: &str) -> Result<(), String> {
+    let b = p.as_bytes();
+    if b.is_empty() {
+        return Err("path must not be empty".into());
+    }
+    if !b[0].is_ascii_alphanumeric() {
+        return Err(format!("path must start with alphanumeric: {p}"));
+    }
+    for seg in p.split('/') {
+        if seg == ".." {
+            return Err(format!("path contains ..: {p}"));
+        }
+    }
+    for c in b {
+        if !(c.is_ascii_alphanumeric() || matches!(c, b'.' | b'_' | b'/' | b' ' | b'-')) {
+            return Err(format!("path contains illegal character: {p}"));
+        }
+    }
+    Ok(())
+}
+
+fn validate_sha256(s: &str) -> Result<(), String> {
+    if s.len() != 64
+        || !s
+            .bytes()
+            .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
+    {
+        return Err(format!("sha256 must be 64 lowercase hex digits: {s}"));
+    }
+    Ok(())
+}
+
+fn parse_ts(s: &str, field: &str) -> Result<chrono::DateTime<chrono::Utc>, String> {
+    chrono::DateTime::parse_from_rfc3339(s)
+        .map(|d| d.with_timezone(&chrono::Utc))
+        .map_err(|e| format!("{field} is not RFC 3339: {e}"))
+}
+
+pub fn load_copy(path: &Path) -> Result<Manifest, String> {
+    let text =
+        std::fs::read_to_string(path).map_err(|e| format!("failed to read {}: {e}", path.display()))?;
+    Manifest::from_json_str(&text)
+}
+
+/// Shared by daemon rescan and `edge0 list`. A broken install is not counted.
+/// `EDGE0_BUNDLED_MODELS` is a read-only second root; a user-home copy of the
+/// same tier wins.
+pub fn scan_installed(home: &Home) -> Vec<InstalledTier> {
+    let mut out = scan_models_root(&home.models(), false);
+    let seen: std::collections::HashSet<String> = out.iter().map(|t| t.tier.clone()).collect();
+    if let Some(root) = bundled_models_root() {
+        for t in scan_models_root(&root, true) {
+            if !seen.contains(&t.tier) {
+                out.push(t);
+            }
+        }
+    }
+    out.sort_by(|a, b| a.tier.cmp(&b.tier).then(a.rev.cmp(&b.rev)));
+    out
+}
+
+fn bundled_models_root() -> Option<std::path::PathBuf> {
+    std::env::var_os("EDGE0_BUNDLED_MODELS").map(std::path::PathBuf::from)
+}
+
+fn scan_models_root(models: &Path, bundled: bool) -> Vec<InstalledTier> {
+    let mut out = Vec::new();
+    let Ok(tier_rd) = std::fs::read_dir(models) else {
+        return out;
+    };
+    for tier_ent in tier_rd.flatten() {
+        let tier_dir = tier_ent.path();
+        if !tier_dir.is_dir() {
+            continue;
+        }
+        let tier = tier_dir.file_name().unwrap().to_string_lossy().to_string();
+        let mut revs: Vec<_> = std::fs::read_dir(&tier_dir)
+            .map(|rd| {
+                rd.filter_map(|e| e.ok())
+                    .map(|e| e.path())
+                    .filter(|p| p.is_dir())
+                    .collect()
+            })
+            .unwrap_or_default();
+        revs.sort();
+        for rev_dir in revs {
+            let rev = rev_dir.file_name().unwrap().to_string_lossy().to_string();
+            if let Ok(text) = std::fs::read_to_string(rev_dir.join("install.json")) {
+                if let Ok(snap) = serde_json::from_str::<crate::state::DirectSnapshot>(&text) {
+                    if snap.commit == rev
+                        && snap
+                            .files
+                            .iter()
+                            .all(|f| rev_dir.join("files").join(&f.path).is_file())
+                    {
+                        out.push(InstalledTier {
+                            tier: tier.clone(),
+                            rev: rev.clone(),
+                            dir: rev_dir.clone(),
+                            bytes_total: snap.bytes_total(),
+                            mode: "direct".to_string(),
+                            install: Some(snap),
+                            bundled,
+                        });
+                    }
+                    continue;
+                }
+            }
+            let Ok(m) = load_copy(&rev_dir.join("manifest.json")) else {
+                continue;
+            };
+            let Some(t) = m.find_tier(&tier).filter(|t| t.rev == rev) else {
+                continue;
+            };
+            let files_ok = t
+                .files
+                .iter()
+                .all(|f| rev_dir.join("files").join(&f.path).is_file());
+            if files_ok {
+                out.push(InstalledTier {
+                    tier: tier.clone(),
+                    rev: rev.clone(),
+                    dir: rev_dir.clone(),
+                    bytes_total: t.files.iter().map(|f| f.size).sum(),
+                    mode: "registry".to_string(),
+                    install: None,
+                    bundled,
+                });
+            }
+        }
+    }
+    out
+}
+
diff --git a/macos/crates/edge0-core/src/paths.rs b/macos/crates/edge0-core/src/paths.rs
new file mode 100644
index 0000000..843a581
--- /dev/null
+++ b/macos/crates/edge0-core/src/paths.rs
@@ -0,0 +1,102 @@
+//! `~/.edge0` layout and the `EDGE0_HOME` instance key.
+
+use std::path::PathBuf;
+
+#[derive(Debug, Clone)]
+pub struct Home {
+    pub root: PathBuf,
+}
+
+impl Home {
+    pub fn new(root: impl Into<PathBuf>) -> Self {
+        Self { root: root.into() }
+    }
+
+    /// `EDGE0_HOME` wins over `~/.edge0`.
+    pub fn from_env() -> anyhow_free::Result<Self> {
+        if let Some(h) = std::env::var_os("EDGE0_HOME") {
+            return Ok(Self {
+                root: PathBuf::from(h),
+            });
+        }
+        let home = std::env::var_os("HOME").ok_or("HOME is unset and EDGE0_HOME was not given")?;
+        Ok(Self {
+            root: PathBuf::from(home).join(".edge0"),
+        })
+    }
+
+    pub fn blobs(&self) -> PathBuf {
+        self.root.join("blobs")
+    }
+    pub fn models(&self) -> PathBuf {
+        self.root.join("models")
+    }
+    pub fn manifests(&self) -> PathBuf {
+        self.root.join("manifests")
+    }
+    pub fn state(&self) -> PathBuf {
+        self.root.join("state")
+    }
+    pub fn logs(&self) -> PathBuf {
+        self.root.join("logs")
+    }
+    pub fn tmp(&self) -> PathBuf {
+        self.root.join("tmp")
+    }
+    pub fn run(&self) -> PathBuf {
+        self.root.join("run")
+    }
+    pub fn daemon_json(&self) -> PathBuf {
+        self.state().join("daemon.json")
+    }
+    pub fn events_json(&self) -> PathBuf {
+        self.state().join("events.json")
+    }
+    pub fn token(&self) -> PathBuf {
+        self.state().join("token")
+    }
+    /// Session DB: only the app process writes this; daemon/CLI must not.
+    pub fn edge0_db(&self) -> PathBuf {
+        self.state().join("edge0.db")
+    }
+    pub fn registry_lock(&self) -> PathBuf {
+        self.state().join("registry.lock")
+    }
+    pub fn blobrefs(&self) -> PathBuf {
+        self.state().join("blobrefs.json")
+    }
+    pub fn config_toml(&self) -> PathBuf {
+        self.root.join("config.toml")
+    }
+    pub fn worker_json(&self, tier: &str) -> PathBuf {
+        self.state().join(format!("worker-{tier}.json"))
+    }
+    pub fn download_json(&self, id: &str) -> PathBuf {
+        self.state().join("downloads").join(format!("{id}.json"))
+    }
+    pub fn worker_sock(&self, tier: &str) -> PathBuf {
+        self.run().join(format!("worker-{tier}.sock"))
+    }
+    /// Installed tier directory; a `manifest.json` copy means the install is verified.
+    pub fn model_dir(&self, tier: &str, rev: &str) -> PathBuf {
+        self.models().join(tier).join(rev)
+    }
+
+    pub fn ensure_dirs(&self) -> anyhow_free::Result<()> {
+        for d in [
+            self.state(),
+            self.state().join("downloads"),
+            self.logs(),
+            self.tmp(),
+            self.run(),
+        ] {
+            std::fs::create_dir_all(&d).map_err(|e| format!("failed to create {}: {e}", d.display()))?;
+        }
+        Ok(())
+    }
+}
+
+// core does not depend on anyhow; keep a tiny error alias.
+pub mod anyhow_free {
+    pub type Result<T> = std::result::Result<T, String>;
+}
diff --git a/macos/crates/edge0-core/src/protocol.rs b/macos/crates/edge0-core/src/protocol.rs
new file mode 100644
index 0000000..8c08beb
--- /dev/null
+++ b/macos/crates/edge0-core/src/protocol.rs
@@ -0,0 +1,177 @@
+//! Wire types for the event stream and `/v1/edge0/system` / catalog snapshots.
+
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+use ts_rs::TS;
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct Event {
+    pub seq: u64,
+    pub ts: DateTime<Utc>,
+    #[serde(rename = "type")]
+    pub etype: String,
+    pub subject: String,
+    pub payload: serde_json::Value,
+}
+
+pub const EVENT_TYPES: &[&str] = &[
+    "daemon.hello",
+    "model.load.started",
+    "model.load.ready",
+    "model.load.failed",
+    "model.unloaded",
+    "download.progress",
+    "download.paused",
+    "download.completed",
+    "download.failed",
+    "worker.crashed",
+    "request.rejected",
+    "api.request.finished",
+    "config.changed",
+    "service.restart",
+    "resync.required",
+];
+
+/// Foldable types may be coalesced; delivery is still in-order.
+pub fn foldable(t: &str) -> bool {
+    matches!(
+        t,
+        "download.progress" | "request.rejected" | "config.changed"
+    )
+}
+
+pub fn known_type(t: &str) -> bool {
+    EVENT_TYPES.contains(&t)
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct HelloPayload {
+    pub session_id: String,
+    pub seq: u64,
+    pub abi_version: u32,
+}
+
+/// After `resync.required`, the client must refetch `/v1/edge0/system`.
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct ResyncPayload {
+    pub gap_from: u64,
+    pub buffer_min: u64,
+    pub session_id: String,
+}
+
+// ---------------------------------------------------------------- `GET /v1/edge0/system`
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct SystemSnapshot {
+    pub daemon: DaemonInfo,
+    pub hardware: HardwareInfo,
+    pub models: Vec<ModelEntry>,
+    pub downloads: Vec<DownloadEntry>,
+    pub events_seq: u64,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct DaemonInfo {
+    pub version: String,
+    pub abi_version: u32,
+    /// Under menu-agent, the launching app version; null for CLI/launchd.
+    pub host_app_version: Option<String>,
+    pub scope: String,
+    pub started_by: String,
+    /// Host app pid under menu-agent (cascade-stop watch); null otherwise.
+    pub owner_pid: Option<u32>,
+    pub session_id: String,
+    pub started_at: String,
+    pub uptime_s: u64,
+    pub throttled: bool,
+    pub port: u16,
+    pub bound: String,
+    /// Resolved `EDGE0_HOME`; the UI must not invent this path.
+    pub home: String,
+    /// Backend identity from the loaded library basename (`native`/`replay`/verbatim).
+    pub engine: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct HardwareInfo {
+    pub chip: String,
+    pub macos: String,
+    pub mem_total_gib: Option<f64>,
+    pub metal_device: bool,
+    pub disk_free_gib: f64,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct ModelEntry {
+    pub id: String,
+    pub rev: String,
+    pub bytes_total: u64,
+    /// `resident|unloaded`. In-flight load is reported via `model.load.*` events.
+    pub state: String,
+    pub unload_at: Option<String>,
+    pub active_requests: u64,
+    pub context_used_tokens: u64,
+    /// True when the files come from the app bundle rather than `~/.edge0`.
+    #[serde(default)]
+    pub bundled: bool,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct DownloadEntry {
+    pub id: String,
+    pub tier: String,
+    pub rev: String,
+    pub phase: String,
+    pub source: Option<String>,
+    pub source_reason: Option<String>,
+    pub bytes_done: u64,
+    pub bytes_total: u64,
+    pub paused: bool,
+    pub failed: Option<DownloadFailure>,
+    /// Sliding-window fetch throughput (B/s). Omitted when there is no sample.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub rate_bps: Option<u64>,
+    /// Remaining seconds. Omitted when it cannot be estimated (not 0).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub eta_s: Option<u64>,
+    /// True if verify failed once and the runner retried from the other source.
+    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
+    pub source_switched: bool,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct DownloadFailure {
+    pub code: String,
+    pub message: String,
+}
+
+// ---------------------------------------------------------------- `GET /v1/edge0/catalog`
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct CatalogTier {
+    pub tier: String,
+    pub rev: String,
+    pub min_cli_version: String,
+    pub bytes_total: u64,
+    pub files_count: u64,
+    pub license: Option<String>,
+    pub notes_url: Option<String>,
+    pub primary_source: String,
+    pub has_modelscope: bool,
+    pub has_huggingface: bool,
+    /// None = not installed; Some(rev) = installed rev.
+    pub installed: Option<String>,
+    /// Manifest rev ≠ installed rev. Always false when not installed.
+    pub update_available: bool,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+pub struct CatalogResponse {
+    pub generated_at: String,
+    pub expires_at: String,
+    pub fetched_at: String,
+    /// True when serving an expired cache after every refetch failed.
+    pub stale: bool,
+    pub signer_key_id: String,
+    pub tiers: Vec<CatalogTier>,
+}
diff --git a/macos/crates/edge0-core/src/state.rs b/macos/crates/edge0-core/src/state.rs
new file mode 100644
index 0000000..fe3266e
--- /dev/null
+++ b/macos/crates/edge0-core/src/state.rs
@@ -0,0 +1,270 @@
+//! On-disk state schemas. `daemon.json` / `worker-*.json` are hints; liveness is `/health`.
+
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+
+use crate::version::E0_ABI_VERSION;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub enum Scope {
+    MenuAgent,
+    LaunchAgent,
+    Foreground,
+}
+
+impl Scope {
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Scope::MenuAgent => "menu-agent",
+            Scope::LaunchAgent => "launchagent",
+            Scope::Foreground => "foreground",
+        }
+    }
+    pub fn parse(s: &str) -> Result<Self, String> {
+        match s {
+            "menu-agent" => Ok(Scope::MenuAgent),
+            "launchagent" | "launch-agent" => Ok(Scope::LaunchAgent),
+            "foreground" => Ok(Scope::Foreground),
+            other => Err(format!("unknown scope: {other}")),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DaemonJson {
+    pub scope: String,
+    pub started_by: String,
+    pub owner_pid: Option<u32>,
+    pub session_id: String,
+    pub version: String,
+    pub abi_version: u32,
+    pub port: u16,
+    pub pid: u32,
+    pub started_at: DateTime<Utc>,
+    /// Backend identity from the loaded library basename (`native`/`replay`/verbatim).
+    pub engine: String,
+}
+
+impl DaemonJson {
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        scope: Scope,
+        started_by: &str,
+        owner_pid: Option<u32>,
+        session_id: String,
+        port: u16,
+        engine: String,
+    ) -> Self {
+        Self {
+            scope: scope.as_str().to_string(),
+            started_by: started_by.to_string(),
+            owner_pid,
+            session_id,
+            version: crate::version::daemon_version().to_string(),
+            abi_version: E0_ABI_VERSION,
+            port,
+            pid: std::process::id(),
+            started_at: Utc::now(),
+            engine,
+        }
+    }
+    pub fn write(&self, path: &std::path::Path) -> Result<(), String> {
+        crate::state::atomic_write_json(path, self)
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WorkerJson {
+    pub tier: String,
+    pub pid: u32,
+    pub engine_lib: String,
+    pub abi_version: u32,
+    pub started_at: DateTime<Utc>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DownloadTask {
+    pub id: String,
+    pub tier: String,
+    pub rev: String,
+    pub phase: String,
+    pub source: Option<String>,
+    pub bytes_total: u64,
+    pub bytes_done: u64,
+    pub paused: bool,
+    /// Resume trusts only these recorded ranges.
+    pub files: Vec<FileTask>,
+    pub created_at: DateTime<Utc>,
+    pub updated_at: DateTime<Utc>,
+    /// Signed-manifest snapshot at task creation; `expires_at` does not abort an in-flight download.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub snapshot: Option<crate::manifest::Manifest>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub source_reason: Option<String>,
+    /// None = active or completed. Cleared on resume.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub failed: Option<TaskFailure>,
+    /// In-memory only: progress files must not pay write amplification for a live rate.
+    #[serde(skip, default)]
+    pub rate_bps: Option<u64>,
+    #[serde(skip, default)]
+    pub eta_s: Option<u64>,
+    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
+    pub source_switched: bool,
+    /// Direct-mode snapshot; mutually exclusive with `snapshot`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub plan: Option<DirectSnapshot>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TaskFailure {
+    pub code: String,
+    pub message: String,
+    pub at: DateTime<Utc>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FileTask {
+    pub path: String,
+    pub size: u64,
+    pub sha256: String,
+    pub done_ranges: Vec<(u64, u64)>,
+}
+
+impl FileTask {
+    /// One stream per file, so the range is always a single `[0, n)` prefix.
+    pub fn done(&self) -> u64 {
+        match self.done_ranges.first() {
+            Some((0, b)) => *b,
+            _ => 0,
+        }
+    }
+    pub fn set_done(&mut self, n: u64) {
+        self.done_ranges = if n == 0 { vec![] } else { vec![(0, n)] };
+    }
+    pub fn complete(&self) -> bool {
+        self.done() >= self.size
+    }
+}
+
+/// Restart seq = max + 10000.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EventsState {
+    pub max_seq: u64,
+    pub updated_at: DateTime<Utc>,
+}
+
+#[derive(Debug, Clone)]
+pub struct InstalledTier {
+    pub tier: String,
+    pub rev: String,
+    pub dir: std::path::PathBuf,
+    pub bytes_total: u64,
+    /// `"registry"` | `"direct"`.
+    pub mode: String,
+    pub install: Option<DirectSnapshot>,
+    /// True when this copy lives under `EDGE0_BUNDLED_MODELS` (read-only).
+    pub bundled: bool,
+}
+
+/// Direct-mode plan and, after commit, the on-disk `install.json`.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DirectSnapshot {
+    pub mode: String,
+    pub source: String,
+    pub base_url: String,
+    pub repo: String,
+    pub commit: String,
+    pub files: Vec<DirectFile>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub recorded_at: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DirectFile {
+    pub path: String,
+    pub size: u64,
+    /// Local sha256 after first download; HF LFS oid is a git sha and is not used as a checksum.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub sha256: Option<String>,
+    /// Head/mid/tail block digests for files >16 MiB; full sha256 cannot be recomputed from blocks.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub probe_blocks: Option<Vec<BlockProbe>>,
+}
+
+pub const SPOT_SMALL_MAX: u64 = 16 * 1024 * 1024;
+pub const SPOT_BLOCK_LEN: u64 = 4 * 1024 * 1024;
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct BlockProbe {
+    pub offset: u64,
+    pub len: u64,
+    pub sha256: String,
+}
+
+/// Record and spot-check must use the same offsets.
+pub fn probe_positions(size: u64) -> Vec<(u64, u64)> {
+    if size <= SPOT_SMALL_MAX {
+        return vec![];
+    }
+    let bl = SPOT_BLOCK_LEN.min(size);
+    let mid = (size / 2).saturating_sub(bl / 2).min(size - bl);
+    vec![(0, bl), (mid, bl), (size - bl, bl)]
+}
+
+impl DirectSnapshot {
+    pub fn url_for(&self, path: &str) -> String {
+        format!(
+            "{}/{}/resolve/{}/{}",
+            self.base_url, self.repo, self.commit, path
+        )
+    }
+    pub fn bytes_total(&self) -> u64 {
+        self.files.iter().map(|f| f.size).sum()
+    }
+}
+
+/// Visibility changes go through rename, not in-place overwrite.
+pub fn atomic_write_json<T: Serialize>(path: &std::path::Path, value: &T) -> Result<(), String> {
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent)
+            .map_err(|e| format!("failed to create {}: {e}", parent.display()))?;
+    }
+    let tmp = path.with_extension("json.tmp");
+    let bytes = serde_json::to_vec_pretty(value).map_err(|e| e.to_string())?;
+    std::fs::write(&tmp, bytes).map_err(|e| format!("failed to write {}: {e}", tmp.display()))?;
+    std::fs::rename(&tmp, path).map_err(|e| format!("failed to rename onto {}: {e}", path.display()))?;
+    Ok(())
+}
+
+pub struct RegistryLock {
+    file: std::fs::File,
+}
+
+impl RegistryLock {
+    pub fn acquire(path: &std::path::Path) -> Result<Self, String> {
+        if let Some(parent) = path.parent() {
+            std::fs::create_dir_all(parent).ok();
+        }
+        let file = std::fs::OpenOptions::new()
+            .create(true)
+            .write(true)
+            .truncate(false)
+            .open(path)
+            .map_err(|e| format!("failed to open registry.lock: {e}"))?;
+        let rc = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
+        if rc != 0 {
+            return Err("registry is busy (registry.lock already held)".into());
+        }
+        Ok(Self { file })
+    }
+}
+
+impl Drop for RegistryLock {
+    fn drop(&mut self) {
+        unsafe { libc::flock(self.file.as_raw_fd(), libc::LOCK_UN) };
+    }
+}
+
+use std::os::fd::AsRawFd;
diff --git a/macos/crates/edge0-core/src/version.rs b/macos/crates/edge0-core/src/version.rs
new file mode 100644
index 0000000..dcab296
--- /dev/null
+++ b/macos/crates/edge0-core/src/version.rs
@@ -0,0 +1,9 @@
+//! ABI and crate version constants mirrored from `engine-abi.h`.
+
+/// Bump this when the ABI semantics change.
+pub const E0_ABI_VERSION: u32 = 2;
+pub const FRAME_PROTO_VERSION: u32 = 1;
+/// Product version reported as `abi_version` on the compatibility API.
+pub fn daemon_version() -> &'static str {
+    env!("CARGO_PKG_VERSION")
+}
diff --git a/macos/crates/edge0-engine/Cargo.toml b/macos/crates/edge0-engine/Cargo.toml
new file mode 100644
index 0000000..2ba953f
--- /dev/null
+++ b/macos/crates/edge0-engine/Cargo.toml
@@ -0,0 +1,22 @@
+[package]
+name = "edge0-engine"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+[[bin]]
+name = "edge0-engine"
+path = "src/main.rs"
+
+[dependencies]
+edge0-core.workspace = true
+tokio.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+libloading.workspace = true
+libc.workspace = true
+thiserror.workspace = true
+anyhow.workspace = true
+tracing.workspace = true
+tracing-subscriber.workspace = true
diff --git a/macos/crates/edge0-engine/src/main.rs b/macos/crates/edge0-engine/src/main.rs
new file mode 100644
index 0000000..7eaecd4
--- /dev/null
+++ b/macos/crates/edge0-engine/src/main.rs
@@ -0,0 +1,340 @@
+//! `edge0-engine` worker: one process per resident tier; dlopen the engine and serve the daemon over frames.
+
+use std::collections::HashMap;
+use std::ffi::c_void;
+use std::os::unix::net::UnixStream;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex};
+
+/// Opaque ABI handle is only touched on the generate thread; cancel on the reader thread is ABI-idempotent.
+#[derive(Clone, Copy)]
+struct SendPtr(*mut std::ffi::c_void);
+unsafe impl Send for SendPtr {}
+
+type LiveStreams = std::sync::Arc<
+    std::sync::Mutex<
+        std::collections::HashMap<u64, (SendPtr, std::sync::Arc<std::sync::atomic::AtomicBool>)>,
+    >,
+>;
+impl SendPtr {
+    /// Must go through this method: capturing `.0` inside a closure treats `*mut` as the capture type.
+    fn get(self) -> *mut std::ffi::c_void {
+        self.0
+    }
+}
+
+use edge0_core::abi_load::EngineLib;
+use edge0_core::frames::{self, Channel, EndReason, Frame};
+
+fn main() {
+    if std::env::args().any(|arg| arg == "--version" || arg == "-V") {
+        println!("edge0-engine {}", env!("CARGO_PKG_VERSION"));
+        return;
+    }
+    if let Err(e) = run() {
+        eprintln!("edge0-engine: {e}");
+        std::process::exit(1);
+    }
+}
+
+fn run() -> Result<(), String> {
+    let sock =
+        std::env::var("EDGE0_WORKER_SOCK").map_err(|_| "missing EDGE0_WORKER_SOCK".to_string())?;
+    let lib_path =
+        std::env::var("EDGE0_ENGINE_LIB").map_err(|_| "missing EDGE0_ENGINE_LIB".to_string())?;
+    let tier = std::env::var("EDGE0_TIER").unwrap_or_else(|_| "unknown".into());
+    let model_dir = std::env::var("EDGE0_MODEL_DIR").unwrap_or_default();
+
+    let lib = Arc::new(EngineLib::load(std::path::Path::new(&lib_path))?);
+    let mut cfg = serde_json::Map::new();
+    if let Some(fwd) = tier.strip_prefix("edge0-") {
+        cfg.insert("forward".into(), serde_json::json!(fwd));
+    }
+    cfg.insert("lora".into(), serde_json::json!(true));
+    cfg.insert(
+        "model_identity".into(),
+        serde_json::json!({
+            "tier": tier.clone(),
+            "revision": std::env::var("EDGE0_MODEL_REV").unwrap_or_else(|_| "unknown".into()),
+            "engine_abi": edge0_core::version::E0_ABI_VERSION,
+            "engine_lib": lib_path.clone(),
+        }),
+    );
+    cfg.insert(
+        "prefix_cache".into(),
+        serde_json::json!(prefix_cache_config()),
+    );
+    let cfg = serde_json::Value::Object(cfg).to_string();
+    let c_tier = std::ffi::CString::new(tier.as_str()).unwrap();
+    let c_dir = std::ffi::CString::new(model_dir.as_str()).unwrap();
+    let c_cfg = std::ffi::CString::new(cfg.as_str()).unwrap();
+    let engine = unsafe { (lib.engine_create)(c_tier.as_ptr(), c_dir.as_ptr(), c_cfg.as_ptr()) };
+    if engine.is_null() {
+        return Err(format!(
+            "e0_engine_create failed: {}",
+            engine_err(&lib, std::ptr::null_mut())
+        ));
+    }
+
+    let stream = connect_with_retry(&sock, std::time::Duration::from_secs(10))?;
+    let writer = Arc::new(Mutex::new(stream.try_clone().map_err(|e| e.to_string())?));
+    send(
+        &writer,
+        &Frame::Hello {
+            v: edge0_core::version::FRAME_PROTO_VERSION,
+            tier,
+            engine_lib: lib_path,
+            abi: edge0_core::version::E0_ABI_VERSION,
+            pid: std::process::id(),
+        },
+    )?;
+
+    let streams: LiveStreams = Arc::new(Mutex::new(HashMap::new()));
+    let (tx, rx) = std::sync::mpsc::channel::<(u64, Vec<i32>, String)>();
+    let lib2 = lib.clone();
+    let w2 = writer.clone();
+    let s2 = streams.clone();
+    let eng2 = engine;
+    // One begin at a time in this process; concurrency is gated by the daemon.
+    let eng2 = SendPtr(eng2);
+    let gen = std::thread::spawn(move || {
+        for (id, prompt, params) in rx {
+            run_begin(&lib2, eng2.get(), &w2, &s2, id, prompt, params);
+        }
+    });
+
+    let mut reader = std::io::BufReader::new(stream);
+    loop {
+        match frames::decode(&mut reader) {
+            Ok(None) => break,
+            Ok(Some(Frame::Begin {
+                id,
+                prompt_tokens,
+                params_json,
+            })) => {
+                if tx.send((id, prompt_tokens, params_json)).is_err() {
+                    break;
+                }
+            }
+            Ok(Some(Frame::Cancel { id })) => {
+                let h = streams.lock().unwrap().get(&id).map(|l| (l.0, l.1.clone()));
+                if let Some((handle, flag)) = h {
+                    flag.store(true, Ordering::SeqCst);
+                    unsafe { (lib.cancel)(engine, handle.0) };
+                }
+            }
+            Ok(Some(Frame::Stats)) => {
+                let json = unsafe { lib.call_json(|b, l, w| (lib.stats)(engine, b, l, w)) }
+                    .unwrap_or_else(|_| "{\"error\":\"stats failed\"}".into());
+                let _ = send(&writer, &Frame::StatsResp { stats_json: json });
+            }
+            Ok(Some(Frame::Ping)) => {
+                let _ = send(&writer, &Frame::Pong);
+            }
+            Ok(Some(Frame::Bye { reason })) => {
+                eprintln!("edge0-engine: daemon Bye({reason}), exiting");
+                break;
+            }
+            Ok(Some(other)) => {
+                eprintln!("edge0-engine: unexpected frame {other:?}");
+            }
+            Err(e) => {
+                eprintln!("edge0-engine: frame protocol error: {e}");
+                break;
+            }
+        }
+    }
+    drop(gen);
+    unsafe { (lib.engine_destroy)(engine) };
+    Ok(())
+}
+
+fn prefix_cache_config() -> serde_json::Value {
+    prefix_cache_config_from_values(
+        std::env::var("EDGE0_PREFIX_CACHE").ok(),
+        std::env::var("EDGE0_PREFIX_CACHE_BUDGET_BYTES").ok(),
+        std::env::var("EDGE0_PREFIX_CACHE_MIN_TOKENS").ok(),
+    )
+}
+
+fn prefix_cache_config_from_values(
+    mode: Option<String>,
+    budget: Option<String>,
+    min_tokens: Option<String>,
+) -> serde_json::Value {
+    let enabled = match mode.as_deref() {
+        Some(v)
+            if matches!(
+                v.trim().to_ascii_lowercase().as_str(),
+                "off" | "0" | "false"
+            ) =>
+        {
+            false
+        }
+        Some(v) if matches!(v.trim().to_ascii_lowercase().as_str(), "on" | "1" | "true") => true,
+        Some(v) => {
+            eprintln!("edge0-engine: invalid EDGE0_PREFIX_CACHE {v:?}, defaulting to on");
+            true
+        }
+        None => true,
+    };
+    let budget_bytes = budget
+        .and_then(|v| v.parse::<u64>().ok())
+        .unwrap_or(512 * 1024 * 1024);
+    let min_prefix_tokens = min_tokens
+        .and_then(|v| v.parse::<u32>().ok())
+        .unwrap_or(128);
+    serde_json::json!({
+        "enabled": enabled && budget_bytes > 0,
+        "budget_bytes": budget_bytes,
+        "min_prefix_tokens": min_prefix_tokens,
+    })
+}
+
+fn run_begin(
+    lib: &Arc<EngineLib>,
+    engine: *mut c_void,
+    writer: &Arc<Mutex<UnixStream>>,
+    streams: &LiveStreams,
+    id: u64,
+    prompt_tokens: Vec<i32>,
+    params_json: String,
+) {
+    // Reset before every begin lives in the ABI, not in policy.
+    unsafe { (lib.turn_reset)(engine) };
+    let pj: serde_json::Value = serde_json::from_str(&params_json).unwrap_or_default();
+    let reasoning_prefix: usize = pj
+        .get("reasoning_prefix")
+        .and_then(|x| x.as_u64())
+        .unwrap_or(0) as usize;
+    // Unknown keys are ignored by the real engine; the key set is append-only.
+    let pace_ms = pj.get("pace_ms").and_then(|x| x.as_u64()).unwrap_or(0);
+    let c_params = match std::ffi::CString::new(params_json.as_str()) {
+        Ok(c) => c,
+        Err(_) => return,
+    };
+    let handle = unsafe {
+        (lib.generate_begin)(
+            engine,
+            prompt_tokens.as_ptr(),
+            prompt_tokens.len() as u32,
+            c_params.as_ptr(),
+        )
+    };
+    if handle.is_null() {
+        let _ = send(
+            writer,
+            &Frame::Error {
+                id: Some(id),
+                e0_status: -1,
+                details_json: engine_err(lib, engine),
+            },
+        );
+        return;
+    }
+    let cancelled = Arc::new(AtomicBool::new(false));
+    streams
+        .lock()
+        .unwrap()
+        .insert(id, (SendPtr(handle), cancelled.clone()));
+    let mut n = 0usize;
+    loop {
+        if cancelled.load(Ordering::SeqCst) {
+            let _ = send(
+                writer,
+                &Frame::Finished {
+                    id,
+                    reason: EndReason::Cancelled,
+                },
+            );
+            break;
+        }
+        let mut out = edge0_core::abi_load::E0TokenOutC {
+            size: std::mem::size_of::<edge0_core::abi_load::E0TokenOutC>() as u32,
+            token: 0,
+            flags: 0,
+        };
+        let st = unsafe { (lib.next)(engine, SendPtr(handle).0, &mut out) };
+        if st < 0 {
+            let _ = send(
+                writer,
+                &Frame::Error {
+                    id: Some(id),
+                    e0_status: st,
+                    details_json: engine_err(lib, engine),
+                },
+            );
+            break;
+        }
+        // This worker is synchronous: any positive status other than NOTIFIED is a protocol violation.
+        if st != edge0_core::abi_load::E0_NOTIFIED {
+            let _ = send(
+                writer,
+                &Frame::Error {
+                    id: Some(id),
+                    e0_status: st,
+                    details_json: format!(
+                        "{{\"message\":\"e0_next returned illegal value {st} (continue values only allow NOTIFIED={}; WOULD_BLOCK is unsupported in this stage)\"}}",
+                        edge0_core::abi_load::E0_NOTIFIED
+                    ),
+                },
+            );
+            break;
+        }
+        if out.flags & 0x1 != 0 {
+            let _ = send(
+                writer,
+                &Frame::Finished {
+                    id,
+                    reason: EndReason::Natural,
+                },
+            );
+            break;
+        }
+        let channel = if n < reasoning_prefix {
+            Channel::Reasoning
+        } else {
+            Channel::Content
+        };
+        n += 1;
+        let _ = send(
+            writer,
+            &Frame::Token {
+                id,
+                token: out.token,
+                channel,
+            },
+        );
+        if pace_ms > 0 {
+            std::thread::sleep(std::time::Duration::from_millis(pace_ms));
+        }
+    }
+    streams.lock().unwrap().remove(&id);
+}
+
+fn engine_err(lib: &EngineLib, engine: *mut c_void) -> String {
+    let mut w = 0usize;
+    let mut buf = vec![0u8; 4096];
+    unsafe { (lib.last_error)(engine, buf.as_mut_ptr() as *mut i8, buf.len(), &mut w) };
+    String::from_utf8_lossy(&buf[..w]).to_string()
+}
+
+fn send(w: &Arc<Mutex<UnixStream>>, f: &Frame) -> Result<(), String> {
+    let mut g = w.lock().unwrap();
+    frames::write_to(&mut *g, f)
+}
+
+fn connect_with_retry(path: &str, timeout: std::time::Duration) -> Result<UnixStream, String> {
+    let start = std::time::Instant::now();
+    loop {
+        match UnixStream::connect(path) {
+            Ok(s) => return Ok(s),
+            Err(e) => {
+                if start.elapsed() >= timeout {
+                    return Err(format!("failed to connect to daemon {path}: {e}"));
+                }
+                std::thread::sleep(std::time::Duration::from_millis(30));
+            }
+        }
+    }
+}
diff --git a/macos/crates/edge0-pull/Cargo.toml b/macos/crates/edge0-pull/Cargo.toml
new file mode 100644
index 0000000..eecd5a3
--- /dev/null
+++ b/macos/crates/edge0-pull/Cargo.toml
@@ -0,0 +1,21 @@
+[package]
+name = "edge0-pull"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+[dependencies]
+edge0-core.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+chrono.workspace = true
+tokio.workspace = true
+reqwest.workspace = true
+futures.workspace = true
+sha2.workspace = true
+ed25519-dalek.workspace = true
+base64.workspace = true
+rand.workspace = true
+libc.workspace = true
+clap.workspace = true
diff --git a/macos/crates/edge0-pull/src/canonical.rs b/macos/crates/edge0-pull/src/canonical.rs
new file mode 100644
index 0000000..70ca959
--- /dev/null
+++ b/macos/crates/edge0-pull/src/canonical.rs
@@ -0,0 +1,32 @@
+//! Canonical JSON bytes for the signed object (keys sorted, no signature fields).
+
+use serde_json::{Map, Value};
+
+pub fn canonical_bytes(manifest: &Value) -> Vec<u8> {
+    let mut v = manifest.clone();
+    if let Some(obj) = v.as_object_mut() {
+        obj.remove("signature");
+        obj.remove("signer_key_id");
+    }
+    let sorted = sort_recursive(v);
+    serde_json::to_vec(&sorted).expect("canonical serialize cannot fail")
+}
+
+fn sort_recursive(v: Value) -> Value {
+    match v {
+        Value::Object(map) => {
+            let mut pairs: Vec<(String, Value)> = map
+                .into_iter()
+                .map(|(k, val)| (k, sort_recursive(val)))
+                .collect();
+            pairs.sort_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()));
+            let mut out = Map::new();
+            for (k, val) in pairs {
+                out.insert(k, val);
+            }
+            Value::Object(out)
+        }
+        Value::Array(arr) => Value::Array(arr.into_iter().map(sort_recursive).collect()),
+        other => other,
+    }
+}
diff --git a/macos/crates/edge0-pull/src/catalog.rs b/macos/crates/edge0-pull/src/catalog.rs
new file mode 100644
index 0000000..dc747a9
--- /dev/null
+++ b/macos/crates/edge0-pull/src/catalog.rs
@@ -0,0 +1,239 @@
+//! Catalog view from a signed manifest, plus direct-mode Hugging Face probe.
+
+use std::collections::BTreeMap;
+
+use edge0_core::config::PullMode;
+use edge0_core::manifest::Manifest;
+use edge0_core::protocol::CatalogTier;
+use edge0_core::state::{DirectFile, DirectSnapshot, InstalledTier};
+
+use crate::error::PullError;
+use crate::mirror;
+use crate::mirror::CheckedList;
+
+pub struct TierInfo {
+    pub tier: &'static str,
+    pub repo: &'static str,
+    pub approx_bytes: u64,
+    pub license: &'static str,
+    pub notes_url: &'static str,
+}
+
+pub const BUILTIN_TIERS: &[TierInfo] = &[
+    TierInfo {
+        tier: "edge0-8b",
+        repo: "Edge0/Edge0-8B-A1B-preview",
+        approx_bytes: 4_600_000_000,
+        license: "Apache-2.0",
+        notes_url: "https://huggingface.co/Edge0/Edge0-8B-A1B-preview",
+    },
+    TierInfo {
+        tier: "edge0-35b",
+        repo: "Edge0/Edge0-35B-A3B-preview",
+        approx_bytes: 19_731_249_723,
+        license: "Apache-2.0",
+        notes_url: "https://huggingface.co/Edge0/Edge0-35B-A3B-preview",
+    },
+];
+
+pub fn builtin_tier(tier: &str) -> Option<&'static TierInfo> {
+    BUILTIN_TIERS.iter().find(|t| t.tier == tier)
+}
+
+pub async fn resolve_files(
+    client: &reqwest::Client,
+    base: &str,
+    tier: &str,
+) -> Result<DirectSnapshot, PullError> {
+    let info = builtin_tier(tier).ok_or_else(|| {
+        PullError::new(
+            "E-MODEL-MISSING",
+            format!(
+                "The direct tier table has no {tier} (available: {})",
+                BUILTIN_TIERS
+                    .iter()
+                    .map(|t| t.tier)
+                    .collect::<Vec<_>>()
+                    .join(", ")
+            ),
+        )
+    })?;
+    let url = format!("{}/api/models/{}?blobs=true", base, info.repo);
+    let resp = client
+        .get(&url)
+        .timeout(std::time::Duration::from_secs(20))
+        .send()
+        .await
+        .and_then(|r| r.error_for_status())
+        .map_err(|e| PullError::new("E-DL-NET", format!("{url} unreachable: {e}")))?;
+    let body: serde_json::Value = resp
+        .json()
+        .await
+        .map_err(|e| PullError::new("E-DL-NET", format!("{url} parse failed: {e}")))?;
+    let commit = body
+        .get("sha")
+        .and_then(|v| v.as_str())
+        .filter(|s| !s.is_empty())
+        .unwrap_or("main")
+        .to_string();
+    let mut files: Vec<DirectFile> = Vec::new();
+    for s in body
+        .get("siblings")
+        .and_then(|v| v.as_array())
+        .cloned()
+        .unwrap_or_default()
+    {
+        let path = match s.get("rfilename").and_then(|v| v.as_str()) {
+            Some(p) if !p.is_empty() => p.to_string(),
+            _ => continue,
+        };
+        if path.starts_with('.') {
+            continue;
+        }
+        files.push(DirectFile {
+            path,
+            size: s.get("size").and_then(|v| v.as_u64()).unwrap_or(0),
+            sha256: None,
+            probe_blocks: None,
+        });
+    }
+    files.sort_by(|a, b| a.path.cmp(&b.path));
+    if files.is_empty() {
+        return Err(PullError::new(
+            "E-DL-NET",
+            format!("{url} returned no usable files (siblings is empty)"),
+        ));
+    }
+    Ok(DirectSnapshot {
+        mode: "direct".to_string(),
+        source: "huggingface".to_string(),
+        base_url: base.trim_end_matches('/').to_string(),
+        repo: info.repo.to_string(),
+        commit,
+        files,
+        recorded_at: None,
+    })
+}
+
+pub enum CreationPlan {
+    Registry(CheckedList, String),
+    Direct(DirectSnapshot),
+}
+
+impl CreationPlan {
+    pub fn rev(&self) -> &str {
+        match self {
+            CreationPlan::Registry(l, _) => &l.rev,
+            CreationPlan::Direct(s) => &s.commit,
+        }
+    }
+    pub fn bytes_total(&self) -> u64 {
+        match self {
+            CreationPlan::Registry(l, tier) => l
+                .manifest
+                .find_tier(tier)
+                .map(|t| t.files.iter().map(|f| f.size).sum())
+                .unwrap_or(0),
+            CreationPlan::Direct(s) => s.bytes_total(),
+        }
+    }
+    pub fn source_reason_direct(base: &str, pin_ignored: bool) -> String {
+        if pin_ignored {
+            format!("Direct mode uses the fixed Hugging Face source (base={base}); request source was ignored")
+        } else {
+            format!("Direct mode uses the fixed Hugging Face source (base={base})")
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+pub async fn plan_for_tier(
+    mode: PullMode,
+    client: &reqwest::Client,
+    urls: &[String],
+    hf_base: &str,
+    tier: &str,
+    home: &edge0_core::paths::Home,
+    now: chrono::DateTime<chrono::Utc>,
+    cli_version: &str,
+) -> Result<CreationPlan, PullError> {
+    match mode {
+        PullMode::Registry => mirror::fetch_manifest(client, urls, tier, home, now, cli_version)
+            .await
+            .map(|l| CreationPlan::Registry(l, tier.to_string())),
+        PullMode::Direct => resolve_files(client, hf_base, tier)
+            .await
+            .map(CreationPlan::Direct),
+    }
+}
+
+pub fn direct_tiers(
+    snaps: &[(String, DirectSnapshot)],
+    installed: &[InstalledTier],
+) -> Vec<CatalogTier> {
+    snaps
+        .iter()
+        .map(|(tier, s)| {
+            let info = builtin_tier(tier);
+            let inst = installed.iter().find(|i| i.tier == *tier);
+            CatalogTier {
+                tier: tier.clone(),
+                rev: s.commit.clone(),
+                min_cli_version: "0.1.0".to_string(),
+                bytes_total: s.bytes_total(),
+                files_count: s.files.len() as u64,
+                license: info.map(|i| i.license.to_string()),
+                notes_url: info.map(|i| i.notes_url.to_string()),
+                primary_source: "huggingface".to_string(),
+                has_modelscope: false,
+                has_huggingface: true,
+                installed: inst.map(|i| i.rev.clone()),
+                update_available: inst.is_some_and(|i| i.rev != s.commit),
+            }
+        })
+        .collect()
+}
+
+pub fn catalog_tiers(manifest: &Manifest, installed: &[InstalledTier]) -> Vec<CatalogTier> {
+    manifest
+        .tiers
+        .iter()
+        .map(|t| {
+            let mut ms = 0u64;
+            let mut hf = 0u64;
+            let mut primary: BTreeMap<&str, u64> = BTreeMap::new();
+            for f in &t.files {
+                if let Some(s) = &f.sources {
+                    if s.modelscope.is_some() {
+                        ms += 1;
+                    }
+                    if s.huggingface.is_some() {
+                        hf += 1;
+                    }
+                    *primary.entry(s.primary.as_str()).or_default() += 1;
+                }
+            }
+            let primary_source = ["modelscope", "huggingface"]
+                .iter()
+                .copied()
+                .max_by_key(|k| (primary.get(*k).copied().unwrap_or(0), *k == "modelscope"))
+                .unwrap_or("modelscope")
+                .to_string();
+            let inst = installed.iter().find(|i| i.tier == t.tier);
+            CatalogTier {
+                tier: t.tier.clone(),
+                rev: t.rev.clone(),
+                min_cli_version: t.min_cli_version.clone(),
+                bytes_total: t.files.iter().map(|f| f.size).sum(),
+                files_count: t.files.len() as u64,
+                license: t.license.clone(),
+                notes_url: t.notes_url.clone(),
+                primary_source,
+                has_modelscope: ms > 0,
+                has_huggingface: hf > 0,
+                installed: inst.map(|i| i.rev.clone()),
+                update_available: inst.is_some_and(|i| i.rev != t.rev),
+            }
+        })
+        .collect()
+}
diff --git a/macos/crates/edge0-pull/src/check.rs b/macos/crates/edge0-pull/src/check.rs
new file mode 100644
index 0000000..0fd7d34
--- /dev/null
+++ b/macos/crates/edge0-pull/src/check.rs
@@ -0,0 +1,206 @@
+//! Install-time gates (signature, expiry, min CLI) and load-time spot checks.
+
+use chrono::{DateTime, Utc};
+
+use edge0_core::manifest::Manifest;
+use edge0_core::manifest::{parse_semver, TierEntry};
+
+use crate::error::PullError;
+use crate::trust;
+
+pub fn check_manifest<'a>(
+    manifest: &'a Manifest,
+    tier: &str,
+    now: DateTime<Utc>,
+    cli_version: &str,
+) -> Result<(&'a TierEntry, &'static str), PullError> {
+    manifest
+        .validate_installable()
+        .map_err(|e| PullError::new("E-MANIFEST-SIG", e))?;
+    let key_id = trust::verify(manifest)?;
+    let entry = manifest.find_tier(tier).ok_or_else(|| {
+        PullError::new(
+            "E-MODEL-MISSING",
+            format!(
+                "Manifest does not contain tier {tier} (available: {})",
+                manifest
+                    .tiers
+                    .iter()
+                    .map(|t| t.tier.as_str())
+                    .collect::<Vec<_>>()
+                    .join(", ")
+            ),
+        )
+    })?;
+    let expires = DateTime::parse_from_rfc3339(&manifest.expires_at)
+        .map(|d| d.with_timezone(&Utc))
+        .map_err(|e| {
+            PullError::new(
+                "E-MANIFEST-SIG",
+                format!("expires_at cannot be parsed: {e}"),
+            )
+        })?;
+    if now >= expires {
+        return Err(PullError::new(
+            "E-MANIFEST-EXPIRED",
+            format!(
+                "Manifest expired at {}; new tasks need a fresh manifest (in-flight downloads are unaffected)",
+                expires.to_rfc3339()
+            ),
+        ));
+    }
+    let min = parse_semver(&entry.min_cli_version).unwrap_or((0, 0, 0));
+    let cur = parse_semver(cli_version).unwrap_or((0, 0, 0));
+    if cur < min {
+        return Err(PullError::new(
+            "E-MANIFEST-MINVER",
+            format!(
+                "Tier {}@{} requires CLI >= {}, current version is {cli_version}; upgrade the app first",
+                entry.tier, entry.rev, entry.min_cli_version
+            ),
+        ));
+    }
+    Ok((entry, key_id))
+}
+
+use edge0_core::state::{BlockProbe, DirectSnapshot};
+
+fn sha256_bytes(b: &[u8]) -> String {
+    use sha2::Digest as _;
+    let mut h = sha2::Sha256::new();
+    sha2::Digest::update(&mut h, b);
+    sha2::Digest::finalize(h)
+        .iter()
+        .map(|x| format!("{x:02x}"))
+        .collect()
+}
+
+pub fn record_probe_blocks(
+    snap: &DirectSnapshot,
+    rev_dir: &std::path::Path,
+) -> Result<DirectSnapshot, PullError> {
+    use std::io::{Read, Seek, SeekFrom};
+    let mut out = snap.clone();
+    for f in &mut out.files {
+        let positions = edge0_core::state::probe_positions(f.size);
+        if positions.is_empty() {
+            continue;
+        }
+        let path = rev_dir.join("files").join(&f.path);
+        let mut file = std::fs::File::open(&path).map_err(|e| {
+            PullError::new(
+                "E-MODEL-INVALID",
+                format!("open {} failed: {e}", path.display()),
+            )
+        })?;
+        let mut blocks = Vec::new();
+        for (off, len) in positions {
+            file.seek(SeekFrom::Start(off))
+                .map_err(|e| PullError::new("E-MODEL-INVALID", format!("seek {off} failed: {e}")))?;
+            let mut buf = vec![0u8; len as usize];
+            file.read_exact(&mut buf).map_err(|e| {
+                PullError::new(
+                    "E-MODEL-INVALID",
+                    format!("read block {}@{off} failed: {e}", f.path),
+                )
+            })?;
+            blocks.push(BlockProbe {
+                offset: off,
+                len,
+                sha256: sha256_bytes(&buf),
+            });
+        }
+        f.probe_blocks = Some(blocks);
+    }
+    Ok(out)
+}
+
+pub fn spot_check(snap: &DirectSnapshot, rev_dir: &std::path::Path) -> Result<f64, PullError> {
+    use std::time::Instant;
+    let t0 = Instant::now();
+    use std::io::{Read, Seek, SeekFrom};
+    for f in &snap.files {
+        let path = rev_dir.join("files").join(&f.path);
+        let meta = std::fs::metadata(&path).map_err(|e| {
+            PullError::new("E-MODEL-INVALID", format!("{} unreadable: {e}", path.display()))
+        })?;
+        if meta.len() != f.size {
+            return Err(PullError::new(
+                "E-MODEL-INVALID",
+                format!("{} size drift: recorded {}, actual {}", f.path, f.size, meta.len()),
+            ));
+        }
+        if f.size <= edge0_core::state::SPOT_SMALL_MAX {
+            let Some(want) = &f.sha256 else { continue };
+            let mut file = std::fs::File::open(&path).map_err(|e| {
+                PullError::new("E-MODEL-INVALID", format!("open {} failed: {e}", f.path))
+            })?;
+            let mut h = {
+                use sha2::Digest as _;
+                sha2::Sha256::new()
+            };
+            let mut buf = vec![0u8; 1 << 20];
+            loop {
+                let n = file.read(&mut buf).map_err(|e| {
+                    PullError::new("E-MODEL-INVALID", format!("read {} failed: {e}", f.path))
+                })?;
+                if n == 0 {
+                    break;
+                }
+                use sha2::Digest as _;
+                h.update(&buf[..n]);
+            }
+            let got: String = sha2::Digest::finalize(h)
+                .iter()
+                .map(|x| format!("{x:02x}"))
+                .collect();
+            if &got != want {
+                return Err(PullError::new(
+                    "E-MODEL-INVALID",
+                    format!(
+                        "{} hash mismatch: expected {}…, got {}…",
+                        f.path,
+                        &want[..12],
+                        &got[..12]
+                    ),
+                ));
+            }
+            continue;
+        }
+        let Some(blocks) = &f.probe_blocks else {
+            continue;
+        };
+        let mut file = std::fs::File::open(&path)
+            .map_err(|e| PullError::new("E-MODEL-INVALID", format!("open {} failed: {e}", f.path)))?;
+        for bp in blocks {
+            file.seek(SeekFrom::Start(bp.offset)).map_err(|e| {
+                PullError::new(
+                    "E-MODEL-INVALID",
+                    format!("seek {}@{} failed: {e}", f.path, bp.offset),
+                )
+            })?;
+            let mut buf = vec![0u8; bp.len as usize];
+            file.read_exact(&mut buf).map_err(|e| {
+                PullError::new(
+                    "E-MODEL-INVALID",
+                    format!("read block {}@{} failed: {e}", f.path, bp.offset),
+                )
+            })?;
+            let got = sha256_bytes(&buf);
+            if got != bp.sha256 {
+                return Err(PullError::new(
+                    "E-MODEL-INVALID",
+                    format!(
+                        "{} probe block @{}/{} hash mismatch: expected {}…, got {}…",
+                        f.path,
+                        bp.offset,
+                        bp.len,
+                        &bp.sha256[..12],
+                        &got[..12]
+                    ),
+                ));
+            }
+        }
+    }
+    Ok(t0.elapsed().as_secs_f64())
+}
diff --git a/macos/crates/edge0-pull/src/error.rs b/macos/crates/edge0-pull/src/error.rs
new file mode 100644
index 0000000..33a20cb
--- /dev/null
+++ b/macos/crates/edge0-pull/src/error.rs
@@ -0,0 +1,30 @@
+//! Pull error with a table code id, mapped to `ApiError` on the daemon.
+
+use std::fmt;
+
+#[derive(Debug, Clone)]
+pub struct PullError {
+    pub code: &'static str,
+    pub message: String,
+}
+
+impl PullError {
+    pub fn new(code: &'static str, message: impl Into<String>) -> Self {
+        Self {
+            code,
+            message: message.into(),
+        }
+    }
+    pub fn into_api_error(self) -> edge0_core::ApiError {
+        let code = edge0_core::errorcodes::lookup(self.code).expect("code must exist in the table");
+        edge0_core::ApiError::new(code, self.message)
+    }
+}
+
+impl fmt::Display for PullError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "[{}] {}", self.code, self.message)
+    }
+}
+
+impl std::error::Error for PullError {}
diff --git a/macos/crates/edge0-pull/src/fetch.rs b/macos/crates/edge0-pull/src/fetch.rs
new file mode 100644
index 0000000..f8882dc
--- /dev/null
+++ b/macos/crates/edge0-pull/src/fetch.rs
@@ -0,0 +1,132 @@
+//! HTTP range fetch used by the download runner.
+
+use std::pin::Pin;
+use std::time::Duration;
+
+use futures::future::BoxFuture;
+use futures::{Stream, StreamExt};
+
+use crate::error::PullError;
+
+pub type ByteStream = Pin<Box<dyn Stream<Item = Result<Vec<u8>, PullError>> + Send + 'static>>;
+
+#[derive(Debug, Clone, Copy)]
+pub struct HeadInfo {
+    pub ok: bool,
+    pub rtt: Duration,
+}
+
+pub trait Fetcher: Send + Sync {
+    fn head(&self, url: String) -> BoxFuture<'static, HeadInfo>;
+    fn open_range(
+        &self,
+        url: String,
+        from: u64,
+    ) -> BoxFuture<'static, Result<ByteStream, PullError>>;
+}
+
+pub struct HttpFetcher {
+    client: reqwest::Client,
+}
+
+impl Default for HttpFetcher {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl HttpFetcher {
+    pub fn new() -> Self {
+        Self {
+            client: reqwest::Client::builder()
+                .no_proxy()
+                .timeout(Duration::from_secs(300))
+                .build()
+                .expect("reqwest client"),
+        }
+    }
+    pub fn with_env_proxy() -> Self {
+        Self {
+            client: reqwest::Client::builder()
+                .timeout(Duration::from_secs(300))
+                .build()
+                .expect("reqwest client"),
+        }
+    }
+}
+
+fn map_err(context: &str, e: &reqwest::Error) -> PullError {
+    if let Some(s) = e.status() {
+        if s.as_u16() == 416 {
+            return PullError::new(
+                "E-DL-SRC",
+                format!("{context}: range is invalid (HTTP 416)"),
+            );
+        }
+        return PullError::new("E-DL-SRC", format!("{context}: source returned HTTP {s}"));
+    }
+    PullError::new("E-DL-NET", format!("{context}: {e}"))
+}
+
+impl Fetcher for HttpFetcher {
+    fn head(&self, url: String) -> BoxFuture<'static, HeadInfo> {
+        let client = self.client.clone();
+        Box::pin(async move {
+            let t0 = std::time::Instant::now();
+            let r = client
+                .head(&url)
+                .timeout(Duration::from_secs(10))
+                .send()
+                .await
+                .and_then(|r| r.error_for_status());
+            HeadInfo {
+                ok: r.is_ok(),
+                rtt: t0.elapsed(),
+            }
+        })
+    }
+
+    fn open_range(
+        &self,
+        url: String,
+        from: u64,
+    ) -> BoxFuture<'static, Result<ByteStream, PullError>> {
+        let client = self.client.clone();
+        Box::pin(async move {
+            let mut req = client.get(&url);
+            if from > 0 {
+                req = req.header("Range", format!("bytes={from}-"));
+            }
+            let resp = req
+                .send()
+                .await
+                .map_err(|e| map_err("failed to open source stream", &e))?;
+            let status = resp.status();
+            if from > 0 && status.as_u16() == 200 {
+                return Err(PullError::new(
+                    "E-DL-SRC",
+                    format!("Source does not support Range ({from}- request returned the full body); switch sources or retry later"),
+                ));
+            }
+            if !status.is_success() {
+                let e_status = status.as_u16();
+                let code = if matches!(e_status, 403 | 429) || (500..=599).contains(&e_status) {
+                    "E-DL-SRC"
+                } else {
+                    "E-DL-NET"
+                };
+                return Err(PullError::new(
+                    code,
+                    format!(
+                        "Source returned HTTP {e_status}; progress remains at the verified range"
+                    ),
+                ));
+            }
+            let stream = resp.bytes_stream().map(|b| {
+                b.map(|chunk| chunk.to_vec())
+                    .map_err(|e| map_err("stream interrupted", &e))
+            });
+            Ok(Box::pin(stream) as ByteStream)
+        })
+    }
+}
diff --git a/macos/crates/edge0-pull/src/lib.rs b/macos/crates/edge0-pull/src/lib.rs
new file mode 100644
index 0000000..9ed07c4
--- /dev/null
+++ b/macos/crates/edge0-pull/src/lib.rs
@@ -0,0 +1,14 @@
+//! Signed-manifest verify/fetch, download state machine, and registry operations.
+
+pub mod canonical;
+pub mod catalog;
+pub mod check;
+pub mod error;
+pub mod fetch;
+pub mod mirror;
+pub mod registry_ops;
+pub mod runner;
+pub mod state_machine;
+pub mod trust;
+
+pub use error::PullError;
diff --git a/macos/crates/edge0-pull/src/mirror.rs b/macos/crates/edge0-pull/src/mirror.rs
new file mode 100644
index 0000000..e460ffc
--- /dev/null
+++ b/macos/crates/edge0-pull/src/mirror.rs
@@ -0,0 +1,241 @@
+//! Fetch a signed manifest from `EDGE0_REGISTRY_URL` mirrors, then verify locally.
+
+use std::time::Duration;
+
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+
+use edge0_core::manifest::Manifest;
+use edge0_core::paths::Home;
+use edge0_core::state::atomic_write_json;
+
+use crate::check::check_manifest;
+use crate::error::PullError;
+
+#[derive(Debug, Clone)]
+pub struct CheckedList {
+    pub manifest: Manifest,
+    pub rev: String,
+    pub signer_key_id: String,
+    pub source_url: String,
+}
+
+pub async fn fetch_manifest(
+    client: &reqwest::Client,
+    urls: &[String],
+    tier: &str,
+    home: &Home,
+    now: DateTime<Utc>,
+    cli_version: &str,
+) -> Result<CheckedList, PullError> {
+    if urls.is_empty() {
+        return Err(PullError::new(
+            "E-MANIFEST-EXPIRED",
+            "EDGE0_REGISTRY_URL is not configured; no manifest source is available",
+        ));
+    }
+    let mut last: Option<PullError> = None;
+    let mut any_network = 0usize;
+    for url in urls {
+        match client
+            .get(url)
+            .timeout(Duration::from_secs(20))
+            .send()
+            .await
+            .and_then(|r| r.error_for_status())
+        {
+            Err(e) => {
+                any_network += 1;
+                last = Some(PullError::new("E-DL-NET", format!("{url} unreachable: {e}")));
+                continue;
+            }
+            Ok(resp) => match resp.text().await {
+                Err(e) => {
+                    any_network += 1;
+                    last = Some(PullError::new("E-DL-NET", format!("{url} read failed: {e}")));
+                    continue;
+                }
+                Ok(text) => {
+                    let parsed = Manifest::from_json_str(&text)
+                        .map_err(|e| PullError::new("E-MANIFEST-SIG", e));
+                    let checked =
+                        parsed.and_then(|m| match check_manifest(&m, tier, now, cli_version) {
+                            Ok((entry, key_id)) => {
+                                let (rev, kid) = (entry.rev.clone(), key_id.to_string());
+                                Ok((m, rev, kid))
+                            }
+                            Err(e) => Err(e),
+                        });
+                    let (manifest, rev, signer_key_id) = match checked {
+                        Ok(v) => v,
+                        Err(e) => {
+                            last = Some(e);
+                            continue;
+                        }
+                    };
+                    cache(
+                        home,
+                        tier,
+                        &rev,
+                        &text,
+                        &Verified {
+                            verified_at: now,
+                            signer_key_id: signer_key_id.clone(),
+                            source_url: url.to_string(),
+                            expires_at: manifest.expires_at.clone(),
+                        },
+                    );
+                    return Ok(CheckedList {
+                        manifest,
+                        rev,
+                        signer_key_id,
+                        source_url: url.clone(),
+                    });
+                }
+            },
+        }
+    }
+    let mut e = last.unwrap_or_else(|| PullError::new("E-DL-NET", "all mirrors unreachable"));
+    if any_network == urls.len() {
+        e.code = "E-DL-NET";
+    }
+    if e.code == "E-MANIFEST-SIG" && any_network < urls.len() {
+        e.message = format!(
+            "No mirror manifest passed local signature or policy checks: {e}",
+            e = e.message
+        );
+    }
+    Err(e)
+}
+
+#[derive(Serialize, Deserialize)]
+struct Verified {
+    verified_at: DateTime<Utc>,
+    signer_key_id: String,
+    source_url: String,
+    expires_at: String,
+}
+
+fn cache(home: &Home, tier: &str, rev: &str, raw: &str, v: &Verified) {
+    let dir = home.manifests();
+    let _ = std::fs::create_dir_all(dir.join(".verified"));
+    let p = dir.join(format!("{tier}-{rev}.json"));
+    let tmp = dir.join(format!(".{tier}-{rev}.json.tmp"));
+    if std::fs::write(&tmp, raw.as_bytes()).is_ok() {
+        let _ = std::fs::rename(&tmp, &p);
+    }
+    let _ = atomic_write_json(&dir.join(".verified").join(format!("{tier}-{rev}.json")), v);
+}
+
+#[derive(Debug, Clone)]
+pub struct CatalogFetch {
+    pub manifest: Manifest,
+    pub fetched_at: DateTime<Utc>,
+    pub signer_key_id: String,
+    pub source_url: String,
+}
+
+pub async fn fetch_catalog(
+    client: &reqwest::Client,
+    urls: &[String],
+    home: &Home,
+    now: DateTime<Utc>,
+) -> Result<CatalogFetch, PullError> {
+    if urls.is_empty() {
+        return Err(PullError::new(
+            "E-DL-NET",
+            "EDGE0_REGISTRY_URL is not configured; no catalog source is available",
+        ));
+    }
+    let mut last: Option<PullError> = None;
+    let mut any_network = 0usize;
+    for url in urls {
+        match client
+            .get(url)
+            .timeout(Duration::from_secs(20))
+            .send()
+            .await
+            .and_then(|r| r.error_for_status())
+        {
+            Err(e) => {
+                any_network += 1;
+                last = Some(PullError::new("E-DL-NET", format!("{url} unreachable: {e}")));
+                continue;
+            }
+            Ok(resp) => match resp.text().await {
+                Err(e) => {
+                    any_network += 1;
+                    last = Some(PullError::new("E-DL-NET", format!("{url} read failed: {e}")));
+                    continue;
+                }
+                Ok(text) => {
+                    let checked: Result<&'static str, PullError> = (|| {
+                        let m = Manifest::from_json_str(&text)
+                            .map_err(|e| PullError::new("E-MANIFEST-SIG", e))?;
+                        m.validate_installable()
+                            .map_err(|e| PullError::new("E-MANIFEST-SIG", e))?;
+                        crate::trust::verify(&m)
+                    })();
+                    match checked {
+                        Ok(kid) => {
+                            let manifest = Manifest::from_json_str(&text)
+                                .map_err(|e| PullError::new("E-MANIFEST-SIG", e))?;
+                            cache(
+                                home,
+                                "catalog",
+                                "latest",
+                                &text,
+                                &Verified {
+                                    verified_at: now,
+                                    signer_key_id: kid.to_string(),
+                                    source_url: url.to_string(),
+                                    expires_at: manifest.expires_at.clone(),
+                                },
+                            );
+                            return Ok(CatalogFetch {
+                                manifest,
+                                fetched_at: now,
+                                signer_key_id: kid.to_string(),
+                                source_url: url.to_string(),
+                            });
+                        }
+                        Err(e) => {
+                            last = Some(e);
+                            continue;
+                        }
+                    }
+                }
+            },
+        }
+    }
+    let mut e = last.unwrap_or_else(|| PullError::new("E-DL-NET", "all mirrors unreachable"));
+    if any_network == urls.len() {
+        e.code = "E-DL-NET";
+    }
+    if e.code == "E-MANIFEST-SIG" && any_network < urls.len() {
+        e.message = format!(
+            "No mirror catalog passed local signature checks: {e}",
+            e = e.message
+        );
+    }
+    Err(e)
+}
+
+pub fn read_catalog_cache(home: &Home) -> Option<CatalogFetch> {
+    use crate::trust;
+    let dir = home.manifests();
+    let raw = std::fs::read_to_string(dir.join("catalog-latest.json")).ok()?;
+    let vtext = std::fs::read_to_string(dir.join(".verified").join("catalog-latest.json")).ok()?;
+    let rec: Verified = serde_json::from_str(&vtext).ok()?;
+    let manifest = Manifest::from_json_str(&raw).ok()?;
+    let kid = trust::verify(&manifest).ok()?;
+    if kid != rec.signer_key_id {
+        return None;
+    }
+    Some(CatalogFetch {
+        manifest,
+        fetched_at: rec.verified_at,
+        signer_key_id: rec.signer_key_id,
+        source_url: rec.source_url,
+    })
+}
diff --git a/macos/crates/edge0-pull/src/registry_ops.rs b/macos/crates/edge0-pull/src/registry_ops.rs
new file mode 100644
index 0000000..eeaf79e
--- /dev/null
+++ b/macos/crates/edge0-pull/src/registry_ops.rs
@@ -0,0 +1,408 @@
+//! Blob commit, `files/` hard links, refcounts, rm, and gc.
+
+use std::collections::BTreeMap;
+use std::os::unix::fs::{MetadataExt, PermissionsExt};
+use std::path::{Path, PathBuf};
+
+use serde::{Deserialize, Serialize};
+
+use edge0_core::manifest::Manifest;
+use edge0_core::paths::Home;
+use edge0_core::state::atomic_write_json;
+
+use crate::error::PullError;
+
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
+pub struct BlobRefs(pub BTreeMap<String, u64>);
+
+impl BlobRefs {
+    pub fn load(home: &Home) -> Self {
+        std::fs::read_to_string(home.blobrefs())
+            .ok()
+            .and_then(|t| serde_json::from_str(&t).ok())
+            .unwrap_or_default()
+    }
+    pub fn save(&self, home: &Home) -> Result<(), PullError> {
+        atomic_write_json(&home.blobrefs(), self).map_err(|e| PullError::new("E-SRV-CONFLICT", e))
+    }
+    pub fn refcount(&self, sha: &str) -> u64 {
+        self.0.get(sha).copied().unwrap_or(0)
+    }
+}
+
+pub fn blob_path(home: &Home, sha: &str) -> PathBuf {
+    home.blobs().join(sha)
+}
+
+pub fn same_volume(home: &Home) -> bool {
+    match (
+        std::fs::metadata(home.tmp()),
+        std::fs::metadata(home.blobs()),
+    ) {
+        (Ok(a), Ok(b)) => a.dev() == b.dev(),
+        _ => false,
+    }
+}
+
+pub fn disk_avail(path: &Path) -> Option<u64> {
+    let c = std::ffi::CString::new(path.as_os_str().as_encoded_bytes().to_vec()).ok()?;
+    let mut st = std::mem::MaybeUninit::<libc::statvfs>::zeroed();
+    if unsafe { libc::statvfs(c.as_ptr(), st.as_mut_ptr()) } != 0 {
+        return None;
+    }
+    let st = unsafe { st.assume_init() };
+    Some((st.f_bavail as u64) * st.f_frsize)
+}
+
+pub fn preflight_need(total: u64, tmp_same_volume: bool) -> u64 {
+    let base = total * 12 / 10 + (2u64 << 30);
+    if tmp_same_volume {
+        base
+    } else {
+        base + total
+    }
+}
+
+pub fn preflight_check(home: &Home, total: u64) -> Result<(), PullError> {
+    let avail = std::env::var("EDGE0_TEST_DISK_AVAIL")
+        .ok()
+        .and_then(|v| v.parse::<u64>().ok())
+        .or_else(|| disk_avail(&home.root))
+        .unwrap_or(u64::MAX);
+    let need = preflight_need(total, same_volume(home));
+    if avail < need {
+        return Err(PullError::new(
+            "E-DL-DISK",
+            format!(
+            "Tier size is {total} B (preflight requires >= {need} B), but only {avail} B is free. Free disk space or move EDGE0_HOME; run edge0 gc to reclaim storage",
+            ),
+        ));
+    }
+    Ok(())
+}
+
+pub fn stage_blob(home: &Home, part: &Path, sha: &str) -> Result<bool, PullError> {
+    std::fs::create_dir_all(home.blobs())
+        .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+    let dst = blob_path(home, sha);
+    if dst.exists() {
+        let _ = std::fs::remove_file(part);
+        return Ok(true);
+    }
+    if same_volume(home) {
+        std::fs::rename(part, &dst).map_err(|e| {
+            PullError::new(
+                "E-SRV-CONFLICT",
+                format!("Failed to commit blob ({}): {e}", dst.display()),
+            )
+        })?;
+    } else {
+        eprintln!(
+            "Warning: tmp and blobs are on different volumes; falling back to copy (use a single-volume EDGE0_HOME for atomic moves)"
+        );
+        std::fs::copy(part, &dst).map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+        let _ = std::fs::remove_file(part);
+    }
+    let _ = std::fs::set_permissions(&dst, std::fs::Permissions::from_mode(0o444));
+    Ok(false)
+}
+
+pub fn link_into_model(
+    home: &Home,
+    tier: &str,
+    rev: &str,
+    rel: &str,
+    sha: &str,
+) -> Result<(), PullError> {
+    let dst = home.model_dir(tier, rev).join("files").join(rel);
+    if dst.exists() {
+        return Ok(());
+    }
+    if let Some(p) = dst.parent() {
+        std::fs::create_dir_all(p).map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+    }
+    let src = blob_path(home, sha);
+    match std::fs::hard_link(&src, &dst) {
+        Ok(()) => {}
+        Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {}
+        Err(_) => {
+            eprintln!(
+                "Warning: hard links unavailable; falling back to copy ({})",
+                dst.display()
+            );
+            std::fs::copy(&src, &dst)
+                .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+        }
+    }
+    let _ = std::fs::set_permissions(&dst, std::fs::Permissions::from_mode(0o444));
+    Ok(())
+}
+
+pub fn write_manifest_copy(
+    home: &Home,
+    tier: &str,
+    rev: &str,
+    manifest: &Manifest,
+) -> Result<bool, PullError> {
+    let dst = home.model_dir(tier, rev).join("manifest.json");
+    let fresh = !dst.exists();
+    if fresh {
+        atomic_write_json(&dst, manifest).map_err(|e| PullError::new("E-SRV-CONFLICT", e))?;
+    }
+    Ok(fresh)
+}
+
+pub fn remove_tier(home: &Home, tier: &str) -> Result<Vec<String>, PullError> {
+    let dir = home.models().join(tier);
+    if !dir.is_dir() {
+        return Err(PullError::new(
+            "E-MODEL-MISSING",
+            format!("Tier {tier} is not installed"),
+        ));
+    }
+    let mut removed_shas = Vec::new();
+    let mut refs = BlobRefs::load(home);
+    for rev_ent in
+        std::fs::read_dir(&dir).map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?
+    {
+        let rev_dir = rev_ent
+            .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?
+            .path();
+        if !rev_dir.is_dir() {
+            continue;
+        }
+        if let Ok(m) = edge0_core::manifest::load_copy(&rev_dir.join("manifest.json")) {
+            if let Some(te) = m.find_tier(tier) {
+                for f in &te.files {
+                    let e = refs.0.entry(f.sha256.clone()).or_insert(0);
+                    *e = e.saturating_sub(1);
+                    removed_shas.push(f.sha256.clone());
+                }
+            }
+        }
+        let _ = std::fs::remove_dir_all(&rev_dir);
+    }
+    let _ = std::fs::remove_dir(&dir);
+    refs.save(home)?;
+    Ok(removed_shas)
+}
+
+#[derive(Debug, Default, Serialize)]
+pub struct GcReport {
+    pub deleted: Vec<String>,
+    pub kept_referenced: usize,
+    pub kept_inflight: Vec<String>,
+}
+
+pub fn gc(home: &Home) -> Result<GcReport, PullError> {
+    let refs = BlobRefs::load(home);
+    let mut inflight: std::collections::HashSet<String> = std::collections::HashSet::new();
+    let dl = home.state().join("downloads");
+    if let Ok(rd) = std::fs::read_dir(&dl) {
+        for ent in rd.flatten() {
+            if let Some(t) = std::fs::read_to_string(ent.path())
+                .ok()
+                .and_then(|s| serde_json::from_str::<edge0_core::state::DownloadTask>(&s).ok())
+            {
+                for f in &t.files {
+                    inflight.insert(f.sha256.clone());
+                }
+            }
+        }
+    }
+    let mut report = GcReport::default();
+    if let Ok(rd) = std::fs::read_dir(home.blobs()) {
+        for ent in rd.flatten() {
+            let p = ent.path();
+            let Some(sha) = p.file_name().map(|s| s.to_string_lossy().to_string()) else {
+                continue;
+            };
+            if sha.len() != 64 {
+                continue;
+            }
+            if inflight.contains(&sha) {
+                report.kept_inflight.push(sha);
+            } else if refs.refcount(&sha) > 0 {
+                report.kept_referenced += 1;
+            } else {
+                if std::fs::remove_file(&p).is_ok() {
+                    report.deleted.push(sha);
+                }
+            }
+        }
+    }
+    let mut refs2 = refs.clone();
+    for sha in &report.deleted {
+        refs2.0.remove(sha);
+    }
+    for (sha, c) in refs2.0.clone() {
+        if c == 0 {
+            refs2.0.remove(&sha);
+        }
+    }
+    refs2.save(home)?;
+    Ok(report)
+}
+
+pub fn staging_dir(home: &Home, tier: &str, rev: &str) -> PathBuf {
+    home.models().join(tier).join(format!("{rev}.staging"))
+}
+
+pub fn staging_file(home: &Home, tier: &str, rev: &str, path: &str) -> PathBuf {
+    staging_dir(home, tier, rev).join("files").join(path)
+}
+
+pub fn commit_direct(
+    home: &Home,
+    tier: &str,
+    snap: &edge0_core::state::DirectSnapshot,
+    now: chrono::DateTime<chrono::Utc>,
+) -> Result<(), PullError> {
+    let st = staging_dir(home, tier, &snap.commit);
+    let dst = home.model_dir(tier, &snap.commit);
+    let mut snap = snap.clone();
+    snap.recorded_at = Some(now.to_rfc3339_opts(chrono::SecondsFormat::Secs, true));
+    if dst.join("install.json").is_file() {
+        let _ = std::fs::remove_dir_all(&st);
+        return Ok(());
+    }
+    if st.is_dir() {
+        if dst.exists() {
+            let _ = std::fs::remove_dir_all(&dst);
+        }
+        std::fs::create_dir_all(dst.parent().unwrap())
+            .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+        std::fs::rename(&st, &dst).map_err(|e| {
+            PullError::new(
+                "E-SRV-CONFLICT",
+                format!(
+                    "commit rename {} → {} failed: {e}",
+                    st.display(),
+                    dst.display()
+                ),
+            )
+        })?;
+    } else if !dst.is_dir() {
+        return Err(PullError::new(
+            "E-MODEL-MISSING",
+            format!("commit: neither staging nor destination exists: {}", st.display()),
+        ));
+    }
+    let mut snap = crate::check::record_probe_blocks(&snap, &dst)?;
+    snap.recorded_at = Some(now.to_rfc3339_opts(chrono::SecondsFormat::Secs, true));
+    atomic_write_json(&dst.join("install.json"), &snap)
+        .map_err(|e| PullError::new("E-SRV-CONFLICT", e))
+}
+
+pub fn prune_old_revs(home: &Home, tier: &str, keep_rev: &str) -> Vec<String> {
+    let mut pruned = Vec::new();
+    let dir = home.models().join(tier);
+    let Ok(rd) = std::fs::read_dir(&dir) else {
+        return pruned;
+    };
+    for ent in rd.flatten() {
+        let p = ent.path();
+        let Some(name) = p.file_name().map(|s| s.to_string_lossy().to_string()) else {
+            continue;
+        };
+        if name == keep_rev {
+            continue;
+        }
+        let stale = name.ends_with(".staging") || p.is_dir();
+        if stale && std::fs::remove_dir_all(&p).is_ok() {
+            pruned.push(name);
+        }
+    }
+    pruned
+}
+
+pub fn read_install(
+    home: &Home,
+    tier: &str,
+    rev: &str,
+) -> Option<edge0_core::state::DirectSnapshot> {
+    let text = std::fs::read_to_string(home.model_dir(tier, rev).join("install.json")).ok()?;
+    serde_json::from_str(&text).ok()
+}
+
+pub fn gc_orphans(home: &Home) -> Vec<String> {
+    let mut removed = Vec::new();
+    let mut claimed: std::collections::HashSet<String> = std::collections::HashSet::new();
+    if let Ok(rd) = std::fs::read_dir(home.state().join("downloads")) {
+        for ent in rd.flatten() {
+            if let Some(t) = std::fs::read_to_string(ent.path())
+                .ok()
+                .and_then(|x| serde_json::from_str::<edge0_core::state::DownloadTask>(&x).ok())
+            {
+                let _ = t;
+                claimed.insert(ent.path().to_string_lossy().to_string());
+            }
+        }
+    }
+    if let Ok(rd) = std::fs::read_dir(home.tmp()) {
+        for ent in rd.flatten() {
+            let p = ent.path();
+            let name = p
+                .file_name()
+                .map(|s| s.to_string_lossy().to_string())
+                .unwrap_or_default();
+            if !name.ends_with(".part") {
+                continue;
+            }
+            let task_side = home
+                .state()
+                .join("downloads")
+                .join(format!("{}.json", name.split('.').next().unwrap_or("")));
+            if task_side.exists() {
+                let _ = claimed;
+                continue;
+            }
+            if std::fs::remove_file(&p).is_ok() {
+                removed.push(format!("tmp/{name}"));
+            }
+        }
+    }
+    if let Ok(rd) = std::fs::read_dir(home.models()) {
+        for tier_ent in rd.flatten() {
+            let tier_dir = tier_ent.path();
+            if !tier_dir.is_dir() {
+                continue;
+            }
+            let Ok(revs) = std::fs::read_dir(&tier_dir) else {
+                continue;
+            };
+            for rev_ent in revs.flatten() {
+                let rp = rev_ent.path();
+                let Some(nm) = rp.file_name().map(|s| s.to_string_lossy().to_string()) else {
+                    continue;
+                };
+                if !nm.ends_with(".staging") || !rp.is_dir() {
+                    continue;
+                }
+                let owner_rev = nm.trim_end_matches(".staging").to_string();
+                let owned = std::fs::read_dir(home.state().join("downloads"))
+                    .map(|rd| {
+                        rd.flatten().any(|e| {
+                            std::fs::read_to_string(e.path())
+                                .ok()
+                                .and_then(|x| {
+                                    serde_json::from_str::<edge0_core::state::DownloadTask>(&x).ok()
+                                })
+                                .is_some_and(|t| t.rev == owner_rev && t.failed.is_none())
+                        })
+                    })
+                    .unwrap_or(false);
+                if owned {
+                    continue;
+                }
+                if std::fs::remove_dir_all(&rp).is_ok() {
+                    removed.push(format!(
+                        "models/{}/{nm}",
+                        tier_dir.file_name().unwrap().to_string_lossy()
+                    ));
+                }
+            }
+        }
+    }
+    removed
+}
diff --git a/macos/crates/edge0-pull/src/runner.rs b/macos/crates/edge0-pull/src/runner.rs
new file mode 100644
index 0000000..37f9f1f
--- /dev/null
+++ b/macos/crates/edge0-pull/src/runner.rs
@@ -0,0 +1,842 @@
+//! Eight-phase download runner used by both the daemon and CLI self-run.
+use std::path::PathBuf;
+use std::sync::{Arc, Mutex};
+
+use futures::{StreamExt, TryStreamExt};
+use sha2::{Digest, Sha256};
+
+use edge0_core::manifest::{FileEntry, Manifest, TierEntry};
+use edge0_core::paths::Home;
+use edge0_core::state::{DownloadTask, FileTask};
+
+use crate::error::PullError;
+use crate::fetch::{Fetcher, HeadInfo};
+use crate::registry_ops;
+use crate::state_machine::PHASES;
+
+pub trait TaskStore: Send + Sync {
+    fn save(&self, task: &DownloadTask) -> Result<(), String>;
+    fn delete(&self, id: &str) -> Result<(), String>;
+}
+
+pub struct FileStore {
+    pub home: Home,
+}
+
+impl TaskStore for FileStore {
+    fn save(&self, task: &DownloadTask) -> Result<(), String> {
+        edge0_core::state::atomic_write_json(&self.home.download_json(&task.id), task)
+    }
+    fn delete(&self, id: &str) -> Result<(), String> {
+        let _ = std::fs::remove_file(self.home.download_json(id));
+        Ok(())
+    }
+}
+
+pub trait DownloadEvents: Send + Sync {
+    fn event(&self, etype: &'static str, task: &DownloadTask, code: Option<&'static str>);
+}
+
+pub struct NoopEvents;
+impl DownloadEvents for NoopEvents {
+    fn event(&self, _: &'static str, _: &DownloadTask, _: Option<&'static str>) {}
+}
+
+pub struct TaskRunner {
+    pub home: Home,
+    pub fetcher: Arc<dyn Fetcher>,
+    pub store: Arc<dyn TaskStore>,
+    pub events: Arc<dyn DownloadEvents>,
+    pub source_pin: Option<String>,
+    pub parallel: usize,
+}
+
+fn file_url_for_src(f: &FileEntry, src: &str) -> Option<String> {
+    let s = f.sources.as_ref()?;
+    match src {
+        "modelscope" => s.modelscope.clone(),
+        "huggingface" => s.huggingface.clone(),
+        _ => None,
+    }
+}
+
+pub fn tier_bytes(manifest: &Manifest, tier: &str) -> Option<u64> {
+    manifest
+        .find_tier(tier)
+        .map(|t| t.files.iter().map(|f| f.size).sum())
+}
+
+#[derive(Clone)]
+struct JobFile {
+    path: String,
+    size: u64,
+    sha256: Option<String>,
+    urls: Vec<(String, String)>,
+    part_key: String,
+}
+
+struct Job {
+    direct: bool,
+    base_url: String,
+    files: Vec<JobFile>,
+}
+
+fn path_key(path: &str) -> String {
+    let mut h = Sha256::new();
+    h.update(path.as_bytes());
+    h.finalize().iter().map(|b| format!("{b:02x}")).collect()
+}
+
+impl TaskRunner {
+    fn part_path(&self, task_id: &str, key: &str) -> PathBuf {
+        self.home.tmp().join(format!("{task_id}.{key}.part"))
+    }
+
+    fn build_job(task: &DownloadTask) -> Result<Job, PullError> {
+        if let Some(snap) = &task.snapshot {
+            let entry: TierEntry = snap.find_tier(&task.tier).cloned().ok_or_else(|| {
+                PullError::new(
+                    "E-MODEL-MISSING",
+                    format!("Snapshot does not contain tier {}", task.tier),
+                )
+            })?;
+            let files = entry
+                .files
+                .iter()
+                .map(|f| JobFile {
+                    path: f.path.clone(),
+                    size: f.size,
+                    sha256: Some(f.sha256.clone()),
+                    urls: {
+                        let mut v = Vec::new();
+                        if let Some(u) = file_url_for_src(f, "modelscope") {
+                            v.push(("modelscope".to_string(), u));
+                        }
+                        if let Some(u) = file_url_for_src(f, "huggingface") {
+                            v.push(("huggingface".to_string(), u));
+                        }
+                        v
+                    },
+                    part_key: f.sha256.clone(),
+                })
+                .collect();
+            return Ok(Job {
+                direct: false,
+                base_url: String::new(),
+                files,
+            });
+        }
+        if let Some(plan) = &task.plan {
+            let files = plan
+                .files
+                .iter()
+                .map(|f| JobFile {
+                    path: f.path.clone(),
+                    size: f.size,
+                    sha256: f.sha256.clone(),
+                    urls: vec![("huggingface".to_string(), plan.url_for(&f.path))],
+                    part_key: path_key(&f.path),
+                })
+                .collect();
+            return Ok(Job {
+                direct: true,
+                base_url: plan.base_url.clone(),
+                files,
+            });
+        }
+        Err(PullError::new(
+            "E-MANIFEST-SIG",
+            "Task has no manifest snapshot or direct plan; refusing to run",
+        ))
+    }
+
+    pub async fn run(&self, initial: DownloadTask) -> Result<DownloadTask, PullError> {
+        let job = Self::build_job(&initial)?;
+        let total: u64 = job.files.iter().map(|f| f.size).sum();
+        let mut task = initial;
+        task.bytes_total = total;
+        reconcile_files(&mut task, &job);
+        let start_done = task.bytes_done;
+        let sh = Arc::new(Mutex::new(Shared {
+            task,
+            saved_bytes: 0,
+            last_sample: Some((std::time::Instant::now(), start_done)),
+            ema: None,
+        }));
+        let out = self.drive(&sh, &job).await;
+        match out {
+            Ok(()) => {
+                let final_task = sh.lock().unwrap().task.clone();
+                Ok(final_task)
+            }
+            Err(e) => {
+                {
+                    let mut g = sh.lock().unwrap();
+                    g.task.failed = Some(edge0_core::state::TaskFailure {
+                        code: e.code.to_string(),
+                        message: e.message.clone(),
+                        at: chrono::Utc::now(),
+                    });
+                    let t = g.task.clone();
+                    drop(g);
+                    let _ = self.store.save(&t);
+                    self.events.event("download.failed", &t, Some(e.code));
+                }
+                Err(e)
+            }
+        }
+    }
+
+    async fn drive(&self, sh: &Arc<Mutex<Shared>>, job: &Job) -> Result<(), PullError> {
+        if sh.lock().unwrap().task.source.is_none() {
+            self.set_phase(sh, "probe")?;
+            let chosen = self.probe_and_select(sh, job).await?;
+            self.set_phase(sh, "select-source")?;
+            let mut g = sh.lock().unwrap();
+            g.task.source = Some(chosen);
+            if job.direct && g.task.source_reason.is_none() {
+                g.task.source_reason = Some(crate::catalog::CreationPlan::source_reason_direct(
+                    &job.base_url,
+                    false,
+                ));
+            }
+            drop(g);
+            self.store
+                .save(&sh.lock().unwrap().task)
+                .map_err(|e| PullError::new("E-SRV-CONFLICT", e))?;
+        }
+        // —— preflight ——
+        self.set_phase(sh, "preflight")?;
+        let total = job.files.iter().map(|f| f.size).sum();
+        registry_ops::preflight_check(&self.home, total)?;
+        for attempt in 0..2u8 {
+            self.set_phase(sh, "fetch")?;
+            self.fetch_pass(sh, job).await?;
+            self.set_phase(sh, "verify")?;
+            let bad = self.verify_pass(sh, job).await?;
+            if bad.is_empty() {
+                break;
+            }
+            if attempt == 1 {
+                for f in job.files.iter().filter(|f| bad.contains(&f.path)) {
+                    let _ = std::fs::remove_file(
+                        self.part_path(&sh.lock().unwrap().task.id, &f.part_key),
+                    );
+                }
+                return Err(PullError::new(
+                    "E-MODEL-HASH",
+                    if job.direct {
+                        format!(
+                            "Files still failed validation after retry (size or recorded hash, {}); temporary files were removed and the task snapshot was kept",
+                            bad.join(", ")
+                        )
+                    } else {
+                        format!(
+                            "File hashes still do not match the manifest after switching sources ({}); temporary files were removed and the manifest snapshot was kept",
+                            bad.join(", ")
+                        )
+                    },
+                ));
+            }
+            if job.direct {
+                let mut g = sh.lock().unwrap();
+                let reason = g.task.source_reason.take().unwrap_or_default().to_string();
+                g.task.source_reason = Some(format!(
+                    "{reason}; recorded hash mismatch, retried once with the same source"
+                ));
+                if let Some(plan) = g.task.plan.as_mut() {
+                    for pf in plan.files.iter_mut().filter(|x| bad.contains(&x.path)) {
+                        pf.sha256 = None;
+                    }
+                }
+                for ft in g.task.files.iter_mut().filter(|ft| bad.contains(&ft.path)) {
+                    ft.set_done(0);
+                }
+                let t = g.task.clone();
+                let _ = self.store.save(&t);
+            } else {
+                let other = {
+                    let g = sh.lock().unwrap();
+                    let cur = g.task.source.clone().unwrap_or_default();
+                    if cur == "modelscope" {
+                        "huggingface"
+                    } else {
+                        "modelscope"
+                    }
+                    .to_string()
+                };
+                {
+                    let mut g = sh.lock().unwrap();
+                    let reason = g.task.source_reason.take().unwrap_or_default().to_string();
+                    g.task.source_reason = Some(format!(
+                        "{reason}; hash mismatch, automatically switched to {other} for one retry"
+                    ));
+                    g.task.source = Some(other.clone());
+                    g.task.source_switched = true;
+                    for ft in g.task.files.iter_mut().filter(|ft| bad.contains(&ft.path)) {
+                        ft.set_done(0);
+                    }
+                    let t = g.task.clone();
+                    let _ = self.store.save(&t);
+                }
+            }
+            for f in job.files.iter().filter(|f| bad.contains(&f.path)) {
+                let _ =
+                    std::fs::remove_file(self.part_path(&sh.lock().unwrap().task.id, &f.part_key));
+            }
+            self.events
+                .event("download.progress", &sh.lock().unwrap().task, None);
+        }
+        if job.direct {
+            return self.drive_direct_finish(sh, job).await;
+        }
+        self.set_phase(sh, "stage")?;
+        let task = sh.lock().unwrap().task.clone();
+        for f in &job.files {
+            let part = self.part_path(&task.id, &f.part_key);
+            if part.exists() {
+                registry_ops::stage_blob(&self.home, &part, &f.sha256.clone().unwrap_or_default())?;
+            }
+        }
+        self.set_phase(sh, "commit")?;
+        let task = sh.lock().unwrap().task.clone();
+        let snap = task.snapshot.as_ref().ok_or_else(|| {
+            PullError::new(
+                "E-MANIFEST-SIG",
+                "Registry commit is missing the manifest snapshot",
+            )
+        })?;
+        let entry = snap.find_tier(&task.tier).ok_or_else(|| {
+            PullError::new(
+                "E-MODEL-MISSING",
+                format!("Snapshot does not contain tier {}", task.tier),
+            )
+        })?;
+        for f in &entry.files {
+            registry_ops::link_into_model(&self.home, &task.tier, &task.rev, &f.path, &f.sha256)?;
+        }
+        self.set_phase(sh, "register")?;
+        let fresh = registry_ops::write_manifest_copy(&self.home, &task.tier, &task.rev, snap)?;
+        if fresh {
+            let mut refs = registry_ops::BlobRefs::load(&self.home);
+            for f in &entry.files {
+                *refs.0.entry(f.sha256.clone()).or_insert(0) += 1;
+            }
+            refs.save(&self.home)?;
+        }
+        self.store
+            .delete(&task.id)
+            .map_err(|e| PullError::new("E-SRV-CONFLICT", e))?;
+        self.events.event("download.completed", &task, None);
+        Ok(())
+    }
+
+    async fn drive_direct_finish(
+        &self,
+        sh: &Arc<Mutex<Shared>>,
+        job: &Job,
+    ) -> Result<(), PullError> {
+        self.set_phase(sh, "stage")?;
+        let task = sh.lock().unwrap().task.clone();
+        for f in &job.files {
+            let part = self.part_path(&task.id, &f.part_key);
+            let dst = registry_ops::staging_file(&self.home, &task.tier, &task.rev, &f.path);
+            if dst.is_file() && dst.metadata().is_ok_and(|m| m.len() == f.size) {
+                let _ = std::fs::remove_file(&part);
+                continue;
+            }
+            if part.exists() {
+                if let Some(parent) = dst.parent() {
+                    std::fs::create_dir_all(parent)
+                        .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+                }
+                std::fs::rename(&part, &dst).map_err(|e| {
+                    PullError::new(
+                        "E-SRV-CONFLICT",
+                        format!("stage {} → {} failed: {e}", part.display(), dst.display()),
+                    )
+                })?;
+            } else {
+                return Err(PullError::new(
+                    "E-MODEL-MISSING",
+                    format!("Staged file is missing: {} (retry the task)", f.path),
+                ));
+            }
+        }
+        self.set_phase(sh, "commit")?;
+        let (mut task, plan) = {
+            let g = sh.lock().unwrap();
+            let plan = g
+                .task
+                .plan
+                .clone()
+                .ok_or_else(|| PullError::new("E-MANIFEST-SIG", "direct commit is missing the plan snapshot"))?;
+            (g.task.clone(), plan)
+        };
+        let mut plan = plan;
+        for pf in &mut plan.files {
+            if pf.sha256.is_none() {
+                pf.sha256 = job
+                    .files
+                    .iter()
+                    .find(|j| j.path == pf.path)
+                    .and_then(|j| j.sha256.clone());
+            }
+        }
+        registry_ops::commit_direct(&self.home, &task.tier, &plan, chrono::Utc::now())?;
+        self.set_phase(sh, "register")?;
+        let pruned = registry_ops::prune_old_revs(&self.home, &task.tier, &task.rev);
+        if !pruned.is_empty() {
+            task.source_reason = Some(format!(
+                "{}; single-revision install replaced old rev {}",
+                task.source_reason.clone().unwrap_or_default(),
+                pruned.join(",")
+            ));
+        }
+        self.store
+            .delete(&task.id)
+            .map_err(|e| PullError::new("E-SRV-CONFLICT", e))?;
+        sh.lock().unwrap().task = task.clone();
+        self.events.event("download.completed", &task, None);
+        Ok(())
+    }
+
+    fn set_phase(&self, sh: &Arc<Mutex<Shared>>, phase: &str) -> Result<(), PullError> {
+        debug_assert!(PHASES.contains(&phase), "phase must be one of the eight download phases");
+        {
+            let mut g = sh.lock().unwrap();
+            g.task.phase = phase.to_string();
+            g.task.updated_at = chrono::Utc::now();
+            let t = g.task.clone();
+            self.store
+                .save(&t)
+                .map_err(|e| PullError::new("E-SRV-CONFLICT", e))?;
+        }
+        Ok(())
+    }
+
+    async fn probe_and_select(
+        &self,
+        sh: &Arc<Mutex<Shared>>,
+        job: &Job,
+    ) -> Result<String, PullError> {
+        if job.direct {
+            let primary = "huggingface".to_string();
+            let url = job
+                .files
+                .first()
+                .and_then(|f| f.urls.first())
+                .map(|(_, u)| u.clone())
+                .unwrap_or_default();
+            let head = self.fetcher.head(url).await;
+            if !head.ok {
+                return Err(PullError::new(
+                    "E-DL-SRC",
+                    format!(
+                        "Direct source is unreachable (base={}); retry later or change networks",
+                        job.base_url
+                    ),
+                ));
+            }
+            return Ok(primary);
+        }
+        let first = &job.files[0];
+        let primary = first
+            .urls
+            .iter()
+            .find(|(s, _)| s == "modelscope")
+            .map(|_| "modelscope")
+            .or(if !first.urls.is_empty() {
+                Some("huggingface")
+            } else {
+                None
+            })
+            .unwrap_or("modelscope")
+            .to_string();
+        if let Some(pin) = &self.source_pin {
+            if pin != "modelscope" && pin != "huggingface" {
+                return Err(PullError::new(
+                    "E-SRV-PARAM",
+                    format!("Unknown source {pin} (use modelscope or huggingface)"),
+                ));
+            }
+            {
+                sh.lock().unwrap().task.source_reason =
+                    Some(format!("Pinned source {pin} (--source/EDGE0_SOURCE)"));
+            }
+            return Ok(pin.clone());
+        }
+        let mut heads: Vec<(String, HeadInfo)> = Vec::new();
+        for src in ["modelscope", "huggingface"] {
+            if let Some((_, url)) = first.urls.iter().find(|(s, _)| s == src) {
+                heads.push((src.to_string(), self.fetcher.head(url.clone()).await));
+            }
+        }
+        let usable: Vec<&(String, HeadInfo)> = heads.iter().filter(|(_, h)| h.ok).collect();
+        if usable.is_empty() {
+            return Err(PullError::new(
+                "E-DL-SRC",
+                "Both sources are unreachable; retry later or change networks",
+            ));
+        }
+        let chosen = usable
+            .iter()
+            .min_by_key(|(_, h)| h.rtt)
+            .map(|(s, _)| s.clone())
+            .unwrap_or_else(|| primary.clone());
+        {
+            let mut g = sh.lock().unwrap();
+            g.task.source_reason = Some(format!(
+                "source probe: modelscope {}ms / huggingface {}ms → {chosen}",
+                head_ms(&heads, "modelscope"),
+                head_ms(&heads, "huggingface"),
+            ));
+        }
+        Ok(chosen)
+    }
+
+    async fn fetch_pass(&self, sh: &Arc<Mutex<Shared>>, job: &Job) -> Result<(), PullError> {
+        let source = sh.lock().unwrap().task.source.clone().unwrap_or_default();
+        let pending: Vec<JobFile> = {
+            let g = sh.lock().unwrap();
+            job.files
+                .iter()
+                .filter(|f| {
+                    let done = g
+                        .task
+                        .files
+                        .iter()
+                        .find(|t| t.path == f.path)
+                        .map(|t| t.done())
+                        .unwrap_or(0);
+                    done < f.size
+                })
+                .cloned()
+                .collect()
+        };
+        let sem = Arc::new(tokio::sync::Semaphore::new(self.parallel.max(1)));
+        let mut jobs = Vec::new();
+        for f in pending {
+            let sh = sh.clone();
+            let sem = sem.clone();
+            let home = self.home.clone();
+            let fetcher = self.fetcher.clone();
+            let store = self.store.clone();
+            let events = self.events.clone();
+            let source = source.clone();
+            jobs.push(async move {
+                let _permit = sem.acquire().await.unwrap();
+                let id = sh.lock().unwrap().task.id.clone();
+                let part_path = home.tmp().join(format!("{id}.{}.part", f.part_key));
+                if !job.direct {
+                    if let Some(sha) = &f.sha256 {
+                        if registry_ops::blob_path(&home, sha).exists() {
+                            update_prefix(&sh, &f.path, f.size, store.as_ref(), events.as_ref())?;
+                            return Ok(());
+                        }
+                    }
+                }
+                let url = f
+                    .urls
+                    .iter()
+                    .find(|(s, _)| *s == source)
+                    .or(f.urls.first())
+                    .map(|(_, u)| u.clone())
+                    .ok_or_else(|| {
+                        PullError::new(
+                            "E-DL-SRC",
+                            format!("No source URL is available for {}", f.path),
+                        )
+                    })?;
+                let start = {
+                    let g = sh.lock().unwrap();
+                    g.task
+                        .files
+                        .iter()
+                        .find(|t| t.path == f.path)
+                        .map(|t| t.done())
+                        .unwrap_or(0)
+                };
+                let mut stream = fetcher.open_range(url.clone(), start).await?;
+                use tokio::io::AsyncWriteExt;
+                let mut file = tokio::fs::OpenOptions::new()
+                    .write(true)
+                    .create(true)
+                    .truncate(false)
+                    .open(&part_path)
+                    .await
+                    .map_err(|e| {
+                        PullError::new(
+                            "E-SRV-CONFLICT",
+                            format!("open {} failed: {e}", part_path.display()),
+                        )
+                    })?;
+                file.set_len(start)
+                    .await
+                    .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+                use tokio::io::AsyncSeekExt;
+                file.seek(std::io::SeekFrom::Start(start))
+                    .await
+                    .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+                let mut pos = start;
+                let mut since_save = 0u64;
+                while let Some(chunk) = stream.next().await {
+                    let chunk = match chunk {
+                        Ok(c) => c,
+                        Err(e) => {
+                            let _ =
+                                update_prefix(&sh, &f.path, pos, store.as_ref(), events.as_ref());
+                            return Err(e);
+                        }
+                    };
+                    file.write_all(&chunk).await.map_err(|e| {
+                        PullError::new("E-SRV-CONFLICT", format!("write staging file failed: {e}"))
+                    })?;
+                    pos += chunk.len() as u64;
+                    since_save += chunk.len() as u64;
+                    if since_save >= 4 * 1024 * 1024 {
+                        file.flush()
+                            .await
+                            .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+                        update_prefix(&sh, &f.path, pos, store.as_ref(), events.as_ref())?;
+                        since_save = 0;
+                    }
+                }
+                file.flush()
+                    .await
+                    .map_err(|e| PullError::new("E-SRV-CONFLICT", e.to_string()))?;
+                if pos != f.size {
+                    let _ = update_prefix(&sh, &f.path, pos, store.as_ref(), events.as_ref());
+                    return Err(PullError::new(
+                        "E-DL-NET",
+                        format!("{} stream ended early ({pos}/{}), stopped at the verified prefix", f.path, f.size),
+                    ));
+                }
+                update_prefix(&sh, &f.path, pos, store.as_ref(), events.as_ref())?;
+                Ok(())
+            });
+        }
+        futures::stream::iter(jobs)
+            .buffer_unordered(self.parallel.max(1))
+            .try_collect::<()>()
+            .await
+    }
+
+    async fn verify_pass(
+        &self,
+        sh: &Arc<Mutex<Shared>>,
+        job: &Job,
+    ) -> Result<Vec<String>, PullError> {
+        let task = sh.lock().unwrap().task.clone();
+        let mut bad = Vec::new();
+        if !job.direct {
+            for f in &job.files {
+                if let Some(sha) = &f.sha256 {
+                    if registry_ops::blob_path(&self.home, sha).exists() {
+                        continue;
+                    }
+                }
+                let part = self.part_path(&task.id, &f.part_key);
+                if !part.exists() {
+                    bad.push(f.path.clone());
+                    continue;
+                }
+                let got = sha256_of(&part).await?;
+                if Some(&got) != f.sha256.as_ref() {
+                    bad.push(f.path.clone());
+                }
+            }
+            return Ok(bad);
+        }
+        let records: Vec<Option<String>> = {
+            let g = sh.lock().unwrap();
+            job.files
+                .iter()
+                .map(|f| {
+                    g.task.plan.as_ref().and_then(|p| {
+                        p.files
+                            .iter()
+                            .find(|x| x.path == f.path)
+                            .and_then(|x| x.sha256.clone())
+                    })
+                })
+                .collect()
+        };
+        let mut recorded: Vec<(String, String)> = Vec::new();
+        for (f, rec) in job.files.iter().zip(records) {
+            let part = self.part_path(&task.id, &f.part_key);
+            let len = match std::fs::metadata(&part) {
+                Ok(m) => m.len(),
+                Err(_) => {
+                    bad.push(f.path.clone());
+                    continue;
+                }
+            };
+            if len != f.size {
+                bad.push(f.path.clone());
+                continue;
+            }
+            let got = sha256_of(&part).await?;
+            if let Some(r) = &rec {
+                if r != &got {
+                    bad.push(f.path.clone());
+                    continue;
+                }
+            }
+            recorded.push((f.path.clone(), got));
+        }
+        let t = {
+            let mut g = sh.lock().unwrap();
+            for (path, got) in &recorded {
+                if let Some(plan) = g.task.plan.as_mut() {
+                    if let Some(pf) = plan.files.iter_mut().find(|x| x.path == *path) {
+                        pf.sha256 = Some(got.clone());
+                    }
+                }
+            }
+            g.task.clone()
+        };
+        if bad.is_empty() {
+            let _ = self.store.save(&t);
+        }
+        Ok(bad)
+    }
+}
+
+fn head_ms(heads: &[(String, HeadInfo)], src: &str) -> u128 {
+    heads
+        .iter()
+        .find(|(s, _)| s == src)
+        .map(|(_, h)| h.rtt.as_millis())
+        .unwrap_or(0)
+}
+
+fn sample_interval() -> std::time::Duration {
+    static I: std::sync::OnceLock<std::time::Duration> = std::sync::OnceLock::new();
+    *I.get_or_init(|| {
+        let ms = std::env::var("EDGE0_TEST_RATE_SAMPLE_MS")
+            .ok()
+            .and_then(|v| v.parse::<u64>().ok())
+            .unwrap_or(50);
+        std::time::Duration::from_millis(ms)
+    })
+}
+
+fn compute_eta(phase: &str, ema: Option<f64>, done: u64, total: u64) -> Option<u64> {
+    let rest = total.checked_sub(done)?;
+    if phase != "fetch" || rest == 0 {
+        return None;
+    }
+    let rate = ema.filter(|r| *r >= 1.0)?;
+    Some((rest as f64 / rate).ceil() as u64)
+}
+
+fn ema_blend(prev: Option<f64>, inst: f64) -> f64 {
+    match prev {
+        Some(e) => 0.7 * e + 0.3 * inst,
+        None => inst,
+    }
+}
+
+fn update_prefix(
+    sh: &Arc<Mutex<Shared>>,
+    path: &str,
+    n: u64,
+    store: &dyn TaskStore,
+    events: &dyn DownloadEvents,
+) -> Result<(), PullError> {
+    let t = {
+        let mut g = sh.lock().unwrap();
+        if let Some(ft) = g.task.files.iter_mut().find(|ft| ft.path == path) {
+            if ft.done() < n {
+                ft.set_done(n);
+            }
+        }
+        g.task.bytes_done = g.task.files.iter().map(|ft| ft.done()).sum();
+        g.task.updated_at = chrono::Utc::now();
+        {
+            let now = std::time::Instant::now();
+            match g.last_sample {
+                Some((t0, b0)) => {
+                    let dt = now.duration_since(t0);
+                    let db = g.task.bytes_done.saturating_sub(b0);
+                    if dt >= sample_interval() && db > 0 {
+                        let inst = db as f64 / dt.as_secs_f64();
+                        g.ema = Some(ema_blend(g.ema, inst));
+                        g.last_sample = Some((now, g.task.bytes_done));
+                    }
+                }
+                None => {
+                    last_sample_set(&mut g, now);
+                }
+            }
+            g.task.rate_bps = g.ema.map(|e| e as u64);
+            g.task.eta_s = compute_eta(&g.task.phase, g.ema, g.task.bytes_done, g.task.bytes_total);
+        }
+        if g.task.bytes_done.saturating_sub(g.saved_bytes) >= 4 * 1024 * 1024
+            || n == g.task.bytes_total
+        {
+            g.saved_bytes = g.task.bytes_done;
+            let t = g.task.clone();
+            drop(g);
+            store
+                .save(&t)
+                .map_err(|e| PullError::new("E-SRV-CONFLICT", e))?;
+            t
+        } else {
+            let t = g.task.clone();
+            drop(g);
+            t
+        }
+    };
+    events.event("download.progress", &t, None);
+    Ok(())
+}
+
+fn last_sample_set(g: &mut Shared, now: std::time::Instant) {
+    g.last_sample = Some((now, g.task.bytes_done));
+}
+
+async fn sha256_of(path: &PathBuf) -> Result<String, PullError> {
+    use tokio::io::AsyncReadExt;
+    let mut f = tokio::fs::File::open(path)
+        .await
+        .map_err(|e| PullError::new("E-MODEL-HASH", format!("read staging file failed: {e}")))?;
+    let mut h = Sha256::new();
+    let mut buf = vec![0u8; 1 << 20];
+    loop {
+        let n = f
+            .read(&mut buf)
+            .await
+            .map_err(|e| PullError::new("E-MODEL-HASH", e.to_string()))?;
+        if n == 0 {
+            break;
+        }
+        h.update(&buf[..n]);
+    }
+    Ok(h.finalize().iter().map(|b| format!("{b:02x}")).collect())
+}
+
+struct Shared {
+    task: DownloadTask,
+    saved_bytes: u64,
+    last_sample: Option<(std::time::Instant, u64)>,
+    ema: Option<f64>,
+}
+
+fn reconcile_files(task: &mut DownloadTask, job: &Job) {
+    let mut out: Vec<FileTask> = Vec::new();
+    for f in &job.files {
+        let existing = task
+            .files
+            .iter()
+            .find(|t| t.path == f.path && t.sha256 == f.part_key);
+        out.push(FileTask {
+            path: f.path.clone(),
+            size: f.size,
+            sha256: f.part_key.clone(),
+            done_ranges: existing.map(|e| e.done_ranges.clone()).unwrap_or_default(),
+        });
+    }
+    task.files = out;
+}
diff --git a/macos/crates/edge0-pull/src/state_machine.rs b/macos/crates/edge0-pull/src/state_machine.rs
new file mode 100644
index 0000000..9e7fa72
--- /dev/null
+++ b/macos/crates/edge0-pull/src/state_machine.rs
@@ -0,0 +1,21 @@
+//! Download phases: `probe → select-source → preflight → fetch → verify → stage → commit → register`.
+
+pub const PHASES: [&str; 8] = [
+    "probe",
+    "select-source",
+    "preflight",
+    "fetch",
+    "verify",
+    "stage",
+    "commit",
+    "register",
+];
+
+pub fn next(phase: &str) -> Option<&'static str> {
+    let i = PHASES.iter().position(|p| *p == phase)?;
+    PHASES.get(i + 1).copied()
+}
+
+pub fn is_valid(phase: &str) -> bool {
+    PHASES.contains(&phase)
+}
diff --git a/macos/crates/edge0-pull/src/trust.rs b/macos/crates/edge0-pull/src/trust.rs
new file mode 100644
index 0000000..46a3ea7
--- /dev/null
+++ b/macos/crates/edge0-pull/src/trust.rs
@@ -0,0 +1,80 @@
+//! Compiled-in public-key anchors; hosting sites are not trusted.
+
+use base64::Engine as _;
+use ed25519_dalek::{Signature, VerifyingKey};
+
+use crate::canonical::canonical_bytes;
+use crate::error::PullError;
+use edge0_core::manifest::Manifest;
+
+pub use base64::engine::general_purpose::STANDARD as B64;
+
+#[derive(Debug, Clone, Copy)]
+pub struct Anchor {
+    pub key_id: &'static str,
+    pub pubkey_b64: &'static str,
+}
+
+pub const ANCHORS: &[Anchor] = &[
+    Anchor {
+        key_id: "+y4LSPkaSzk=", // edge0-registry-2026
+        pubkey_b64: "VIq4ssoyjcpM+PcOFN6k42YZASdocY3tFYxilcUYW0o=",
+    },
+    Anchor {
+        key_id: "pvsAPXUAzqE=", // dev-test-2026
+        pubkey_b64: "ACNtc1fOxqUBHUWOpTp0jBQT9+FB8eglCuEeFOwu0DM=",
+    },
+];
+
+pub fn verify(manifest: &Manifest) -> Result<&'static str, PullError> {
+    let value = serde_json::to_value(manifest)
+        .map_err(|e| PullError::new("E-MANIFEST-SIG", format!("manifest reserialize failed: {e}")))?;
+    let msg = canonical_bytes(&value);
+    let sig_bytes = B64
+        .decode(&manifest.signature)
+        .map_err(|_| PullError::new("E-MANIFEST-SIG", "signature is not valid base64 (64 bytes)"))?;
+    if sig_bytes.len() != 64 {
+        return Err(PullError::new("E-MANIFEST-SIG", "signature length is not 64 bytes"));
+    }
+    let sig = Signature::from_bytes(sig_bytes.as_slice().try_into().unwrap());
+    let candidates: Vec<&Anchor> = match manifest.signer_key_id.as_deref() {
+        Some(kid) => ANCHORS.iter().filter(|a| a.key_id == kid).collect(),
+        None => ANCHORS.iter().collect(),
+    };
+    let n_cand = candidates.len();
+    for a in candidates {
+        let Ok(raw) = B64.decode(a.pubkey_b64) else {
+            continue;
+        };
+        let Ok(bytes): Result<[u8; 32], _> = raw.try_into() else {
+            continue;
+        };
+        if let Ok(vk) = VerifyingKey::from_bytes(&bytes) {
+            if vk.verify_strict(&msg, &sig).is_ok() {
+                return Ok(a.key_id);
+            }
+        }
+    }
+    Err(PullError::new(
+        "E-MANIFEST-SIG",
+        format!("signature check failed: untrusted source or tampered content ({n_cand} candidate anchors)"),
+    ))
+}
+
+pub fn sign_value(manifest: &mut serde_json::Value, secret: &[u8; 32], key_id: &str) {
+    use ed25519_dalek::{Signer, SigningKey};
+    if let Some(obj) = manifest.as_object_mut() {
+        obj.remove("signature");
+        obj.remove("signer_key_id");
+    }
+    let msg = canonical_bytes(manifest);
+    let sk = SigningKey::from_bytes(secret);
+    let sig = sk.sign(&msg);
+    if let Some(obj) = manifest.as_object_mut() {
+        obj.insert(
+            "signature".into(),
+            serde_json::json!(B64.encode(sig.to_bytes())),
+        );
+        obj.insert("signer_key_id".into(), serde_json::json!(key_id));
+    }
+}
diff --git a/macos/crates/edge0d/Cargo.toml b/macos/crates/edge0d/Cargo.toml
new file mode 100644
index 0000000..75717ec
--- /dev/null
+++ b/macos/crates/edge0d/Cargo.toml
@@ -0,0 +1,31 @@
+[package]
+name = "edge0d"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+[dependencies]
+edge0-core.workspace = true
+edge0-pull.workspace = true
+reqwest.workspace = true
+tokio.workspace = true
+axum.workspace = true
+tower.workspace = true
+http.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+chrono.workspace = true
+uuid.workspace = true
+rand.workspace = true
+sha2.workspace = true
+thiserror.workspace = true
+anyhow.workspace = true
+tracing.workspace = true
+tracing-subscriber.workspace = true
+futures.workspace = true
+tokio-stream.workspace = true
+clap.workspace = true
+libloading.workspace = true
+tokio-util.workspace = true
+libc.workspace = true
diff --git a/macos/crates/edge0d/src/app.rs b/macos/crates/edge0d/src/app.rs
new file mode 100644
index 0000000..2f12d22
--- /dev/null
+++ b/macos/crates/edge0d/src/app.rs
@@ -0,0 +1,385 @@
+//! Daemon app state: single-instance check, token lifecycle, startup, and HTTP middleware.
+
+use std::collections::HashMap;
+use std::net::{IpAddr, SocketAddr};
+use std::os::unix::fs::PermissionsExt;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
+
+use axum::response::IntoResponse;
+use chrono::{DateTime, Utc};
+use rand::Rng;
+
+use edge0_core::abi_load::{EngineLib, TokHandle};
+use edge0_core::config::Config;
+use edge0_core::errorcodes::{ApiError, Code};
+use edge0_core::paths::Home;
+use edge0_core::state::{DaemonJson, Scope};
+
+use crate::bus::EventBus;
+use crate::registry::Registry;
+use crate::workers::WorkerPool;
+
+pub struct App {
+    pub cfg: Config,
+    pub home: Home,
+    pub scope: Scope,
+    pub started_by: String,
+    pub owner_pid: Option<u32>,
+    pub session_id: String,
+    pub started_at: DateTime<Utc>,
+    pub uptime_from: Instant,
+    pub token: Option<String>,
+    pub bus: EventBus,
+    pub registry: Arc<Mutex<Registry>>,
+    pub downloads: Arc<crate::registry::DownloadStore>,
+    pub pool: Arc<WorkerPool>,
+    pub lib: Arc<EngineLib>,
+    /// Per-tier tokenizer handles; independent of worker lifetime.
+    pub tok_handles: Mutex<HashMap<String, Arc<Mutex<TokHandle>>>>,
+    pub gates: Mutex<HashMap<String, Arc<tokio::sync::Semaphore>>>,
+    pub request_counter: AtomicU64,
+    pub host_app_version: Option<String>,
+    /// App Nap has no public probe API; `throttled` is always false (doctor explains why).
+    pub throttled: bool,
+    pub stop: Arc<std::sync::atomic::AtomicBool>,
+    /// Held for the daemon lifetime so CLI cannot bypass the registry lock.
+    pub _reglock: crate::registry::HeldLock,
+    pub runners: Mutex<HashMap<String, tokio::task::JoinHandle<()>>>,
+    pub catalog: Mutex<Option<crate::http::Cached>>,
+    /// UI double-clicks / parallel fetches share one network round-trip.
+    pub catalog_flight: tokio::sync::Mutex<()>,
+}
+
+pub type AppRef = Arc<App>;
+
+impl App {
+    pub fn new(
+        cfg: Config,
+        home: Home,
+        scope: Scope,
+        started_by: String,
+        owner_pid: Option<u32>,
+    ) -> Result<Self, ApiError> {
+        if let Some(peer) = Self::peer_daemon(cfg.port) {
+            return Err(ApiError::new(code("E-SRV-PORT"), format!(
+                "A service is already running (version {}, session_id {}, scope {}, pid {}, uptime {}s). A second daemon is not allowed for the same {}, and the port is not changed.",
+                peer.version, peer.session_id, peer.scope, peer.pid, peer.uptime_s, home.root.display()
+            )));
+        }
+        if std::net::TcpListener::bind(cfg.socket_addr()).is_err() {
+            return Err(ApiError::new(code("E-SRV-PORT"), format!(
+                "Port {} is in use by another process: free it or set EDGE0_PORT (the port is not changed automatically, so base_url stays predictable).",
+                cfg.port
+            )));
+        }
+        let token = if cfg.is_loopback_bind() {
+            None
+        } else {
+            match read_token(&home) {
+                Some(t) => Some(t),
+                None => {
+                    if cfg.enable_lan {
+                        let t = gen_token();
+                        write_token(&home, &t).map_err(|e| ApiError::new(code("E-SRV-AUTH"), e))?;
+                        Some(t)
+                    } else {
+                        return Err(ApiError::new(code("E-SRV-AUTH"), format!(
+                            "Refusing to start: bound to {} (non-loopback) with no existing token. After accepting the risk, start with EDGE0_ENABLE_LAN=1 (first launch writes a token to {}); the GUI “Allow LAN connections” control is the same action.",
+                            cfg.bind, home.token().display()
+                        )));
+                    }
+                }
+            }
+        };
+
+        let reglock = crate::registry::acquire_lock(&home).map_err(|e| {
+            ApiError::new(
+                code("E-SRV-CONFLICT"),
+                format!("{e} — daemon cannot become the sole registry writer: retry later or stop the holder"),
+            )
+        })?;
+        let bus = EventBus::open(&home);
+        let registry = Registry::rescan(&home).map_err(|e| ApiError::new(code("E-SRV-PORT"), e))?;
+        let downloads = Arc::new(crate::registry::DownloadStore::new(
+            &home,
+            registry.downloads.clone(),
+        ));
+        let engine_lib_path = std::env::var("EDGE0_ENGINE_LIB")
+            .map(std::path::PathBuf::from)
+            .unwrap_or_else(|_| default_engine_lib());
+        let lib = Arc::new(
+            EngineLib::load(&engine_lib_path)
+                .map_err(|e| ApiError::new(code("E-SRV-WORKER"), e))?,
+        );
+        eprintln!(
+            "edge0d: engine={} lib={}",
+            lib.backend(),
+            engine_lib_path.display()
+        );
+        let registry = Arc::new(Mutex::new(registry));
+        let pool = Arc::new(WorkerPool::new(
+            home.clone(),
+            bus.clone(),
+            engine_lib_path,
+            registry.clone(),
+        ));
+        Ok(Self {
+            cfg,
+            home,
+            scope,
+            started_by,
+            owner_pid,
+            session_id: format!("e0-sess-{}", uuid::Uuid::new_v4()),
+            started_at: Utc::now(),
+            uptime_from: Instant::now(),
+            token,
+            bus,
+            registry,
+            downloads,
+            pool,
+            lib,
+            tok_handles: Mutex::new(HashMap::new()),
+            gates: Mutex::new(HashMap::new()),
+            request_counter: AtomicU64::new(0),
+            host_app_version: std::env::var("EDGE0_HOST_APP_VERSION").ok(),
+            throttled: false,
+            stop: Arc::new(std::sync::atomic::AtomicBool::new(false)),
+            _reglock: reglock,
+            runners: Mutex::new(HashMap::new()),
+            catalog: Mutex::new(None),
+            catalog_flight: tokio::sync::Mutex::new(()),
+        })
+    }
+
+    /// Liveness is `/health`; `daemon.json` is only a hint. Bare TCP so a timeout means "no peer".
+    fn peer_daemon(port: u16) -> Option<PeerInfo> {
+        use std::io::{Read, Write};
+        let addr = SocketAddr::new(IpAddr::V4(std::net::Ipv4Addr::LOCALHOST), port);
+        let mut sock =
+            std::net::TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(400))
+                .ok()?;
+        sock.set_read_timeout(Some(std::time::Duration::from_millis(400)))
+            .ok()?;
+        sock.write_all(b"GET /health HTTP/1.1\r\nHost: edge0-probe\r\nConnection: close\r\n\r\n")
+            .ok()?;
+        let mut buf = Vec::new();
+        let _ = sock.read_to_end(&mut buf);
+        let text = String::from_utf8_lossy(&buf).into_owned();
+        let body = text.split("\r\n\r\n").nth(1)?;
+        let v: serde_json::Value = serde_json::from_str(body).ok()?;
+        Some(PeerInfo {
+            scope: v["scope"].as_str()?.to_string(),
+            session_id: v["session_id"].as_str()?.to_string(),
+            version: v["daemon_version"].as_str()?.to_string(),
+            pid: v["pid"].as_u64()? as u32,
+            uptime_s: v["uptime_s"].as_u64()?,
+        })
+    }
+
+    pub fn uptime_s(&self) -> u64 {
+        self.uptime_from.elapsed().as_secs()
+    }
+
+    pub fn next_request_id(&self) -> u64 {
+        self.request_counter.fetch_add(1, Ordering::Relaxed)
+    }
+
+    /// Lazy resident tokenizer; missing tier → E-MODEL-MISSING.
+    pub fn tok_of(&self, tier: &str) -> Result<Arc<Mutex<TokHandle>>, ApiError> {
+        {
+            let g = self.tok_handles.lock().unwrap();
+            if let Some(h) = g.get(tier) {
+                return Ok(h.clone());
+            }
+        }
+        let dir = {
+            let r = self.registry.lock().unwrap();
+            r.engine_dir(tier).ok_or_else(|| {
+                ApiError::new(
+                    code("E-MODEL-MISSING"),
+                    format!("tier {tier} is not installed: run `edge0 pull {tier}`"),
+                )
+            })
+        }?;
+        let handle = TokHandle::open(self.lib.clone(), &dir)
+            .map_err(|e| ApiError::new(code("E-MODEL-MISSING"), e))?;
+        let arc = Arc::new(Mutex::new(handle));
+        self.tok_handles
+            .lock()
+            .unwrap()
+            .insert(tier.to_string(), arc.clone());
+        Ok(arc)
+    }
+
+    pub fn try_acquire(&self, tier: &str) -> Option<tokio::sync::OwnedSemaphorePermit> {
+        let sem = {
+            let mut g = self.gates.lock().unwrap();
+            g.entry(tier.to_string())
+                .or_insert_with(|| {
+                    Arc::new(tokio::sync::Semaphore::new(self.cfg.concurrency as usize))
+                })
+                .clone()
+        };
+        sem.try_acquire_owned().ok()
+    }
+
+    pub fn write_daemon_json(&self) {
+        let d = DaemonJson::new(
+            self.scope,
+            &self.started_by,
+            self.owner_pid,
+            self.session_id.clone(),
+            self.cfg.port,
+            self.lib.backend(),
+        );
+        let _ = d.write(&self.home.daemon_json());
+    }
+    pub fn clear_daemon_json(&self) {
+        let _ = std::fs::remove_file(self.home.daemon_json());
+    }
+}
+
+struct PeerInfo {
+    scope: String,
+    session_id: String,
+    version: String,
+    pid: u32,
+    uptime_s: u64,
+}
+
+pub fn code(id: &str) -> &'static Code {
+    edge0_core::errorcodes::lookup(id).expect("code must exist in the table")
+}
+
+fn default_engine_lib() -> std::path::PathBuf {
+    // Packaged sibling first, then workspace `target/engine/` in debug.
+    resolve_engine_lib()
+}
+
+/// Same implementation (native), multiple paths: packaged sibling first, then workspace `target/engine/` in debug.
+fn native_engine_candidates() -> Vec<std::path::PathBuf> {
+    let mut v = Vec::new();
+    if let Ok(exe) = std::env::current_exe() {
+        if let Some(dir) = exe.parent() {
+            v.push(dir.join("libedge0_engine_native.dylib"));
+        }
+    }
+    if cfg!(debug_assertions) {
+        v.push(
+            std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+                .join("../../target/engine/libedge0_engine_native.dylib"),
+        );
+    }
+    v
+}
+
+/// First candidate that exists; if none, return the preferred path so dlopen reports it honestly.
+fn resolve_engine_lib() -> std::path::PathBuf {
+    native_engine_candidates()
+        .into_iter()
+        .find(|p| p.is_file())
+        .unwrap_or_else(|| {
+            native_engine_candidates()
+                .first()
+                .cloned()
+                .expect("candidates are non-empty")
+        })
+}
+
+fn read_token(home: &Home) -> Option<String> {
+    std::fs::read_to_string(home.token())
+        .ok()
+        .map(|s| s.trim().to_string())
+        .filter(|s| !s.is_empty())
+}
+
+fn write_token(home: &Home, token: &str) -> Result<(), String> {
+    home.ensure_dirs()?;
+    std::fs::write(home.token(), token.as_bytes()).map_err(|e| e.to_string())?;
+    std::fs::set_permissions(home.token(), std::fs::Permissions::from_mode(0o600))
+        .map_err(|e| e.to_string())?;
+    Ok(())
+}
+
+fn gen_token() -> String {
+    let mut rng = rand::thread_rng();
+    (0..43)
+        .map(|_| {
+            const A: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
+            let i: u8 = rng.gen_range(0..62);
+            A[i as usize] as char
+        })
+        .collect()
+}
+
+/// Origin → path exemption → Bearer.
+pub async fn guard_mw(
+    axum::extract::State(app): axum::extract::State<AppRef>,
+    req: axum::http::Request<axum::body::Body>,
+    next: axum::middleware::Next,
+) -> Result<axum::response::Response, axum::response::Response> {
+    let path = req.uri().path().to_string();
+    let peer: SocketAddr = req
+        .extensions()
+        .get::<axum::extract::connect_info::ConnectInfo<SocketAddr>>()
+        .map(|c| c.0)
+        .unwrap_or(SocketAddr::from(([127, 0, 0, 1], 0)));
+
+    if let Some(origin) = req.headers().get("origin").and_then(|v| v.to_str().ok()) {
+        if !app.cfg.origins.iter().any(|o| o == origin) {
+            return Err(render_error(ApiError::cors(format!(
+                "Origin {origin} is not on the allowlist"
+            ))));
+        }
+    }
+    if path != "/health" && !peer.ip().is_loopback() {
+        let want = app.token.clone().unwrap_or_default();
+        let got = req
+            .headers()
+            .get(axum::http::header::AUTHORIZATION)
+            .and_then(|v| v.to_str().ok())
+            .and_then(|v| v.strip_prefix("Bearer "))
+            .unwrap_or("");
+        if got != want {
+            return Err(render_error(ApiError::new(
+                code("E-SRV-AUTH"),
+                "Non-loopback access requires a Bearer token (state/token or the app service panel)",
+            )));
+        }
+    }
+    if let Some(cv) = req
+        .headers()
+        .get("x-edge0-client-version")
+        .and_then(|v| v.to_str().ok())
+    {
+        if cv != edge0_core::version::daemon_version() {
+            return Err(render_error(ApiError::new(
+                code("E-SRV-VERSION"),
+                format!(
+                    "Service version {} ≠ client version {cv}: restart the service or upgrade the app as prompted; no silent compatibility.",
+                    edge0_core::version::daemon_version()
+                ),
+            )));
+        }
+    }
+    Ok(next.run(req).await)
+}
+
+/// Sole error renderer: four envelope fields, no stack trace.
+pub fn render_error(e: ApiError) -> axum::response::Response {
+    let base = e.status_override.unwrap_or(e.code.http);
+    let status = if base == 0 {
+        axum::http::StatusCode::INTERNAL_SERVER_ERROR
+    } else {
+        axum::http::StatusCode::from_u16(base).unwrap_or(axum::http::StatusCode::BAD_REQUEST)
+    };
+    let mut resp = axum::Json(e.envelope()).into_response();
+    *resp.status_mut() = status;
+    if let Some(ra) = e.retry_after {
+        resp.headers_mut()
+            .insert("retry-after", ra.to_string().parse().unwrap());
+    }
+    resp
+}
diff --git a/macos/crates/edge0d/src/bus.rs b/macos/crates/edge0d/src/bus.rs
new file mode 100644
index 0000000..5133c51
--- /dev/null
+++ b/macos/crates/edge0d/src/bus.rs
@@ -0,0 +1,149 @@
+//! EventBus: coalescable, droppable, never reordered. Slow consumers resync rather than see out-of-order events.
+
+use std::collections::VecDeque;
+use std::sync::{Arc, Mutex};
+
+use chrono::Utc;
+use tokio::sync::watch;
+
+use edge0_core::paths::Home;
+use edge0_core::state::EventsState;
+
+pub const RING_CAP: usize = 4096;
+pub const RESTART_OFFSET: u64 = 10000;
+const PERSIST_EVERY: u64 = 100;
+
+pub use edge0_core::protocol::{foldable, Event};
+
+struct Inner {
+    ring: VecDeque<Event>,
+    base: u64, // absolute seq of ring[0]
+    seq: u64,
+    since_boot: u64, // publishes since last persist
+}
+
+#[derive(Clone)]
+pub struct EventBus {
+    inner: Arc<Mutex<Inner>>,
+    tx: watch::Sender<u64>,
+    home: Home,
+}
+
+impl EventBus {
+    /// After restart, seq = persisted max + 10000; a fresh install starts at 1.
+    pub fn open(home: &Home) -> Self {
+        let start = match std::fs::read_to_string(home.events_json())
+            .ok()
+            .and_then(|t| serde_json::from_str::<EventsState>(&t).ok())
+        {
+            Some(s) => s.max_seq + RESTART_OFFSET,
+            None => 1,
+        };
+        let (tx, _) = watch::channel(0u64);
+        Self {
+            inner: Arc::new(Mutex::new(Inner {
+                ring: VecDeque::new(),
+                base: start,
+                seq: start - 1,
+                since_boot: 0,
+            })),
+            tx,
+            home: home.clone(),
+        }
+    }
+
+    pub fn publish(&self, etype: &'static str, subject: &str, payload: serde_json::Value) -> u64 {
+        let ev = {
+            let mut g = self.inner.lock().unwrap();
+            g.seq += 1;
+            let ev = Event {
+                seq: g.seq,
+                ts: Utc::now(),
+                etype: etype.to_string(),
+                subject: subject.to_string(),
+                payload,
+            };
+            if foldable(etype) {
+                // In-place replace: an index table would break after the head seq shifts.
+                if let Some(slot) = g
+                    .ring
+                    .iter()
+                    .position(|e| e.etype == etype && e.subject == ev.subject)
+                {
+                    g.ring[slot] = ev.clone();
+                    self.tx.send_replace(g.seq);
+                    return ev.seq;
+                }
+            }
+            if g.ring.len() >= RING_CAP {
+                g.ring.pop_front();
+                g.base += 1;
+            }
+            g.ring.push_back(ev.clone());
+            g.since_boot += 1;
+            if g.since_boot >= PERSIST_EVERY {
+                let max = g.seq;
+                g.since_boot = 0;
+                let _ = persist(&self.home, max);
+            }
+            self.tx.send_replace(g.seq);
+            ev
+        };
+        ev.seq
+    }
+
+    pub fn current_seq(&self) -> u64 {
+        self.inner.lock().unwrap().seq
+    }
+
+    /// New subscription: (seq cursor, watch rx). `daemon.hello` is sent on connect and does not occupy the ring.
+    pub fn subscribe(&self) -> (u64, watch::Receiver<u64>) {
+        let g = self.inner.lock().unwrap();
+        (g.seq, self.tx.subscribe())
+    }
+
+    pub fn window(&self) -> (u64, u64) {
+        let g = self.inner.lock().unwrap();
+        (g.base, g.seq)
+    }
+
+    /// Replay ring events after `since`. Overflow (out of window or cursor ahead of this session) forces resync.
+    pub fn catchup(&self, since: u64) -> (Vec<Event>, bool) {
+        let g = self.inner.lock().unwrap();
+        if since + 1 < g.base || since > g.seq {
+            return (Vec::new(), true);
+        }
+        let mut out: Vec<Event> = Vec::new();
+        for ev in &g.ring {
+            if ev.seq > since {
+                out.push(ev.clone());
+            }
+        }
+        (out, false)
+    }
+
+    pub fn flush_persist(&self) {
+        let max = self.inner.lock().unwrap().seq;
+        let _ = persist(&self.home, max);
+    }
+}
+
+fn persist(home: &Home, max_seq: u64) -> Result<(), String> {
+    edge0_core::state::atomic_write_json(
+        &home.events_json(),
+        &EventsState {
+            max_seq,
+            updated_at: Utc::now(),
+        },
+    )
+}
+
+pub fn spawn_persist_loop(bus: EventBus, stop: Arc<std::sync::atomic::AtomicBool>) {
+    std::thread::spawn(move || {
+        while !stop.load(std::sync::atomic::Ordering::Relaxed) {
+            std::thread::sleep(std::time::Duration::from_secs(5));
+            bus.flush_persist();
+        }
+        bus.flush_persist();
+    });
+}
diff --git a/macos/crates/edge0d/src/http.rs b/macos/crates/edge0d/src/http.rs
new file mode 100644
index 0000000..96db077
--- /dev/null
+++ b/macos/crates/edge0d/src/http.rs
@@ -0,0 +1,910 @@
+//! HTTP routes for health, OpenAI-compatible chat, control plane, and events.
+
+use std::sync::Arc;
+
+use axum::extract::{Path, State};
+use axum::response::sse::{Event, Sse};
+use axum::response::{IntoResponse, Response};
+use axum::routing::{delete, get, post};
+use axum::{Json, Router};
+use tokio::sync::mpsc;
+
+use edge0_core::errorcodes::ApiError;
+use edge0_core::keepalive;
+use edge0_core::version::{daemon_version, E0_ABI_VERSION};
+
+use crate::app::{code, render_error, AppRef};
+use crate::workers::StreamItem;
+use edge0_core::frames::{Channel, EndReason};
+
+pub fn router(app: AppRef) -> Router {
+    Router::new()
+        .route("/health", get(health))
+        .route("/v1/models", get(models_list))
+        .route("/v1/models/{id}", get(model_get))
+        .route("/v1/chat/completions", post(post_chat))
+        .route("/v1/completions", get(unsupported).post(unsupported))
+        .route("/v1/embeddings", get(unsupported).post(unsupported))
+        .route("/v1/edge0/tokenize", post(tokenize))
+        .route("/v1/edge0/apply-template", post(apply_template))
+        .route("/v1/edge0/events", get(events))
+        .route("/v1/edge0/system", get(system))
+        .route("/v1/edge0/catalog", get(catalog))
+        .route("/v1/edge0/doctor", get(doctor))
+        .route("/v1/edge0/models/{id}/load", post(load))
+        .route("/v1/edge0/models/{id}/unload", post(unload))
+        .route("/v1/edge0/models/{id}", delete(models_delete))
+        .route(
+            "/v1/edge0/downloads",
+            get(downloads_list).post(downloads_create),
+        )
+        .route("/v1/edge0/downloads/{id}/pause", post(downloads_pause))
+        .route("/v1/edge0/downloads/{id}/resume", post(downloads_resume))
+        .route("/v1/edge0/downloads/{id}", delete(downloads_delete))
+        .fallback(not_found)
+        .layer(axum::middleware::from_fn_with_state(
+            app.clone(),
+            crate::app::guard_mw,
+        ))
+        .with_state(app)
+}
+
+async fn not_found() -> Response {
+    render_error(ApiError::new(code("E-SRV-PARAM"), "Unknown path").with_status(404))
+}
+
+async fn unsupported() -> Response {
+    render_error(ApiError::new(
+        code("E-MODEL-MISSING"),
+        "This release does not provide /v1/completions or /v1/embeddings. Use /v1/chat/completions.",
+    ))
+}
+
+// ---------------------------------------------------------------- /health
+
+async fn health(State(app): State<AppRef>) -> Json<serde_json::Value> {
+    Json(serde_json::json!({
+        "daemon_version": daemon_version(),
+        "abi_version": E0_ABI_VERSION,
+        "host_app_version": app.host_app_version,
+        "session_id": app.session_id,
+        "scope": app.scope.as_str(),
+        "started_by": app.started_by,
+        "started_at": app.started_at.to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
+        "uptime_s": app.uptime_s(),
+        "pid": std::process::id(),
+        "throttled": app.throttled,
+    }))
+}
+
+async fn models_list(State(app): State<AppRef>) -> Json<serde_json::Value> {
+    let residents = app.pool.resident();
+    let created = app.started_at.timestamp();
+    let data: Vec<_> = app
+        .registry
+        .lock()
+        .unwrap()
+        .tiers
+        .iter()
+        .map(|t| {
+            serde_json::json!({
+                "id": t.tier,
+                "object": "model",
+                "created": created,
+                "owned_by": "edge0",
+            })
+        })
+        .collect();
+    let _ = residents;
+    Json(serde_json::json!({ "object": "list", "data": data }))
+}
+
+async fn model_get(State(app): State<AppRef>, Path(id): Path<String>) -> Response {
+    let r = app.registry.lock().unwrap();
+    match r.find(&id) {
+        Some(t) => {
+            let resident = app.pool.is_running(&id);
+            Json(serde_json::json!({
+                "id": t.tier, "object": "model", "created": app.started_at.timestamp(),
+                "owned_by": "edge0",
+                "edge0": { "rev": t.rev, "resident": resident }
+            }))
+            .into_response()
+        }
+        None => render_error(ApiError::new(
+            code("E-MODEL-MISSING"),
+            format!("Model {id} is not installed. Run `edge0 pull {id}`."),
+        )),
+    }
+}
+
+#[derive(serde::Deserialize)]
+struct ChatIn {
+    model: String,
+    messages: Vec<serde_json::Value>,
+    #[serde(default)]
+    stream: bool,
+    max_tokens: Option<usize>,
+    temperature: Option<f32>,
+    top_p: Option<f32>,
+    top_k: Option<usize>,
+    #[allow(dead_code)]
+    stop: Option<serde_json::Value>,
+    keep_alive: Option<String>,
+    n: Option<usize>,
+    #[serde(flatten)]
+    _extra: serde_json::Map<String, serde_json::Value>,
+}
+
+fn build_params(c: &ChatIn) -> String {
+    let mut p = serde_json::Map::new();
+    p.insert(
+        "max_new_tokens".into(),
+        serde_json::json!(c.max_tokens.unwrap_or(64)),
+    );
+    if let Some(stop) = &c.stop {
+        let norm = match stop {
+            serde_json::Value::String(s) => serde_json::json!([s]),
+            v => v.clone(),
+        };
+        p.insert("stop".into(), norm);
+    }
+    for (k, v) in [
+        ("temperature", c.temperature.map(|x| x as f64)),
+        ("top_p", c.top_p.map(|x| x as f64)),
+    ] {
+        if let Some(v) = v {
+            p.insert(k.into(), serde_json::json!(v));
+        }
+    }
+    if let Some(k) = c.top_k {
+        p.insert("top_k".into(), serde_json::json!(k));
+    }
+    serde_json::Value::Object(p).to_string()
+}
+
+struct PeerAddr(Option<std::net::SocketAddr>);
+
+impl<S: Send + Sync> axum::extract::FromRequestParts<S> for PeerAddr {
+    type Rejection = std::convert::Infallible;
+
+    fn from_request_parts(
+        parts: &mut axum::http::request::Parts,
+        _state: &S,
+    ) -> impl std::future::Future<Output = Result<Self, Self::Rejection>> + Send {
+        let got = parts
+            .extensions
+            .get::<axum::extract::connect_info::ConnectInfo<std::net::SocketAddr>>()
+            .map(|c| c.0);
+        std::future::ready(Ok(Self(got)))
+    }
+}
+
+async fn post_chat(
+    State(app): State<AppRef>,
+    peer: PeerAddr,
+    headers: axum::http::HeaderMap,
+    body: axum::body::Bytes,
+) -> Response {
+    let trace = ChatTrace::new(app.next_request_id(), classify_source(peer.0, &headers));
+    let raw_hint = body.clone();
+    match chat_inner(app.clone(), trace.clone(), body).await {
+        Ok(r) => r,
+        Err(e) => {
+            if !trace.emitted() {
+                trace.finish(
+                    &app,
+                    &model_hint(&raw_hint),
+                    "error",
+                    Some(e.code.id),
+                    None,
+                    0,
+                );
+            }
+            render_error(e)
+        }
+    }
+}
+
+fn classify_source(
+    peer: Option<std::net::SocketAddr>,
+    headers: &axum::http::HeaderMap,
+) -> &'static str {
+    let bearer = headers
+        .get(axum::http::header::AUTHORIZATION)
+        .and_then(|v| v.to_str().ok())
+        .is_some_and(|v| {
+            v.trim_start()
+                .get(0..7)
+                .is_some_and(|p| p.eq_ignore_ascii_case("bearer "))
+        });
+    match peer.map(|p| p.ip()) {
+        Some(ip) if !ip.is_loopback() => "lan",
+        _ if bearer => "cli",
+        _ => "loopback",
+    }
+}
+
+fn model_hint(body: &[u8]) -> String {
+    serde_json::from_slice::<serde_json::Value>(body)
+        .ok()
+        .and_then(|v| v["model"].as_str().map(str::to_string))
+        .unwrap_or_default()
+}
+
+#[derive(Clone)]
+struct ChatTrace {
+    rid: u64,
+    source: &'static str,
+    t0: std::time::Instant,
+    emitted: std::sync::Arc<std::sync::atomic::AtomicBool>,
+}
+
+impl ChatTrace {
+    fn new(rid: u64, source: &'static str) -> Self {
+        Self {
+            rid,
+            source,
+            t0: std::time::Instant::now(),
+            emitted: Default::default(),
+        }
+    }
+
+    fn emitted(&self) -> bool {
+        self.emitted.load(std::sync::atomic::Ordering::Relaxed)
+    }
+
+    fn ms(&self) -> u64 {
+        self.t0.elapsed().as_millis() as u64
+    }
+
+    fn finish(
+        &self,
+        app: &AppRef,
+        model: &str,
+        status: &str,
+        code: Option<&str>,
+        ttft_from_t0: Option<std::time::Duration>,
+        tokens: u64,
+    ) {
+        if self
+            .emitted
+            .swap(true, std::sync::atomic::Ordering::Relaxed)
+        {
+            return;
+        }
+        let mut payload = serde_json::json!({
+            "status": status,
+            "model": model,
+            "latency_ms": self.ms(),
+            "tokens": tokens,
+            "source": self.source,
+        });
+        if let Some(c) = code {
+            payload["code"] = serde_json::json!(c);
+        }
+        if let Some(d) = ttft_from_t0 {
+            payload["ttft_ms"] = serde_json::json!(d.as_millis() as u64);
+        }
+        app.bus.publish(
+            "api.request.finished",
+            &format!("req-{}", self.rid),
+            payload,
+        );
+    }
+}
+
+async fn chat_inner(
+    app: AppRef,
+    trace: ChatTrace,
+    body: axum::body::Bytes,
+) -> Result<Response, ApiError> {
+    let rid = trace.rid;
+    let raw = String::from_utf8(body.to_vec())
+        .map_err(|_| ApiError::new(code("E-SRV-PARAM"), "Request body is not valid UTF-8"))?;
+    let strict: serde_json::Value = serde_json::from_str(&raw)
+        .map_err(|_| ApiError::new(code("E-SRV-PARAM"), "Request body is not valid JSON"))?;
+    if let Some(obj) = strict.as_object() {
+        const ALLOWED: &[&str] = &[
+            "model",
+            "messages",
+            "stream",
+            "max_tokens",
+            "temperature",
+            "top_p",
+            "top_k",
+            "stop",
+            "keep_alive",
+            "n",
+        ];
+        if let Some(bad) = obj.keys().find(|k| !ALLOWED.contains(&k.as_str())) {
+            return Err(ApiError::new(
+                code("E-SRV-PARAM"),
+                format!("Private extension parameters are not accepted: {bad}"),
+            )
+            .with_param(bad.clone()));
+        }
+    }
+    let c: ChatIn = serde_json::from_value(strict)
+        .map_err(|e| ApiError::new(code("E-SRV-AUTH"), e.to_string()))?;
+    if c.n.unwrap_or(1) != 1 {
+        return Err(
+            ApiError::new(code("E-SRV-PARAM"), "This release only supports n=1").with_param("n"),
+        );
+    }
+    if c.messages.is_empty()
+        || !c
+            .messages
+            .iter()
+            .all(|m| m["role"].is_string() && m["content"].is_string())
+    {
+        return Err(ApiError::new(
+            code("E-SRV-PARAM"),
+            "messages must be non-empty and role/content must be strings (multimodal input is not supported)",
+        )
+        .with_param("messages"));
+    }
+    match &c.stop {
+        None | Some(serde_json::Value::String(_)) => {}
+        Some(serde_json::Value::Array(v))
+            if !v.is_empty() && v.iter().all(|x| x.as_str().is_some_and(|s| !s.is_empty())) => {}
+        Some(_) => {
+            return Err(ApiError::new(
+                code("E-SRV-PARAM"),
+                "stop must be a non-empty string or an array of non-empty strings",
+            )
+            .with_param("stop"))
+        }
+    }
+    let installed = app.registry.lock().unwrap().find(&c.model).is_some();
+    if !installed {
+        app.bus.publish(
+            "request.rejected",
+            "*",
+            serde_json::json!({"code":"E-SRV-NOTLOAD","model":c.model}),
+        );
+        return Err(ApiError::new(
+            code("E-SRV-NOTLOAD"),
+            format!(
+                "Model {} is not installed. Run `edge0 pull {}`.",
+                c.model, c.model
+            ),
+        )
+        .with_param("model"));
+    }
+    if !app.pool.is_running(&c.model) {
+        let app2 = app.clone();
+        let tier = c.model.clone();
+        let h = tokio::spawn(async move { app2.pool.ensure_loaded(&tier).await });
+        match tokio::time::timeout(std::time::Duration::from_secs(3), h).await {
+            Ok(Ok(Ok(()))) => {}
+            Ok(Ok(Err(e))) => return Err(ApiError::new(code("E-MEM-LOAD"), e)),
+            Ok(Err(_)) | Err(_) => {
+                return Err(ApiError::new(
+                    code("E-SRV-NOTLOAD"),
+                    "Model is still loading; retry shortly",
+                )
+                .with_param("model")
+                .with_retry_after(2));
+            }
+        }
+    }
+    let permit = match app.try_acquire(&c.model) {
+        Some(p) => p,
+        None => {
+            app.bus.publish(
+                "request.rejected",
+                "*",
+                serde_json::json!({"code":"E-SRV-BUSY","model":c.model}),
+            );
+            return Err(ApiError::new(
+                code("E-SRV-BUSY"),
+                "The model concurrency limit has been reached",
+            )
+            .with_retry_after(1));
+        }
+    };
+    let tokh = app.tok_of(&c.model)?;
+    let msgs = serde_json::json!({
+        "messages": c.messages,
+        "add_generation_prompt": true,
+        "enable_thinking": false,
+    })
+    .to_string();
+    let rendered: serde_json::Value = {
+        let h = tokh.lock().unwrap();
+        serde_json::from_str(
+            &h.apply_template(&msgs)
+                .map_err(|_| ApiError::new(code("E-MODEL-MISSING"), "Template rendering failed"))?,
+        )
+        .map_err(|e| ApiError::new(code("E-SRV-WORKER"), e.to_string()))?
+    };
+    let prompt = rendered["prompt"].as_str().unwrap_or_default().to_string();
+    let ids: Vec<i32> = {
+        let h = tokh.lock().unwrap();
+        let out = h
+            .tokenize(&serde_json::json!({"text": prompt}).to_string())
+            .map_err(|_| ApiError::new(code("E-MODEL-MISSING"), "Tokenization failed"))?;
+        let v: serde_json::Value = serde_json::from_str(&out).unwrap_or_default();
+        v["tokens"][0]
+            .as_array()
+            .map(|a| {
+                a.iter()
+                    .filter_map(|x| x.as_i64().map(|i| i as i32))
+                    .collect()
+            })
+            .unwrap_or_default()
+    };
+    let params = build_params(&c);
+    let prompt_len = ids.len();
+    let (sid, mut rx) = app
+        .pool
+        .begin(&c.model, ids, &params)
+        .map_err(|e| ApiError::new(code("E-SRV-WORKER"), e))?;
+
+    let keep: Option<std::time::Duration> = match &c.keep_alive {
+        Some(s) => keepalive::parse(s).map_err(|e| {
+            ApiError::new(code("E-SRV-PARAM"), format!("Invalid keep_alive: {e}"))
+                .with_param("keep_alive")
+        })?,
+        None => keepalive::parse(&app.cfg.keep_alive_default).unwrap_or(None),
+    };
+    if c.stream {
+        let stream = TokenizerStream::new(
+            app.clone(),
+            c.model.clone(),
+            sid,
+            trace,
+            prompt_len,
+            StreamEnd {
+                rx,
+                _permit: permit,
+                keep,
+            },
+        );
+        let created = chrono::Utc::now().timestamp();
+        let body_stream = futures::stream::unfold(stream, move |mut s| async move {
+            let item = s.next_chunk(created).await?;
+            Some((item, s))
+        });
+        Ok(Sse::new(body_stream).into_response())
+    } else {
+        let mut ids = [Vec::<i32>::new(), Vec::<i32>::new()];
+        let mut ntok = 0usize;
+        let mut first_content: Option<std::time::Instant> = None;
+        let mut first_token: Option<std::time::Instant> = None;
+        loop {
+            match rx.recv().await {
+                Some(StreamItem::Token { token, channel }) => {
+                    ntok += 1;
+                    if first_token.is_none() {
+                        first_token = Some(std::time::Instant::now());
+                    }
+                    ids[chan_idx(channel)].push(token);
+                }
+                Some(StreamItem::Finished(_)) | None => break,
+                Some(StreamItem::Gone) => {
+                    app.pool.finish(&c.model, sid);
+                    trace.finish(
+                        &app,
+                        &c.model,
+                        "error",
+                        Some("E-SRV-WORKER"),
+                        ttft(&trace, first_content),
+                        ntok as u64,
+                    );
+                    return Err(ApiError::new(
+                        code("E-SRV-WORKER"),
+                        "The worker crashed and generation stopped; retry to reload it",
+                    ));
+                }
+            }
+        }
+        let content = decode_full(&tokh, &ids[chan_idx(Channel::Content)]);
+        if !content.is_empty() {
+            first_content = Some(std::time::Instant::now());
+        }
+        app.pool.finish(&c.model, sid);
+        app.pool.touch_keepalive(&c.model, keep);
+        crate::http::control::note_ctx(&c.model, (prompt_len + ntok) as u64);
+        let id = format!("chatcmpl-e0-{rid}");
+        trace.finish(
+            &app,
+            &c.model,
+            "done",
+            None,
+            ttft(&trace, first_content),
+            ntok as u64,
+        );
+        let perf = edge0_perf(
+            &app,
+            &c.model,
+            &trace,
+            &TurnPerf {
+                first_content,
+                first_token,
+                prompt_len,
+                ntok,
+            },
+        )
+        .await;
+        Ok(Json(serde_json::json!({
+            "id": id, "object": "chat.completion", "created": chrono::Utc::now().timestamp(),
+            "model": c.model,
+            "choices": [{ "index": 0, "message": { "role": "assistant", "content": content }, "finish_reason": "stop" }],
+            "usage": {
+                "prompt_tokens": prompt_len, "completion_tokens": ntok,
+                "total_tokens": prompt_len + ntok, "edge0": perf.edge0,
+                "prompt_tokens_details": { "cached_tokens": perf.cached_tokens },
+            },
+        }))
+        .into_response())
+    }
+}
+
+fn ttft(
+    trace: &ChatTrace,
+    first_content: Option<std::time::Instant>,
+) -> Option<std::time::Duration> {
+    first_content.map(|t| t.duration_since(trace.t0))
+}
+
+struct TurnPerf {
+    first_content: Option<std::time::Instant>,
+    first_token: Option<std::time::Instant>,
+    prompt_len: usize,
+    ntok: usize,
+}
+
+struct PerfReport {
+    edge0: serde_json::Value,
+    cached_tokens: usize,
+}
+
+async fn edge0_perf(app: &AppRef, tier: &str, trace: &ChatTrace, t: &TurnPerf) -> PerfReport {
+    let now = std::time::Instant::now();
+    let mut o = serde_json::Map::new();
+    let mut cached_tokens = 0usize;
+    let mut prefill_tokens = t.prompt_len;
+    if let Some(fc) = t.first_content {
+        o.insert(
+            "ttft_ms".into(),
+            serde_json::json!(fc.duration_since(trace.t0).as_millis()),
+        );
+    }
+    if let Some(ft) = t.first_token {
+        let pre = ft.duration_since(trace.t0);
+        if prefill_tokens > 0 && pre.as_micros() > 0 {
+            let v = prefill_tokens as f64 * 1e6 / pre.as_micros() as f64;
+            o.insert(
+                "prefill_tps".into(),
+                serde_json::json!((v * 10.0).round() / 10.0),
+            );
+        }
+        let dec = now.duration_since(ft);
+        if t.ntok > 1 && dec.as_micros() > 0 {
+            let v = (t.ntok - 1) as f64 * 1e6 / dec.as_micros() as f64;
+            o.insert(
+                "decode_tps".into(),
+                serde_json::json!((v * 10.0).round() / 10.0),
+            );
+        }
+    }
+    if let Ok(Some(s)) =
+        tokio::time::timeout(std::time::Duration::from_millis(300), app.pool.stats(tier)).await
+    {
+        if let Ok(v) = serde_json::from_str::<serde_json::Value>(&s) {
+            if let Some(b) = v["memory"]["footprint_bytes"].as_u64() {
+                o.insert("footprint_mb".into(), serde_json::json!(b / (1024 * 1024)));
+            }
+            if let Some(pc) = v.get("prefix_cache") {
+                cached_tokens = pc
+                    .get("last_cached_tokens")
+                    .and_then(|x| x.as_u64())
+                    .unwrap_or(0) as usize;
+                let reported_prefill = pc
+                    .get("last_prefill_tokens")
+                    .and_then(|x| x.as_u64())
+                    .unwrap_or(t.prompt_len.saturating_sub(cached_tokens) as u64)
+                    as usize;
+                if pc.get("enabled").and_then(|x| x.as_bool()).unwrap_or(true)
+                    || reported_prefill > 0
+                {
+                    prefill_tokens = reported_prefill;
+                }
+                o.insert("cached_tokens".into(), serde_json::json!(cached_tokens));
+                o.insert("prefill_tokens".into(), serde_json::json!(prefill_tokens));
+                o.insert(
+                    "prefix_cache_status".into(),
+                    pc.get("last_reason")
+                        .cloned()
+                        .unwrap_or_else(|| serde_json::json!("unknown")),
+                );
+            }
+        }
+    }
+    if let Some(ft) = t.first_token {
+        let pre = ft.duration_since(trace.t0);
+        if prefill_tokens > 0 && pre.as_micros() > 0 {
+            let v = prefill_tokens as f64 * 1e6 / pre.as_micros() as f64;
+            o.insert(
+                "prefill_tps".into(),
+                serde_json::json!((v * 10.0).round() / 10.0),
+            );
+        }
+    }
+    PerfReport {
+        edge0: serde_json::Value::Object(o),
+        cached_tokens,
+    }
+}
+
+struct StreamEnd {
+    rx: mpsc::UnboundedReceiver<StreamItem>,
+    _permit: tokio::sync::OwnedSemaphorePermit,
+    keep: Option<std::time::Duration>,
+}
+
+struct TokenizerStream {
+    app: AppRef,
+    tier: String,
+    sid: u64,
+    end: StreamEnd,
+    done: bool,
+    perf: TurnPerf,
+    trace: ChatTrace,
+    tail: std::collections::VecDeque<String>,
+    acc: [Vec<i32>; 2],
+    sent: [String; 2],
+}
+
+fn chan_idx(channel: Channel) -> usize {
+    match channel {
+        Channel::Content => 0,
+        Channel::Reasoning => 1,
+    }
+}
+
+fn safe_prefix_len(text: &str) -> usize {
+    let t = text.trim_end_matches('\u{FFFD}');
+    t.len()
+}
+
+impl TokenizerStream {
+    fn new(
+        app: AppRef,
+        tier: String,
+        sid: u64,
+        trace: ChatTrace,
+        prompt_len: usize,
+        end: StreamEnd,
+    ) -> Self {
+        let rid = trace.rid;
+        let mut me = Self {
+            app,
+            tier,
+            sid,
+            end,
+            done: false,
+            perf: TurnPerf {
+                first_content: None,
+                first_token: None,
+                prompt_len,
+                ntok: 0,
+            },
+            trace,
+            tail: Default::default(),
+            acc: [Vec::new(), Vec::new()],
+            sent: [String::new(), String::new()],
+        };
+        me.tail.push_back(
+            serde_json::json!({
+                "id": format!("chatcmpl-e0-{rid}"),
+                "object": "chat.completion.chunk",
+                "created": chrono::Utc::now().timestamp(),
+                "model": me.tier,
+                "choices": [{ "index": 0, "delta": {"role": "assistant"}, "finish_reason": null }],
+            })
+            .to_string(),
+        );
+        me
+    }
+
+    async fn next_chunk(&mut self, created: i64) -> Option<Result<Event, std::io::Error>> {
+        if let Some(data) = self.tail.pop_front() {
+            if data == "[DONE]" {
+                self.done = true;
+            }
+            return Some(Ok(Event::default().data(data)));
+        }
+        if self.done {
+            return None;
+        }
+        let tokh = match self.app.tok_of(&self.tier) {
+            Ok(t) => t,
+            Err(e) => {
+                self.done = true;
+                self.app.pool.finish(&self.tier, self.sid);
+                self.trace.finish(
+                    &self.app,
+                    &self.tier,
+                    "error",
+                    Some(e.code.id),
+                    ttft(&self.trace, self.perf.first_content),
+                    self.perf.ntok as u64,
+                );
+                return Some(Err(std::io::Error::other("Tokenizer unavailable")));
+            }
+        };
+        let item = self.end.rx.recv().await;
+        match item {
+            Some(StreamItem::Token { token, channel }) => {
+                self.perf.ntok += 1;
+                if self.perf.first_token.is_none() {
+                    self.perf.first_token = Some(std::time::Instant::now());
+                }
+                let i = chan_idx(channel);
+                self.acc[i].push(token);
+                let text = decode_full(&tokh, &self.acc[i]);
+                let safe = &text[..safe_prefix_len(&text)];
+                let ch = channel;
+                let text = if safe.len() >= self.sent[i].len() && safe.starts_with(&self.sent[i]) {
+                    let d = safe[self.sent[i].len()..].to_string();
+                    self.sent[i] = safe.to_string();
+                    d
+                } else {
+                    debug_assert!(false, "decoded prefix is not monotonic");
+                    self.sent[i] = safe.to_string();
+                    safe.to_string()
+                };
+                let mut delta = serde_json::Map::new();
+                match ch {
+                    Channel::Content => {
+                        if !text.is_empty() && self.perf.first_content.is_none() {
+                            self.perf.first_content = Some(std::time::Instant::now());
+                        }
+                        delta.insert("content".into(), serde_json::json!(text));
+                    }
+                    Channel::Reasoning => {
+                        if !text.is_empty() {
+                            delta.insert("reasoning_content".into(), serde_json::json!(text));
+                        }
+                    }
+                }
+                Some(Ok(Event::default().data(
+                    serde_json::json!({
+                        "id": format!("chatcmpl-e0-{}", self.sid),
+                        "object": "chat.completion.chunk",
+                        "created": created,
+                        "model": self.tier,
+                        "choices": [{ "index": 0, "delta": delta, "finish_reason": null }],
+                    })
+                    .to_string(),
+                )))
+            }
+            Some(StreamItem::Finished(reason)) => {
+                for (i, key) in [(0usize, "content"), (1, "reasoning_content")] {
+                    if self.acc[i].is_empty() {
+                        continue;
+                    }
+                    let full = decode_full(&tokh, &self.acc[i]);
+                    if full.len() > self.sent[i].len() && full.starts_with(&self.sent[i]) {
+                        let rest = full[self.sent[i].len()..].to_string();
+                        self.sent[i] = full;
+                        self.tail.push_back(
+                            serde_json::json!({
+                                "id": format!("chatcmpl-e0-{}", self.sid),
+                                "object": "chat.completion.chunk",
+                                "created": created,
+                                "model": self.tier,
+                                "choices": [{ "index": 0, "delta": { key: rest }, "finish_reason": null }],
+                            })
+                            .to_string(),
+                        );
+                    }
+                }
+                let (reason, status) = match reason {
+                    EndReason::Cancelled => ("cancel", "cancelled"),
+                    EndReason::Natural => ("stop", "done"),
+                };
+                self.trace.finish(
+                    &self.app,
+                    &self.tier,
+                    status,
+                    None,
+                    ttft(&self.trace, self.perf.first_content),
+                    self.perf.ntok as u64,
+                );
+                self.app.pool.finish(&self.tier, self.sid);
+                self.app.pool.touch_keepalive(&self.tier, self.end.keep);
+                crate::http::control::note_ctx(&self.tier, self.perf.ntok as u64);
+                self.tail.push_back(
+                    serde_json::json!({
+                        "id": format!("chatcmpl-e0-{}", self.trace.rid),
+                        "object": "chat.completion.chunk",
+                        "created": created,
+                        "model": self.tier,
+                        "choices": [{ "index": 0, "delta": {}, "finish_reason": reason }],
+                    })
+                    .to_string(),
+                );
+                let perf = edge0_perf(&self.app, &self.tier, &self.trace, &self.perf).await;
+                self.tail.push_back(
+                    serde_json::json!({
+                        "id": format!("chatcmpl-e0-{}", self.trace.rid),
+                        "object": "chat.completion.chunk",
+                        "created": created, "model": self.tier,
+                        "choices": [],
+                        "usage": {
+                            "prompt_tokens": self.perf.prompt_len,
+                            "completion_tokens": self.perf.ntok,
+                            "total_tokens": self.perf.prompt_len + self.perf.ntok,
+                            "edge0": perf.edge0,
+                            "prompt_tokens_details": { "cached_tokens": perf.cached_tokens },
+                        },
+                    })
+                    .to_string(),
+                );
+                self.tail.push_back("[DONE]".into());
+                if let Some(data) = self.tail.pop_front() {
+                    return Some(Ok(Event::default().data(data)));
+                }
+                None
+            }
+            Some(StreamItem::Gone) | None => {
+                self.done = true;
+                self.app.pool.finish(&self.tier, self.sid);
+                self.trace.finish(
+                    &self.app,
+                    &self.tier,
+                    "error",
+                    Some("E-SRV-WORKER"),
+                    ttft(&self.trace, self.perf.first_content),
+                    self.perf.ntok as u64,
+                );
+                Some(Ok(Event::default().event("error").data(
+                    serde_json::json!({"error": {"message":"The worker crashed and generation stopped; retry to reload it","type":"service","code":"E-SRV-WORKER","param":null}}).to_string(),
+                )))
+            }
+        }
+    }
+}
+
+impl Drop for TokenizerStream {
+    fn drop(&mut self) {
+        if !self.done {
+            self.app.pool.cancel(&self.tier, self.sid);
+            self.app.pool.finish(&self.tier, self.sid);
+            self.trace.finish(
+                &self.app,
+                &self.tier,
+                "cancelled",
+                None,
+                ttft(&self.trace, self.perf.first_content),
+                self.perf.ntok as u64,
+            );
+        }
+    }
+}
+
+fn decode_full(
+    tokh: &Arc<std::sync::Mutex<edge0_core::abi_load::TokHandle>>,
+    ids: &[i32],
+) -> String {
+    let out = {
+        let h = tokh.lock().unwrap();
+        h.decode(&serde_json::json!({ "tokens": ids }).to_string())
+    };
+    match out {
+        Ok(j) => serde_json::from_str::<serde_json::Value>(&j)
+            .ok()
+            .and_then(|v| v["text"].as_str().map(String::from))
+            .unwrap_or_default(),
+        Err(_) => String::new(),
+    }
+}
+
+mod control;
+pub use control::*;
+mod catalog;
+pub use catalog::*;
diff --git a/macos/crates/edge0d/src/http/catalog.rs b/macos/crates/edge0d/src/http/catalog.rs
new file mode 100644
index 0000000..6164150
--- /dev/null
+++ b/macos/crates/edge0d/src/http/catalog.rs
@@ -0,0 +1,216 @@
+//! `GET /v1/edge0/catalog`: fresh check, single-flight refetch, then stale fallback.
+
+use chrono::{DateTime, Duration, SecondsFormat, Utc};
+
+use axum::extract::State;
+use axum::response::{IntoResponse, Response};
+use serde::{Deserialize, Serialize};
+
+use edge0_core::config::PullMode;
+use edge0_core::errorcodes::ApiError;
+use edge0_core::protocol::CatalogResponse;
+use edge0_core::state::{DirectSnapshot, InstalledTier};
+use edge0_pull::catalog::{catalog_tiers, direct_tiers, resolve_files, BUILTIN_TIERS};
+use edge0_pull::mirror;
+
+use crate::app::{code, render_error, AppRef};
+
+#[derive(Clone)]
+pub enum Cached {
+    Signed(mirror::CatalogFetch),
+    Direct(DirectState),
+}
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct DirectState {
+    pub fetched_at: DateTime<Utc>,
+    pub snaps: Vec<(String, DirectSnapshot)>,
+}
+
+fn ttl() -> Duration {
+    std::env::var("EDGE0_CATALOG_TTL_S")
+        .ok()
+        .and_then(|v| v.parse::<i64>().ok())
+        .map(Duration::seconds)
+        .unwrap_or_else(|| Duration::seconds(300))
+}
+
+fn fresh(f: &mirror::CatalogFetch, now: DateTime<Utc>) -> bool {
+    !expired(f, now) && now.signed_duration_since(f.fetched_at) < ttl()
+}
+
+fn expired(f: &mirror::CatalogFetch, now: DateTime<Utc>) -> bool {
+    DateTime::parse_from_rfc3339(&f.manifest.expires_at)
+        .map(|d| d.with_timezone(&Utc) <= now)
+        .unwrap_or(true)
+}
+
+fn build(f: &mirror::CatalogFetch, stale: bool, installed: &[InstalledTier]) -> CatalogResponse {
+    CatalogResponse {
+        generated_at: f.manifest.generated_at.clone(),
+        expires_at: f.manifest.expires_at.clone(),
+        fetched_at: f.fetched_at.to_rfc3339_opts(SecondsFormat::Secs, true),
+        stale,
+        signer_key_id: f.signer_key_id.clone(),
+        tiers: catalog_tiers(&f.manifest, installed),
+    }
+}
+
+fn build_direct(d: &DirectState, stale: bool, installed: &[InstalledTier]) -> CatalogResponse {
+    let snaps: Vec<(String, DirectSnapshot)> = d.snaps.clone();
+    CatalogResponse {
+        generated_at: d.fetched_at.to_rfc3339_opts(SecondsFormat::Secs, true),
+        // Direct has no manifest expiry; a rolling window keeps the response shape. Clients use `stale`.
+        expires_at: (d.fetched_at + Duration::hours(24)).to_rfc3339_opts(SecondsFormat::Secs, true),
+        fetched_at: d.fetched_at.to_rfc3339_opts(SecondsFormat::Secs, true),
+        stale,
+        signer_key_id: "direct".to_string(),
+        tiers: direct_tiers(&snaps, installed),
+    }
+}
+
+fn serve(app: &AppRef, f: &mirror::CatalogFetch, stale: bool) -> Response {
+    let installed = app.registry.lock().unwrap().tiers.clone();
+    axum::Json(build(f, stale, &installed)).into_response()
+}
+
+fn serve_direct(app: &AppRef, d: &DirectState, stale: bool) -> Response {
+    let installed = app.registry.lock().unwrap().tiers.clone();
+    axum::Json(build_direct(d, stale, &installed)).into_response()
+}
+
+fn direct_disk_path(home: &edge0_core::paths::Home) -> std::path::PathBuf {
+    home.manifests().join("catalog-direct.json")
+}
+
+fn read_direct_cache(app: &AppRef) -> Option<DirectState> {
+    let text = std::fs::read_to_string(direct_disk_path(&app.home)).ok()?;
+    serde_json::from_str(&text).ok()
+}
+
+fn write_direct_cache(app: &AppRef, d: &DirectState) {
+    let _ = edge0_core::state::atomic_write_json(&direct_disk_path(&app.home), d);
+}
+
+pub async fn catalog(State(app): State<AppRef>) -> Response {
+    if app.cfg.pull_mode == PullMode::Direct {
+        return catalog_direct(app).await;
+    }
+    let now = Utc::now();
+    if let Some(Cached::Signed(f)) = app.catalog.lock().unwrap().as_ref() {
+        if fresh(f, now) {
+            return serve(&app, f, false);
+        }
+    }
+    if let Some(f) = mirror::read_catalog_cache(&app.home) {
+        if fresh(&f, now) {
+            *app.catalog.lock().unwrap() = Some(Cached::Signed(f.clone()));
+            return serve(&app, &f, false);
+        }
+    }
+    let _flight = app.catalog_flight.lock().await;
+    if let Some(Cached::Signed(f)) = app.catalog.lock().unwrap().as_ref() {
+        if fresh(f, now) {
+            let f = f.clone();
+            return serve(&app, &f, false);
+        }
+    }
+    let client = match reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(20))
+        .build()
+    {
+        Ok(c) => c,
+        Err(e) => return render_error(ApiError::new(code("E-SRV-PARAM"), e.to_string())),
+    };
+    match mirror::fetch_catalog(&client, &app.cfg.registry_urls, &app.home, now).await {
+        Ok(f) => {
+            // Successful refetch of an already-expired manifest is still stale for browse; create-task still rejects.
+            let stale = expired(&f, now);
+            *app.catalog.lock().unwrap() = Some(Cached::Signed(f.clone()));
+            serve(&app, &f, stale)
+        }
+        Err(e) => {
+            let fallback = {
+                let mem = app.catalog.lock().unwrap().clone();
+                match mem {
+                    Some(Cached::Signed(f)) => Some(f),
+                    _ => {
+                        let disk = mirror::read_catalog_cache(&app.home);
+                        if let Some(d) = disk.clone() {
+                            *app.catalog.lock().unwrap() = Some(Cached::Signed(d));
+                        }
+                        disk
+                    }
+                }
+            };
+            match fallback {
+                Some(f) => serve(&app, &f, true),
+                None => render_error(e.into_api_error()),
+            }
+        }
+    }
+}
+
+async fn catalog_direct(app: AppRef) -> Response {
+    let now = Utc::now();
+    if let Some(Cached::Direct(d)) = app.catalog.lock().unwrap().as_ref() {
+        if now.signed_duration_since(d.fetched_at) < ttl() {
+            return serve_direct(&app, d, false);
+        }
+    }
+    let _flight = app.catalog_flight.lock().await;
+    if let Some(Cached::Direct(d)) = app.catalog.lock().unwrap().as_ref() {
+        if now.signed_duration_since(d.fetched_at) < ttl() {
+            let d = d.clone();
+            return serve_direct(&app, &d, false);
+        }
+    }
+    let client = match reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(20))
+        .build()
+    {
+        Ok(c) => c,
+        Err(e) => return render_error(ApiError::new(code("E-SRV-PARAM"), e.to_string())),
+    };
+    let mut snaps = Vec::new();
+    let mut last_err = None;
+    for info in BUILTIN_TIERS {
+        match resolve_files(&client, &app.cfg.hf_base, info.tier).await {
+            Ok(s) => snaps.push((info.tier.to_string(), s)),
+            Err(e) => last_err = Some(e),
+        }
+    }
+    if !snaps.is_empty() {
+        let d = DirectState {
+            fetched_at: now,
+            snaps,
+        };
+        *app.catalog.lock().unwrap() = Some(Cached::Direct(d.clone()));
+        write_direct_cache(&app, &d);
+        // Partial probe failure: omit missing tiers rather than invent rows.
+        return serve_direct(&app, &d, last_err.is_some());
+    }
+    let fallback = {
+        let mem = app.catalog.lock().unwrap().clone();
+        match mem {
+            Some(Cached::Direct(d)) => Some(d),
+            _ => {
+                let disk = read_direct_cache(&app);
+                if let Some(d) = disk.clone() {
+                    *app.catalog.lock().unwrap() = Some(Cached::Direct(d));
+                }
+                disk
+            }
+        }
+    };
+    match fallback {
+        Some(d) => serve_direct(&app, &d, true),
+        None => render_error(
+            last_err
+                .unwrap_or_else(|| {
+                    edge0_pull::error::PullError::new("E-DL-NET", "direct catalog unreachable")
+                })
+                .into_api_error(),
+        ),
+    }
+}
diff --git a/macos/crates/edge0d/src/http/control.rs b/macos/crates/edge0d/src/http/control.rs
new file mode 100644
index 0000000..5309234
--- /dev/null
+++ b/macos/crates/edge0d/src/http/control.rs
@@ -0,0 +1,897 @@
+//! Control plane `/v1/edge0/*` and the event stream.
+
+use std::collections::HashMap;
+use std::sync::{Mutex, OnceLock};
+use std::time::UNIX_EPOCH;
+
+use axum::extract::{Path, Query, State};
+use axum::http::HeaderMap;
+use axum::response::sse::{Event, KeepAlive, Sse};
+use axum::response::{IntoResponse, Response};
+use axum::Json;
+use futures::StreamExt as _;
+use sha2::Digest as _;
+
+use edge0_core::errorcodes::ApiError;
+use edge0_core::keepalive;
+use edge0_core::protocol::{
+    DaemonInfo, DownloadEntry, DownloadFailure, HardwareInfo, HelloPayload, ModelEntry,
+    ResyncPayload, SystemSnapshot,
+};
+use edge0_core::version::{daemon_version, E0_ABI_VERSION};
+use edge0_pull::runner::TaskStore as _;
+
+use crate::app::{code, render_error, AppRef};
+use crate::bus::Event as BusEvent;
+
+// ---------------------------------------------------------------- tokenize / apply-template
+
+#[derive(serde::Deserialize)]
+pub struct TokIn {
+    model: String,
+    input: String,
+}
+
+pub async fn tokenize(State(app): State<AppRef>, Json(body): Json<TokIn>) -> Response {
+    match app.tok_of(&body.model) {
+        Err(e) => render_error(e),
+        Ok(h) => {
+            let out = {
+                let g = h.lock().unwrap();
+                g.tokenize(&serde_json::json!({ "text": body.input }).to_string())
+            };
+            match out {
+                Ok(j) => {
+                    let v: serde_json::Value = serde_json::from_str(&j).unwrap_or_default();
+                    let batch = &v["tokens"][0];
+                    let count = batch.as_array().map(|a| a.len()).unwrap_or(0);
+                    Json(serde_json::json!({
+                        "model": body.model,
+                        "data": [{ "index": 0, "tokens": batch, "count": count }],
+                    }))
+                    .into_response()
+                }
+                Err(st) => render_error(ApiError::new(
+                    code("E-MODEL-MISSING"),
+                    format!("Tokenization failed (ABI status {st})"),
+                )),
+            }
+        }
+    }
+}
+
+#[derive(serde::Deserialize)]
+pub struct TmplIn {
+    model: String,
+    messages: Vec<serde_json::Value>,
+    #[serde(default = "default_true")]
+    add_generation_prompt: bool,
+}
+
+fn default_true() -> bool {
+    true
+}
+
+pub async fn apply_template(State(app): State<AppRef>, Json(body): Json<TmplIn>) -> Response {
+    let h = match app.tok_of(&body.model) {
+        Ok(h) => h,
+        Err(e) => return render_error(e),
+    };
+    let arg = serde_json::json!({ "messages": body.messages, "add_generation_prompt": body.add_generation_prompt }).to_string();
+    let out = {
+        let g = h.lock().unwrap();
+        g.apply_template(&arg)
+    };
+    match out {
+        Ok(j) => {
+            let mut v: serde_json::Value = serde_json::from_str(&j).unwrap_or_default();
+            if let Some(o) = v.as_object_mut() {
+                o.insert("model".into(), serde_json::json!(body.model));
+            }
+            Json(v).into_response()
+        }
+        Err(st) => render_error(ApiError::new(
+            code("E-MODEL-MISSING"),
+            format!("Template application failed (ABI status {st})"),
+        )),
+    }
+}
+
+#[derive(serde::Deserialize)]
+pub struct EventsQ {
+    since: Option<u64>,
+}
+
+pub async fn events(
+    State(app): State<AppRef>,
+    headers: HeaderMap,
+    Query(q): Query<EventsQ>,
+) -> Response {
+    let since_hdr = headers
+        .get("last-event-id")
+        .and_then(|v| v.to_str().ok())
+        .and_then(|s| s.parse::<u64>().ok());
+    let since = since_hdr.or(q.since);
+    let (cur, mut watch) = app.bus.subscribe();
+    let app2 = app.clone();
+    let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
+    tokio::spawn(async move {
+        let hello = BusEvent {
+            seq: cur,
+            ts: chrono::Utc::now(),
+            etype: "daemon.hello".to_string(),
+            subject: "*".to_string(),
+            payload: serde_json::json!(HelloPayload {
+                session_id: app2.session_id.clone(),
+                seq: cur,
+                abi_version: E0_ABI_VERSION,
+            }),
+        };
+        let _ = tx.send(sse_event(&hello));
+        let mut cursor = since.unwrap_or(cur);
+        if since.is_some() {
+            let (evs, overflow) = app2.bus.catchup(since.unwrap());
+            if overflow {
+                let (base, top) = app2.bus.window();
+                let _ = tx.send(sse_resync(since.unwrap(), base, top, &app2.session_id));
+                cursor = top;
+            } else {
+                for e in evs {
+                    cursor = e.seq;
+                    let _ = tx.send(sse_event(&e));
+                }
+            }
+        }
+        while watch.changed().await.is_ok() {
+            let (evs, overflow) = app2.bus.catchup(cursor);
+            if overflow {
+                let (base, top) = app2.bus.window();
+                let _ = tx.send(sse_resync(cursor, base, top, &app2.session_id));
+                cursor = top;
+                continue;
+            }
+            for e in evs {
+                cursor = e.seq;
+                if tx.send(sse_event(&e)).is_err() {
+                    return;
+                }
+            }
+        }
+    });
+    Sse::new(tokio_stream::wrappers::UnboundedReceiverStream::new(rx).map(Ok::<_, std::io::Error>))
+        .keep_alive(KeepAlive::new().interval(std::time::Duration::from_secs(10)))
+        .into_response()
+}
+
+fn sse_event(e: &BusEvent) -> Event {
+    Event::default()
+        .event(e.etype.as_str())
+        .id(e.seq.to_string())
+        .data(serde_json::to_string(e).unwrap_or_else(|_| "{}".into()))
+}
+
+fn sse_resync(gap_from: u64, buffer_min: u64, top: u64, session_id: &str) -> Event {
+    let ev = BusEvent {
+        seq: top,
+        ts: chrono::Utc::now(),
+        etype: "resync.required".to_string(),
+        subject: "*".to_string(),
+        payload: serde_json::json!(ResyncPayload {
+            gap_from,
+            buffer_min,
+            session_id: session_id.to_string(),
+        }),
+    };
+    sse_event(&ev)
+}
+
+pub fn ctx_map() -> &'static Mutex<HashMap<String, u64>> {
+    static C: OnceLock<Mutex<HashMap<String, u64>>> = OnceLock::new();
+    C.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+pub async fn system(State(app): State<AppRef>) -> Json<SystemSnapshot> {
+    let residents = app.pool.resident();
+    let registry = app.registry.lock().unwrap();
+    let models: Vec<ModelEntry> = registry
+        .tiers
+        .iter()
+        .map(|t| {
+            let (state, unload_at, active) = match residents.iter().find(|(r, _, _)| *r == t.tier) {
+                Some((_, ua, a)) => (
+                    "resident",
+                    ua.map(|d| d.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)),
+                    *a as u64,
+                ),
+                None => ("unloaded", None, 0),
+            };
+            ModelEntry {
+                id: t.tier.clone(),
+                rev: t.rev.clone(),
+                bytes_total: t.bytes_total,
+                state: state.to_string(),
+                unload_at,
+                active_requests: active,
+                context_used_tokens: ctx_used(&t.tier),
+                bundled: t.bundled,
+            }
+        })
+        .collect();
+    drop(registry);
+    let downloads: Vec<DownloadEntry> = app.downloads.list().iter().map(dl_json).collect();
+    Json(SystemSnapshot {
+        daemon: DaemonInfo {
+            version: daemon_version().to_string(),
+            abi_version: E0_ABI_VERSION,
+            host_app_version: app.host_app_version.clone(),
+            scope: app.scope.as_str().to_string(),
+            started_by: app.started_by.clone(),
+            owner_pid: app.owner_pid,
+            session_id: app.session_id.clone(),
+            started_at: app
+                .started_at
+                .to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
+            uptime_s: app.uptime_s(),
+            throttled: app.throttled,
+            port: app.cfg.port,
+            bound: app.cfg.bind.to_string(),
+            home: app.home.root.display().to_string(),
+            engine: app.lib.backend(),
+        },
+        hardware: hardware_snapshot(),
+        models,
+        downloads,
+        events_seq: app.bus.current_seq(),
+    })
+}
+
+fn ctx_used(tier: &str) -> u64 {
+    ctx_map().lock().unwrap().get(tier).copied().unwrap_or(0)
+}
+
+pub fn note_ctx(tier: &str, tokens: u64) {
+    *ctx_map()
+        .lock()
+        .unwrap()
+        .entry(tier.to_string())
+        .or_default() = tokens;
+}
+
+fn dl_json(t: &edge0_core::state::DownloadTask) -> DownloadEntry {
+    DownloadEntry {
+        id: t.id.clone(),
+        tier: t.tier.clone(),
+        rev: t.rev.clone(),
+        phase: t.phase.clone(),
+        source: t.source.clone(),
+        source_reason: t.source_reason.clone(),
+        bytes_done: t.bytes_done,
+        bytes_total: t.bytes_total,
+        paused: t.paused,
+        failed: t.failed.as_ref().map(|f| DownloadFailure {
+            code: f.code.clone(),
+            message: f.message.clone(),
+        }),
+        rate_bps: t.rate_bps,
+        eta_s: t.eta_s,
+        source_switched: t.source_switched,
+    }
+}
+
+fn hardware_snapshot() -> HardwareInfo {
+    let chip = sysctl_string("machdep.cpu.brand_string").unwrap_or_else(|| "unknown".into());
+    let macos = sysctl_string("kern.osproductVersion").unwrap_or_else(|| "unknown".into());
+    HardwareInfo {
+        chip,
+        macos,
+        mem_total_gib: sysctl_u64("hw.memsize")
+            .map(|b| (b as f64 / (1u64 << 30) as f64 * 10.0).round() / 10.0),
+        metal_device: false,
+        disk_free_gib: disk_free_gib(&std::env::current_dir().unwrap_or_default()),
+    }
+}
+
+pub async fn doctor(State(app): State<AppRef>) -> Json<serde_json::Value> {
+    let mut checks = Vec::new();
+    let chip = sysctl_string("machdep.cpu.brand_string").unwrap_or_default();
+    let m_gen = extract_m_gen(&chip);
+    checks.push(serde_json::json!({
+        "id": "chip",
+        "verdict": if m_gen.is_some_and(|g| g >= 3) { "pass" } else { "warn" },
+        "detail": format!("{chip} (baseline Apple silicon M3 or newer)"),
+    }));
+    let macos = sysctl_string("kern.osproductVersion").unwrap_or_default();
+    let v_ok = macos
+        .split('.')
+        .next()
+        .and_then(|x| x.parse::<u32>().ok())
+        .is_some_and(|x| x >= 14);
+    checks.push(serde_json::json!({ "id": "macos", "verdict": if v_ok { "pass" } else { "warn" }, "detail": format!("{macos} (>= 14)") }));
+    let metal_detail = match app.pool.resident().first().cloned() {
+        None => "No resident worker; engine evidence unavailable. Replay does not touch Metal; hardware check requires the native engine."
+            .to_string(),
+        Some((tier, _, _)) => match app.pool.stats(&tier).await {
+            Some(j) => {
+                let v: serde_json::Value = serde_json::from_str(&j).unwrap_or_default();
+                format!(
+                    "engine={} (stats evidence; Metal check requires the native engine)",
+                    v["engine"].as_str().unwrap_or("?")
+                )
+            }
+            None => "Engine stats unavailable (reported as unknown).".to_string(),
+        },
+    };
+    checks.push(serde_json::json!({ "id": "metal", "verdict": "pass", "detail": metal_detail, }));
+    checks.push(serde_json::json!({
+        "id": "port", "verdict": "pass",
+        "detail": format!("{}:{} is owned by this service (no automatic port fallback)", app.cfg.bind, app.cfg.port),
+    }));
+    let disk = disk_free_gib(app.home.root.as_path());
+    let (d_verdict, d_detail, d_code, d_next) = if disk >= 25.7 {
+        (
+            "pass",
+            format!("{disk:.1} GiB free >= the 35b preflight threshold (about 25.7 GiB)"),
+            None,
+            None,
+        )
+    } else if disk >= 7.6 {
+        (
+            "warn",
+            format!("{disk:.1} GiB free < the 35b preflight threshold (about 25.7 GiB)"),
+            Some("E-DL-DISK"),
+            Some("Free disk space, or install only edge0-8b (about 7.5 GiB required)"),
+        )
+    } else {
+        (
+            "warn",
+            format!("{disk:.1} GiB free < the 8b preflight threshold (about 7.5 GiB)"),
+            Some("E-DL-DISK"),
+            Some(
+                "Free disk space or move EDGE0_HOME to another volume; edge0 gc can reclaim blobs",
+            ),
+        )
+    };
+    let mut chk = serde_json::json!({ "id": "disk", "verdict": d_verdict, "detail": d_detail });
+    if let Some(c) = d_code {
+        chk["code"] = serde_json::json!(c);
+        chk["next"] = serde_json::json!(d_next.unwrap_or(""));
+    }
+    checks.push(chk);
+    let agent = std::process::Command::new("/bin/launchctl")
+        .args([
+            "print",
+            &format!("gui/{}/com.edge0.daemon", unsafe { libc::getuid() }),
+        ])
+        .output();
+    let agent_running = agent.map(|o| o.status.success()).unwrap_or(false);
+    checks.push(serde_json::json!({
+        "id": "autostart",
+        "verdict": if agent_running { "warn" } else { "pass" },
+        "detail": if agent_running {
+            String::from("LaunchAgent is registered while scope=menu-agent: two startup paths are present")
+        } else {
+            String::from("Startup matches the current scope; no LaunchAgent residue")
+        },
+    }));
+    let mut models_detail = Vec::new();
+    let mut models_ok = true;
+    for t in app.registry.lock().unwrap().tiers.iter() {
+        if let Err(detail) = verify_installed_tier(t) {
+            models_ok = false;
+            models_detail.push(detail);
+        }
+    }
+    checks.push(serde_json::json!({
+        "id": "models",
+        "verdict": if models_ok { "pass" } else { "warn" },
+        "detail": if models_ok {
+            format!("Installed tier integrity OK ({} tiers)", app.registry.lock().unwrap().tiers.len())
+        } else {
+            models_detail.join("; ")
+        },
+        "code": if models_ok { serde_json::Value::Null } else { serde_json::json!("E-MODEL-HASH") },
+    }));
+    let overall = if checks.iter().any(|c| c["verdict"] == "warn") {
+        "warn"
+    } else {
+        "pass"
+    };
+    Json(serde_json::json!({ "overall": overall, "checks": checks }))
+}
+
+/// Verify the installed-file record according to the installation mode.
+///
+/// Registry installs keep a signed manifest; direct HF installs keep an
+/// install.json snapshot instead. Treating the latter as a missing manifest
+/// produces a false E-MODEL-HASH even though the same files are loadable.
+fn verify_installed_tier(t: &edge0_core::state::InstalledTier) -> Result<(), String> {
+    match t.mode.as_str() {
+        "direct" => {
+            let snap = t
+                .install
+                .as_ref()
+                .ok_or_else(|| format!("{} install record is unreadable", t.tier))?;
+            for f in &snap.files {
+                let path = t.dir.join("files").join(&f.path);
+                let meta = std::fs::metadata(&path)
+                    .map_err(|_| format!("{}@{} file missing ({})", t.tier, t.rev, f.path))?;
+                if !meta.is_file() || meta.len() != f.size {
+                    return Err(format!(
+                        "{}@{} integrity check failed ({})",
+                        t.tier, t.rev, f.path
+                    ));
+                }
+            }
+            Ok(())
+        }
+        "registry" => {
+            let manifest = edge0_core::manifest::load_copy(&t.dir.join("manifest.json"))
+                .map_err(|_| format!("{} manifest is unreadable", t.tier))?;
+            let entry = manifest
+                .find_tier(&t.tier)
+                .filter(|entry| entry.rev == t.rev)
+                .ok_or_else(|| {
+                    format!("{} manifest does not contain current rev {}", t.tier, t.rev)
+                })?;
+            for f in &entry.files {
+                let path = t.dir.join("files").join(&f.path);
+                let bytes = std::fs::read(&path)
+                    .map_err(|_| format!("{}@{} file missing ({})", t.tier, t.rev, f.path))?;
+                if hex(sha2::Sha256::digest(&bytes)) != f.sha256 {
+                    return Err(format!(
+                        "{}@{} integrity check failed ({})",
+                        t.tier, t.rev, f.path
+                    ));
+                }
+            }
+            Ok(())
+        }
+        mode => Err(format!(
+            "{} install record has unknown mode ({})",
+            t.tier, mode
+        )),
+    }
+}
+
+// ---------------------------------------------------------------- load / unload / downloads
+
+#[derive(serde::Deserialize)]
+pub struct LoadIn {
+    keep_alive: Option<String>,
+}
+
+pub async fn load(
+    State(app): State<AppRef>,
+    Path(id): Path<String>,
+    Json(body): Json<LoadIn>,
+) -> Response {
+    if app.registry.lock().unwrap().find(&id).is_none() {
+        return render_error(ApiError::new(
+            code("E-MODEL-MISSING"),
+            format!("Tier {id} is not installed. Run `edge0 pull {id}`."),
+        ));
+    }
+    if let Some(k) = &body.keep_alive {
+        if keepalive::parse(k).is_err() {
+            return render_error(
+                ApiError::new(
+                    code("E-SRV-PARAM"),
+                    format!("Invalid keep_alive: {k} (use 5m/120s/1h/inf)"),
+                )
+                .with_param("keep_alive"),
+            );
+        }
+    }
+    let app2 = app.clone();
+    let tier = id.clone();
+    let keep = body
+        .keep_alive
+        .as_deref()
+        .and_then(|s| keepalive::parse(s).ok())
+        .flatten();
+    tokio::spawn(async move {
+        if let Ok(()) = app2.pool.ensure_loaded(&tier).await {
+            app2.pool.touch_keepalive(&tier, keep);
+        }
+    });
+    (
+        axum::http::StatusCode::ACCEPTED,
+        Json(serde_json::json!({ "model": id, "status": "loading" })),
+    )
+        .into_response()
+}
+
+pub async fn unload(State(app): State<AppRef>, Path(id): Path<String>) -> Response {
+    if !app.pool.is_running(&id) {
+        return Json(serde_json::json!({ "model": id, "unloaded": false })).into_response();
+    }
+    match app.pool.unload(&id) {
+        Ok(_) => Json(serde_json::json!({ "model": id, "unloaded": true })).into_response(),
+        Err(active_msg) => render_error(ApiError::new(
+            code("E-SRV-CONFLICT"),
+            format!("Unload refused: {active_msg}; wait for or cancel active requests and retry"),
+        )),
+    }
+}
+
+pub async fn downloads_list(State(app): State<AppRef>) -> Json<serde_json::Value> {
+    Json(serde_json::json!({
+        "object": "list",
+        "data": app.downloads.list().iter().map(dl_json).collect::<Vec<_>>(),
+    }))
+}
+
+struct BusEvents {
+    bus: crate::bus::EventBus,
+    app: AppRef,
+}
+
+impl edge0_pull::runner::DownloadEvents for BusEvents {
+    fn event(
+        &self,
+        etype: &'static str,
+        task: &edge0_core::state::DownloadTask,
+        code: Option<&'static str>,
+    ) {
+        if etype == "download.completed" {
+            if let Ok(scan) = crate::registry::Registry::rescan(&self.app.home) {
+                self.app.registry.lock().unwrap().tiers = scan.tiers;
+            }
+        }
+        let mut payload = serde_json::json!({
+            "id": task.id, "tier": task.tier, "rev": task.rev, "phase": task.phase,
+            "source": task.source, "bytes_done": task.bytes_done, "bytes_total": task.bytes_total,
+            "paused": task.paused, "code": code,
+        });
+        if let Some(r) = task.rate_bps {
+            payload["rate_bps"] = serde_json::json!(r);
+        }
+        if let Some(e) = task.eta_s {
+            payload["eta_s"] = serde_json::json!(e);
+        }
+        if task.source_switched {
+            payload["source_switched"] = serde_json::json!(true);
+        }
+        self.bus.publish(etype, &task.id, payload);
+    }
+}
+
+pub fn spawn_download(app: AppRef, task: edge0_core::state::DownloadTask) {
+    let id = task.id.clone();
+    let pin = task.source.clone().or_else(|| app.cfg.source.clone());
+    let store: std::sync::Arc<dyn edge0_pull::runner::TaskStore> = app.downloads.clone();
+    let events: std::sync::Arc<dyn edge0_pull::runner::DownloadEvents> =
+        std::sync::Arc::new(BusEvents {
+            bus: app.bus.clone(),
+            app: app.clone(),
+        });
+    let runner = edge0_pull::runner::TaskRunner {
+        home: app.home.clone(),
+        fetcher: std::sync::Arc::new(edge0_pull::fetch::HttpFetcher::with_env_proxy()),
+        store,
+        events,
+        source_pin: pin,
+        parallel: 4,
+    };
+    let app2 = app.clone();
+    let id2 = id.clone();
+    let h = tokio::spawn(async move {
+        let _ = runner.run(task).await;
+        app2.runners.lock().unwrap().remove(&id2);
+    });
+    app.runners.lock().unwrap().insert(id, h);
+}
+
+#[derive(serde::Deserialize)]
+pub struct CreateIn {
+    tier: String,
+    #[serde(default)]
+    source: Option<String>,
+}
+
+pub async fn downloads_create(State(app): State<AppRef>, Json(body): Json<CreateIn>) -> Response {
+    let existing = app.downloads.list();
+    if let Some(t) = existing.iter().find(|t| t.failed.is_none()) {
+        return render_error(
+            ApiError::new(
+                code("E-SRV-CONFLICT"),
+                format!(
+                    "Download task {} is running (phase={}); wait for it to finish or DELETE it before creating another",
+                    t.id, t.phase
+                ),
+            )
+            .with_status(409),
+        );
+    }
+    let client = match reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(30))
+        .build()
+    {
+        Ok(c) => c,
+        Err(e) => return render_error(ApiError::new(code("E-SRV-PARAM"), e.to_string())),
+    };
+    let direct = app.cfg.pull_mode == edge0_core::config::PullMode::Direct;
+    let plan = match edge0_pull::catalog::plan_for_tier(
+        app.cfg.pull_mode,
+        &client,
+        &app.cfg.registry_urls,
+        &app.cfg.hf_base,
+        &body.tier,
+        &app.home,
+        chrono::Utc::now(),
+        daemon_version(),
+    )
+    .await
+    {
+        Ok(v) => v,
+        Err(e) => return render_error(e.into_api_error()),
+    };
+    let total = plan.bytes_total();
+    if let Err(e) = edge0_pull::registry_ops::preflight_check(&app.home, total) {
+        return render_error(e.into_api_error());
+    }
+    let now = chrono::Utc::now();
+    let (snapshot, task_plan, source) = match plan {
+        edge0_pull::catalog::CreationPlan::Registry(list, _) => {
+            (Some(list.manifest), None, body.source.clone())
+        }
+        edge0_pull::catalog::CreationPlan::Direct(snap) => (None, Some(snap), None),
+    };
+    let task = edge0_core::state::DownloadTask {
+        id: format!("dl-{}", &uuid::Uuid::new_v4().simple().to_string()[..12]),
+        tier: body.tier.clone(),
+        rev: snapshot
+            .as_ref()
+            .and_then(|m| m.find_tier(&body.tier).map(|t| t.rev.clone()))
+            .or_else(|| task_plan.as_ref().map(|p| p.commit.clone()))
+            .unwrap_or_default(),
+        phase: "probe".into(),
+        source,
+        bytes_total: total,
+        bytes_done: 0,
+        paused: false,
+        files: vec![],
+        created_at: now,
+        updated_at: now,
+        snapshot,
+        source_reason: direct.then(|| {
+            edge0_pull::catalog::CreationPlan::source_reason_direct(
+                &app.cfg.hf_base,
+                body.source.is_some(),
+            )
+        }),
+        failed: None,
+        rate_bps: None,
+        eta_s: None,
+        source_switched: false,
+        plan: task_plan,
+    };
+    if let Err(e) = app.downloads.upsert(&task) {
+        return render_error(ApiError::new(code("E-SRV-CONFLICT"), e).with_status(500));
+    }
+    spawn_download(app.clone(), task.clone());
+    (
+        axum::http::StatusCode::CREATED,
+        Json(serde_json::json!({
+            "id": task.id, "tier": task.tier, "rev": task.rev,
+            "bytes_total": total, "phase": "probe",
+            "task_state": format!("state/downloads/{}.json", task.id),
+            "manifest_snapshot_taken": task.snapshot.is_some(),
+        })),
+    )
+        .into_response()
+}
+
+fn task_404(id: &str) -> Response {
+    render_error(
+        ApiError::new(
+            code("E-SRV-PARAM"),
+            format!("Download task does not exist: {id}"),
+        )
+        .with_status(404),
+    )
+}
+
+pub async fn downloads_pause(State(app): State<AppRef>, Path(id): Path<String>) -> Response {
+    let Some(mut t) = app.downloads.get(&id) else {
+        return task_404(&id);
+    };
+    if let Some(h) = app.runners.lock().unwrap().remove(&id) {
+        h.abort();
+    }
+    if !t.paused {
+        t.paused = true;
+        t.updated_at = chrono::Utc::now();
+        if let Err(e) = app.downloads.upsert(&t) {
+            return render_error(ApiError::new(code("E-SRV-CONFLICT"), e).with_status(500));
+        }
+        app.bus.publish(
+            "download.paused",
+            &id,
+            serde_json::json!({ "id": id, "tier": t.tier, "phase": t.phase, "bytes_done": t.bytes_done }),
+        );
+    }
+    Json(serde_json::json!({ "id": id, "paused": true })).into_response()
+}
+
+pub async fn downloads_resume(State(app): State<AppRef>, Path(id): Path<String>) -> Response {
+    let Some(mut t) = app.downloads.get(&id) else {
+        return task_404(&id);
+    };
+    let running = {
+        let g = app.runners.lock().unwrap();
+        g.get(&id).map(|h| !h.is_finished()).unwrap_or(false)
+    };
+    if running {
+        return Json(serde_json::json!({ "id": id, "running": true })).into_response();
+    }
+    if let Some(other) = app
+        .downloads
+        .list()
+        .iter()
+        .find(|o| o.id != id && o.failed.is_none())
+    {
+        return render_error(
+            ApiError::new(
+                code("E-SRV-CONFLICT"),
+                format!(
+                    "Task {} is already running; only one task may run at a time",
+                    other.id
+                ),
+            )
+            .with_status(409),
+        );
+    }
+    t.paused = false;
+    t.failed = None;
+    t.updated_at = chrono::Utc::now();
+    if let Err(e) = app.downloads.upsert(&t) {
+        return render_error(ApiError::new(code("E-SRV-CONFLICT"), e).with_status(500));
+    }
+    spawn_download(app.clone(), t);
+    (
+        axum::http::StatusCode::ACCEPTED,
+        Json(serde_json::json!({ "id": id, "running": true })),
+    )
+        .into_response()
+}
+
+pub async fn downloads_delete(State(app): State<AppRef>, Path(id): Path<String>) -> Response {
+    if app.downloads.get(&id).is_none() {
+        return task_404(&id);
+    }
+    if let Some(h) = app.runners.lock().unwrap().remove(&id) {
+        h.abort();
+    }
+    let _ = app.downloads.delete(&id);
+    let prefix = format!("{id}.");
+    if let Ok(rd) = std::fs::read_dir(app.home.tmp()) {
+        for ent in rd.flatten() {
+            if ent.file_name().to_string_lossy().starts_with(&prefix) {
+                let _ = std::fs::remove_file(ent.path());
+            }
+        }
+    }
+    Json(serde_json::json!({ "id": id, "deleted": true })).into_response()
+}
+
+pub async fn models_delete(State(app): State<AppRef>, Path(id): Path<String>) -> Response {
+    let found = app.registry.lock().unwrap().find(&id).cloned();
+    let Some(installed) = found else {
+        return render_error(ApiError::new(
+            code("E-MODEL-MISSING"),
+            format!("Tier {id} is not installed"),
+        ));
+    };
+    if installed.bundled {
+        return render_error(ApiError::new(
+            code("E-SRV-CONFLICT"),
+            format!("Tier {id} is bundled with the app and cannot be deleted"),
+        ));
+    }
+    if app.pool.is_running(&id) {
+        if let Err(active) = app.pool.unload(&id) {
+            return render_error(ApiError::new(
+                code("E-SRV-CONFLICT"),
+                format!("Delete refused: {active}; wait for or cancel active requests and retry"),
+            ));
+        }
+    }
+    match edge0_pull::registry_ops::remove_tier(&app.home, &id) {
+        Err(e) => render_error(e.into_api_error()),
+        Ok(shas) => {
+            app.registry.lock().unwrap().tiers.retain(|t| t.tier != id);
+            Json(serde_json::json!({
+                "id": id, "removed": true, "refs_decremented": shas.len(),
+                "note": "Blob references were decremented; run edge0 gc to reclaim storage",
+            }))
+            .into_response()
+        }
+    }
+}
+
+fn sysctl_string(name: &str) -> Option<String> {
+    let c = std::ffi::CString::new(name).ok()?;
+    let mut len = 0usize;
+    unsafe {
+        if libc::sysctlbyname(
+            c.as_ptr(),
+            std::ptr::null_mut(),
+            &mut len,
+            std::ptr::null_mut(),
+            0,
+        ) != 0
+        {
+            return None;
+        }
+        let mut buf = vec![0u8; len];
+        if libc::sysctlbyname(
+            c.as_ptr(),
+            buf.as_mut_ptr() as *mut libc::c_void,
+            &mut len,
+            std::ptr::null_mut(),
+            0,
+        ) != 0
+        {
+            return None;
+        }
+        Some(String::from_utf8_lossy(&buf[..len.saturating_sub(1)]).into_owned())
+    }
+}
+
+fn sysctl_u64(name: &str) -> Option<u64> {
+    let c = std::ffi::CString::new(name).ok()?;
+    let mut val: u64 = 0;
+    let mut len = std::mem::size_of::<u64>();
+    unsafe {
+        if libc::sysctlbyname(
+            c.as_ptr(),
+            &mut val as *mut _ as *mut libc::c_void,
+            &mut len,
+            std::ptr::null_mut(),
+            0,
+        ) != 0
+        {
+            return None;
+        }
+        Some(val)
+    }
+}
+
+fn disk_free_gib(path: &std::path::Path) -> f64 {
+    let c = match std::ffi::CString::new(path.as_os_str().as_encoded_bytes().to_vec()) {
+        Ok(c) => c,
+        Err(_) => return 0.0,
+    };
+    let mut st = std::mem::MaybeUninit::<libc::statfs>::zeroed();
+    let rc = unsafe { libc::statfs(c.as_ptr(), st.as_mut_ptr()) };
+    if rc != 0 {
+        return 0.0;
+    }
+    let st = unsafe { st.assume_init() };
+    let bsize = st.f_bsize as f64;
+    let avail = st.f_bavail as f64;
+    avail * bsize / (1u64 << 30) as f64
+}
+
+fn extract_m_gen(chip: &str) -> Option<u32> {
+    let after = chip.rfind("M")?;
+    chip[after + 1..]
+        .chars()
+        .take_while(|c| c.is_ascii_digit())
+        .collect::<String>()
+        .parse()
+        .ok()
+}
+
+fn hex(d: impl AsRef<[u8]>) -> String {
+    d.as_ref().iter().map(|b| format!("{b:02x}")).collect()
+}
+
+#[allow(dead_code)]
+fn unix_now() -> u64 {
+    std::time::SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_secs())
+        .unwrap_or(0)
+}
diff --git a/macos/crates/edge0d/src/lifecycle.rs b/macos/crates/edge0d/src/lifecycle.rs
new file mode 100644
index 0000000..a69eae4
--- /dev/null
+++ b/macos/crates/edge0d/src/lifecycle.rs
@@ -0,0 +1,126 @@
+//! Cascade-stop watch, startup orphan sweep, and App Nap prevention.
+
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+
+use libloading::{Library, Symbol};
+
+use crate::app::AppRef;
+
+/// Under menu-agent: host app pid gone → reclaim workers and exit.
+pub fn spawn_owner_monitor(app: AppRef, stop: Arc<AtomicBool>) {
+    let Some(owner) = app.owner_pid else { return };
+    if owner == 0 {
+        return;
+    }
+    std::thread::spawn(move || {
+        while !stop.load(Ordering::Relaxed) {
+            std::thread::sleep(std::time::Duration::from_millis(200));
+            let alive = unsafe { libc::kill(owner as i32, 0) } == 0;
+            if !alive {
+                eprintln!("edge0d: host app (pid {owner}) gone → cascade stop");
+                app.pool.stop_all("cascade");
+                app.bus.flush_persist();
+                app.clear_daemon_json();
+                std::process::exit(0);
+            }
+        }
+    });
+}
+
+/// Kill only processes named `edge0-engine` whose path is in this daemon's directory.
+pub fn orphan_sweep() -> Vec<u32> {
+    let mut killed = Vec::new();
+    let exe_dir = match std::env::current_exe() {
+        Ok(p) => p.parent().map(|d| d.to_path_buf()).unwrap_or_default(),
+        Err(_) => return killed,
+    };
+    let mut pids = vec![0i32; 4096];
+    let sz = std::mem::size_of_val(pids.as_slice());
+    let n = unsafe { libc::proc_listallpids(pids.as_mut_ptr() as *mut libc::c_void, sz as i32) };
+    if n <= 0 {
+        return killed;
+    }
+    pids.truncate(n as usize / 4);
+    for pid in pids {
+        if pid <= 1 {
+            continue;
+        }
+        let mut buf = [0i8; 1024];
+        let len = unsafe { libc::proc_pidpath(pid, buf.as_mut_ptr() as *mut libc::c_void, 1024) };
+        if len <= 0 {
+            continue;
+        }
+        let path = String::from_utf8_lossy(unsafe {
+            std::slice::from_raw_parts(buf.as_ptr() as *const u8, len as usize)
+        })
+        .into_owned();
+        let name = std::path::Path::new(&path)
+            .file_name()
+            .map(|f| f.to_string_lossy().into_owned())
+            .unwrap_or_default();
+        if name == "edge0-engine" && std::path::Path::new(&path).parent() == Some(exe_dir.as_path())
+        {
+            if unsafe { libc::kill(pid, libc::SIGKILL) } == 0 {
+                killed.push(pid as u32);
+            }
+        }
+    }
+    killed
+}
+
+/// Hold an activity assertion for the process lifetime. On arm64, `objc_msgSend` must be
+/// called through a typed function pointer matching the selector arity (variadic SIGSEGVs).
+pub fn begin_activity_assertion() -> bool {
+    use std::ffi::c_void;
+    type Id = *mut c_void;
+    type Sel = *mut c_void;
+    type Msg0 = unsafe extern "C" fn(Id, Sel) -> Id;
+    type MsgStr = unsafe extern "C" fn(Id, Sel, *const i8) -> Id;
+    type MsgAct = unsafe extern "C" fn(Id, Sel, u64, Id) -> Id;
+
+    let Ok(lib) = (unsafe { Library::new("/usr/lib/libobjc.A.dylib") }) else {
+        return false;
+    };
+    let Ok(get_class): Result<Symbol<unsafe extern "C" fn(*const i8) -> Id>, _> =
+        (unsafe { lib.get(b"objc_getClass\0") })
+    else {
+        return false;
+    };
+    let Ok(sel_reg): Result<Symbol<unsafe extern "C" fn(*const i8) -> Sel>, _> =
+        (unsafe { lib.get(b"sel_registerName\0") })
+    else {
+        return false;
+    };
+    let Ok(msg): Result<Symbol<Msg0>, _> = (unsafe { lib.get(b"objc_msgSend\0") }) else {
+        return false;
+    };
+    let msg = *msg;
+    unsafe {
+        let info_cls = get_class(c"NSProcessInfo".as_ptr());
+        if info_cls.is_null() {
+            return false;
+        }
+        let process_info_sel = sel_reg(c"processInfo".as_ptr());
+        let info = msg(info_cls, process_info_sel);
+        let str_cls = get_class(c"NSString".as_ptr());
+        if str_cls.is_null() || info.is_null() {
+            return false;
+        }
+        let sws_sel = sel_reg(c"stringWithUTF8String:".as_ptr());
+        let reason_c =
+            std::ffi::CString::new("edge0d resident service: disable App Nap throttling").unwrap();
+        let sws: MsgStr = std::mem::transmute::<Msg0, MsgStr>(msg);
+        let reason_str = sws(str_cls, sws_sel, reason_c.as_ptr());
+        let begin_sel = sel_reg(c"beginActivityWithOptions:reason:".as_ptr());
+        let begin: MsgAct = std::mem::transmute::<Msg0, MsgAct>(msg);
+        // NSActivityUserInitiatedAllowingIdleSystemSleep = 0x00FFFFFF
+        let activity = begin(info, begin_sel, 0x00FFFFFFu64, reason_str);
+        if activity.is_null() {
+            return false;
+        }
+        // Retain forever: a leaked handle keeps the assertion until process exit.
+        let _ = activity;
+        true
+    }
+}
diff --git a/macos/crates/edge0d/src/main.rs b/macos/crates/edge0d/src/main.rs
new file mode 100644
index 0000000..f461520
--- /dev/null
+++ b/macos/crates/edge0d/src/main.rs
@@ -0,0 +1,332 @@
+//! `edge0d` entry: assemble App, attach lifecycle hooks, serve HTTP.
+
+mod app;
+mod bus;
+mod http;
+mod lifecycle;
+mod registry;
+mod workers;
+
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use clap::Parser;
+
+use edge0_core::config::{Config, Overrides};
+use edge0_core::paths::Home;
+use edge0_core::state::Scope;
+
+#[derive(Parser, Debug)]
+#[command(name = "edge0d", version)]
+struct Args {
+    #[arg(long, default_value = "foreground")]
+    scope: String,
+    /// Host app pid under menu-agent (defaults to getppid).
+    #[arg(long)]
+    owner_pid: Option<u32>,
+    #[arg(long, default_value = "cli")]
+    started_by: String,
+    #[arg(long)]
+    port: Option<u16>,
+    #[arg(long)]
+    bind: Option<String>,
+    #[arg(long)]
+    origins: Option<String>,
+    #[arg(long)]
+    concurrency: Option<u32>,
+    #[arg(long)]
+    keep_alive: Option<String>,
+    #[arg(long)]
+    home: Option<String>,
+    /// Print the token once after bind (for copying to a phone; only set on non-loopback).
+    #[arg(long)]
+    print_token: bool,
+    /// Headless opt-in: install/uninstall LaunchAgent then exit (does not start the service).
+    #[arg(long)]
+    enable_agent: bool,
+    #[arg(long)]
+    disable_agent: bool,
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let args = Args::parse();
+    if args.enable_agent {
+        return agent_install(true);
+    }
+    if args.disable_agent {
+        return agent_install(false);
+    }
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()?;
+    rt.block_on(run(args))
+}
+
+async fn run(args: Args) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let scope = Scope::parse(&args.scope)
+        .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { e.into() })?;
+    let overrides = Overrides {
+        port: args.port,
+        bind: args.bind.clone(),
+        origins: args.origins.clone(),
+        concurrency: args.concurrency,
+        keep_alive: args.keep_alive.clone(),
+        home: args.home.clone(),
+        ..Default::default()
+    };
+    let home = match &args.home {
+        Some(h) => Home::new(h),
+        None => Home::from_env()?,
+    };
+    home.ensure_dirs()?;
+    init_tracing(&home);
+
+    let cfg = Config::load(&overrides, Some(&home.config_toml()))
+        .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { e.into() })?;
+
+    let swept = lifecycle::orphan_sweep();
+    if !swept.is_empty() {
+        tracing::warn!(pids = ?swept, "reaped leaked worker orphans from a previous run");
+    }
+
+    let activity_held = lifecycle::begin_activity_assertion();
+    tracing::info!(activity_held, "App Nap activity assertion registered");
+
+    let owner_pid = match scope {
+        Scope::MenuAgent => Some(
+            args.owner_pid
+                .unwrap_or_else(|| unsafe { libc::getppid() } as u32),
+        ),
+        _ => args.owner_pid,
+    };
+
+    let app = match app::App::new(
+        cfg.clone(),
+        home.clone(),
+        scope,
+        args.started_by.clone(),
+        owner_pid,
+    ) {
+        Ok(a) => Arc::new(a),
+        Err(e) => {
+            eprintln!("{}", serde_json::to_string(&e.envelope())?);
+            std::process::exit(1);
+        }
+    };
+    app.write_daemon_json();
+    if args.print_token {
+        if let Some(t) = &app.token {
+            println!("{t}");
+        }
+    }
+
+    let stop = app.stop.clone();
+    workers::spawn_maintenance(app.pool.clone(), stop.clone());
+    bus::spawn_persist_loop(app.bus.clone(), stop.clone());
+    lifecycle::spawn_owner_monitor(app.clone(), stop.clone());
+
+    app.bus.publish(
+        "service.restart",
+        "*",
+        serde_json::json!({ "by": app.started_by, "session_id": app.session_id }),
+    );
+
+    let listener = {
+        let sa = cfg.socket_addr();
+        match tokio::net::TcpListener::bind(sa).await {
+            Ok(l) => l,
+            Err(e) => {
+                eprintln!("edge0d: bind {sa} failed (port is not changed automatically): {e}");
+                std::process::exit(1);
+            }
+        }
+    };
+    tracing::info!(port = cfg.port, bind = %cfg.bind, scope = args.scope.as_str(), session = %app.session_id, "edge0d ready");
+
+    let server_app = app.clone();
+    let shutdown_app = app.clone();
+    let shutdown_stop = stop.clone();
+    let axum_app =
+        http::router(server_app).into_make_service_with_connect_info::<std::net::SocketAddr>();
+    let serve = axum::serve(listener, axum_app).with_graceful_shutdown(async move {
+        shutdown_signal().await;
+        shutdown_stop.store(true, std::sync::atomic::Ordering::Relaxed);
+        shutdown_app.pool.stop_all("shutdown");
+        shutdown_app.bus.flush_persist();
+        shutdown_app.clear_daemon_json();
+    });
+    serve.await?;
+    Ok(())
+}
+
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        let _ = tokio::signal::ctrl_c().await;
+    };
+    #[cfg(unix)]
+    let term = async {
+        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
+            Ok(mut sig) => {
+                sig.recv().await;
+            }
+            Err(_) => std::future::pending::<()>().await,
+        }
+    };
+    #[cfg(unix)]
+    tokio::select! { _ = ctrl_c => {}, _ = term => {} }
+    #[cfg(not(unix))]
+    ctrl_c.await;
+}
+
+fn init_tracing(home: &Home) {
+    use tracing_subscriber::layer::SubscriberExt;
+    let file = RotatingFile {
+        path: home.logs().join(format!(
+            "edge0d.{}.log",
+            chrono::Local::now().format("%Y-%m-%d")
+        )),
+        state: Arc::new(std::sync::Mutex::new(None)),
+    };
+    let sub = tracing_subscriber::registry()
+        .with(
+            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
+        )
+        .with(tracing_subscriber::fmt::layer().with_writer(file))
+        .with(tracing_subscriber::fmt::layer().with_writer(std::io::stderr));
+    let _ = tracing::subscriber::set_global_default(sub);
+}
+
+/// 10 MB rotate, keep 5 files. MakeWriter shares Arc state with the file write.
+#[derive(Clone)]
+struct RotatingFile {
+    path: PathBuf,
+    state: Arc<std::sync::Mutex<Option<std::fs::File>>>,
+}
+
+impl std::io::Write for RotatingFile {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let mut g = self.state.lock().unwrap();
+        let f = match &mut *g {
+            Some(f) => f,
+            None => {
+                if let Some(parent) = self.path.parent() {
+                    std::fs::create_dir_all(parent).ok();
+                }
+                let f = std::fs::OpenOptions::new()
+                    .create(true)
+                    .append(true)
+                    .open(&self.path)?;
+                *g = Some(f);
+                g.as_mut().unwrap()
+            }
+        };
+        let n = f.write(buf)?;
+        if f.metadata().map(|m| m.len() > 10_485_760).unwrap_or(false) {
+            rotate(&self.path);
+            *g = None;
+        }
+        Ok(n)
+    }
+    fn flush(&mut self) -> std::io::Result<()> {
+        if let Some(f) = &mut *self.state.lock().unwrap() {
+            f.flush()?;
+        }
+        Ok(())
+    }
+}
+
+impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for RotatingFile {
+    type Writer = RotatingFile;
+    fn make_writer(&'a self) -> Self::Writer {
+        self.clone()
+    }
+}
+
+fn rotate(path: &std::path::Path) {
+    for i in (1..5).rev() {
+        let from = suffix(path, i);
+        let to = suffix(path, i + 1);
+        if from.exists() {
+            let _ = std::fs::rename(&from, &to);
+        }
+    }
+    let _ = std::fs::rename(path, suffix(path, 1));
+}
+
+fn suffix(path: &std::path::Path, i: u32) -> PathBuf {
+    let name = path.file_name().unwrap().to_string_lossy().to_string();
+    path.with_file_name(format!("{name}.{i}"))
+}
+
+fn agent_install(enable: bool) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let uid = unsafe { libc::getuid() };
+    let label = "com.edge0.daemon";
+    let plist_dir = {
+        let home = std::env::var_os("HOME").ok_or("HOME is unset")?;
+        PathBuf::from(home).join("Library/LaunchAgents")
+    };
+    if enable {
+        std::fs::create_dir_all(&plist_dir)?;
+        let exe = std::env::current_exe()?;
+        let bundle_env = exe
+            .parent()
+            .and_then(|bin| bin.parent().map(|resources| (bin, resources)))
+            .filter(|(_, resources)| {
+                resources.file_name().and_then(|n| n.to_str()) == Some("Resources")
+            })
+            .map(|(bin, resources)| {
+                format!(
+                    r#"  <key>EnvironmentVariables</key><dict>
+    <key>EDGE0_ENGINE_LIB</key><string>{}</string>
+    <key>EDGE0_WORKER_BIN</key><string>{}</string>
+  </dict>
+"#,
+                    resources.join("lib/libedge0_engine_native.dylib").display(),
+                    bin.join("edge0-engine").display()
+                )
+            })
+            .unwrap_or_default();
+        let plist = format!(
+            r#"<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0"><dict>
+  <key>Label</key><string>{label}</string>
+  <key>ProgramArguments</key><array>
+    <string>{}</string><string>--scope</string><string>launchagent</string><string>--started-by</string><string>launchd</string>
+  </array>
+  <key>RunAtLoad</key><true/><key>KeepAlive</key><true/>
+{bundle_env}  <key>StandardOutPath</key><string>/tmp/edge0d.launchagent.out</string>
+  <key>StandardErrorPath</key><string>/tmp/edge0d.launchagent.err</string>
+</dict></plist>
+"#,
+            exe.display()
+        );
+        let path = plist_dir.join(format!("{label}.plist"));
+        std::fs::write(&path, plist)?;
+        // bootstrap fails if already loaded; bootout then bootstrap is idempotent.
+        let _ = std::process::Command::new("/bin/launchctl")
+            .args(["bootout", &format!("gui/{uid}/{label}")])
+            .status();
+        let st = std::process::Command::new("/bin/launchctl")
+            .args([
+                "bootstrap",
+                &format!("gui/{uid}"),
+                path.to_string_lossy().as_ref(),
+            ])
+            .status()?;
+        if !st.success() {
+            return Err("launchctl bootstrap failed".into());
+        }
+        println!("LaunchAgent installed (idempotent): {path:?}");
+    } else {
+        let _ = std::process::Command::new("/bin/launchctl")
+            .args(["bootout", &format!("gui/{uid}/{label}")])
+            .status();
+        let path = plist_dir.join(format!("{label}.plist"));
+        if path.exists() {
+            std::fs::remove_file(&path)?;
+        }
+        println!("LaunchAgent removed (idempotent)");
+    }
+    Ok(())
+}
diff --git a/macos/crates/edge0d/src/registry.rs b/macos/crates/edge0d/src/registry.rs
new file mode 100644
index 0000000..163313d
--- /dev/null
+++ b/macos/crates/edge0d/src/registry.rs
@@ -0,0 +1,149 @@
+//! Registry and download-task store. Startup rescan builds the install view; unowned `tmp/` leftovers are marked, not deleted.
+
+use std::sync::Mutex;
+
+use chrono::{DateTime, Utc};
+use serde::Serialize;
+
+use edge0_core::paths::Home;
+use edge0_core::state::{DownloadTask, InstalledTier};
+
+#[derive(Default)]
+pub struct Registry {
+    pub tiers: Vec<InstalledTier>,
+    pub downloads: Vec<DownloadTask>,
+}
+
+impl Registry {
+    pub fn rescan(home: &Home) -> Result<Self, String> {
+        let tiers = edge0_core::manifest::scan_installed(home);
+        let mut downloads = Vec::new();
+        let dl_dir = home.state().join("downloads");
+        if dl_dir.is_dir() {
+            for ent in std::fs::read_dir(&dl_dir).map_err(|e| e.to_string())? {
+                let p = ent.map_err(|e| e.to_string())?.path();
+                if p.extension().and_then(|x| x.to_str()) != Some("json") {
+                    continue;
+                }
+                match std::fs::read_to_string(&p)
+                    .ok()
+                    .and_then(|t| serde_json::from_str::<DownloadTask>(&t).ok())
+                {
+                    Some(mut task) => {
+                        // Interrupted mid-fetch: keep ranges, leave phase visible, do not delete.
+                        if task.phase == "fetch" {
+                            task.paused = true;
+                        }
+                        downloads.push(task);
+                    }
+                    None => {
+                        let mut orphans = OrphanMark {
+                            paths: vec![],
+                            marked_at: Utc::now(),
+                        };
+                        orphans.paths.push(p.display().to_string());
+                        let _ = write_orphans(home, &orphans);
+                    }
+                }
+            }
+        }
+        let tmp = home.tmp();
+        if tmp.is_dir() {
+            let live: Vec<String> = downloads.iter().map(|d| d.id.clone()).collect();
+            let mut orphans = OrphanMark {
+                paths: vec![],
+                marked_at: Utc::now(),
+            };
+            if let Ok(rd) = std::fs::read_dir(&tmp) {
+                for ent in rd.flatten() {
+                    let name = ent.file_name().to_string_lossy().to_string();
+                    let owned = name.starts_with("dl-")
+                        && live
+                            .iter()
+                            .any(|id| name.starts_with(&format!("{id}.")) || name == *id);
+                    if !owned {
+                        orphans.paths.push(ent.path().display().to_string());
+                    }
+                }
+            }
+            if !orphans.paths.is_empty() {
+                write_orphans(home, &orphans)?;
+            }
+        }
+        downloads.sort_by_key(|d| d.id.clone());
+        Ok(Self { tiers, downloads })
+    }
+
+    pub fn find(&self, tier: &str) -> Option<&InstalledTier> {
+        self.tiers.iter().find(|t| t.tier == tier)
+    }
+
+    /// Shared model directory: current installed rev's `files/`. None if not installed.
+    pub fn engine_dir(&self, tier: &str) -> Option<std::path::PathBuf> {
+        self.find(tier).map(|t| t.dir.join("files"))
+    }
+}
+
+#[derive(Serialize)]
+struct OrphanMark {
+    paths: Vec<String>,
+    marked_at: DateTime<Utc>,
+}
+
+fn write_orphans(home: &Home, mark: &OrphanMark) -> Result<(), String> {
+    edge0_core::state::atomic_write_json(&home.state().join("tmp-orphans.json"), mark)
+}
+
+/// Held for the daemon lifetime; CLI cannot bypass the lock while the daemon is up.
+pub struct HeldLock {
+    _lock: edge0_core::state::RegistryLock,
+}
+
+pub fn acquire_lock(home: &Home) -> Result<HeldLock, String> {
+    edge0_core::state::RegistryLock::acquire(&home.registry_lock()).map(|l| HeldLock { _lock: l })
+}
+
+pub struct DownloadStore {
+    home: Home,
+    pub inner: Mutex<Vec<DownloadTask>>,
+}
+
+impl DownloadStore {
+    pub fn new(home: &Home, tasks: Vec<DownloadTask>) -> Self {
+        Self {
+            home: home.clone(),
+            inner: Mutex::new(tasks),
+        }
+    }
+    pub fn list(&self) -> Vec<DownloadTask> {
+        self.inner.lock().unwrap().clone()
+    }
+    pub fn get(&self, id: &str) -> Option<DownloadTask> {
+        self.inner
+            .lock()
+            .unwrap()
+            .iter()
+            .find(|t| t.id == id)
+            .cloned()
+    }
+    pub fn upsert(&self, task: &DownloadTask) -> Result<(), String> {
+        edge0_core::state::atomic_write_json(&self.home.download_json(&task.id), task)?;
+        let mut g = self.inner.lock().unwrap();
+        match g.iter_mut().find(|t| t.id == task.id) {
+            Some(slot) => *slot = task.clone(),
+            None => g.push(task.clone()),
+        }
+        Ok(())
+    }
+}
+
+impl edge0_pull::runner::TaskStore for DownloadStore {
+    fn save(&self, task: &DownloadTask) -> Result<(), String> {
+        self.upsert(task)
+    }
+    fn delete(&self, id: &str) -> Result<(), String> {
+        let _ = std::fs::remove_file(self.home.download_json(id));
+        self.inner.lock().unwrap().retain(|t| t.id != id);
+        Ok(())
+    }
+}
diff --git a/macos/crates/edge0d/src/workers.rs b/macos/crates/edge0d/src/workers.rs
new file mode 100644
index 0000000..c4c02fe
--- /dev/null
+++ b/macos/crates/edge0d/src/workers.rs
@@ -0,0 +1,530 @@
+//! WorkerPool: spawn one worker per resident tier, route frames, heartbeat, and keep_alive reclaim.
+
+use std::collections::HashMap;
+use std::os::unix::net::{UnixListener, UnixStream};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use chrono::{DateTime, Utc};
+use tokio::sync::mpsc;
+
+use edge0_core::frames::{Channel, EndReason, Frame};
+use edge0_core::paths::Home;
+use edge0_core::state::WorkerJson;
+
+use crate::bus::EventBus;
+
+#[derive(Debug, Clone)]
+pub enum StreamItem {
+    Token {
+        token: i32,
+        channel: Channel,
+    },
+    Finished(EndReason),
+    /// Worker death / protocol error ends the SSE stream with a coded error.
+    Gone,
+}
+
+struct TierWorker {
+    child_pid: u32,
+    child: Mutex<std::process::Child>,
+    writer: Arc<Mutex<UnixStream>>,
+    streams: Arc<Mutex<HashMap<u64, mpsc::UnboundedSender<StreamItem>>>>,
+    next_id: Arc<std::sync::atomic::AtomicU64>,
+    last_pong: Arc<Mutex<Instant>>,
+    unload_at: Arc<Mutex<Option<Instant>>>,
+    active: Arc<std::sync::atomic::AtomicUsize>,
+    engine_lib: String,
+    started_at: DateTime<Utc>,
+    #[allow(dead_code)]
+    listener: UnixListener,
+}
+
+#[derive(Clone)]
+pub struct WorkerPool {
+    inner: Arc<Mutex<HashMap<String, Arc<TierWorker>>>>,
+    home: Home,
+    bus: EventBus,
+    engine_lib: std::path::PathBuf,
+    registry: Arc<Mutex<crate::registry::Registry>>,
+}
+
+impl WorkerPool {
+    pub fn new(
+        home: Home,
+        bus: EventBus,
+        engine_lib: std::path::PathBuf,
+        registry: Arc<Mutex<crate::registry::Registry>>,
+    ) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(HashMap::new())),
+            home,
+            bus,
+            engine_lib,
+            registry,
+        }
+    }
+
+    pub fn is_running(&self, tier: &str) -> bool {
+        self.inner.lock().unwrap().contains_key(tier)
+    }
+
+    pub fn resident(&self) -> Vec<(String, Option<DateTime<Utc>>, usize)> {
+        self.inner
+            .lock()
+            .unwrap()
+            .iter()
+            .map(|(t, w)| {
+                (
+                    t.clone(),
+                    w.unload_at.lock().unwrap().map(|i| {
+                        let left = i.saturating_duration_since(Instant::now());
+                        Utc::now()
+                            + chrono::Duration::from_std(left).unwrap_or(chrono::Duration::zero())
+                    }),
+                    w.active.load(std::sync::atomic::Ordering::Relaxed),
+                )
+            })
+            .collect()
+    }
+
+    /// Spawn + hello handshake + reader. Events: `model.load.started` → ready|failed.
+    /// Model dir is the current installed rev's `files/`. Missing install → E-MODEL-MISSING, no worker.
+    pub async fn ensure_loaded(self: &Arc<Self>, tier: &str) -> Result<(), String> {
+        if self.is_running(tier) {
+            return Ok(());
+        }
+        self.bus
+            .publish("model.load.started", tier, serde_json::json!({}));
+        // Direct installs: full hash for small files, three-block probe for large ones.
+        // Mismatch → E-MODEL-INVALID, no worker, no auto-redownload.
+        let pre_err = {
+            let inst = self.registry.lock().unwrap().find(tier).cloned();
+            match inst {
+                Some(it) if it.mode == "direct" && it.install.is_some() => {
+                    let snap = it.install.clone().unwrap();
+                    let rev_dir = it.dir.clone();
+                    match tokio::task::spawn_blocking(move || {
+                        edge0_pull::check::spot_check(&snap, &rev_dir)
+                    })
+                    .await
+                    {
+                        Ok(Ok(secs)) => {
+                            tracing::debug!("spot_check {tier} passed in {secs:.2}s");
+                            None
+                        }
+                        Ok(Err(e)) => Some((e.code, e.message)),
+                        Err(_) => Some(("E-INTERNAL", "spot-check JoinError".to_string())),
+                    }
+                }
+                _ => None,
+            }
+        };
+        let engine_dir = self.registry.lock().unwrap().engine_dir(tier);
+        let spawned = if let Some((c, m)) = pre_err {
+            Err((c, m))
+        } else {
+            match engine_dir {
+                Some(dir) => self.spawn(tier, &dir),
+                None => Err((
+                    "E-MODEL-MISSING",
+                    format!("tier {tier} is not in the install view (not installed or files/ missing)"),
+                )),
+            }
+        };
+        match spawned {
+            Ok(w) => {
+                self.inner
+                    .lock()
+                    .unwrap()
+                    .insert(tier.to_string(), Arc::new(w));
+                self.attach_reader(tier);
+                self.bus.publish(
+                    "model.load.ready",
+                    tier,
+                    serde_json::json!({ "unload_at": null }),
+                );
+                Ok(())
+            }
+            Err((code, e)) => {
+                self.bus.publish(
+                    "model.load.failed",
+                    tier,
+                    serde_json::json!({ "code": code, "detail": e }),
+                );
+                Err(e)
+            }
+        }
+    }
+
+    /// Spawn failures (socket/bind/handshake/process) map to E-MEM-LOAD; missing install is classified in `ensure_loaded`.
+    fn spawn(
+        &self,
+        tier: &str,
+        model_dir: &std::path::Path,
+    ) -> Result<TierWorker, (&'static str, String)> {
+        self.spawn_inner(tier, model_dir)
+            .map_err(|e| ("E-MEM-LOAD", e))
+    }
+
+    fn spawn_inner(&self, tier: &str, model_dir: &std::path::Path) -> Result<TierWorker, String> {
+        let sock = self.home.worker_sock(tier);
+        let _ = std::fs::remove_file(&sock);
+        std::fs::create_dir_all(self.home.run()).map_err(|e| e.to_string())?;
+        let listener = UnixListener::bind(&sock).map_err(|e| format!("bind {sock:?}: {e}"))?;
+        listener.set_nonblocking(true).map_err(|e| e.to_string())?;
+        let lib = std::env::var("EDGE0_ENGINE_LIB")
+            .map(std::path::PathBuf::from)
+            .unwrap_or_else(|_| self.engine_lib.clone());
+        let bin = std::env::var("EDGE0_WORKER_BIN")
+            .map(std::path::PathBuf::from)
+            .unwrap_or_else(|_| {
+                std::env::current_exe()
+                    .map(|p| p.parent().unwrap().join("edge0-engine"))
+                    .unwrap_or_else(|_| "edge0-engine".into())
+            });
+        let mut cmd = std::process::Command::new(&bin);
+        cmd.env("EDGE0_WORKER_SOCK", &sock)
+            .env("EDGE0_ENGINE_LIB", &lib)
+            .env("EDGE0_TIER", tier)
+            .env("EDGE0_MODEL_DIR", model_dir)
+            .stdout(std::process::Stdio::inherit())
+            .stderr(std::process::Stdio::inherit());
+        if let Some(rev) = self
+            .registry
+            .lock()
+            .unwrap()
+            .find(tier)
+            .map(|it| it.rev.clone())
+        {
+            cmd.env("EDGE0_MODEL_REV", rev);
+        }
+        let mut child = cmd.spawn().map_err(|e| format!("failed to spawn worker: {e}"))?;
+        let pid = child.id();
+
+        // Deadline on accept: a worker that never connects (or exits immediately) must fail, not hang.
+        let raw = loop {
+            match listener.accept() {
+                Ok((s, _)) => break s,
+                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
+                    if let Ok(Some(_)) = child.try_wait() {
+                        return Err("worker exited before connecting back".into());
+                    }
+                    std::thread::sleep(Duration::from_millis(20));
+                }
+                Err(e) => return Err(format!("accept failed: {e}")),
+            }
+        };
+        raw.set_nonblocking(false).map_err(|e| e.to_string())?;
+        let mut reader = std::io::BufReader::new(raw.try_clone().map_err(|e| e.to_string())?);
+        match frames_decode(&mut reader) {
+            Ok(Some(Frame::Hello {
+                abi, engine_lib, ..
+            })) if abi == edge0_core::version::E0_ABI_VERSION => {
+                let w = TierWorker {
+                    child_pid: pid,
+                    child: Mutex::new(child),
+                    writer: Arc::new(Mutex::new(raw)),
+                    streams: Arc::new(Mutex::new(HashMap::new())),
+                    next_id: Arc::new(std::sync::atomic::AtomicU64::new(1)),
+                    last_pong: Arc::new(Mutex::new(Instant::now())),
+                    unload_at: Arc::new(Mutex::new(None)),
+                    active: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
+                    engine_lib,
+                    started_at: Utc::now(),
+                    listener,
+                };
+                let _ = std::fs::write(
+                    self.home.worker_json(tier),
+                    serde_json::to_vec_pretty(&WorkerJson {
+                        tier: tier.to_string(),
+                        pid,
+                        engine_lib: w.engine_lib.clone(),
+                        abi_version: abi,
+                        started_at: w.started_at,
+                    })
+                    .unwrap(),
+                );
+                Ok(w)
+            }
+            Ok(Some(other)) => {
+                let _ = child.kill();
+                Err(format!("expected Hello frame, got {other:?}"))
+            }
+            Ok(None) => {
+                let _ = child.kill();
+                Err("worker disconnected before handshake".into())
+            }
+            Err(e) => {
+                let _ = child.kill();
+                Err(format!("frame protocol error: {e}"))
+            }
+        }
+    }
+
+    pub fn attach_reader(self: &Arc<Self>, tier: &str) {
+        let w = self
+            .inner
+            .lock()
+            .unwrap()
+            .get(tier)
+            .cloned()
+            .expect("already resident");
+        let pool = self.clone();
+        let tier_s = tier.to_string();
+        let raw = w.writer.lock().unwrap().try_clone().unwrap();
+        std::thread::spawn(move || {
+            let mut r = std::io::BufReader::new(raw);
+            loop {
+                match frames_decode(&mut r) {
+                    Ok(None) | Err(_) => break,
+                    Ok(Some(f)) => match f {
+                        Frame::Token { id, token, channel } => {
+                            route(&w, id, StreamItem::Token { token, channel })
+                        }
+                        Frame::Finished { id, reason } => {
+                            route(&w, id, StreamItem::Finished(reason))
+                        }
+                        Frame::Error {
+                            id: Some(id),
+                            e0_status,
+                            details_json,
+                        } => {
+                            tracing::error!(id, e0_status, details = %details_json, "worker generate error");
+                            route(&w, id, StreamItem::Gone);
+                        }
+                        Frame::StatsResp { stats_json } => {
+                            if let Some(tx) = STATS_WAIT.lock().unwrap().take() {
+                                let _ = tx.send(stats_json);
+                            }
+                        }
+                        Frame::Pong => *w.last_pong.lock().unwrap() = Instant::now(),
+                        Frame::Bye { .. } => break,
+                        _ => {}
+                    },
+                }
+            }
+            pool.handle_worker_dead(&tier_s, &w);
+        });
+    }
+
+    fn handle_worker_dead(&self, tier: &str, w: &TierWorker) {
+        let alive = self.inner.lock().unwrap().remove(tier).is_some();
+        if !alive {
+            return;
+        }
+        let affected = w.streams.lock().unwrap().len();
+        for tx in w.streams.lock().unwrap().values() {
+            let _ = tx.send(StreamItem::Gone);
+        }
+        w.streams.lock().unwrap().clear();
+        let _ = std::fs::remove_file(self.home.worker_json(tier));
+        let _ = std::fs::remove_file(self.home.worker_sock(tier));
+        self.bus.publish(
+            "worker.crashed",
+            tier,
+            serde_json::json!({ "affected_requests": affected }),
+        );
+    }
+
+    pub fn begin(
+        &self,
+        tier: &str,
+        prompt_tokens: Vec<i32>,
+        params_json: &str,
+    ) -> Result<(u64, mpsc::UnboundedReceiver<StreamItem>), String> {
+        let w = self
+            .inner
+            .lock()
+            .unwrap()
+            .get(tier)
+            .cloned()
+            .ok_or("worker is not resident")?;
+        let id = w.next_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        let (tx, rx) = mpsc::unbounded_channel();
+        w.streams.lock().unwrap().insert(id, tx);
+        w.active.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        let frame = Frame::Begin {
+            id,
+            prompt_tokens,
+            params_json: params_json.to_string(),
+        };
+        {
+            let mut g = w.writer.lock().unwrap();
+            edge0_core::frames::write_to(&mut *g, &frame).map_err(|e| e.to_string())?;
+        }
+        Ok((id, rx))
+    }
+
+    pub fn finish(&self, tier: &str, id: u64) {
+        if let Some(w) = self.inner.lock().unwrap().get(tier).cloned() {
+            w.streams.lock().unwrap().remove(&id);
+            w.active.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
+        }
+    }
+
+    pub fn cancel(&self, tier: &str, id: u64) {
+        if let Some(w) = self.inner.lock().unwrap().get(tier).cloned() {
+            let mut g = w.writer.lock().unwrap();
+            let _ = edge0_core::frames::write_to(&mut *g, &Frame::Cancel { id });
+        }
+    }
+
+    pub async fn stats(&self, tier: &str) -> Option<String> {
+        let w = self.inner.lock().unwrap().get(tier).cloned()?;
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        *STATS_WAIT.lock().unwrap() = Some(tx);
+        {
+            let mut g = w.writer.lock().unwrap();
+            edge0_core::frames::write_to(&mut *g, &Frame::Stats).ok()?;
+        }
+        tokio::time::timeout(Duration::from_secs(5), rx)
+            .await
+            .ok()
+            .and_then(|r| r.ok())
+    }
+
+    /// `inf` → None; otherwise records unload_at for `/system`.
+    pub fn touch_keepalive(&self, tier: &str, keep: Option<Duration>) {
+        if let Some(w) = self.inner.lock().unwrap().get(tier).cloned() {
+            *w.unload_at.lock().unwrap() = keep.map(|d| Instant::now() + d);
+        }
+    }
+
+    pub fn unload(&self, tier: &str) -> Result<usize, String> {
+        let w = self.inner.lock().unwrap().get(tier).cloned();
+        match w {
+            None => Err("not-loaded".into()),
+            Some(w) => {
+                let active = w.active.load(std::sync::atomic::Ordering::Relaxed);
+                if active > 0 {
+                    return Err(format!("{active} in-flight request(s)"));
+                }
+                self.terminate(tier, &w, "explicit-unload");
+                Ok(0)
+            }
+        }
+    }
+
+    fn terminate(&self, tier: &str, w: &TierWorker, reason: &str) {
+        // If a blocked write holds the writer lock, shutdown must not wait — go to SIGTERM.
+        if let Ok(mut g) = w.writer.try_lock() {
+            let _ = edge0_core::frames::write_to(
+                &mut *g,
+                &Frame::Bye {
+                    reason: reason.into(),
+                },
+            );
+        }
+        let _ = std::process::Command::new("/bin/kill")
+            .arg("-TERM")
+            .arg(w.child_pid.to_string())
+            .status();
+        let pid = w.child_pid;
+        let home = self.home.clone();
+        let tier_s = tier.to_string();
+        let reason_s = reason.to_string();
+        let bus = self.bus.clone();
+        let pool_streams = w.streams.clone();
+        std::thread::spawn(move || {
+            let deadline = Instant::now() + Duration::from_secs(5);
+            loop {
+                let alive = unsafe { libc::kill(pid as i32, 0) } == 0;
+                if !alive {
+                    break;
+                }
+                if Instant::now() >= deadline {
+                    let _ = std::process::Command::new("/bin/kill")
+                        .arg("-KILL")
+                        .arg(pid.to_string())
+                        .status();
+                    break;
+                }
+                std::thread::sleep(Duration::from_millis(50));
+            }
+            for tx in pool_streams.lock().unwrap().values() {
+                let _ = tx.send(StreamItem::Gone);
+            }
+            let _ = std::fs::remove_file(home.worker_json(&tier_s));
+            let _ = std::fs::remove_file(home.worker_sock(&tier_s));
+            bus.publish(
+                "model.unloaded",
+                &tier_s,
+                serde_json::json!({ "by": reason_s }),
+            );
+        });
+        self.inner.lock().unwrap().remove(tier);
+    }
+
+    /// Snapshot then drop the lock before terminate: terminate takes the same lock (nested lock would deadlock).
+    pub fn stop_all(&self, reason: &str) {
+        let victims: Vec<(String, Arc<TierWorker>)> = {
+            let g = self.inner.lock().unwrap();
+            g.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
+        };
+        for (t, w) in victims {
+            self.terminate(&t, &w, reason);
+        }
+    }
+}
+
+fn route(w: &TierWorker, id: u64, item: StreamItem) {
+    if let Some(tx) = w.streams.lock().unwrap().get(&id) {
+        let _ = tx.send(item);
+    }
+}
+
+fn frames_decode<R: std::io::BufRead>(r: &mut R) -> Result<Option<Frame>, String> {
+    edge0_core::frames::decode(r)
+}
+
+static STATS_WAIT: Mutex<Option<tokio::sync::oneshot::Sender<String>>> = Mutex::new(None);
+
+pub fn spawn_maintenance(pool: Arc<WorkerPool>, stop: Arc<std::sync::atomic::AtomicBool>) {
+    std::thread::spawn(move || {
+        while !stop.load(std::sync::atomic::Ordering::Relaxed) {
+            std::thread::sleep(Duration::from_secs(1));
+            let now = Instant::now();
+            let all: Vec<(String, Arc<TierWorker>)> = pool
+                .inner
+                .lock()
+                .unwrap()
+                .iter()
+                .map(|(k, v)| (k.clone(), v.clone()))
+                .collect();
+            for (tier, w) in all {
+                if let Ok(mut g) = w.writer.try_lock() {
+                    let _ = edge0_core::frames::write_to(&mut *g, &Frame::Ping);
+                }
+                if now.duration_since(*w.last_pong.lock().unwrap()) > Duration::from_secs(15) {
+                    let _ = std::process::Command::new("/bin/kill")
+                        .arg("-9")
+                        .arg(w.child_pid.to_string())
+                        .status();
+                    pool.handle_worker_dead(&tier, &w);
+                    continue;
+                }
+                if let Ok(Some(_)) = w.child.lock().unwrap().try_wait() {
+                    pool.handle_worker_dead(&tier, &w);
+                    continue;
+                }
+                let due = {
+                    let mut g = w.unload_at.lock().unwrap();
+                    match *g {
+                        Some(at)
+                            if now >= at
+                                && w.active.load(std::sync::atomic::Ordering::Relaxed) == 0 =>
+                        {
+                            *g = None;
+                            true
+                        }
+                        _ => false,
+                    }
+                };
+                if due {
+                    pool.terminate(&tier, &w, "keep_alive-expired");
+                }
+            }
+        }
+    });
+}
diff --git a/macos/crates/engine-native/CMakeLists.txt b/macos/crates/engine-native/CMakeLists.txt
new file mode 100644
index 0000000..d8de97d
--- /dev/null
+++ b/macos/crates/engine-native/CMakeLists.txt
@@ -0,0 +1,88 @@
+# edge0 native inference engine (C++/Metal on top of mlx::core).
+# Built by the top-level Makefile (`make engine`); not part of the cargo workspace.
+cmake_minimum_required(VERSION 3.24)
+project(edge0_engine_native LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(E0N_VERSION "0.1.0" CACHE STRING "edge0 product version embedded in the native engine")
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/out)
+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/out)
+
+# ---- mlx::core (vendored, v0.30.6). mlx defaults to a static build; the engine needs the dylib.
+set(BUILD_SHARED_LIBS ON CACHE BOOL "" FORCE)
+set(MLX_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+set(MLX_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
+set(MLX_BUILD_BENCHMARKS OFF CACHE BOOL "" FORCE)
+set(MLX_BUILD_PYTHON_BINDINGS OFF CACHE BOOL "" FORCE)
+# mlx 0.30.6 turns on NAX GEMM/qmm on Apple M5 (arch gen >= 17). Those kernels
+# skew edge0-8b first-token logits (fixture cosine ~0.85 vs ~0.9999 with NAX
+# off). The define must be set before add_subdirectory so both mlx and this
+# tree see it — device.h inlines is_nax_available().
+option(E0N_MLX_NAX "Enable mlx NAX kernels (numerically wrong for edge0-8b on M5)" OFF)
+if(NOT E0N_MLX_NAX)
+  add_compile_definitions(MLX_METAL_NO_NAX)
+endif()
+add_subdirectory(
+  ${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/mlx
+  ${CMAKE_BINARY_DIR}/mlx
+  EXCLUDE_FROM_ALL)
+
+# nlohmann_json is reused from the copy mlx already fetches.
+
+set(EDGE0_ABI_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/include)
+
+# ---- Tokenizer: Rust staticlib from crates/engine-tok.
+set(EDGE0_TOK_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../engine-tok/Cargo.toml)
+file(GLOB EDGE0_TOK_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/../engine-tok/src/*.rs)
+string(TOLOWER ${CMAKE_BUILD_TYPE} EDGE0_TOK_PROFILE)
+add_custom_command(
+  OUTPUT ${CMAKE_BINARY_DIR}/rust-${EDGE0_TOK_PROFILE}/libedge0_engine_tok.a
+  COMMAND ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${CMAKE_BINARY_DIR}/rust
+          cargo build --manifest-path ${EDGE0_TOK_MANIFEST} --${EDGE0_TOK_PROFILE}
+  COMMAND ${CMAKE_COMMAND} -E copy
+          ${CMAKE_BINARY_DIR}/rust/${EDGE0_TOK_PROFILE}/libedge0_engine_tok.a
+          ${CMAKE_BINARY_DIR}/rust-${EDGE0_TOK_PROFILE}/libedge0_engine_tok.a
+  DEPENDS ${EDGE0_TOK_SRCS} ${EDGE0_TOK_MANIFEST}
+  COMMENT "engine-tok (rust staticlib, ${EDGE0_TOK_PROFILE})")
+add_custom_target(engine_tok_rust
+  DEPENDS ${CMAKE_BINARY_DIR}/rust-${EDGE0_TOK_PROFILE}/libedge0_engine_tok.a)
+add_library(edge0_engine_tok STATIC IMPORTED GLOBAL)
+add_dependencies(edge0_engine_tok engine_tok_rust)
+set_target_properties(edge0_engine_tok PROPERTIES
+  IMPORTED_LOCATION ${CMAKE_BINARY_DIR}/rust-${EDGE0_TOK_PROFILE}/libedge0_engine_tok.a)
+
+# ---- Engine core.
+add_library(engine_core STATIC
+  src/memstat.cpp
+  src/safio.cpp
+  src/expertmap.cpp
+  src/hotstack.cpp
+  src/routing_replay.cpp
+  src/stepprefetch.cpp
+  src/qgemm.cpp
+  src/router.cpp
+  src/registry.cpp
+  src/dense.cpp
+  src/prerouter.cpp
+  src/fwd_common.cpp
+  src/lora.cpp
+  src/sampling.cpp
+  src/expertio.cpp
+  src/net35.cpp
+  src/net8b.cpp)
+target_include_directories(engine_core PUBLIC src ${EDGE0_ABI_INCLUDE_DIR}
+  ${CMAKE_BINARY_DIR}/_deps/metal_cpp-src)
+target_link_libraries(engine_core PUBLIC mlx nlohmann_json::nlohmann_json)
+
+# ---- Engine dylib exporting every symbol in engine-abi.h.
+add_library(edge0_engine_native SHARED src/abiface.cpp)
+target_compile_definitions(edge0_engine_native PRIVATE E0N_VERSION=\"${E0N_VERSION}\")
+target_link_libraries(edge0_engine_native PRIVATE engine_core edge0_engine_tok)
+set_target_properties(edge0_engine_native PROPERTIES
+  OUTPUT_NAME edge0_engine_native
+  BUILD_RPATH "@loader_path"
+  BUILD_RPATH_USE_ORIGIN ON
+  INSTALL_RPATH "@loader_path"
+  MACOSX_RPATH ON)
diff --git a/macos/crates/engine-native/include/engine-abi.h b/macos/crates/engine-native/include/engine-abi.h
new file mode 100644
index 0000000..ce4bce9
--- /dev/null
+++ b/macos/crates/engine-native/include/engine-abi.h
@@ -0,0 +1,146 @@
+/*
+ * C ABI between the daemon and the inference engine.
+ *
+ * Opaque handles; extensible structs start with `size` (caller fills sizeof).
+ * Compile-time E0_ABI_VERSION must match e0_abi_version() or load is refused.
+ * JSON outputs use size-query then write: buf == NULL returns needed length
+ * (excluding NUL); a short buf returns E0_ERR_BUFFER without truncation.
+ * Only pointers, integers, and UTF-8 JSON cross the boundary.
+ */
+#ifndef EDGE0_ENGINE_ABI_H
+#define EDGE0_ENGINE_ABI_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* Semantic changes bump this; appending fields is a size-compatible addition. */
+#define E0_ABI_VERSION 2u
+
+/* E0Engine binds weights and a Metal device; E0Tokenizer does not load weights. */
+typedef struct E0Engine    E0Engine;
+typedef struct E0Stream    E0Stream;
+typedef struct E0Tokenizer E0Tokenizer;
+
+/*
+ * Negative values are errors. E0_WOULD_BLOCK / E0_NOTIFIED are for e0_next
+ * poll/wait; the wait strategy (fd vs callback) is the daemon's concern.
+ */
+typedef enum E0Status {
+  E0_OK            = 0,
+  E0_NOTIFIED      = 1,  /* e0_next produced a token */
+  E0_WOULD_BLOCK   = 2,  /* no token yet; not an error */
+  E0_ERR_INVALID   = -1, /* bad argument, handle, or state */
+  E0_ERR_BUFFER    = -2, /* output buffer too small (size-query then write) */
+  E0_ERR_MODEL     = -3, /* weight / layout / hash problem */
+  E0_ERR_MEM       = -4, /* load or activation out of memory */
+  E0_ERR_GPU       = -5, /* Metal device or command error */
+  E0_ERR_INTERNAL  = -6, /* other; details_json has a readable message */
+  E0_ERR_CANCELLED = -7  /* stream was stopped by e0_cancel */
+} E0Status;
+
+typedef struct E0TokenOut {
+  uint32_t size;    /* caller fills sizeof(E0TokenOut) */
+  int32_t  token;   /* token id; undefined when FINISHED is set */
+  uint32_t flags;
+} E0TokenOut;
+
+#define E0_TOKEN_FINISHED 0x1u /* stream ended (natural EOS or cancel) */
+
+uint32_t e0_abi_version(void);
+const char *e0_engine_version(void);
+
+/* Engine surface: GB-scale, resident for the worker / keep_alive lifetime. */
+
+/* config_json: {"max_layers":..., "hot_n":..., ...}; keys may be added, not redefined. */
+E0Engine *e0_engine_create(const char *tier_id,
+                           const char *model_dir,
+                           const char *config_json);
+
+/* prompt_tokens are tokenized ids; params_json holds sampling and stop criteria.
+ * NULL on failure; details via e0_last_error. */
+E0Stream *e0_generate_begin(E0Engine *h,
+                            const int32_t *prompt_tokens,
+                            uint32_t n_tokens,
+                            const char *params_json);
+
+/* Non-blocking next token: E0_NOTIFIED (out valid) / E0_WOULD_BLOCK / E0_ERR_*.
+ * End of stream sets E0_TOKEN_FINISHED. */
+E0Status e0_next(E0Engine *h, E0Stream *stream, E0TokenOut *out);
+
+/* Cancel one generation. Idempotent. A cancelled stream ends with E0_TOKEN_FINISHED. */
+E0Status e0_cancel(E0Engine *h, E0Stream *stream);
+
+/*
+ * Runtime stats as JSON (size-query then write). Key names are frozen:
+ * adding a key must not change the meaning of existing keys.
+ * Required keys:
+ *   "route_source_by_layer" — who chose each layer's expert set:
+ *        {"predicted"|"routed"|"replayed", ...}; "replayed" only when the
+ *        layer's set comes from a reference-sample replay;
+ *   "loaded_experts" / "cache_hits" / "cache_misses";
+ *   "memory" — object with three orthogonal fields:
+ *        "resident_peak_bytes"   peak process resident_size, including
+ *            file-backed weight pages;
+ *        "expert_cache_heap_bytes"  explicit heap bytes of the streaming
+ *            expert cache (numerator of the resident-budget check);
+ *        "file_backed_bytes"     mmap-backed weight pages reclaimable by the OS.
+ *        Tier and doctor thresholds use
+ *        committed_resident = resident_peak − file_backed;
+ *   "prefix_cache" — independent of the expert-weight cache: enabled/budget_bytes,
+ *        hits/misses/cached_tokens, entries/bytes/evictions, and
+ *        last_cached_tokens/last_prefill_tokens/last_reason for the last request.
+ */
+E0Status e0_stats(E0Engine *h, char *buf, size_t buf_len, size_t *written);
+
+/*
+ * Per-request reset: drop the active stream, prerouter cross-token state
+ * (previous-token exec one-hot, linear_init, stager double-buffer), request
+ * KV/conv/delta, and the prefetch queue. Keep the bounded prefix_cache snapshot
+ * for the current model identity. The daemon calls this before each
+ * e0_generate_begin. Reset belongs on the ABI, not in policy. Snapshots from a
+ * different model identity, ABI generation, or worker must not be reused.
+ */
+E0Status e0_turn_reset(E0Engine *h);
+
+/* Last error as readable JSON (size-query then write). */
+E0Status e0_last_error(E0Engine *h, char *buf, size_t buf_len, size_t *written);
+
+void e0_engine_destroy(E0Engine *h);
+
+/* Tokenizer surface: MB-scale, no weights, independent of the worker lifetime.
+ * Must not take an engine handle — UI token counts (real ids, not char estimates)
+ * must not pin a GB-scale worker. Reads tokenizer.json, chat_template.jinja,
+ * and config.json under model_dir. */
+
+E0Tokenizer *e0_tok_create(const char *model_dir);
+
+/* text_json: {"text": "..."} or {"texts": [...]}; output {"tokens":[[...]]}. */
+E0Status e0_tok_tokenize(E0Tokenizer *t, const char *text_json,
+                         char *buf, size_t buf_len, size_t *written);
+
+/* messages_json: OpenAI messages array (or an object with messages and template
+ * variables); output is the rendered prompt text and token count. */
+E0Status e0_tok_apply_template(E0Tokenizer *t, const char *messages_json,
+                               char *buf, size_t buf_len, size_t *written);
+
+/*
+ * tokens_json: {"tokens": [id, ...]}; output {"text": "..."} (size-query then write).
+ * Decodes a complete token sequence. Incremental streaming render is the
+ * caller's job (decode the full sequence, diff against the previous text, and
+ * wait for the next token when a multi-byte sequence is split). Both the real
+ * engine and the replay engine must export this symbol (missing symbol = load fail).
+ */
+E0Status e0_tok_decode(E0Tokenizer *t, const char *tokens_json,
+                       char *buf, size_t buf_len, size_t *written);
+
+void e0_tok_destroy(E0Tokenizer *t);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* EDGE0_ENGINE_ABI_H */
diff --git a/macos/crates/engine-native/src/abiface.cpp b/macos/crates/engine-native/src/abiface.cpp
new file mode 100644
index 0000000..b9423bd
--- /dev/null
+++ b/macos/crates/engine-native/src/abiface.cpp
@@ -0,0 +1,946 @@
+#include "engine-abi.h"
+
+#include <algorithm>
+#include <cstring>
+#include <cassert>
+#include <cstdlib>
+#include <fstream>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <vector>
+
+#include <sys/stat.h>
+#include <dirent.h>
+
+#include "memstat.hpp"
+#include "fwd_common.hpp"
+#include "safio.hpp"
+#include "expertmap.hpp"
+#include "hotstack.hpp"
+#include "expertio.hpp"
+#include "net35.hpp"
+#include "net8b.hpp"
+#include "genstream.hpp"
+#include "prerouter.hpp"
+#include "prefixcache.hpp"
+#include "sampling.hpp"
+
+#include <nlohmann/json.hpp>
+
+extern "C" {
+typedef struct ETok ETok;
+void *etok_create(const char *model_dir);
+void etok_destroy(void *t);
+int etok_last_error(char *buf, size_t buf_len, size_t *written);
+int etok_tokenize(void *t, const char *text, int32_t *ids_out, size_t ids_cap,
+                  size_t *written);
+int etok_decode(void *t, const int32_t *ids, size_t n_ids, char *buf, size_t buf_len,
+                size_t *written);
+int etok_render(void *t, const char *ctx_json, char *buf, size_t buf_len, size_t *written);
+}
+
+namespace {
+
+std::string esc(const std::string &s) {
+  std::string o;
+  for (char c : s) {
+    if (c == '"' || c == '\\') { o += '\\'; o += c; }
+    else if (c == '\n') o += "\\n";
+    else o += c;
+  }
+  return o;
+}
+
+E0Status json_out(const std::string &json, char *buf, size_t buf_len, size_t *written) {
+  if (written) *written = json.size();
+  if (buf == nullptr) return E0_OK;
+  if (buf_len < json.size()) return E0_ERR_BUFFER;
+  std::memcpy(buf, json.data(), json.size());
+  return E0_OK;
+}
+
+bool is_dir(const char *p) {
+  struct stat st;
+  return p && stat(p, &st) == 0 && S_ISDIR(st.st_mode);
+}
+
+}  // namespace
+
+std::string etok_err_json() {
+  size_t w = 0;
+  etok_last_error(nullptr, 0, &w);
+  std::string msg(w, '\0');
+  if (w) etok_last_error(msg.data(), w, &w);
+  return std::string("{\"error\":\"E-INTERNAL\",\"message\":\"") + esc(msg) + "\"}";
+}
+
+namespace e0n { namespace gen { class StreamBase; } }
+
+struct E0Engine {
+  std::string tier;
+  std::string model_dir;
+  std::string config_json;
+  std::mutex mu;
+  std::string last_error;
+  std::shared_ptr<e0n::ModelFiles> files;
+  std::shared_ptr<e0n::ExpertMap> experts;
+  std::unique_ptr<e0n::PrefixCache> prefix_cache;
+  std::string model_identity;
+  std::unique_ptr<e0n::Net35> net35;
+  std::unique_ptr<e0n::State35> st35;
+  std::unique_ptr<e0n::Net8b> net8b;
+  std::unique_ptr<e0n::State8b> st8b;
+  std::vector<e0n::RouteRec> last_recs;
+  bool has_recs = false;
+  std::vector<int32_t> eos_ids;
+  e0n::gen::StreamBase *active = nullptr;
+};
+
+struct CachedState35 {
+  e0n::State35 state;
+  std::vector<float> logits;
+};
+
+struct CachedState8 {
+  e0n::State8b state;
+  std::vector<float> logits;
+};
+
+namespace {
+
+mlx::core::array snapshot_array(const mlx::core::array &a) {
+  auto out = mlx::core::copy(a);
+  mlx::core::eval(out);
+  return out;
+}
+
+template <typename T>
+std::vector<mlx::core::array> snapshot_arrays(const std::vector<mlx::core::array> &src) {
+  std::vector<mlx::core::array> out;
+  out.reserve(src.size());
+  for (const auto &a : src) out.push_back(snapshot_array(a));
+  return out;
+}
+
+std::shared_ptr<CachedState35> capture_state35(const e0n::State35 &src,
+                                                const std::vector<float> &logits) {
+  auto out = std::make_shared<CachedState35>();
+  out->state.pr = src.pr;
+  out->state.conv = snapshot_arrays<float>(src.conv);
+  out->state.recur = snapshot_arrays<float>(src.recur);
+  out->state.kc = snapshot_arrays<float>(src.kc);
+  out->state.vc = snapshot_arrays<float>(src.vc);
+  out->state.is_linear = src.is_linear;
+  out->state.position = src.position;
+  out->state.started = src.started;
+  out->logits = logits;
+  return out;
+}
+
+std::shared_ptr<CachedState8> capture_state8(const e0n::State8b &src,
+                                               const std::vector<float> &logits) {
+  auto out = std::make_shared<CachedState8>();
+  out->state.pr = src.pr;
+  out->state.cq = snapshot_arrays<float>(src.cq);
+  out->state.ck = snapshot_arrays<float>(src.ck);
+  out->state.cv = snapshot_arrays<float>(src.cv);
+  out->state.ssm = snapshot_arrays<float>(src.ssm);
+  out->state.kc = snapshot_arrays<float>(src.kc);
+  out->state.vc = snapshot_arrays<float>(src.vc);
+  out->state.is_mla = src.is_mla;
+  out->state.position = src.position;
+  out->state.started = src.started;
+  out->logits = logits;
+  return out;
+}
+
+void restore_state35(e0n::State35 &dst, const e0n::State35 &src) {
+  dst.pr = src.pr;
+  dst.conv = snapshot_arrays<float>(src.conv);
+  dst.recur = snapshot_arrays<float>(src.recur);
+  dst.kc = snapshot_arrays<float>(src.kc);
+  dst.vc = snapshot_arrays<float>(src.vc);
+  dst.is_linear = src.is_linear;
+  dst.position = src.position;
+  dst.started = src.started;
+}
+
+void restore_state8(e0n::State8b &dst, const e0n::State8b &src) {
+  dst.pr = src.pr;
+  dst.cq = snapshot_arrays<float>(src.cq);
+  dst.ck = snapshot_arrays<float>(src.ck);
+  dst.cv = snapshot_arrays<float>(src.cv);
+  dst.ssm = snapshot_arrays<float>(src.ssm);
+  dst.kc = snapshot_arrays<float>(src.kc);
+  dst.vc = snapshot_arrays<float>(src.vc);
+  dst.is_mla = src.is_mla;
+  dst.position = src.position;
+  dst.started = src.started;
+}
+
+uint64_t state_bytes(const e0n::State35 &s) {
+  uint64_t n = sizeof(s);
+  for (const auto *v : {&s.conv, &s.recur, &s.kc, &s.vc})
+    for (const auto &a : *v) n += a.nbytes();
+  for (const auto &v : s.pr.cap_hidden) n += v.size() * sizeof(float);
+  for (const auto &v : s.pr.cap_exec) n += v.size() * sizeof(int64_t);
+  for (const auto &v : s.pr.prev_exec) n += v.size() * sizeof(int64_t);
+  return n;
+}
+
+uint64_t state_bytes(const e0n::State8b &s) {
+  uint64_t n = sizeof(s);
+  for (const auto *v : {&s.cq, &s.ck, &s.cv, &s.ssm, &s.kc, &s.vc})
+    for (const auto &a : *v) n += a.nbytes();
+  for (const auto &v : s.pr.cap_hidden) n += v.size() * sizeof(float);
+  for (const auto &v : s.pr.cap_exec) n += v.size() * sizeof(int64_t);
+  for (const auto &v : s.pr.prev_exec) n += v.size() * sizeof(int64_t);
+  return n;
+}
+
+uint64_t snapshot_bytes(const CachedState35 &s) {
+  return state_bytes(s.state) + s.logits.size() * sizeof(float);
+}
+
+uint64_t snapshot_bytes(const CachedState8 &s) {
+  return state_bytes(s.state) + s.logits.size() * sizeof(float);
+}
+
+}  // namespace
+
+struct E0Tokenizer {
+  std::unique_ptr<void, void (*)(void *)> raw;
+  explicit E0Tokenizer(void *h)
+      : raw(h, [](void *p) { if (p) etok_destroy(p); }) {}
+};
+
+static std::string g_create_error;
+static std::mutex g_tok_err_mu;
+
+#include <unordered_map>
+
+static std::mutex g_pool_mu;
+static std::unordered_map<E0Engine *, std::unique_ptr<E0Engine>> g_engines;
+static std::unordered_map<E0Tokenizer *, std::unique_ptr<E0Tokenizer>> g_toks;
+static std::unordered_map<E0Stream *, std::unique_ptr<e0n::gen::StreamBase>> g_streams;
+
+#include "genstream.hpp"
+
+template <class MapT, class PtrT>
+static PtrT pool_new(MapT &m, PtrT p) {
+  std::lock_guard<std::mutex> g(g_pool_mu);
+  using V = typename MapT::mapped_type::element_type;
+  m.emplace(p, std::unique_ptr<V>(static_cast<V *>((void *)p)));
+  return p;
+}
+template <class MapT, class PtrT>
+static void pool_free(MapT &m, PtrT p) {
+  std::lock_guard<std::mutex> g(g_pool_mu);
+  m.erase(p);
+}
+
+namespace {
+uint64_t logits_hash(const std::vector<float> &v) {
+  uint64_t h = 1469598103934665603ull;
+  const uint8_t *p = reinterpret_cast<const uint8_t *>(v.data());
+  for (size_t i = 0; i < v.size() * sizeof(float); i++) { h ^= p[i]; h *= 1099511628211ull; }
+  return h;
+}
+}  // namespace
+
+template <class NetT, class StateT>
+E0Status test_forward_run(std::unique_ptr<NetT> &net, std::unique_ptr<StateT> &st, E0Engine *h,
+                          const std::vector<int32_t> &tok, uint32_t steps, char *buf,
+                          size_t buf_len, size_t *written) {
+  std::string err;
+  std::vector<float> lg;
+  std::vector<e0n::RouteRec> recs;
+  if (!net->prefill(tok, *st, lg, recs, err)) {
+    h->last_error = "{\"error\":\"E-MODEL-INVALID\",\"message\":\"" + esc(err) + "\"}";
+    return E0_ERR_INTERNAL;
+  }
+  h->last_recs = recs;
+  h->has_recs = true;
+  auto argmax = [&](const std::vector<float> &v) {
+    size_t b = 0;
+    for (size_t i = 1; i < v.size(); i++) if (v[i] > v[b]) b = i;
+    return static_cast<int32_t>(b);
+  };
+  std::string json = "{\"prefill_logits_finite\":" +
+                     std::string(std::all_of(lg.begin(), lg.end(), [](float x) { return std::isfinite(x); }) ? "true" : "false") +
+                     ",\"lh\":" + std::to_string(logits_hash(lg)) +
+                     ",\"argmax\":" + std::to_string(argmax(lg)) + ",\"steps\":[";
+  int32_t cur = argmax(lg);
+  for (uint32_t s = 0; s < steps; s++) {
+    if (!net->step(cur, *st, lg, recs, err)) {
+      h->last_error = "{\"error\":\"E-MODEL-INVALID\",\"message\":\"" + esc(err) + "\"}";
+      return E0_ERR_INTERNAL;
+    }
+    h->last_recs = recs;
+    h->has_recs = true;
+    bool fin = std::all_of(lg.begin(), lg.end(), [](float x) { return std::isfinite(x); });
+    cur = argmax(lg);
+    json += (s ? ",{" : "{") + std::string("\"finite\":") + (fin ? "true" : "false") +
+            ",\"lh\":" + std::to_string(logits_hash(lg)) +
+            ",\"argmax\":" + std::to_string(cur) + "}";
+  }
+  json += "]}";
+  return json_out(json, buf, buf_len, written);
+}
+
+static void sweep_terminal(const char *what) {
+  (void)what;
+  std::lock_guard<std::mutex> pg(g_pool_mu);
+  for (auto it = g_streams.begin(); it != g_streams.end();) {
+    if (it->second && it->second->terminal()) it = g_streams.erase(it);
+    else ++it;
+  }
+}
+
+extern "C" {
+
+#ifndef E0N_VERSION
+#define E0N_VERSION "0.1.0"
+#endif
+
+uint32_t e0_abi_version(void) { return E0_ABI_VERSION; }
+const char *e0_engine_version(void) { return E0N_VERSION; }
+
+E0Engine *e0_engine_create(const char *tier_id, const char *model_dir,
+                           const char *config_json) {
+  auto note = [](const std::string &e) { g_create_error = e; };
+  if (!tier_id || !*tier_id) {
+    note("{\"error\":\"E-MODEL-INVALID\",\"message\":\"tier_id is empty\"}");
+    return nullptr;
+  }
+  auto files = std::make_shared<e0n::ModelFiles>();
+  std::string err;
+  if (!files->open(model_dir ? model_dir : "", err)) {
+    note(std::string("{\"error\":\"E-MODEL-PATH\",\"message\":\"") + esc(err) + "\"}");
+    return nullptr;
+  }
+  auto experts = std::make_shared<e0n::ExpertMap>();
+  if (!e0n::ExpertMap::build(*files, *experts, err)) {
+    note(std::string("{\"error\":\"E-MODEL-LAYOUT\",\"message\":\"") + esc(err) + "\"}");
+    return nullptr;
+  }
+  uint64_t budget = 2ull << 30;
+  int io_threads = 3;
+  try {
+    auto cj = nlohmann::json::parse(config_json ? config_json : "{}", nullptr, false);
+    if (!cj.is_discarded()) {
+      if (cj.contains("cache_budget_bytes") && cj["cache_budget_bytes"].is_number_unsigned())
+        budget = cj["cache_budget_bytes"].get<uint64_t>();
+      if (cj.contains("io_threads") && cj["io_threads"].is_number_integer())
+        io_threads = std::max(1, cj["io_threads"].get<int>());
+    }
+  } catch (...) {}
+  auto *h = new E0Engine();
+  h->tier = tier_id;
+  h->model_dir = model_dir ? model_dir : "";
+  h->config_json = config_json ? config_json : "{}";
+  h->files = std::move(files);
+  h->experts = std::move(experts);
+  e0n::PrefixCacheConfig pcc;
+  pcc.enabled = true;
+  std::string identity = std::string(tier_id) + "|" + (model_dir ? model_dir : "|") +
+                         "|abi=" + std::to_string(E0_ABI_VERSION);
+  try {
+    auto cj = nlohmann::json::parse(h->config_json, nullptr, false);
+    if (!cj.is_discarded()) {
+      if (cj.contains("model_identity")) identity = cj["model_identity"].dump();
+      if (cj.contains("prefix_cache") && cj["prefix_cache"].is_object()) {
+        const auto &pc = cj["prefix_cache"];
+        if (pc.contains("enabled") && pc["enabled"].is_boolean())
+          pcc.enabled = pc["enabled"].get<bool>();
+        if (pc.contains("budget_bytes") && pc["budget_bytes"].is_number_unsigned())
+          pcc.budget_bytes = pc["budget_bytes"].get<uint64_t>();
+        if (pc.contains("min_prefix_tokens") && pc["min_prefix_tokens"].is_number_unsigned())
+          pcc.min_prefix_tokens = std::max<uint32_t>(1, pc["min_prefix_tokens"].get<uint32_t>());
+      }
+    }
+  } catch (...) {
+    identity = std::string(tier_id) + "|" + (model_dir ? model_dir : "|") +
+               "|abi=" + std::to_string(E0_ABI_VERSION);
+  }
+  h->model_identity = std::move(identity);
+  h->prefix_cache = std::make_unique<e0n::PrefixCache>(pcc);
+  try {
+    auto cj = nlohmann::json::parse(h->config_json, nullptr, false);
+    if (!cj.is_discarded() && cj.contains("forward") && cj["forward"].is_string()) {
+      const std::string fwd = cj["forward"].get<std::string>();
+      if (fwd == "35b" || fwd == "8b") {
+        const std::string want_tier = "edge0-" + fwd;
+        if (h->tier != want_tier) {
+          pool_free(g_engines, h);
+          note("{\"error\":\"E-MODEL-INVALID\",\"message\":\"forward=" + fwd +
+               " expected tier=" + want_tier + "\"}");
+          return nullptr;
+        }
+        std::shared_ptr<const e0n::lora::Bundle> lo;
+        if (cj.contains("lora") && cj["lora"].is_boolean() && cj["lora"].get<bool>()) {
+          e0n::TierRouting treg;
+          if (!e0n::tier_routing(h->tier, treg)) {
+            pool_free(g_engines, h);
+            note("{\"error\":\"E-MODEL-INVALID\",\"message\":\"tier missing from registry\"}");
+            return nullptr;
+          }
+          auto b = std::make_shared<e0n::lora::Bundle>();
+          std::string lerr;
+          if (!b->load_file(h->model_dir + "/" + treg.lora_file, fwd == "8b",
+                            treg.num_layers, lerr)) {
+            pool_free(g_engines, h);
+            const bool missing = lerr.rfind("lora file unreadable", 0) == 0;
+            note(std::string("{\"error\":\"") + (missing ? "E-MODEL-INVALID" : "E-MODEL-LAYOUT") +
+                 "\",\"message\":\"lora: " + esc(lerr) + "\"}");
+            return nullptr;
+          }
+          lo = b;
+        }
+        std::string ferr;
+        if (fwd == "35b") {
+          h->net35 = e0n::Net35::create(h->files, h->experts, ferr, lo);
+          if (!h->net35) {
+            pool_free(g_engines, h);
+            note(std::string("{\"error\":\"E-MODEL-LAYOUT\",\"message\":\"forward ctx: ") +
+                 esc(ferr) + "\"}");
+            return nullptr;
+          }
+          h->st35 = std::make_unique<e0n::State35>();
+          h->net35->init_state(*h->st35);
+        } else {
+          h->net8b = e0n::Net8b::create(h->files, h->experts, ferr, lo);
+          if (!h->net8b) {
+            pool_free(g_engines, h);
+            note(std::string("{\"error\":\"E-MODEL-LAYOUT\",\"message\":\"forward ctx: ") +
+                 esc(ferr) + "\"}");
+            return nullptr;
+          }
+          h->st8b = std::make_unique<e0n::State8b>();
+          h->net8b->init_state(*h->st8b);
+        }
+        std::ifstream cf(h->model_dir + "/config.json");
+        if (cf) {
+          auto cj2 = nlohmann::json::parse(cf, nullptr, false);
+          if (!cj2.is_discarded() && cj2.contains("eos_token_id")) {
+            const auto &x = cj2["eos_token_id"];
+            if (x.is_number_integer()) {
+              h->eos_ids.push_back(x.get<int32_t>());
+            } else if (x.is_array()) {
+              for (const auto &e : x)
+                if (e.is_number_integer()) h->eos_ids.push_back(e.get<int32_t>());
+            }
+          }
+        }
+      }
+    }
+  } catch (const std::exception &e) {
+    note(std::string("{\"error\":\"E-INTERNAL\",\"message\":\"create: ") +
+         esc(e.what()) + "\"}");
+    delete h;
+    return nullptr;
+  }
+  return pool_new(g_engines, h);
+}
+
+E0Stream *e0_generate_begin(E0Engine *h, const int32_t *prompt_tokens, uint32_t n_tokens,
+                            const char *params_json) {
+  if (!h) return nullptr;
+  {
+    std::lock_guard<std::mutex> pg(g_pool_mu);
+    if (!g_engines.count(h)) return nullptr;
+  }
+  sweep_terminal("begin");
+  std::lock_guard<std::mutex> g(h->mu);
+  e0n::sampling::TierDefaults td = e0n::sampling::tier_defaults(h->tier);
+  auto p = e0n::sampling::parse_params(params_json ? params_json : "{}", td, h->eos_ids);
+  bool transport_only = false;
+  if (p.teacher_tokens.empty()) {
+    auto pv0 = nlohmann::json::parse(params_json ? params_json : "{}", nullptr, false);
+    if (!pv0.is_discarded() && pv0.contains("replay_tokens") && pv0["replay_tokens"].is_array()) {
+      for (const auto &it : pv0["replay_tokens"])
+        if (it.is_number_integer()) p.teacher_tokens.push_back(it.get<int32_t>());
+      transport_only = !p.teacher_tokens.empty();
+    }
+  }
+  if (params_json) {
+    auto pv = nlohmann::json::parse(params_json, nullptr, false);
+    if (!pv.is_discarded() && pv.contains("stop") && pv["stop"].is_array()) {
+      for (const auto &it : pv["stop"]) {
+        if (!it.is_string() || it.get_ref<const std::string &>().empty()) continue;
+        void *tk = etok_create(h->model_dir.c_str());
+        if (!tk) {
+          h->last_error = etok_err_json();
+          return nullptr;
+        }
+        std::vector<int32_t> ids;
+        size_t w = 0;
+        const std::string &str = it.get_ref<const std::string &>();
+        if (etok_tokenize(tk, str.c_str(), nullptr, 0, &w) == 0) {
+          ids.resize(w);
+          if (w) etok_tokenize(tk, str.c_str(), ids.data(), w, &w);
+        }
+        etok_destroy(tk);
+        p.stop_seqs.push_back(std::move(ids));
+      }
+    }
+  }
+  std::vector<int32_t> prompt(prompt_tokens, prompt_tokens + n_tokens);
+  const size_t prompt_len = prompt.size();
+  e0n::gen::StreamBase *ns;
+#ifdef E0N_TEST_SCAFFOLD
+  if (const char *sc = getenv("E0N_SCRIPT")) {
+    std::vector<std::vector<float>> script;
+    const int n = atoi(sc);
+    for (int i = 0; i < n; i++) {
+      std::vector<float> lg(8, -10.f);
+      lg[i % 8] = 5.f;
+      script.push_back(std::move(lg));
+    }
+    ns = new e0n::gen::ScriptedStream(std::move(prompt), std::move(p), std::move(script));
+    ns->set_owner(h);
+    if (h->active) h->active->seal();
+    h->active = ns;
+    return pool_new(g_streams, (E0Stream *)ns);
+  }
+#endif
+  if (transport_only && (h->net35 || h->net8b)) {
+    ns = new e0n::gen::StreamBase(std::move(prompt), std::move(p));
+    ns->set_transport();
+    ns->set_owner(h);
+    if (h->active) h->active->seal();
+    h->active = ns;
+    return pool_new(g_streams, (E0Stream *)ns);
+  }
+  if (!h->net35 && !h->net8b) {
+    const bool empty_turn = p.teacher_tokens.empty() && prompt.empty();
+    ns = new e0n::gen::StreamBase(std::move(prompt), std::move(p));
+    ns->set_transport();
+    ns->set_owner(h);
+    if (empty_turn) ns->set_empty_done();
+    if (h->active) h->active->seal();
+    h->active = ns;
+    return pool_new(g_streams, (E0Stream *)ns);
+  }
+  std::vector<int32_t> prefill_prompt = prompt;
+  size_t cached_tokens = 0;
+  if (h->net35) {
+    auto hit = h->prefix_cache->lookup(h->model_identity, prompt);
+    if (hit.snapshot) {
+      try {
+        auto snap = std::static_pointer_cast<CachedState35>(hit.snapshot);
+        restore_state35(*h->st35, snap->state);
+        cached_tokens = hit.matched_tokens;
+        prefill_prompt.assign(prompt.begin() + static_cast<std::ptrdiff_t>(cached_tokens), prompt.end());
+      } catch (const std::exception &e) {
+        h->prefix_cache->clear();
+        h->last_error = std::string("{\"error\":\"E-INTERNAL\",\"message\":\"") +
+                        esc(e.what()) + "\"}";
+        return nullptr;
+      }
+    } else {
+      h->net35->reset_state(*h->st35);
+    }
+    ns = new e0n::gen::Stream<e0n::Net35, e0n::State35>(
+        h->net35.get(), h->st35.get(),
+        [net = h->net35.get()](const e0n::Step &plan) { net->prefetch(plan); }, h,
+        std::move(prompt), std::move(p));
+    ns->set_prefill_prompt(std::move(prefill_prompt));
+    if (cached_tokens == prompt_len && hit.snapshot) {
+      auto snap = std::static_pointer_cast<CachedState35>(hit.snapshot);
+      ns->set_cached_logits(snap->logits);
+    }
+    ns->set_on_prefill([h](const std::vector<int32_t> &full, const std::vector<float> &logits) {
+      try {
+        auto snap = capture_state35(*h->st35, logits);
+        h->prefix_cache->put(h->model_identity, full, snap, snapshot_bytes(*snap));
+      } catch (const std::exception &e) {
+        h->last_error = std::string("{\"error\":\"E-INTERNAL\",\"message\":\"") +
+                        esc(e.what()) + "\"}";
+      }
+    });
+    ns->set_on_complete([h](const std::vector<int32_t> &full, const std::vector<float> &logits) {
+      try {
+        auto snap = capture_state35(*h->st35, logits);
+        h->prefix_cache->put(h->model_identity, full, snap, snapshot_bytes(*snap));
+      } catch (const std::exception &e) {
+        h->last_error = std::string("{\"error\":\"E-INTERNAL\",\"message\":\"") +
+                        esc(e.what()) + "\"}";
+      }
+    });
+  } else {
+    auto hit = h->prefix_cache->lookup(h->model_identity, prompt);
+    if (hit.snapshot) {
+      try {
+        auto snap = std::static_pointer_cast<CachedState8>(hit.snapshot);
+        restore_state8(*h->st8b, snap->state);
+        cached_tokens = hit.matched_tokens;
+        prefill_prompt.assign(prompt.begin() + static_cast<std::ptrdiff_t>(cached_tokens), prompt.end());
+      } catch (const std::exception &e) {
+        h->prefix_cache->clear();
+        h->last_error = std::string("{\"error\":\"E-INTERNAL\",\"message\":\"") +
+                        esc(e.what()) + "\"}";
+        return nullptr;
+      }
+    } else {
+      h->net8b->reset_state(*h->st8b);
+    }
+    ns = new e0n::gen::Stream<e0n::Net8b, e0n::State8b>(
+        h->net8b.get(), h->st8b.get(),
+        [net = h->net8b.get()](const e0n::Step &plan) { net->prefetch(plan); }, h,
+        std::move(prompt), std::move(p));
+    ns->set_prefill_prompt(std::move(prefill_prompt));
+    if (cached_tokens == prompt_len && hit.snapshot) {
+      auto snap = std::static_pointer_cast<CachedState8>(hit.snapshot);
+      ns->set_cached_logits(snap->logits);
+    }
+    ns->set_on_prefill([h](const std::vector<int32_t> &full, const std::vector<float> &logits) {
+      try {
+        auto snap = capture_state8(*h->st8b, logits);
+        h->prefix_cache->put(h->model_identity, full, snap, snapshot_bytes(*snap));
+      } catch (const std::exception &e) {
+        h->last_error = std::string("{\"error\":\"E-INTERNAL\",\"message\":\"") +
+                        esc(e.what()) + "\"}";
+      }
+    });
+    ns->set_on_complete([h](const std::vector<int32_t> &full, const std::vector<float> &logits) {
+      try {
+        auto snap = capture_state8(*h->st8b, logits);
+        h->prefix_cache->put(h->model_identity, full, snap, snapshot_bytes(*snap));
+      } catch (const std::exception &e) {
+        h->last_error = std::string("{\"error\":\"E-INTERNAL\",\"message\":\"") +
+                        esc(e.what()) + "\"}";
+      }
+    });
+  }
+  h->prefix_cache->set_last_prefill_tokens(prompt_len - cached_tokens);
+  if (h->active) h->active->seal();
+  h->active = ns;
+  return pool_new(g_streams, (E0Stream *)ns);
+}
+
+namespace {
+e0n::gen::StreamBase *checked_stream(E0Engine *h, E0Stream *stream) {
+  if (!h || !stream) return nullptr;
+  std::lock_guard<std::mutex> pg(g_pool_mu);
+  if (!g_engines.count(h)) return nullptr;
+  auto it = g_streams.find(stream);
+  if (it == g_streams.end() || !it->second || !it->second->owner_of(h)) return nullptr;
+  return it->second.get();
+}
+}  // namespace
+
+E0Status e0_next(E0Engine *h, E0Stream *stream, E0TokenOut *out) {
+  auto *st = checked_stream(h, stream);
+  if (!st) return E0_ERR_INVALID;
+  std::lock_guard<std::mutex> g(st->mu);
+  if (out) {
+    if (out->size < sizeof(E0TokenOut)) return E0_ERR_INVALID;
+    out->token = 0;
+    out->flags = 0;
+  }
+  int32_t tok = 0;
+  const bool got = st->next_token(tok);
+  if (got) {
+    if (out) out->token = tok;
+    return E0_NOTIFIED;
+  }
+  const bool errored = st->errored();
+  const std::string em = st->err_msg();
+  {
+    std::lock_guard<std::mutex> pg(g_pool_mu);
+    if (st->owner()) {
+      st->owner()->last_recs = st->last_recs();
+      st->owner()->has_recs = true;
+      if (errored) {
+        st->owner()->last_error =
+            "{\"error\":\"E-INTERNAL\",\"message\":" + esc(em) + "}";
+        if (st->owner()->active == st) st->owner()->active = nullptr;
+        g_streams.erase(stream);
+      }
+    }
+  }
+  if (errored) return E0_ERR_INTERNAL;
+  if (out) out->flags = E0_TOKEN_FINISHED;
+  return E0_NOTIFIED;
+}
+
+E0Status e0_cancel(E0Engine *h, E0Stream *stream) {
+  if (!h) return E0_ERR_INVALID;
+  auto *st = checked_stream(h, stream);
+  if (!st) return E0_OK;
+  std::lock_guard<std::mutex> g(st->mu);
+  st->cancel();
+  return E0_OK;
+}
+
+E0Status e0_stats(E0Engine *h, char *buf, size_t buf_len, size_t *written) {
+  if (!h) return E0_ERR_INVALID;
+  e0n::MemSample m = e0n::mem_sample();
+  e0n::ExpertIo::Metrics c{};
+  if (h->net35) c = h->net35->io_metrics();
+  else if (h->net8b) c = h->net8b->io_metrics();
+  std::string route = "{\"route_source_by_layer\":{";
+  {
+    std::lock_guard<std::mutex> g(h->mu);
+    if (h->has_recs) {
+      for (size_t i = 0; i < h->last_recs.size(); i++) {
+        if (i) route += ",";
+        route += "\"" + std::to_string(h->last_recs[i].layer) + "\":\"" +
+                 (h->last_recs[i].predicted ? "predicted" : "routed") + "\"";
+      }
+    }
+  }
+  route += "}";
+  auto cx = e0n::exec_counters().snapshot();
+  std::string diag = ",\"diag\":{\"qmm_calls\":" + std::to_string(cx.qmm_calls) +
+                     ",\"qmm_experts\":" + std::to_string(cx.qmm_experts) +
+                     ",\"qmm_weight_bytes\":" + std::to_string(cx.qmm_weight_bytes) +
+                     ",\"host_syncs\":" + std::to_string(cx.host_syncs) +
+                     ",\"host_wraps\":" + std::to_string(cx.host_wraps) +
+                     ",\"leases\":" + std::to_string(cx.leases) + "}";
+  const auto pc = h->prefix_cache ? h->prefix_cache->metrics() : e0n::PrefixCacheMetrics{};
+  std::string prefix = ",\"prefix_cache\":{\"enabled\":" +
+                       std::string(pc.enabled ? "true" : "false") +
+                       ",\"budget_bytes\":" + std::to_string(pc.budget_bytes) +
+                       ",\"min_prefix_tokens\":" + std::to_string(pc.min_prefix_tokens) +
+                       ",\"hits\":" + std::to_string(pc.hits) +
+                       ",\"misses\":" + std::to_string(pc.misses) +
+                       ",\"cached_tokens\":" + std::to_string(pc.cached_tokens) +
+                       ",\"entries\":" + std::to_string(pc.entries) +
+                       ",\"bytes\":" + std::to_string(pc.bytes) +
+                       ",\"evictions\":" + std::to_string(pc.evictions) +
+                       ",\"last_cached_tokens\":" + std::to_string(pc.last_cached_tokens) +
+                       ",\"last_prefill_tokens\":" + std::to_string(pc.last_prefill_tokens) +
+                       ",\"last_reason\":\"" + esc(pc.last_reason) + "\"}";
+  std::string json =
+      route + diag + prefix + ",\"loaded_experts\":" + std::to_string(c.loads) +
+      ",\"cache_hits\":" + std::to_string(c.hits) +
+      ",\"cache_misses\":" + std::to_string(c.misses) +
+      ",\"cache_evictions\":" + std::to_string(c.evictions) +
+      ",\"memory\":{\"resident_peak_bytes\":" + std::to_string(m.resident_peak_bytes) +
+      ",\"footprint_bytes\":" + std::to_string(m.footprint_bytes) +
+      ",\"expert_cache_heap_bytes\":" + std::to_string(c.heap_bytes) +
+      ",\"file_backed_bytes\":0},"
+      "\"tier\":\"" + esc(h->tier) + "\"}";
+  return json_out(json, buf, buf_len, written);
+}
+
+E0Status e0_turn_reset(E0Engine *h) {
+  if (!h) return E0_ERR_INVALID;
+  std::lock_guard<std::mutex> g(h->mu);
+  if (h->active) {
+    h->active->seal();
+    h->active = nullptr;
+  }
+  sweep_terminal("turn_reset");
+  if (h->net35 && h->st35) {
+    h->net35->reset_state(*h->st35);
+    h->last_recs.clear();
+    h->has_recs = false;
+  }
+  if (h->net8b && h->st8b) {
+    h->net8b->reset_state(*h->st8b);
+    h->last_recs.clear();
+    h->has_recs = false;
+  }
+  h->last_error.clear();
+  return E0_OK;
+}
+
+E0Status e0_last_error(E0Engine *h, char *buf, size_t buf_len, size_t *written) {
+  std::string msg;
+  if (h) {
+    std::lock_guard<std::mutex> g(h->mu);
+    msg = h->last_error.empty()
+              ? std::string("{\"error\":\"E-INTERNAL\",\"message\":\"no error recorded\"}")
+              : h->last_error;
+  } else {
+    std::lock_guard<std::mutex> g(g_tok_err_mu);
+    msg = g_create_error.empty()
+              ? std::string("{\"error\":\"E-INTERNAL\",\"message\":\"no error recorded\"}")
+              : g_create_error;
+  }
+  return json_out(msg, buf, buf_len, written);
+}
+
+void e0_engine_destroy(E0Engine *h) {
+  if (!h) return;
+  {
+    std::lock_guard<std::mutex> pg(g_pool_mu);
+    for (auto it = g_streams.begin(); it != g_streams.end();) {
+      if (it->second && it->second->owner_of(h))
+        it = g_streams.erase(it);
+      else
+        ++it;
+    }
+  }
+  pool_free(g_engines, h);
+}
+
+extern "C++" {
+namespace {
+
+void tok_err(const std::string &msg) {
+  std::lock_guard<std::mutex> g(g_tok_err_mu);
+  g_create_error = msg;
+}
+
+template <class F>
+bool etok_out(F f, std::string &out, const char *fail_what) {
+  size_t w = 0;
+  if (f(nullptr, 0, &w) != 0) {
+    tok_err(etok_err_json());
+    return false;
+  }
+  out.assign(w + 1, '\0');
+  if (w && f(out.data(), w, &w) != 0) {
+    tok_err(etok_err_json());
+    return false;
+  }
+  out.resize(w);
+  return true;
+}
+
+bool etok_ids(void *raw, const std::string &text, std::vector<int32_t> &ids) {
+  size_t n = 0;
+  if (etok_tokenize(raw, text.c_str(), nullptr, 0, &n) != 0) {
+    tok_err(etok_err_json());
+    return false;
+  }
+  ids.assign(n, 0);
+  if (n && etok_tokenize(raw, text.c_str(), ids.data(), n, &n) != 0) {
+    tok_err(etok_err_json());
+    return false;
+  }
+  return true;
+}
+
+}  // namespace
+}  // extern "C++"
+
+E0Tokenizer *e0_tok_create(const char *model_dir) {
+  if (!is_dir(model_dir)) {
+    tok_err(std::string("{\"error\":\"E-MODEL-PATH\",\"message\":\"directory does not exist: ") +
+            esc(model_dir ? model_dir : "") + "\"}");
+    return nullptr;
+  }
+  void *raw = etok_create(model_dir);
+  if (!raw) {
+    tok_err(etok_err_json());
+    return nullptr;
+  }
+  return pool_new(g_toks, new E0Tokenizer(raw));
+}
+
+E0Status e0_tok_tokenize(E0Tokenizer *t, const char *text_json, char *buf, size_t buf_len,
+                         size_t *written) {
+  if (!t) return E0_ERR_INVALID;
+  nlohmann::json v = nlohmann::json::parse(text_json ? text_json : "", nullptr, false);
+  if (v.is_discarded() || !v.is_object()) {
+    tok_err("{\"error\":\"E-INVALID-ARG\",\"message\":\"invalid text_json\"}");
+    return E0_ERR_INTERNAL;
+  }
+  std::vector<std::vector<int32_t>> batches;
+  if (v.contains("text") && v["text"].is_string()) {
+    std::vector<int32_t> ids;
+    if (!etok_ids(t->raw.get(), v["text"].get<std::string>(), ids)) return E0_ERR_INTERNAL;
+    batches.push_back(std::move(ids));
+  } else if (v.contains("texts") && v["texts"].is_array()) {
+    for (const auto &item : v["texts"]) {
+      if (!item.is_string()) {
+        tok_err("{\"error\":\"E-INVALID-ARG\",\"message\":\"texts contains a non-string element\"}");
+        return E0_ERR_INTERNAL;
+      }
+      std::vector<int32_t> ids;
+      if (!etok_ids(t->raw.get(), item.get<std::string>(), ids)) return E0_ERR_INTERNAL;
+      batches.push_back(std::move(ids));
+    }
+  } else {
+    tok_err("{\"error\":\"E-INVALID-ARG\",\"message\":\"need text or texts\"}");
+    return E0_ERR_INTERNAL;
+  }
+  nlohmann::json out = {{"tokens", batches}};
+  return json_out(out.dump(), buf, buf_len, written);
+}
+
+E0Status e0_tok_apply_template(E0Tokenizer *t, const char *messages_json, char *buf,
+                               size_t buf_len, size_t *written) {
+  if (!t) return E0_ERR_INVALID;
+  nlohmann::json v = nlohmann::json::parse(messages_json ? messages_json : "", nullptr, false);
+  if (v.is_discarded()) {
+    tok_err("{\"error\":\"E-INVALID-ARG\",\"message\":\"invalid messages_json\"}");
+    return E0_ERR_INTERNAL;
+  }
+  nlohmann::json ctx = nlohmann::json::object();
+  bool agp = true;
+  if (v.is_array()) {
+    ctx["messages"] = v;
+  } else {
+    for (auto it = v.begin(); it != v.end(); ++it) {
+      if (it.key() == "add_generation_prompt" && it.value().is_boolean())
+        agp = it.value().get<bool>();
+      else
+        ctx[it.key()] = it.value();
+    }
+    if (!ctx.contains("messages")) ctx["messages"] = nlohmann::json::array();
+  }
+  ctx["add_generation_prompt"] = agp;
+  std::string cj = ctx.dump();
+  void *raw = t->raw.get();
+  std::string prompt;
+  if (!etok_out([&](char *b, size_t l, size_t *w) {
+        return etok_render(raw, cj.c_str(), b, l, w);
+      },
+      prompt, "render"))
+    return E0_ERR_INTERNAL;
+  std::vector<int32_t> ids;
+  if (!etok_ids(raw, prompt, ids)) return E0_ERR_INTERNAL;
+  nlohmann::json out = {{"prompt", prompt}, {"token_count", ids.size()}};
+  return json_out(out.dump(), buf, buf_len, written);
+}
+
+E0Status e0_tok_decode(E0Tokenizer *t, const char *tokens_json, char *buf, size_t buf_len,
+                       size_t *written) {
+  if (!t) return E0_ERR_INVALID;
+  nlohmann::json v = nlohmann::json::parse(tokens_json ? tokens_json : "", nullptr, false);
+  if (v.is_discarded() || !v.is_object() || !v.contains("tokens") || !v["tokens"].is_array()) {
+    tok_err("{\"error\":\"E-INVALID-ARG\",\"message\":\"need {\\\"tokens\\\":[...]} \"}");
+    return E0_ERR_INTERNAL;
+  }
+  std::vector<int32_t> ids;
+  for (const auto &x : v["tokens"]) {
+    if (!x.is_number_integer()) {
+      tok_err("{\"error\":\"E-INVALID-ARG\",\"message\":\"tokens contains a non-integer\"}");
+      return E0_ERR_INTERNAL;
+    }
+    ids.push_back(x.get<int32_t>());
+  }
+  void *raw = t->raw.get();
+  std::string text;
+  if (!etok_out([&](char *b, size_t l, size_t *w) {
+        return etok_decode(raw, ids.data(), ids.size(), b, l, w);
+      },
+      text, "decode"))
+    return E0_ERR_INTERNAL;
+  nlohmann::json out = {{"text", text}};
+  return json_out(out.dump(), buf, buf_len, written);
+}
+
+void e0_tok_destroy(E0Tokenizer *t) { pool_free(g_toks, t); }
+
+E0Status e0_test_forward(E0Engine *h, const int32_t *tokens, uint32_t n_tokens,
+                         uint32_t steps, char *buf, size_t buf_len, size_t *written) {
+  if (!h) return E0_ERR_INVALID;
+  std::lock_guard<std::mutex> g(h->mu);
+  if ((!h->net35 || !h->st35) && (!h->net8b || !h->st8b)) {
+    h->last_error = "{\"error\":\"E-INVALID-ARG\",\"message\":\"forward context not enabled (config {\"forward\":\"35b\"|\"8b\"})\"}";
+    return E0_ERR_INVALID;
+  }
+  std::vector<int32_t> tok(tokens, tokens + n_tokens);
+  if (h->net35) return test_forward_run(h->net35, h->st35, h, tok, steps, buf, buf_len, written);
+  return test_forward_run(h->net8b, h->st8b, h, tok, steps, buf, buf_len, written);
+}
+
+}  // extern "C"
diff --git a/macos/crates/engine-native/src/dense.cpp b/macos/crates/engine-native/src/dense.cpp
new file mode 100644
index 0000000..79d778d
--- /dev/null
+++ b/macos/crates/engine-native/src/dense.cpp
@@ -0,0 +1,142 @@
+#include "dense.hpp"
+
+#include <fcntl.h>
+#include <unistd.h>
+
+#include <cmath>
+#include <cstring>
+
+namespace e0n {
+namespace {
+
+bool read_range(const ModelFiles& files, const TensorInfo& t, std::vector<uint8_t>& out,
+                std::string& err) {
+  if (t.end < t.begin) { err = "tensor range inverted: " + t.shard; return false; }
+  out.resize(static_cast<size_t>(t.end - t.begin));
+  int fd = ::open(files.shard_path(t.shard).c_str(), O_RDONLY);
+  if (fd < 0) { err = "open failed: " + files.shard_path(t.shard); return false; }
+  size_t done = 0;
+  while (done < out.size()) {
+    ssize_t got = ::pread(fd, out.data() + done, out.size() - done,
+                          static_cast<off_t>(t.begin + done));
+    if (got <= 0) {
+      ::close(fd);
+      err = "pread failed: " + t.shard;
+      return false;
+    }
+    done += static_cast<size_t>(got);
+  }
+  ::close(fd);  // one-shot resident read; page-cache open is fine (not the hot-stack F_NOCACHE path)
+  return true;
+}
+
+}  // namespace
+
+float f16_bits_to_f(uint16_t b) {
+  _Float16 h;
+  __builtin_memcpy(&h, &b, 2);
+  return static_cast<float>(h);
+}
+
+bool read_dense_f32(const ModelFiles& files, const std::string& tensor, DenseF32& out,
+                    std::string& err) {
+  const TensorInfo* t = files.find(tensor);
+  if (!t) { err = "missing tensor: " + tensor; return false; }
+  if (t->shape.empty()) { err = "dense rejects scalar tensors: " + tensor; return false; }
+  std::vector<uint8_t> bytes;
+  if (!read_range(files, *t, bytes, err)) return false;
+  int64_t rows = 1;
+  for (size_t i = 0; i + 1 < t->shape.size(); i++) rows *= t->shape[i];
+  int64_t cols = t->shape.back();
+  size_t n = static_cast<size_t>(rows * cols);
+  out.rows = rows;
+  out.cols = cols;
+  out.v.resize(n);
+  if (t->dtype == "F32") {
+    if (bytes.size() != n * 4) { err = "F32 byte size mismatch: " + tensor; return false; }
+    std::memcpy(out.v.data(), bytes.data(), bytes.size());
+  } else if (t->dtype == "BF16") {
+    if (bytes.size() != n * 2) { err = "BF16 byte size mismatch: " + tensor; return false; }
+    const uint16_t* p = reinterpret_cast<const uint16_t*>(bytes.data());
+    for (size_t i = 0; i < n; i++) out.v[i] = bf16_bits_to_f(p[i]);
+  } else if (t->dtype == "F16") {
+    if (bytes.size() != n * 2) { err = "F16 byte size mismatch: " + tensor; return false; }
+    const uint16_t* p = reinterpret_cast<const uint16_t*>(bytes.data());
+    for (size_t i = 0; i < n; i++) out.v[i] = f16_bits_to_f(p[i]);
+  } else if (t->dtype == "I32" || t->dtype == "I64") {
+    size_t esz = t->dtype == "I32" ? 4 : 8;
+    if (bytes.size() != n * esz) { err = "integer byte size mismatch: " + tensor; return false; }
+    for (size_t i = 0; i < n; i++) {
+      if (esz == 4) { int32_t v; std::memcpy(&v, bytes.data() + 4 * i, 4); out.v[i] = (float)v; }
+      else { int64_t v; std::memcpy(&v, bytes.data() + 8 * i, 8); out.v[i] = (float)v; }
+    }
+  } else {
+    err = "unsupported dense dtype: " + t->dtype + " (" + tensor + ")";
+    return false;
+  }
+  return true;
+}
+
+void dequant_row(const uint32_t* packed, const uint16_t* scales_row,
+                 const uint16_t* biases_row, int64_t cols_packed, int bits,
+                 int64_t group_size, std::vector<float>& dst_row) {
+  const int per_word = 32 / bits;          // int4→8, int8→4
+  const uint32_t mask = (1u << bits) - 1u;  // 0xF / 0xFF
+  const int64_t cols = cols_packed * per_word;
+  dst_row.resize(static_cast<size_t>(cols));
+  for (int64_t c = 0; c < cols; c++) {
+    int64_t word = c / per_word;
+    int shift = static_cast<int>((c % per_word) * bits);
+    uint32_t q = (packed[word] >> shift) & mask;  // little-endian grouping
+    int64_t g = c / group_size;
+    double s = bf16_bits_to_f(scales_row[g]);
+    double b = bf16_bits_to_f(biases_row[g]);
+    dst_row[c] = static_cast<float>(static_cast<double>(q) * s + b);  // q*scale + bias
+  }
+}
+
+bool read_dense_dequant(const ModelFiles& files, const std::string& tensor, int64_t bits,
+                        int64_t group_size, DenseF32& out, std::string& err) {
+  if (bits != 4 && bits != 8) { err = "dense dequant supports only 4/8 bit"; return false; }
+  const int per_word = static_cast<int>(32 / bits);
+  // tensor = "...gate.weight" → module path = strip last segment; scales/biases hang there.
+  std::string mod = tensor;
+  auto dot = mod.rfind('.');
+  if (dot != std::string::npos && mod.substr(dot + 1) == "weight") mod = mod.substr(0, dot);
+  const TensorInfo *w = files.find(tensor), *s = files.find(mod + ".scales"),
+                   *b = files.find(mod + ".biases");
+  if (!w || !s || !b) { err = "missing weight/scales/biases: " + tensor; return false; }
+  if (w->dtype != "U32" || s->dtype != "BF16" || b->dtype != "BF16") {
+    err = "quantized triple dtype mismatch: " + tensor;
+    return false;
+  }
+  std::vector<uint8_t> wb, sb, bb;
+  if (!read_range(files, *w, wb, err) || !read_range(files, *s, sb, err) ||
+      !read_range(files, *b, bb, err))
+    return false;
+  int64_t rows = 1;
+  for (size_t i = 0; i + 1 < w->shape.size(); i++) rows *= w->shape[i];
+  int64_t cp = w->shape.back();                 // packed last dim
+  int64_t cols = cp * per_word;                 // logical K
+  int64_t groups = cols / group_size;
+  if (s->shape.back() != groups || static_cast<int64_t>(sb.size()) != rows * groups * 2 ||
+      static_cast<int64_t>(bb.size()) != rows * groups * 2) {
+    err = "scales/biases shape mismatch: " + tensor;
+    return false;
+  }
+  out.rows = rows;
+  out.cols = cols;
+  out.v.resize(static_cast<size_t>(rows * cols));
+  const uint32_t* wp = reinterpret_cast<const uint32_t*>(wb.data());
+  const uint16_t* sp = reinterpret_cast<const uint16_t*>(sb.data());
+  const uint16_t* bp = reinterpret_cast<const uint16_t*>(bb.data());
+  for (int64_t r = 0; r < rows; r++) {
+    std::vector<float> row;
+    dequant_row(wp + r * cp, sp + r * groups, bp + r * groups, cp, static_cast<int>(bits),
+                group_size, row);
+    std::memcpy(out.v.data() + r * cols, row.data(), row.size() * sizeof(float));
+  }
+  return true;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/dense.hpp b/macos/crates/engine-native/src/dense.hpp
new file mode 100644
index 0000000..767b224
--- /dev/null
+++ b/macos/crates/engine-native/src/dense.hpp
@@ -0,0 +1,36 @@
+// Whole-tensor load of non-expert weights and fp32 dequant (q*scale + bias, little-endian).
+#pragma once
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "safio.hpp"
+
+namespace e0n {
+
+struct DenseF32 {
+  int64_t rows = 0, cols = 0;
+  std::vector<float> v;  // row-major [rows, cols]
+};
+
+bool read_dense_f32(const ModelFiles& files, const std::string& tensor, DenseF32& out,
+                    std::string& err);
+
+// Quantized tensor (last dim is packed: bits=4 → ×8, bits=8 → ×4).
+// weight name like "...gate.weight"; scales/biases take the same suffix. Output [R, C] fp32.
+bool read_dense_dequant(const ModelFiles& files, const std::string& tensor, int64_t bits,
+                        int64_t group_size, DenseF32& out, std::string& err);
+
+inline float bf16_bits_to_f(uint16_t b) {
+  uint32_t u = static_cast<uint32_t>(b) << 16;
+  float f;
+  __builtin_memcpy(&f, &u, 4);
+  return f;
+}
+float f16_bits_to_f(uint16_t b);
+
+void dequant_row(const uint32_t* packed, const uint16_t* scales_row,
+                 const uint16_t* biases_row, int64_t cols_packed, int bits,
+                 int64_t group_size, std::vector<float>& dst_row);
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/expertio.cpp b/macos/crates/engine-native/src/expertio.cpp
new file mode 100644
index 0000000..a7516f3
--- /dev/null
+++ b/macos/crates/engine-native/src/expertio.cpp
@@ -0,0 +1,882 @@
+#include "expertio.hpp"
+
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <fcntl.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <set>
+#include <vector>
+
+#include "fwd_common.hpp"  // mkshape
+
+namespace e0n {
+namespace mx = mlx::core;
+// ---------------- ShardMaps ----------------
+
+ExpertIo::ShardMaps::~ShardMaps() {
+  std::lock_guard<std::mutex> lk(mu_);
+  for (auto &[p, m] : maps_) munmap(m.first, size_t(m.second));
+  for (auto &[p, fd] : fds_) ::close(fd);
+}
+
+const uint8_t *ExpertIo::ShardMaps::raw(const std::string &shard_path, uint64_t off,
+                                        uint64_t len) {
+  uint8_t *base = nullptr;
+  uint64_t total = 0;
+  {
+    std::lock_guard<std::mutex> lk(mu_);
+    auto it = maps_.find(shard_path);
+    if (it == maps_.end()) {
+      int fd = ::open(shard_path.c_str(), O_RDONLY);
+      if (fd < 0) return nullptr;
+      struct stat sb;
+      if (fstat(fd, &sb) != 0) { ::close(fd); return nullptr; }
+      void *m = mmap(nullptr, size_t(sb.st_size), PROT_READ, MAP_SHARED, fd, 0);
+      if (m == MAP_FAILED) { ::close(fd); return nullptr; }
+      int keep = fd;
+      it = maps_.emplace(shard_path,
+                         std::pair<uint8_t *, uint64_t>{static_cast<uint8_t *>(m),
+                                                        uint64_t(sb.st_size)}).first;
+      fds_.emplace(shard_path, keep);
+    }
+    base = it->second.first;
+    total = it->second.second;
+  }
+  if (off + len > total) return nullptr;
+  return base + off;
+}
+
+void ExpertIo::ShardMaps::advise_range(const std::string &shard_path, uint64_t off,
+                                       uint64_t len) {
+  if (len == 0) return;
+  if (!raw(shard_path, 0, 1)) return;
+  std::lock_guard<std::mutex> lk(mu_);
+  auto it = maps_.find(shard_path);
+  if (it == maps_.end()) return;
+  uint64_t page = uint64_t(sysconf(_SC_PAGESIZE));
+  uint64_t start = off / page * page;
+  uint64_t end = (off + len + page - 1) / page * page;
+  if (end > it->second.second) end = it->second.second;
+  if (end > start) madvise(it->second.first + start, size_t(end - start), MADV_WILLNEED);
+}
+
+void ExpertIo::ShardMaps::advise_whole(const std::string &shard_path) {  // mmap.py:43-49
+  if (!raw(shard_path, 0, 1)) return;
+  std::lock_guard<std::mutex> lk(mu_);
+  auto it = maps_.find(shard_path);
+  if (it != maps_.end()) madvise(it->second.first, size_t(it->second.second), MADV_WILLNEED);
+}
+
+void ExpertIo::ShardMaps::seq_read(const std::string &shard_path) {
+  if (!raw(shard_path, 0, 1)) return;
+  uint64_t total;
+  {
+    std::lock_guard<std::mutex> lk(mu_);
+    total = maps_[shard_path].second;
+  }
+  volatile uint64_t sink = 0;
+  constexpr uint64_t kChunk = 1ull << 24;
+  for (uint64_t off = 0; off < total; off += kChunk) {
+    uint64_t n = std::min(kChunk, total - off);
+    const uint8_t *p = raw(shard_path, off, n);
+    for (uint64_t o = 0; o + 8 <= n; o += 8) {
+      uint64_t v;
+      std::memcpy(&v, p + o, 8);
+      sink += v;
+    }
+  }
+  (void)sink;
+}
+
+std::shared_ptr<Bundle> SharedExpertCache::get(uint64_t key) {  // cache.py:24-29
+  std::lock_guard<std::mutex> lk(mu_);
+  auto it = map_.find(key);
+  if (it == map_.end()) return nullptr;
+  for (auto o = order_.begin(); o != order_.end(); ++o)
+    if (*o == key) { order_.erase(o); order_.push_back(key); break; }
+  return it->second;
+}
+std::shared_ptr<Bundle> SharedExpertCache::peek(uint64_t key) {  // 31-33
+  std::lock_guard<std::mutex> lk(mu_);
+  auto it = map_.find(key);
+  return it == map_.end() ? nullptr : it->second;
+}
+void SharedExpertCache::put(uint64_t key, std::shared_ptr<Bundle> b) {  // 34-39
+  std::lock_guard<std::mutex> lk(mu_);
+  if (map_.count(key)) {
+    map_[key] = std::move(b);
+    for (auto o = order_.begin(); o != order_.end(); ++o)
+      if (*o == key) { order_.erase(o); order_.push_back(key); break; }
+    return;
+  }
+  map_[key] = std::move(b);
+  order_.push_back(key);
+  if (slots_ > 0 && int(map_.size()) > slots_) {  // popitem(last=False)
+    uint64_t old = order_.front();
+    order_.pop_front();
+    map_.erase(old);
+    evictions_++;
+  }
+}
+std::shared_ptr<Bundle> SharedExpertCache::pop(uint64_t key) {  // 41-46
+  std::lock_guard<std::mutex> lk(mu_);
+  auto it = map_.find(key);
+  if (it == map_.end()) return nullptr;
+  auto b = it->second;
+  map_.erase(it);
+  for (auto o = order_.begin(); o != order_.end(); ++o)
+    if (*o == key) { order_.erase(o); break; }
+  return b;
+}
+size_t SharedExpertCache::size() {
+  std::lock_guard<std::mutex> lk(mu_);
+  return map_.size();
+}
+bool SharedExpertCache::contains(uint64_t key) {
+  std::lock_guard<std::mutex> lk(mu_);
+  return map_.count(key) > 0;
+}
+
+bool PrefetchBuffer::put(uint64_t key, std::shared_ptr<Bundle> b) {  // cache.py:66-72
+  std::lock_guard<std::mutex> lk(mu_);
+  if (map_.count(key)) { map_[key] = std::move(b); return true; }
+  map_[key] = std::move(b);
+  order_.push_back(key);
+  if (cap_ > 0 && int(map_.size()) > cap_) {
+    uint64_t old = order_.front();
+    order_.pop_front();
+    map_.erase(old);
+    return false;
+  }
+  return true;
+}
+std::shared_ptr<Bundle> PrefetchBuffer::pop(uint64_t key) {  // 74-75
+  std::lock_guard<std::mutex> lk(mu_);
+  auto it = map_.find(key);
+  if (it == map_.end()) return nullptr;
+  auto b = it->second;
+  map_.erase(it);
+  for (auto o = order_.begin(); o != order_.end(); ++o)
+    if (*o == key) { order_.erase(o); break; }
+  return b;
+}
+bool PrefetchBuffer::contains(uint64_t key) {
+  std::lock_guard<std::mutex> lk(mu_);
+  return map_.count(key) > 0;
+}
+
+// ---------------- IoPool ----------------
+
+IoPool::IoPool(int n) {
+  for (int i = 0; i < n; i++)
+    th_.emplace_back([this] {
+      for (;;) {
+        std::function<void()> job;
+        {
+          std::unique_lock<std::mutex> lk(mu_);
+          cv_.wait(lk, [this] { return stop_ || !q_.empty(); });
+          if (stop_ && q_.empty()) return;
+          job = std::move(q_.front());
+          q_.pop_front();
+        }
+        job();
+      }
+    });
+}
+IoPool::~IoPool() {
+  {
+    std::lock_guard<std::mutex> lk(mu_);
+    stop_ = true;
+  }
+  cv_.notify_all();
+  for (auto &t : th_)
+    if (t.joinable()) t.join();
+}
+
+// ---------------- ExpertIo ----------------
+
+ExpertIo::ExpertIo(const ModelFiles *files, int cache_slots, int prefetch_cap, int load_threads,
+                   int prefetch_threads)
+    : files_(files), cache_(cache_slots), pbuf_(prefetch_cap), load_pool_(load_threads),
+      prefetch_pool_(prefetch_threads) {}
+
+namespace {
+void materialize(const std::shared_ptr<Bundle> &b) {
+  std::vector<mx::array> v;
+  for (int i = 0; i < 9; i++) v.push_back(b->at(i));
+  mx::eval(v);
+}
+
+const char *kProjFile[3] = {"up_proj", "gate_proj", "down_proj"};
+const char *kPartFile[3] = {"weight", "scales", "biases"};
+}  // namespace
+
+void ExpertIo::add_layer(const LayerCfg &cfg) {
+  li2idx_[cfg.layer] = int(layers_.size());
+  layers_.emplace_back();
+  Layer &L = layers_.back();
+  L.cfg = cfg;
+  for (int p = 0; p < 3; p++)
+    for (int q = 0; q < 3; q++) {
+      const int i = p * 3 + q;
+      const std::string name = cfg.prefix + "." + kProjFile[p] + "." + kPartFile[q];
+      const TensorInfo *ti = files_->find(name);
+      if (!ti) {
+        std::fprintf(stderr, "expertio: missing tensor %s\n", name.c_str());
+        continue;
+      }
+      const int64_t O = (p == 2) ? cfg.H : cfg.F;
+      const int64_t K = (q == 0) ? ((p == 2) ? cfg.F / 8 : cfg.H / 8)
+                                 : ((p == 2) ? cfg.F / cfg.group : cfg.H / cfg.group);
+      L.reg[i] = {ti, O, K, (ti->end - ti->begin) / uint64_t(cfg.E), q == 0 ? 4u : 2u};
+      if (ti->dtype == "BF16") L.reg[i].esize = 2;
+      else if (ti->dtype == "U32") L.reg[i].esize = 4;
+      else std::fprintf(stderr, "expertio: unexpected dtype %s @%s\n", ti->dtype.c_str(), name.c_str());
+    }
+}
+
+int ExpertIo::layer_idx(int li) const {
+  auto it = li2idx_.find(li);
+  return it == li2idx_.end() ? -1 : it->second;
+}
+
+std::shared_ptr<Bundle> ExpertIo::build(int li, int64_t expert) {
+  int idx = layer_idx(li);
+  if (idx < 0) return nullptr;
+  Layer &L = layers_[size_t(idx)];
+  auto b = std::make_shared<Bundle>();
+  uint64_t tot = 0;
+  for (int i = 0; i < 9; i++) {
+      const auto &r = L.reg[i];
+      if (!r.ti) return nullptr;
+      const uint8_t *src = shards_.raw(files_->shard_path(r.ti->shard),
+                                       r.ti->begin + uint64_t(expert) * r.per, r.per);
+      if (!src) {
+        std::fprintf(stderr, "expertio: raw out of range %s e%lld off=%llu per=%llu\n",
+                     files_->shard_path(r.ti->shard).c_str(), (long long)expert,
+                     (unsigned long long)(r.ti->begin + uint64_t(expert) * r.per),
+                     (unsigned long long)r.per);
+        return nullptr;
+      }
+      if (!r.ti->shape.empty() && r.per != uint64_t(r.rows) * uint64_t(r.cols) * r.esize) {
+        std::fprintf(stderr, "expertio: region size mismatch %s per=%llu vs %lldx%lldx%u\n",
+                     r.ti->dtype.c_str(), (unsigned long long)r.per, (long long)r.rows,
+                     (long long)r.cols, r.esize);
+        return nullptr;
+      }
+      if (r.esize == 4) {
+        b->at_mut(i) = mx::array(reinterpret_cast<const uint32_t *>(src),
+                                 mkshape({r.rows, r.cols}), mx::uint32);
+      } else {
+        auto a16 = mx::array(reinterpret_cast<const uint16_t *>(src), mkshape({r.rows, r.cols}),
+                             mx::uint16);
+        b->at_mut(i) = mx::view(a16, mx::bfloat16);
+      }
+      tot += r.per;
+  }
+  b->bytes = tot;
+  heap_->add(tot);
+  auto heap = heap_;
+  auto keep = b;
+  return std::shared_ptr<Bundle>(b.get(), [heap, tot, keep](Bundle *p) mutable {
+    heap->sub(tot);
+    keep.reset();
+  });
+}
+
+std::map<int64_t, std::shared_ptr<Bundle>> ExpertIo::get_bundles(
+    int li, const std::vector<int64_t> &experts) {
+  std::map<int64_t, std::shared_ptr<Bundle>> out;
+  std::vector<int64_t> missing;
+  int lidx = layer_idx(li);
+  Layer *L = (lidx >= 0) ? &layers_[size_t(lidx)] : nullptr;
+  const double decay = L ? L->cfg.hot_decay : 0.75;
+
+  for (auto e : experts) {  // 430-435
+    auto c = cache_.get(key_of(li, e));
+    if (c) { out[e] = std::move(c); hits_++; }
+    else missing.push_back(e);
+  }
+
+  std::vector<int64_t> still;  // 437-449
+  for (auto e : missing) {
+    auto pb = pbuf_.pop(key_of(li, e));
+    if (pb) {
+      phi_++; hits_++;
+      if (L) L->bump_count(e, decay);
+      cache_.put(key_of(li, e), pb);
+      out[e] = std::move(pb);
+    } else still.push_back(e);
+  }
+
+  std::vector<int64_t> cold;  // 451-467 inflight
+  for (auto e : still) {
+    std::shared_future<std::shared_ptr<Bundle>> fut;
+    {
+      std::lock_guard<std::mutex> lk(infl_mu_);
+      auto it = inflight_.find(key_of(li, e));
+      if (it != inflight_.end()) fut = it->second;
+    }
+    if (fut.valid()) {
+        auto b = fut.get();
+        {
+        std::lock_guard<std::mutex> lk(infl_mu_);
+        inflight_.erase(key_of(li, e));
+      }
+      if (!b) {
+        std::fprintf(stderr, "expertio: inflight build failed L%d E%lld\n", li, (long long)e);
+        continue;
+      }
+      phi_++; hits_++;
+      if (L) L->bump_count(e, decay);
+      cache_.put(key_of(li, e), b);
+      out[e] = std::move(b);
+    } else cold.push_back(e);
+  }
+  if (!cold.empty()) {
+    std::vector<std::shared_future<std::shared_ptr<Bundle>>> fs;
+    for (auto e : cold) fs.push_back(load_pool_.submit([this, li, e] { return build(li, e); }));
+    for (size_t i = 0; i < cold.size(); i++) {
+      auto b = fs[i].get();
+      if (!b) {
+        std::fprintf(stderr, "expertio: build failed L%d E%lld (missing region key or map OOB)\n", li,
+                     (long long)cold[i]);
+        continue;
+      }
+      loads_++; misses_++;
+      if (L) L->bump_count(cold[i], decay);
+      cache_.put(key_of(li, cold[i]), b);
+      out[cold[i]] = std::move(b);
+    }
+  }
+  for (auto &[e, b] : out) materialize(b);
+  return out;
+}
+
+void ExpertIo::prefetch(int li, const std::vector<int64_t> &experts) {
+  std::vector<int64_t> uniq(experts.begin(), experts.end());
+  std::sort(uniq.begin(), uniq.end());
+  uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+  if (uniq.empty()) return;
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  std::vector<int64_t> missing;
+  for (auto e : uniq) {
+    {
+      std::lock_guard<std::mutex> lk(infl_mu_);
+      if (inflight_.count(key_of(li, e))) continue;
+    }
+    if (cache_.peek(key_of(li, e)) || pbuf_.contains(key_of(li, e))) continue;
+    missing.push_back(e);
+  }
+  if (missing.empty()) return;
+  if (L.cfg.warm_willneed) warm_willneed(li, missing);  // 502-503
+  std::vector<std::pair<int64_t, std::shared_future<std::shared_ptr<Bundle>>>> sub;
+  for (auto e : missing) {
+    auto fut = prefetch_pool_.submit([this, li, e] { return build(li, e); });
+    {
+      std::lock_guard<std::mutex> lk(infl_mu_);
+      inflight_[key_of(li, e)] = fut;
+    }
+    sub.emplace_back(e, fut);
+  }
+  psub_ += missing.size();
+  for (auto &[e, fut] : sub) {
+    uint64_t key = key_of(li, e);
+    prefetch_pool_.submit([this, e, key, fut] {
+      auto b = fut.get();
+      {
+        std::lock_guard<std::mutex> lk(infl_mu_);
+        auto it = inflight_.find(key);
+        if (it != inflight_.end()) inflight_.erase(it);  // 530
+      }
+      if (cache_.peek(key)) return;                       // 531
+      if (!pbuf_.put(key, b)) pwaste_++;                  // 533-535
+    });
+  }
+}
+
+void ExpertIo::warm_willneed(int li, const std::vector<int64_t> &experts) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  for (auto e : experts)
+    for (int i = 0; i < 9; i++) {
+      if (!L.reg[i].ti) continue;
+      shards_.advise_range(files_->shard_path(L.reg[i].ti->shard),
+                           L.reg[i].ti->begin + uint64_t(e) * L.reg[i].per, L.reg[i].per);
+    }
+}
+
+void ExpertIo::warm_pages(int li, const std::vector<int64_t> &experts) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  volatile uint8_t sink = 0;
+  uint64_t page = uint64_t(sysconf(_SC_PAGESIZE));
+  for (auto e : experts)
+    for (int i = 0; i < 9; i++) {
+      if (!L.reg[i].ti) continue;
+      const uint8_t *src = shards_.raw(files_->shard_path(L.reg[i].ti->shard),
+                                       L.reg[i].ti->begin + uint64_t(e) * L.reg[i].per,
+                                       L.reg[i].per);
+      if (!src) continue;
+      for (uint64_t o = 0; o < L.reg[i].per; o += page) sink ^= src[o];
+    }
+  (void)sink;
+}
+
+void ExpertIo::shard_advise_whole(const std::string &shard_name) {
+  shards_.advise_whole(files_->shard_path(shard_name));
+}
+void ExpertIo::shard_seq_read(const std::string &shard_name) {
+  shards_.seq_read(files_->shard_path(shard_name));
+}
+
+void ExpertIo::stage_experts(int li, const std::vector<int64_t> &experts) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  std::vector<int64_t> uniq(experts.begin(), experts.end());
+  std::sort(uniq.begin(), uniq.end());
+  uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+  if (uniq.empty()) return;  // 721-722
+  if (L.cfg.warm_willneed) warm_willneed(li, uniq);  // 723-724
+  std::lock_guard<std::mutex> lk(*L.mu);
+  if (L.cfg.incr_mode) {
+    stage_incr_locked(L, uniq);
+    return;
+  }
+  // 733-738:next = _build_staged_bundles(unique)(539-585).
+  const int n = L.cfg.staged_n;
+  StagedState st;
+  st.exp.assign(uniq.begin(), uniq.begin() + std::min<int>(size_t(n), uniq.size()));
+  st.active = true;
+  for (auto e : st.exp) {
+    std::shared_ptr<Bundle> b;
+    for (size_t i = 0; i < L.state.exp.size(); i++)
+      if (L.state.exp[i] == e) { b = L.state.bundles[i]; break; }
+    if (!b) b = cache_.get(key_of(li, e));
+    if (!b) {
+      b = build(li, e);
+      if (!b) return;
+      cache_.put(key_of(li, e), b);
+      sbuilds_++;
+    }
+    st.bundles.push_back(std::move(b));
+  }
+  L.next = std::move(st);
+}
+
+void ExpertIo::wait_staged(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  std::lock_guard<std::mutex> lk(*L.mu);
+  if (L.next.active) {
+    L.state = L.next;
+    L.next = StagedState{};
+  }
+}
+const ExpertIo::StagedState &ExpertIo::staged(int li) const {
+  return layers_[size_t(layer_idx(li))].state;
+}
+const ExpertIo::StagedState &ExpertIo::next_staged(int li) const {
+  return layers_[size_t(layer_idx(li))].next;
+}
+
+void ExpertIo::stage_incr_locked(Layer &L, const std::vector<int64_t> &uniq) {
+  const int li = L.cfg.layer;
+  const int n = L.cfg.staged_n;
+  std::vector<int64_t> exp(uniq.begin(), uniq.begin() + std::min<int>(size_t(n), uniq.size()));
+  std::set<int64_t> exp_set(exp.begin(), exp.end());
+  if (L.incr_tensors.empty()) {
+    for (int i = 0; i < 9; i++) {
+      auto t = mx::zeros(mkshape({int64_t(n) + 1, L.reg[i].rows, L.reg[i].cols}),
+                         L.reg[i].esize == 4 ? mx::uint32 : mx::bfloat16);
+      mx::eval(t);
+      L.incr_tensors.push_back(std::move(t));
+    }
+    L.incr_table = mx::full(mkshape({L.cfg.E}), int32_t(n), mx::int32);  // core.full(614-616)
+    mx::eval(L.incr_table);
+    L.occ.assign(n, -1);
+  }
+  auto &occ = L.occ;
+  for (int i = 0; i < n; i++)
+    if (occ[i] >= 0 && !exp_set.count(occ[i])) occ[i] = -1;
+  std::vector<int> free;
+  for (int i = 0; i < n; i++) if (occ[i] < 0) free.push_back(i);
+  std::vector<int64_t> changed;  // 633-637
+  for (auto e : exp) if (!L.incr_bundles.count(e)) changed.push_back(e);
+  std::map<int64_t, std::shared_ptr<Bundle>> bundles;
+  for (auto e : exp) if (L.incr_bundles.count(e)) bundles[e] = L.incr_bundles[e];
+  struct Write { int slot; std::shared_ptr<Bundle> b; };
+  std::vector<Write> writes;
+  for (size_t j = 0; j < changed.size(); j++) {  // 646-668
+    int64_t e = changed[j];
+    if (int(j) >= int(free.size())) continue;
+    std::shared_ptr<Bundle> b = cache_.get(key_of(li, e));
+    if (!b) {
+      b = build(li, e);
+      if (!L.cfg.incr_writeback && b) cache_.put(key_of(li, e), b);
+      sbuilds_++;
+    }
+    int slot = free[j];
+    materialize(b);
+    bundles[e] = b;
+    occ[slot] = e;
+    writes.push_back({slot, std::move(b)});
+  }
+  std::vector<int64_t> evicted;  // 669-670
+  for (auto &[e, b] : L.incr_bundles)
+    if (!exp_set.count(e)) evicted.push_back(e);
+  if (L.cfg.incr_writeback)
+    for (auto e : evicted) cache_.put(key_of(li, e), L.incr_bundles[e]);
+  L.incr_bundles = std::move(bundles);
+  for (auto &w : writes)
+    for (int i = 0; i < 9; i++) {
+      uint64_t row = uint64_t(L.reg[i].rows) * uint64_t(L.reg[i].cols);
+      if (L.reg[i].esize == 4)
+        std::memcpy(L.incr_tensors[i].data<uint32_t>() + uint64_t(w.slot) * row,
+                    w.b->at(i).data<uint32_t>(), row * 4);
+      else
+        std::memcpy(L.incr_tensors[i].data<uint16_t>() + uint64_t(w.slot) * row,
+                    w.b->at(i).data<uint16_t>(), row * 2);
+    }
+  int32_t *tbl = L.incr_table.data<int32_t>();
+  for (auto e : evicted) tbl[e] = int32_t(n);
+  for (int i = 0; i < n; i++)
+    if (occ[i] >= 0) tbl[occ[i]] = int32_t(i);
+  if (!writes.empty() || !evicted.empty() || !changed.empty()) mx::eval(L.incr_table);
+  StagedState st;
+  st.active = true;
+  st.exp = exp;
+  for (auto e : exp) st.bundles.push_back(L.incr_bundles[e]);
+  L.next = std::move(st);
+}
+
+bool ExpertIo::incr_mode(int li) const {
+  int idx = layer_idx(li);
+  return idx >= 0 && layers_[size_t(idx)].cfg.incr_mode;
+}
+bool ExpertIo::incr_ready(int li) const {
+  int idx = layer_idx(li);
+  return idx >= 0 && layers_[size_t(idx)].cfg.incr_mode &&
+         !layers_[size_t(idx)].incr_tensors.empty();
+}
+const std::vector<mx::array> &ExpertIo::incr_tensors(int li) {
+  return layers_[size_t(layer_idx(li))].incr_tensors;
+}
+const mx::array &ExpertIo::incr_slot_table(int li) const {
+  return layers_[size_t(layer_idx(li))].incr_table;
+}
+
+static std::string exp_key(const std::vector<int64_t> &exp) {
+  std::string s;
+  for (auto e : exp) { s += std::to_string(e); s += ','; }
+  return s;
+}
+
+ExpertIo::Asm ExpertIo::asm_for(int li, const StagedState &st) {
+  Layer &L = layers_[size_t(layer_idx(li))];
+  const std::string key = exp_key(st.exp);
+  if (L.cfg.asm_cache) {
+    auto it = L.asm_cache.find(key);
+    if (it != L.asm_cache.end() && it->second.exp == st.exp) return it->second;
+  }
+  const int n = L.cfg.staged_n;
+  if (L.zero_rows.empty()) {
+    for (int i = 0; i < 9; i++)
+      L.zero_rows.push_back(mx::zeros(mkshape({L.reg[i].rows, L.reg[i].cols}),
+                                      L.reg[i].esize == 4 ? mx::uint32 : mx::bfloat16));
+  }
+  Asm a;
+  a.exp = st.exp;
+  auto sit = L.slot_cache.find(key);  // _slot_table_cache(1269-1276)
+  if (sit == L.slot_cache.end()) {
+    std::vector<int32_t> tab(size_t(L.cfg.E), int32_t(n));  // 557-559
+    for (size_t i = 0; i < st.exp.size(); i++) tab[size_t(st.exp[i])] = int32_t(i);
+    mx::array tof(tab.begin(), mkshape({L.cfg.E}), mx::int32);  // core.array(list)
+    mx::eval(tof);  // 1271-1272 force eval
+    L.slot_cache.insert_or_assign(key, tof);
+    L.slot_keys.push_back(key);
+    if (L.slot_keys.size() > 8) {
+      L.slot_cache.erase(L.slot_keys.front());
+      L.slot_keys.pop_front();
+    }
+    a.slot_of = L.slot_cache.at(key);
+  } else {
+    a.slot_of = sit->second;
+  }
+  for (int i = 0; i < 9; i++) {
+    std::vector<mx::array> rows;
+    for (auto &b : st.bundles) rows.push_back(b->at(i));
+    while (int(rows.size()) < n + 1) rows.push_back(L.zero_rows[i]);
+    a.wargs.push_back(mx::stack(rows));
+  }
+  sused_ += st.exp.size();
+  if (L.cfg.asm_cache) {
+    if (!L.asm_cache.empty()) L.asm_cache.clear();  // cap=1
+    L.asm_cache.insert_or_assign(key, a);
+  }
+  return a;
+}
+
+void ExpertIo::note_call(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  L.calls++;
+  if (L.cfg.hot_update_interval && L.calls % L.cfg.hot_update_interval == 0) {
+    (void)L.cfg.hot_per_layer;
+  }
+}
+
+void ExpertIo::refresh_hot_pins(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  if (layers_[size_t(idx)].cfg.hot_per_layer <= 0) return;
+}
+
+void ExpertIo::load_hot_layer(int li, int n_hot) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  std::vector<int64_t> hot;
+  if (!L.hc_order.empty()) {
+    hot = L.hc_order;
+    std::stable_sort(hot.begin(), hot.end(),
+                     [&](int64_t a, int64_t b) { return L.hc_val[a] > L.hc_val[b]; });
+    if ((int64_t)hot.size() > n_hot) hot.resize(n_hot);  // hot[:n_hot](1007)
+  } else {
+    for (int64_t e = 0; e < n_hot; e++) hot.push_back(e);  // 1009
+  }
+  if (L.hot_backing) {
+    if (L.hot_key == hot) return;  // 1012-1013
+    std::set<int64_t> old_set(L.hot_key.begin(), L.hot_key.end());
+    int changed = 0;
+    for (auto e : hot)
+      if (!old_set.count(e)) changed++;
+    if (changed <= std::max(1, n_hot / 8)) return;  // 1016-1017
+  }
+  auto st = std::make_shared<Layer::HotStack>();
+  st->key = hot;
+  st->sorted = st->key;
+  std::sort(st->sorted.begin(), st->sorted.end());
+  {
+    std::lock_guard<std::mutex> lk(*L.backing_mu);
+    L.hot_backing = st;
+  }
+  L.hot_key = st->key;
+  if (!L.hot_weights.empty()) {
+    heap_->sub(L.hot_heap);
+    L.hot_heap = 0;
+    L.hot_weights.clear();
+  }
+}
+
+void ExpertIo::materialize_hot(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  if (!L.hot_weights.empty() || !L.hot_backing) return;  // 1053
+  std::shared_ptr<const Layer::HotStack> snap;
+  {
+    std::lock_guard<std::mutex> lk(*L.backing_mu);
+    snap = L.hot_backing;
+  }
+  const int64_t n_e = (int64_t)snap->key.size() + 1;
+  L.hot_weights.assign(9, mx::array(0.f));
+  for (int i = 0; i < 9; i++) {
+    const auto &r = L.reg[i];
+    if (!r.ti) return;
+    std::vector<uint8_t> rows(static_cast<size_t>(n_e) * r.per, 0);
+    const std::string path = files_->shard_path(r.ti->shard);
+    for (size_t j = 0; j < snap->key.size(); j++) {
+      const uint8_t *src = shards_.raw(path, r.ti->begin + uint64_t(snap->key[j]) * r.per,
+                                       r.per);
+      if (!src) return;
+      std::memcpy(rows.data() + uint64_t(j) * r.per, src, r.per);
+    }
+    const uint8_t *base = rows.data();
+    if (r.esize == 4) {
+      L.hot_weights[i] = mx::array(reinterpret_cast<const uint32_t *>(base),
+                                   mkshape({n_e, r.rows, r.cols}), mx::uint32);
+    } else {
+      auto a16 = mx::array(reinterpret_cast<const uint16_t *>(base),
+                           mkshape({n_e, r.rows, r.cols}), mx::uint16);
+      L.hot_weights[i] = mx::view(a16, mx::bfloat16);
+    }
+    L.hot_heap += rows.size();
+  }
+  heap_->add(L.hot_heap);
+}
+
+void ExpertIo::dematerialize_hot(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  if (L.hot_weights.empty()) return;
+  heap_->sub(L.hot_heap);
+  L.hot_heap = 0;
+  L.hot_weights.clear();
+}
+
+void ExpertIo::clear_hot_layer(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  if (!L.hot_weights.empty()) {
+    heap_->sub(L.hot_heap);
+    L.hot_heap = 0;
+    L.hot_weights.clear();
+  }
+  std::lock_guard<std::mutex> lk(*L.backing_mu);
+  L.hot_backing = nullptr;
+  L.hot_key.clear();
+}
+
+void ExpertIo::load_full_layer(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  if (!L.full_w.empty()) return;  // 934-935
+  const int64_t E = L.cfg.E;
+  L.full_w.assign(9, mx::array(0.f));
+  uint64_t tot = 0;
+  for (int i = 0; i < 9; i++) {
+    const auto &r = L.reg[i];
+    if (!r.ti) {
+      std::fprintf(stderr, "expertio: load_full_layer missing region L%d r%d\n", li, i);
+      return;
+    }
+    const uint64_t span = uint64_t(E) * r.per;
+    const uint8_t *raw = shards_.raw(files_->shard_path(r.ti->shard), r.ti->begin, span);
+    if (!raw) {
+      std::fprintf(stderr, "expertio: load_full_layer raw OOB L%d span=%llu\n", li,
+                   (unsigned long long)span);
+      return;
+    }
+    if (r.esize == 4) {  // 974-975 u32_view(raw,shape) → core.array
+      L.full_w[i] =
+          mx::array(reinterpret_cast<const uint32_t *>(raw), mkshape({E, r.rows, r.cols}),
+                    mx::uint32);
+    } else {  // 976-978 raw.view(<u2) → core.array → .view(bf16)
+      auto a16 = mx::array(reinterpret_cast<const uint16_t *>(raw),
+                           mkshape({E, r.rows, r.cols}), mx::uint16);
+      L.full_w[i] = mx::view(a16, mx::bfloat16);
+    }
+    tot += span;
+  }
+  L.full_heap = tot;
+  heap_->add(tot);
+}
+
+void ExpertIo::clear_full_layer(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  if (L.full_w.empty()) return;
+  heap_->sub(L.full_heap);
+  L.full_heap = 0;
+  L.full_w.clear();
+}
+
+bool ExpertIo::hot_ready(int li) const {
+  int idx = layer_idx(li);
+  return idx >= 0 && !layers_[size_t(idx)].hot_weights.empty();  // 1121 `_hot_weights is not None`
+}
+bool ExpertIo::full_ready(int li) const {
+  int idx = layer_idx(li);
+  return idx >= 0 && !layers_[size_t(idx)].full_w.empty();  // 1231 `_full_weights is not None`
+}
+const std::vector<mx::array> &ExpertIo::hot_weights(int li) {
+  static const std::vector<mx::array> kEmpty;
+  int idx = layer_idx(li);
+  return idx >= 0 ? layers_[size_t(idx)].hot_weights : kEmpty;
+}
+const std::vector<mx::array> &ExpertIo::full_weights(int li) {
+  static const std::vector<mx::array> kEmpty;
+  int idx = layer_idx(li);
+  return idx >= 0 ? layers_[size_t(idx)].full_w : kEmpty;
+}
+const std::vector<int64_t> &ExpertIo::hot_key(int li) const {
+  static const std::vector<int64_t> kEmpty;
+  int idx = layer_idx(li);
+  return idx >= 0 ? layers_[size_t(idx)].hot_key : kEmpty;
+}
+const std::vector<int64_t> &ExpertIo::last_used(int li) const {
+  static const std::vector<int64_t> kEmpty;
+  int idx = layer_idx(li);
+  return idx >= 0 ? layers_[size_t(idx)].last_used : kEmpty;
+}
+
+void ExpertIo::note_prefill_use(int li, const std::vector<int64_t> &flat, int k,
+                                bool staged_mode) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  std::vector<int64_t> uniq = flat;
+  std::sort(uniq.begin(), uniq.end());
+  uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+  for (auto e : uniq) L.bump_count(e, L.cfg.hot_decay);  // 1125-1127
+  if (staged_mode && (int)flat.size() >= k)
+    L.last_prefill_topk.assign(flat.end() - k, flat.end());  // 1128-1130
+}
+
+void ExpertIo::prefetch_from_prefill(int li) {
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  if (L.last_prefill_topk.empty()) return;  // 887-888
+  std::vector<int64_t> uniq = L.last_prefill_topk;
+  std::sort(uniq.begin(), uniq.end());
+  uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+  if (uniq.empty()) return;  // 890
+  L.last_used = uniq;        // 891
+  prefetch(li, uniq);        // 892
+}
+
+// Cross-request memory plus incremental slot table cache.
+void ExpertIo::reset_layer(int li) {  // layer.py:1429-1441
+  int idx = layer_idx(li);
+  if (idx < 0) return;
+  Layer &L = layers_[size_t(idx)];
+  std::lock_guard<std::mutex> lk(*L.mu);
+  L.state = StagedState{};   // _staged_state = None(1436)
+  L.next = StagedState{};    // _staged_state_next = None(1437)
+  L.last_prefill_topk.clear();  // 1435
+  L.last_used.clear();          // 1440
+  if (!L.full_w.empty()) {
+    heap_->sub(L.full_heap);
+    L.full_heap = 0;
+    L.full_w.clear();
+  }
+}
+
+ExpertIo::Metrics ExpertIo::snapshot() const {
+  Metrics m;
+  m.loads = loads_.load();
+  m.hits = hits_.load();
+  m.misses = misses_.load();
+  m.evictions = cache_.evictions();
+  m.prefetch_submitted = psub_.load();
+  m.prefetch_hits = phi_.load();
+  m.prefetch_wasted = pwaste_.load();
+  m.staged_used = sused_.load();
+  m.staged_fallback = sfallback_.load();
+  m.stage_builds = sbuilds_.load();
+  m.heap_bytes = heap_->live.load();
+  m.heap_peak = heap_->peak.load();
+  m.cache_size = const_cast<SharedExpertCache &>(cache_).size();
+  m.inflight = inflight_size();
+  return m;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/expertio.hpp b/macos/crates/engine-native/src/expertio.hpp
new file mode 100644
index 0000000..1f4d17e
--- /dev/null
+++ b/macos/crates/engine-native/src/expertio.hpp
@@ -0,0 +1,279 @@
+// Expert load/staging: mmap shards, global LRU, prefetch buffer, and staged double-buffer.
+// Stack/lazy-graph construction stays on the main thread; pool threads only produce Bundles.
+// Prefetch done-callbacks are registered outside the lock (deadlock otherwise).
+#pragma once
+#include <atomic>
+#include <condition_variable>
+#include <cstdint>
+#include <deque>
+#include <functional>
+#include <future>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <thread>
+#include <unordered_map>
+#include <vector>
+
+#include <mlx/mlx.h>
+
+#include "safio.hpp"
+
+namespace e0n {
+namespace mx = mlx::core;
+
+// 9-region bundle. Field order is the moe_math signature: (up,gate,down)×(w,s,b),
+// not file key order (gate_proj before up_proj). Shape is the stacked tensor
+// without the expert axis [O,K]; elements live in mlx-owned buffers.
+struct Bundle {
+  Bundle()
+      : f{mx::array(0.f), mx::array(0.f), mx::array(0.f), mx::array(0.f), mx::array(0.f),
+         mx::array(0.f), mx::array(0.f), mx::array(0.f), mx::array(0.f)} {}
+  mx::array f[9];  // 0..2=up w/s/b, 3..5=gate w/s/b, 6..8=down w/s/b
+  uint64_t bytes = 0;
+  const mx::array &at(int p) const { return f[p]; }
+  mx::array &at_mut(int p) { return f[p]; }
+};
+
+class SharedExpertCache {
+ public:
+  explicit SharedExpertCache(int slots) : slots_(slots) {}
+  std::shared_ptr<Bundle> get(uint64_t key);
+  std::shared_ptr<Bundle> peek(uint64_t key);
+  void put(uint64_t key, std::shared_ptr<Bundle> b);
+  std::shared_ptr<Bundle> pop(uint64_t key);
+  bool contains(uint64_t key);
+  size_t size();
+  uint64_t evictions() const { return evictions_; }
+
+ private:
+  int slots_;
+  uint64_t evictions_ = 0;
+  std::mutex mu_;
+  std::deque<uint64_t> order_;  // front = LRU
+  std::map<uint64_t, std::shared_ptr<Bundle>> map_;
+};
+
+class PrefetchBuffer {
+ public:
+  explicit PrefetchBuffer(int cap) : cap_(cap) {}
+  bool put(uint64_t key, std::shared_ptr<Bundle> b);     // true = kept, false = overflow drop
+  std::shared_ptr<Bundle> pop(uint64_t key);
+  bool contains(uint64_t key);
+  void set_cap(int c) { cap_ = c; }
+
+ private:
+  int cap_;
+  std::mutex mu_;
+  std::deque<uint64_t> order_;
+  std::map<uint64_t, std::shared_ptr<Bundle>> map_;
+};
+
+class IoPool {
+ public:
+  explicit IoPool(int n);
+  ~IoPool();
+  template <class F>
+  auto submit(F &&f) -> std::shared_future<decltype(f())> {
+    using R = decltype(f());
+    auto pk = std::make_shared<std::packaged_task<R()>>(std::forward<F>(f));
+    auto fut = pk->get_future().share();
+    {
+      std::lock_guard<std::mutex> lk(mu_);
+      q_.emplace_back([pk] { (*pk)(); });
+    }
+    cv_.notify_one();
+    return fut;
+  }
+
+ private:
+  std::mutex mu_;
+  std::condition_variable cv_;
+  std::deque<std::function<void()>> q_;
+  bool stop_ = false;
+  std::vector<std::thread> th_;
+};
+
+class ExpertIo {
+ public:
+  struct LayerCfg {
+    int layer = 0;
+    std::string prefix;   // "model.layers.N.mlp.experts" / "...mlp.switch_mlp"
+    int64_t E = 0;
+    int64_t H = 0, F = 0;
+    int group = 64;
+    int staged_n = 8;            // prod_k8=8 / staged_k4=4
+    bool incr_mode = false;      // 8b prod_k8 false; 35b staged_k4 true
+    bool incr_writeback = false; // 35b true
+    bool warm_willneed = false;  // 35b true / 8b false
+    bool asm_cache = true;
+    int prefill_hot = 0;             // 35b staged_k4=32 / 8b prod_k8=0
+    bool full_layer_prefill = false; // 35b false / 8b true
+    int prefill_full_layers = 0;
+    int hot_per_layer = 0;
+    int hot_update_interval = 4;
+    double hot_decay = 0.75;
+  };
+
+  ExpertIo(const ModelFiles *files, int cache_slots, int prefetch_cap, int load_threads,
+           int prefetch_threads);
+  void add_layer(const LayerCfg &cfg);
+
+  std::shared_ptr<Bundle> build(int li, int64_t expert);
+  // LRU → prefetch_buf.pop (promote) → wait inflight (promote) → pool sync + LRU.
+  std::map<int64_t, std::shared_ptr<Bundle>> get_bundles(int li,
+                                                         const std::vector<int64_t> &experts);
+  void prefetch(int li, const std::vector<int64_t> &experts);
+  void warm_pages(int li, const std::vector<int64_t> &experts);
+  void warm_willneed(int li, const std::vector<int64_t> &experts);
+  void shard_advise_whole(const std::string &shard_name);
+  void shard_seq_read(const std::string &shard_name);
+
+  struct StagedState {
+    std::vector<std::shared_ptr<Bundle>> bundles;
+    std::vector<int64_t> exp;  // sorted(set)[:staged_n]
+    bool active = false;
+  };
+  void stage_experts(int li, const std::vector<int64_t> &experts);  // main thread
+  void wait_staged(int li);
+  const StagedState &staged(int li) const;
+  const StagedState &next_staged(int li) const;
+
+  bool incr_mode(int li) const;
+  bool incr_ready(int li) const;
+  const std::vector<mx::array> &incr_tensors(int li);
+  const mx::array &incr_slot_table(int li) const;
+
+  // Non-incr consume path: slot table (set-key cache, cap 8) + 9×[n+1,...] stack (cap 1).
+  // Returns a cheap copy (mx::array is a node pointer). asm_cache=false rebuilds each call.
+  struct Asm {
+    std::vector<mx::array> wargs;  // 9×[n+1,...]
+    mx::array slot_of{0.f};        // [E] int32, eval'd GPU-resident table
+    std::vector<int64_t> exp;
+  };
+  Asm asm_for(int li, const StagedState &st);
+
+  void load_hot_layer(int li, int n_hot);
+  void materialize_hot(int li);
+  void dematerialize_hot(int li);
+  void clear_hot_layer(int li);
+  void load_full_layer(int li);
+  void clear_full_layer(int li);
+  bool hot_ready(int li) const;
+  bool full_ready(int li) const;
+  const std::vector<mx::array> &hot_weights(int li);
+  const std::vector<mx::array> &full_weights(int li);
+  const std::vector<int64_t> &hot_key(int li) const;
+  void note_call(int li);
+  void refresh_hot_pins(int li);
+  void note_prefill_use(int li, const std::vector<int64_t> &flat, int k, bool staged_mode);
+  void prefetch_from_prefill(int li);
+  const std::vector<int64_t> &last_used(int li) const;
+
+  void reset_layer(int li);  // only _staged_state/_next; the rest survives
+
+  struct Metrics {
+    uint64_t loads = 0, hits = 0, misses = 0, evictions = 0;
+    uint64_t prefetch_submitted = 0, prefetch_hits = 0, prefetch_wasted = 0;
+    uint64_t staged_used = 0, staged_fallback = 0, stage_builds = 0;
+    uint64_t heap_bytes = 0, heap_peak = 0;
+    size_t cache_size = 0, inflight = 0;
+  };
+  Metrics snapshot() const;
+  size_t inflight_size() const {
+    std::lock_guard<std::mutex> lk(infl_mu_);
+    return inflight_.size();
+  }
+
+  static uint64_t key_of(int li, int64_t e) {
+    return (uint64_t(li) << 32) | uint64_t(e);
+  }
+
+ private:
+  struct Layer {
+    LayerCfg cfg;
+    struct Region {
+      const TensorInfo *ti = nullptr;
+      int64_t rows = 0, cols = 0;
+      uint64_t per = 0;
+      uint32_t esize = 4;          // element bytes (u32=4 / bf16=2)
+    };
+    Region reg[9];
+    std::unique_ptr<std::mutex> mu = std::make_unique<std::mutex>();
+    StagedState state, next;
+    std::deque<std::string> slot_keys;
+    std::map<std::string, mx::array> slot_cache;
+    std::map<std::string, Asm> asm_cache;
+    std::vector<mx::array> zero_rows;
+    std::vector<mx::array> incr_tensors;
+    mx::array incr_table{0.f};
+    std::vector<int64_t> occ;
+    std::map<int64_t, std::shared_ptr<Bundle>> incr_bundles;
+    // Hot expert keys only; shard mmap holds the bytes. shared_ptr snapshot so
+    // pool threads can read keys without UAF.
+    struct HotStack {
+      std::vector<int64_t> key;
+      std::vector<int64_t> sorted;
+    };
+    std::shared_ptr<const HotStack> hot_backing;
+    // Separate from L.mu: build is also called from stage_experts while L.mu is held
+    // (non-recursive lock). Writers (load/clear_hot_layer) vs pool-thread build readers.
+    std::unique_ptr<std::mutex> backing_mu = std::make_unique<std::mutex>();
+    std::vector<int64_t> hot_key;
+    std::vector<mx::array> hot_weights;
+    uint64_t hot_heap = 0;
+    std::vector<mx::array> full_w;
+    uint64_t full_heap = 0;
+    std::vector<int64_t> hc_order;
+    std::unordered_map<int64_t, double> hc_val;
+    std::vector<int64_t> last_prefill_topk;
+    std::vector<int64_t> last_used;
+    int64_t calls = 0;
+    // New keys append; existing keys update in place (dict insertion order).
+    void bump_count(int64_t e, double decay) {
+      auto it = hc_val.find(e);
+      if (it == hc_val.end()) { hc_order.push_back(e); hc_val.emplace(e, 1.0); }
+      else it->second = it->second * decay + 1.0;
+    }
+  };
+
+  void stage_incr_locked(Layer &L, const std::vector<int64_t> &uniq);
+  int layer_idx(int li) const;
+
+  struct HeapAcc {
+    std::atomic<uint64_t> live{0}, peak{0};
+    void add(uint64_t b) {
+      uint64_t v = live.fetch_add(b) + b;
+      uint64_t p = peak.load();
+      while (v > p && !peak.compare_exchange_weak(p, v)) {
+      }
+    }
+    void sub(uint64_t b) { live.fetch_sub(b); }
+  };
+
+  const ModelFiles *files_;
+  struct ShardMaps {
+    ~ShardMaps();
+    const uint8_t *raw(const std::string &shard_path, uint64_t off, uint64_t len);
+    void advise_range(const std::string &shard_path, uint64_t off, uint64_t len);
+    void advise_whole(const std::string &shard_path);
+    void seq_read(const std::string &shard_path);
+    std::mutex mu_;
+    std::map<std::string, std::pair<uint8_t *, uint64_t>> maps_;
+    std::map<std::string, int> fds_;
+  } shards_;
+
+  std::vector<Layer> layers_;
+  std::map<int, int> li2idx_;
+  SharedExpertCache cache_;
+  PrefetchBuffer pbuf_;
+  IoPool load_pool_, prefetch_pool_;
+  mutable std::mutex infl_mu_;
+  std::map<uint64_t, std::shared_future<std::shared_ptr<Bundle>>> inflight_;
+  std::shared_ptr<HeapAcc> heap_{std::make_shared<HeapAcc>()};
+  std::atomic<uint64_t> loads_{0}, hits_{0}, misses_{0}, psub_{0}, phi_{0}, pwaste_{0},
+      sused_{0}, sfallback_{0}, sbuilds_{0};
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/expertmap.cpp b/macos/crates/engine-native/src/expertmap.cpp
new file mode 100644
index 0000000..924a62b
--- /dev/null
+++ b/macos/crates/engine-native/src/expertmap.cpp
@@ -0,0 +1,119 @@
+#include "expertmap.hpp"
+
+#include <cstdio>
+
+namespace e0n {
+namespace {
+
+uint64_t dtype_size(const std::string& d) {
+  if (d == "BF16" || d == "F16") return 2;
+  if (d == "F32" || d == "U32" || d == "I32") return 4;
+  if (d == "F64") return 8;
+  if (d == "U8" || d == "I8" || d == "BOOL") return 1;
+  return 0;
+}
+
+constexpr const char* kParts[] = {"gate_proj", "up_proj", "down_proj"};
+constexpr const char* kFields[] = {"weight", "scales", "biases"};
+
+}  // namespace
+
+bool ExpertMap::build(const ModelFiles& files, ExpertMap& out, std::string& err) {
+  int64_t E = 0, L = 0;
+  if (!files.config_int("num_experts", E) || E <= 0) {
+    err = "config.json missing num_experts";
+    return false;
+  }
+  if (!files.config_int("num_hidden_layers", L) || L <= 0) {
+    err = "config.json missing num_hidden_layers";
+    return false;
+  }
+  // qwen (35b): presence of language_model.model.layers.0.mlp.gate.weight.
+  const bool qwen = files.find("language_model.model.layers.0.mlp.gate.weight") != nullptr;
+
+  out.num_layers_ = static_cast<int>(L);
+  out.num_experts_ = E;
+  out.prefixes_.assign(static_cast<size_t>(L), "");
+  out.layers_.clear();
+  out.layers_.resize(static_cast<size_t>(L));
+
+  for (int l = 0; l < static_cast<int>(L); l++) {
+    const std::string base =
+        qwen ? "language_model.model.layers." + std::to_string(l)
+             : "model.layers." + std::to_string(l);
+    std::string prefix;
+    for (const std::string& cand :
+         {base + ".mlp.experts", base + ".mlp.switch_mlp", base + ".mlp"}) {
+      if (files.find(cand + ".gate_proj.weight")) { prefix = cand; break; }
+    }
+    if (prefix.empty()) {
+      err = "layer " + std::to_string(l) + " has no MoE fused prefix";
+      return false;
+    }
+    out.prefixes_[static_cast<size_t>(l)] = prefix;
+
+    auto& slices = out.layers_[static_cast<size_t>(l)];
+    for (const char* p : kParts) {
+      for (const char* f : kFields) {
+        const std::string name = std::string(p) + "." + f;
+        const TensorInfo* t = files.find(prefix + "." + name);
+        if (!t) {
+          err = "missing tensor " + prefix + "." + name;
+          return false;
+        }
+        const uint64_t es = dtype_size(t->dtype);
+        if (es == 0) {
+          err = "unknown dtype " + t->dtype + " for " + name;
+          return false;
+        }
+        Slice s;
+        if (t->shape.size() >= 3 && t->shape[0] == E) {
+          // First dim = expert count (35b): one expert is a contiguous stride.
+          s.stride = static_cast<uint64_t>(t->shape[1]) * static_cast<uint64_t>(t->shape[2]) * es;
+        } else if (t->shape.size() == 2 && t->shape[0] % E == 0) {
+          // E folded into the row dim (8b and 8b-L0 dense).
+          s.stride = (static_cast<uint64_t>(t->shape[0]) / static_cast<uint64_t>(E)) *
+                     static_cast<uint64_t>(t->shape[1]) * es;
+        } else {
+          err = "unsupported layout " + prefix + "." + name;
+          return false;
+        }
+        s.base = t->begin;
+        s.shard = t->shard;
+        s.name = name;
+        slices.emplace(name, s);
+      }
+    }
+  }
+  return true;
+}
+
+std::vector<Region> ExpertMap::regions(int layer, int64_t expert) const {
+  std::vector<Region> out;
+  if (layer < 0 || layer >= num_layers_ || expert < 0 || expert >= num_experts_) return out;
+  out.reserve(9);
+  for (const char* p : kParts) {
+    for (const char* f : kFields) {
+      auto it = layers_[static_cast<size_t>(layer)].find(std::string(p) + "." + f);
+      if (it == layers_[static_cast<size_t>(layer)].end()) return {};
+      const Slice& s = it->second;
+      out.push_back(Region{s.base + static_cast<uint64_t>(expert) * s.stride, s.stride,
+                           s.shard, s.name});
+    }
+  }
+  return out;
+}
+
+uint64_t ExpertMap::expert_bytes(int layer, int64_t expert) const {
+  uint64_t total = 0;
+  for (auto& r : regions(layer, expert)) total += r.len;
+  return total;
+}
+
+const std::string& ExpertMap::layer_prefix(int layer) const {
+  static const std::string empty;
+  if (layer < 0 || layer >= num_layers_) return empty;
+  return prefixes_[static_cast<size_t>(layer)];
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/expertmap.hpp b/macos/crates/engine-native/src/expertmap.hpp
new file mode 100644
index 0000000..de9c6af
--- /dev/null
+++ b/macos/crates/engine-native/src/expertmap.hpp
@@ -0,0 +1,42 @@
+// (layer, expert) → contiguous row slices inside fused tensors. 35b: shape[0]=E;
+// 8b: E folded into the row dim. Nine regions: {gate,up,down}_proj × {weight,scales,biases}.
+#pragma once
+#include <cstdint>
+#include <map>
+#include <string>
+#include <vector>
+
+#include "safio.hpp"
+
+namespace e0n {
+
+struct Region {
+  uint64_t off, len;
+  std::string shard;
+  std::string name;   // "gate_proj.weight" etc., diagnostics only
+};
+
+class ExpertMap {
+ public:
+  static bool build(const ModelFiles& files, ExpertMap& out, std::string& err);
+
+  int num_layers() const { return num_layers_; }
+  int64_t num_experts() const { return num_experts_; }
+
+  std::vector<Region> regions(int layer, int64_t expert) const;
+  uint64_t expert_bytes(int layer, int64_t expert) const;
+
+  const std::string& layer_prefix(int layer) const;
+
+ private:
+  struct Slice {
+    uint64_t base, stride;
+    std::string shard, name;
+  };
+  int num_layers_ = 0;
+  int64_t num_experts_ = 0;
+  std::vector<std::string> prefixes_;
+  std::vector<std::map<std::string, Slice>> layers_;
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/fwd_common.cpp b/macos/crates/engine-native/src/fwd_common.cpp
new file mode 100644
index 0000000..1cd0139
--- /dev/null
+++ b/macos/crates/engine-native/src/fwd_common.cpp
@@ -0,0 +1,213 @@
+#include "fwd_common.hpp"
+
+#include <chrono>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+#include "dense.hpp"
+#include "hotstack.hpp"  // Entry (9-region expert blob) lives here, not in expertmap.hpp
+
+namespace e0n {
+namespace mx = mlx::core;
+
+mlx::core::Shape mkshape(std::initializer_list<int64_t> dims) {
+  mlx::core::Shape s;
+  for (auto d : dims) s.push_back(static_cast<int32_t>(d));
+  return s;
+}
+
+bool QuantLin::load(const ModelFiles &files, const std::string &weight_name, std::string &err) {
+  auto dot = weight_name.rfind('.');
+  const std::string mod = dot == std::string::npos ? weight_name : weight_name.substr(0, dot);
+  if (!files.read_tensor(weight_name, wb, err) || !files.read_tensor(mod + ".scales", sb, err))
+    return false;
+  std::string b_err;
+  if (!files.read_tensor(mod + ".biases", bb, b_err)) bb.clear();
+  const TensorInfo *tw = files.find(weight_name), *ts = files.find(mod + ".scales");
+  if (tw->dtype != "U32" || ts->dtype != "BF16") {
+    err = "QuantLin expected U32 weight + BF16 scales: " + weight_name;
+    return false;
+  }
+  if (!files.config_quant_for(weight_name, bits, group)) {
+    err = "QuantLin missing quantization config: " + weight_name;
+    return false;
+  }
+  const int64_t per_word = 32 / bits;
+  int64_t r = 1;
+  for (size_t i = 0; i + 1 < tw->shape.size(); i++) r *= tw->shape[i];
+  rows = r;
+  cols = tw->shape.back() * per_word;
+  if (ts->shape.back() * group != cols) {
+    err = "QuantLin scale group count mismatch: " + weight_name;
+    return false;
+  }
+  mlx::core::array wa(wb.data(), mkshape({rows, tw->shape.back()}), mlx::core::uint32,
+                      [](void *) {});
+  mlx::core::array sa(sb.data(), mkshape({rows, ts->shape.back()}), mlx::core::bfloat16,
+                      [](void *) {});
+  w = wa;
+  s = sa;
+  if (!bb.empty())
+    b = mlx::core::array(bb.data(), mkshape({rows, ts->shape.back()}), mlx::core::bfloat16,
+                         [](void *) {});
+  loaded = true;
+  return true;
+}
+
+mlx::core::array QuantLin::apply(const mlx::core::array &x) const {
+  return mlx::core::quantized_matmul(x, w, s, biases(), /*transpose=*/true, group, bits,
+                                     "affine");
+}
+
+bool RawLin::load(const ModelFiles &files, const std::string &name, std::string &err) {
+  const TensorInfo *t = files.find(name);
+  if (!t) { err = "missing tensor: " + name; return false; }
+  if (t->dtype != "BF16" && t->dtype != "F16" && t->dtype != "F32") {
+    err = "RawLin accepts only BF16/F16/F32: " + name;
+    return false;
+  }
+  if (!files.read_tensor(name, wb, err)) return false;
+  int64_t r = 1;
+  for (size_t i = 0; i + 1 < t->shape.size(); i++) r *= t->shape[i];
+  if (t->shape.size() == 1) r = 1;
+  rows = r;
+  cols = t->shape.back();
+  auto dt = t->dtype == "BF16"   ? mlx::core::bfloat16
+            : t->dtype == "F16"  ? mlx::core::float16
+                                 : mlx::core::float32;
+  w = mlx::core::array(wb.data(), mkshape(t->shape), dt, [](void *) {});
+  return true;
+}
+
+mlx::core::array RawLin::apply(const mlx::core::array &x) const {
+  return mlx::core::matmul(x, mlx::core::transpose(w));
+}
+
+// Upstream nn.silu is x * sigmoid(x); x/(1+exp(-x)) is algebraically equal but
+// rounds differently (one division vs reciprocal+mul).
+mlx::core::array e0_silu(const mlx::core::array &x) {
+  return x * mlx::core::sigmoid(x);
+}
+
+mlx::core::array e0_swiglu_precise(mlx::core::Dtype out_dt, const mlx::core::array &gate,
+                                   const mlx::core::array &up) {
+  auto g = e0_silu(mx::astype(gate, mx::float32));
+  auto u = mx::astype(up, mx::float32);
+  return mx::astype(g * u, out_dt);
+}
+
+mlx::core::array e0_rms_norm(const mlx::core::array &x,
+                             const std::optional<mlx::core::array> &weight, double eps) {
+  return mx::fast::rms_norm(x, weight, static_cast<float>(eps));
+}
+
+ExecCounters &exec_counters() {
+  static ExecCounters c;
+  return c;
+}
+
+uint64_t TimeSink::monotonic_ns() {
+  return std::chrono::duration_cast<std::chrono::nanoseconds>(
+             std::chrono::steady_clock::now().time_since_epoch())
+      .count();
+}
+bool TimingCounters::enabled() {
+  static const bool on = [] {
+    const char *e = getenv("EDGE0_TIMING");
+    return e && !strcmp(e, "1");
+  }();
+  return on;
+}
+TimingCounters &timing_counters() {
+  static TimingCounters c;
+  return c;
+}
+
+ExecMode exec_mode() {
+  static const ExecMode m = [] {
+    const char *e = getenv("EDGE0_EXEC");
+    if (!e || !*e) return ExecMode::Batched;
+    if (!strcmp(e, "batched")) return ExecMode::Batched;
+    if (!strcmp(e, "legacy")) return ExecMode::Legacy;
+    std::fprintf(stderr, "edge0: unknown EDGE0_EXEC=%s, using legacy (batched|legacy)\n", e);
+    return ExecMode::Legacy;
+  }();
+  return m;
+}
+
+bool prefill_ondemand() {
+  static const bool on = [] {
+    const char *e = getenv("EDGE0_PREFILL");
+    if (!e || !*e) return true;
+    if (!strcmp(e, "ondemand")) return true;
+    if (!strcmp(e, "production")) return false;
+    std::fprintf(stderr, "edge0: unknown EDGE0_PREFILL=%s, using ondemand (ondemand|production)\n", e);
+    return true;
+  }();
+  return on;
+}
+
+bool to_host32(const mx::array &a, std::vector<float> &out, std::string &err) {
+  E0_TIME(timing_counters().host_sync_ns);
+  exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+  try {
+    auto cpu = mx::copy(mx::astype(a, mx::float32), mx::default_stream(mx::Device::cpu));
+    mx::eval(cpu);
+    const float *p = cpu.data<float>();
+    out.assign(p, p + cpu.size());
+  } catch (const std::exception &ex) {
+    err = std::string("to_host32: ") + ex.what();
+    return false;
+  }
+  return true;
+}
+
+mx::array wrap_f32(const std::vector<float> &v, mx::Shape shape) {
+  exec_counters().host_wraps.fetch_add(1, std::memory_order_relaxed);
+  return mx::array(v.begin(), std::move(shape), mx::float32);
+}
+
+bool expert_ptrs(const Entry &e, ExpertPtrs &p, std::string &err) {
+  if (e.regions.size() != 9) { err = "expert region count != 9"; return false; }
+  for (size_t i = 0; i < 9; i++) {
+    const uint8_t *base = e.blob.data() + e.blob_offsets[i];
+    const std::string &nm = e.regions[i].name;
+    const uint32_t **wp = nullptr;
+    const uint16_t **sp = nullptr, **bp = nullptr;
+    if (nm == "gate_proj.weight") wp = &p.gw;
+    else if (nm == "up_proj.weight") wp = &p.uw;
+    else if (nm == "down_proj.weight") wp = &p.dw;
+    else if (nm == "gate_proj.scales") sp = &p.gs;
+    else if (nm == "gate_proj.biases") bp = &p.gb;
+    else if (nm == "up_proj.scales") sp = &p.us;
+    else if (nm == "up_proj.biases") bp = &p.ub;
+    else if (nm == "down_proj.scales") sp = &p.ds;
+    else if (nm == "down_proj.biases") bp = &p.db;
+    else { err = "unknown region name " + nm; return false; }
+    if (wp) *wp = reinterpret_cast<const uint32_t *>(base);
+    if (sp) *sp = reinterpret_cast<const uint16_t *>(base);
+    if (bp) *bp = reinterpret_cast<const uint16_t *>(base);
+  }
+  return true;
+}
+
+bool embed_lookup(const QuantLin &emb, const std::vector<int32_t> &ids, int64_t H, int64_t V,
+                  std::vector<float> &out, std::string &err) {
+  const int64_t groups = H / 64, packed = H / 8;
+  if (emb.rows != V) { err = "embed rows != vocab"; return false; }
+  out.assign(ids.size() * static_cast<size_t>(H), 0.f);
+  for (size_t i = 0; i < ids.size(); i++) {
+    int64_t id = ids[i];
+    if (id < 0 || id >= V) { err = "token id out of range " + std::to_string(id); return false; }
+    const uint32_t *wp = reinterpret_cast<const uint32_t *>(emb.wb.data()) + id * packed;
+    const uint16_t *sp = reinterpret_cast<const uint16_t *>(emb.sb.data()) + id * groups;
+    const uint16_t *bp = reinterpret_cast<const uint16_t *>(emb.bb.data()) + id * groups;
+    std::vector<float> row;
+    dequant_row(wp, sp, bp, packed, 4, 64, row);
+    std::memcpy(out.data() + i * H, row.data(), H * 4);
+  }
+  return true;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/fwd_common.hpp b/macos/crates/engine-native/src/fwd_common.hpp
new file mode 100644
index 0000000..b552378
--- /dev/null
+++ b/macos/crates/engine-native/src/fwd_common.hpp
@@ -0,0 +1,144 @@
+// Resident non-expert weights and shared forward ops (safio read → heap → mlx no-op deleter).
+#pragma once
+#include <atomic>
+#include <cmath>
+#include <cstdint>
+#include <initializer_list>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <mlx/mlx.h>
+
+#include "safio.hpp"
+
+namespace e0n {
+
+mlx::core::Shape mkshape(std::initializer_list<int64_t> dims);
+inline mlx::core::Shape mkshape(const std::vector<int64_t> &dims) {
+  mlx::core::Shape s;
+  for (auto d : dims) s.push_back(static_cast<int32_t>(d));
+  return s;
+}
+
+// Quantized dense linear (attention / shared expert / lm_head / gate int8):
+// weight [O, I/pack] U32 + scales/biases [O, I/g] BF16; affine; bits/group per module.
+struct QuantLin {
+  std::vector<uint8_t> wb, sb, bb;  // resident bytes (stable array pointers)
+  mlx::core::array w{0.f}, s{0.f}, b{0.f};
+  int64_t rows = 0, cols = 0;       // logical [O, I]
+  int64_t bits = 4, group = 64;
+  bool loaded = false;
+
+  bool load(const ModelFiles &files, const std::string &weight_name, std::string &err);
+  std::optional<mlx::core::array> biases() const {
+    if (bb.empty()) return std::nullopt;
+    return b;
+  }
+  mlx::core::array apply(const mlx::core::array &x) const;
+};
+
+struct RawLin {
+  std::vector<uint8_t> wb;
+  mlx::core::array w{0.f};  // original dtype [rows, cols]
+  int64_t rows = 0, cols = 0;
+  bool load(const ModelFiles &files, const std::string &name, std::string &err);
+  std::optional<mlx::core::array> biases() const { return std::nullopt; }
+  mlx::core::array apply(const mlx::core::array &x) const;  // x @ wᵀ
+};
+
+mlx::core::array e0_silu(const mlx::core::array &x);
+// Qwen3Next._precise_swiglu: gate/up to f32, silu(gate)*up, cast back to out dtype.
+mlx::core::array e0_swiglu_precise(mlx::core::Dtype out_dt, const mlx::core::array &gate,
+                                   const mlx::core::array &up);
+// Empty weight → unscaled rms_norm (GDN q/k).
+mlx::core::array e0_rms_norm(const mlx::core::array &x,
+                             const std::optional<mlx::core::array> &weight, double eps);
+
+bool to_host32(const mlx::core::array &a, std::vector<float> &out, std::string &err);
+// Host→graph: iterator ctor copies into mlx-owned buffers. no-op-deleter wrap is
+// only safe for member weights that outlive the graph (QuantLin::load); wrapping a
+// local vector that a lazy graph still holds after eval is a dangling read.
+mlx::core::array wrap_f32(const std::vector<float> &v, mlx::core::Shape shape);
+
+inline float silu_f(float x) { return x * (1.0f / (1.0f + std::exp(-x))); }
+
+struct ExecCounters {
+  std::atomic<uint64_t> qmm_calls{0};
+  std::atomic<uint64_t> qmm_experts{0};
+  std::atomic<uint64_t> qmm_weight_bytes{0};
+  std::atomic<uint64_t> host_syncs{0};
+  std::atomic<uint64_t> host_wraps{0};
+  std::atomic<uint64_t> leases{0};
+  struct Snap {
+    uint64_t qmm_calls, qmm_experts, qmm_weight_bytes, host_syncs, host_wraps, leases;
+  };
+  Snap snapshot() const {
+    return {qmm_calls.load(std::memory_order_acquire), qmm_experts.load(std::memory_order_acquire),
+            qmm_weight_bytes.load(std::memory_order_acquire),
+            host_syncs.load(std::memory_order_acquire),
+            host_wraps.load(std::memory_order_acquire),
+            leases.load(std::memory_order_acquire)};
+  }
+  static Snap diff(const Snap &a, const Snap &b) {  // b − a
+    return {b.qmm_calls - a.qmm_calls,       b.qmm_experts - a.qmm_experts,
+            b.qmm_weight_bytes - a.qmm_weight_bytes, b.host_syncs - a.host_syncs,
+            b.host_wraps - a.host_wraps,     b.leases - a.leases};
+  }
+  void reset() {
+    qmm_calls = 0; qmm_experts = 0; qmm_weight_bytes = 0;
+    host_syncs = 0; host_wraps = 0; leases = 0;
+  }
+};
+ExecCounters &exec_counters();
+
+struct TimingCounters {
+  std::atomic<uint64_t> eval_ns{0};
+  std::atomic<uint64_t> host_sync_ns{0};
+  std::atomic<uint64_t> gate_cpu_ns{0};
+  std::atomic<uint64_t> staging_ns{0};
+  struct Snap { uint64_t eval_ns, host_sync_ns, gate_cpu_ns, staging_ns; };
+  Snap snapshot() const {
+    return {eval_ns.load(std::memory_order_relaxed), host_sync_ns.load(std::memory_order_relaxed),
+            gate_cpu_ns.load(std::memory_order_relaxed), staging_ns.load(std::memory_order_relaxed)};
+  }
+  static bool enabled();  // EDGE0_TIMING=1 (once per process)
+};
+TimingCounters &timing_counters();
+struct TimeSink {
+  std::atomic<uint64_t> &sink;
+  uint64_t t0;
+  explicit TimeSink(std::atomic<uint64_t> &s)
+      : sink(s), t0(TimingCounters::enabled() ? monotonic_ns() : 0) {}
+  ~TimeSink() { stop(); }
+  void stop() {
+    if (t0) {
+      sink.fetch_add(monotonic_ns() - t0, std::memory_order_relaxed);
+      t0 = 0;
+    }
+  }
+  static uint64_t monotonic_ns();
+};
+#define E0_TIME_CAT2(a, b) a##b
+#define E0_TIME_CAT(a, b) E0_TIME_CAT2(a, b)
+#define E0_TIME(field) ::e0n::TimeSink E0_TIME_CAT(_e0_ts_, __LINE__)(field)
+
+enum class ExecMode { Legacy, Batched };
+ExecMode exec_mode();  // first call reads env (default batched; unknown → warn + legacy)
+
+// EDGE0_PREFILL=ondemand (default) vs production. production is an env-only hybrid
+// (hot fast path is not built); it is not a faithful production-shape replica.
+bool prefill_ondemand();
+
+struct Entry;
+struct ExpertPtrs {
+  const uint32_t *gw = nullptr, *uw = nullptr, *dw = nullptr;
+  const uint16_t *gs = nullptr, *gb = nullptr, *us = nullptr, *ub = nullptr, *ds = nullptr,
+                 *db = nullptr;
+};
+bool expert_ptrs(const Entry &e, ExpertPtrs &p, std::string &err);
+
+bool embed_lookup(const QuantLin &emb, const std::vector<int32_t> &ids, int64_t H, int64_t V,
+                  std::vector<float> &out, std::string &err);
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/genstream.hpp b/macos/crates/engine-native/src/genstream.hpp
new file mode 100644
index 0000000..337f1fa
--- /dev/null
+++ b/macos/crates/engine-native/src/genstream.hpp
@@ -0,0 +1,252 @@
+// Generation stream: ABI wrapper around net prefill/step plus sampling.
+#pragma once
+#include <cstdint>
+#include <mutex>
+#include <functional>
+#include <string>
+#include <vector>
+
+#include <algorithm>
+
+#include "net35.hpp"
+#include "prerouter.hpp"
+#include "sampling.hpp"
+#include "provider.hpp"
+
+struct E0Engine;
+
+namespace e0n { namespace gen {
+
+class StreamBase {
+ public:
+  StreamBase(const StreamBase &) = delete;
+  StreamBase &operator=(const StreamBase &) = delete;
+  virtual ~StreamBase() = default;
+  std::mutex mu;  // stream-level; the engine handle may die first so the lock is not on E0Engine
+  bool next_token(int32_t &tok_out) {
+    if (finished_) return false;
+    if (cancelled_) { finished_ = true; return false; }
+    if (generated_ >= p_.max_new_tokens) { finish_natural(); return false; }
+    std::string err;
+    std::vector<float> lg;
+    std::vector<RouteRec> recs;
+    if (exhausted()) { finished_ = true; return false; }
+    if (transport_) {
+      if (generated_ >= p_.teacher_tokens.size() || generated_ >= p_.max_new_tokens) {
+        finish_natural();
+        return false;
+      }
+      tok_out = p_.teacher_tokens[generated_++];
+      hist_.push_back(tok_out);
+      emitted_ = true;
+      return true;
+    }
+    if (!started_) {
+      if (prefill_prompt_.empty() && cached_logits_set_) {
+        lg = cached_logits_;
+      } else {
+        if (!prefill(prefill_prompt_, lg, recs, err)) { note_err(err); return false; }
+        if (on_prefill_) on_prefill_(prompt_, lg);
+      }
+      last_logits_ = lg;
+      last_recs_ = recs;
+      cur_ = p_.teacher_tokens.empty() ? host_argmax(lg) : p_.teacher_tokens[0];
+      started_ = true;
+    } else {
+      if (warm_) {
+        Step plan = live_next_step(plan_state_());
+        if (!plan.experts.empty()) warm_(plan);
+      }
+      if (!step(cur_, lg, recs, err)) { note_err(err); return false; }
+      last_logits_ = lg;
+      last_recs_ = recs;
+    }
+    int32_t t;
+    if (!p_.teacher_tokens.empty()) {
+      if (generated_ >= p_.teacher_tokens.size()) { finish_natural(); return false; }
+      t = p_.teacher_tokens[generated_];
+    } else {
+      t = sampling::sample_step(lg, p_, !emitted_, hist_, rng_);
+    }
+    emitted_ = true;
+    // EOS / stop-id: terminate without emitting that token.
+    for (int32_t e : p_.eos_ids)
+      if (t == e) { finish_natural(); return false; }
+    for (int32_t e : p_.stop_token_ids)
+      if (t == e) { finish_natural(); return false; }
+    hist_.push_back(t);
+    // String stop: subsequence ending at this token, wholly in the generated region,
+    // rolls back the matched span. hist_ = prompt + emitted (incl. this t);
+    // generated-region start gp = hist_.size() - generated_ - 1.
+    {
+      const size_t gp = hist_.size() - generated_ - 1;
+      for (const auto &seq : p_.stop_seqs) {
+        if (seq.empty() || seq.size() > generated_ + 1) continue;
+        const size_t base = hist_.size() - seq.size();
+        if (base < gp) continue;
+        bool hit = true;
+        for (size_t i = 0; i < seq.size(); i++)
+          if (hist_[base + i] != seq[i]) { hit = false; break; }
+        if (hit) {
+          generated_ -= (seq.size() - 1);
+          hist_.resize(base);
+          finish_natural();
+          return false;
+        }
+      }
+    }
+    generated_++;
+    cur_ = t;
+    tok_out = t;
+    return true;
+  }
+  void cancel() { cancelled_ = true; }
+  void set_transport() { transport_ = true; }
+  void set_empty_done() { finished_ = true; }
+  void set_prefill_prompt(std::vector<int32_t> ids) { prefill_prompt_ = std::move(ids); }
+  void set_on_prefill(std::function<void(const std::vector<int32_t> &)> cb) {
+    on_prefill_ = [cb = std::move(cb)](const std::vector<int32_t> &ids,
+                                       const std::vector<float> &) { cb(ids); };
+  }
+  void set_on_prefill(std::function<void(const std::vector<int32_t> &,
+                                         const std::vector<float> &)> cb) {
+    on_prefill_ = std::move(cb);
+  }
+  void set_cached_logits(std::vector<float> logits) {
+    cached_logits_ = std::move(logits);
+    cached_logits_set_ = true;
+  }
+  void set_on_complete(std::function<void(const std::vector<int32_t> &,
+                                          const std::vector<float> &)> cb) {
+    on_complete_ = std::move(cb);
+  }
+  void set_on_complete(std::function<void(const std::vector<int32_t> &)> cb) {
+    on_complete_ = [cb = std::move(cb)](const std::vector<int32_t> &ids,
+                                        const std::vector<float> &) { cb(ids); };
+  }
+  void set_owner(E0Engine *h) { owner_ = h; }
+  bool finished() const { return finished_; }
+  bool errored() const { return err_; }
+  const std::string &err_msg() const { return err_msg_; }
+  void seal() { finished_ = true; }
+  const std::vector<RouteRec> &last_recs() const { return last_recs_; }
+  bool owner_of(E0Engine *h) const { return owner_ == h; }
+  E0Engine *owner() { return owner_; }
+  bool started() const { return started_; }
+  size_t generated() const { return generated_; }
+  bool terminal() const { return finished_ && !err_; }
+  const std::vector<int32_t> &history() const { return hist_; }
+  StreamBase(std::vector<int32_t> prompt, sampling::Params p)
+      : prompt_(std::move(prompt)), prefill_prompt_(prompt_), p_(std::move(p)), hist_(prompt_) {}
+
+ protected:
+  virtual bool prefill(const std::vector<int32_t> &, std::vector<float> &,
+                       std::vector<RouteRec> &, std::string &err) {
+    err = "stream without forward context";
+    return false;
+  }
+  virtual bool step(int32_t, std::vector<float> &, std::vector<RouteRec> &,
+                    std::string &err) {
+    err = "stream without forward context";
+    return false;
+  }
+  virtual const PrerouterState &plan_state_() const { static PrerouterState k; return k; }
+  virtual bool exhausted() const { return false; }
+  void note_err(std::string m) { err_ = true; err_msg_ = std::move(m); finished_ = true; }
+  void finish_natural() {
+    if (!finished_) {
+      finished_ = true;
+      if (!completion_called_ && on_complete_) {
+        completion_called_ = true;
+        on_complete_(hist_, last_logits_);
+      }
+    }
+  }
+  static int32_t host_argmax(const std::vector<float> &v) {
+    size_t b = 0;
+    for (size_t i = 1; i < v.size(); i++)
+      if (v[i] > v[b]) b = i;  // ties take the smaller id (strict >)
+    return (int32_t)b;
+  }
+
+  sampling::Params p_;
+  std::vector<int32_t> prompt_;
+  std::vector<int32_t> prefill_prompt_;
+  std::vector<int32_t> hist_;
+  size_t generated_ = 0;
+  bool emitted_ = false;
+  bool cancelled_ = false;
+  bool finished_ = false;
+  bool err_ = false;
+  std::string err_msg_;
+  sampling::Rng rng_{};
+  std::function<void(const Step &)> warm_;
+  std::function<void(const std::vector<int32_t> &, const std::vector<float> &)> on_prefill_;
+  std::function<void(const std::vector<int32_t> &, const std::vector<float> &)> on_complete_;
+  std::vector<float> cached_logits_;
+  std::vector<float> last_logits_;
+  bool cached_logits_set_ = false;
+  bool completion_called_ = false;
+  E0Engine *owner_ = nullptr;  // destroy scans the pool and drops streams of a dead engine
+  std::vector<RouteRec> last_recs_;
+
+ private:
+  bool started_ = false;
+  bool transport_ = false;
+  int32_t cur_ = 0;
+};
+
+template <class NetT, class StateT>
+class Stream : public StreamBase {
+ public:
+  Stream(NetT *net, StateT *st, std::function<void(const Step &)> warm, E0Engine *owner,
+         std::vector<int32_t> prompt, sampling::Params p)
+      : StreamBase(std::move(prompt), std::move(p)), net_(net), st_(st) {
+    this->warm_ = std::move(warm);
+    this->owner_ = owner;
+  }
+
+ protected:
+  bool prefill(const std::vector<int32_t> &ids, std::vector<float> &lg,
+               std::vector<RouteRec> &recs, std::string &err) override {
+    return net_->prefill(ids, *st_, lg, recs, err);
+  }
+  bool step(int32_t tok, std::vector<float> &lg, std::vector<RouteRec> &recs,
+            std::string &err) override {
+    return net_->step(tok, *st_, lg, recs, err);
+  }
+  const PrerouterState &plan_state_() const override { return st_->pr; }
+
+ private:
+  NetT *net_;
+  StateT *st_;
+};
+
+class ScriptedStream : public StreamBase {
+ public:
+  ScriptedStream(std::vector<int32_t> prompt, sampling::Params p,
+                 std::vector<std::vector<float>> script)
+      : StreamBase(std::move(prompt), std::move(p)), script_(std::move(script)) {}
+
+ protected:
+  bool prefill(const std::vector<int32_t> &, std::vector<float> &lg,
+               std::vector<RouteRec> &recs, std::string &) override {
+    recs.clear();
+    lg = script_[0];
+    return true;
+  }
+  bool step(int32_t, std::vector<float> &lg, std::vector<RouteRec> &recs,
+            std::string &) override {
+    recs.clear();
+    idx_++;
+    lg = script_[idx_];
+    return true;
+  }
+  bool exhausted() const override { return idx_ + 1 >= script_.size(); }
+
+ private:
+  std::vector<std::vector<float>> script_;
+  size_t idx_ = 0;
+};
+
+}}  // namespace e0n::gen
diff --git a/macos/crates/engine-native/src/hotstack.cpp b/macos/crates/engine-native/src/hotstack.cpp
new file mode 100644
index 0000000..050ede1
--- /dev/null
+++ b/macos/crates/engine-native/src/hotstack.cpp
@@ -0,0 +1,131 @@
+#include "hotstack.hpp"
+
+#include <fcntl.h>
+#include <unistd.h>
+
+#include <algorithm>
+
+namespace e0n {
+
+HotStack::HotStack(std::shared_ptr<const ModelFiles> files,
+                   std::shared_ptr<const ExpertMap> map, uint64_t budget_bytes)
+    : files_(std::move(files)), map_(std::move(map)), budget_(budget_bytes) {}
+
+HotStack::~HotStack() {
+  std::lock_guard<std::mutex> g(mu_);
+  for (auto &[s, fd] : fds_) ::close(fd);
+}
+
+bool HotStack::load_entry(int layer, int64_t expert, EntryPtr &out, std::string &err) {
+  auto rs = map_->regions(layer, expert);
+  if (rs.size() != 9) {
+    err = "region count != 9 (L" + std::to_string(layer) + " E" + std::to_string(expert) + ")";
+    return false;
+  }
+  auto e = std::make_shared<Entry>();
+  e->layer = layer;
+  e->expert = expert;
+  e->regions = rs;
+  uint64_t total = 0;
+  for (auto &r : rs) total += r.len;
+  e->bytes = total;
+  e->blob.resize(total);
+  e->blob_offsets.resize(rs.size());
+  uint64_t done_logical = 0;
+  for (size_t i = 0; i < rs.size(); i++) {
+    const Region &r = rs[i];
+    e->blob_offsets[i] = done_logical;
+    int fd;
+    {
+      // fds_ is used outside the caller's lock (this runs in acquire's unlocked
+      // section); opening a new fd still needs a short lock on fds_.
+      std::lock_guard<std::mutex> g(mu_);
+      auto it = fds_.find(r.shard);
+      if (it == fds_.end()) {
+        fd = ::open(files_->shard_path(r.shard).c_str(), O_RDONLY);
+        if (fd < 0) {
+          err = "open shard failed: " + r.shard;
+          return false;
+        }
+        // Load fds are *not* F_NOCACHE: an evicted reload should reuse the OS page cache.
+        fds_[r.shard] = fd;
+      } else {
+        fd = it->second;
+      }
+    }
+    uint8_t *dst = e->blob.data() + done_logical;
+    uint64_t got_total = 0;
+    while (got_total < r.len) {
+      uint64_t want = std::min<uint64_t>(r.len - got_total, 1u << 20);
+      ssize_t got = ::pread(fd, dst + got_total, want, r.off + got_total);
+      if (got <= 0) {
+        err = "pread failed " + r.shard;
+        return false;
+      }
+      got_total += static_cast<uint64_t>(got);
+    }
+    done_logical += r.len;
+  }
+  out = std::move(e);
+  return true;
+}
+
+HotStack::Lease HotStack::acquire(int layer, int64_t expert, std::string &err) {
+  const uint64_t k = key_of(layer, expert);
+  // 1) hit (fast path).
+  {
+    std::lock_guard<std::mutex> g(mu_);
+    auto it = items_.find(k);
+    if (it != items_.end()) {
+      lru_.splice(lru_.begin(), lru_, it->second.lru);  // MRU
+      metrics_.hits++;
+      return it->second.entry;
+    }
+    metrics_.misses++;
+  }
+  // 2) disk read outside the lock.
+  EntryPtr entry;
+  if (!load_entry(layer, expert, entry, err)) return nullptr;
+  // 3) insert + evict (double-check: a concurrent load of the same key discards ours).
+  {
+    std::lock_guard<std::mutex> g(mu_);
+    if (auto it = items_.find(k); it != items_.end()) {
+      lru_.splice(lru_.begin(), lru_, it->second.lru);
+      return it->second.entry;  // drop this thread's duplicate load (bytes not counted)
+    }
+    lru_.push_front(k);
+    metrics_.heap_bytes += entry->bytes;
+    metrics_.heap_peak = std::max(metrics_.heap_peak, metrics_.heap_bytes);
+    metrics_.loaded++;
+    items_.emplace(k, Item{entry, lru_.begin()});
+    Lease ret = entry;  // hold a lease first so eviction cannot drop the entry we just inserted
+    // LRU eviction from the tail; skip entries still referenced by a lease (use_count>1).
+    if (metrics_.heap_bytes > budget_) {
+      auto it = lru_.end();
+      while (metrics_.heap_bytes > budget_ && it != lru_.begin()) {
+        --it;
+        auto iit = items_.find(*it);
+        if (iit == items_.end()) { it = lru_.begin(); continue; }
+        if (iit->second.entry.use_count() > 1) continue;  // in use: skip this round
+        metrics_.heap_bytes -= iit->second.entry->bytes;
+        metrics_.evictions++;
+        lru_.erase(iit->second.lru);
+        items_.erase(iit);
+        it = lru_.end();  // iterator invalid after erase; restart from the tail
+      }
+    }
+    return ret;
+  }
+}
+
+CacheMetrics HotStack::snapshot() const {
+  std::lock_guard<std::mutex> g(mu_);
+  return metrics_;
+}
+
+size_t HotStack::size() const {
+  std::lock_guard<std::mutex> g(mu_);
+  return items_.size();
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/hotstack.hpp b/macos/crates/engine-native/src/hotstack.hpp
new file mode 100644
index 0000000..cae4548
--- /dev/null
+++ b/macos/crates/engine-native/src/hotstack.hpp
@@ -0,0 +1,72 @@
+// Cross-layer LRU hot stack: explicit heap buffers + F_NOCACHE pread (eviction is testable).
+#pragma once
+#include <cstdint>
+#include <list>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "expertmap.hpp"
+#include "safio.hpp"
+
+namespace e0n {
+
+struct Entry {
+  std::vector<uint8_t> blob;
+  std::vector<Region> regions;
+  std::vector<uint64_t> blob_offsets;
+  uint64_t bytes = 0;
+  int layer = 0;
+  int64_t expert = -1;
+};
+
+struct CacheMetrics {
+  uint64_t loaded = 0;
+  uint64_t hits = 0;
+  uint64_t misses = 0;
+  uint64_t heap_bytes = 0;
+  uint64_t heap_peak = 0;
+  uint64_t evictions = 0;
+};
+
+class HotStack {
+ public:
+  using EntryPtr = std::shared_ptr<const Entry>;
+  using Lease = EntryPtr;
+
+  HotStack(std::shared_ptr<const ModelFiles> files, std::shared_ptr<const ExpertMap> map,
+           uint64_t budget_bytes);
+  ~HotStack();
+
+  // Miss → F_NOCACHE cold read on this thread. Failure returns nullptr + err.
+  // A live lease keeps the entry allocated (eviction only drops the stack ref).
+  Lease acquire(int layer, int64_t expert, std::string &err);
+
+  CacheMetrics snapshot() const;
+  uint64_t budget() const { return budget_; }
+  size_t size() const;
+
+ private:
+  struct Item {
+    EntryPtr entry;
+    std::list<uint64_t>::iterator lru;
+  };
+  static uint64_t key_of(int layer, int64_t expert) {
+    return (static_cast<uint64_t>(layer) << 32) | static_cast<uint64_t>(expert);
+  }
+  bool load_entry(int layer, int64_t expert, EntryPtr &out, std::string &err);
+
+  std::shared_ptr<const ModelFiles> files_;
+  std::shared_ptr<const ExpertMap> map_;
+  uint64_t budget_;
+  mutable std::mutex mu_;
+  std::unordered_map<uint64_t, Item> items_;
+  std::list<uint64_t> lru_;  // front = MRU
+  CacheMetrics metrics_;
+  std::map<std::string, int> fds_;  // shard → F_NOCACHE fd
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/lora.cpp b/macos/crates/engine-native/src/lora.cpp
new file mode 100644
index 0000000..ebc91a4
--- /dev/null
+++ b/macos/crates/engine-native/src/lora.cpp
@@ -0,0 +1,410 @@
+#include "lora.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <map>
+#include <set>
+
+#include <nlohmann/json.hpp>
+
+namespace e0n {
+namespace lora {
+namespace mx = mlx::core;
+using json = nlohmann::json;
+
+namespace {
+
+void jsonl_capture(const std::string &path, const std::string &key, int layer,
+                   const mx::array &base, const mx::array &delta) {
+  try {
+    mx::array be = mx::astype(base, mx::float32), de = mx::astype(delta, mx::float32);
+    mx::eval(be, de);
+    const size_t n = be.size();
+    if (de.size() != n) return;
+    auto side = [](const float *p, size_t n) {
+      double ss = 0;
+      for (size_t i = 0; i < n; i++) ss += double(p[i]) * p[i];
+      json j;
+      json pref = json::array();
+      for (size_t i = 0; i < 64 && i < n; i++) pref.push_back(p[i]);
+      j["prefix"] = pref;
+      j["summary"] = {{"l2", std::sqrt(ss)}};
+      return j;
+    };
+    json rec;
+    rec["key"] = key;
+    rec["layer"] = layer;
+    rec["step"] = -1;
+    rec["n"] = static_cast<int64_t>(n);
+    rec["base"] = side(be.data<float>(), n);
+    rec["delta"] = side(de.data<float>(), n);
+    std::ofstream f(path, std::ios::app);
+    f << rec.dump() << "\n";
+  } catch (const std::exception &) {
+  }
+}
+
+bool parse_meta_field(const json &m, const char *key, int64_t &out) {
+  if (!m.contains(key)) return false;
+  const json &v = m[key];
+  if (v.is_number_integer()) { out = v.get<int64_t>(); return true; }
+  if (v.is_string()) {
+    try { out = std::stoll(v.get<std::string>()); return true; } catch (...) { return false; }
+  }
+  return false;
+}
+
+int slot_of_35(int layer, const std::string &tail, std::string &err) {
+  const bool full = (layer + 1) % 4 == 0;
+  auto need = [&](bool cond) {
+    if (cond) return true;
+    err = "lora layer family mismatch: layer=" + std::to_string(layer) + " tail=" + tail;
+    return false;
+  };
+  if (tail == "linear_attn.in_proj_a") return need(!full) ? S35_IN_A : -1;
+  if (tail == "linear_attn.in_proj_b") return need(!full) ? S35_IN_B : -1;
+  if (tail == "linear_attn.in_proj_qkv") return need(!full) ? S35_IN_QKV : -1;
+  if (tail == "linear_attn.in_proj_z") return need(!full) ? S35_IN_Z : -1;
+  if (tail == "linear_attn.out_proj") return need(!full) ? S35_OUT_PROJ : -1;
+  if (tail == "self_attn.q_proj") return need(full) ? S35_Q : -1;
+  if (tail == "self_attn.k_proj") return need(full) ? S35_K : -1;
+  if (tail == "self_attn.v_proj") return need(full) ? S35_V : -1;
+  if (tail == "self_attn.o_proj") return need(full) ? S35_O : -1;
+  if (tail == "mlp.shared_expert.gate_proj") return S35_SH_GATE;
+  if (tail == "mlp.shared_expert.up_proj") return S35_SH_UP;
+  if (tail == "mlp.shared_expert.down_proj") return S35_SH_DOWN;
+  err = "lora unknown target: " + tail;
+  return -1;
+}
+
+int slot_of_8b(int layer, const std::string &tail, std::string &err) {
+  const bool mla = (layer + 1) % 4 == 0;
+  auto need = [&](bool cond) {
+    if (cond) return true;
+    err = "lora layer family mismatch: layer=" + std::to_string(layer) + " tail=" + tail;
+    return false;
+  };
+  if (tail == "attention.q_proj") return need(!mla) ? S8B_Q : -1;
+  if (tail == "attention.k_proj") return need(!mla) ? S8B_K : -1;
+  if (tail == "attention.v_proj") return need(!mla) ? S8B_V : -1;
+  if (tail == "attention.o_proj") return need(!mla) ? S8B_O : -1;
+  if (tail == "attention.b_proj") return need(!mla) ? S8B_B : -1;
+  if (tail == "attention.f_proj") return need(!mla) ? S8B_F : -1;
+  if (tail == "attention.g_proj") return mla ? S8B_M_G : S8B_G;
+  if (tail == "attention.q_b_proj") return need(mla) ? S8B_M_QB : -1;
+  if (tail == "attention.kv_b_proj") return need(mla) ? S8B_M_KVB : -1;
+  if (tail == "attention.dense") return need(mla) ? S8B_M_DENSE : -1;
+  if (tail == "mlp.gate_proj") return need(layer == 0) ? S8B_D_GATE : -1;
+  if (tail == "mlp.up_proj") return need(layer == 0) ? S8B_D_UP : -1;
+  if (tail == "mlp.down_proj") return need(layer == 0) ? S8B_D_DOWN : -1;
+  err = "lora unknown target: " + tail;
+  return -1;
+}
+
+}  // namespace
+
+bool Bundle::load_raw(std::map<std::string, std::pair<std::string, std::vector<uint8_t>>> tensors,
+                      const std::string &meta_json, bool is_8b, int num_layers,
+                      std::string &err) {
+  try {
+    json m = json::parse(meta_json);
+    int64_t r = 0, a = 0, k = 0;
+    if (!parse_meta_field(m, "r", r) || !parse_meta_field(m, "alpha", a) ||
+        !parse_meta_field(m, "K", k) || r <= 0 || a <= 0) {
+      err = "lora metadata missing r/alpha/K or not positive";
+      return false;
+    }
+    meta_.r = r;
+    meta_.alpha = a;
+    meta_.k = k;
+    meta_.scale = static_cast<float>(a) / static_cast<float>(r);
+    meta_.ok = true;
+  } catch (const std::exception &ex) {
+    err = std::string("lora metadata parse: ") + ex.what();
+    return false;
+  }
+
+  struct Half {
+    std::pair<std::string, std::vector<uint8_t>> A, B;
+  };
+  std::map<std::string, Half> halves;
+  for (auto &kv : tensors) {
+    const std::string &key = kv.first;
+    const size_t sfx = 7;  // ".lora_A"
+    bool is_a = key.size() > sfx && key.compare(key.size() - sfx, sfx, ".lora_A") == 0;
+    bool is_b = key.size() > sfx && key.compare(key.size() - sfx, sfx, ".lora_B") == 0;
+    if (!is_a && !is_b) {
+      err = "lora unexpected key (not .lora_A/.lora_B): " + key;
+      return false;
+    }
+    std::string base = key.substr(0, key.size() - sfx);
+    if (is_a) halves[base].A = kv.second;
+    else halves[base].B = kv.second;
+  }
+
+  const std::string pfx35 = "language_model.model.layers.";
+  const std::string pfx8 = "model.layers.";
+  ns_ = is_8b ? static_cast<int>(S8B_NSLOT) : static_cast<int>(S35_NSLOT);
+  slots_.assign(num_layers, std::vector<int>(ns_, -1));
+
+  for (auto &h : halves) {
+    const std::string &base = h.first;
+    const bool has_a = !h.second.A.second.empty();
+    const bool has_b = !h.second.B.second.empty();
+    if (!has_a || !has_b) {
+      err = "lora pairing incomplete: " + base;
+      return false;
+    }
+    const size_t n_a = h.second.A.second.size(), n_b = h.second.B.second.size();
+    std::string dta = h.second.A.first, dtb = h.second.B.first;
+    auto tag = [](std::string &d) {
+      auto pos = d.find('|');
+      std::string shape;
+      if (pos != std::string::npos) {
+        shape = d.substr(pos + 1);
+        d = d.substr(0, pos);
+      }
+      return shape;
+    };
+    std::string sha = tag(dta), shb = tag(dtb);
+    if (dta != "F16" || dtb != "F16") {
+      err = "lora expected F16 domain: " + base;
+      return false;
+    }
+    auto parse_dims = [](const std::string &s, std::vector<int64_t> &out) {
+      // "r,in" / "out,r"
+      size_t c = s.find(',');
+      if (c == std::string::npos) return false;
+      try {
+        out.push_back(std::stoll(s.substr(0, c)));
+        out.push_back(std::stoll(s.substr(c + 1)));
+      } catch (...) { return false; }
+      return true;
+    };
+    std::vector<int64_t> da, db;
+    if (!parse_dims(sha, da) || !parse_dims(shb, db) || da.size() != 2 || db.size() != 2) {
+      err = "lora shape markers missing: " + base;
+      return false;
+    }
+    const int64_t r = da[0], in = da[1], out = db[0];
+    if (db[1] != r || r != meta_.r || in <= 0 || out <= 0) {
+      err = "lora shape mismatch (r/in/out): " + base + " A=[" + std::to_string(r) + "," +
+            std::to_string(in) + "] B=[" + std::to_string(out) + "," + std::to_string(db[1]) + "]";
+      return false;
+    }
+    if (n_a != static_cast<size_t>(2 * r * in) || n_b != static_cast<size_t>(2 * out * r)) {
+      err = "lora byte size does not match shape: " + base;
+      return false;
+    }
+    std::string prefix = is_8b ? pfx8 : pfx35;
+    if (base.compare(0, prefix.size(), prefix) != 0) {
+      std::string other = is_8b ? pfx35 : pfx8;
+      if (base.compare(0, other.size(), other) == 0) {
+        err = "lora tier prefix mismatch: " + base;
+        return false;
+      }
+      err = "lora unknown prefix: " + base;
+      return false;
+    }
+    std::string rest = base.substr(prefix.size());
+    size_t dot = rest.find('.');
+    int layer = -1;
+    if (dot == std::string::npos) {
+      err = "lora key shape invalid: " + base;
+      return false;
+    }
+    try {
+      layer = std::stoi(rest.substr(0, dot));
+    } catch (...) {
+      err = "lora layer is not a number: " + base;
+      return false;
+    }
+    if (layer < 0 || layer >= num_layers) {
+      err = "lora layer out of range: " + base;
+      return false;
+    }
+    std::string tail = rest.substr(dot + 1);
+    std::string serr;
+    int slot = is_8b ? slot_of_8b(layer, tail, serr) : slot_of_35(layer, tail, serr);
+    if (slot < 0) {
+      err = serr.empty() ? "lora target parse failed: " + base : serr;
+      return false;
+    }
+    if (slots_[layer][slot] >= 0) {
+      err = "lora slot already bound: " + base;
+      return false;
+    }
+    auto p = std::make_unique<Pair>();
+    p->layer = layer;
+    p->slot = slot;
+    p->r = r;
+    p->in = in;
+    p->out = out;
+    p->ab = std::move(h.second.A.second);
+    p->bb = std::move(h.second.B.second);
+    p->key = base;
+    slots_[layer][slot] = static_cast<int>(pairs_.size());
+    bytes_ += static_cast<int64_t>(p->ab.size() + p->bb.size());
+    pairs_.push_back(std::move(p));
+  }
+
+  for (auto &p : pairs_) {
+    p->A = mx::array(p->ab.data(), mx::Shape{static_cast<int32_t>(p->r), static_cast<int32_t>(p->in)},
+                     mx::float16, [](void *) {});
+    p->B = mx::array(p->bb.data(), mx::Shape{static_cast<int32_t>(p->out), static_cast<int32_t>(p->r)},
+                     mx::float16, [](void *) {});
+  }
+  for (auto &p : pairs_) keys_sorted_.push_back(p->key);
+  std::sort(keys_sorted_.begin(), keys_sorted_.end());
+  hit_.assign(pairs_.size(), 0);
+  return true;
+}
+
+bool Bundle::load_file(const std::string &path, bool is_8b, int num_layers, std::string &err) {
+  std::ifstream f(path, std::ios::binary);
+  if (!f) { err = "lora file unreadable: " + path; return false; }
+  uint64_t hl = 0;
+  f.read(reinterpret_cast<char *>(&hl), 8);
+  if (!f || hl > (64u << 20)) { err = "lora header length invalid: " + path; return false; }
+  std::vector<char> hs(hl, '\0');
+  f.read(hs.data(), static_cast<std::streamsize>(hl));
+  if (!f) { err = "lora header truncated: " + path; return false; }
+  json hdr;
+  try {
+    hdr = json::parse(std::string(hs.begin(), hs.end()));
+  } catch (const std::exception &ex) {
+    err = std::string("lora header parse: ") + ex.what();
+    return false;
+  }
+  std::string meta = "{}";
+  if (hdr.contains("__metadata__")) {
+    const json &mv = hdr["__metadata__"];
+    if (mv.is_string()) meta = mv.get<std::string>();
+    else if (mv.is_object()) {
+      if (mv.size() == 1 && mv.contains("__metadata__") && mv["__metadata__"].is_string())
+        meta = mv["__metadata__"].get<std::string>();
+      else
+        meta = mv.dump();
+    }
+  }
+  std::map<std::string, std::pair<std::string, std::vector<uint8_t>>> tensors;
+  const uint64_t base_off = 8 + hl;
+  for (auto &el : hdr.items()) {
+    if (el.key() == "__metadata__") continue;
+    const json &v = el.value();
+    if (!v.contains("dtype") || !v.contains("data_offsets")) continue;
+    std::string dt = v["dtype"].get<std::string>();
+    std::vector<int64_t> shape;
+    for (auto &d : v["shape"]) shape.push_back(d.get<int64_t>());
+    std::string tag = dt;
+    if (shape.size() == 2) tag += "|" + std::to_string(shape[0]) + "," + std::to_string(shape[1]);
+    uint64_t b = v["data_offsets"][0].get<uint64_t>() + base_off;
+    uint64_t e = v["data_offsets"][1].get<uint64_t>() + base_off;
+    std::vector<uint8_t> buf(e - b);
+    f.seekg(static_cast<std::streamoff>(b));
+    f.read(reinterpret_cast<char *>(buf.data()), static_cast<std::streamsize>(buf.size()));
+    if (!f) { err = "lora tensor truncated: " + el.key(); return false; }
+    tensors[el.key()] = {tag, std::move(buf)};
+  }
+  if (!load_raw(std::move(tensors), meta, is_8b, num_layers, err)) return false;
+  if (const char *sc = std::getenv("E0_LORA_SCAFFOLD")) cap_path_ = sc;
+  return true;
+}
+
+bool Bundle::has_pair(int layer, int slot) const { return find_pair(layer, slot) != nullptr; }
+
+int Bundle::n_hits() const {
+  int n = 0;
+  for (char c : hit_) n += c != 0;
+  return n;
+}
+
+bool Bundle::pair_raw(int layer, int slot, std::vector<uint8_t> &a, std::vector<uint8_t> &b,
+                      int64_t &r, int64_t &in, int64_t &out) const {
+  const Pair *p = find_pair(layer, slot);
+  if (!p) return false;
+  a = p->ab;
+  b = p->bb;
+  r = p->r;
+  in = p->in;
+  out = p->out;
+  return true;
+}
+
+mlx::core::array Bundle::dressed(const QuantLin &lin, const mlx::core::array &x, int layer,
+                                  int slot) const {
+  mx::array y = lin.apply(x);
+  const Pair *pp = find_pair(layer, slot);
+  if (!pp) return y;
+  const Pair &p = *pp;
+  hit_[static_cast<size_t>(slots_[layer][slot])] = 1;
+  if (p.in != lin.cols || p.out != lin.rows) {
+    throw std::runtime_error("lora wiring rank mismatch: " + p.key + " file=[in=" + std::to_string(p.in) +
+                             ",out=" + std::to_string(p.out) + "] lin=[" + std::to_string(lin.cols) +
+                             "," + std::to_string(lin.rows) + "]");
+  }
+  mx::array xf = mx::astype(x, mx::float16);
+  mx::array d = mx::matmul(mx::matmul(xf, mx::transpose(p.A)), mx::transpose(p.B));
+  d = mx::astype(meta_.scale * d, y.dtype());
+  mx::array out = y + d;
+  if (!cap_path_.empty() && cap_done_.insert(p.key).second)
+    jsonl_capture(cap_path_, p.key, p.layer, y, d);
+  return out;
+}
+
+bool Bundle::dressed_variant(const QuantLin &lin, const mlx::core::array &x, int layer, int slot,
+                             int which, mx::array &out, std::string &err) const {
+  const Pair *p = find_pair(layer, slot);
+  if (!p) { err = "dressed_variant: no pairing"; return false; }
+  try {
+    if (which == 1) {
+      mx::array xf = mx::astype(x, mx::float16);
+      mx::array d = mx::matmul(mx::matmul(xf, mx::transpose(p->A)), mx::transpose(p->B));
+      out = mx::astype(meta_.scale * d, lin.apply(x).dtype());
+      return true;
+    }
+    if (which == 0) {
+      float wrong = static_cast<float>(meta_.alpha);
+      mx::array xf = mx::astype(x, mx::float16);
+      mx::array d = mx::matmul(mx::matmul(xf, mx::transpose(p->A)), mx::transpose(p->B));
+      out = lin.apply(x) + mx::astype(wrong * d, lin.apply(x).dtype());
+      return true;
+    }
+    if (which == 2) {
+      mx::array xf = mx::astype(x, mx::float32);
+      mx::array a32 = mx::astype(p->A, mx::float32), b32 = mx::astype(p->B, mx::float32);
+      mx::array d = mx::matmul(mx::matmul(xf, mx::transpose(a32)), mx::transpose(b32));
+      out = lin.apply(x) + mx::astype(meta_.scale * d, lin.apply(x).dtype());
+      return true;
+    }
+  } catch (const std::exception &ex) {
+    err = std::string("dressed_variant: ") + ex.what();
+    return false;
+  }
+  err = "dressed_variant: unknown which";
+  return false;
+}
+
+bool Bundle::delta_bits(int layer, int slot, const mx::array &x, std::vector<uint16_t> &bits,
+                        std::string &err) const {
+  const Pair *p = find_pair(layer, slot);
+  if (!p) { err = "delta_bits: no pairing"; return false; }
+  try {
+    mx::array xf = mx::astype(x, mx::float16);
+    mx::array d = mx::matmul(mx::matmul(xf, mx::transpose(p->A)), mx::transpose(p->B));
+    d = mx::astype(meta_.scale * d, mx::float16);
+    mx::eval(d);
+    const uint16_t *hp = d.data<uint16_t>();
+    bits.assign(hp, hp + d.size());
+  } catch (const std::exception &ex) {
+    err = std::string("delta_bits: ") + ex.what();
+    return false;
+  }
+  return true;
+}
+
+}  // namespace lora
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/lora.hpp b/macos/crates/engine-native/src/lora.hpp
new file mode 100644
index 0000000..8be0d28
--- /dev/null
+++ b/macos/crates/engine-native/src/lora.hpp
@@ -0,0 +1,103 @@
+// Unmerged parallel LoRA: y = base(x) + (scale * ((x.f16 @ A^T) @ B^T)).astype(y.dtype), scale = alpha/r.
+#pragma once
+#include <algorithm>
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include <mlx/mlx.h>
+
+#include "fwd_common.hpp"
+
+namespace e0n {
+namespace lora {
+
+enum Slot35 {
+  S35_IN_A = 0, S35_IN_B, S35_IN_QKV, S35_IN_Z, S35_OUT_PROJ,  // linear_attn (30 layers)
+  S35_Q, S35_K, S35_V, S35_O,                                   // self_attn (10 layers)
+  S35_SH_GATE, S35_SH_UP, S35_SH_DOWN,                          // mlp.shared_expert (40 layers)
+  S35_NSLOT = 12
+};
+enum Slot8b {
+  S8B_Q = 0, S8B_K, S8B_V, S8B_O, S8B_B, S8B_F, S8B_G,  // KDA attention (18 layers)
+  S8B_M_QB, S8B_M_KVB, S8B_M_DENSE, S8B_M_G,             // MLA attention (6 layers; g by family)
+  S8B_D_GATE, S8B_D_UP, S8B_D_DOWN,                      // L0 dense mlp only
+  S8B_NSLOT = 14
+};
+
+struct Meta {
+  int64_t r = 0, alpha = 0, k = 0;
+  float scale = 0.f;
+  bool ok = false;  // metadata complete and self-consistent
+};
+
+class Bundle {
+ public:
+  // Single safetensors file (self-parsed header + __metadata__).
+  // num_layers / family predicate: (li+1)%4==0 is full/MLA on both tiers.
+  bool load_file(const std::string& path, bool is_8b, int num_layers, std::string& err);
+  // In-memory tensors for tests; same core as load_file.
+  bool load_raw(std::map<std::string, std::pair<std::string, std::vector<uint8_t>>> tensors,
+                const std::string& meta_json, bool is_8b, int num_layers, std::string& err);
+
+  bool empty() const { return pairs_.empty(); }
+  const Meta& meta() const { return meta_; }
+  int n_bound() const { return static_cast<int>(pairs_.size()); }
+  int64_t bytes() const { return bytes_; }
+  const std::vector<std::string>& bound_keys() const { return keys_sorted_; }
+  int n_hits() const;
+  void reset_hits() const { std::fill(hit_.begin(), hit_.end(), (char)0); }
+
+  // y = base.apply(x); add delta if this (layer, slot) has a pair.
+  // In/out mismatch vs QuantLin is a wiring bug and throws.
+  mlx::core::array dressed(const QuantLin& lin, const mlx::core::array& x, int layer,
+                           int slot) const;
+
+  bool has_pair(int layer, int slot) const;
+  bool pair_raw(int layer, int slot, std::vector<uint8_t>& a, std::vector<uint8_t>& b,
+                int64_t& r, int64_t& in, int64_t& out) const;
+  // Test-only error shapes: 0 = scale uses α (32); 1 = delta replaces base; 2 = f32 path.
+  bool dressed_variant(const QuantLin& lin, const mlx::core::array& x, int layer, int slot,
+                       int which, mlx::core::array& out, std::string& err) const;
+  // Test-only: export the f16 delta bit pattern.
+  bool delta_bits(int layer, int slot, const mlx::core::array& x, std::vector<uint16_t>& bits,
+                  std::string& err) const;
+
+ private:
+  struct Pair {
+    int layer = -1, slot = -1;
+    int64_t r = 0, in = 0, out = 0;
+    std::vector<uint8_t> ab, bb;  // stable host buffers (unique_ptr-owned, no realloc)
+    mlx::core::array A{0.f}, B{0.f};
+    std::string key;
+  };
+  const Pair* find_pair(int layer, int slot) const {
+    if (layer < 0 || layer >= static_cast<int>(slots_.size()) || slot < 0 || slot >= ns_)
+      return nullptr;
+    int i = slots_[layer][slot];
+    return i < 0 ? nullptr : pairs_[i].get();
+  }
+
+  Meta meta_;
+  std::vector<std::unique_ptr<Pair>> pairs_;
+  std::vector<std::vector<int>> slots_;  // slots_[layer][slot] → pair index, -1 unbound
+  std::vector<std::string> keys_sorted_;
+  mutable std::vector<char> hit_;
+  // If E0_LORA_SCAFFOLD=<path.jsonl>, write base/delta prefix-64+summary on first dressed
+  // call per key. Observe-only; does not change numeric paths.
+  mutable std::string cap_path_;
+  mutable std::set<std::string> cap_done_;
+  int ns_ = 0;
+  int64_t bytes_ = 0;
+};
+
+inline mlx::core::array apply_lora(const Bundle* lo, const QuantLin& lin,
+                                   const mlx::core::array& x, int layer, int slot) {
+  return lo ? lo->dressed(lin, x, layer, slot) : lin.apply(x);
+}
+
+}  // namespace lora
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/memstat.cpp b/macos/crates/engine-native/src/memstat.cpp
new file mode 100644
index 0000000..a7bd851
--- /dev/null
+++ b/macos/crates/engine-native/src/memstat.cpp
@@ -0,0 +1,23 @@
+#include "memstat.hpp"
+
+#include <mach/mach.h>
+
+namespace e0n {
+
+MemSample mem_sample() {
+  MemSample s;
+  // resident_size includes file-backed pages; phys_footprint is dirty-page
+  // residency (Activity Monitor "Memory"). This SDK has no ri_resident_size_peak.
+  task_vm_info_data_t vi{};
+  mach_msg_type_number_t cnt = TASK_VM_INFO_COUNT;
+  kern_return_t kr = task_info(mach_task_self(), TASK_VM_INFO,
+                               reinterpret_cast<task_info_t>(&vi), &cnt);
+  if (kr == KERN_SUCCESS) {
+    s.resident_peak_bytes = vi.resident_size_peak;
+    s.resident_bytes = vi.resident_size;
+    s.footprint_bytes = vi.phys_footprint;
+  }
+  return s;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/memstat.hpp b/macos/crates/engine-native/src/memstat.hpp
new file mode 100644
index 0000000..f466d0e
--- /dev/null
+++ b/macos/crates/engine-native/src/memstat.hpp
@@ -0,0 +1,16 @@
+// Process memory sample for e0_stats: resident (incl. file-backed) vs phys_footprint.
+#pragma once
+#include <cstdint>
+
+namespace e0n {
+
+struct MemSample {
+  uint64_t resident_peak_bytes = 0;  // peak resident_size, including file-backed pages
+  uint64_t resident_bytes = 0;       // current resident_size (not in ABI stats)
+  uint64_t footprint_bytes = 0;      // current phys_footprint (Activity Monitor "Memory")
+};
+
+// Current process; zeros if rusage fails (caller must tell "unmeasured" from "measured 0").
+MemSample mem_sample();
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/net35.cpp b/macos/crates/engine-native/src/net35.cpp
new file mode 100644
index 0000000..a4e499c
--- /dev/null
+++ b/macos/crates/engine-native/src/net35.cpp
@@ -0,0 +1,1230 @@
+#include "net35.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <cstdlib>
+#include <functional>
+#include <numeric>  // iota (batched selk)
+#include <set>
+
+#include "dense.hpp"
+#include "qgemm.hpp"
+
+namespace e0n {
+namespace mx = mlx::core;
+
+namespace {
+
+using S = mx::Shape;
+
+namespace {
+inline uint16_t f_to_bf16(float x) {
+  uint32_t u;
+  std::memcpy(&u, &x, 4);
+  u += 0x7FFFu + ((u >> 16) & 1u);
+  return static_cast<uint16_t>(u >> 16);
+}
+inline void bf_round(std::vector<float> &v) {
+  for (float &x : v) {
+    uint32_t u = static_cast<uint32_t>(f_to_bf16(x)) << 16;
+    std::memcpy(&x, &u, 4);
+  }
+}
+}  // namespace
+
+mx::array first_k(const mx::array &a, int64_t k) {
+  mx::Shape s1 = a.shape();
+  mx::Shape s0 = a.shape();
+  for (auto &v : s0) v = 0;
+  s1[s1.size() - 1] = static_cast<mx::Shape::value_type>(k);
+  return mx::slice(a, s0, s1);
+}
+// heads.py select_from_logits(44-51) ≡ moe/routing.py select_from_logits(13-24):
+void select_softmax_gpu(const mx::array &logits, int64_t E, int64_t k, bool norm,
+                        mx::array &idx, mx::array &w) {
+  (void)E;  // shape is carried by logits; retained in the call signature for router symmetry.
+  auto gates = mx::softmax(logits, std::vector<int>{-1}, true /*precise*/);  // heads.py:47
+  idx = first_k(mx::argpartition(-gates, static_cast<int>(k - 1), -1), k);  // 48
+  w = mx::take_along_axis(gates, idx, -1);                                // 49
+  if (norm) w = w / mx::sum(w, std::vector<int>{-1}, true);               // 50(norm_topk_prob)
+}
+
+mx::array topk_onehot(const mx::array &inds, int64_t E) {
+  auto eq = mx::equal(mx::expand_dims(inds, -1), mx::arange(0, static_cast<int>(E)));
+  return mx::astype(mx::sum(eq, std::vector<int>{-2}, false), mx::float32);
+}
+
+mx::array gelu_erf16(const mx::array &x) {
+  return mx::array(0.5f, mx::float16) * x *
+         (mx::array(1.0f, mx::float16) +
+          mx::erf(x / mx::array(static_cast<float>(std::sqrt(2.0)), mx::float16)));
+}
+
+std::vector<mx::array> moe_math_body(const std::vector<mx::array> &a) {
+  auto xe = mx::reshape(a[0], mkshape({a[0].shape(0), a[0].shape(1), 1, 1, a[0].shape(2)}));  // layer.py:159
+  auto gq = [&](const mx::array &xx, int o) {
+    return mx::gather_qmm(xx, a[o], a[o + 1], a[o + 2], std::nullopt, a[10], true, 64, 4,
+                          "affine", false);
+  };
+  auto x_up = gq(xe, 1);
+  auto x_gate = gq(xe, 4);
+  auto yh = (x_gate * mx::sigmoid(x_gate)) * x_up;  // _swiglu(x_up, x_gate)(layer.py:57-58,169)
+  auto z = gq(yh, 7);
+  return {mx::squeeze(z, -2)};                      // 172
+}
+
+std::vector<mx::array> moe_math_sorted_body(const std::vector<mx::array> &a) {
+  auto gq = [&](const mx::array &xx, int o) {
+    return mx::gather_qmm(xx, a[o], a[o + 1], a[o + 2], std::nullopt, a[10], true, 64, 4,
+                          "affine", true);
+  };
+  auto x_up = gq(a[0], 1);
+  auto x_gate = gq(a[0], 4);
+  auto yh = (x_gate * mx::sigmoid(x_gate)) * x_up;             // _swiglu
+  return {gq(yh, 7)};                                          // z=[T*k,1,H]
+}
+
+ExpertPtrs bundle_ptrs(const Bundle &b) {
+  ExpertPtrs p;
+  p.uw = b.at(0).data<uint32_t>();
+  p.us = reinterpret_cast<const uint16_t *>(b.at(1).data<mx::bfloat16_t>());
+  p.ub = reinterpret_cast<const uint16_t *>(b.at(2).data<mx::bfloat16_t>());
+  p.gw = b.at(3).data<uint32_t>();
+  p.gs = reinterpret_cast<const uint16_t *>(b.at(4).data<mx::bfloat16_t>());
+  p.gb = reinterpret_cast<const uint16_t *>(b.at(5).data<mx::bfloat16_t>());
+  p.dw = b.at(6).data<uint32_t>();
+  p.ds = reinterpret_cast<const uint16_t *>(b.at(7).data<mx::bfloat16_t>());
+  p.db = reinterpret_cast<const uint16_t *>(b.at(8).data<mx::bfloat16_t>());
+  return p;
+}
+
+}  // namespace
+
+std::unique_ptr<Net35> Net35::create(std::shared_ptr<const ModelFiles> files,
+                                     std::shared_ptr<const ExpertMap> emap, std::string &err,
+                                     std::shared_ptr<const lora::Bundle> lo) {
+  std::unique_ptr<Net35> n(new Net35());
+  n->files_ = files;
+  n->emap_ = emap;
+  n->lo_ = std::move(lo);
+  if (!tier_routing("edge0-35b", n->reg_)) { err = "registry missing edge0-35b"; return nullptr; }
+  int64_t cfg_layers = 0;
+  if (files->config_int("num_hidden_layers", cfg_layers) && cfg_layers != n->reg_.num_layers) {
+    err = "config layer count does not match registry";
+    return nullptr;
+  }
+  auto ci = [&](const char *k, int64_t def, int64_t &o) {
+    int64_t v;
+    o = files->config_int(k, v) ? v : def;
+  };
+  auto cd = [&](const char *k, double def, double &o) {
+    double v;
+    o = files->config_double(k, v) ? v : def;
+  };
+  ci("hidden_size", 2048, n->H_);
+  ci("vocab_size", 248320, n->V_);
+  ci("moe_intermediate_size", 512, n->moe_ffn_);
+  ci("shared_expert_intermediate_size", 512, n->sh_ffn_);
+  ci("linear_num_key_heads", 16, n->hk_);
+  ci("linear_num_value_heads", 32, n->hv_);
+  ci("linear_key_head_dim", 128, n->dk_);
+  ci("linear_value_head_dim", 128, n->dv_);
+  ci("linear_conv_kernel_dim", 4, n->conv_k_);
+  ci("num_attention_heads", 16, n->nq_);
+  ci("num_key_value_heads", 2, n->nkv_);
+  ci("head_dim", 256, n->hd_);
+  int64_t interval = 4;
+  ci("full_attention_interval", 4, interval);
+  n->full_interval_ = static_cast<double>(interval);
+  cd("rms_norm_eps", 1e-6, n->eps_);
+  cd("rope_theta", 100000.0, n->rope_base_);
+  cd("partial_rotary_factor", 0.25, n->partial_rot_);
+  bool og = true;
+  if (!files->config_bool("attn_output_gate", og)) og = true;
+  n->attn_gate_ = og;
+
+  const int64_t keyd = n->hk_ * n->dk_, vald = n->hv_ * n->dv_;
+  n->conv_dim_ = 2 * keyd + vald;
+  {
+    static const std::string kGdnSource = R"gdn(
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }
+
+        // g: [B, T, Hv]
+        auto g_ = g + b_idx * T * Hv;
+        auto beta_ = beta + b_idx * T * Hv;
+
+        for (int t = 0; t < T; ++t) {
+          if (true) {
+            float kv_mem = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] * g_[hv_idx];
+              kv_mem += state[i] * k_[s_idx];
+            }
+            kv_mem = simd_sum(kv_mem);
+
+            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+
+            float out = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] + k_[s_idx] * delta;
+              out += state[i] * q_[s_idx];
+            }
+            out = simd_sum(out);
+            if (thread_index_in_simdgroup == 0) {
+              y[dv_idx] = static_cast<InT>(out);
+            }
+          }
+          // Increment data pointers to next time step
+          q_ += Hk * Dk;
+          k_ += Hk * Dk;
+          v_ += Hv * Dv;
+          y += Hv * Dv;
+          g_ += Hv;
+          beta_ += Hv;
+        }
+        for (int i = 0; i < n_per_t; ++i) {
+          auto s_idx = n_per_t * dk_idx + i;
+          o_state[s_idx] = static_cast<InT>(state[i]);
+        }
+    )gdn";
+    n->gdn_kernel_ = mx::fast::metal_kernel(
+        "gated_delta_step", {"q", "k", "v", "g", "beta", "state_in", "T"}, {"y", "state_out"},
+        kGdnSource);
+  }
+
+  n->layers_.resize(static_cast<size_t>(n->reg_.num_layers));
+  const std::string M = "language_model.model.layers.";
+  for (int64_t li = 0; li < n->reg_.num_layers; li++) {
+    Layer &L = n->layers_[static_cast<size_t>(li)];
+    const std::string P = M + std::to_string(li);
+    L.is_linear = ((li + 1) % interval) != 0;
+    auto q = [&](const std::string &name, QuantLin &dst) {
+      if (!dst.load(*files, P + name, err)) { err += " @" + std::to_string(li); return false; }
+      n->dense_bytes_ += dst.wb.size() + dst.sb.size() + dst.bb.size();
+      return true;
+    };
+    auto vec = [&](const std::string &name, std::vector<float> &dstv) {
+      DenseF32 d;
+      if (!read_dense_f32(*files, P + name, d, err)) {
+        err += " @" + std::to_string(li);
+        return false;
+      }
+      dstv = std::move(d.v);
+      n->dense_bytes_ += dstv.size() * 4;
+      return true;
+    };
+    if (!vec(".input_layernorm.weight", L.in_ln_w) ||
+        !vec(".post_attention_layernorm.weight", L.post_ln_w))
+      return nullptr;
+    if (L.is_linear) {
+      if (!q(".linear_attn.in_proj_qkv.weight", L.in_proj_qkv) ||
+          !q(".linear_attn.in_proj_z.weight", L.in_proj_z) ||
+          !q(".linear_attn.in_proj_b.weight", L.in_proj_b) ||
+          !q(".linear_attn.in_proj_a.weight", L.in_proj_a) ||
+          !q(".linear_attn.out_proj.weight", L.out_proj) ||
+          !vec(".linear_attn.A_log", L.A_log) || !vec(".linear_attn.dt_bias", L.dt_bias) ||
+          !vec(".linear_attn.norm.weight", L.gnorm_w))
+        return nullptr;
+      if (static_cast<int64_t>(L.A_log.size()) != n->hv_ ||
+          static_cast<int64_t>(L.dt_bias.size()) != n->hv_ ||
+          static_cast<int64_t>(L.gnorm_w.size()) != n->dv_) {
+        err = "GDN parameter vector length mismatch";
+        return nullptr;
+      }
+      if (!files->read_tensor(P + ".linear_attn.conv1d.weight", L.conv_wb, err)) return nullptr;
+      L.conv_w = mx::array(L.conv_wb.data(), mkshape({n->conv_dim_, n->conv_k_, 1}),
+                           mx::bfloat16, [](void *) {});
+      n->dense_bytes_ += L.conv_wb.size();
+    } else {
+      if (!q(".self_attn.q_proj.weight", L.q_proj) || !q(".self_attn.k_proj.weight", L.k_proj) ||
+          !q(".self_attn.v_proj.weight", L.v_proj) || !q(".self_attn.o_proj.weight", L.o_proj) ||
+          !vec(".self_attn.q_norm.weight", L.qn_w) || !vec(".self_attn.k_norm.weight", L.kn_w))
+        return nullptr;
+    }
+    DenseF32 g;
+    int64_t gbits = 4, ggs = 64;
+    const std::string gw = P + ".mlp.gate.weight";
+    if (!files->config_quant_for(gw, gbits, ggs)) { err = "gate has no quantization config"; return nullptr; }
+    if (!read_dense_dequant(*files, gw, gbits, ggs, g, err)) return nullptr;
+    if (g.rows != n->reg_.num_experts || g.cols != n->H_) {
+      err = "gate shape mismatch";
+      return nullptr;
+    }
+    L.gate_h = std::move(g.v);
+    n->dense_bytes_ += L.gate_h.size() * 4;
+    if (!L.gate_q.load(*files, gw, err)) { err += " @gate_q"; return nullptr; }
+    n->dense_bytes_ += L.gate_q.wb.size() + L.gate_q.sb.size() + L.gate_q.bb.size();
+    if (!q(".mlp.shared_expert.gate_proj.weight", L.sh_gate) ||
+        !q(".mlp.shared_expert.up_proj.weight", L.sh_up) ||
+        !q(".mlp.shared_expert.down_proj.weight", L.sh_down) ||
+        !q(".mlp.shared_expert_gate.weight", L.sh_egate))
+      return nullptr;
+  }
+  if (!n->embed_.load(*files, "language_model.model.embed_tokens.weight", err)) return nullptr;
+  if (!n->lm_head_.load(*files, "language_model.lm_head.weight", err)) return nullptr;
+  n->dense_bytes_ += n->embed_.wb.size() + n->embed_.sb.size() + n->embed_.bb.size() +
+                     n->lm_head_.wb.size() + n->lm_head_.sb.size() + n->lm_head_.bb.size();
+  {
+    DenseF32 f;
+    if (!read_dense_f32(*files, "language_model.model.norm.weight", f, err)) return nullptr;
+    n->final_norm_w_ = std::move(f.v);
+  }
+  n->pr_ = std::make_unique<Prerouter>();
+  if (!Prerouter::load(*files, n->reg_, n->H_, *n->pr_, err)) return nullptr;
+  {
+    const int64_t NL2 = n->reg_.num_layers;
+    n->io_ = std::make_unique<ExpertIo>(files.get(), 64, 48, 8, 4);
+    for (int64_t li = 0; li < NL2; li++) {
+      ExpertIo::LayerCfg c;
+      c.layer = int(li);
+      c.prefix = emap->layer_prefix(int(li));
+      c.E = n->reg_.num_experts;
+      c.H = n->H_;
+      c.F = n->moe_ffn_;
+      c.group = 64;
+      c.staged_n = int(n->reg_.top_k);
+      c.incr_mode = true;
+      c.incr_writeback = true;
+      c.warm_willneed = true;
+      c.prefill_hot = (int)n->reg_.prefill_hot;
+      c.full_layer_prefill = n->reg_.full_layer_prefill;
+      n->io_->add_layer(c);
+    }
+  }
+  if (exec_mode() == ExecMode::Batched) {
+    const char *cl = getenv("MLX_CACHE_LIMIT_MB");
+    mx::set_cache_limit(static_cast<size_t>(cl ? atoi(cl) : 256) << 20);
+    const int64_t NL = n->reg_.num_layers;
+    n->m_in_cache_.assign(NL, mx::array(0.f));
+    n->oh_cur_.assign(NL, mx::array(0.f));
+    n->oh_prev_.assign(NL, mx::array(0.f));
+    n->pred_inds_gpu_.assign(NL, mx::array(0.f));
+    n->pred_scores_gpu_.assign(NL, mx::array(0.f));
+    const int64_t nH = n->pr_->n_heads(), f_in = n->pr_->stacked_f_in(),
+                  prh = n->pr_->stacked_prh(), hE = n->pr_->stacked_E();
+    if (nH > 0) {
+      auto wrap16 = [&](const std::vector<float> &v, std::initializer_list<int64_t> dims) {
+        return mx::astype(mx::array((void *)v.data(), mkshape(dims), mx::float32, [](void *) {}),
+                          mx::float16);
+      };
+      n->hw1_ = wrap16(n->pr_->stacked_fc1(), {nH, f_in, prh});
+      n->hw2_ = wrap16(n->pr_->stacked_fc2(), {nH, prh, hE});
+      n->hwl_ = wrap16(n->pr_->stacked_lin(), {nH, f_in, hE});
+      mx::eval(n->hw1_, n->hw2_, n->hwl_);
+    }
+    // @core.compile(layer.py:157;staged_k4 use_compile=True，options.py:129).
+    try {
+      auto fn = mx::compile(&moe_math_body);
+      const int64_t F = n->moe_ffn_;
+      std::vector<mx::array> d = {
+          mx::zeros(mkshape({1, 1, n->H_}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, F / 8}), mx::uint32),
+          mx::zeros(mkshape({2, n->H_, F / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, F / 64}), mx::bfloat16),
+          mx::zeros(mkshape({1, 1, 2}), mx::int32)};
+      auto out = fn(d);
+      mx::eval(out[0]);
+      n->moe_math_ = std::move(fn);
+    } catch (const std::exception &ex) {
+      std::fprintf(stderr, "edge0-35b: mlx compile unavailable (%s), MoE math falling back to eager\n",
+                   ex.what());
+      n->moe_math_ = &moe_math_body;
+    }
+    try {
+      auto fn = mx::compile(&moe_math_sorted_body);
+      const int64_t F = n->moe_ffn_;
+      std::vector<mx::array> d = {
+          mx::zeros(mkshape({16, 1, n->H_}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, F / 8}), mx::uint32),
+          mx::zeros(mkshape({2, n->H_, F / 64}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, F / 64}), mx::bfloat16),
+          mx::zeros(mkshape({16}), mx::int32)};
+      auto out = fn(d);
+      mx::eval(out[0]);
+      n->moe_math_sorted_ = std::move(fn);
+    } catch (const std::exception &ex) {
+      std::fprintf(stderr, "edge0-35b: mlx compile unavailable (%s), sorted MoE math falling back to eager\n",
+                   ex.what());
+      n->moe_math_sorted_ = &moe_math_sorted_body;
+    }
+  }
+  if (n->lo_) {
+    const int nl = n->reg_.num_layers;
+    const int full = nl / 4, lin = nl - full;
+    const int want = 5 * lin + 4 * full + 3 * nl;
+    if (n->lo_->n_bound() != want) {
+      err = "LoRA bound count " + std::to_string(n->lo_->n_bound()) + " != tier structural count " +
+            std::to_string(want);
+      return nullptr;
+    }
+    n->dense_bytes_ += static_cast<uint64_t>(n->lo_->bytes());
+  }
+  return n;
+}
+
+void Net35::init_state(State35 &st) const {
+  st.pr.init(reg_);
+  const size_t NL = layers_.size();
+  st.conv.assign(NL, mx::array(0.f));
+  st.recur.assign(NL, mx::array(0.f));
+  st.kc.assign(NL, mx::array(0.f));
+  st.vc.assign(NL, mx::array(0.f));
+  st.is_linear.assign(NL, 0);
+  for (size_t i = 0; i < NL; i++) st.is_linear[i] = layers_[i].is_linear ? 1 : 0;
+  st.position = 0;
+  st.started = true;
+  reset_graph_caches();
+}
+
+void Net35::reset_state(State35 &st) const {
+  st.pr.reset();
+  st.conv.assign(layers_.size(), mx::array(0.f));
+  st.recur.assign(layers_.size(), mx::array(0.f));
+  st.kc.assign(layers_.size(), mx::array(0.f));
+  st.vc.assign(layers_.size(), mx::array(0.f));
+  st.position = 0;
+  reset_graph_caches();
+}
+
+bool Net35::layer_fwd(Layer &L, int li, mx::array &x, State35 &st, bool decode, RouteRec &rec,
+                      mx::array *m_in_out, std::string &err) {
+  const int64_t T = x.shape(1);
+  const auto dt = x.dtype();
+  const int64_t keyd = hk_ * dk_, vald = hv_ * dv_;
+  auto norm_w = [&](const std::vector<float> &v) {
+    return mx::astype(wrap_f32(v, mkshape({static_cast<int64_t>(v.size())})), dt);
+  };
+  mx::array xn = e0_rms_norm(x, norm_w(L.in_ln_w), eps_);
+
+  mx::array r{0.f};
+  if (L.is_linear) {
+    const auto *lo = lo_.get();
+    using LS = e0n::lora::Slot35;
+    mx::array qkv = e0n::lora::apply_lora(lo, L.in_proj_qkv, xn, li, LS::S35_IN_QKV);  // [1,T,conv_dim]
+    mx::array z = mx::reshape(e0n::lora::apply_lora(lo, L.in_proj_z, xn, li, LS::S35_IN_Z),
+                              mkshape({1, T, hv_, dv_}));
+    mx::array b = e0n::lora::apply_lora(lo, L.in_proj_b, xn, li, LS::S35_IN_B);  // [1,T,hv]
+    mx::array a = e0n::lora::apply_lora(lo, L.in_proj_a, xn, li, LS::S35_IN_A);
+    if (st.conv[li].ndim() == 0)
+      st.conv[li] = mx::zeros(mkshape({1, conv_k_ - 1, conv_dim_}), dt);
+    mx::array cat = mx::concatenate({st.conv[li], qkv}, 1);  // [1, K-1+T, C]
+    st.conv[li] = mx::slice(cat, mkshape({0, T, 0}), mkshape({1, T + conv_k_ - 1, conv_dim_}));
+    mx::array convout = e0_silu(mx::conv1d(
+        cat, L.conv_w, 1, 0, 1, static_cast<int>(conv_dim_)));  // [1,T,C]
+    auto parts = mx::split(convout, S{static_cast<int32_t>(keyd), static_cast<int32_t>(2 * keyd)},
+                           -1);
+    mx::array qf = mx::reshape(parts[0], mkshape({1, T, hk_, dk_}));
+    mx::array kf = mx::reshape(parts[1], mkshape({1, T, hk_, dk_}));
+    mx::array vf = mx::reshape(parts[2], mkshape({1, T, hv_, dv_}));
+    float inv = 1.0f / std::sqrt(static_cast<float>(dk_));
+    auto scalar_bf = [&](float v) { return mx::astype(mx::array(v), dt); };
+    qf = scalar_bf(inv * inv) * e0_rms_norm(qf, std::nullopt, 1e-6);
+    kf = scalar_bf(inv) * e0_rms_norm(kf, std::nullopt, 1e-6);
+    mx::array beta = mx::sigmoid(b);
+    mx::array alog = wrap_f32(L.A_log, mkshape({hv_}));
+    mx::array dtb = mx::astype(wrap_f32(L.dt_bias, mkshape({hv_})), dt);  // bf16
+    mx::array xab = a + dtb;
+    mx::array sp = mx::log1p(mx::exp(-mx::abs(xab))) +
+                   mx::maximum(xab, mx::astype(mx::array(0.f), dt));
+    mx::array g = mx::astype(mx::exp(-mx::exp(alog) * mx::astype(sp, mx::float32)), dt);
+    if (st.recur[li].ndim() == 0)
+      st.recur[li] = mx::zeros(mkshape({1, hv_, dv_, dk_}), dt);  // gated_delta_update:275-278
+    auto gouts = gdn_kernel_(
+        {qf, kf, vf, g, beta, st.recur[li], mx::array((int)T, mx::int32)},
+        {mkshape({1, T, hv_, dv_}), st.recur[li].shape()},
+        {dt, dt},
+        {32, static_cast<int>(dv_), static_cast<int>(hv_)},  // grid=(32, Dv, B*Hv)，B=1
+        {32, 4, 1},                                          // threadgroup(gated_delta.py:206)
+        {{"InT", dt},
+         {"Dk", static_cast<int>(dk_)},
+         {"Dv", static_cast<int>(dv_)},
+         {"Hk", static_cast<int>(hk_)},
+         {"Hv", static_cast<int>(hv_)}},
+        std::nullopt, false, {});
+    mx::array y_arr = gouts[0];
+    st.recur[li] = gouts[1];
+    mx::array ng = e0_rms_norm(y_arr, norm_w(L.gnorm_w), eps_);
+    mx::array gated = e0_swiglu_precise(dt, z, ng);  // _precise_swiglu(out, z, x):silu(z)*x
+    r = e0n::lora::apply_lora(lo_.get(), L.out_proj,
+                              mx::reshape(gated, mkshape({1, T, vald})), li,
+                              e0n::lora::S35_OUT_PROJ);
+  } else {
+    mx::array qo = e0n::lora::apply_lora(lo_.get(), L.q_proj, xn, li, e0n::lora::S35_Q);
+    auto qs = mx::split(mx::reshape(qo, mkshape({1, T, nq_, hd_ * 2})),
+                        S{static_cast<int32_t>(hd_)}, -1);
+    mx::array qg = qs[0], gate = mx::reshape(qs[1], mkshape({1, T, nq_ * hd_}));
+    mx::array q = mx::transpose(e0_rms_norm(qg, norm_w(L.qn_w), eps_), {0, 2, 1, 3});
+    mx::array k = mx::transpose(
+        e0_rms_norm(mx::reshape(e0n::lora::apply_lora(lo_.get(), L.k_proj, xn, li,
+                                                   e0n::lora::S35_K),
+                                   mkshape({1, T, nkv_, hd_})),
+                    norm_w(L.kn_w), eps_),
+        {0, 2, 1, 3});
+    mx::array v = mx::transpose(mx::reshape(e0n::lora::apply_lora(lo_.get(), L.v_proj, xn, li,
+                                                                        e0n::lora::S35_V),
+                                                        mkshape({1, T, nkv_, hd_})),
+                                {0, 2, 1, 3});
+    const int rdim = static_cast<int>(hd_ * partial_rot_);
+    q = mx::fast::rope(q, rdim, false, static_cast<float>(rope_base_), 1.0f,
+                       static_cast<int>(st.position));
+    k = mx::fast::rope(k, rdim, false, static_cast<float>(rope_base_), 1.0f,
+                       static_cast<int>(st.position));
+    if (st.kc[li].ndim() == 0) {
+      st.kc[li] = k;
+      st.vc[li] = v;
+    } else {
+      st.kc[li] = mx::concatenate({st.kc[li], k}, 2);
+      st.vc[li] = mx::concatenate({st.vc[li], v}, 2);
+    }
+    mx::array o = mx::fast::scaled_dot_product_attention(
+        q, st.kc[li], st.vc[li], 1.0f / std::sqrt(static_cast<float>(hd_)), "causal");
+    o = mx::reshape(mx::transpose(o, {0, 2, 1, 3}), mkshape({1, T, nq_ * hd_}));
+    if (attn_gate_) o = o * mx::sigmoid(gate);
+    r = e0n::lora::apply_lora(lo_.get(), L.o_proj, o, li, e0n::lora::S35_O);
+  }
+
+  mx::array h = x + r;
+  mx::array m_in = e0_rms_norm(h, norm_w(L.post_ln_w), eps_);
+  if (m_in_out) *m_in_out = m_in;
+  if (exec_mode() == ExecMode::Batched && decode && T == 1) {
+    const int64_t E = reg_.num_experts, k = reg_.top_k;
+    io_->note_call(li);
+    const bool owner = li >= reg_.first_owner && li <= reg_.last_owner;   // 6..38
+    const bool consumer = li >= reg_.first_head_consumer && li <= reg_.last_head_consumer;  // 7..38
+    mx::array idx(0.f), w(0.f), routed(0.f);
+    rec.layer = li;
+    if (consumer && pred_inds_gpu_[static_cast<size_t>(li)].ndim() > 0) {
+      idx = pred_inds_gpu_[static_cast<size_t>(li)];
+      w = pred_scores_gpu_[static_cast<size_t>(li)];
+      rec.predicted = true;
+      rec.sel = st.pr.pred[static_cast<size_t>(li)];
+      st.pr.pred_avail[static_cast<size_t>(li)] = false;
+      io_->wait_staged(li);
+      auto local = mx::take(io_->incr_slot_table(li), idx);  // layer.py:1338 take(_incr_slot_table, indices)
+      std::vector<mx::array> args{m_in};
+      for (auto &a : io_->incr_tensors(li)) args.push_back(a);
+      args.push_back(local);
+      auto &cx = exec_counters();
+      cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+      cx.qmm_experts.fetch_add(3 * static_cast<uint64_t>(k), std::memory_order_relaxed);
+      cx.qmm_weight_bytes.fetch_add(
+          static_cast<uint64_t>(k) *
+              (2 * (moe_ffn_ * (H_ / 8) * 4 + 2 * moe_ffn_ * (H_ / 64) * 2) +
+               (H_ * (moe_ffn_ / 8) * 4 + 2 * H_ * (moe_ffn_ / 64) * 2)),
+          std::memory_order_relaxed);
+      routed = moe_math_(args)[0];  // layer.py:1341
+    } else {
+      const int F = static_cast<int>(moe_ffn_);
+      const int gsh = static_cast<int>(H_ / 64), gsf = static_cast<int>(moe_ffn_ / 64);
+      auto gl = L.gate_q.apply(m_in);
+      select_softmax_gpu(gl, E, k, reg_.norm_topk_prob, idx, w);
+      mx::eval(idx, w);
+      exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+      auto idxh = mx::copy(mx::astype(idx, mx::int32), mx::default_stream(mx::Device::cpu));
+      auto wh = mx::copy(mx::astype(w, mx::float32), mx::default_stream(mx::Device::cpu));
+      mx::eval(idxh, wh);
+      const int32_t *ip = idxh.data<int32_t>();
+      const float *wp = wh.data<float>();
+      std::vector<int64_t> flat(ip, ip + k);
+      std::vector<int64_t> uniq = flat;
+      std::sort(uniq.begin(), uniq.end());
+      uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+      std::vector<int32_t> local_h(flat.size());
+      for (size_t j = 0; j < flat.size(); j++)
+        local_h[j] = static_cast<int32_t>(
+            std::lower_bound(uniq.begin(), uniq.end(), flat[j]) - uniq.begin());
+      auto bundles = io_->get_bundles(li, uniq);
+      if (bundles.size() != uniq.size()) {
+        err = "expert load missing L" + std::to_string(li);
+        return false;
+      }
+      exec_counters().leases.fetch_add(uniq.size(), std::memory_order_relaxed);
+      for (auto &[e, b] : bundles) step_bundles_.push_back(b);
+      std::vector<mx::array> args{m_in};
+      for (int p = 0; p < 9; p++) {
+        std::vector<mx::array> rows;
+        for (auto e : uniq) rows.push_back(bundles[e]->at(p));
+        args.push_back(mx::stack(rows));
+      }
+      args.push_back(mx::array(local_h.begin(), mkshape({1, 1, k}), mx::int32));  // 1371
+      auto &cx = exec_counters();
+      cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+      cx.qmm_experts.fetch_add(3 * static_cast<uint64_t>(uniq.size()), std::memory_order_relaxed);
+      cx.qmm_weight_bytes.fetch_add(
+          static_cast<uint64_t>(uniq.size()) *
+              (2 * (F * (H_ / 8) * 4 + 2 * F * gsh * 2) + (H_ * (F / 8) * 4 + 2 * H_ * gsf * 2)),
+          std::memory_order_relaxed);
+      routed = moe_math_(args)[0];
+      rec.predicted = false;
+      rec.sel.inds = uniq;
+      rec.sel.scores.clear();
+      for (int64_t e : uniq) {
+        double s = 0.0;
+        for (int64_t j = 0; j < k; j++)
+          if (flat[j] == e) s = wp[j];
+        rec.sel.scores.push_back(s);
+      }
+    }
+    if (owner) {
+      m_in_cache_[static_cast<size_t>(li)] = m_in;
+      oh_cur_[static_cast<size_t>(li)] = topk_onehot(idx, E);
+    }
+    // mix+shared(install.py:190-194):(y*scores[...,None]).sum(-2) + sigmoid(egate(x))*shared(x)
+    auto out = mx::sum(routed * mx::expand_dims(w, -1), std::vector<int>{-2}, false);
+    auto shg = e0n::lora::apply_lora(lo_.get(), L.sh_gate, m_in, li, e0n::lora::S35_SH_GATE);
+    auto shu = e0n::lora::apply_lora(lo_.get(), L.sh_up, m_in, li, e0n::lora::S35_SH_UP);
+    mx::array sh =
+        e0n::lora::apply_lora(lo_.get(), L.sh_down, e0_silu(shg) * shu, li, e0n::lora::S35_SH_DOWN);
+    mx::array eg = mx::sigmoid(L.sh_egate.apply(m_in));
+    x = h + out + sh * eg;
+    if (clip_ > 0)
+      x = mx::where(mx::isnan(x), mx::zeros_like(x),
+                    mx::clip(x, mx::array(static_cast<float>(-clip_)),
+                             mx::array(static_cast<float>(clip_))));
+    return true;
+  }
+  if (exec_mode() == ExecMode::Batched) {
+    const int64_t E = reg_.num_experts, k = reg_.top_k;
+    io_->note_call(li);
+    mx::array idx(0.f), w(0.f);
+    rec.layer = li;
+    rec.predicted = false;
+    {
+      auto gl = L.gate_q.apply(m_in);
+      select_softmax_gpu(gl, E, k, reg_.norm_topk_prob, idx, w);  // qwen3_next.py:335-341
+    }
+    mx::eval(idx, w);
+    exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+    auto idxh = mx::copy(mx::astype(idx, mx::int32), mx::default_stream(mx::Device::cpu));
+    auto wh = mx::copy(mx::astype(w, mx::float32), mx::default_stream(mx::Device::cpu));
+    mx::eval(idxh, wh);
+    const int32_t *ip = idxh.data<int32_t>();
+    const float *wp = wh.data<float>();
+    std::vector<int64_t> flat(static_cast<size_t>(T * k));
+    for (int64_t i = 0; i < T * k; i++) flat[static_cast<size_t>(i)] = ip[i];
+    mx::array routed(0.f);
+    const bool hot = !prefill_ondemand() && T * k > 8 && io_->hot_ready(li) &&
+                     !io_->full_ready(li);  // __call__ guard(layer.py:1121-1122)
+    if (hot) {
+      io_->note_prefill_use(li, flat, (int)k,
+                            li >= reg_.first_head_consumer && li <= reg_.last_head_consumer);
+      if (!moe_hot_branch_(li, m_in, T, flat, routed, err)) return false;
+    }
+    if (!hot) {
+      std::vector<int64_t> uniq = flat;                     // 1364 sorted(set(...))
+      std::sort(uniq.begin(), uniq.end());
+      uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+      std::vector<int32_t> local_h(static_cast<size_t>(T * k));  // 1370-1373 host remap
+      for (int64_t i = 0; i < T * k; i++)
+        local_h[static_cast<size_t>(i)] = static_cast<int32_t>(
+            std::lower_bound(uniq.begin(), uniq.end(), flat[static_cast<size_t>(i)]) - uniq.begin());
+      auto bundles = io_->get_bundles(li, uniq);            // 1368
+      if (bundles.size() != uniq.size()) {
+        err = "expert load missing L" + std::to_string(li);
+        return false;
+      }
+      exec_counters().leases.fetch_add(uniq.size(), std::memory_order_relaxed);
+      for (auto &[e, b] : bundles) step_bundles_.push_back(b);
+      std::vector<mx::array> wargs;                         // 1376-1380 core.stack(rows)
+      for (int p = 0; p < 9; p++) {
+        std::vector<mx::array> rows;
+        for (auto e : uniq) rows.push_back(bundles[e]->at(p));
+        wargs.push_back(mx::stack(rows));
+      }
+      auto local_arr = mx::array(local_h.begin(), mkshape({1, T, k}), mx::int32);  // 1371
+      const int F = static_cast<int>(moe_ffn_);
+      auto &cx = exec_counters();
+      cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+      cx.qmm_experts.fetch_add(3 * static_cast<uint64_t>(uniq.size()), std::memory_order_relaxed);
+      cx.qmm_weight_bytes.fetch_add(
+          static_cast<uint64_t>(uniq.size()) *
+              (2 * (F * (H_ / 8) * 4 + 2 * F * (H_ / 64) * 2) +
+               (H_ * (F / 8) * 4 + 2 * H_ * (F / 64) * 2)),
+          std::memory_order_relaxed);
+      if (T * k >= 64) {                            // 1384 do_sort = local2d.size >= 64
+        auto flat_local = mx::reshape(local_arr, mkshape({T * k}));
+        auto order = mx::argsort(flat_local, -1);   // _gather_sort(switch_layers.py:12-17)
+        auto inv_order = mx::argsort(order, -1);
+        auto xe = mx::reshape(m_in, mkshape({1, T, 1, 1, H_}));  // 1383 expand_dims((-2,-3))
+        auto xs = mx::flatten(xe, 0, -3);           // [T,1,H]
+        auto tok_rows = mx::take(
+            xs, mx::floor_divide(order, mx::array((uint32_t)k, mx::uint32)), 0);  // x[order//M]
+        auto local_sorted = mx::take(flat_local, order, 0);  // indices[order]
+        std::vector<mx::array> sargs{tok_rows};
+        for (auto &wa : wargs) sargs.push_back(wa);
+        sargs.push_back(local_sorted);
+        auto z = moe_math_sorted_(sargs)[0];        // [T*k,1,H]
+        auto z_un = mx::take(z, inv_order, 0);
+        routed = mx::squeeze(mx::reshape(z_un, mkshape({1, T, k, 1, H_})), -2);
+      } else {
+        std::vector<mx::array> args{m_in};
+        for (auto &wa : wargs) args.push_back(wa);
+        args.push_back(local_arr);
+        routed = moe_math_(args)[0];
+      }
+    }
+    // mix+shared(install.py:190-194):(y*scores[...,None]).sum(-2) + sigmoid(egate(x))*shared(x)
+    auto out = mx::sum(routed * mx::expand_dims(w, -1), std::vector<int>{-2}, false);
+    auto shg = e0n::lora::apply_lora(lo_.get(), L.sh_gate, m_in, li, e0n::lora::S35_SH_GATE);
+    auto shu = e0n::lora::apply_lora(lo_.get(), L.sh_up, m_in, li, e0n::lora::S35_SH_UP);
+    mx::array sh =
+        e0n::lora::apply_lora(lo_.get(), L.sh_down, e0_silu(shg) * shu, li, e0n::lora::S35_SH_DOWN);
+    mx::array eg = mx::sigmoid(L.sh_egate.apply(m_in));
+    x = h + out + sh * eg;
+    if (clip_ > 0)
+      x = mx::where(mx::isnan(x), mx::zeros_like(x),
+                    mx::clip(x, mx::array(static_cast<float>(-clip_)),
+                             mx::array(static_cast<float>(clip_))));
+    {
+      std::vector<std::pair<int64_t, double>> lv;
+      for (int64_t j = 0; j < k; j++)
+        lv.push_back({ip[(T - 1) * k + j], wp[(T - 1) * k + j]});
+      std::sort(lv.begin(), lv.end());
+      for (auto &[e, s] : lv) {
+        rec.sel.inds.push_back(e);
+        rec.sel.scores.push_back(s);
+      }
+    }
+    if (!hot) {
+      mx::eval(x);
+      step_bundles_.clear();
+    }
+    return true;
+  }
+
+  std::vector<float> mhost;
+  if (!to_host32(m_in, mhost, err)) return false;  // [T,H]
+
+  const int64_t E = reg_.num_experts;
+  std::vector<Selection> sels(static_cast<size_t>(T));
+  bool predicted = false;
+  if (decode && st.pr.take_pred(li, sels[0])) {
+    predicted = true;
+  } else {
+    std::vector<double> lg;
+    {
+      mx::array glh = L.gate_q.apply(m_in);
+      std::vector<float> glf;
+      if (!to_host32(glh, glf, err)) return false;
+      lg.assign(glf.begin(), glf.end());
+    }
+    for (int64_t t = 0; t < T; t++)
+      sels[static_cast<size_t>(t)] = select_softmax_topk(
+          lg.data() + t * E, static_cast<int>(E), reg_.top_k, reg_.norm_topk_prob);
+  }
+  rec.layer = li;
+  rec.predicted = predicted;
+  rec.sel = sels[static_cast<size_t>(T - 1)];
+  std::vector<float> last_row(mhost.begin() + (T - 1) * H_, mhost.end());
+  st.pr.capture(li, std::move(last_row), rec.sel.inds);
+
+  std::vector<int64_t> uni;
+  for (auto &s : sels)
+    for (auto e : s.inds)
+      if (uni.empty() || uni.back() != e) uni.push_back(e);
+  std::sort(uni.begin(), uni.end());
+  uni.erase(std::unique(uni.begin(), uni.end()), uni.end());
+  auto bundles = io_->get_bundles(li, uni);
+  exec_counters().leases.fetch_add(uni.size(), std::memory_order_relaxed);
+  std::vector<ExpertPtrs> ptrs(uni.size());
+  for (size_t i = 0; i < uni.size(); i++)
+    ptrs[i] = bundle_ptrs(*bundles[uni[i]]);
+  std::vector<float> routed(static_cast<size_t>(T) * H_, 0.f);
+  const int F = static_cast<int>(moe_ffn_);
+  {
+  for (size_t ei = 0; ei < uni.size(); ei++) {
+    int64_t e = uni[ei];
+    std::vector<int> toks;
+    for (int64_t t = 0; t < T; t++)
+      for (auto id : sels[static_cast<size_t>(t)].inds)
+        if (id == e) toks.push_back(static_cast<int>(t));
+    if (toks.empty()) continue;
+    std::vector<float> rows(toks.size() * H_);
+    for (size_t k = 0; k < toks.size(); k++)
+      std::memcpy(rows.data() + k * H_, mhost.data() + static_cast<size_t>(toks[k]) * H_,
+                  H_ * 4);
+    std::vector<int64_t> one{0};
+    std::vector<float> yg, yu;
+    if (!gather_qmm_mlx(ptrs[ei].gw, ptrs[ei].gs, ptrs[ei].gb, 1, F, static_cast<int>(H_),
+                        static_cast<int>(H_ / 64), rows.data(), static_cast<int>(toks.size()), one,
+                        yg, err) ||
+        !gather_qmm_mlx(ptrs[ei].uw, ptrs[ei].us, ptrs[ei].ub, 1, F, static_cast<int>(H_),
+                        static_cast<int>(H_ / 64), rows.data(), static_cast<int>(toks.size()), one,
+                        yu, err))
+      return false;
+    std::vector<float> yh(toks.size() * moe_ffn_);
+    {
+      bf_round(yg);
+      bf_round(yu);
+      auto bf = [](float v) {
+        uint32_t u; std::memcpy(&u, &v, 4);
+        u += 0x7FFFu + ((u >> 16) & 1u); u >>= 16; u <<= 16;
+        float o; std::memcpy(&o, &u, 4); return o;
+      };
+      for (size_t k = 0; k < yh.size(); k++) {
+        float g = yg[k];
+        float sg = bf(1.0f / (1.0f + std::exp(-g)));
+        float sl = bf(g * sg);
+        yh[k] = bf(sl * yu[k]);
+      }
+    }
+    std::vector<float> yd;
+    if (!gather_qmm_mlx(ptrs[ei].dw, ptrs[ei].ds, ptrs[ei].db, 1, static_cast<int>(H_), F,
+                        static_cast<int>(moe_ffn_ / 64), yh.data(),
+                        static_cast<int>(toks.size()), one, yd, err))
+      return false;
+    bf_round(yd);
+    for (size_t k = 0; k < toks.size(); k++) {
+      int t = toks[k];
+      double w = 0.0;
+      auto &s = sels[static_cast<size_t>(t)];
+      for (size_t j = 0; j < s.inds.size(); j++)
+        if (s.inds[j] == e) w = s.scores[j];
+      for (int64_t o = 0; o < H_; o++)
+        routed[static_cast<size_t>(t) * H_ + o] += static_cast<float>(w * yd[k * H_ + o]);
+    }
+  }
+  }
+
+  auto shg = e0n::lora::apply_lora(lo_.get(), L.sh_gate, m_in, li, e0n::lora::S35_SH_GATE);
+  auto shu = e0n::lora::apply_lora(lo_.get(), L.sh_up, m_in, li, e0n::lora::S35_SH_UP);
+  mx::array sh = e0n::lora::apply_lora(
+      lo_.get(), L.sh_down, e0_silu(shg) * shu, li, e0n::lora::S35_SH_DOWN);  // [1,T,H]
+  mx::array eg = mx::sigmoid(L.sh_egate.apply(m_in));                      // [1,T,1]
+  std::vector<float> shost;
+  if (!to_host32(sh * eg, shost, err)) return false;
+  for (int64_t t = 0; t < T; t++)
+    for (int64_t o = 0; o < H_; o++)
+      routed[static_cast<size_t>(t) * H_ + o] += shost[static_cast<size_t>(t) * H_ + o];
+  mx::array y_arr = mx::astype(mx::reshape(wrap_f32(routed, mkshape({1, T, H_})),
+                                           mkshape({1, T, H_})),
+                               dt);
+  x = h + y_arr;
+  if (clip_ > 0) {
+    x = mx::where(mx::isnan(x), mx::zeros_like(x),
+                  mx::clip(x, mx::array(static_cast<float>(-clip_)),
+                           mx::array(static_cast<float>(clip_))));
+  }
+  mx::eval(x);
+  return true;
+}
+
+bool Net35::stage_all_gpu(State35 &st, std::string &err, DecodeProbe *probe) {
+  const int64_t E = reg_.num_experts;
+  std::vector<int> metas;
+  std::vector<mx::array> ms, curs, prevs;
+  for (int o : pr_->owners()) {
+    if (m_in_cache_[static_cast<size_t>(o)].ndim() < 3 ||
+        oh_cur_[static_cast<size_t>(o)].ndim() == 0)
+      continue;
+    metas.push_back(o);
+    const auto &m = m_in_cache_[static_cast<size_t>(o)];
+    mx::Shape s0 = m.shape(), s1 = m.shape();
+    for (auto &v : s0) v = 0;
+    s0[s0.size() - 2] = static_cast<mx::Shape::value_type>(m.shape(-2) - 1);
+    ms.push_back(mx::slice(m, s0, s1));
+    curs.push_back(oh_cur_[static_cast<size_t>(o)]);  // oh = executed(install:183-188)
+    mx::array po(0.f);
+    if (oh_prev_[static_cast<size_t>(o)].ndim() > 0)
+      po = oh_prev_[static_cast<size_t>(o)];
+    else
+      po = mx::zeros(mkshape({1, 1, E}), mx::float32);
+    prevs.push_back(po);
+  }
+  if (metas.empty()) return true;
+  const int64_t n = static_cast<int64_t>(metas.size());
+  auto feats = mx::astype(
+      mx::concatenate({mx::stack(ms, 0), mx::stack(curs, 0), mx::stack(prevs, 0)}, -1),
+      mx::float16);                            // stager.py:158-161
+  auto f2 = mx::reshape(feats, mkshape({n, H_ + 2 * E}));  // 162
+  auto h1 = mx::einsum("ni,nij->nj", {f2, hw1_});          // 163
+  auto act = gelu_erf16(h1);
+  auto lin = mx::einsum("ni,nij->nj", {f2, hwl_});
+  auto fc2 = mx::einsum("ni,nij->nj", {act, hw2_});
+  auto out = lin + fc2;                                      // 164-166
+  auto logits_all = mx::reshape(out, mkshape({n, 1, 1, E}));    // 167
+  mx::array inds_all(0.f), scores_all(0.f);
+  select_softmax_gpu(logits_all, E, reg_.top_k, reg_.norm_topk_prob, inds_all, scores_all);  // _select:241-242
+  // `first_k` is a strided view over the E-wide row. Materialize both outputs before
+  // the single host read; reading data() from the view as if it were [n,k] walks into
+  // the next row and corrupts the staged expert set (same fix as net8b).
+  auto inds_h = mx::add(inds_all, mx::array((uint32_t)0, mx::uint32));
+  auto scores_h = mx::add(scores_all, mx::array(0.0f, mx::float32));
+  mx::eval(inds_h, scores_h);                  // stager:193 ONE eval
+  exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+  const int k = static_cast<int>(reg_.top_k);
+  const uint32_t *ip = inds_h.data<uint32_t>();  // 194 ONE tolist
+  const float *sp = scores_h.data<float>();
+  for (int64_t i = 0; i < n; i++) {
+    const int o = metas[i], consumer = o + 1;  // 196-197
+    Selection s;                                                     // _store:244-252
+    std::vector<int64_t> eids;
+    for (int j = 0; j < k; j++) eids.push_back(ip[i * k + j]);
+    std::sort(eids.begin(), eids.end());
+    eids.erase(std::unique(eids.begin(), eids.end()), eids.end());
+    for (auto e : eids) {
+      double v = 0.0;
+      for (int j = 0; j < k; j++)
+        if (ip[i * k + j] == static_cast<uint32_t>(e)) v = sp[i * k + j];
+      s.inds.push_back(e);
+      s.scores.push_back(v);
+    }
+    if (consumer >= reg_.first_head_consumer && consumer <= reg_.last_head_consumer) {
+      st.pr.pred[static_cast<size_t>(consumer)] = s;
+      st.pr.pred_avail[static_cast<size_t>(consumer)] = true;
+      io_->stage_experts(consumer, s.inds);
+      mx::Shape a0 = inds_all.shape(), a1 = inds_all.shape();  // pred_inds[p]/pred_scores[p]
+      for (auto &v : a0) v = 0;
+      a0[0] = static_cast<mx::Shape::value_type>(i);
+      a1[0] = static_cast<mx::Shape::value_type>(i + 1);
+      pred_inds_gpu_[static_cast<size_t>(consumer)] = mx::squeeze(mx::slice(inds_all, a0, a1), 0);     // [1,1,k]
+      pred_scores_gpu_[static_cast<size_t>(consumer)] = mx::squeeze(mx::slice(scores_all, a0, a1), 0);
+    }
+    oh_prev_[static_cast<size_t>(o)] = curs[i];
+  }
+  st.pr.step++;
+  if (probe) {
+    auto lg32 = mx::astype(logits_all, mx::float32);
+    mx::eval(lg32);
+    exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+    const float *lp = lg32.data<float>();
+    probe->owners.clear();
+    probe->m_in.clear();
+    probe->logits.clear();
+    for (int64_t i = 0; i < n; i++) {
+      probe->owners.push_back(metas[i]);
+      probe->logits.emplace_back(lp + i * E, lp + (i + 1) * E);
+      std::vector<float> mv;
+      auto r32 = mx::astype(mx::reshape(ms[i], mkshape({H_})), mx::float32);
+      if (!to_host32(r32, mv, err)) return false;
+      probe->m_in.push_back(std::move(mv));
+    }
+  }
+  return true;
+}
+
+void Net35::reset_graph_caches() const {
+  const size_t NL = layers_.size();
+  m_in_cache_.assign(NL, mx::array(0.f));
+  oh_cur_.assign(NL, mx::array(0.f));
+  oh_prev_.assign(NL, mx::array(0.f));
+  pred_inds_gpu_.assign(NL, mx::array(0.f));
+  pred_scores_gpu_.assign(NL, mx::array(0.f));
+  for (size_t li = 0; li < layers_.size(); li++) io_->reset_layer(static_cast<int>(li));
+  step_bundles_.clear();
+}
+
+void Net35::prefetch(const Step &plan) {
+  if (pr_ && !pr_->owners().empty()) return;  // qwen.py:195(pg_stager is not None)
+  std::map<int, std::vector<int64_t>> by_layer;
+  for (auto &[li, e] : plan.experts) by_layer[li].push_back(e);
+  for (auto &[li, es] : by_layer) io_->prefetch(li, es);
+}
+
+void Net35::shard_warm(const std::vector<std::string> &shard_names, bool seq) {
+  for (auto &sn : shard_names) {
+    io_->shard_advise_whole(sn);
+    if (seq) io_->shard_seq_read(sn);
+  }
+}
+
+bool Net35::prefill(const std::vector<int32_t> &tokens, State35 &st,
+                    std::vector<float> &logits_last, std::vector<RouteRec> &recs,
+                    std::string &err, PrefillProbe *probe) {
+  if (exec_mode() == ExecMode::Batched)
+    return prefill_batched_(tokens, st, logits_last, recs, err, probe);
+  try {
+    std::vector<float> xh;
+    if (!embed_lookup(embed_, tokens, H_, V_, xh, err)) return false;
+    const int64_t T = static_cast<int64_t>(tokens.size());
+    if (probe) probe->embed_last.assign(xh.begin() + (T - 1) * H_, xh.end());
+    mx::array x = mx::astype(mx::reshape(wrap_f32(xh, mkshape({1, T, H_})), mkshape({1, T, H_})),
+                             mx::bfloat16);
+    recs.clear();
+    for (size_t li = 0; li < layers_.size(); li++) {
+      RouteRec rec;
+      if (!layer_fwd(layers_[li], static_cast<int>(li), x, st, false, rec, nullptr, err))
+        return false;
+      if (probe) {
+        mx::array tail = mx::astype(
+            mx::reshape(mx::slice(x, mkshape({0, T - 1, 0}), mkshape({1, T, H_})),
+                        mkshape({1, H_})),
+            mx::float32);
+        std::vector<float> hv;
+        if (!to_host32(tail, hv, err)) return false;
+        probe->per_block.push_back(std::move(hv));
+      }
+      recs.push_back(std::move(rec));
+    }
+    mx::array xs = e0_rms_norm(
+        x, mx::astype(wrap_f32(final_norm_w_, mkshape({H_})), mx::bfloat16), eps_);
+    mx::array last = mx::reshape(mx::slice(xs, mkshape({0, T - 1, 0}), mkshape({1, T, H_})),
+                                 mkshape({1, H_}));
+    if (!to_host32(lm_head_.apply(last), logits_last, err)) return false;  // [V]
+    st.position += T;
+  } catch (const std::exception &ex) {
+    err = std::string("prefill: ") + ex.what();
+    return false;
+  }
+  return true;
+}
+
+bool Net35::prefill_batched_(const std::vector<int32_t> &tokens, State35 &st,
+                             std::vector<float> &logits_last, std::vector<RouteRec> &recs,
+                             std::string &err, PrefillProbe *probe) {
+  try {
+    const int64_t total = static_cast<int64_t>(tokens.size());
+    const int64_t chunk = reg_.prefill_chunk > 0 ? reg_.prefill_chunk : total;  // base.py:39
+    recs.clear();
+    std::vector<float> lg;
+    for (int64_t start = 0; start < total; start += chunk) {  // base.py:80-84
+      const int64_t n = std::min(chunk, total - start);
+      const bool last = start + n >= total;
+      step_bundles_.clear();
+      if (!forward_chunk_batched_(tokens.data() + start, n, st, lg, last, recs, err,
+                                  last ? probe : nullptr))
+        return false;
+      st.position += n;  // base.py:83 self.pos += len(chunk)
+      if (!stage_all_gpu(st, err, nullptr)) return false;
+    }
+    prefill_end_(st);
+    logits_last = std::move(lg);
+    return true;
+  } catch (const std::exception &ex) {
+    err = std::string("prefill: ") + ex.what();
+    return false;
+  }
+}
+
+bool Net35::forward_chunk_batched_(const int32_t *tok, int64_t T, State35 &st,
+                                   std::vector<float> &logits_out, bool to_host,
+                                   std::vector<RouteRec> &recs, std::string &err,
+                                   PrefillProbe *probe) {
+  std::vector<float> xh;
+  if (!embed_lookup(embed_, std::vector<int32_t>(tok, tok + T), H_, V_, xh, err)) return false;
+  if (probe) probe->embed_last.assign(xh.begin() + (T - 1) * H_, xh.end());
+  mx::array x = mx::astype(mx::reshape(wrap_f32(xh, mkshape({1, T, H_})), mkshape({1, T, H_})),
+                           mx::bfloat16);
+  const bool cb = T > 1 && reg_.prefill_hot > 0 && !prefill_ondemand();
+  std::vector<RouteRec> local_recs;
+  for (size_t li = 0; li < layers_.size(); li++) {
+    if (cb) prefill_before_layer(static_cast<int>(li));  // qwen3_5.py:285-287 before_layer_cb
+    RouteRec rec;
+    if (!layer_fwd(layers_[li], static_cast<int>(li), x, st, false, rec, nullptr, err))
+      return false;
+    if (probe) {
+      mx::array tail = mx::astype(
+          mx::reshape(mx::slice(x, mkshape({0, T - 1, 0}), mkshape({1, T, H_})),
+                      mkshape({1, H_})),
+          mx::float32);
+      std::vector<float> hv;
+      if (!to_host32(tail, hv, err)) return false;
+      probe->per_block.push_back(std::move(hv));
+    }
+    local_recs.push_back(std::move(rec));
+  }
+  mx::array xs = e0_rms_norm(
+      x, mx::astype(wrap_f32(final_norm_w_, mkshape({H_})), mx::bfloat16), eps_);
+  mx::array lastrow =
+      mx::reshape(mx::slice(xs, mkshape({0, T - 1, 0}), mkshape({1, T, H_})), mkshape({1, H_}));
+  mx::array lg = lm_head_.apply(lastrow);  // [1,V](qwen.py:177)
+  if (to_host) {
+    if (!to_host32(lg, logits_out, err)) return false;
+    recs = std::move(local_recs);
+  } else {
+    mx::eval(lg);
+  }
+  return true;
+}
+
+void Net35::prefill_before_layer(int li) {
+  const int w = std::max(1, reg_.hot_window);  // hooks.py:31
+  const int ahead = std::max(1, w / 2);        // :32
+  io_->clear_full_layer(li - 1);
+  if (reg_.prefill_hot) {                       // :40
+    const int nL = static_cast<int>(layers_.size());
+    for (int lj = li; lj < std::min(li + ahead + 1, nL); lj++) {
+      io_->load_hot_layer(lj, reg_.prefill_hot);
+      io_->materialize_hot(lj);
+    }
+    for (int lj = 0; lj < std::max(0, li - (w - ahead)); lj++)
+      io_->dematerialize_hot(lj);
+    return;
+  }
+}
+
+// prefill_end_ = qwen.py:233-250 `_prefill_end`.
+void Net35::prefill_end_(State35 &st) {
+  for (size_t li = 0; li < layers_.size(); li++) io_->clear_full_layer(static_cast<int>(li));
+  for (size_t li = 0; li < layers_.size(); li++) io_->refresh_hot_pins(static_cast<int>(li));
+  for (int li = reg_.first_head_consumer; li <= reg_.last_head_consumer; li++)
+    io_->prefetch_from_prefill(li);
+  (void)st;
+}
+
+bool Net35::moe_hot_branch_(int li, const mx::array &m_in, int64_t T,
+                            const std::vector<int64_t> &flat, mx::array &routed,
+                            std::string &err) {
+  const int64_t k = reg_.top_k, TK = static_cast<int64_t>(flat.size());
+  const auto &key = io_->hot_key(li);         // _hot_key(1132)
+  const auto &wst = io_->hot_weights(li);     // 9×[n+1,...](1131)
+  if ((int)wst.size() != 9) { err = "hot window not materialized L" + std::to_string(li); return false; }
+  std::map<int64_t, int32_t> pos;             // :1132 pos={e:i}
+  for (size_t j = 0; j < key.size(); j++) pos[key[j]] = (int32_t)j;
+  const int32_t n_rows = (int32_t)key.size();  // :1133
+  std::vector<int32_t> g(TK);
+  std::vector<int64_t> miss_rows;              // :1143
+  for (int64_t i = 0; i < TK; i++) {
+    auto it = pos.find(flat[i]);
+    if (it == pos.end()) { g[i] = n_rows; miss_rows.push_back(i); }
+    else g[i] = it->second;
+  }
+  auto remapped = mx::array(g.begin(), mkshape({T, k}), mx::int32);  // :1151-1152
+  auto x_flat = mx::reshape(m_in, mkshape({T, H_}));                 // :1136
+  auto xe = mx::reshape(x_flat, mkshape({T, 1, 1, H_}));             // :1153 expand_dims((-2,-3))
+  auto flatr = mx::reshape(remapped, mkshape({TK}));
+  auto order = mx::argsort(flatr, -1);                               // _gather_sort 1154
+  auto inv_order = mx::argsort(order, -1);
+  auto xs = mx::flatten(xe, 0, -3);                                  // [T,1,H]
+  auto tok_rows = mx::take(xs, mx::floor_divide(order, mx::array((uint32_t)k, mx::uint32)), 0);
+  auto local = mx::take(flatr, order, 0);                            // indices[order]
+  const int F = static_cast<int>(moe_ffn_);
+  auto gq = [&](const mx::array &xx, int o) {  // _gather_gate_up/down unfused+sorted(1095-1104,1156-1161)
+    return mx::gather_qmm(xx, wst[o], wst[o + 1], wst[o + 2], std::nullopt, local, true, 64, 4,
+                          "affine", true);
+  };
+  auto x_up = gq(tok_rows, 0);
+  auto x_gate = gq(tok_rows, 3);
+  auto yh = (x_gate * mx::sigmoid(x_gate)) * x_up;                   // _swiglu(1157;layer.py:57-58)
+  auto z = mx::gather_qmm(yh, wst[6], wst[7], wst[8], std::nullopt, local, true, 64, 4,
+                          "affine", true);
+  auto out = mx::take(z, inv_order, 0);                              // _scatter_unsort(1162)
+  if (!miss_rows.empty()) {
+    std::map<int64_t, std::vector<int64_t>> by_expert;
+    for (auto i : miss_rows) by_expert[flat[i]].push_back(i);
+    auto &cx = exec_counters();
+    cx.qmm_calls.fetch_add(3 * (uint64_t)by_expert.size(), std::memory_order_relaxed);
+    cx.qmm_experts.fetch_add(3 * (uint64_t)by_expert.size(), std::memory_order_relaxed);
+    cx.leases.fetch_add(by_expert.size(), std::memory_order_relaxed);
+    for (auto &[e, tk] : by_expert) {
+      auto b = io_->build(li, e);
+      if (!b) { err = "hot miss build L" + std::to_string(li); return false; }
+      std::vector<int32_t> rows, fpos;
+      for (auto i : tk) { rows.push_back((int32_t)(i / k)); fpos.push_back((int32_t)i); }  // :1169,1217-1218
+      const int64_t m = (int64_t)rows.size();
+      auto rarr = mx::array(rows.begin(), mkshape({m}), mx::int32);
+      auto xe_m = mx::reshape(mx::take(x_flat, rarr, 0), mkshape({m, 1, 1, H_}));  // :1176-1177
+      auto one = mx::zeros(mkshape({1}), mx::int32);                                // :1182
+      auto e1 = [&](int p) {  // _e1 view-like reshape(:1180-1181)
+        return mx::reshape(b->at(p), mkshape({1, b->at(p).shape(0), b->at(p).shape(1)}));
+      };
+      auto gqm = [&](const mx::array &xx, int o) {
+        return mx::gather_qmm(xx, e1(o), e1(o + 1), e1(o + 2), std::nullopt, one, true, 64, 4,
+                              "affine", true);
+      };
+      auto gu = gqm(xe_m, 0);
+      auto gg = gqm(xe_m, 3);
+      auto zm = mx::gather_qmm((gg * mx::sigmoid(gg)) * gu, e1(6), e1(7), e1(8), std::nullopt,
+                               one, true, 64, 4, "affine", true);                    // :1207-1215
+      auto ze = mx::squeeze(mx::squeeze(zm, -2), -2);                                // :1216 [m,H]
+      auto farr = mx::array(fpos.begin(), mkshape({m}), mx::int32);
+      out = mx::scatter_add(
+          out, {farr}, mx::reshape(ze, mkshape({m, 1, 1, H_})), {0});
+    }
+  }
+  routed = mx::reshape(out, mkshape({1, T, k, H_}));
+  auto &cx = exec_counters();
+  cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+  cx.qmm_experts.fetch_add(3 * (uint64_t)std::min<int64_t>(n_rows, TK),
+                           std::memory_order_relaxed);
+  cx.qmm_weight_bytes.fetch_add(
+      (uint64_t)std::min<int64_t>(n_rows, TK) *
+          (2 * (F * (H_ / 8) * 4 + 2 * F * (H_ / 64) * 2) +
+           (H_ * (F / 8) * 4 + 2 * H_ * (F / 64) * 2)),
+      std::memory_order_relaxed);
+  return true;
+}
+
+bool Net35::step(int32_t token, State35 &st, std::vector<float> &logits,
+                 std::vector<RouteRec> &recs, std::string &err, DecodeProbe *probe) {
+  try {
+    step_bundles_.clear();
+    std::vector<float> xh;
+    if (!embed_lookup(embed_, {token}, H_, V_, xh, err)) return false;
+    mx::array x = mx::astype(mx::reshape(wrap_f32(xh, mkshape({1, 1, H_})), mkshape({1, 1, H_})),
+                             mx::bfloat16);
+    recs.clear();
+    for (size_t li = 0; li < layers_.size(); li++) {
+      RouteRec rec;
+      if (!layer_fwd(layers_[li], static_cast<int>(li), x, st, true, rec, nullptr, err))
+        return false;
+      recs.push_back(std::move(rec));
+    }
+    mx::array xs = e0_rms_norm(
+        x, mx::astype(wrap_f32(final_norm_w_, mkshape({H_})), mx::bfloat16), eps_);
+    if (!to_host32(lm_head_.apply(mx::reshape(xs, mkshape({1, H_}))), logits, err)) return false;
+    st.position += 1;
+    if (exec_mode() == ExecMode::Batched) {
+      if (!stage_all_gpu(st, err, probe)) return false;
+    } else if (probe) {
+      probe->owners.clear(); probe->m_in.clear(); probe->logits.clear();
+      for (int h = 0; h < pr_->n_heads(); h++) {
+        int own = pr_->owner(h);
+        if (!st.pr.cap_hidden[own].empty()) {
+          probe->owners.push_back(own);
+          probe->m_in.push_back(st.pr.cap_hidden[own]);
+        }
+      }
+      std::vector<std::vector<double>> diag;
+      if (!st.pr.run_step_boundary(*pr_, &diag, err)) return false;
+      probe->logits = std::move(diag);
+    } else if (!st.pr.run_step_boundary(*pr_, nullptr, err)) return false;
+  } catch (const std::exception &ex) {
+    err = std::string("step: ") + ex.what();
+    return false;
+  }
+  return true;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/net35.hpp b/macos/crates/engine-native/src/net35.hpp
new file mode 100644
index 0000000..fee1cf2
--- /dev/null
+++ b/macos/crates/engine-native/src/net35.hpp
@@ -0,0 +1,146 @@
+// 35b (Qwen3.5-MoE) forward graph: mixed GatedDeltaNet and full attention.
+#pragma once
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <mlx/mlx.h>
+
+#include "expertio.hpp"
+#include "expertmap.hpp"
+#include "fwd_common.hpp"
+#include "hotstack.hpp"
+#include "lora.hpp"
+#include "provider.hpp"
+#include "prerouter.hpp"
+#include "registry.hpp"
+#include "router.hpp"
+#include "safio.hpp"
+
+namespace e0n {
+
+struct RouteRec {
+  int layer = -1;
+  bool predicted = false;  // false = live router (incl. pos-0 fallback and routed layers)
+  Selection sel;
+};
+
+struct DecodeProbe {
+  std::vector<int> owners;
+  std::vector<std::vector<float>> m_in;
+  std::vector<std::vector<double>> logits;
+};
+
+struct PrefillProbe {
+  std::vector<std::vector<float>> per_block;  // [num_layers][H], block output (residual sum)
+  std::vector<float> embed_last;              // layer-0 input (last embedding position)
+};
+
+struct State35 {
+  PrerouterState pr;
+  std::vector<mlx::core::array> conv;
+  std::vector<mlx::core::array> recur;
+  std::vector<mlx::core::array> kc;
+  std::vector<mlx::core::array> vc;
+  std::vector<char> is_linear;
+  int64_t position = 0;
+  bool started = false;
+};
+
+class Net35 {
+ public:
+  // lo = unmerged LoRA (nullptr = base). Loading is owned by the net (ExpertIo).
+  static std::unique_ptr<Net35> create(std::shared_ptr<const ModelFiles> files,
+                                       std::shared_ptr<const ExpertMap> emap, std::string &err,
+                                       std::shared_ptr<const lora::Bundle> lo = nullptr);
+
+  const TierRouting &reg() const { return reg_; }
+  Prerouter &prerouter() { return *pr_; }
+  const lora::Bundle *lora() const { return lo_.get(); }
+  uint64_t dense_resident_bytes() const { return dense_bytes_; }
+  void prefetch(const Step &plan);
+  ExpertIo::Metrics io_metrics() const { return io_->snapshot(); }
+  void shard_warm(const std::vector<std::string> &shard_names, bool seq);
+
+  void init_state(State35 &st) const;
+  void reset_state(State35 &st) const;
+
+  // Batched: chunked (2048) graph, hot-32 stack. Returns last-position logits
+  // (host f32 [V]) + last-chunk per-layer routing. Does not seed decode; first
+  // decode uses the live router. Consumer layers are warmed via prefetch_from_prefill.
+  bool prefill(const std::vector<int32_t> &tokens, State35 &st, std::vector<float> &logits_last,
+               std::vector<RouteRec> &recs, std::string &err, PrefillProbe *probe = nullptr);
+  // Decode T=1: previous token → logits [V] + routing; runs run_step_boundary at step end.
+  bool step(int32_t token, State35 &st, std::vector<float> &logits, std::vector<RouteRec> &recs,
+            std::string &err, DecodeProbe *probe = nullptr);
+
+ private:
+  Net35() = default;
+  struct Layer {
+    bool is_linear = true;
+    std::vector<float> in_ln_w, post_ln_w;
+    QuantLin in_proj_qkv, in_proj_z, in_proj_b, in_proj_a, out_proj;
+    std::vector<uint8_t> conv_wb;
+    mlx::core::array conv_w{0.f};
+    std::vector<float> A_log, dt_bias, gnorm_w;
+    QuantLin q_proj, k_proj, v_proj, o_proj;
+    std::vector<float> qn_w, kn_w;
+    std::vector<float> gate_h;
+    QuantLin gate_q;
+    QuantLin sh_gate, sh_up, sh_down, sh_egate;
+  };
+
+  bool layer_fwd(Layer &L, int li, mlx::core::array &x, State35 &st, bool decode,
+                 RouteRec &rec, mlx::core::array *m_in_out, std::string &err);
+  bool moe_host(int li, const std::vector<float> &m_in, int T, const Selection &sel,
+                std::vector<float> &out, std::string &err);
+
+  mutable std::shared_ptr<const ModelFiles> files_;
+  std::shared_ptr<const ExpertMap> emap_;
+  std::unique_ptr<ExpertIo> io_;
+  TierRouting reg_;
+  std::vector<Layer> layers_;
+  std::unique_ptr<Prerouter> pr_;
+  std::shared_ptr<const lora::Bundle> lo_;
+  mx::fast::CustomKernelFunction gdn_kernel_;
+
+  int64_t H_ = 0, V_ = 0, E_ = 0;
+  int64_t hk_ = 0, hv_ = 0, dk_ = 0, dv_ = 0, conv_k_ = 0, conv_dim_ = 0;
+  int64_t nq_ = 0, nkv_ = 0, hd_ = 0;
+  int64_t moe_ffn_ = 0, sh_ffn_ = 0;
+  double eps_ = 1e-6, rope_base_ = 100000.0, full_interval_ = 4.0, partial_rot_ = 0.25;
+  bool attn_gate_ = true;
+  double clip_ = 1000.0;
+
+  QuantLin embed_;
+  QuantLin lm_head_;
+  std::vector<float> final_norm_w_;
+  uint64_t dense_bytes_ = 0;
+
+  bool stage_all_gpu(State35 &st, std::string &err, DecodeProbe *probe);
+  void reset_graph_caches() const;
+
+  bool prefill_batched_(const std::vector<int32_t> &tokens, State35 &st,
+                        std::vector<float> &logits_last, std::vector<RouteRec> &recs,
+                        std::string &err, PrefillProbe *probe);
+  bool forward_chunk_batched_(const int32_t *tok, int64_t T, State35 &st,
+                              std::vector<float> &logits_out, bool to_host,
+                              std::vector<RouteRec> &recs, std::string &err,
+                              PrefillProbe *probe);
+  void prefill_before_layer(int li);
+  void prefill_end_(State35 &st);
+  bool moe_hot_branch_(int li, const mlx::core::array &m_in, int64_t T,
+                       const std::vector<int64_t> &flat, mlx::core::array &routed,
+                       std::string &err);
+
+  mutable std::vector<mlx::core::array> m_in_cache_, oh_cur_, oh_prev_;
+  mutable std::vector<mlx::core::array> pred_inds_gpu_, pred_scores_gpu_;
+  mutable std::vector<std::shared_ptr<Bundle>> step_bundles_;
+  mlx::core::array hw1_{0.f}, hw2_{0.f}, hwl_{0.f};
+  std::function<std::vector<mlx::core::array>(const std::vector<mlx::core::array> &)> moe_math_;
+  std::function<std::vector<mlx::core::array>(const std::vector<mlx::core::array> &)>
+      moe_math_sorted_;
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/net8b.cpp b/macos/crates/engine-native/src/net8b.cpp
new file mode 100644
index 0000000..8ed310a
--- /dev/null
+++ b/macos/crates/engine-native/src/net8b.cpp
@@ -0,0 +1,1342 @@
+#include "net8b.hpp"
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <functional>
+#include <limits>
+#include <numeric>
+#include <random>
+#include <string>
+
+#include <mlx/linalg.h>
+
+#include "dense.hpp"
+#include "qgemm.hpp"
+
+namespace e0n {
+namespace mx = mlx::core;
+
+namespace {
+
+using S = mx::Shape;
+
+mx::array conv_from_torch(std::vector<uint8_t> &bytes, int64_t C, int64_t K) {
+  mx::array w(bytes.data(), mkshape({C, 1, K}), mx::bfloat16, [](void *) {});
+  return mx::swapaxes(w, 1, 2);
+}
+
+mx::array first_k(const mx::array &a, int64_t k) {
+  mx::Shape s1 = a.shape();
+  mx::Shape s0 = a.shape();
+  for (auto &v : s0) v = 0;
+  s1[s1.size() - 1] = static_cast<mx::Shape::value_type>(k);
+  return mx::slice(a, s0, s1);
+}
+
+// moe/routing.py group_select_from_logits(30-70) ≡ bailing_hybrid.py BailingGate.__call__
+void select_group_gpu(const mx::array &logits, const mx::array *bias, int64_t k, int64_t n_group,
+                      int64_t topk_group, double scaling, bool norm, mx::array &idx,
+                      mx::array &w) {
+  auto scores = mx::sigmoid(mx::astype(logits, mx::float32));  // routing.py:41
+  mx::array select = scores;                                   // 42
+  if (bias) select = scores + *bias;
+  const int64_t k_drop = n_group - topk_group;                 // 47
+  if (k_drop > 0) {
+    mx::Shape g = scores.shape();                              // 49-50 (*b_shape, n_group, E/ng)
+    g[g.size() - 1] =
+        static_cast<mx::Shape::value_type>(g[g.size() - 1] / n_group);
+    g.insert(g.end() - 1, static_cast<mx::Shape::value_type>(n_group));
+    auto grouped = mx::reshape(select, g);
+    auto gs = mx::sum(mx::topk(grouped, 2, -1), std::vector<int>{-1}, false);
+    auto drop = first_k(mx::argpartition(gs, static_cast<int>(k_drop - 1), -1), k_drop);  // 54-55
+    auto masked =
+        mx::put_along_axis(grouped, mx::expand_dims(drop, -1),
+                           mx::array(-std::numeric_limits<float>::infinity(), mx::float32), -2);
+    select = mx::reshape(masked, scores.shape());              // 62
+  }
+  idx = first_k(mx::argpartition(-select, static_cast<int>(k - 1), -1), k);  // 65
+  w = mx::take_along_axis(scores, idx, -1);
+  if (norm) w = w / (mx::sum(w, std::vector<int>{-1}, true) + mx::array(1e-20f, mx::float32));  // 67-68
+  w = w * mx::array(static_cast<float>(scaling), mx::float32);  // 69
+}
+
+mx::array gate_logits32(const mx::array &m_in, const mx::array &gate_w) {
+  return mx::astype(mx::matmul(m_in, mx::transpose(gate_w)), mx::float32);
+}
+
+bool onehot_pos(const mx::array &idx, int64_t E, int64_t pos, mx::array &oh) {
+  if (idx.ndim() < 3) return false;                            // stager.py:45-46
+  const int64_t t = idx.shape(-2);
+  const int64_t lo = pos >= 0 ? pos : t + pos;
+  if (lo < 0 || lo >= t) return false;                         // 47-49
+  mx::Shape s0 = idx.shape(), s1 = idx.shape();
+  for (auto &v : s0) v = 0;
+  s0[s0.size() - 2] = static_cast<mx::Shape::value_type>(lo);
+  s1[s1.size() - 2] = static_cast<mx::Shape::value_type>(lo + 1);
+  auto row = mx::slice(idx, s0, s1);                           // 51 [..., 1, k]
+  auto eq = mx::equal(mx::expand_dims(row, -1), mx::arange(0, static_cast<int>(E)));  // 52-53
+  oh = mx::astype(mx::sum(eq, std::vector<int>{-2}, false), mx::float32);  // 53-54
+  return true;
+}
+
+mx::array rope_interleave_torch(const mx::array &x, int64_t Tpos_start, double theta,
+                                int64_t D) {
+  // 98: freqs = 1.0 / (theta ** (arange(0,D,2,f32)/D))
+  auto ar = mx::divide(mx::arange(0, static_cast<int>(D), 2, mx::float32),
+                       mx::array(static_cast<float>(D), mx::float32));
+  auto freqs = mx::divide(mx::array(1.0f, mx::float32),
+                          mx::power(mx::array(static_cast<float>(theta), mx::float32), ar));
+  // 99: angles = outer(positions.f32, freqs);100: emb = cat(angles,angles)
+  auto positions =
+      mx::astype(mx::add(mx::arange(0, static_cast<int>(x.shape(2)), mx::int32),
+                         mx::array(static_cast<int32_t>(Tpos_start), mx::int32)),
+                 mx::float32);
+  auto angles = mx::outer(positions, freqs);                   // [L,D/2]
+  auto emb = mx::concatenate({angles, angles}, -1);            // 100 [L,D]
+  auto cos = mx::cos(emb);                                     // 101
+  auto sin = mx::sin(emb);
+  // 102-103: xi = x.reshape(B,H,L,D/2,2).transpose(0,1,2,4,3).reshape(B,H,L,D)
+  const int64_t B = x.shape(0), H = x.shape(1), Ld = x.shape(2);
+  auto xi = mx::reshape(
+      mx::transpose(mx::reshape(x, mkshape({B, H, Ld, D / 2, 2})), {0, 1, 2, 4, 3}),
+      mkshape({B, H, Ld, D}));
+  // 105: rh = cat(-xi[...,D/2:], xi[...,:D/2])
+  mx::Shape s0 = xi.shape(), s1 = xi.shape();
+  for (auto &v : s0) v = 0;
+  s1[3] = static_cast<mx::Shape::value_type>(D / 2);
+  auto lo = mx::slice(xi, s0, s1);                 // xi[..., :h]
+  s0[3] = static_cast<mx::Shape::value_type>(D / 2);
+  s1[3] = static_cast<mx::Shape::value_type>(D);
+  auto hi = mx::slice(xi, s0, s1);                 // xi[..., h:]
+  auto rh = mx::concatenate({mx::negative(hi), lo}, -1);
+  auto cosb = mx::reshape(cos, mkshape({1, 1, Ld, D}));
+  auto sinb = mx::reshape(sin, mkshape({1, 1, Ld, D}));
+  return mx::add(mx::multiply(xi, cosb), mx::multiply(rh, sinb));
+}
+
+std::vector<mx::array> moe_math_body(const std::vector<mx::array> &a) {
+  auto xe = mx::reshape(a[0], mkshape({a[0].shape(0), a[0].shape(1), 1, 1, a[0].shape(2)}));  // layer.py:159
+  auto gq = [&](const mx::array &xx, int o) {                  // 160-167/169-171
+    return mx::gather_qmm(xx, a[o], a[o + 1], a[o + 2], std::nullopt, a[10], true, 64, 4,
+                          "affine", false);
+  };
+  auto x_up = gq(xe, 1);
+  auto x_gate = gq(xe, 4);
+  auto yh = (x_gate * mx::sigmoid(x_gate)) * x_up;             // _swiglu(x_up, x_gate)
+  auto z = gq(yh, 7);
+  return {mx::squeeze(z, -2)};                                 // 172
+}
+
+std::vector<mx::array> moe_math_sorted_body(const std::vector<mx::array> &a) {
+  auto gq = [&](const mx::array &xx, int o) {                  // 181-192
+    return mx::gather_qmm(xx, a[o], a[o + 1], a[o + 2], std::nullopt, a[10], true, 64, 4,
+                          "affine", true);
+  };
+  auto x_up = gq(a[0], 1);
+  auto x_gate = gq(a[0], 4);
+  auto yh = (x_gate * mx::sigmoid(x_gate)) * x_up;             // _swiglu
+  return {gq(yh, 7)};
+}
+
+ExpertPtrs bundle_ptrs(const Bundle &b) {
+  ExpertPtrs p;
+  p.uw = b.at(0).data<uint32_t>();
+  p.us = reinterpret_cast<const uint16_t *>(b.at(1).data<mx::bfloat16_t>());
+  p.ub = reinterpret_cast<const uint16_t *>(b.at(2).data<mx::bfloat16_t>());
+  p.gw = b.at(3).data<uint32_t>();
+  p.gs = reinterpret_cast<const uint16_t *>(b.at(4).data<mx::bfloat16_t>());
+  p.gb = reinterpret_cast<const uint16_t *>(b.at(5).data<mx::bfloat16_t>());
+  p.dw = b.at(6).data<uint32_t>();
+  p.ds = reinterpret_cast<const uint16_t *>(b.at(7).data<mx::bfloat16_t>());
+  p.db = reinterpret_cast<const uint16_t *>(b.at(8).data<mx::bfloat16_t>());
+  return p;
+}
+
+}  // namespace
+
+std::unique_ptr<Net8b> Net8b::create(std::shared_ptr<const ModelFiles> files,
+                                     std::shared_ptr<const ExpertMap> emap, std::string &err,
+                                     std::shared_ptr<const lora::Bundle> lo) {
+  std::unique_ptr<Net8b> n(new Net8b());
+  n->lo_ = lo;
+  n->files_ = files;
+  n->emap_ = emap;
+  if (!tier_routing("edge0-8b", n->reg_)) { err = "registry missing edge0-8b"; return nullptr; }
+  int64_t cfg_layers = 0;
+  if (files->config_int("num_hidden_layers", cfg_layers) && cfg_layers != n->reg_.num_layers) {
+    err = "config layer count does not match registry";
+    return nullptr;
+  }
+  auto ci = [&](const char *k, int64_t def, int64_t &o) {
+    int64_t v;
+    o = files->config_int(k, v) ? v : def;
+  };
+  auto cd = [&](const char *k, double def, double &o) {
+    double v;
+    o = files->config_double(k, v) ? v : def;
+  };
+  ci("hidden_size", 1536, n->H_);
+  ci("vocab_size", 157184, n->V_);
+  ci("num_attention_heads", 16, n->nh_);
+  ci("head_dim", 128, n->hd_);
+  ci("short_conv_kernel_size", 4, n->conv_k_);
+  ci("moe_intermediate_size", 512, n->moe_ffn_);
+  {
+    int64_t shi = 512, shn = 1;
+    ci("moe_shared_expert_intermediate_size", 512, shi);
+    ci("num_shared_experts", 1, shn);
+    n->sh_ffn_ = shi * shn;
+  }
+  ci("intermediate_size", 4608, n->dense_ffn_);
+  ci("q_lora_rank", 256, n->qlora_);
+  ci("kv_lora_rank", 512, n->kvlora_);
+  ci("qk_nope_head_dim", 128, n->nope_);
+  ci("qk_rope_head_dim", 64, n->roped_);
+  ci("v_head_dim", 128, n->vhd_);
+  cd("rms_norm_eps", 1e-6, n->eps_);
+  cd("rope_theta", 6000000.0, n->rope_theta_);
+  cd("kda_lower_bound", -5.0, n->kda_lb_);
+  n->proj_ = n->nh_ * n->hd_;
+  int64_t qbits = 4, qgroup = 64;
+  if (files->config_quant(qbits, qgroup) && qgroup > 0) n->group_ = qgroup;
+
+  int64_t interval = 4, dense_prefix_cfg = 1;
+  ci("layer_group_size", 4, interval);
+  ci("first_k_dense_replace", 1, dense_prefix_cfg);
+  if (dense_prefix_cfg != n->reg_.dense_prefix) { err = "config dense prefix does not match registry"; return nullptr; }
+
+  const int64_t NL = n->reg_.num_layers;
+  {
+    static const std::string kKdaSource = R"gdn(
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }
+
+        // g: [B, T, Hv, Dk]
+        auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;
+        auto beta_ = beta + b_idx * T * Hv;
+
+        for (int t = 0; t < T; ++t) {
+          if (true) {
+            float kv_mem = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] * g_[s_idx];
+              kv_mem += state[i] * k_[s_idx];
+            }
+            kv_mem = simd_sum(kv_mem);
+
+            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+
+            float out = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] + k_[s_idx] * delta;
+              out += state[i] * q_[s_idx];
+            }
+            out = simd_sum(out);
+            if (thread_index_in_simdgroup == 0) {
+              y[dv_idx] = static_cast<InT>(out);
+            }
+          }
+          // Increment data pointers to next time step
+          q_ += Hk * Dk;
+          k_ += Hk * Dk;
+          v_ += Hv * Dv;
+          y += Hv * Dv;
+          g_ += Hv * Dk;
+          beta_ += Hv;
+        }
+        for (int i = 0; i < n_per_t; ++i) {
+          auto s_idx = n_per_t * dk_idx + i;
+          o_state[s_idx] = static_cast<InT>(state[i]);
+        }
+    )gdn";
+    n->kda_kernel_ = mx::fast::metal_kernel(
+        "gated_delta_step_vec", {"q", "k", "v", "g", "beta", "state_in", "T"},
+        {"y", "state_out"}, kKdaSource);
+  }
+  n->layers_.resize(static_cast<size_t>(NL));
+  const std::string M = "model.layers.";
+  for (int64_t li = 0; li < NL; li++) {
+    Layer &L = n->layers_[static_cast<size_t>(li)];
+    const std::string P = M + std::to_string(li);
+    L.is_mla = ((li + 1) % interval) == 0 || li >= NL / interval * interval;
+    L.has_moe = li >= n->reg_.dense_prefix;
+    auto q = [&](const std::string &name, QuantLin &dst) {
+      if (!dst.load(*files, P + name, err)) { err += " @" + std::to_string(li); return false; }
+      n->dense_bytes_ += dst.wb.size() + dst.sb.size() + dst.bb.size();
+      return true;
+    };
+    auto vec = [&](const std::string &name, std::vector<float> &dstv) {
+      DenseF32 d;
+      if (!read_dense_f32(*files, P + name, d, err)) { err += " @" + std::to_string(li); return false; }
+      dstv = std::move(d.v);
+      n->dense_bytes_ += dstv.size() * 4;
+      return true;
+    };
+    if (!vec(".input_layernorm.weight", L.in_ln_w) ||
+        !vec(".post_attention_layernorm.weight", L.post_ln_w))
+      return nullptr;
+    if (!L.is_mla) {
+      if (!q(".attention.q_proj.weight", L.k_q) || !q(".attention.k_proj.weight", L.k_k) ||
+          !q(".attention.v_proj.weight", L.k_v) || !q(".attention.f_proj.weight", L.k_f) ||
+          !q(".attention.g_proj.weight", L.k_g) || !q(".attention.b_proj.weight", L.k_b) ||
+          !q(".attention.o_proj.weight", L.k_o) ||
+          !vec(".attention.A_log", L.A_log) || !vec(".attention.dt_bias", L.dt_bias) ||
+          !vec(".attention.o_norm.weight", L.onorm_w))
+        return nullptr;
+      if (static_cast<int64_t>(L.A_log.size()) != n->nh_ ||
+          static_cast<int64_t>(L.dt_bias.size()) != n->proj_ ||
+          static_cast<int64_t>(L.onorm_w.size()) != n->hd_) {
+        err = "KDA parameter vector length mismatch";
+        return nullptr;
+      }
+      if (!files->read_tensor(P + ".attention.q_conv1d.weight", L.conv_qb, err) ||
+          !files->read_tensor(P + ".attention.k_conv1d.weight", L.conv_kb, err) ||
+          !files->read_tensor(P + ".attention.v_conv1d.weight", L.conv_vb, err))
+        return nullptr;
+      auto ct = files->find(P + ".attention.q_conv1d.weight")->shape;
+      if (ct.size() != 3 || ct[0] != n->proj_ || ct[2] != n->conv_k_) {
+        err = "KDA conv shape is not [proj,1,ksize]";
+        return nullptr;
+      }
+      L.conv_q = conv_from_torch(L.conv_qb, n->proj_, n->conv_k_);
+      L.conv_k = conv_from_torch(L.conv_kb, n->proj_, n->conv_k_);
+      L.conv_v = conv_from_torch(L.conv_vb, n->proj_, n->conv_k_);
+      n->dense_bytes_ += L.conv_qb.size() + L.conv_kb.size() + L.conv_vb.size();
+    } else {
+      if (!q(".attention.q_a_proj.weight", L.m_qa) || !q(".attention.q_b_proj.weight", L.m_qb) ||
+          !q(".attention.kv_a_proj_with_mqa.weight", L.m_kva) ||
+          !q(".attention.kv_b_proj.weight", L.m_kvb) || !q(".attention.g_proj.weight", L.m_g) ||
+          !q(".attention.dense.weight", L.m_dense) ||
+          !vec(".attention.q_a_layernorm.weight", L.m_qaln) ||
+          !vec(".attention.kv_a_layernorm.weight", L.m_kvaln))
+        return nullptr;
+      if (static_cast<int64_t>(L.m_qaln.size()) != n->qlora_ ||
+          static_cast<int64_t>(L.m_kvaln.size()) != n->kvlora_) {
+        err = "MLA LoRA normalize length mismatch";
+        return nullptr;
+      }
+    }
+    if (!L.has_moe) {
+      if (!q(".mlp.gate_proj.weight", L.d_gate) || !q(".mlp.up_proj.weight", L.d_up) ||
+          !q(".mlp.down_proj.weight", L.d_down))
+        return nullptr;
+      continue;
+    }
+    DenseF32 g;
+    const std::string gw = P + ".mlp.gate.weight";
+    const TensorInfo *gi = files->find(gw);
+    if (!gi) { err = "missing gate"; return nullptr; }
+    if (gi->dtype == "U32") {
+      int64_t bits2 = 4, gs2 = 64;
+      if (!files->config_quant_for(gw, bits2, gs2) ||
+          !read_dense_dequant(*files, gw, bits2, gs2, g, err))
+        return nullptr;
+    } else if (!read_dense_f32(*files, gw, g, err)) {
+      return nullptr;
+    }
+    if (g.rows != n->reg_.num_experts || g.cols != n->H_) { err = "gate shape mismatch"; return nullptr; }
+    L.gate_h = std::move(g.v);
+    L.gate_bf16 = gi->dtype != "U32";
+    n->dense_bytes_ += L.gate_h.size() * 4;
+    {
+      DenseF32 eb;
+      if (!read_dense_f32(*files, P + ".mlp.gate.expert_bias", eb, err)) return nullptr;
+      if (static_cast<int64_t>(eb.v.size()) != n->reg_.num_experts) {
+        err = "expert_bias length mismatch";
+        return nullptr;
+      }
+      L.gate_bias.assign(eb.v.begin(), eb.v.end());
+      n->dense_bytes_ += L.gate_bias.size() * 4;
+    }
+    if (!q(".mlp.shared_experts.gate_proj.weight", L.sh_gate) ||
+        !q(".mlp.shared_experts.up_proj.weight", L.sh_up) ||
+        !q(".mlp.shared_experts.down_proj.weight", L.sh_down))
+      return nullptr;
+  }
+  if (!n->embed_.load(*files, "model.word_embeddings.weight", err)) return nullptr;
+  if (!n->lm_head_.load(*files, "lm_head.weight", err)) return nullptr;
+  n->dense_bytes_ += n->embed_.wb.size() + n->embed_.sb.size() + n->embed_.bb.size() +
+                     n->lm_head_.wb.size() + n->lm_head_.sb.size() + n->lm_head_.bb.size();
+  {
+    DenseF32 f;
+    if (!read_dense_f32(*files, "model.norm.weight", f, err)) return nullptr;
+    n->final_norm_w_ = std::move(f.v);
+  }
+  n->pr_ = std::make_unique<Prerouter>();
+  if (!Prerouter::load(*files, n->reg_, n->H_, *n->pr_, err)) return nullptr;
+  {
+    const int64_t NL2 = n->reg_.num_layers;
+    n->io_ = std::make_unique<ExpertIo>(files.get(), 64, 48, 8, 4);
+    for (int64_t li = 0; li < NL2; li++) {
+      if (!n->layers_[size_t(li)].has_moe) continue;
+      ExpertIo::LayerCfg c;
+      c.layer = int(li);
+      c.prefix = emap->layer_prefix(int(li));
+      c.E = n->reg_.num_experts;
+      c.H = n->H_;
+      c.F = n->moe_ffn_;
+      c.group = n->group_;
+      c.staged_n = int(n->reg_.top_k);
+      c.prefill_hot = (int)n->reg_.prefill_hot;
+      c.full_layer_prefill = n->reg_.full_layer_prefill;
+      n->io_->add_layer(c);
+    }
+  }
+  if (exec_mode() == ExecMode::Batched) {
+    const char *cl = getenv("MLX_CACHE_LIMIT_MB");
+    mx::set_cache_limit(static_cast<size_t>(cl ? atoi(cl) : 256) << 20);
+    const int64_t E = n->reg_.num_experts, NL = n->reg_.num_layers;
+    n->m_in_cache_.assign(NL, mx::array(0.f));
+    n->last_topk_.assign(NL, mx::array(0.f));
+    n->prev_topk_oh_.assign(NL, mx::array(0.f));
+    n->pg_logits_.assign(NL, mx::array(0.f));
+    n->gate_w32_.assign(NL, mx::array(0.f));
+    n->gate_b32_.assign(NL, mx::array(0.f));
+    for (int64_t li = 0; li < NL; li++) {
+      Layer &L = n->layers_[static_cast<size_t>(li)];
+      if (!L.has_moe) continue;
+      auto gw =
+          mx::array(L.gate_h.data(), mkshape({E, n->H_}), mx::float32, [](void *) {});
+      n->gate_w32_[li] = L.gate_bf16 ? mx::astype(gw, mx::bfloat16) : gw;
+      std::vector<float> b32(static_cast<size_t>(E));
+      for (int64_t e = 0; e < E; e++)
+        b32[static_cast<size_t>(e)] = static_cast<float>(L.gate_bias[static_cast<size_t>(e)]);
+      n->gate_b32_[li] = mx::array(b32.begin(), mkshape({E}), mx::float32);
+    }
+    const int64_t nH = n->pr_->n_heads(), f_in = n->pr_->stacked_f_in(),
+                  prh = n->pr_->stacked_prh(), hE = n->pr_->stacked_E();
+    if (nH > 0) {
+      auto wrap16 = [&](const std::vector<float> &v, std::initializer_list<int64_t> dims) {
+        return mx::astype(mx::array((void *)v.data(), mkshape(dims), mx::float32, [](void *) {}),
+                          mx::float16);
+      };
+      n->hw1_ = wrap16(n->pr_->stacked_fc1(), {nH, f_in, prh});
+      n->hw2_ = wrap16(n->pr_->stacked_fc2(), {nH, prh, hE});
+      n->hwl_ = wrap16(n->pr_->stacked_lin(), {nH, f_in, hE});
+      mx::eval(n->hw1_, n->hw2_, n->hwl_);
+    }
+    try {
+      auto fn = mx::compile(&moe_math_body);
+      const int64_t F = n->moe_ffn_, gsh = n->H_ / 64, gsf = F / 64;
+      std::vector<mx::array> d = {
+          mx::zeros(mkshape({1, 1, n->H_}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, F / 8}), mx::uint32),
+          mx::zeros(mkshape({2, n->H_, gsf}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, gsf}), mx::bfloat16),
+          mx::zeros(mkshape({1, 1, 2}), mx::int32)};
+      auto out = fn(d);
+      mx::eval(out[0]);
+      n->moe_math_ = std::move(fn);
+    } catch (const std::exception &ex) {
+      std::fprintf(stderr, "edge0-8b: mlx compile unavailable (%s), MoE math falls back to eager\n",
+                   ex.what());
+      n->moe_math_ = &moe_math_body;
+    }
+    try {
+      auto fn = mx::compile(&moe_math_sorted_body);
+      const int64_t F = n->moe_ffn_, gsh = n->H_ / 64, gsf = F / 64;
+      std::vector<mx::array> d = {
+          mx::zeros(mkshape({16, 1, n->H_}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, n->H_ / 8}), mx::uint32),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, F, gsh}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, F / 8}), mx::uint32),
+          mx::zeros(mkshape({2, n->H_, gsf}), mx::bfloat16),
+          mx::zeros(mkshape({2, n->H_, gsf}), mx::bfloat16),
+          mx::zeros(mkshape({16}), mx::int32)};
+      auto out = fn(d);
+      mx::eval(out[0]);
+      n->moe_math_sorted_ = std::move(fn);
+    } catch (const std::exception &ex) {
+      std::fprintf(stderr, "edge0-8b: mlx compile unavailable (%s), sorted MoE math falls back to eager\n",
+                   ex.what());
+      n->moe_math_sorted_ = &moe_math_sorted_body;
+    }
+  }
+    {
+      const char *pw = getenv("EDGE0_PREWARM");
+      if (!pw) pw = getenv("LING_PREWARM");
+      if (pw && !strcmp(pw, "1")) {
+        auto tp0 = std::chrono::steady_clock::now();
+        try {
+          for (const auto &sn : files->shard_names()) {  // shard.advise_willneed + seq_read
+            n->io_->shard_advise_whole(sn);
+            n->io_->shard_seq_read(sn);
+          }
+          std::string perr;
+          std::vector<float> lg;
+          std::vector<RouteRec> recs;
+          State8b pst;
+          n->init_state(pst);
+          const std::vector<int32_t> dummy(4, 0);  // prefill(dummy*4)
+          n->prefill(dummy, pst, lg, recs, perr);
+          n->step(0, pst, lg, recs, perr);          // step(dummy[0])
+          n->reset_state(pst);
+          std::fprintf(stderr, "[edge0-8b] prewarm done in %.1fs\n",
+                       std::chrono::duration<double>(std::chrono::steady_clock::now() - tp0).count());
+        } catch (...) { /* advisory */ }
+      }
+    }
+  if (n->lo_) {
+    const int nl = n->reg_.num_layers, mla = nl / 4;
+    const int want = 7 * (nl - mla) + 4 * mla + 3;
+    if (n->lo_->n_bound() != want) {
+      err = "LoRA bind count " + std::to_string(n->lo_->n_bound()) + " != structural count " +
+            std::to_string(want);
+      return nullptr;
+    }
+    n->dense_bytes_ += static_cast<uint64_t>(n->lo_->bytes());
+  }
+  return n;
+}
+
+void Net8b::init_state(State8b &st) const {
+  st.pr.init(reg_);
+  const size_t NL = layers_.size();
+  st.cq.assign(NL, mx::array(0.f));
+  st.ck.assign(NL, mx::array(0.f));
+  st.cv.assign(NL, mx::array(0.f));
+  st.ssm.assign(NL, mx::array(0.f));
+  st.kc.assign(NL, mx::array(0.f));
+  st.vc.assign(NL, mx::array(0.f));
+  st.is_mla.assign(NL, 0);
+  for (size_t i = 0; i < NL; i++) st.is_mla[i] = layers_[i].is_mla ? 1 : 0;
+  st.position = 0;
+  st.started = true;
+  reset_graph_caches();
+}
+
+void Net8b::reset_state(State8b &st) const {
+  st.pr.reset();
+  const size_t NL = layers_.size();
+  st.cq.assign(NL, mx::array(0.f));
+  st.ck.assign(NL, mx::array(0.f));
+  st.cv.assign(NL, mx::array(0.f));
+  st.ssm.assign(NL, mx::array(0.f));
+  st.kc.assign(NL, mx::array(0.f));
+  st.vc.assign(NL, mx::array(0.f));
+  st.position = 0;
+  reset_graph_caches();
+}
+
+bool Net8b::layer_fwd(Layer &L, int li, mx::array &x, State8b &st, bool decode, bool prefill_seed,
+                      RouteRec &rec, std::string &err) {
+  const int64_t T = x.shape(1);
+  const auto dt = x.dtype();
+  auto norm_w = [&](const std::vector<float> &v) {
+    return mx::astype(wrap_f32(v, mkshape({static_cast<int64_t>(v.size())})), dt);
+  };
+  mx::array xn = e0_rms_norm(x, norm_w(L.in_ln_w), eps_);
+
+  mx::array r{0.f};
+  if (!L.is_mla) {
+    const auto *lo = lo_.get();
+    mx::array q_lin = e0n::lora::apply_lora(lo, L.k_q, xn, li, e0n::lora::S8B_Q),
+        k_lin = e0n::lora::apply_lora(lo, L.k_k, xn, li, e0n::lora::S8B_K),
+        v_lin = e0n::lora::apply_lora(lo, L.k_v, xn, li, e0n::lora::S8B_V);
+    auto conv_step = [&](const mx::array &lin, mx::array &state, const mx::array &w) {
+      if (state.ndim() == 0)
+        state = mx::zeros(mkshape({1, conv_k_ - 1, proj_}), dt);
+      mx::array cat = mx::concatenate({state, lin}, 1);  // [1, K-1+T, proj]
+      state = mx::slice(cat, mkshape({0, T, 0}), mkshape({1, T + conv_k_ - 1, proj_}));
+      return e0_silu(mx::conv1d(cat, w, 1, 0, 1, static_cast<int>(proj_)));  // [1,T,proj]
+    };
+    mx::array qc = conv_step(q_lin, st.cq[li], L.conv_q);
+    mx::array kc = conv_step(k_lin, st.ck[li], L.conv_k);
+    mx::array vc = conv_step(v_lin, st.cv[li], L.conv_v);
+    mx::array qf = mx::astype(mx::reshape(qc, mkshape({1, T, nh_, hd_})), mx::float32);
+    mx::array kf = mx::astype(mx::reshape(kc, mkshape({1, T, nh_, hd_})), mx::float32);
+    mx::array vf = mx::astype(mx::reshape(vc, mkshape({1, T, nh_, hd_})), mx::float32);
+    float sc = 1.0f / std::sqrt(static_cast<float>(hd_));
+    qf = sc * qf / (mx::linalg::norm(qf, 2.0, -1, true) + 1e-6f);
+    kf = kf / (mx::linalg::norm(kf, 2.0, -1, true) + 1e-6f);
+    mx::array gate = mx::reshape(
+        e0n::lora::apply_lora(lo, L.k_g, xn, li, e0n::lora::S8B_G), mkshape({1, T, nh_, hd_}));
+    mx::array f = mx::astype(mx::reshape(e0n::lora::apply_lora(lo, L.k_f, xn, li, e0n::lora::S8B_F),
+                                         mkshape({1, T, nh_, hd_})),
+                               mx::float32);
+    mx::array dtb = mx::reshape(wrap_f32(L.dt_bias, mkshape({nh_, hd_})), mkshape({1, 1, nh_, hd_}));
+    mx::array alog = wrap_f32(L.A_log, mkshape({nh_, 1}));
+    mx::array glog = static_cast<float>(kda_lb_) * mx::sigmoid(mx::exp(alog) * (f + dtb));
+    mx::array gdec = mx::exp(glog);
+    mx::array beta = mx::sigmoid(mx::astype(
+        e0n::lora::apply_lora(lo, L.k_b, xn, li, e0n::lora::S8B_B), mx::float32));  // [1,T,H]
+    if (st.ssm[li].ndim() == 0)
+      st.ssm[li] = mx::zeros(mkshape({1, nh_, hd_, hd_}), mx::float32);
+    auto kouts = kda_kernel_(
+        {qf, kf, vf, gdec, beta, st.ssm[li], mx::array((int)T, mx::int32)},
+        {mkshape({1, T, nh_, hd_}), st.ssm[li].shape()},
+        {mx::float32, mx::float32},
+        {32, static_cast<int>(hd_), static_cast<int>(nh_)},  // grid=(32, Dv, B*Hv)，B=1
+        {32, 4, 1},                                          // threadgroup(gated_delta.py:206)
+        {{"InT", mx::float32},
+         {"Dk", static_cast<int>(hd_)},
+         {"Dv", static_cast<int>(hd_)},
+         {"Hk", static_cast<int>(nh_)},
+         {"Hv", static_cast<int>(nh_)}},
+        std::nullopt, false, {});
+    mx::array y = kouts[0];  // [1,T,nh,hd] f32
+    st.ssm[li] = kouts[1];
+    mx::array gated =
+        e0_rms_norm(mx::astype(y, dt), norm_w(L.onorm_w), eps_) * mx::sigmoid(gate);
+    r = e0n::lora::apply_lora(lo, L.k_o, mx::reshape(gated, mkshape({1, T, proj_})), li,
+                              e0n::lora::S8B_O);
+  } else {
+    mx::array q_lat = e0_rms_norm(L.m_qa.apply(xn), norm_w(L.m_qaln), eps_);
+    mx::array qfull = mx::transpose(
+        mx::reshape(e0n::lora::apply_lora(lo_.get(), L.m_qb, q_lat, li, e0n::lora::S8B_M_QB),
+                    mkshape({1, T, nh_, nope_ + roped_})),
+        {0, 2, 1, 3});
+    mx::eval(qfull);
+    auto qs = mx::split(qfull, S{static_cast<int32_t>(nope_)}, -1);
+    mx::array q_nope = qs[0], q_pe = qs[1];
+    mx::array comp = L.m_kva.apply(xn);  // [1,T,kvlora+rope]
+    mx::eval(comp);
+    auto cs = mx::split(comp, S{static_cast<int32_t>(kvlora_)}, -1);
+    mx::array kv_lat = e0_rms_norm(cs[0], norm_w(L.m_kvaln), eps_);
+    mx::array kv = mx::transpose(
+        mx::reshape(
+            e0n::lora::apply_lora(lo_.get(), L.m_kvb, kv_lat, li, e0n::lora::S8B_M_KVB),
+            mkshape({1, T, nh_, nope_ + vhd_})),
+        {0, 2, 1, 3});
+    auto ks = mx::split(kv, S{static_cast<int32_t>(nope_)}, -1);
+    mx::array k_nope = ks[0], values = ks[1];
+    mx::array k_pe = mx::broadcast_to(
+        rope_interleave_torch(
+            mx::reshape(mx::slice(comp, mkshape({0, 0, kvlora_}),
+                                  mkshape({1, T, kvlora_ + roped_})),
+                        mkshape({1, 1, T, roped_})),
+            st.position, rope_theta_, roped_),
+        mkshape({1, nh_, T, roped_}));
+    q_pe = rope_interleave_torch(q_pe, st.position, rope_theta_, roped_);
+    mx::array queries = mx::concatenate({q_nope, q_pe}, -1);   // [1,nh,T,192]
+    mx::array keys = mx::concatenate({k_nope, k_pe}, -1);
+    if (st.kc[li].ndim() == 0) {
+      st.kc[li] = keys;
+      st.vc[li] = values;
+    } else {
+      st.kc[li] = mx::concatenate({st.kc[li], keys}, 2);
+      st.vc[li] = mx::concatenate({st.vc[li], values}, 2);
+    }
+    float scale = 1.0f / std::sqrt(static_cast<float>(nope_ + roped_));
+    mx::array o = mx::fast::scaled_dot_product_attention(
+        queries, st.kc[li], st.vc[li], scale, "causal");  // [1,nh,T,vhd]
+    o = mx::reshape(mx::transpose(o, {0, 2, 1, 3}), mkshape({1, T, nh_ * vhd_}));
+    mx::array gate = mx::sigmoid(
+        e0n::lora::apply_lora(lo_.get(), L.m_g, xn, li, e0n::lora::S8B_M_G));  // [1,T,nh] head-wise
+    o = mx::reshape(mx::reshape(o, mkshape({1, T, nh_, vhd_})) *
+                        mx::expand_dims(gate, -1),
+                    mkshape({1, T, nh_ * vhd_}));
+    r = e0n::lora::apply_lora(lo_.get(), L.m_dense, o, li, e0n::lora::S8B_M_DENSE);
+  }
+
+  mx::array h = x + r;
+  mx::array m_in = e0_rms_norm(h, norm_w(L.post_ln_w), eps_);
+  if (!L.has_moe) {
+    mx::array y =
+        e0n::lora::apply_lora(
+            lo_.get(), L.d_down,
+            e0_silu(e0n::lora::apply_lora(lo_.get(), L.d_gate, m_in, li, e0n::lora::S8B_D_GATE)) *
+                e0n::lora::apply_lora(lo_.get(), L.d_up, m_in, li, e0n::lora::S8B_D_UP),
+            li, e0n::lora::S8B_D_DOWN);
+    x = h + y;
+    if (clip_ > 0)
+      x = mx::where(mx::isnan(x), mx::zeros_like(x),
+                    mx::clip(x, mx::array(static_cast<float>(-clip_)),
+                             mx::array(static_cast<float>(clip_))));
+    if (!(exec_mode() == ExecMode::Batched && decode && T == 1))
+      mx::eval(x);
+    return true;
+  }
+
+  // bailing_hybrid.py BailingSparseMoE.__call__(679-695)/DecoderLayer.__call__(739-762):
+  if (exec_mode() == ExecMode::Batched && decode && T == 1) {
+    const int64_t k = reg_.top_k;
+    io_->note_call(li);
+    const bool owner = li >= reg_.first_owner && li <= reg_.last_owner;
+    const bool consumer = li >= reg_.first_head_consumer && li <= reg_.last_head_consumer;
+    if (owner) m_in_cache_[static_cast<size_t>(li)] = m_in;
+    io_->wait_staged(li);
+    const auto &st_st = io_->staged(li);
+    const bool use_pg =
+        consumer && pg_logits_[static_cast<size_t>(li)].ndim() > 0 &&
+        !st.pr.pred[static_cast<size_t>(li)].inds.empty() &&
+        st_st.active && st_st.exp == st.pr.pred[static_cast<size_t>(li)].inds;
+    mx::array idx(0.f), w(0.f), routed(0.f), teach(0.f);
+    rec.layer = li;
+    if (use_pg) {
+      select_group_gpu(pg_logits_[static_cast<size_t>(li)], nullptr, k, reg_.n_group,
+                       reg_.topk_group, reg_.routed_scaling, reg_.norm_topk_prob, idx, w);
+      if (owner) {
+        auto gl = gate_logits32(m_in, gate_w32_[static_cast<size_t>(li)]);
+        mx::array w_unused(0.f);
+        select_group_gpu(gl, &gate_b32_[static_cast<size_t>(li)], k, reg_.n_group,
+                         reg_.topk_group, reg_.routed_scaling, reg_.norm_topk_prob, teach,
+                         w_unused);
+      }
+      rec.predicted = true;
+      rec.sel = st.pr.pred[static_cast<size_t>(li)];
+      st.pr.pred_avail[static_cast<size_t>(li)] = false;
+      auto asm_ = io_->asm_for(li, st_st);
+      auto local = mx::take(asm_.slot_of, idx);
+      std::vector<mx::array> args{m_in};
+      for (auto &a : asm_.wargs) args.push_back(a);
+      args.push_back(local);
+      auto &cx = exec_counters();
+      cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+      cx.qmm_experts.fetch_add(3 * static_cast<uint64_t>(asm_.exp.size()), std::memory_order_relaxed);
+      cx.qmm_weight_bytes.fetch_add(
+          static_cast<uint64_t>(asm_.exp.size()) *
+              (2 * (moe_ffn_ * (H_ / 8) * 4 + 2 * moe_ffn_ * (H_ / group_) * 2) +
+               (H_ * (moe_ffn_ / 8) * 4 + 2 * H_ * (moe_ffn_ / group_) * 2)),
+          std::memory_order_relaxed);
+      routed = moe_math_(args)[0];  // layer.py:1359 self._moe_math(x, *wargs, local2d)
+    } else {
+      const int F = static_cast<int>(moe_ffn_);
+      const int gsh = static_cast<int>(H_ / group_), gsf = static_cast<int>(moe_ffn_ / group_);
+      auto gl = gate_logits32(m_in, gate_w32_[static_cast<size_t>(li)]);
+      select_group_gpu(gl, &gate_b32_[static_cast<size_t>(li)], k, reg_.n_group, reg_.topk_group,
+                       reg_.routed_scaling, reg_.norm_topk_prob, idx, w);
+      if (owner) teach = idx;
+      mx::eval(idx, w);
+      exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+      auto idxh = mx::copy(mx::astype(idx, mx::int32), mx::default_stream(mx::Device::cpu));
+      auto wh = mx::copy(mx::astype(w, mx::float32), mx::default_stream(mx::Device::cpu));
+      mx::eval(idxh, wh);
+      const int32_t *ip = idxh.data<int32_t>();
+      const float *wp = wh.data<float>();
+      std::vector<int64_t> flat(ip, ip + k);
+      std::vector<int64_t> uniq = flat;
+      std::sort(uniq.begin(), uniq.end());
+      uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+      std::vector<int64_t> local_h(flat.size());
+      for (size_t j = 0; j < flat.size(); j++)
+        local_h[j] = static_cast<int32_t>(
+            std::lower_bound(uniq.begin(), uniq.end(), flat[j]) - uniq.begin());
+      auto bundles = io_->get_bundles(li, uniq);
+      if (bundles.size() != uniq.size()) {
+        err = "expert load missing items L" + std::to_string(li);
+        return false;
+      }
+      exec_counters().leases.fetch_add(uniq.size(), std::memory_order_relaxed);
+      for (auto &[e, b] : bundles) step_bundles_.push_back(b);
+      std::vector<mx::array> args{m_in};
+      for (int p = 0; p < 9; p++) {
+        std::vector<mx::array> rows;
+        for (auto e : uniq) rows.push_back(bundles[e]->at(p));
+        args.push_back(mx::stack(rows));  // [U,...](1380)
+      }
+      args.push_back(mx::array(local_h.begin(), mkshape({1, 1, k}), mx::int32));  // 1371 core.array(list)
+      auto &cx = exec_counters();
+      cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+      cx.qmm_experts.fetch_add(3 * static_cast<uint64_t>(uniq.size()), std::memory_order_relaxed);
+      cx.qmm_weight_bytes.fetch_add(
+          static_cast<uint64_t>(uniq.size()) *
+              (2 * (F * (H_ / 8) * 4 + 2 * F * gsh * 2) + (H_ * (F / 8) * 4 + 2 * H_ * gsf * 2)),
+          std::memory_order_relaxed);
+      routed = moe_math_(args)[0];
+      rec.predicted = false;
+      rec.sel.inds = uniq;
+      rec.sel.scores.clear();
+      for (int64_t e : uniq) {
+        double s = 0.0;
+        for (int64_t j = 0; j < k; j++)
+          if (flat[j] == e) s = wp[j];
+        rec.sel.scores.push_back(s);
+      }
+    }
+    if (owner) last_topk_[static_cast<size_t>(li)] = use_pg ? teach : idx;
+    auto out = mx::sum(routed * mx::astype(mx::expand_dims(w, -1), routed.dtype()),
+                       std::vector<int>{-2}, false);
+    mx::array sh =
+        L.sh_down.apply(e0_silu(L.sh_gate.apply(m_in)) * L.sh_up.apply(m_in));  // BailingMLP:503
+    x = h + out + sh;
+    if (clip_ > 0)
+      x = mx::where(mx::isnan(x), mx::zeros_like(x),
+                    mx::clip(x, mx::array(static_cast<float>(-clip_)),
+                             mx::array(static_cast<float>(clip_))));
+    return true;
+  }
+
+  //     before_layer_cb load-drop hooks.py:35-58 window=1 + async_eval_per_layer
+  if (exec_mode() == ExecMode::Batched) {
+    const int64_t k = reg_.top_k;
+    const int64_t E = reg_.num_experts;
+    io_->note_call(li);
+    const bool owner = li >= reg_.first_owner && li <= reg_.last_owner;
+    if (owner) m_in_cache_[static_cast<size_t>(li)] = m_in;
+    e0n::TimeSink gt_cpu(e0n::timing_counters().gate_cpu_ns);
+    auto gl = gate_logits32(m_in, gate_w32_[static_cast<size_t>(li)]);  // :527
+    mx::array idx(0.f), w(0.f);
+    select_group_gpu(gl, &gate_b32_[static_cast<size_t>(li)], k, reg_.n_group, reg_.topk_group,
+                     reg_.routed_scaling, reg_.norm_topk_prob, idx, w);
+    if (owner) last_topk_[static_cast<size_t>(li)] = idx;
+    mx::eval(idx, w);
+    exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+    auto idxh = mx::copy(mx::astype(idx, mx::int32), mx::default_stream(mx::Device::cpu));
+    auto wh = mx::copy(mx::astype(w, mx::float32), mx::default_stream(mx::Device::cpu));
+    mx::eval(idxh, wh);
+    const int32_t *ip = idxh.data<int32_t>();
+    const float *wp = wh.data<float>();
+    std::vector<int64_t> flat(static_cast<size_t>(T * k));
+    for (int64_t i = 0; i < T * k; i++) flat[static_cast<size_t>(i)] = ip[i];
+    mx::array routed(0.f);
+    const bool full = !prefill_ondemand() && io_->full_ready(li) && T * k > 8 &&
+                      !io_->hot_ready(li);
+    if (full) {
+      io_->note_prefill_use(li, flat, (int)k,
+                            li >= reg_.first_head_consumer && li <= reg_.last_head_consumer);
+      const auto &fw = io_->full_weights(li);
+      if ((int)fw.size() != 9) { err = "full-layer stack not loaded L" + std::to_string(li); return false; }
+      auto xe = mx::reshape(m_in, mkshape({1, T, 1, 1, H_}));           // 1246 expand_dims((-2,-3))
+      auto iflat = mx::reshape(idx, mkshape({T * k}));
+      auto order = mx::argsort(iflat, -1);                              // _gather_sort(switch_layers.py:12-17)
+      auto inv_order = mx::argsort(order, -1);
+      auto xs = mx::flatten(xe, 0, -3);                                 // [T,1,H]
+      auto tok_rows = mx::take(
+          xs, mx::floor_divide(order, mx::array((uint32_t)k, mx::uint32)), 0);  // x[order//M]
+      auto local = mx::take(iflat, order, 0);
+      auto gq = [&](const mx::array &xx, int o) {
+        return mx::gather_qmm(xx, fw[o], fw[o + 1], fw[o + 2], std::nullopt, local, true, 64, 4,
+                              "affine", true);
+      };
+      auto x_up = gq(tok_rows, 0);
+      auto x_gate = gq(tok_rows, 3);
+      auto yh = (x_gate * mx::sigmoid(x_gate)) * x_up;                   // _swiglu(1250;layer.py:57-58)
+      auto z = gq(yh, 6);
+      auto z_un = mx::take(z, inv_order, 0);                             // _scatter_unsort(1260)
+      routed = mx::squeeze(mx::reshape(z_un, mkshape({1, T, k, 1, H_})), -2);
+      auto &cx = exec_counters();
+      cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+      cx.qmm_experts.fetch_add(3 * static_cast<uint64_t>(E), std::memory_order_relaxed);
+      cx.qmm_weight_bytes.fetch_add(
+          static_cast<uint64_t>(E) *
+              (2 * (moe_ffn_ * (H_ / 8) * 4 + 2 * moe_ffn_ * (H_ / group_) * 2) +
+               (H_ * (moe_ffn_ / 8) * 4 + 2 * H_ * (moe_ffn_ / group_) * 2)),
+          std::memory_order_relaxed);
+    }
+    if (!full) {
+      std::vector<int64_t> uniq = flat;                       // 1364 sorted(set(int(v) for v in flat))
+    std::sort(uniq.begin(), uniq.end());
+    uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end());
+    std::vector<int32_t> local_h(static_cast<size_t>(T * k));  // 1370-1373 host remap
+    for (int64_t i = 0; i < T * k; i++)
+      local_h[static_cast<size_t>(i)] = static_cast<int32_t>(
+          std::lower_bound(uniq.begin(), uniq.end(), flat[static_cast<size_t>(i)]) - uniq.begin());
+    auto bundles = io_->get_bundles(li, uniq);
+    if (bundles.size() != uniq.size()) {
+      err = "expert load missing items L" + std::to_string(li);
+      return false;
+    }
+    exec_counters().leases.fetch_add(uniq.size(), std::memory_order_relaxed);
+    for (auto &[e, b] : bundles) step_bundles_.push_back(b);
+    std::vector<mx::array> wargs;                            // 1376-1380 wargs=core.stack(rows)
+    for (int p = 0; p < 9; p++) {
+      std::vector<mx::array> rows;
+      for (auto e : uniq) rows.push_back(bundles[e]->at(p));
+      wargs.push_back(mx::stack(rows));
+    }
+    auto local_arr = mx::array(local_h.begin(), mkshape({1, T, k}), mx::int32);  // 1371 core.array(list)
+    const int F = static_cast<int>(moe_ffn_);
+    const int gsh = static_cast<int>(H_ / group_), gsf = static_cast<int>(moe_ffn_ / group_);
+    auto &cx = exec_counters();
+    cx.qmm_calls.fetch_add(3, std::memory_order_relaxed);
+    cx.qmm_experts.fetch_add(3 * static_cast<uint64_t>(uniq.size()), std::memory_order_relaxed);
+    cx.qmm_weight_bytes.fetch_add(
+        static_cast<uint64_t>(uniq.size()) *
+            (2 * (F * (H_ / 8) * 4 + 2 * F * gsh * 2) + (H_ * (F / 8) * 4 + 2 * H_ * gsf * 2)),
+        std::memory_order_relaxed);
+    if (T * k >= 64) {
+      auto flat_local = mx::reshape(local_arr, mkshape({T * k}));
+      auto order = mx::argsort(flat_local, -1);             // _gather_sort: order=argsort(flat)
+      auto inv_order = mx::argsort(order, -1);              //                     inv=argsort(order)
+      auto xe = mx::reshape(m_in, mkshape({1, T, 1, 1, H_}));  // 1383 expand_dims((-2,-3)) [B,T,1,1,H]
+      auto xs = mx::flatten(xe, 0, -3);                     // 13-15 [T,1,H]
+      auto tok_rows = mx::take(
+          xs, mx::floor_divide(order, mx::array((uint32_t)k, mx::uint32)), 0);
+      auto local_sorted = mx::take(flat_local, order, 0);   // indices[order]
+      std::vector<mx::array> sargs{tok_rows};
+      for (auto &wa : wargs) sargs.push_back(wa);
+      sargs.push_back(local_sorted);
+      auto z = moe_math_sorted_(sargs)[0];
+      auto z_un = mx::take(z, inv_order, 0);
+      routed = mx::squeeze(mx::reshape(z_un, mkshape({1, T, k, 1, H_})), -2);  // unflatten+squeeze(-2)
+    } else {
+      std::vector<mx::array> args{m_in};
+      for (auto &wa : wargs) args.push_back(wa);
+      args.push_back(local_arr);
+      routed = moe_math_(args)[0];
+    }
+    }
+    auto out = mx::sum(routed * mx::astype(mx::expand_dims(w, -1), routed.dtype()),
+                       std::vector<int>{-2}, false);
+    mx::array sh =
+        L.sh_down.apply(e0_silu(L.sh_gate.apply(m_in)) * L.sh_up.apply(m_in));  // BailingMLP:503
+    gt_cpu.stop();
+    x = h + out + sh;  // :762
+    if (clip_ > 0)
+      x = mx::where(mx::isnan(x), mx::zeros_like(x),
+                    mx::clip(x, mx::array(static_cast<float>(-clip_)),
+                             mx::array(static_cast<float>(clip_))));
+    rec.layer = li;
+    rec.predicted = false;
+    {
+      std::vector<std::pair<int64_t, double>> lv;
+      for (int64_t j = 0; j < k; j++)
+        lv.push_back({ip[(T - 1) * k + j], wp[(T - 1) * k + j]});
+      std::sort(lv.begin(), lv.end());
+      for (auto &[e, s] : lv) {
+        rec.sel.inds.push_back(e);
+        rec.sel.scores.push_back(s);
+      }
+    }
+    if (full) {
+      mx::async_eval(x);
+    } else {
+      mx::eval(x);
+      step_bundles_.clear();
+    }
+    return true;
+  }
+
+  std::vector<float> mhost;
+  if (!to_host32(m_in, mhost, err)) return false;  // [T,H]
+
+  e0n::TimeSink gt_cpu(e0n::timing_counters().gate_cpu_ns);
+  const int64_t E = reg_.num_experts;
+  const double *bias = L.gate_bias.empty() ? nullptr : L.gate_bias.data();
+  std::vector<Selection> sels(static_cast<size_t>(T));
+  bool predicted = false;
+  if (decode && st.pr.take_pred(li, sels[0])) {
+    predicted = true;
+  } else {
+    std::vector<double> lg;
+    gate_logits(mhost.data(), static_cast<int>(T), static_cast<int>(H_), L.gate_h.data(),
+                static_cast<int>(E), lg);
+    for (int64_t t = 0; t < T; t++)
+      sels[static_cast<size_t>(t)] = select_sigmoid_group(
+          lg.data() + t * E, bias, static_cast<int>(E), reg_.top_k, reg_.n_group,
+          reg_.topk_group, reg_.routed_scaling, reg_.norm_topk_prob);
+  }
+  rec.layer = li;
+  rec.predicted = predicted;
+  rec.sel = sels[static_cast<size_t>(T - 1)];
+
+  const bool owner = li >= reg_.first_owner && li <= reg_.last_owner;
+  if (owner) {
+    std::vector<int64_t> feat = rec.sel.inds;
+    if (predicted && reg_.head_feature_teacher) {
+      std::vector<double> lg;
+      gate_logits(mhost.data() + (T - 1) * H_, 1, static_cast<int>(H_), L.gate_h.data(),
+                  static_cast<int>(E), lg);
+      feat = select_sigmoid_group(lg.data(), bias, static_cast<int>(E), reg_.top_k, reg_.n_group,
+                                  reg_.topk_group, reg_.routed_scaling, reg_.norm_topk_prob)
+                 .inds;
+    }
+    std::vector<float> last_row(mhost.begin() + (T - 1) * H_, mhost.end());
+    st.pr.capture(li, std::move(last_row), std::move(feat));
+    if (prefill_seed && reg_.head_feature_teacher) {
+      std::vector<int64_t> prev;
+      if (T >= 2) {
+        std::vector<double> lg;
+        gate_logits(mhost.data() + (T - 2) * H_, 1, static_cast<int>(H_), L.gate_h.data(),
+                    static_cast<int>(E), lg);
+        prev = select_sigmoid_group(lg.data(), bias, static_cast<int>(E), reg_.top_k,
+                                    reg_.n_group, reg_.topk_group, reg_.routed_scaling,
+                                    reg_.norm_topk_prob)
+                   .inds;
+      }
+      st.pr.prev_exec[li] = std::move(prev);
+    }
+  }
+  gt_cpu.stop();
+
+  std::vector<int64_t> uni;
+  for (auto &s : sels)
+    for (auto e : s.inds)
+      if (uni.empty() || uni.back() != e) uni.push_back(e);
+  std::sort(uni.begin(), uni.end());
+  uni.erase(std::unique(uni.begin(), uni.end()), uni.end());
+  auto bundles = io_->get_bundles(li, uni);
+  exec_counters().leases.fetch_add(uni.size(), std::memory_order_relaxed);
+  std::vector<ExpertPtrs> ptrs(uni.size());
+  for (size_t i = 0; i < uni.size(); i++)
+    ptrs[i] = bundle_ptrs(*bundles[uni[i]]);
+  std::vector<float> routed(static_cast<size_t>(T) * H_, 0.f);
+  const int F = static_cast<int>(moe_ffn_);
+  const int gsh = static_cast<int>(H_ / group_), gsf = static_cast<int>(moe_ffn_ / group_);
+  {
+  for (size_t ei = 0; ei < uni.size(); ei++) {
+    int64_t e = uni[ei];
+    std::vector<int> toks;
+    for (int64_t t = 0; t < T; t++)
+      for (auto id : sels[static_cast<size_t>(t)].inds)
+        if (id == e) toks.push_back(static_cast<int>(t));
+    if (toks.empty()) continue;
+    std::vector<float> rows(toks.size() * H_);
+    for (size_t k = 0; k < toks.size(); k++)
+      std::memcpy(rows.data() + k * H_, mhost.data() + static_cast<size_t>(toks[k]) * H_,
+                  H_ * 4);
+    std::vector<int64_t> one{0};
+    std::vector<float> yg, yu;
+    if (!gather_qmm_mlx(ptrs[ei].gw, ptrs[ei].gs, ptrs[ei].gb, 1, F, static_cast<int>(H_), gsh,
+                        rows.data(), static_cast<int>(toks.size()), one, yg, err) ||
+        !gather_qmm_mlx(ptrs[ei].uw, ptrs[ei].us, ptrs[ei].ub, 1, F, static_cast<int>(H_), gsh,
+                        rows.data(), static_cast<int>(toks.size()), one, yu, err))
+      return false;
+    std::vector<float> yh(toks.size() * moe_ffn_);
+    for (size_t k = 0; k < yh.size(); k++) yh[k] = silu_f(yg[k]) * yu[k];
+    std::vector<float> yd;
+    if (!gather_qmm_mlx(ptrs[ei].dw, ptrs[ei].ds, ptrs[ei].db, 1, static_cast<int>(H_), F, gsf,
+                        yh.data(), static_cast<int>(toks.size()), one, yd, err))
+      return false;
+    for (size_t k = 0; k < toks.size(); k++) {
+      int t = toks[k];
+      double w = 0.0;
+      auto &s = sels[static_cast<size_t>(t)];
+      for (size_t j = 0; j < s.inds.size(); j++)
+        if (s.inds[j] == e) w = s.scores[j];
+      for (int64_t o = 0; o < H_; o++)
+        routed[static_cast<size_t>(t) * H_ + o] += static_cast<float>(w * yd[k * H_ + o]);
+    }
+  }
+  }
+
+  mx::array sh =
+      L.sh_down.apply(e0_silu(L.sh_gate.apply(m_in)) * L.sh_up.apply(m_in));  // [1,T,H]
+  std::vector<float> shost;
+  if (!to_host32(sh, shost, err)) return false;
+  for (int64_t t = 0; t < T; t++)
+    for (int64_t o = 0; o < H_; o++)
+      routed[static_cast<size_t>(t) * H_ + o] += shost[static_cast<size_t>(t) * H_ + o];
+  mx::array y_arr = mx::astype(wrap_f32(routed, mkshape({1, T, H_})), dt);
+  x = h + y_arr;
+  if (clip_ > 0) {
+    x = mx::where(mx::isnan(x), mx::zeros_like(x),
+                  mx::clip(x, mx::array(static_cast<float>(-clip_)),
+                           mx::array(static_cast<float>(clip_))));
+  }
+  E0_TIME(e0n::timing_counters().eval_ns);
+  mx::eval(x);
+  return true;
+}
+
+bool Net8b::stage_all_gpu(State8b &st, std::string &err, DecodeProbe *probe) {
+  const int64_t E = reg_.num_experts;
+  std::vector<int> metas;
+  std::vector<mx::array> ms, curs, prevs;
+  for (int o : pr_->owners()) {
+    const auto &m = m_in_cache_[static_cast<size_t>(o)];
+    if (m.ndim() < 3) continue;
+    metas.push_back(o);
+    mx::Shape s0 = m.shape(), s1 = m.shape();
+    for (auto &v : s0) v = 0;
+    s0[s0.size() - 2] = static_cast<mx::Shape::value_type>(m.shape(-2) - 1);
+    ms.push_back(mx::slice(m, s0, s1));
+    mx::array co(0.f);                                  // 284 cur=pos-1 one-hot
+    if (!onehot_pos(last_topk_[static_cast<size_t>(o)], E, -1, co))
+      co = mx::zeros(mkshape({1, 1, E}), mx::float32);
+    curs.push_back(co);
+    mx::array po(0.f);
+    if (prev_topk_oh_[static_cast<size_t>(o)].ndim() > 0)
+      po = prev_topk_oh_[static_cast<size_t>(o)];
+    else if (!onehot_pos(last_topk_[static_cast<size_t>(o)], E, -2, po))
+      po = mx::zeros(mkshape({1, 1, E}), mx::float32);  // 292-293
+    prevs.push_back(po);
+  }
+  if (metas.empty()) return true;                  // 185-186
+  const int64_t n = static_cast<int64_t>(metas.size());
+  auto s_ms = mx::stack(ms, 0), s_curs = mx::stack(curs, 0), s_prevs = mx::stack(prevs, 0);
+  auto feats = mx::astype(mx::concatenate({s_ms, s_curs, s_prevs}, -1), mx::float16);  // 158-161
+  auto f2 = mx::reshape(feats, mkshape({n, H_ + 2 * E}));  // 162
+  auto gel = [](const mx::array &x) {
+    return mx::array(0.5f, mx::float16) * x *
+           (mx::array(1.0f, mx::float16) +
+            mx::erf(x / mx::array(static_cast<float>(std::sqrt(2.0)), mx::float16)));
+  };
+  auto h1 = mx::einsum("ni,nij->nj", {f2, hw1_});  // 163
+  auto out = mx::einsum("ni,nij->nj", {f2, hwl_}) +
+             mx::einsum("ni,nij->nj", {gel(h1), hw2_});  // 165-166
+  auto logits_all = mx::reshape(out, mkshape({n, 1, 1, E}));  // 167
+  mx::array inds_all(0.f), scores_all(0.f);
+  select_group_gpu(logits_all, nullptr, reg_.top_k, reg_.n_group, reg_.topk_group,
+                   reg_.routed_scaling, reg_.norm_topk_prob, inds_all, scores_all);  // _select 296-300
+  auto inds_h = mx::add(inds_all, mx::array((uint32_t)0, mx::uint32));
+  auto scores_h = mx::add(scores_all, mx::array(0.0f, mx::float32));
+  mx::eval(inds_h, scores_h);                      // 193 ONE eval
+  exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+  const int k = static_cast<int>(reg_.top_k);
+  const uint32_t *ip = inds_h.data<uint32_t>();  // 194 ONE tolist
+  const float *sp = scores_h.data<float>();
+  for (int64_t i = 0; i < n; i++) {
+    const int o = metas[i], consumer = static_cast<int>(o + 1);
+    Selection s;
+    std::vector<int64_t> eids;
+    for (int j = 0; j < k; j++) eids.push_back(ip[i * k + j]);
+    std::sort(eids.begin(), eids.end());
+    eids.erase(std::unique(eids.begin(), eids.end()), eids.end());
+    for (auto e : eids) {
+      double v = 0.0;
+      for (int j = 0; j < k; j++)
+        if (ip[i * k + j] == static_cast<uint32_t>(e)) v = sp[i * k + j];
+      s.inds.push_back(e);
+      s.scores.push_back(v);
+    }
+    if (consumer <= reg_.last_head_consumer) {
+      st.pr.pred[static_cast<size_t>(consumer)] = s;
+      st.pr.pred_avail[static_cast<size_t>(consumer)] = true;
+      io_->stage_experts(consumer, s.inds);
+      mx::Shape a0 = logits_all.shape(), a1 = logits_all.shape();
+      for (auto &v : a0) v = 0;
+      a0[0] = static_cast<mx::Shape::value_type>(i);
+      a1[0] = static_cast<mx::Shape::value_type>(i + 1);
+      pg_logits_[static_cast<size_t>(consumer)] =
+          mx::squeeze(mx::slice(logits_all, a0, a1), 0);
+    }
+    prev_topk_oh_[static_cast<size_t>(o)] = curs[i];
+  }
+  st.pr.step++;
+  if (probe) {
+    auto lg32 = mx::astype(logits_all, mx::float32);
+    mx::eval(lg32);
+    exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+    const float *lp = lg32.data<float>();
+    probe->owners.clear();
+    probe->m_in.clear();
+    probe->logits.clear();
+    for (int64_t i = 0; i < n; i++) {
+      probe->owners.push_back(metas[i]);
+      probe->logits.emplace_back(lp + i * E, lp + (i + 1) * E);
+      std::vector<float> mv;
+      auto r32 = mx::astype(mx::reshape(ms[i], mkshape({H_})), mx::float32);
+      if (!to_host32(r32, mv, err)) return false;
+      probe->m_in.push_back(std::move(mv));
+    }
+  }
+  return true;
+}
+
+void Net8b::reset_graph_caches() const {
+  const size_t NL = layers_.size();
+  m_in_cache_.assign(NL, mx::array(0.f));
+  last_topk_.assign(NL, mx::array(0.f));
+  prev_topk_oh_.assign(NL, mx::array(0.f));
+  pg_logits_.assign(NL, mx::array(0.f));
+  for (auto &L : layers_)
+    if (L.has_moe) io_->reset_layer(static_cast<int>(&L - layers_.data()));  // layer.py:1429-1441
+  step_bundles_.clear();
+}
+
+void Net8b::prefetch(const Step &plan) {
+  std::map<int, std::vector<int64_t>> by_layer;
+  for (auto &[li, e] : plan.experts) by_layer[li].push_back(e);
+  for (auto &[li, es] : by_layer) io_->prefetch(li, es);
+}
+
+void Net8b::shard_warm(const std::vector<std::string> &shard_names, bool seq) {
+  for (auto &s : shard_names) {
+    io_->shard_advise_whole(s);
+    if (seq) io_->shard_seq_read(s);
+  }
+}
+
+bool Net8b::prefill(const std::vector<int32_t> &tokens, State8b &st,
+                    std::vector<float> &logits_last, std::vector<RouteRec> &recs,
+                    std::string &err, PrefillProbe *probe) {
+  if (exec_mode() == ExecMode::Batched)
+    return prefill_batched_(tokens, st, logits_last, recs, err, probe);
+  try {
+    std::vector<float> xh;
+    if (!embed_lookup(embed_, tokens, H_, V_, xh, err)) return false;
+    const int64_t T = static_cast<int64_t>(tokens.size());
+    if (probe) probe->embed_last.assign(xh.begin() + (T - 1) * H_, xh.end());
+    mx::array x = mx::astype(wrap_f32(xh, mkshape({1, T, H_})), mx::bfloat16);
+    recs.clear();
+    for (size_t li = 0; li < layers_.size(); li++) {
+      RouteRec rec;
+      if (!layer_fwd(layers_[li], static_cast<int>(li), x, st, false, true, rec, err))
+        return false;
+      if (probe) {
+        mx::array tail = mx::astype(
+            mx::reshape(mx::slice(x, mkshape({0, T - 1, 0}), mkshape({1, T, H_})),
+                        mkshape({1, H_})),
+            mx::float32);
+        std::vector<float> hv;
+        if (!to_host32(tail, hv, err)) return false;
+        probe->per_block.push_back(std::move(hv));
+      }
+      if (layers_[li].has_moe) recs.push_back(std::move(rec));
+    }
+    mx::array xs = e0_rms_norm(
+        x, mx::astype(wrap_f32(final_norm_w_, mkshape({H_})), mx::bfloat16), eps_);
+    mx::array last = mx::reshape(mx::slice(xs, mkshape({0, T - 1, 0}), mkshape({1, T, H_})),
+                                 mkshape({1, H_}));
+    if (!to_host32(lm_head_.apply(last), logits_last, err)) return false;  // [V]
+    st.position += T;
+    if (!st.pr.run_step_boundary(*pr_, nullptr, err)) return false;
+  } catch (const std::exception &ex) {
+    err = std::string("prefill: ") + ex.what();
+    return false;
+  }
+  return true;
+}
+
+bool Net8b::prefill_batched_(const std::vector<int32_t> &tokens, State8b &st,
+                             std::vector<float> &logits_last, std::vector<RouteRec> &recs,
+                             std::string &err, PrefillProbe *probe) {
+  try {
+    const int64_t total = static_cast<int64_t>(tokens.size());
+    const int64_t chunk = reg_.prefill_chunk > 0 ? reg_.prefill_chunk : total;  // base.py:39(wrapper :71=2048)
+    recs.clear();
+    std::vector<float> lg;
+    for (int64_t start = 0; start < total; start += chunk) {
+      const int64_t n = std::min(chunk, total - start);
+      const bool last = start + n >= total;
+      step_bundles_.clear();
+      if (!forward_chunk_batched_(tokens.data() + start, n, st, lg, last, recs, err,
+                                  last ? probe : nullptr))
+        return false;
+      st.position += n;  // base.py:83
+    }
+    if (!prefill_end_(st, err)) return false;
+    logits_last = std::move(lg);
+    return true;
+  } catch (const std::exception &ex) {
+    err = std::string("prefill: ") + ex.what();
+    return false;
+  }
+}
+
+bool Net8b::forward_chunk_batched_(const int32_t *tok, int64_t T, State8b &st,
+                                   std::vector<float> &logits_out, bool to_host,
+                                   std::vector<RouteRec> &recs, std::string &err,
+                                   PrefillProbe *probe) {
+  std::vector<float> xh;
+  if (!embed_lookup(embed_, std::vector<int32_t>(tok, tok + T), H_, V_, xh, err)) return false;
+  if (probe) probe->embed_last.assign(xh.begin() + (T - 1) * H_, xh.end());
+  mx::array x = mx::astype(wrap_f32(xh, mkshape({1, T, H_})), mx::bfloat16);
+  const bool cb = T > 1 && !prefill_ondemand() &&
+                  (reg_.full_layer_prefill || reg_.prefill_hot > 0);
+  std::vector<RouteRec> local_recs;
+  for (size_t li = 0; li < layers_.size(); li++) {
+    if (cb) prefill_before_layer(static_cast<int>(li));
+    RouteRec rec;
+    if (!layer_fwd(layers_[li], static_cast<int>(li), x, st, false, true, rec, err))
+      return false;
+    if (probe) {
+      mx::array tail = mx::astype(
+          mx::reshape(mx::slice(x, mkshape({0, T - 1, 0}), mkshape({1, T, H_})),
+                      mkshape({1, H_})),
+          mx::float32);
+      std::vector<float> hv;
+      if (!to_host32(tail, hv, err)) return false;
+      probe->per_block.push_back(std::move(hv));
+    }
+    if (layers_[li].has_moe) local_recs.push_back(std::move(rec));
+  }
+  mx::array xs = e0_rms_norm(
+      x, mx::astype(wrap_f32(final_norm_w_, mkshape({H_})), mx::bfloat16), eps_);
+  mx::array lastrow =
+      mx::reshape(mx::slice(xs, mkshape({0, T - 1, 0}), mkshape({1, T, H_})), mkshape({1, H_}));
+  mx::array lg = lm_head_.apply(lastrow);  // ling.py:220
+  if (to_host) {
+    if (!to_host32(lg, logits_out, err)) return false;
+    recs = std::move(local_recs);
+  } else {
+    mx::eval(lg);
+  }
+  return true;
+}
+
+void Net8b::prefill_before_layer(int li) {
+  io_->clear_full_layer(li - 1);
+  if (li >= 0 && li < (int)layers_.size() && layers_[size_t(li)].has_moe)
+    io_->load_full_layer(li);
+}
+
+bool Net8b::prefill_end_(State8b &st, std::string &err) {
+  for (size_t li = 0; li < layers_.size(); li++) io_->clear_full_layer(static_cast<int>(li));  // :278-279
+  return stage_all_gpu(st, err, nullptr);
+}
+
+bool Net8b::step(int32_t token, State8b &st, std::vector<float> &logits,
+                 std::vector<RouteRec> &recs, std::string &err, DecodeProbe *probe) {
+  try {
+    step_bundles_.clear();
+    std::vector<float> xh;
+    if (!embed_lookup(embed_, {token}, H_, V_, xh, err)) return false;
+    mx::array x = mx::astype(wrap_f32(xh, mkshape({1, 1, H_})), mx::bfloat16);
+    recs.clear();
+    for (size_t li = 0; li < layers_.size(); li++) {
+      RouteRec rec;
+      if (!layer_fwd(layers_[li], static_cast<int>(li), x, st, true, false, rec, err))
+        return false;
+      if (layers_[li].has_moe) recs.push_back(std::move(rec));
+    }
+    mx::array xs = e0_rms_norm(
+        x, mx::astype(wrap_f32(final_norm_w_, mkshape({H_})), mx::bfloat16), eps_);
+    if (!to_host32(lm_head_.apply(mx::reshape(xs, mkshape({1, H_}))), logits, err)) return false;
+    st.position += 1;
+    if (exec_mode() == ExecMode::Batched) {
+      if (!stage_all_gpu(st, err, probe)) return false;
+    } else if (probe) {
+      probe->owners.clear(); probe->m_in.clear(); probe->logits.clear();
+      for (int h = 0; h < pr_->n_heads(); h++) {
+        int own = pr_->owner(h);
+        if (!st.pr.cap_hidden[own].empty()) {
+          probe->owners.push_back(own);
+          probe->m_in.push_back(st.pr.cap_hidden[own]);
+        }
+      }
+      std::vector<std::vector<double>> diag;
+      if (!st.pr.run_step_boundary(*pr_, &diag, err)) return false;
+      probe->logits = std::move(diag);
+    } else if (!st.pr.run_step_boundary(*pr_, nullptr, err)) return false;
+  } catch (const std::exception &ex) {
+    err = std::string("step: ") + ex.what();
+    return false;
+  }
+  return true;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/net8b.hpp b/macos/crates/engine-native/src/net8b.hpp
new file mode 100644
index 0000000..26fb1bb
--- /dev/null
+++ b/macos/crates/engine-native/src/net8b.hpp
@@ -0,0 +1,124 @@
+// 8b (Ling 3.0 hybrid) forward graph: KDA/MLA mix, L0 dense FFN, MoE 128-choose-8 + shared.
+#pragma once
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <mlx/mlx.h>
+
+#include "expertio.hpp"
+#include "expertmap.hpp"
+#include "fwd_common.hpp"
+#include "hotstack.hpp"
+#include "lora.hpp"
+#include "net35.hpp"  // RouteRec (shared per-layer routing record)
+#include "prerouter.hpp"
+#include "provider.hpp"
+#include "registry.hpp"
+#include "router.hpp"
+#include "safio.hpp"
+
+namespace e0n {
+
+struct State8b {
+  PrerouterState pr;
+  std::vector<mlx::core::array> cq, ck, cv, ssm;
+  std::vector<mlx::core::array> kc, vc;
+  std::vector<char> is_mla;
+  int64_t position = 0;
+  bool started = false;
+};
+
+class Net8b {
+ public:
+  // lo = unmerged LoRA (nullptr = base). Loading is owned by the net (ExpertIo).
+  static std::unique_ptr<Net8b> create(std::shared_ptr<const ModelFiles> files,
+                                       std::shared_ptr<const ExpertMap> emap, std::string &err,
+                                       std::shared_ptr<const lora::Bundle> lo = nullptr);
+
+  const TierRouting &reg() const { return reg_; }
+  Prerouter &prerouter() { return *pr_; }
+  const lora::Bundle *lora() const { return lo_.get(); }
+  uint64_t dense_resident_bytes() const { return dense_bytes_; }
+  void prefetch(const Step &plan);
+  ExpertIo::Metrics io_metrics() const { return io_->snapshot(); }
+  void shard_warm(const std::vector<std::string> &shard_names, bool seq);
+
+  void init_state(State8b &st) const;
+  void reset_state(State8b &st) const;
+
+  bool prefill(const std::vector<int32_t> &tokens, State8b &st, std::vector<float> &logits_last,
+               std::vector<RouteRec> &recs, std::string &err, PrefillProbe *probe = nullptr);
+  bool step(int32_t token, State8b &st, std::vector<float> &logits, std::vector<RouteRec> &recs,
+            std::string &err, DecodeProbe *probe = nullptr);
+
+ private:
+  Net8b() = default;
+  struct Layer {
+    bool is_mla = false;
+    bool has_moe = false;  // L0 (dense_prefix) is false
+    std::vector<float> in_ln_w, post_ln_w;
+    QuantLin k_q, k_k, k_v, k_f, k_g, k_b, k_o;
+    std::vector<uint8_t> conv_qb, conv_kb, conv_vb;  // BF16 [proj,1,4]
+    mlx::core::array conv_q{0.f}, conv_k{0.f}, conv_v{0.f};  // load-time swap to [proj,4,1]
+    std::vector<float> A_log, dt_bias, onorm_w;
+    QuantLin m_qa, m_qb, m_kva, m_kvb, m_dense, m_g;
+    std::vector<float> m_qaln, m_kvaln;
+    std::vector<float> gate_h;
+    bool gate_bf16 = false;       // file is BF16 — wrap back to bf16 for same-domain mul
+    std::vector<double> gate_bias;
+    QuantLin sh_gate, sh_up, sh_down;
+    QuantLin d_gate, d_up, d_down;
+  };
+
+  bool layer_fwd(Layer &L, int li, mlx::core::array &x, State8b &st, bool decode, bool prefill_seed,
+                 RouteRec &rec, std::string &err);
+
+  mutable std::shared_ptr<const ModelFiles> files_;
+  std::shared_ptr<const ExpertMap> emap_;
+  std::unique_ptr<ExpertIo> io_;
+  TierRouting reg_;
+  std::vector<Layer> layers_;
+  std::unique_ptr<Prerouter> pr_;
+  std::shared_ptr<const lora::Bundle> lo_;
+
+  int64_t H_ = 0, V_ = 0, E_ = 0;
+  int64_t nh_ = 0, hd_ = 0, proj_ = 0, conv_k_ = 0;
+  int64_t qlora_ = 0, kvlora_ = 0, nope_ = 0, roped_ = 0, vhd_ = 0;
+  int64_t moe_ffn_ = 0, sh_ffn_ = 0, dense_ffn_ = 0;
+  int64_t group_ = 64;
+  double eps_ = 1e-6, rope_theta_ = 6000000.0, kda_lb_ = -5.0;
+  double clip_ = 1000.0;
+
+  QuantLin embed_;
+  QuantLin lm_head_;
+  std::vector<float> final_norm_w_;
+  uint64_t dense_bytes_ = 0;
+
+  bool stage_all_gpu(State8b &st, std::string &err, DecodeProbe *probe);
+  void reset_graph_caches() const;
+
+  bool prefill_batched_(const std::vector<int32_t> &tokens, State8b &st,
+                        std::vector<float> &logits_last, std::vector<RouteRec> &recs,
+                        std::string &err, PrefillProbe *probe);
+  bool forward_chunk_batched_(const int32_t *tok, int64_t T, State8b &st,
+                              std::vector<float> &logits_out, bool to_host,
+                              std::vector<RouteRec> &recs, std::string &err,
+                              PrefillProbe *probe);
+  void prefill_before_layer(int li);
+  bool prefill_end_(State8b &st, std::string &err);
+
+  mutable std::vector<mlx::core::array> m_in_cache_, last_topk_, prev_topk_oh_, pg_logits_;
+  mutable std::vector<std::shared_ptr<Bundle>> step_bundles_;
+  std::vector<mlx::core::array> gate_w32_, gate_b32_;
+  mlx::core::array hw1_{0.f}, hw2_{0.f}, hwl_{0.f};
+  std::function<std::vector<mlx::core::array>(const std::vector<mlx::core::array> &)> moe_math_;
+  std::function<std::vector<mlx::core::array>(const std::vector<mlx::core::array> &)>
+      moe_math_sorted_;
+
+  mlx::core::fast::CustomKernelFunction kda_kernel_;
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/prefixcache.hpp b/macos/crates/engine-native/src/prefixcache.hpp
new file mode 100644
index 0000000..ae72667
--- /dev/null
+++ b/macos/crates/engine-native/src/prefixcache.hpp
@@ -0,0 +1,264 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+#include <list>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace e0n {
+
+struct PrefixCacheConfig {
+  bool enabled = true;
+  uint64_t budget_bytes = 512ull * 1024ull * 1024ull;
+  uint32_t min_prefix_tokens = 128;
+};
+
+struct PrefixCacheLookup {
+  std::shared_ptr<void> snapshot;
+  size_t matched_tokens = 0;
+  std::string reason;
+};
+
+struct PrefixCacheMetrics {
+  uint64_t hits = 0;
+  uint64_t misses = 0;
+  uint64_t cached_tokens = 0;
+  uint64_t entries = 0;
+  uint64_t bytes = 0;
+  uint64_t evictions = 0;
+  uint64_t budget_bytes = 0;
+  uint32_t min_prefix_tokens = 0;
+  bool enabled = false;
+  size_t last_cached_tokens = 0;
+  size_t last_prefill_tokens = 0;
+  std::string last_reason = "cold";
+};
+
+// Prefix index and LRU tracking are independent from backend-specific KV state.
+// Hashes only narrow candidates; the complete token vector is always rechecked.
+class PrefixCache {
+ public:
+  explicit PrefixCache(PrefixCacheConfig cfg = {}) : cfg_(cfg) {}
+
+  PrefixCacheLookup lookup(const std::string &identity,
+                          const std::vector<int32_t> &tokens) {
+    std::lock_guard<std::mutex> g(mu_);
+    PrefixCacheLookup out;
+    out.reason = "prefix_miss";
+    if (!cfg_.enabled || cfg_.budget_bytes == 0) {
+      out.reason = "disabled";
+      metrics_.last_cached_tokens = 0;
+      metrics_.last_reason = out.reason;
+      return out;
+    }
+    if (tokens.size() < cfg_.min_prefix_tokens) {
+      out.reason = "too_short";
+      metrics_.last_cached_tokens = 0;
+      metrics_.last_reason = out.reason;
+      return out;
+    }
+
+    uint64_t h1 = kFnvOffset;
+    uint64_t h2 = kMixOffset;
+    uint64_t best_id = 0;
+    size_t best_len = 0;
+    for (size_t i = 0; i < tokens.size(); ++i) {
+      mix_token(tokens[i], h1, h2);
+      const size_t len = i + 1;
+      if (len < cfg_.min_prefix_tokens) continue;
+      Key key{identity, len, h1, h2};
+      auto it = index_.find(key);
+      if (it == index_.end()) continue;
+      for (uint64_t id : it->second) {
+        auto ent = entries_.find(id);
+        if (ent == entries_.end()) continue;
+        if (ent->second.identity != identity || ent->second.tokens.size() != len ||
+            !std::equal(ent->second.tokens.begin(), ent->second.tokens.end(), tokens.begin()))
+          continue;
+        if (len >= best_len) {
+          best_len = len;
+          best_id = id;
+        }
+      }
+    }
+    if (best_id == 0) {
+      ++metrics_.misses;
+      metrics_.last_cached_tokens = 0;
+      metrics_.last_reason = out.reason;
+      return out;
+    }
+    auto it = entries_.find(best_id);
+    touch_locked(it->second);
+    ++metrics_.hits;
+    metrics_.cached_tokens += best_len;
+    metrics_.last_cached_tokens = best_len;
+    metrics_.last_reason = "hit";
+    out.snapshot = it->second.snapshot;
+    out.matched_tokens = best_len;
+    out.reason = "hit";
+    return out;
+  }
+
+  bool put(const std::string &identity, const std::vector<int32_t> &tokens,
+           std::shared_ptr<void> snapshot, uint64_t bytes, std::string *reason = nullptr) {
+    std::lock_guard<std::mutex> g(mu_);
+    if (!cfg_.enabled || cfg_.budget_bytes == 0) {
+      if (reason) *reason = "disabled";
+      return false;
+    }
+    if (tokens.size() < cfg_.min_prefix_tokens) {
+      if (reason) *reason = "too_short";
+      return false;
+    }
+    if (bytes > cfg_.budget_bytes) {
+      metrics_.last_reason = "oversize";
+      if (reason) *reason = "oversize";
+      return false;
+    }
+    uint64_t h1 = kFnvOffset;
+    uint64_t h2 = kMixOffset;
+    for (int32_t tok : tokens) mix_token(tok, h1, h2);
+    Key key{identity, tokens.size(), h1, h2};
+    auto &bucket = index_[key];
+    for (uint64_t id : bucket) {
+      auto it = entries_.find(id);
+      if (it == entries_.end()) continue;
+      if (it->second.identity == identity && it->second.tokens == tokens) {
+        bytes_ -= it->second.bytes;
+        it->second.snapshot = std::move(snapshot);
+        it->second.bytes = bytes;
+        bytes_ += bytes;
+        touch_locked(it->second);
+        if (reason) *reason = "updated";
+        return true;
+      }
+    }
+    while (bytes_ + bytes > cfg_.budget_bytes && !lru_.empty()) {
+      const uint64_t old = lru_.back();
+      lru_.pop_back();
+      auto it = entries_.find(old);
+      if (it == entries_.end()) continue;
+      remove_from_index_locked(it->second);
+      bytes_ -= it->second.bytes;
+      entries_.erase(it);
+      ++metrics_.evictions;
+    }
+    Entry ent;
+    ent.id = next_id_++;
+    ent.identity = identity;
+    ent.tokens = tokens;
+    ent.h1 = h1;
+    ent.h2 = h2;
+    ent.snapshot = std::move(snapshot);
+    ent.bytes = bytes;
+    lru_.push_front(ent.id);
+    ent.lru = lru_.begin();
+    bytes_ += bytes;
+    bucket.push_back(ent.id);
+    entries_.emplace(ent.id, std::move(ent));
+    if (reason) *reason = "stored";
+    return true;
+  }
+
+  void set_last_prefill_tokens(size_t n) {
+    std::lock_guard<std::mutex> g(mu_);
+    metrics_.last_prefill_tokens = n;
+  }
+
+  PrefixCacheMetrics metrics() const {
+    std::lock_guard<std::mutex> g(mu_);
+    PrefixCacheMetrics m = metrics_;
+    m.entries = entries_.size();
+    m.bytes = bytes_;
+    m.budget_bytes = cfg_.budget_bytes;
+    m.min_prefix_tokens = cfg_.min_prefix_tokens;
+    m.enabled = cfg_.enabled && cfg_.budget_bytes != 0;
+    return m;
+  }
+
+  void clear() {
+    std::lock_guard<std::mutex> g(mu_);
+    index_.clear();
+    entries_.clear();
+    lru_.clear();
+    bytes_ = 0;
+    metrics_.last_cached_tokens = 0;
+    metrics_.last_reason = "cleared";
+  }
+
+ private:
+  static constexpr uint64_t kFnvOffset = 1469598103934665603ull;
+  static constexpr uint64_t kFnvPrime = 1099511628211ull;
+  static constexpr uint64_t kMixOffset = 0x9e3779b97f4a7c15ull;
+
+  struct Key {
+    std::string identity;
+    size_t length = 0;
+    uint64_t h1 = 0;
+    uint64_t h2 = 0;
+    bool operator==(const Key &o) const {
+      return length == o.length && h1 == o.h1 && h2 == o.h2 && identity == o.identity;
+    }
+  };
+  struct KeyHash {
+    size_t operator()(const Key &k) const {
+      size_t h = std::hash<std::string>{}(k.identity);
+      h ^= std::hash<size_t>{}(k.length) + 0x9e3779b9 + (h << 6) + (h >> 2);
+      h ^= std::hash<uint64_t>{}(k.h1) + 0x9e3779b9 + (h << 6) + (h >> 2);
+      h ^= std::hash<uint64_t>{}(k.h2) + 0x9e3779b9 + (h << 6) + (h >> 2);
+      return h;
+    }
+  };
+  struct Entry {
+    uint64_t id = 0;
+    std::string identity;
+    std::vector<int32_t> tokens;
+    uint64_t h1 = 0;
+    uint64_t h2 = 0;
+    std::shared_ptr<void> snapshot;
+    uint64_t bytes = 0;
+    std::list<uint64_t>::iterator lru;
+  };
+
+  static void mix_token(int32_t token, uint64_t &h1, uint64_t &h2) {
+    uint32_t v = static_cast<uint32_t>(token);
+    for (int i = 0; i < 4; ++i) {
+      const uint8_t b = static_cast<uint8_t>(v >> (i * 8));
+      h1 ^= b;
+      h1 *= kFnvPrime;
+      h2 ^= static_cast<uint64_t>(b) + 0x100000001b3ull;
+      h2 *= 0x9e3779b185ebca87ull;
+      h2 ^= h2 >> 29;
+    }
+  }
+
+  void touch_locked(Entry &ent) {
+    lru_.erase(ent.lru);
+    lru_.push_front(ent.id);
+    ent.lru = lru_.begin();
+  }
+
+  void remove_from_index_locked(const Entry &ent) {
+    Key key{ent.identity, ent.tokens.size(), ent.h1, ent.h2};
+    auto it = index_.find(key);
+    if (it == index_.end()) return;
+    auto &v = it->second;
+    v.erase(std::remove(v.begin(), v.end(), ent.id), v.end());
+    if (v.empty()) index_.erase(it);
+  }
+
+  mutable std::mutex mu_;
+  PrefixCacheConfig cfg_;
+  uint64_t next_id_ = 1;
+  uint64_t bytes_ = 0;
+  PrefixCacheMetrics metrics_;
+  std::list<uint64_t> lru_;
+  std::unordered_map<uint64_t, Entry> entries_;
+  std::unordered_map<Key, std::vector<uint64_t>, KeyHash> index_;
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/prerouter.cpp b/macos/crates/engine-native/src/prerouter.cpp
new file mode 100644
index 0000000..a90ecbe
--- /dev/null
+++ b/macos/crates/engine-native/src/prerouter.cpp
@@ -0,0 +1,297 @@
+#include "prerouter.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+
+#include <mlx/mlx.h>
+
+#include "dense.hpp"
+
+using mlx::core::array;
+using mlx::core::Device;
+
+namespace e0n {
+namespace {
+
+uint16_t to_f16_bits1(float v) {
+  _Float16 h = static_cast<_Float16>(v);
+  uint16_t b;
+  std::memcpy(&b, &h, 2);
+  return b;
+}
+
+float f16_bits_to_f(uint16_t b) {
+  _Float16 h;
+  std::memcpy(&h, &b, 2);
+  return static_cast<float>(h);
+}
+
+std::vector<float> expand_f16(const std::vector<float> &v) {
+  std::vector<float> o(v.size());
+  for (size_t i = 0; i < v.size(); i++) o[i] = f16_bits_to_f(to_f16_bits1(v[i]));
+  return o;
+}
+
+}  // namespace
+
+struct Prerouter::Impl {
+  TierRouting reg;
+  int64_t hidden = 0, f_in = 0;
+  std::vector<int> owners;
+  std::vector<float> fc1, fc2, lin;
+  bool ready = false;
+};
+
+Prerouter::Prerouter() : impl_(std::make_unique<Impl>()) {}
+Prerouter::~Prerouter() = default;
+Prerouter::Prerouter(Prerouter &&) noexcept = default;
+Prerouter &Prerouter::operator=(Prerouter &&) noexcept = default;
+
+int Prerouter::n_heads() const { return impl_ ? static_cast<int>(impl_->owners.size()) : 0; }
+int Prerouter::owner(int i) const { return impl_->owners[i]; }
+
+const std::vector<int> &Prerouter::owners() const { return impl_->owners; }
+const std::vector<float> &Prerouter::stacked_fc1() const { return impl_->fc1; }
+const std::vector<float> &Prerouter::stacked_fc2() const { return impl_->fc2; }
+const std::vector<float> &Prerouter::stacked_lin() const { return impl_->lin; }
+int64_t Prerouter::stacked_f_in() const { return impl_->f_in; }
+int64_t Prerouter::stacked_prh() const { return impl_->reg.prerouter_hidden; }
+int64_t Prerouter::stacked_E() const { return impl_->reg.num_experts; }
+
+bool Prerouter::build(const std::vector<HeadW> &heads, const TierRouting &reg, int64_t hidden,
+                      Prerouter &out, std::string &err) {
+  if (heads.empty()) { err = "prerouter: zero heads"; return false; }
+  const int64_t prh = reg.prerouter_hidden, E = reg.num_experts;
+  int64_t f_in = -1;
+  for (auto &h : heads) {
+    if (h.fc1.empty() || h.fc2.empty() || h.lin.empty() ||
+        static_cast<int64_t>(h.fc2.size()) != E * prh ||
+        static_cast<int64_t>(h.lin.size()) != static_cast<int64_t>(h.fc1.size()) / prh * E) {
+      err = "prerouter: header shape mismatch owner=" + std::to_string(h.owner);
+      return false;
+    }
+    int64_t fi = static_cast<int64_t>(h.fc1.size()) / prh;
+    if (fi != hidden + 2 * E) {
+      err = "prerouter: f_in != hidden+2E (" + std::to_string(fi) + " vs " +
+            std::to_string(hidden + 2 * E) + ") owner=" + std::to_string(h.owner);
+      return false;
+    }
+    if (f_in < 0) f_in = fi;
+    else if (f_in != fi) { err = "prerouter: f_in disagrees across heads"; return false; }
+  }
+  for (size_t i = 1; i < heads.size(); i++)
+    if (heads[i].owner <= heads[i - 1].owner) { err = "prerouter: owners are not sorted"; return false; }
+
+  auto im = std::make_unique<Prerouter::Impl>();
+  im->reg = reg;
+  im->hidden = hidden;
+  im->f_in = f_in;
+  const int O = static_cast<int>(heads.size());
+  for (auto &h : heads) im->owners.push_back(h.owner);
+
+  std::vector<float> fc1t(static_cast<size_t>(O) * f_in * prh),
+      fc2s(static_cast<size_t>(O) * prh * E), lin_s(static_cast<size_t>(O) * f_in * E);
+  for (int i = 0; i < O; i++) {
+    const HeadW &h = heads[i];
+    for (int64_t r = 0; r < prh; r++)
+      for (int64_t c = 0; c < f_in; c++)
+        fc1t[(static_cast<size_t>(i) * f_in + c) * prh + r] = h.fc1[r * f_in + c];
+    for (int64_t e = 0; e < E; e++) {
+      for (int64_t r = 0; r < prh; r++)
+        fc2s[(static_cast<size_t>(i) * prh + r) * E + e] = h.fc2[e * prh + r];
+      for (int64_t c = 0; c < f_in; c++)
+        lin_s[(static_cast<size_t>(i) * f_in + c) * E + e] = h.lin[e * f_in + c];
+    }
+  }
+  im->fc1 = expand_f16(fc1t);
+  im->fc2 = expand_f16(fc2s);
+  im->lin = expand_f16(lin_s);
+  im->ready = true;
+  out.impl_ = std::move(im);
+  return true;
+}
+
+bool Prerouter::load(const ModelFiles &files, const TierRouting &reg, int64_t hidden,
+                     Prerouter &out, std::string &err) {
+  std::vector<HeadW> heads;
+  for (int o = reg.first_owner; o <= reg.last_owner; o++) {
+    HeadW h;
+    h.owner = o;
+    const std::string base = reg.head_file;
+    DenseF32 fc1, fc2, lin;
+    std::string k1 = "layers." + std::to_string(o) + ".fc1.weight";
+    std::string k2 = "layers." + std::to_string(o) + ".fc2.weight";
+    std::string k3 = "layers." + std::to_string(o) + ".linear_init.weight";
+    if (!read_dense_f32(files, k1, fc1, err) || !read_dense_f32(files, k2, fc2, err) ||
+        !read_dense_f32(files, k3, lin, err)) {
+      if (err.find("missing tensor") == 0) err += " (header file " + base + " or missing from the directory)";
+      return false;
+    }
+    if (fc1.rows != reg.prerouter_hidden || fc2.cols != reg.prerouter_hidden ||
+        fc2.rows != reg.num_experts || lin.rows != reg.num_experts || fc1.cols != lin.cols) {
+      err = "prerouter: header shape does not match registry owner=" + std::to_string(o);
+      return false;
+    }
+    h.fc1 = std::move(fc1.v);
+    h.fc2 = std::move(fc2.v);
+    h.lin = std::move(lin.v);
+    heads.push_back(std::move(h));
+  }
+  return build(heads, reg, hidden, out, err);
+}
+
+bool Prerouter::predict(const std::vector<Input> &in, std::vector<Selection> &out_sel,
+                        std::vector<std::vector<double>> *out_logits, std::string &err) {
+  if (!impl_ || !impl_->ready) { err = "prerouter: not built"; return false; }
+  const int O = n_heads();
+  if (static_cast<int>(in.size()) != O) { err = "prerouter: head count does not match input count"; return false; }
+  const auto &reg = impl_->reg;
+  const int64_t prh = reg.prerouter_hidden, E = reg.num_experts, f_in = impl_->f_in;
+
+  std::vector<float> feats(static_cast<size_t>(O) * f_in, 0.f);
+  for (int i = 0; i < O; i++) {
+    if (!in[i].hidden) { err = "prerouter: missing header hidden input owner=" + std::to_string(owner(i)); return false; }
+    float *row = feats.data() + static_cast<size_t>(i) * f_in;
+    std::memcpy(row, in[i].hidden, sizeof(float) * impl_->hidden);
+    for (int64_t e : in[i].cur_exec)
+      if (e >= 0 && e < E) row[impl_->hidden + e] = 1.f;
+    for (int64_t e : in[i].prev_exec)
+      if (e >= 0 && e < E) row[impl_->hidden + E + e] = 1.f;
+  }
+
+  try {
+    auto wrap = [](void *p, mlx::core::Shape shape, mlx::core::Dtype dt) {
+      return array(p, std::move(shape), dt, [](void *) {});
+    };
+    auto a_feats = wrap(feats.data(), {O, 1, static_cast<int32_t>(f_in)}, mlx::core::float32);
+    auto a_fc1 = wrap(impl_->fc1.data(),
+                      {O, static_cast<int32_t>(f_in), static_cast<int32_t>(prh)},
+                      mlx::core::float32);
+    auto a_fc2 = wrap(impl_->fc2.data(),
+                      {O, static_cast<int32_t>(prh), static_cast<int32_t>(E)},
+                      mlx::core::float32);
+    auto a_lin = wrap(impl_->lin.data(),
+                      {O, static_cast<int32_t>(f_in), static_cast<int32_t>(E)},
+                      mlx::core::float32);
+    const auto f16 = mlx::core::float16;
+    auto qf = mlx::core::astype(a_feats, f16);
+    auto h = mlx::core::matmul(qf, mlx::core::astype(a_fc1, f16));  // [O,1,prh] f16
+    auto gel = h * 0.5f * (1.0f + mlx::core::erf(h / static_cast<float>(std::sqrt(2.0))));
+    auto logits = mlx::core::astype(
+        mlx::core::matmul(gel, mlx::core::astype(a_fc2, f16)) +
+            mlx::core::matmul(qf, mlx::core::astype(a_lin, f16)),
+        mlx::core::float32);
+    auto host = mlx::core::copy(logits, mlx::core::default_stream(Device::cpu));
+    mlx::core::eval(host);
+    const float *lp = host.data<float>();
+    out_sel.assign(O, Selection{});
+    if (out_logits) out_logits->assign(O, {});
+    for (int i = 0; i < O; i++) {
+      const double *rowp = nullptr;
+      std::vector<double> drow(E);
+      for (int64_t e = 0; e < E; e++) drow[e] = lp[static_cast<size_t>(i) * E + e];
+      rowp = drow.data();
+      if (out_logits) (*out_logits)[i] = drow;
+      out_sel[i] = reg.family == RouterFamily::kSoftmaxTopk
+                       ? select_softmax_topk(rowp, static_cast<int>(E), reg.top_k,
+                                             reg.norm_topk_prob)
+                       : select_sigmoid_group(rowp, nullptr, static_cast<int>(E), reg.top_k,
+                                              reg.n_group, reg.topk_group, reg.routed_scaling,
+                                              reg.norm_topk_prob);
+    }
+  } catch (const std::exception &ex) {
+    err = std::string("prerouter mlx: ") + ex.what();
+    return false;
+  }
+  return true;
+}
+
+// ---------------- PrerouterState ----------------
+
+Step live_next_step(const PrerouterState &st) {
+  Step plan;
+  for (int l = st.reg.first_head_consumer; l <= st.reg.last_head_consumer; l++)
+    if (!st.pred[l].inds.empty())
+      for (auto e : st.pred[l].inds) plan.experts.emplace_back(l, e);
+  return plan;
+}
+
+void PrerouterState::init(const TierRouting &r) {
+  reg = r;
+  cap_hidden.assign(r.num_layers, {});
+  cap_exec.assign(r.num_layers, {});
+  prev_exec.assign(r.num_layers, {});
+  pred.assign(r.num_layers, Selection{});
+  pred_avail.assign(r.num_layers, false);
+  step = 0;
+}
+
+void PrerouterState::capture(int layer, std::vector<float> m_in, std::vector<int64_t> exec) {
+  cap_hidden[layer] = std::move(m_in);
+  std::sort(exec.begin(), exec.end());
+  cap_exec[layer] = std::move(exec);
+}
+
+bool PrerouterState::run_step_boundary(Prerouter &pr,
+                                       std::vector<std::vector<double>> *diag_logits,
+                                       std::string &err) {
+  const int O = pr.n_heads();
+  for (int l = reg.first_head_consumer; l <= reg.last_head_consumer; l++) pred_avail[l] = false;
+  std::vector<Prerouter::Input> in(O);
+  for (int i = 0; i < O; i++) {
+    int own = pr.owner(i);
+    if (cap_hidden[own].empty()) continue;
+    in[i].hidden = cap_hidden[own].data();
+    in[i].cur_exec = cap_exec[own];
+    in[i].prev_exec = prev_exec[own];
+  }
+  std::vector<Prerouter::Input> compact;
+  std::vector<int> idx_of;
+  for (int i = 0; i < O; i++)
+    if (in[i].hidden) { compact.push_back(in[i]); idx_of.push_back(i); }
+  if (compact.empty()) { swap(); return true; }
+
+  std::vector<Selection> sel;
+  std::vector<std::vector<double>> lg;
+  if (!pr.predict(compact, sel, diag_logits ? &lg : nullptr, err)) return false;
+  if (diag_logits) *diag_logits = std::move(lg);
+  for (size_t j = 0; j < idx_of.size(); j++) {
+    int own = pr.owner(idx_of[j]);
+    int consumer = own + 1;
+    if (consumer >= reg.first_head_consumer && consumer <= reg.last_head_consumer) {
+      pred[consumer] = sel[j];
+      pred_avail[consumer] = true;
+    }
+  }
+  if (diag_logits) {
+  }
+  swap();
+  return true;
+}
+
+void PrerouterState::swap() {
+  prev_exec = cap_exec;
+  for (auto &v : cap_hidden) v.clear();
+  for (auto &v : cap_exec) v.clear();
+  step++;
+}
+
+bool PrerouterState::take_pred(int layer, Selection &out) {
+  if (layer < 0 || layer >= reg.num_layers || !pred_avail[layer] || pred[layer].inds.empty())
+    return false;
+  out = pred[layer];
+  pred_avail[layer] = false;
+  return true;
+}
+
+void PrerouterState::reset() {
+  for (auto &v : cap_hidden) v.clear();
+  for (auto &v : cap_exec) v.clear();
+  for (auto &v : prev_exec) v.clear();
+  for (auto &p : pred) p = Selection{};
+  std::fill(pred_avail.begin(), pred_avail.end(), false);
+  step = 0;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/prerouter.hpp b/macos/crates/engine-native/src/prerouter.hpp
new file mode 100644
index 0000000..93f0fad
--- /dev/null
+++ b/macos/crates/engine-native/src/prerouter.hpp
@@ -0,0 +1,92 @@
+// Prerouter heads and step-boundary batched prediction (stacked matmul, no per-head sync).
+#pragma once
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "provider.hpp"
+#include "registry.hpp"
+#include "router.hpp"
+#include "safio.hpp"
+
+namespace e0n {
+
+struct HeadW {
+  int owner = -1;
+  std::vector<float> fc1;  // [prh, f_in]
+  std::vector<float> fc2;  // [E, prh]
+  std::vector<float> lin;  // [E, f_in]
+};
+
+class Prerouter {
+ public:
+  // Load from the tier dir (heads from the registry; f_in = hidden + 2×E per head).
+  // hidden is the tier hidden_size (35b 2048 / 8b 1536; same as gate weight columns).
+  static bool load(const ModelFiles& files, const TierRouting& reg, int64_t hidden,
+                   Prerouter& out, std::string& err);
+  // Direct weights (CI / onboarding). Head order is the input order (must be ascending).
+  static bool build(const std::vector<HeadW>& heads, const TierRouting& reg, int64_t hidden,
+                    Prerouter& out, std::string& err);
+
+  Prerouter();
+  ~Prerouter();
+  Prerouter(Prerouter&&) noexcept;
+  Prerouter& operator=(Prerouter&&) noexcept;
+  Prerouter(const Prerouter&) = delete;
+
+  int n_heads() const;
+  int owner(int i) const;  // source layer of head i (ascending)
+
+  // Same stacked weights as host predict, layout = stager._stacked_weights:
+  // fc1 [n,f_in,prh], fc2 [n,prh,E], lin [n,f_in,E].
+  const std::vector<int> &owners() const;
+  const std::vector<float> &stacked_fc1() const;
+  const std::vector<float> &stacked_fc2() const;
+  const std::vector<float> &stacked_lin() const;
+  int64_t stacked_f_in() const;
+  int64_t stacked_prh() const;
+  int64_t stacked_E() const;
+
+  struct Input {
+    const float* hidden = nullptr;           // length = hidden_size (m_in)
+    std::vector<int64_t> cur_exec;           // this token's exec set (ascending)
+    std::vector<int64_t> prev_exec;          // previous token (empty = zero feature)
+  };
+  bool predict(const std::vector<Input>& in, std::vector<Selection>& out_sel,
+               std::vector<std::vector<double>>* out_logits, std::string& err);
+
+  struct Impl;
+
+ private:
+  std::unique_ptr<Impl> impl_;
+};
+
+// Cross-token double-buffer. Fields are public for the forward glue; only
+// reset()/swap() mutate the batch structure.
+struct PrerouterState {
+  TierRouting reg;
+  std::vector<std::vector<float>> cap_hidden;  // per layer: m_in (decode = last position)
+  std::vector<std::vector<int64_t>> cap_exec;  // per layer: this step's exec set
+  std::vector<std::vector<int64_t>> prev_exec;
+  std::vector<Selection> pred;
+  std::vector<bool> pred_avail;  // one-shot consume (cleared after take)
+  long long step = 0;
+
+  void init(const TierRouting& r);
+  void capture(int layer, std::vector<float> m_in, std::vector<int64_t> exec);
+  // Batch predict in ascending source-layer order → pred[layer+1] only if that
+  // layer is in predicted(); 35b layer 38 → consumer 39 is computed then discarded.
+  bool run_step_boundary(Prerouter& pr, std::vector<std::vector<double>>* diag_logits,
+                         std::string& err);
+  void swap();
+  bool take_pred(int layer, Selection& out);
+  void reset();
+};
+
+// Prefetch plan: predicted consumer layers that already have a prediction,
+// (layer, expert) with layer then expert ascending. Live-router layers are omitted
+// (their set is computed at the gate).
+Step live_next_step(const PrerouterState &st);
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/provider.hpp b/macos/crates/engine-native/src/provider.hpp
new file mode 100644
index 0000000..10a22a8
--- /dev/null
+++ b/macos/crates/engine-native/src/provider.hpp
@@ -0,0 +1,37 @@
+// Per-step (layer, expert) load plan. Implementations may swap replay for live routing.
+#pragma once
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace e0n {
+
+struct Step {
+  // Load order: (layer, expert), experts ascending within a layer.
+  std::vector<std::pair<int, int64_t>> experts;
+};
+
+class StepPlanProvider {
+ public:
+  virtual ~StepPlanProvider() = default;
+  // Next step; nullptr = end. The Step stays valid until the next next()/reset().
+  virtual const Step *next() = 0;
+  virtual void reset() = 0;
+  virtual const char *source() const = 0;  // "routing-replay" | "prerouter" | ...
+};
+
+// Rebuild decode order from fixtures/<tier>/routing/records.jsonl (`kind=router`)
+// grouped by (prompt_id, step). 35b samples cover layers 0–6; remaining layers
+// reuse the matching-size expert set via L%7. Counted in steps_total().
+bool make_routing_replay(const std::string &fixtures_tier_dir, const std::string &prompt_id,
+                         int num_layers, int64_t num_experts,
+                         std::unique_ptr<StepPlanProvider> &out, std::string &err);
+
+namespace replay_meta {
+uint64_t stale_regions(const StepPlanProvider *p);  // intervals that reuse a prefill set
+uint64_t total_regions(const StepPlanProvider *p);
+const std::vector<int> *step_ids(const StepPlanProvider *p);  // original step numbers
+}  // namespace replay_meta
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/qgemm.cpp b/macos/crates/engine-native/src/qgemm.cpp
new file mode 100644
index 0000000..06b3c8e
--- /dev/null
+++ b/macos/crates/engine-native/src/qgemm.cpp
@@ -0,0 +1,364 @@
+#include "qgemm.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+
+#include <fcntl.h>
+#include <unistd.h>
+
+#include <Foundation/Foundation.hpp>
+#include <Metal/Metal.hpp>
+#include "mlx/backend/metal/device.h"
+
+#include "expertmap.hpp"
+#include "fwd_common.hpp"
+
+using mlx::core::array;
+using mlx::core::Device;
+
+namespace e0n {
+namespace {
+
+constexpr const char *kMslKernel = R"metal(
+#include <metal_stdlib>
+using namespace metal;
+inline float bf2f(ushort v) {
+  uint bits = uint(v) << 16;
+  return *(thread float*)&bits;
+}
+kernel void e0n_dequant_gemm(
+    const device float* x [[buffer(0)]],
+    const device uint* w [[buffer(1)]],
+    const device ushort* s [[buffer(2)]],
+    const device ushort* b [[buffer(3)]],
+    const device uint* idxs [[buffer(4)]],
+    device float* y [[buffer(5)]],
+    constant uint& O [[buffer(6)]],
+    constant uint& I [[buffer(7)]],
+    constant uint& T [[buffer(8)]],
+    uint3 g [[threadgroup_position_in_grid]],
+    uint3 t [[thread_position_in_threadgroup]]) {
+  const uint G = I / 64;
+  const uint o = g.x * 32u + t.x;
+  const uint tk = g.y;
+  const uint e = g.z;
+  if (o >= O || tk >= T) return;
+  const uint expert = idxs[e];
+  const uint wrow = (expert * O + o) * (I / 8);
+  const uint srow = (expert * O + o) * G;
+  float acc = 0.f;
+  for (uint k = 0; k < I; k++) {
+    const uint word = w[wrow + (k >> 3)];
+    const uint nib = (word >> (4u * (k & 7u))) & 0xFu;
+    acc += x[tk * I + k] * (bf2f(s[srow + (k >> 6)]) * float(nib) + bf2f(b[srow + (k >> 6)]));
+  }
+  y[(e * T + tk) * O + o] = acc;
+}
+)metal";
+
+inline float bf16f(uint16_t v) {
+  uint32_t bits = static_cast<uint32_t>(v) << 16;
+  float f;
+  std::memcpy(&f, &bits, 4);
+  return f;
+}
+
+bool read_full_tensor(const ModelFiles &files, const std::string &prefix,
+                      const std::string &name, std::vector<uint8_t> &out, std::string &err) {
+  const TensorInfo *t = files.find(prefix + "." + name);
+  if (!t) { err = "missing tensor " + prefix + "." + name; return false; }
+  out.resize(t->end - t->begin);
+  int fd = ::open(files.shard_path(t->shard).c_str(), O_RDONLY);
+  if (fd < 0) { err = "open failed " + t->shard; return false; }
+  size_t done = 0;
+  while (done < out.size()) {
+    ssize_t got = ::pread(fd, out.data() + done, out.size() - done, t->begin + done);
+    if (got <= 0) { ::close(fd); err = "pread failed " + name; return false; }
+    done += got;
+  }
+  ::close(fd);
+  return true;
+}
+
+std::string fused_prefix(const ModelFiles &files, int layer) {
+  const bool qwen = files.find("language_model.model.layers.0.mlp.gate.weight") != nullptr;
+  const std::string base =
+      qwen ? "language_model.model.layers." + std::to_string(layer)
+           : "model.layers." + std::to_string(layer);
+  for (const std::string &cand : {base + ".mlp.experts", base + ".mlp.switch_mlp", base + ".mlp"}) {
+    if (files.find(cand + ".gate_proj.weight")) return cand;
+  }
+  return "";
+}
+
+bool msl_proj(MTL::CommandQueue *q, MTL::ComputePipelineState *pso, MTL::Buffer *xbuf,
+              MTL::Buffer *w, MTL::Buffer *s, MTL::Buffer *b, MTL::Buffer *idx, MTL::Buffer *y,
+              int ES, int O, int I, std::string &err) {
+  uint32_t uO = O, uI = I, uT = 1;
+  MTL::CommandBuffer *cb = q->commandBuffer();
+  MTL::ComputeCommandEncoder *enc = cb->computeCommandEncoder();
+  enc->setComputePipelineState(pso);
+  enc->setBuffer(xbuf, 0, 0);
+  enc->setBuffer(w, 0, 1);
+  enc->setBuffer(s, 0, 2);
+  enc->setBuffer(b, 0, 3);
+  enc->setBuffer(idx, 0, 4);
+  enc->setBuffer(y, 0, 5);
+  enc->setBytes(&uO, sizeof(uO), 6);
+  enc->setBytes(&uI, sizeof(uI), 7);
+  enc->setBytes(&uT, sizeof(uT), 8);
+  enc->dispatchThreads(MTL::Size::Make((O + 31) / 32 * 32, 1, ES), MTL::Size::Make(32, 1, 1));
+  enc->endEncoding();
+  cb->commit();
+  cb->waitUntilCompleted();
+  if (cb->error() != nullptr) {
+    err = "MSL dispatch error";
+    return false;
+  }
+  return true;
+}
+
+}  // namespace
+
+bool dims_of(const ModelFiles &files, int layer, ExpertMapDims &out, std::string &err) {
+  const std::string prefix = fused_prefix(files, layer);
+  if (prefix.empty()) { err = "layer has no fused prefix"; return false; }
+  const TensorInfo *g = files.find(prefix + ".gate_proj.weight");
+  const TensorInfo *d = files.find(prefix + ".down_proj.weight");
+  if (!g || !d) { err = "missing gate/down weights"; return false; }
+  int64_t E = 0;
+  if (!files.config_int("num_experts", E)) { err = "missing num_experts"; return false; }
+  auto rows_of = [&](const TensorInfo *t) {
+    return t->shape.size() == 3 ? t->shape[1] : t->shape[0] / E;
+  };
+  out.E = E;
+  out.hid = static_cast<int>(g->shape.back() * 8);
+  out.ffn = static_cast<int>(rows_of(g));
+  return true;
+}
+
+struct QuantLayer::Impl {
+  GemmPath path = GemmPath::kMlxGatherQmm;
+  std::vector<array> w, s, b;
+  std::vector<std::shared_ptr<std::vector<uint8_t>>> keepers;
+  MTL::Device *mtl = nullptr;
+  MTL::CommandQueue *q = nullptr;
+  MTL::ComputePipelineState *pso = nullptr;
+  MTL::Buffer *bw[3] = {}, *bs[3] = {}, *bb[3] = {};
+  ~Impl() {
+    for (int i = 0; i < 3; i++) {
+      if (bw[i]) bw[i]->release();
+      if (bs[i]) bs[i]->release();
+      if (bb[i]) bb[i]->release();
+    }
+  }
+};
+
+QuantLayer::QuantLayer(std::unique_ptr<Impl> i) : impl_(std::move(i)) {}
+QuantLayer::~QuantLayer() = default;
+
+std::unique_ptr<QuantLayer> QuantLayer::load(const ModelFiles &files, const ExpertMapDims &dims,
+                                             int layer, GemmPath path, std::string &err) {
+  const std::string prefix = fused_prefix(files, layer);
+  if (prefix.empty()) { err = "layer has no fused prefix"; return nullptr; }
+
+  auto l = std::unique_ptr<QuantLayer>(new QuantLayer(std::unique_ptr<Impl>(new Impl())));
+  l->e_ = static_cast<int>(dims.E);
+  l->hid_ = dims.hid;
+  l->ffn_ = dims.ffn;
+  auto &im = *l->impl_;
+  im.path = path;
+
+  const int E = static_cast<int>(dims.E), HID = dims.hid, FFN = dims.ffn;
+  const char *parts[3] = {"gate_proj", "up_proj", "down_proj"};
+  const int in_dim[3] = {HID, HID, FFN};
+  const int out_rows[3] = {FFN, FFN, HID};
+  std::vector<std::vector<uint8_t>> wb(3), sb(3), bb(3);
+  for (int p = 0; p < 3; p++) {
+    if (!read_full_tensor(files, prefix, std::string(parts[p]) + ".weight", wb[p], err) ||
+        !read_full_tensor(files, prefix, std::string(parts[p]) + ".scales", sb[p], err) ||
+        !read_full_tensor(files, prefix, std::string(parts[p]) + ".biases", bb[p], err))
+      return nullptr;
+  }
+  if (path == GemmPath::kMlxGatherQmm) {
+    for (int p = 0; p < 3; p++) {
+      auto kw = std::make_shared<std::vector<uint8_t>>(std::move(wb[p]));
+      auto ks = std::make_shared<std::vector<uint8_t>>(std::move(sb[p]));
+      auto kb = std::make_shared<std::vector<uint8_t>>(std::move(bb[p]));
+      const int Gi = in_dim[p] / 64;
+      im.w.push_back(array(kw->data(), {E, out_rows[p], in_dim[p] / 8}, mlx::core::uint32,
+                           [kw](void *) {}));
+      im.s.push_back(array(ks->data(), {E, out_rows[p], Gi}, mlx::core::bfloat16, [ks](void *) {}));
+      im.b.push_back(array(kb->data(), {E, out_rows[p], Gi}, mlx::core::bfloat16, [kb](void *) {}));
+      im.keepers.insert(im.keepers.end(), {kw, ks, kb});
+    }
+  } else {
+    NS::AutoreleasePool *pool = NS::AutoreleasePool::alloc()->init();
+    auto &mdev = mlx::core::metal::device(Device::gpu);
+    im.mtl = mdev.mtl_device();
+    im.q = mdev.get_queue(mlx::core::default_stream(Device::gpu));
+    MTL::Library *lib = mdev.get_library("e0n_dequant_gemm", [] { return std::string(kMslKernel); });
+    im.pso = lib ? mdev.get_kernel("e0n_dequant_gemm", lib) : nullptr;
+    if (!im.pso) { pool->release(); err = "MSL kernel compile failed"; return nullptr; }
+    auto nbuf = [&](const void *d, size_t bytes) {
+      return im.mtl->newBuffer(d, bytes, MTL::ResourceStorageModeShared);
+    };
+    for (int p = 0; p < 3; p++) {
+      im.bw[p] = nbuf(wb[p].data(), wb[p].size());
+      im.bs[p] = nbuf(sb[p].data(), sb[p].size());
+      im.bb[p] = nbuf(bb[p].data(), bb[p].size());
+    }
+    pool->release();
+  }
+  return l;
+}
+
+bool QuantLayer::moe_proxy(std::vector<float> &x,
+                           const std::vector<std::vector<uint32_t>> &layer_indices,
+                           std::string &err) {
+  if (layer_indices.empty()) return true;
+  const int HID = hid_, FFN = ffn_;
+  auto &im = *impl_;
+  if (im.path == GemmPath::kMlxGatherQmm) {
+    auto xg = array(x.begin(), {1, HID});
+    for (const auto &indices : layer_indices) {
+      const int ES = (int)indices.size();
+      auto ia = array(indices.begin(), {ES});
+      auto a = mlx::core::gather_qmm(xg, im.w[0], im.s[0], im.b[0], std::nullopt, ia, true, 64,
+                                     4, "affine");
+      auto b = mlx::core::gather_qmm(xg, im.w[1], im.s[1], im.b[1], std::nullopt, ia, true, 64,
+                                     4, "affine");
+      auto ga = mlx::core::sum(a, {0}, false);
+      auto h = (ga * mlx::core::sigmoid(ga)) * mlx::core::sum(b, {0}, false);
+      auto d = mlx::core::gather_qmm(h, im.w[2], im.s[2], im.b[2], std::nullopt, ia, true, 64,
+                                     4, "affine");
+      xg = xg + mlx::core::sum(d, {0}, false);
+    }
+    auto cpu = mlx::core::copy(xg, mlx::core::default_stream(Device::cpu));
+    mlx::core::eval(cpu);
+    const float *p = cpu.data<float>();
+    x.assign(p, p + HID);
+    return true;
+  }
+  const int MAX_ES = (int)layer_indices[0].size();
+  NS::AutoreleasePool *pool = NS::AutoreleasePool::alloc()->init();
+  MTL::Buffer *bidx = im.mtl->newBuffer(MAX_ES * 4, MTL::ResourceStorageModeShared);
+  MTL::Buffer *bx = im.mtl->newBuffer((size_t)HID * 4, MTL::ResourceStorageModeShared);
+  MTL::Buffer *bg = im.mtl->newBuffer((size_t)MAX_ES * FFN * 4, MTL::ResourceStorageModeShared);
+  MTL::Buffer *bu = im.mtl->newBuffer((size_t)MAX_ES * FFN * 4, MTL::ResourceStorageModeShared);
+  MTL::Buffer *bd = im.mtl->newBuffer((size_t)MAX_ES * HID * 4, MTL::ResourceStorageModeShared);
+  bool ok = true;
+  std::vector<float> h(FFN);
+  for (const auto &indices : layer_indices) {
+    const int ES = (int)indices.size();
+    ::memcpy(bidx->contents(), indices.data(), ES * 4);
+    ::memcpy(bx->contents(), x.data(), HID * 4);
+    ok = msl_proj(im.q, im.pso, bx, im.bw[0], im.bs[0], im.bb[0], bidx, bg, ES, FFN, HID, err) &&
+         msl_proj(im.q, im.pso, bx, im.bw[1], im.bs[1], im.bb[1], bidx, bu, ES, FFN, HID, err);
+    if (!ok) break;
+    const float *pg = (const float *)bg->contents();
+    const float *pu = (const float *)bu->contents();
+    for (int o = 0; o < FFN; o++) {
+      double gs = 0, us = 0;
+      for (int e = 0; e < ES; e++) {
+        float gv = pg[(size_t)e * FFN + o];
+        gs += (double)gv / (1.0 + std::exp(-(double)gv));  // silu
+        us += pu[(size_t)e * FFN + o];
+      }
+      h[o] = (float)(gs * us);
+    }
+    ::memcpy(bx->contents(), h.data(), FFN * 4);
+    ok = msl_proj(im.q, im.pso, bx, im.bw[2], im.bs[2], im.bb[2], bidx, bd, ES, HID, FFN, err);
+    if (!ok) break;
+    const float *pd = (const float *)bd->contents();
+    for (int k = 0; k < HID; k++) {
+      double acc = 0;
+      for (int e = 0; e < ES; e++) acc += pd[(size_t)e * HID + k];
+      x[k] += (float)acc;
+    }
+  }
+  bidx->release(); bx->release(); bg->release(); bu->release(); bd->release();
+  pool->release();
+  return ok;
+}
+
+std::vector<double> fp64_qmm_ref(const float *x, int T, int I, const uint32_t *W,
+                                 const uint16_t *S, const uint16_t *B, int O, int G,
+                                 const std::vector<int64_t> &sel) {
+  std::vector<double> ref;
+  ref.reserve(sel.size() * T * O);
+  for (int64_t e : sel) {
+    const uint32_t *we = W + static_cast<size_t>(e) * O * (I / 8);
+    const uint16_t *se = S + static_cast<size_t>(e) * O * G;
+    const uint16_t *be = B + static_cast<size_t>(e) * O * G;
+    for (int t = 0; t < T; t++) {
+      for (int o = 0; o < O; o++) {
+        double acc = 0;
+        for (int k = 0; k < I; k++) {
+          const uint32_t word = we[static_cast<size_t>(o) * (I / 8) + k / 8];
+          const int nib = (word >> (4 * (k % 8))) & 0xF;
+          acc += (double)x[(size_t)t * I + k] *
+                 ((double)bf16f(se[o * G + k / 64]) * nib + (double)bf16f(be[o * G + k / 64]));
+        }
+        ref.push_back(acc);
+      }
+    }
+  }
+  return ref;
+}
+
+bool gather_qmm_mlx(const uint32_t *W, const uint16_t *S, const uint16_t *B, int E_all, int O,
+                    int I, int G, const float *x, int T, const std::vector<int64_t> &sel,
+                    std::vector<float> &y, std::string &err) {
+  {
+    auto &cx = exec_counters();
+    cx.qmm_calls.fetch_add(1, std::memory_order_relaxed);
+    cx.qmm_experts.fetch_add(sel.size(), std::memory_order_relaxed);
+    cx.qmm_weight_bytes.fetch_add(sel.size() * ((size_t)O * (I / 8) * 4 + 2 * (size_t)O * G * 2),
+                                  std::memory_order_relaxed);
+  }
+  std::vector<uint32_t> idx(sel.begin(), sel.end());
+  auto wa = array((void *)W, {E_all, O, I / 8}, mlx::core::uint32, [](void *) {});
+  auto sa = array((void *)S, {E_all, O, G}, mlx::core::bfloat16, [](void *) {});
+  auto ba = array((void *)B, {E_all, O, G}, mlx::core::bfloat16, [](void *) {});
+  auto xa = array(x, {T, I});
+  auto ia = array(idx.begin(), {static_cast<int>(idx.size())});
+  try {
+    auto ya = mlx::core::gather_qmm(xa, wa, sa, ba, std::nullopt, ia, true, 64, 4, "affine");
+    auto ycpu = mlx::core::copy(ya, mlx::core::default_stream(Device::cpu));
+    exec_counters().host_syncs.fetch_add(1, std::memory_order_relaxed);
+    mlx::core::eval(ycpu);
+    const float *p = ycpu.data<float>();
+    y.assign(p, p + sel.size() * T * O);
+  } catch (const std::exception &e) {
+    err = e.what();
+    return false;
+  }
+  return true;
+}
+
+mlx::core::array gather_qmm_graph(const uint32_t *W, const uint16_t *S, const uint16_t *B,
+                                   int E_all, int O, int I, int G, const mlx::core::array &xa,
+                                   int T, const std::vector<int64_t> &sel) {
+  auto &cx = exec_counters();
+  cx.qmm_calls.fetch_add(1, std::memory_order_relaxed);
+  cx.qmm_experts.fetch_add(sel.size(), std::memory_order_relaxed);
+  cx.qmm_weight_bytes.fetch_add(sel.size() * ((size_t)O * (I / 8) * 4 + 2 * (size_t)O * G * 2),
+                                std::memory_order_relaxed);
+  std::vector<uint32_t> idx(sel.begin(), sel.end());
+  auto wa = array((void *)W, {E_all, O, I / 8}, mlx::core::uint32, [](void *) {});
+  auto sa = array((void *)S, {E_all, O, G}, mlx::core::bfloat16, [](void *) {});
+  auto ba = array((void *)B, {E_all, O, G}, mlx::core::bfloat16, [](void *) {});
+  auto ia = array(idx.begin(), {static_cast<int>(idx.size())});
+  return mlx::core::gather_qmm(xa, wa, sa, ba, std::nullopt, ia, true, 64, 4, "affine");
+}
+
+double max_rel(const std::vector<float> &got, const std::vector<double> &ref) {  double norm = 0, mr = 0;
+  for (double v : ref) norm = std::max(norm, std::fabs(v));
+  for (size_t i = 0; i < ref.size() && i < got.size(); i++)
+    mr = std::max(mr, std::fabs((double)got[i] - ref[i]) / (norm + 1e-12));
+  return mr;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/qgemm.hpp b/macos/crates/engine-native/src/qgemm.hpp
new file mode 100644
index 0000000..eef693a
--- /dev/null
+++ b/macos/crates/engine-native/src/qgemm.hpp
@@ -0,0 +1,62 @@
+// Grouped quantized matmul via mlx gather_qmm, with an optional custom MSL path.
+#pragma once
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <mlx/mlx.h>
+
+#include "safio.hpp"
+
+namespace e0n {
+
+enum class GemmPath { kMlxGatherQmm, kMslCustom };
+
+struct ExpertMapDims {
+  int64_t E;
+  int hid;
+  int ffn;
+};
+bool dims_of(const ModelFiles &files, int layer, ExpertMapDims &out, std::string &err);
+
+class QuantLayer {
+ public:
+  static std::unique_ptr<QuantLayer> load(const ModelFiles &files, const ExpertMapDims &dims,
+                                          int layer, GemmPath path, std::string &err);
+  ~QuantLayer();
+
+  int E() const { return e_; }
+  int HID() const { return hid_; }
+  int FFN() const { return ffn_; }
+
+  // swiglu + down + residual; x is host [HID] float, updated in place.
+  // layer_indices = per-layer expert sets (one proxy compute per layer).
+  bool moe_proxy(std::vector<float> &x, const std::vector<std::vector<uint32_t>> &layer_indices,
+                 std::string &err);
+
+ private:
+  struct Impl;
+  explicit QuantLayer(std::unique_ptr<Impl> i);
+  std::unique_ptr<Impl> impl_;
+  int e_ = 0, hid_ = 0, ffn_ = 0;
+};
+
+std::vector<double> fp64_qmm_ref(const float *x, int T, int I, const uint32_t *W,
+                                 const uint16_t *S, const uint16_t *B, int O, int G,
+                                 const std::vector<int64_t> &sel);
+
+bool gather_qmm_mlx(const uint32_t *W, const uint16_t *S, const uint16_t *B, int E_all, int O,
+                    int I, int G, const float *x, int T, const std::vector<int64_t> &sel,
+                    std::vector<float> &y, std::string &err);
+
+// Lazy graph form: x and the return are mlx arrays; this call does not sync.
+// Host weight pointers use a no-op deleter; the caller must keep staging alive until
+// the layer-end eval. Output shape {sel.size(), T, O}.
+mlx::core::array gather_qmm_graph(const uint32_t *W, const uint16_t *S, const uint16_t *B,
+                                   int E_all, int O, int I, int G, const mlx::core::array &xa,
+                                   int T, const std::vector<int64_t> &sel);
+
+double max_rel(const std::vector<float> &got, const std::vector<double> &ref);
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/registry.cpp b/macos/crates/engine-native/src/registry.cpp
new file mode 100644
index 0000000..7f22c80
--- /dev/null
+++ b/macos/crates/engine-native/src/registry.cpp
@@ -0,0 +1,85 @@
+#include "registry.hpp"
+
+namespace e0n {
+
+std::vector<int> TierRouting::predicted() const {
+  std::vector<int> v;
+  for (int l = first_head_consumer; l <= last_head_consumer; l++) v.push_back(l);
+  return v;
+}
+
+std::vector<int> TierRouting::routed() const {
+  std::vector<int> v;
+  for (int l = dense_prefix; l < first_head_consumer; l++) v.push_back(l);
+  // 35b: last layer L39 does not consume predictions; 8b L23 == last_head_consumer
+  // is consumed. The last-layer exception is qwen-family only.
+  if (family == RouterFamily::kSoftmaxTopk) v.push_back(num_layers - 1);
+  return v;
+}
+
+bool tier_routing(const std::string &tier_id, TierRouting &out) {
+  if (tier_id == "edge0-35b") {
+    out = TierRouting{};
+    out.tier = tier_id;
+    out.family = RouterFamily::kSoftmaxTopk;
+    out.num_experts = 256;
+    out.top_k = 4;
+    out.norm_topk_prob = true;
+    out.n_group = 0; out.topk_group = 0; out.routed_scaling = 1.0;
+    out.expert_bias = false;
+    out.num_layers = 40;
+    out.first_head_consumer = 7; out.last_head_consumer = 38;
+    out.first_owner = 6; out.last_owner = 38;
+    out.dense_prefix = 0;
+    out.prerouter_hidden = 512;
+    out.head_file = "prerouter_edge0_35b.safetensors";
+    out.lora_file = "lora_edge0_35b.safetensors";
+    out.heads_baked_in_model = false;
+    out.head_feature_teacher = false;  // qwen: block patch captures the executed set
+    out.prefill_chunk = 2048;          // models/edge0_35b/__init__.py:63
+    out.hot_window = 4;                // :64
+    out.prefill_hot = 32;              // staged_k4(options.py:134)
+    out.full_layer_prefill = false;    // options.py:133
+    out.async_eval_prefill = false;    // qwen.py:176 gate=full_layer ⇒ False
+    return true;
+  }
+  if (tier_id == "edge0-8b") {
+    out = TierRouting{};
+    out.tier = tier_id;
+    out.family = RouterFamily::kSigmoidGroup;
+    out.num_experts = 128;
+    out.top_k = 8;
+    out.norm_topk_prob = true;
+    out.n_group = 8; out.topk_group = 4; out.routed_scaling = 2.5;
+    out.expert_bias = true;
+    out.num_layers = 24;
+    out.first_head_consumer = 8; out.last_head_consumer = 23;
+    out.first_owner = 7; out.last_owner = 22;
+    out.dense_prefix = 1;
+    out.prerouter_hidden = 512;
+    out.head_file = "prerouter_edge0_8b.safetensors";
+    out.lora_file = "lora_edge0_8b.safetensors";
+    out.heads_baked_in_model = true;
+    out.head_feature_teacher = true;  // ling: teacher top-8 feeds features
+    out.prefill_chunk = 2048;          // models/edge0_8b/__init__.py:71
+    out.hot_window = 1;
+    out.prefill_hot = 0;               // prod_k8(options.py:189)
+    out.full_layer_prefill = true;     // options.py:189
+    out.async_eval_prefill = true;     // ling.py:217 gate=full_layer ⇒ True
+    return true;
+  }
+  return false;
+}
+
+std::string routing_config_mismatch(const TierRouting &r, int64_t cfg_num_experts,
+                                    int64_t cfg_num_layers) {
+  if (cfg_num_experts > 0 && cfg_num_experts != r.num_experts)
+    return "num_experts: config=" + std::to_string(cfg_num_experts) +
+           " registry=" + std::to_string(r.num_experts);
+  if (cfg_num_layers > 0 && cfg_num_layers != r.num_layers)
+    return "num_layers: config=" + std::to_string(cfg_num_layers) +
+           " registry=" + std::to_string(r.num_layers);
+  return "";
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/registry.hpp b/macos/crates/engine-native/src/registry.hpp
new file mode 100644
index 0000000..faaf602
--- /dev/null
+++ b/macos/crates/engine-native/src/registry.hpp
@@ -0,0 +1,48 @@
+// Runtime routing facts that are not in config.json (config is only a cross-check).
+#pragma once
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace e0n {
+
+enum class RouterFamily { kSoftmaxTopk, kSigmoidGroup };
+
+struct TierRouting {
+  std::string tier;              // "edge0-8b" / "edge0-35b" (= tier_id)
+  RouterFamily family;
+  int64_t num_experts;
+  int top_k;                     // runtime (35b=4, not config's 8)
+  bool norm_topk_prob;           // true on both tiers
+  int n_group, topk_group;
+  double routed_scaling;
+  bool expert_bias;              // live gate scores include bias; head path does not
+  int num_layers;
+  int first_head_consumer;       // 35b 7 / 8b 8
+  int last_head_consumer;        // 35b 38 / 8b 23 (35b layer 39 is not a consumer)
+  int first_owner, last_owner;   // 35b 6/38 / 8b 7/22
+  int dense_prefix;              // leading dense (no MoE) layers: 8b=1 (L0); 35b=0
+  int prerouter_hidden;          // 512 on both live-head tiers
+  std::string head_file;
+  std::string lora_file;
+  bool heads_baked_in_model;
+  // Head-feature one-hot: false = executed set (qwen); true = teacher top-k (ling).
+  bool head_feature_teacher;
+
+  int prefill_chunk;          // 2048 both wrappers
+  int hot_window;             // 35b=4 / 8b=1
+  int prefill_hot;            // 35b staged_k4=32 / 8b prod_k8=0
+  bool full_layer_prefill;    // 35b=false / 8b=true
+  bool async_eval_prefill;    // = prefill_multi and full_layer
+
+  std::vector<int> predicted() const;
+  std::vector<int> routed() const;     // decode-time live-router layers, ascending
+};
+
+bool tier_routing(const std::string &tier_id, TierRouting &out);
+
+// Empty string = match; non-empty = readable mismatch (caller decides whether to reject).
+std::string routing_config_mismatch(const TierRouting &r, int64_t cfg_num_experts,
+                                    int64_t cfg_num_layers);
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/router.cpp b/macos/crates/engine-native/src/router.cpp
new file mode 100644
index 0000000..0ff110a
--- /dev/null
+++ b/macos/crates/engine-native/src/router.cpp
@@ -0,0 +1,112 @@
+#include "router.hpp"
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+
+namespace e0n {
+namespace {
+
+// Deterministic top-K: score descending, ties by id ascending (argpartition order is not the contract).
+std::vector<int64_t> topk_ids(const std::vector<double> &v, int k) {
+  std::vector<int64_t> ids(v.size());
+  for (size_t i = 0; i < ids.size(); i++) ids[i] = static_cast<int64_t>(i);
+  std::sort(ids.begin(), ids.end(), [&](int64_t a, int64_t b) {
+    if (v[a] != v[b]) return v[a] > v[b];
+    return a < b;
+  });
+  ids.resize(static_cast<size_t>(k));
+  std::sort(ids.begin(), ids.end());  // canonical: ascending
+  return ids;
+}
+
+}  // namespace
+
+Selection select_softmax_topk(const double *logits, int E, int k, bool norm) {
+  // precise softmax (fp64) → top-K → renormalize.
+  double mx = logits[0];
+  for (int i = 1; i < E; i++) mx = std::max(mx, logits[i]);
+  std::vector<double> gates(static_cast<size_t>(E));
+  double sum = 0.0;
+  for (int i = 0; i < E; i++) {
+    gates[i] = std::exp(logits[i] - mx);
+    sum += gates[i];
+  }
+  for (auto &g : gates) g /= sum;
+  Selection out;
+  out.inds = topk_ids(gates, k);
+  out.scores.resize(out.inds.size());
+  double ssum = 0.0;
+  for (size_t j = 0; j < out.inds.size(); j++) {
+    out.scores[j] = gates[out.inds[j]];
+    ssum += out.scores[j];
+  }
+  if (norm && ssum > 0.0)
+    for (auto &s : out.scores) s /= ssum;
+  return out;
+}
+
+Selection select_sigmoid_group(const double *logits, const double *bias, int E, int k,
+                               int n_group, int topk_group, double scaling, bool norm) {
+  std::vector<double> scores(static_cast<size_t>(E));
+  for (int i = 0; i < E; i++) scores[i] = 1.0 / (1.0 + std::exp(-logits[i]));
+  std::vector<double> select = scores;
+  if (bias)
+    for (int i = 0; i < E; i++) select[i] += bias[i];
+
+  // Group mask: group score = sum of in-group top-2 select scores; drop the weakest k_drop groups (-inf).
+  const int k_drop = n_group - topk_group;
+  if (k_drop > 0) {
+    const int per = E / n_group;
+    std::vector<double> gscore(static_cast<size_t>(n_group));
+    for (int g = 0; g < n_group; g++) {
+      double m1 = -INFINITY, m2 = -INFINITY;  // top-2 (ties: two max passes)
+      for (int j = 0; j < per; j++) {
+        double s = select[static_cast<size_t>(g * per + j)];
+        if (s >= m1) { m2 = m1; m1 = s; }
+        else if (s > m2) m2 = s;
+      }
+      gscore[g] = m1 + m2;
+    }
+    // Drop the k_drop groups with the smallest scores; ties by group id ascending.
+    std::vector<int> gids(static_cast<size_t>(n_group));
+    for (int g = 0; g < n_group; g++) gids[g] = g;
+    std::sort(gids.begin(), gids.end(), [&](int a, int b) {
+      if (gscore[a] != gscore[b]) return gscore[a] < gscore[b];
+      return a < b;
+    });
+    for (int d = 0; d < k_drop; d++) {
+      int g = gids[d];
+      for (int j = 0; j < per; j++) select[static_cast<size_t>(g * per + j)] = -INFINITY;
+    }
+  }
+
+  Selection out;
+  out.inds = topk_ids(select, k);
+  out.scores.resize(out.inds.size());
+  double ssum = 0.0;
+  for (size_t j = 0; j < out.inds.size(); j++) {
+    out.scores[j] = scores[out.inds[j]];  // weight = raw sigmoid (bias not in weights)
+    ssum += out.scores[j];
+  }
+  if (norm)
+    for (auto &s : out.scores) s /= (ssum + 1e-20);
+  for (auto &s : out.scores) s *= scaling;
+  return out;
+}
+
+void gate_logits(const float *x, int T, int H, const float *W, int E,
+                 std::vector<double> &out) {
+  out.assign(static_cast<size_t>(T) * E, 0.0);
+  for (int t = 0; t < T; t++) {
+    const float *xr = x + static_cast<size_t>(t) * H;
+    for (int e = 0; e < E; e++) {
+      const float *wr = W + static_cast<size_t>(e) * H;
+      double acc = 0.0;
+      for (int h = 0; h < H; h++) acc += static_cast<double>(xr[h]) * wr[h];
+      out[static_cast<size_t>(t) * E + e] = acc;
+    }
+  }
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/router.hpp b/macos/crates/engine-native/src/router.hpp
new file mode 100644
index 0000000..76ed1d3
--- /dev/null
+++ b/macos/crates/engine-native/src/router.hpp
@@ -0,0 +1,25 @@
+// Host fp64 expert selection for the two router families (bit-identical to routing.py).
+#pragma once
+#include <cstdint>
+#include <vector>
+
+namespace e0n {
+
+struct Selection {
+  std::vector<int64_t> inds;
+  std::vector<double> scores;
+};
+
+// logits length = E. k ≤ E. norm = norm_topk_prob.
+Selection select_softmax_topk(const double *logits, int E, int k, bool norm);
+
+// bias may be nullptr (exclude bias). E % n_group == 0 (8b: 128/8=16).
+// topk_group == n_group means no group is dropped.
+Selection select_sigmoid_group(const double *logits, const double *bias, int E, int k,
+                               int n_group, int topk_group, double scaling, bool norm);
+
+// Dense gate: x[T,H] × W[E,H]ᵀ → out[T,E], fp64 accumulate. Live graphs run gate on mlx.
+void gate_logits(const float *x, int T, int H, const float *W, int E,
+                 std::vector<double> &out);
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/routing_replay.cpp b/macos/crates/engine-native/src/routing_replay.cpp
new file mode 100644
index 0000000..e951d2f
--- /dev/null
+++ b/macos/crates/engine-native/src/routing_replay.cpp
@@ -0,0 +1,111 @@
+// Replay provider: `kind=router` records used as-is. 8b: layers 1–23 (L0 is dense);
+// 35b: layers 0–39; top_k comes from the record.
+#include "provider.hpp"
+
+#include <algorithm>
+#include <fstream>
+#include <map>
+
+#include <nlohmann/json.hpp>
+
+using json = nlohmann::json;
+
+namespace e0n {
+namespace {
+
+class RoutingReplay final : public StepPlanProvider {
+ public:
+  std::vector<Step> steps;               // decode step ascending (step 0..N)
+  std::vector<int> step_ids;             // original step numbers
+  size_t i = 0;
+  Step cur;
+  uint64_t stale_regions = 0;
+  uint64_t total_regions = 0;
+
+  const Step *next() override {
+    if (i >= steps.size()) return nullptr;
+    cur = steps[i++];
+    return &cur;
+  }
+  void reset() override { i = 0; }
+  const char *source() const override { return "routing-replay"; }
+};
+
+}  // namespace
+
+bool make_routing_replay(const std::string &fixtures_tier_dir, const std::string &prompt_id,
+                         int num_layers, int64_t num_experts,
+                         std::unique_ptr<StepPlanProvider> &out, std::string &err) {
+  std::ifstream f(fixtures_tier_dir + "/routing/records.jsonl");
+  if (!f) { err = "no routing records: " + fixtures_tier_dir; return false; }
+
+  // (step, layer) → experts; step=-1 is prefill and is not in the decode sequence.
+  std::map<std::pair<int, int>, std::vector<int64_t>> sets;
+  std::map<std::pair<int, int>, bool> is_stale;
+  std::string line;
+  while (std::getline(f, line)) {
+    if (line.empty()) continue;
+    json r = json::parse(line, nullptr, false);
+    if (r.is_discarded() || r.value("kind", "") != "router") continue;
+    if (r.value("prompt_id", "") != prompt_id) continue;
+    int step = r.value("step", -2);
+    if (step < 0) continue;  // decode path only
+    int layer = r.value("layer", -1);
+    if (layer < 0 || layer >= num_layers) continue;  // 8b L0 has no records
+    std::vector<int64_t> es;
+    for (auto &e : r.at("experts")) es.push_back(e.get<int64_t>());
+    std::sort(es.begin(), es.end());
+    bool dup = es.size() > 1 && std::adjacent_find(es.begin(), es.end()) != es.end();
+    if (dup) es.erase(std::unique(es.begin(), es.end()), es.end());
+    auto key = std::make_pair(step, layer);
+    if (sets.count(key)) continue;  // first record for a key wins
+    sets[key] = std::move(es);
+    is_stale[key] = r.value("routing_source", "").find("stale") != std::string::npos;
+  }
+  if (sets.empty()) {
+    err = "prompt " + prompt_id + " has no decode routing records for this tier";
+    return false;
+  }
+
+  auto rep = std::make_unique<RoutingReplay>();
+  int cur_step = -1;
+  for (auto &[key, es] : sets) {
+    auto [step, layer] = key;
+    for (int64_t e : es) {
+      if (e < 0 || e >= num_experts) {
+        err = "recorded expert id out of range: " + prompt_id + " L" + std::to_string(layer) +
+              " E" + std::to_string(e);
+        return false;
+      }
+    }
+    if (step != cur_step) {
+      rep->steps.emplace_back();
+      rep->step_ids.push_back(step);
+      cur_step = step;
+    }
+    auto &st = rep->steps.back();
+    for (int64_t e : es) st.experts.emplace_back(layer, e);
+    rep->total_regions += es.size();
+    if (is_stale[key]) rep->stale_regions += es.size();
+  }
+  // map is ordered by (step, layer); experts within a layer are already sorted.
+  out = std::move(rep);
+  return true;
+}
+
+namespace replay_meta {
+uint64_t stale_regions(const StepPlanProvider *p) {
+  auto *r = dynamic_cast<const RoutingReplay *>(p);
+  return r ? r->stale_regions : 0;
+}
+uint64_t total_regions(const StepPlanProvider *p) {
+  auto *r = dynamic_cast<const RoutingReplay *>(p);
+  return r ? r->total_regions : 0;
+}
+const std::vector<int> *step_ids(const StepPlanProvider *p) {
+  auto *r = dynamic_cast<const RoutingReplay *>(p);
+  return r ? &r->step_ids : nullptr;
+}
+}  // namespace replay_meta
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/safio.cpp b/macos/crates/engine-native/src/safio.cpp
new file mode 100644
index 0000000..2658e86
--- /dev/null
+++ b/macos/crates/engine-native/src/safio.cpp
@@ -0,0 +1,242 @@
+#include "safio.hpp"
+
+#include <fcntl.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <sstream>
+
+#include <nlohmann/json.hpp>
+
+using json = nlohmann::json;
+
+namespace e0n {
+namespace {
+
+uint64_t file_size_or(const std::string& p, uint64_t fallback) {
+  struct stat st;
+  return ::stat(p.c_str(), &st) == 0 ? static_cast<uint64_t>(st.st_size) : fallback;
+}
+
+// u64 length + JSON; data-region offsets include the prefix.
+bool parse_shard_header(const std::string& path, std::map<std::string, TensorInfo>& out,
+                        std::string& err) {
+  int fd = ::open(path.c_str(), O_RDONLY);
+  if (fd < 0) { err = "open failed: " + path; return false; }
+  uint64_t n = 0;
+  if (::pread(fd, &n, 8, 0) != 8 || n == 0 || n > (1ull << 28)) {
+    ::close(fd);
+    err = "invalid header length: " + path;
+    return false;
+  }
+  std::string js(n, '\0');
+  uint64_t done = 0;
+  while (done < n) {
+    ssize_t got = ::pread(fd, js.data() + done, n - done, 8 + done);
+    if (got <= 0) { ::close(fd); err = "header read failed: " + path; return false; }
+    done += got;
+  }
+  ::close(fd);
+  json h;
+  try {
+    h = json::parse(js);
+  } catch (const std::exception& e) {
+    err = "header JSON parse failed " + path + ": " + e.what();
+    return false;
+  }
+  const uint64_t data_start = 8 + n;
+  std::string base = path.substr(path.find_last_of('/') + 1);
+  for (auto& [k, v] : h.items()) {
+    if (k == "__metadata__" || k == "version" || !v.is_object()) continue;
+    if (!v.contains("data_offsets")) continue;
+    TensorInfo t;
+    t.dtype = v.at("dtype").get<std::string>();
+    t.shape = v.at("shape").get<std::vector<int64_t>>();
+    auto offs = v.at("data_offsets").get<std::vector<uint64_t>>();
+    t.begin = data_start + offs[0];
+    t.end = data_start + offs[1];
+    t.shard = base;
+    out.emplace(k, std::move(t));
+  }
+  return true;
+}
+
+}  // namespace
+
+bool ModelFiles::open(const std::string& dir, std::string& err) {
+  dir_ = dir.empty() || dir.back() == '/' ? dir : dir + "/";
+  std::error_code ec;
+  if (!std::filesystem::is_directory(dir_, ec)) { err = "not a directory: " + dir_; return false; }
+
+  std::vector<std::string> shards;
+  for (auto& e : std::filesystem::directory_iterator(dir_, ec)) {
+    std::string p = e.path().string();
+    if (p.size() > 12 && p.compare(p.size() - 12, 12, ".safetensors") == 0) shards.push_back(p);
+  }
+  std::sort(shards.begin(), shards.end());
+  if (shards.empty()) { err = "no .safetensors in directory: " + dir_; return false; }
+
+  for (auto& s : shards) {
+    std::string one_err;
+    if (!parse_shard_header(s, tensors_, one_err)) { err = one_err; return false; }
+    shard_sizes_[s.substr(s.find_last_of('/') + 1)] = file_size_or(s, 0);
+  }
+
+  std::ifstream idx(dir_ + "model.safetensors.index.json");
+  if (idx) {
+    try {
+      json wm = json::parse(idx).at("weight_map");
+      for (auto& [name, shard] : wm.items()) {
+        auto it = tensors_.find(name);
+        if (it != tensors_.end()) it->second.shard = shard.get<std::string>();
+      }
+    } catch (...) { /* corrupt index: keep header self-assignment */ }
+  }
+
+  std::ifstream cfg(dir_ + "config.json");
+  if (cfg) {
+    std::stringstream ss;
+    ss << cfg.rdbuf();
+    config_json_ = ss.str();
+  }
+  return true;
+}
+
+const TensorInfo* ModelFiles::find(const std::string& name) const {
+  auto it = tensors_.find(name);
+  return it == tensors_.end() ? nullptr : &it->second;
+}
+
+std::string ModelFiles::shard_path(const std::string& shard_name) const {
+  return dir_ + shard_name;
+}
+
+std::vector<std::string> ModelFiles::shard_names() const {
+  std::vector<std::string> v;
+  for (auto &[k, n] : shard_sizes_) v.push_back(k);
+  return v;
+}
+
+bool ModelFiles::read_tensor(const std::string& tensor_name, std::vector<uint8_t>& out,
+                             std::string& err) const {
+  const TensorInfo* t = find(tensor_name);
+  if (!t) { err = "missing tensor: " + tensor_name; return false; }
+  out.resize(static_cast<size_t>(t->end - t->begin));
+  int fd = ::open(shard_path(t->shard).c_str(), O_RDONLY);
+  if (fd < 0) { err = "open failed: " + shard_path(t->shard); return false; }
+  size_t done = 0;
+  while (done < out.size()) {
+    ssize_t got = ::pread(fd, out.data() + done, out.size() - done,
+                          static_cast<off_t>(t->begin + done));
+    if (got <= 0) {
+      ::close(fd);
+      err = "pread failed: " + tensor_name;
+      return false;
+    }
+    done += static_cast<size_t>(got);
+  }
+  ::close(fd);
+  return true;
+}
+
+uint64_t ModelFiles::shard_size(const std::string& shard_name) const {
+  auto it = shard_sizes_.find(shard_name);
+  return it == shard_sizes_.end() ? 0 : it->second;
+}
+
+bool ModelFiles::config_int(const std::string& key, int64_t& out) const {
+  if (config_json_.empty()) return false;
+  try {
+    json c = json::parse(config_json_);
+    if (c.contains(key) && c[key].is_number_integer()) { out = c[key].get<int64_t>(); return true; }
+    if (c.contains("text_config")) {
+      const json& tc = c["text_config"];
+      if (tc.contains(key) && tc[key].is_number_integer()) { out = tc[key].get<int64_t>(); return true; }
+    }
+  } catch (...) {}
+  return false;
+}
+
+bool ModelFiles::config_bool(const std::string& key, bool& out) const {
+  if (config_json_.empty()) return false;
+  try {
+    json c = json::parse(config_json_);
+    if (c.contains(key) && c[key].is_boolean()) { out = c[key].get<bool>(); return true; }
+    if (c.contains("text_config")) {
+      const json& tc = c["text_config"];
+      if (tc.contains(key) && tc[key].is_boolean()) { out = tc[key].get<bool>(); return true; }
+    }
+  } catch (...) {}
+  return false;
+}
+
+bool ModelFiles::config_double(const std::string& key, double& out) const {
+  if (config_json_.empty()) return false;
+  try {
+    json c = json::parse(config_json_);
+    if (c.contains(key) && c[key].is_number()) { out = c[key].get<double>(); return true; }
+    if (c.contains("text_config")) {
+      const json& tc = c["text_config"];
+      if (tc.contains(key) && tc[key].is_number()) { out = tc[key].get<double>(); return true; }
+    }
+  } catch (...) {}
+  return false;
+}
+
+bool ModelFiles::config_quant(int64_t& bits, int64_t& group_size) const {
+  if (config_json_.empty()) return false;
+  try {
+    json c = json::parse(config_json_);
+    const json* q = nullptr;
+    if (c.contains("quantization_config")) q = &c["quantization_config"];
+    else if (c.contains("text_config") && c["text_config"].contains("quantization_config"))
+      q = &c["text_config"]["quantization_config"];
+    if (q && q->contains("bits") && q->contains("group_size")) {
+      bits = q->at("bits").get<int64_t>();
+      group_size = q->at("group_size").get<int64_t>();
+      return true;
+    }
+  } catch (...) {}
+  return false;
+}
+
+bool ModelFiles::config_quant_for(const std::string& tensor_name, int64_t& bits,
+                                  int64_t& group_size) const {
+  if (config_json_.empty()) return false;
+  try {
+    json c = json::parse(config_json_);
+    const json* q = nullptr;
+    if (c.contains("quantization_config")) q = &c["quantization_config"];
+    else if (c.contains("text_config") && c["text_config"].contains("quantization_config"))
+      q = &c["text_config"]["quantization_config"];
+    if (!q) return false;
+    // Override keys are module paths; strip the last segment (.weight etc.) then
+    // take the longest prefix match ("...mlp.gate" ⊂ "...mlp.gate.weight").
+    auto dot = tensor_name.rfind('.');
+    std::string mod = dot == std::string::npos ? tensor_name : tensor_name.substr(0, dot);
+    std::string best;
+    for (auto& [k, v] : q->items()) {
+      if (!v.is_object() || !v.contains("bits") || !v.contains("group_size")) continue;
+      if (mod.size() >= k.size() && mod.compare(0, k.size(), k) == 0 &&
+          (k.size() > best.size()))
+        best = k;
+    }
+    const json* use = q;
+    if (!best.empty()) use = &(*q)[best];
+    bits = use->at("bits").get<int64_t>();
+    group_size = use->at("group_size").get<int64_t>();
+    return true;
+  } catch (...) {}
+  return false;
+}
+
+bool ModelFiles::has_file(const std::string& file_name) const {
+  struct stat st;
+  return ::stat((dir_ + file_name).c_str(), &st) == 0;
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/safio.hpp b/macos/crates/engine-native/src/safio.hpp
new file mode 100644
index 0000000..cff9543
--- /dev/null
+++ b/macos/crates/engine-native/src/safio.hpp
@@ -0,0 +1,47 @@
+// Safetensors multi-shard header parse. A tensor lives in one shard; an expert's
+// 9 regions may span shards — weight_map in model.safetensors.index.json is authoritative.
+#pragma once
+#include <cstdint>
+#include <map>
+#include <string>
+#include <vector>
+
+namespace e0n {
+
+struct TensorInfo {
+  std::string dtype;
+  std::vector<int64_t> shape;
+  uint64_t begin = 0, end = 0;  // absolute offset in the shard file (includes 8+header_len)
+  std::string shard;
+};
+
+class ModelFiles {
+ public:
+  bool open(const std::string& dir, std::string& err);
+
+  const TensorInfo* find(const std::string& tensor_name) const;
+  std::string shard_path(const std::string& shard_name) const;
+  std::vector<std::string> shard_names() const;
+  uint64_t shard_size(const std::string& shard_name) const;  // 0 on failure
+
+  bool read_tensor(const std::string& tensor_name, std::vector<uint8_t>& out,
+                   std::string& err) const;
+
+  bool config_int(const std::string& key, int64_t& out) const;
+  bool config_bool(const std::string& key, bool& out) const;
+  bool config_double(const std::string& key, double& out) const;
+  bool config_quant(int64_t& bits, int64_t& group_size) const;
+  // Per-module override (e.g. 35b mlp.gate is int8/g64 while the root is int4).
+  // Longest-prefix match on the module path; else root quantization_config.
+  bool config_quant_for(const std::string& tensor_name, int64_t& bits,
+                        int64_t& group_size) const;
+  bool has_file(const std::string& file_name) const;
+
+ private:
+  std::string dir_;
+  std::map<std::string, TensorInfo> tensors_;
+  std::map<std::string, uint64_t> shard_sizes_;
+  std::string config_json_;
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/sampling.cpp b/macos/crates/engine-native/src/sampling.cpp
new file mode 100644
index 0000000..741a12f
--- /dev/null
+++ b/macos/crates/engine-native/src/sampling.cpp
@@ -0,0 +1,198 @@
+#include "sampling.hpp"
+
+#include <cmath>
+#include <cstring>
+#include <limits>
+#include <set>
+
+#include <nlohmann/json.hpp>
+
+#include <mlx/mlx.h>
+
+#include "fwd_common.hpp"
+
+namespace mx = mlx::core;
+
+namespace e0n {
+namespace sampling {
+
+namespace {
+constexpr float kNegInf = -std::numeric_limits<float>::infinity();
+
+// Single-row f32 mlx pipeline (upstream _mask_logits, penalty already applied).
+// Vendored mlx topk returns the k largest values; min(topk) is the ascending-topk first element.
+std::vector<float> mask_ops(const std::vector<float> &in, double temperature, int top_k,
+                            double top_p) {
+  if (in.empty()) return in;
+  const int64_t V = (int64_t)in.size();
+  mx::array l = wrap_f32(in, mkshape({1, V}));
+  if (temperature > 0.0) {
+    l = mx::divide(l, mx::array((float)temperature, mx::float32));
+  }
+  if (top_k > 0) {
+    int k = top_k < V ? top_k : (int)V;  // min(top_k, V)
+    auto top = mx::topk(l, k, -1);
+    auto threshold = mx::min(top, -1, /*keepdims=*/true);  // k-th largest = first of ascending top-k
+    l = mx::where(mx::less(l, threshold), mx::array(kNegInf, mx::float32), l);
+  }
+  if (top_p < 1.0) {
+    // Descending cumsum without flip/negative stride (those produce a 0-size slice here):
+    // descending_cumsum = total − ascending exclusive-cumsum. Threshold is the
+    // (k-1)-th largest, located as ascending index (V−k).
+    auto asc = mx::sort(l, -1);
+    auto p_asc = mx::softmax(asc, -1);
+    auto cum_exc = mx::subtract(mx::cumsum(p_asc, -1), p_asc);
+    auto total = mx::sum(p_asc, -1, /*keepdims=*/true);
+    auto cum_desc = mx::subtract(total, cum_exc);
+    auto cutmask = mx::less_equal(cum_desc, mx::array((float)top_p, mx::float32));
+    auto counts = mx::sum(mx::astype(cutmask, mx::int32), -1, /*keepdims=*/true);
+    auto k = mx::maximum(counts, mx::array(1, mx::int32));
+    auto pos = mx::subtract(mx::array((int32_t)V, mx::int32), k);
+    auto threshold = mx::take_along_axis(asc, pos, -1);
+    l = mx::where(mx::less(l, threshold), mx::array(kNegInf, mx::float32), l);
+  }
+  if (temperature <= 0.0) {
+    auto m = mx::max(l, -1, /*keepdims=*/true);
+    l = mx::where(mx::less(l, m), mx::array(kNegInf, mx::float32), l);
+  }
+  std::vector<float> out;
+  std::string err;
+  if (!to_host32(l, out, err)) return in;  // pipeline failure → leave unmasked
+  return out;
+}
+}  // namespace
+
+std::vector<float> mask_logits(const std::vector<float> &logits, double temperature,
+                               int top_k, double top_p) {
+  return mask_ops(logits, temperature, top_k, top_p);
+}
+
+std::vector<float> apply_repetition_penalty(const std::vector<float> &logits, double penalty,
+                                            const std::vector<int32_t> &history) {
+  std::vector<float> out = logits;
+  if (penalty == 1.0 || history.empty()) return out;
+  std::set<int32_t> uniq(history.begin(), history.end());
+  for (int32_t id : uniq) {
+    if (id < 0 || (size_t)id >= out.size()) continue;
+    float v = out[id];
+    out[id] = (v > 0.f) ? (float)((double)v / penalty) : v * (float)penalty;
+  }
+  return out;
+}
+
+void rng_seed(Rng &r, uint64_t seed) {
+  uint64_t z = seed + 0x9E3779B97F4A7C15ull;
+  auto sm = [](uint64_t &x) {
+    x += 0x9E3779B97F4A7C15ull;
+    uint64_t z2 = x;
+    z2 = (z2 ^ (z2 >> 30)) * 0xBF58476D1CE4E5B9ull;
+    z2 = (z2 ^ (z2 >> 27)) * 0x94D049BB133111EBull;
+    return z2 ^ (z2 >> 31);
+  };
+  r.s[0] = sm(z);
+  r.s[1] = sm(z) ^ 0xD1B54A32D192ED03ull;
+}
+
+uint64_t rng_next(Rng &r) {
+  uint64_t z = (r.s[0] += 0x9E3779B97F4A7C15ull);
+  z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ull;
+  z = (z ^ (z >> 27)) * 0x94D049BB133111EBull;
+  uint64_t o = z ^ (z >> 31);
+  r.s[1] += 0x9E3779B97F4A7C15ull;
+  return o ^ r.s[1];
+}
+
+int32_t gumbel_draw(const std::vector<float> &masked, uint64_t seed, Rng &rng) {
+  rng_seed(rng, seed);
+  size_t b = 0;
+  double best = -std::numeric_limits<double>::infinity();
+  for (size_t i = 0; i < masked.size(); i++) {
+    if (std::isinf(masked[i])) continue;
+    uint64_t u64 = rng_next(rng);
+    double u = ((double)(u64 >> 11) + 1.0) / 9007199254740993.0;
+    double g = masked[i] - std::log(-std::log(u));
+    if (g > best) { best = g; b = i; }
+  }
+  return (int32_t)b;
+}
+
+int32_t sample_step(const std::vector<float> &logits, const Params &p, bool first,
+                    const std::vector<int32_t> &history, Rng &rng) {
+  size_t V = logits.size();
+  if (V == 0) return -1;
+  const bool ftg = first && p.first_token_greedy && p.temperature > 0.0;
+  const bool greedy = p.temperature <= 0.0 || ftg;
+  if (greedy) {
+    std::vector<float> l = apply_repetition_penalty(logits, p.repetition_penalty, history);
+    size_t b = 0;
+    for (size_t i = 1; i < V; i++)
+      if (l[i] > l[b]) b = i;  // ties take the smaller id (strict >)
+    return (int32_t)b;
+  }
+  std::vector<float> l = apply_repetition_penalty(logits, p.repetition_penalty, history);
+  l = mask_ops(l, p.temperature, p.top_k, p.top_p);
+  if (p.has_seed) return gumbel_draw(l, p.seed, rng);
+  return gumbel_draw(l, rng_next(rng), rng);  // no seed = stream-local (not stable across runs)
+}
+
+TierDefaults tier_defaults(const std::string &tier) {
+  TierDefaults td;
+  td.repetition_penalty = (tier == "edge0-8b") ? 1.1 : 1.0;
+  return td;
+}
+
+Params parse_params(const std::string &params_json, const TierDefaults &td,
+                    const std::vector<int32_t> &eos_from_config) {
+  Params p;
+  p.temperature = td.temperature;
+  p.top_p = td.top_p;
+  p.top_k = td.top_k;
+  p.repetition_penalty = td.repetition_penalty;
+  p.max_new_tokens = td.max_new_tokens;
+  p.first_token_greedy = td.first_token_greedy;
+  p.eos_ids = eos_from_config;
+
+  nlohmann::json v = nlohmann::json::parse(params_json.empty() ? "{}" : params_json,
+                                           nullptr, /*allow_exceptions=*/false);
+  if (v.is_discarded() || !v.is_object()) return p;
+  auto get = [&v](const char *k) -> const nlohmann::json * {
+    auto it = v.find(k);
+    return it == v.end() ? nullptr : &*it;
+  };
+  if (auto x = get("max_new_tokens"); x && x->is_number()) {
+    int64_t n = x->get<int64_t>();
+    if (n < 0) n = 0;
+    if (n > 4096) n = 4096;
+    p.max_new_tokens = (size_t)n;
+  }
+  if (auto x = get("temperature"); x && x->is_number()) p.temperature = x->get<double>();
+  if (auto x = get("top_p"); x && x->is_number()) p.top_p = x->get<double>();
+  if (auto x = get("top_k"); x && x->is_number()) p.top_k = (int)x->get<int64_t>();
+  if (auto x = get("repetition_penalty"); x && x->is_number())
+    p.repetition_penalty = x->get<double>();
+  if (auto x = get("first_token_greedy"); x && x->is_boolean())
+    p.first_token_greedy = x->get<bool>();
+  if (auto x = get("seed"); x && (x->is_number_integer() || x->is_null())) {
+    if (x->is_number_integer()) {
+      p.has_seed = true;
+      int64_t s = x->get<int64_t>();
+      p.seed = s < 0 ? 0 : (uint64_t)s;
+    } else {
+      p.has_seed = false;
+    }
+  }
+  if (auto x = get("teacher_tokens"); x && x->is_array()) {
+    for (const auto &it : *x)
+      if (it.is_number_integer()) p.teacher_tokens.push_back(it.get<int32_t>());
+  }
+  if (auto x = get("stop"); x && x->is_array()) {
+    for (const auto &it : *x) {
+      if (it.is_number_integer())
+        p.stop_token_ids.push_back(it.get<int32_t>());
+    }
+  }
+  return p;
+}
+
+}  // namespace sampling
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/sampling.hpp b/macos/crates/engine-native/src/sampling.hpp
new file mode 100644
index 0000000..6149785
--- /dev/null
+++ b/macos/crates/engine-native/src/sampling.hpp
@@ -0,0 +1,61 @@
+// Sampling pipeline: mlx f32 mask ops plus host Gumbel-max categorical.
+#pragma once
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace e0n {
+namespace sampling {
+
+struct TierDefaults {
+  double temperature = 0.7;
+  double top_p = 0.95;
+  int top_k = 64;
+  double repetition_penalty = 1.0;  // 8b=1.1, 35b=1.0
+  size_t max_new_tokens = 2048;
+  bool first_token_greedy = true;
+};
+
+struct Params {
+  double temperature = 0.7;
+  double top_p = 0.95;
+  int top_k = 64;
+  double repetition_penalty = 1.0;
+  size_t max_new_tokens = 64;
+  std::vector<int32_t> eos_ids;           // from config.json at create (35b is an array)
+  std::vector<int32_t> stop_token_ids;    // params "stop" as ids
+  std::vector<std::vector<int32_t>> stop_seqs;  // params "stop" strings after tokenize
+  bool has_seed = false;
+  uint64_t seed = 0;
+  std::vector<int32_t> teacher_tokens;  // oracle replay; step in and out follow this
+  bool first_token_greedy = true;
+};
+
+TierDefaults tier_defaults(const std::string &tier);  // edge0-8b / edge0-35b; else generic
+
+// Missing keys fall back to tier defaults; unknown keys are ignored.
+Params parse_params(const std::string &params_json, const TierDefaults &td,
+                    const std::vector<int32_t> &eos_from_config);
+
+// logits = this step's host f32 [V]; history = prompt + emitted (same as upstream).
+// rng is owned by the stream (seed is replayed before each draw).
+struct Rng {
+  uint64_t s[2];  // splitmix64; seed-determined, advanced per step
+};
+void rng_seed(Rng &r, uint64_t seed);
+uint64_t rng_next(Rng &r);
+
+// One sample. Greedy shortcut: temp≤0 (except first-token-greedy off) skips the
+// mask chain and host-argnmaxes. first_token_greedy uses temperature=0 masks then argmax.
+int32_t sample_step(const std::vector<float> &logits, const Params &p, bool first,
+                    const std::vector<int32_t> &history, Rng &rng);
+
+std::vector<float> mask_logits(const std::vector<float> &logits, double temperature,
+                               int top_k, double top_p);
+int32_t gumbel_draw(const std::vector<float> &masked, uint64_t seed, Rng &rng);
+std::vector<float> apply_repetition_penalty(const std::vector<float> &logits,
+                                            double penalty,
+                                            const std::vector<int32_t> &history);
+
+}  // namespace sampling
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/stepprefetch.cpp b/macos/crates/engine-native/src/stepprefetch.cpp
new file mode 100644
index 0000000..cbf2d89
--- /dev/null
+++ b/macos/crates/engine-native/src/stepprefetch.cpp
@@ -0,0 +1,87 @@
+#include "stepprefetch.hpp"
+
+#include <algorithm>
+#include <unordered_set>
+
+namespace e0n {
+
+StepPrefetch::StepPrefetch(HotStack &stack, int n_threads) : stack_(stack) {
+  n_threads = std::max(1, n_threads);
+  for (int i = 0; i < n_threads; i++) workers_.emplace_back([this] { worker_loop(); });
+}
+
+StepPrefetch::~StepPrefetch() {
+  {
+    std::lock_guard<std::mutex> g(qmu_);
+    stop_ = true;
+  }
+  qcv_.notify_all();
+  for (auto &t : workers_) t.join();
+}
+
+std::shared_ptr<StepPrefetch::Ticket> StepPrefetch::submit(const Step &plan) {
+  auto t = std::make_shared<Ticket>();
+  std::lock_guard<std::mutex> g(qmu_);
+  t->pending = static_cast<int>(plan.experts.size());
+  if (plan.experts.empty()) return t;
+  for (auto [layer, expert] : plan.experts) queue_.push_back(Job{t, layer, expert});
+  qcv_.notify_all();
+  return t;
+}
+
+void StepPrefetch::cancel_pending() {
+  // Do not decrement ticket counts here (in-flight finish would double-decrement).
+  // Mark only: a dequeued worker skips the load and still calls finish so counts stay balanced.
+  std::lock_guard<std::mutex> g(qmu_);
+  std::unordered_set<Ticket *> seen;
+  for (auto &j : queue_) {
+    if (seen.insert(j.t.get()).second) j.t->cancelled.store(true);
+  }
+}
+
+void StepPrefetch::drain() {
+  std::unique_lock<std::mutex> g(qmu_);
+  qcv_.wait(g, [&] { return queue_.empty() && active_ == 0; });
+}
+
+size_t StepPrefetch::queue_depth() const {
+  std::lock_guard<std::mutex> g(qmu_);
+  return queue_.size();
+}
+
+void StepPrefetch::finish(std::shared_ptr<Ticket> t, std::string err) {
+  std::lock_guard<std::mutex> g(t->mu);
+  if (!err.empty()) { t->bad = true; if (t->err.empty()) t->err = std::move(err); }
+  if (--t->pending == 0) t->cv.notify_all();
+}
+
+void StepPrefetch::worker_loop() {
+  for (;;) {
+    Job j;
+    {
+      std::unique_lock<std::mutex> g(qmu_);
+      qcv_.wait(g, [&] { return stop_ || !queue_.empty(); });
+      if (stop_ && queue_.empty()) return;
+      j = queue_.front();
+      queue_.pop_front();
+      active_++;
+    }
+    std::string err;
+    if (j.t->cancelled.load()) {
+      // Cancelled: skip the load but still finish (count balance; see cancel_pending).
+      err.clear();
+    } else {
+      std::string load_err;
+      auto lease = stack_.acquire(j.layer, j.expert, load_err);
+      err = load_err;
+    }
+    finish(j.t, std::move(err));
+    {
+      std::lock_guard<std::mutex> g(qmu_);
+      active_--;
+      if (queue_.empty() && active_ == 0) qcv_.notify_all();
+    }
+  }
+}
+
+}  // namespace e0n
diff --git a/macos/crates/engine-native/src/stepprefetch.hpp b/macos/crates/engine-native/src/stepprefetch.hpp
new file mode 100644
index 0000000..baef838
--- /dev/null
+++ b/macos/crates/engine-native/src/stepprefetch.hpp
@@ -0,0 +1,63 @@
+// Async per-step prefetch: submit(step i+1) returns a ticket; IO threads load into
+// the hot stack while compute GEMMs step i. turn_reset → cancel_pending + drain.
+#pragma once
+#include <atomic>
+#include <condition_variable>
+#include <cstdint>
+#include <deque>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "hotstack.hpp"
+#include "provider.hpp"
+
+namespace e0n {
+
+class StepPrefetch {
+ public:
+  struct Ticket {
+    std::mutex mu;
+    std::condition_variable cv;
+    int pending = 0;
+    std::atomic<bool> cancelled{false};  // dequeued worker skips load but still finishes the count
+    bool bad = false;
+    std::string err;
+    bool done() const { return pending == 0; }
+    bool wait() {
+      std::unique_lock<std::mutex> g(mu);
+      cv.wait(g, [&] { return pending == 0; });
+      return !bad;
+    }
+  };
+
+  StepPrefetch(HotStack &stack, int n_threads);
+  ~StepPrefetch();  // cancel then join; no in-flight load at destructor is the caller's duty
+
+  std::shared_ptr<Ticket> submit(const Step &plan);
+
+  void cancel_pending();
+  void drain();
+  size_t queue_depth() const;
+
+ private:
+  struct Job {
+    std::shared_ptr<Ticket> t;
+    int layer;
+    int64_t expert;
+  };
+  void finish(std::shared_ptr<Ticket> t, std::string err);
+  void worker_loop();
+
+  HotStack &stack_;
+  std::vector<std::thread> workers_;
+  mutable std::mutex qmu_;
+  std::condition_variable qcv_;
+  std::deque<Job> queue_;
+  bool stop_ = false;
+  int active_ = 0;
+};
+
+}  // namespace e0n
diff --git a/macos/crates/engine-tok/Cargo.toml b/macos/crates/engine-tok/Cargo.toml
new file mode 100644
index 0000000..8180914
--- /dev/null
+++ b/macos/crates/engine-tok/Cargo.toml
@@ -0,0 +1,17 @@
+[package]
+name = "engine-tok"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+# Tokenizer staticlib linked into the native engine. Callers own the JSON envelope.
+[lib]
+name = "edge0_engine_tok"
+crate-type = ["staticlib", "rlib"]
+
+[dependencies]
+serde_json.workspace = true
+tokenizers.workspace = true
+minijinja.workspace = true
+minijinja-contrib.workspace = true
diff --git a/macos/crates/engine-tok/src/lib.rs b/macos/crates/engine-tok/src/lib.rs
new file mode 100644
index 0000000..155bd1f
--- /dev/null
+++ b/macos/crates/engine-tok/src/lib.rs
@@ -0,0 +1,261 @@
+//! Tokenizer C ABI: string ↔ token ids and chat-template render. Callers own the JSON envelope.
+use std::cell::RefCell;
+use std::os::raw::{c_char, c_int, c_void};
+use std::path::Path;
+
+use serde_json::{Map, Value};
+use tokenizers::Tokenizer;
+
+pub const ETOK_OK: c_int = 0;
+pub const ETOK_ERR_INVALID: c_int = -1;
+pub const ETOK_ERR_BUFFER: c_int = -2;
+
+pub struct TokFace {
+    tok: Tokenizer,
+    template: String,
+}
+
+impl TokFace {
+    pub fn open(model_dir: &Path) -> Result<Self, String> {
+        let tpath = model_dir.join("tokenizer.json");
+        let cpath = model_dir.join("chat_template.jinja");
+        let gpath = model_dir.join("config.json");
+        for p in [&tpath, &cpath, &gpath] {
+            if !p.is_file() {
+                return Err(format!(
+                    "tokenizer requires {p:?} (read-only tokenizer files; weights are not loaded here)"
+                ));
+            }
+        }
+        let tok =
+            Tokenizer::from_file(&tpath).map_err(|e| format!("failed to parse tokenizer.json: {e}"))?;
+        let template = std::fs::read_to_string(&cpath).map_err(|e| e.to_string())?;
+        let _ = gpath;
+        Ok(Self { tok, template })
+    }
+
+    pub fn encode(&self, text: &str) -> Vec<i32> {
+        self.tok
+            .encode(text, false)
+            .map(|e| e.get_ids().iter().map(|i| *i as i32).collect())
+            .unwrap_or_default()
+    }
+
+    pub fn decode(&self, tokens: &[i32]) -> Result<String, String> {
+        let ids: Vec<u32> = tokens.iter().map(|t| *t as u32).collect();
+        self.tok
+            .decode(&ids, false)
+            .map_err(|e| format!("decode failed: {e}"))
+    }
+
+    pub fn render(&self, ctx: &Map<String, Value>) -> Result<String, String> {
+        let mut env = minijinja::Environment::new();
+        env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
+        env.add_function(
+            "raise_exception",
+            |msg: minijinja::Value| -> Result<(), minijinja::Error> {
+                Err(minijinja::Error::new(
+                    minijinja::ErrorKind::InvalidOperation,
+                    msg.to_string(),
+                ))
+            },
+        );
+        let tmpl = env
+            .template_from_str(&self.template)
+            .map_err(|e| format!("template compile failed: {e:#}"))?;
+        let value = minijinja::Value::from_serialize(ctx);
+        tmpl.render(value)
+            .map_err(|e| format!("template render failed: {e:#}"))
+    }
+}
+
+thread_local! {
+    static LAST_ERR: RefCell<String> = const { RefCell::new(String::new()) };
+}
+
+fn set_err(msg: impl ToString) {
+    LAST_ERR.with(|e| e.replace(msg.to_string()));
+}
+
+fn clear_err() {
+    LAST_ERR.with(|e| e.borrow_mut().clear());
+}
+
+fn as_face<'a>(t: *mut c_void) -> Option<&'a TokFace> {
+    if t.is_null() {
+        set_err("etok: null handle");
+        return None;
+    }
+    Some(unsafe { &*(t as *const TokFace) })
+}
+
+fn write_str(s: &str, buf: *mut c_char, buf_len: usize, written: *mut usize) -> c_int {
+    let need = s.len();
+    if !written.is_null() {
+        unsafe { *written = need };
+    }
+    if buf.is_null() {
+        return ETOK_OK;
+    }
+    if buf_len < need {
+        return ETOK_ERR_BUFFER;
+    }
+    unsafe {
+        std::ptr::copy_nonoverlapping(s.as_ptr() as *const c_char, buf, need);
+    }
+    ETOK_OK
+}
+
+/// # Safety
+#[no_mangle]
+pub unsafe extern "C" fn etok_create(model_dir: *const c_char) -> *mut c_void {
+    if model_dir.is_null() {
+        set_err("etok_create: model_dir is empty");
+        return std::ptr::null_mut();
+    }
+    let dir = std::ffi::CStr::from_ptr(model_dir);
+    let dir = match dir.to_str() {
+        Ok(s) => s,
+        Err(_) => {
+            set_err("etok_create: model_dir is not UTF-8");
+            return std::ptr::null_mut();
+        }
+    };
+    match TokFace::open(Path::new(dir)) {
+        Ok(face) => {
+            clear_err();
+            Box::into_raw(Box::new(face)) as *mut c_void
+        }
+        Err(e) => {
+            set_err(e);
+            std::ptr::null_mut()
+        }
+    }
+}
+
+/// # Safety
+#[no_mangle]
+pub extern "C" fn etok_destroy(t: *mut c_void) {
+    if !t.is_null() {
+        drop(unsafe { Box::from_raw(t as *mut TokFace) });
+    }
+}
+
+/// # Safety
+#[no_mangle]
+pub unsafe extern "C" fn etok_last_error(
+    buf: *mut c_char,
+    buf_len: usize,
+    written: *mut usize,
+) -> c_int {
+    let msg = LAST_ERR.with(|e| e.borrow().clone());
+    write_str(&msg, buf, buf_len, written)
+}
+
+/// # Safety
+#[no_mangle]
+pub unsafe extern "C" fn etok_tokenize(
+    t: *mut c_void,
+    text: *const c_char,
+    ids_out: *mut i32,
+    ids_cap: usize,
+    written: *mut usize,
+) -> c_int {
+    let Some(face) = as_face(t) else {
+        return ETOK_ERR_INVALID;
+    };
+    if text.is_null() {
+        set_err("etok_tokenize: text is empty");
+        return ETOK_ERR_INVALID;
+    }
+    let Ok(s) = std::ffi::CStr::from_ptr(text).to_str() else {
+        set_err("etok_tokenize: text is not UTF-8");
+        return ETOK_ERR_INVALID;
+    };
+    let ids = face.encode(s);
+    let n = ids.len();
+    if !written.is_null() {
+        *written = n;
+    }
+    if ids_out.is_null() {
+        return ETOK_OK;
+    }
+    if ids_cap < n {
+        return ETOK_ERR_BUFFER;
+    }
+    std::ptr::copy_nonoverlapping(ids.as_ptr(), ids_out, n);
+    ETOK_OK
+}
+
+/// # Safety
+#[no_mangle]
+pub unsafe extern "C" fn etok_decode(
+    t: *mut c_void,
+    ids: *const i32,
+    n_ids: usize,
+    buf: *mut c_char,
+    buf_len: usize,
+    written: *mut usize,
+) -> c_int {
+    let Some(face) = as_face(t) else {
+        return ETOK_ERR_INVALID;
+    };
+    if ids.is_null() && n_ids > 0 {
+        set_err("etok_decode: ids is null but length is non-zero");
+        return ETOK_ERR_INVALID;
+    }
+    let tokens = if n_ids == 0 {
+        Vec::new()
+    } else {
+        std::slice::from_raw_parts(ids, n_ids).to_vec()
+    };
+    match face.decode(&tokens) {
+        Ok(s) => {
+            clear_err();
+            write_str(&s, buf, buf_len, written)
+        }
+        Err(e) => {
+            set_err(e);
+            ETOK_ERR_INVALID
+        }
+    }
+}
+
+/// # Safety
+#[no_mangle]
+pub unsafe extern "C" fn etok_render(
+    t: *mut c_void,
+    ctx_json: *const c_char,
+    buf: *mut c_char,
+    buf_len: usize,
+    written: *mut usize,
+) -> c_int {
+    let Some(face) = as_face(t) else {
+        return ETOK_ERR_INVALID;
+    };
+    if ctx_json.is_null() {
+        set_err("etok_render: ctx_json is empty");
+        return ETOK_ERR_INVALID;
+    }
+    let Ok(s) = std::ffi::CStr::from_ptr(ctx_json).to_str() else {
+        set_err("etok_render: ctx_json is not UTF-8");
+        return ETOK_ERR_INVALID;
+    };
+    let ctx: Map<String, Value> = match serde_json::from_str(s) {
+        Ok(m) => m,
+        Err(e) => {
+            set_err(format!("invalid ctx_json: {e}"));
+            return ETOK_ERR_INVALID;
+        }
+    };
+    match face.render(&ctx) {
+        Ok(out) => {
+            clear_err();
+            write_str(&out, buf, buf_len, written)
+        }
+        Err(e) => {
+            set_err(e);
+            ETOK_ERR_INVALID
+        }
+    }
+}
diff --git a/macos/rust-toolchain.toml b/macos/rust-toolchain.toml
new file mode 100644
index 0000000..e88baf1
--- /dev/null
+++ b/macos/rust-toolchain.toml
@@ -0,0 +1,2 @@
+[toolchain]
+channel = "1.88.0"
diff --git a/macos/scripts/bundle-check.sh b/macos/scripts/bundle-check.sh
new file mode 100755
index 0000000..13beecf
--- /dev/null
+++ b/macos/scripts/bundle-check.sh
@@ -0,0 +1,90 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+BUNDLE="${1:?usage: bundle-check.sh <edge0.app>}"
+RES="$BUNDLE/Contents/Resources"
+BIN="$RES/bin"
+LIB="$RES/lib"
+
+[[ -d "$BUNDLE/Contents/MacOS" ]] || { echo "bundle-check: missing Contents/MacOS" >&2; exit 1; }
+[[ -d "$BIN" && -d "$LIB" ]] || { echo "bundle-check: missing Resources/bin or Resources/lib" >&2; exit 1; }
+for path in "$BIN/edge0" "$BIN/edge0d" "$BIN/edge0-engine" "$LIB/libedge0_engine_native.dylib" "$LIB/libmlx.dylib" "$LIB/mlx.metallib"; do
+  [[ -f "$path" ]] || { echo "bundle-check: missing $path" >&2; exit 1; }
+done
+
+PLIST="$BUNDLE/Contents/Info.plist"
+[[ "$(/usr/libexec/PlistBuddy -c 'Print :CFBundleIdentifier' "$PLIST")" == "app.edge0.client" ]] || {
+  echo "bundle-check: CFBundleIdentifier must be app.edge0.client" >&2
+  exit 1
+}
+if /usr/libexec/PlistBuddy -c 'Print :LSUIElement' "$PLIST" >/dev/null 2>&1; then
+  echo "bundle-check: LSUIElement must be absent (Dock-resident app)" >&2
+  exit 1
+fi
+VERSION="$(/usr/libexec/PlistBuddy -c 'Print :CFBundleShortVersionString' "$PLIST")"
+
+for bin in edge0 edge0d edge0-engine; do
+  output="$("$BIN/$bin" --version 2>/dev/null || true)"
+  echo "$output" | grep -Eq "(^| )$VERSION($|[[:space:]])" || {
+    echo "bundle-check: $bin version mismatch: $output vs $VERSION" >&2
+    exit 1
+  }
+done
+
+ENGINE_VERSION_TOOL="$(mktemp "/tmp/edge0-engine-version.XXXXXX")"
+ENGINE_VERSION_SRC="$ENGINE_VERSION_TOOL.c"
+trap 'rm -f "$ENGINE_VERSION_TOOL" "$ENGINE_VERSION_SRC"' EXIT
+cat > "$ENGINE_VERSION_SRC" <<'EOF'
+#include <dlfcn.h>
+#include <stdio.h>
+typedef const char *(*version_fn)(void);
+int main(int argc, char **argv) {
+  void *h = dlopen(argv[1], RTLD_NOW);
+  if (!h) return 2;
+  version_fn f = (version_fn)dlsym(h, "e0_engine_version");
+  if (!f) return 3;
+  puts(f());
+  return 0;
+}
+EOF
+cc "$ENGINE_VERSION_SRC" -o "$ENGINE_VERSION_TOOL"
+ENGINE_VERSION_VALUE="$("$ENGINE_VERSION_TOOL" "$LIB/libedge0_engine_native.dylib")"
+[[ "$ENGINE_VERSION_VALUE" == "$VERSION" ]] || {
+  echo "bundle-check: native engine version mismatch: $ENGINE_VERSION_VALUE vs $VERSION" >&2
+  exit 1
+}
+
+MACHOS=()
+while IFS= read -r macho; do
+  MACHOS+=("$macho")
+done < <(
+  find "$BUNDLE/Contents/MacOS" "$BIN" "$LIB" -type f -print0 |
+    xargs -0 file 2>/dev/null |
+    awk -F: '$2 ~ /Mach-O/ {print $1}' |
+    sort
+)
+for macho in "${MACHOS[@]}"; do
+  otool -l "$macho" | awk '{
+    line=$0
+    sub(/^[[:space:]]+/, "", line)
+    if (index(line, "path ") == 1 && $2 !~ /^@/ && $2 !~ /^\/System\/Library/ && $2 !~ /^\/usr\/lib/) {
+      print
+      bad=1
+    }
+  } END {exit bad}' || {
+    echo "bundle-check: absolute non-system rpath in $macho" >&2
+    exit 1
+  }
+  while IFS= read -r dep; do
+    [[ "$dep" == @* || "$dep" == /System/Library/* || "$dep" == /usr/lib/* ]] || {
+      echo "bundle-check: unsupported dependency $dep in $macho" >&2
+      exit 1
+    }
+  done < <(otool -L "$macho" | awk 'NR>1 {print $1}')
+done
+
+codesign --verify --deep --strict "$BUNDLE" || {
+  echo "bundle-check: codesign verify failed" >&2
+  exit 1
+}
+echo "bundle-check: PASS $BUNDLE"
diff --git a/macos/scripts/check-prereqs.sh b/macos/scripts/check-prereqs.sh
new file mode 100755
index 0000000..6dc8cf4
--- /dev/null
+++ b/macos/scripts/check-prereqs.sh
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+fail() { echo "check-prereqs: $*" >&2; exit 1; }
+
+[[ "$(uname -s)" == "Darwin" ]] || fail "macOS is required"
+[[ "$(uname -m)" == "arm64" ]] || fail "Apple Silicon (arm64) is required"
+
+macos_major="$(sw_vers -productVersion | cut -d. -f1)"
+[[ "$macos_major" -ge 14 ]] || fail "macOS 14 or later is required (found $(sw_vers -productVersion))"
+
+command -v rustc >/dev/null || fail "Rust 1.88+ is required (https://rustup.rs)"
+command -v cargo >/dev/null || fail "cargo is required"
+command -v cmake >/dev/null || fail "CMake 3.24+ is required"
+command -v node >/dev/null || fail "Node 22 or later is required"
+node_major="$(node -p 'process.versions.node.split(".")[0]')"
+[[ "$node_major" -ge 22 ]] || fail "Node 22 or later is required (found $(node --version))"
+# yarn 1.22: use PATH if present, otherwise Corepack (ships with Node).
+export PATH="${HOME}/.local/bin:${PATH}"
+if ! command -v yarn >/dev/null 2>&1; then
+  command -v corepack >/dev/null || fail "yarn 1.22 is required (install Corepack or: npm i -g yarn@1.22.22)"
+  mkdir -p "${HOME}/.local/bin"
+  corepack enable --install-directory "${HOME}/.local/bin"
+  corepack prepare yarn@1.22.22 --activate
+fi
+command -v yarn >/dev/null || fail "yarn 1.22 is required"
+command -v xcrun >/dev/null || fail "Xcode Command Line Tools are required (Metal compiler)"
+xcrun --find metal >/dev/null 2>&1 || fail "the Metal compiler is missing (install Xcode or CLT)"
+
+echo "check-prereqs: ok"
+echo "  macOS $(sw_vers -productVersion) $(uname -m)"
+echo "  rustc $(rustc --version)"
+echo "  cmake $(cmake --version | head -1)"
+echo "  node  $(node --version)"
+echo "  yarn  $(yarn --version)"
diff --git a/macos/scripts/package-macos.sh b/macos/scripts/package-macos.sh
new file mode 100755
index 0000000..a1cda59
--- /dev/null
+++ b/macos/scripts/package-macos.sh
@@ -0,0 +1,105 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export PATH="${HOME}/.local/bin:${PATH}"
+
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+export CARGO_TARGET_DIR="$ROOT/target"
+APP_DIR="$ROOT/app"
+BUILD_DIR="${EDGE0_BUILD_DIR:-$CARGO_TARGET_DIR/release}"
+ENGINE_DIR="${EDGE0_ENGINE_DIR:-$ROOT/target/engine}"
+OUT_DIR="${EDGE0_DIST_DIR:-$ROOT/dist}"
+WORK_DIR="$(mktemp -d "${TMPDIR:-/tmp}/edge0-package.XXXXXX")"
+trap 'rm -rf "$WORK_DIR"' EXIT
+
+VERSION="${EDGE0_VERSION:-$(awk '
+  $0 == "[workspace.package]" { in_pkg=1; next }
+  /^\[/ { in_pkg=0 }
+  in_pkg && $1 == "version" { gsub(/"/, "", $3); print $3; exit }
+' "$ROOT/Cargo.toml")}"
+ARCH="$(uname -m)"
+if [[ "$ARCH" != "arm64" && "${EDGE0_ALLOW_NON_ARM64:-0}" != "1" ]]; then
+  echo "package-macos: expected arm64 host, got $ARCH" >&2
+  exit 2
+fi
+
+if [[ "${EDGE0_SKIP_TAURI_BUILD:-0}" != "1" ]]; then
+  (cd "$APP_DIR" && yarn tauri build --bundles app --config "{\"version\":\"$VERSION\"}")
+fi
+TAURI_APP="${EDGE0_TAURI_APP:-$CARGO_TARGET_DIR/release/bundle/macos/edge0.app}"
+[[ -d "$TAURI_APP" ]] || { echo "package-macos: Tauri app not found: $TAURI_APP" >&2; exit 2; }
+
+for bin in edge0 edge0d edge0-engine; do
+  [[ -x "$BUILD_DIR/$bin" ]] || { echo "package-macos: missing $BUILD_DIR/$bin" >&2; exit 2; }
+done
+for lib in libedge0_engine_native.dylib libmlx.dylib mlx.metallib; do
+  [[ -f "$ENGINE_DIR/$lib" ]] || { echo "package-macos: missing $ENGINE_DIR/$lib" >&2; exit 2; }
+done
+
+STAGE="$WORK_DIR/edge0.app"
+cp -R "$TAURI_APP" "$STAGE"
+mkdir -p "$STAGE/Contents/Resources/bin" "$STAGE/Contents/Resources/lib"
+rm -f "$STAGE/Contents/Resources/bin/"* "$STAGE/Contents/Resources/lib/"*
+install -m 755 "$BUILD_DIR/edge0" "$STAGE/Contents/Resources/bin/edge0"
+install -m 755 "$BUILD_DIR/edge0d" "$STAGE/Contents/Resources/bin/edge0d"
+install -m 755 "$BUILD_DIR/edge0-engine" "$STAGE/Contents/Resources/bin/edge0-engine"
+install -m 644 "$ENGINE_DIR/libedge0_engine_native.dylib" "$STAGE/Contents/Resources/lib/libedge0_engine_native.dylib"
+install -m 644 "$ENGINE_DIR/libmlx.dylib" "$STAGE/Contents/Resources/lib/libmlx.dylib"
+install -m 644 "$ENGINE_DIR/mlx.metallib" "$STAGE/Contents/Resources/lib/mlx.metallib"
+
+# Build an icns from the checked-in 512px source when the macOS toolchain is
+# available. Tauri's existing icon remains the fallback for inspection hosts.
+if command -v sips >/dev/null 2>&1 && command -v iconutil >/dev/null 2>&1; then
+  ICONSET="$WORK_DIR/AppIcon.iconset"
+  mkdir -p "$ICONSET"
+  for spec in "16:icon_16x16" "32:icon_16x16@2x" "32:icon_32x32" "64:icon_32x32@2x" "128:icon_128x128" "256:icon_128x128@2x" "256:icon_256x256" "512:icon_256x256@2x" "512:icon_512x512" "1024:icon_512x512@2x"; do
+    size="$(printf '%s' "$spec" | cut -d: -f1)"
+    name="$(printf '%s' "$spec" | cut -d: -f2-)"
+    sips -z "$size" "$size" "$APP_DIR/src-tauri/icons/edge0_logo.png" --out "$ICONSET/$name.png" >/dev/null
+  done
+  if ! iconutil -c icns "$ICONSET" -o "$STAGE/Contents/Resources/AppIcon.icns"; then
+    echo "package-macos: warning: iconutil rejected iconset; keeping Tauri icon fallback" >&2
+    rm -f "$STAGE/Contents/Resources/AppIcon.icns"
+  fi
+fi
+
+# Repair legacy build-tree rpaths in copied dylibs. New CMake builds emit
+# @loader_path directly; this keeps local packaging deterministic for older
+# artifacts while bundle-check remains the final enforcement point.
+if command -v install_name_tool >/dev/null 2>&1; then
+  for dylib in "$STAGE/Contents/Resources/lib/"*.dylib; do
+    old_rpaths="$(otool -l "$dylib" | awk '/LC_RPATH/{getline; getline; if ($1 == "path") print $2}')"
+    if [[ -n "$old_rpaths" ]]; then
+      while IFS= read -r old; do
+        [[ "$old" == @* ]] && continue
+        install_name_tool -delete_rpath "$old" "$dylib" || true
+      done <<< "$old_rpaths"
+      if ! printf '%s\n' "$old_rpaths" | grep -qx '@loader_path'; then
+        install_name_tool -add_rpath "@loader_path" "$dylib" || true
+      fi
+    fi
+  done
+fi
+
+PLIST="$STAGE/Contents/Info.plist"
+/usr/libexec/PlistBuddy -c "Set :CFBundleIdentifier app.edge0.client" "$PLIST"
+/usr/libexec/PlistBuddy -c "Set :CFBundleShortVersionString $VERSION" "$PLIST"
+/usr/libexec/PlistBuddy -c "Set :CFBundleVersion $VERSION" "$PLIST"
+/usr/libexec/PlistBuddy -c "Set :CFBundleName edge0" "$PLIST"
+if [[ -f "$STAGE/Contents/Resources/AppIcon.icns" ]]; then
+  /usr/libexec/PlistBuddy -c "Set :CFBundleIconFile AppIcon.icns" "$PLIST"
+else
+  /usr/libexec/PlistBuddy -c "Set :CFBundleIconFile edge0.icns" "$PLIST"
+fi
+# Keep LSUIElement unset so the app is a normal Dock-resident GUI.
+
+command -v xattr >/dev/null 2>&1 && xattr -cr "$STAGE" || true
+
+"$ROOT/scripts/sign-macos.sh" "$STAGE"
+"$ROOT/scripts/bundle-check.sh" "$STAGE"
+
+NAME="edge0-$VERSION-$ARCH.app"
+mkdir -p "$OUT_DIR"
+rm -rf "$OUT_DIR/$NAME"
+mv "$STAGE" "$OUT_DIR/$NAME"
+echo "package-macos: $OUT_DIR/$NAME"
diff --git a/macos/scripts/sign-macos.sh b/macos/scripts/sign-macos.sh
new file mode 100755
index 0000000..0dcb88b
--- /dev/null
+++ b/macos/scripts/sign-macos.sh
@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+BUNDLE="${1:?usage: sign-macos.sh <edge0.app>}"
+IDENTITY="${EDGE0_SIGN_IDENTITY:-}"
+ENTITLEMENTS="${EDGE0_ENTITLEMENTS:-$ROOT/app/src-tauri/Entitlements.plist}"
+
+if [[ ! -d "$BUNDLE" ]]; then
+  echo "sign-macos: bundle not found: $BUNDLE" >&2
+  exit 2
+fi
+command -v codesign >/dev/null 2>&1 || { echo "sign-macos: codesign is required on macOS" >&2; exit 2; }
+
+command -v xattr >/dev/null 2>&1 && xattr -cr "$BUNDLE" || true
+MACHOS=()
+while IFS= read -r macho; do
+  MACHOS+=("$macho")
+done < <(
+  find "$BUNDLE/Contents/MacOS" "$BUNDLE/Contents/Resources/bin" "$BUNDLE/Contents/Resources/lib" -type f -print0 2>/dev/null |
+    xargs -0 file 2>/dev/null |
+    awk -F: '$2 ~ /Mach-O/ {print $1}' |
+    sort -r
+)
+
+SIGN_ARGS=()
+if [[ -n "$IDENTITY" ]]; then
+  SIGN_ARGS+=(--options runtime --timestamp)
+  [[ -f "$ENTITLEMENTS" ]] && SIGN_ARGS+=(--entitlements "$ENTITLEMENTS")
+  echo "sign-macos: signing with Developer ID identity $IDENTITY"
+else
+  echo "sign-macos: signing ad-hoc"
+  IDENTITY="-"
+fi
+
+sign_one() {
+  local path="$1"
+  if [[ ${#SIGN_ARGS[@]} -gt 0 ]]; then
+    codesign --force --sign "$IDENTITY" "${SIGN_ARGS[@]}" "$path"
+  else
+    codesign --force --sign "$IDENTITY" "$path"
+  fi
+}
+
+for macho in "${MACHOS[@]}"; do
+  sign_one "$macho"
+done
+sign_one "$BUNDLE"
+codesign --verify --deep --strict "$BUNDLE"
+
+if [[ -n "${EDGE0_NOTARY_PROFILE:-}" && -n "${EDGE0_SIGN_IDENTITY:-}" ]]; then
+  echo "sign-macos: submitting for notarization"
+  xcrun notarytool submit "$BUNDLE" --keychain-profile "$EDGE0_NOTARY_PROFILE" --wait
+  xcrun stapler staple "$BUNDLE"
+else
+  echo "sign-macos: notarization skipped (identity/profile not set)"
+fi
diff --git a/macos/third_party/mlx/.clang-format b/macos/third_party/mlx/.clang-format
new file mode 100644
index 0000000..eab4576
--- /dev/null
+++ b/macos/third_party/mlx/.clang-format
@@ -0,0 +1,87 @@
+---
+AccessModifierOffset: -1
+AlignAfterOpenBracket: AlwaysBreak
+AlignConsecutiveAssignments: false
+AlignConsecutiveDeclarations: false
+AlignEscapedNewlinesLeft: true
+AlignOperands:   false
+AlignTrailingComments: false
+AllowAllParametersOfDeclarationOnNextLine: false
+AllowShortBlocksOnASingleLine: false
+AllowShortCaseLabelsOnASingleLine: false
+AllowShortFunctionsOnASingleLine: Empty
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AlwaysBreakAfterReturnType: None
+AlwaysBreakBeforeMultilineStrings: true
+AlwaysBreakTemplateDeclarations: true
+BinPackArguments: false
+BinPackParameters: false
+BraceWrapping:
+  AfterClass:      false
+  AfterControlStatement: false
+  AfterEnum:       false
+  AfterFunction:   false
+  AfterNamespace:  false
+  AfterObjCDeclaration: false
+  AfterStruct:     false
+  AfterUnion:      false
+  BeforeCatch:     false
+  BeforeElse:      false
+  IndentBraces:    false
+BreakBeforeBinaryOperators: None
+BreakBeforeBraces: Attach
+BreakBeforeTernaryOperators: true
+BreakConstructorInitializersBeforeComma: false
+BreakAfterJavaFieldAnnotations: false
+BreakStringLiterals: false
+ColumnLimit:     80
+CommentPragmas:  '^ IWYU pragma:'
+ConstructorInitializerAllOnOneLineOrOnePerLine: true
+ConstructorInitializerIndentWidth: 4
+ContinuationIndentWidth: 4
+Cpp11BracedListStyle: true
+DerivePointerAlignment: false
+DisableFormat:   false
+ForEachMacros:   [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ]
+IncludeCategories:
+  - Regex:           '^<.*\.h(pp)?>'
+    Priority:        1
+  - Regex:           '^<.*'
+    Priority:        2
+  - Regex:           '.*'
+    Priority:        3
+IndentCaseLabels: true
+IndentWidth:     2
+IndentWrappedFunctionNames: false
+KeepEmptyLinesAtTheStartOfBlocks: false
+MacroBlockBegin: ''
+MacroBlockEnd:   ''
+MaxEmptyLinesToKeep: 1
+NamespaceIndentation: None
+ObjCBlockIndentWidth: 2
+ObjCSpaceAfterProperty: false
+ObjCSpaceBeforeProtocolList: false
+PenaltyBreakBeforeFirstCallParameter: 1
+PenaltyBreakComment: 300
+PenaltyBreakFirstLessLess: 120
+PenaltyBreakString: 1000
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 200
+PointerAlignment: Left
+ReflowComments:  true
+SortIncludes:    true
+SpaceAfterCStyleCast: false
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeParens: ControlStatements
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles:  false
+SpacesInContainerLiterals: true
+SpacesInCStyleCastParentheses: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Standard:        Cpp11
+TabWidth:        8
+UseTab:          Never
+...
diff --git a/macos/third_party/mlx/.gitignore b/macos/third_party/mlx/.gitignore
new file mode 100644
index 0000000..1daaa46
--- /dev/null
+++ b/macos/third_party/mlx/.gitignore
@@ -0,0 +1,81 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# tensor files
+*.safe
+*.safetensors
+
+# Metal libraries
+*.metallib
+
+# Distribution / packaging
+python/mlx/core
+python/mlx/share
+python/mlx/include
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+venv/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+uv.lock
+.DS_Store
+
+# Prerequisites
+*.d
+
+# Compiled Object files
+*.slo
+*.lo
+*.o
+*.obj
+*.ilk
+
+# Precompiled Headers
+*.gch
+*.pch
+
+# Compiled Dynamic libraries
+*.so
+*.dylib
+*.dll
+
+# Fortran module files
+*.mod
+*.smod
+
+# Compiled Static libraries
+*.lai
+*.la
+*.a
+*.lib
+
+# Executables
+*.exe
+*.out
+*.app
+
+# Debug symbols
+*.pdb
+
+# VSCode
+.vscode/
+# Jetbrains
+.cache/
+# vim
+*.swp
diff --git a/macos/third_party/mlx/ACKNOWLEDGMENTS.md b/macos/third_party/mlx/ACKNOWLEDGMENTS.md
new file mode 100644
index 0000000..186908f
--- /dev/null
+++ b/macos/third_party/mlx/ACKNOWLEDGMENTS.md
@@ -0,0 +1,268 @@
+# Individual Contributors
+
+If you wish to be acknowledged for your contributions, please list your name
+with a short description of your contribution(s) below. For example:
+
+- Jane Smith: Added the `foo` and `bar` ops.
+
+MLX was developed with contributions from the following individuals:
+
+- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops. Added `clip_grad_norm` along with `tree_reduce`. Added `cross`. Added `orthogonal` initializer.
+- Juarez Bochi: Fixed bug in cross attention.
+- Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example.
+- Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream`, safetensors support, `einsum`, and `einsum_path`.
+- Gabrijel Boduljak: Added `mlx.core.linalg`, implemented `norm` method and `InstanceNorm` layer. Implemented pooling layers and ``Upsample``.
+- Hinrik Snær Guðmundsson: Added `atleast_1d`, `atleast_2d`, `atleast_3d` ops.
+- Luca Arnaboldi: Added `Ceil` and `Floor` ops; implemented pickling, copy and deepcopy for mlx arrays.
+- Brian Keene & Atila Orhon, with Argmax Inc.: Added `fast.scaled_dot_product_attention`
+- AmirHossein Razlighi: Added chaining support for some of the ops in `nn.Module`. Comparison works for non array objects in `mlx.core.array`. Exception handling for invalid operations in `mlx.core.array`.
+- Gleb Pobudzey: Added the `where` primitive, and groups in 1D and 2D convolutions.
+- Paul Paczuski: Improved stability of BCE loss calculation
+- Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops.
+- Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function.
+
+<a href="https://github.com/ml-explore/mlx/graphs/contributors">
+  <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />
+</a>
+
+# Organizations
+
+MLX has received contributions from the following companies:
+- NVIDIA Corporation & Affiliates
+
+# Third-Party Software
+
+MLX leverages several third-party software, listed here together with
+their license copied verbatim.
+
+## PocketFFT
+
+Copyright (C) 2010-2018 Max-Planck-Society
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+* Redistributions in binary form must reproduce the above copyright notice, this
+  list of conditions and the following disclaimer in the documentation and/or
+  other materials provided with the distribution.
+* Neither the name of the copyright holder nor the names of its contributors may
+  be used to endorse or promote products derived from this software without
+  specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+## metal-cpp
+
+                              Apache License
+                        Version 2.0, January 2004
+                    http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+  "License" shall mean the terms and conditions for use, reproduction,
+  and distribution as defined by Sections 1 through 9 of this document.
+
+  "Licensor" shall mean the copyright owner or entity authorized by
+  the copyright owner that is granting the License.
+
+  "Legal Entity" shall mean the union of the acting entity and all
+  other entities that control, are controlled by, or are under common
+  control with that entity. For the purposes of this definition,
+  "control" means (i) the power, direct or indirect, to cause the
+  direction or management of such entity, whether by contract or
+  otherwise, or (ii) ownership of fifty percent (50%) or more of the
+  outstanding shares, or (iii) beneficial ownership of such entity.
+
+  "You" (or "Your") shall mean an individual or Legal Entity
+  exercising permissions granted by this License.
+
+  "Source" form shall mean the preferred form for making modifications,
+  including but not limited to software source code, documentation
+  source, and configuration files.
+
+  "Object" form shall mean any form resulting from mechanical
+  transformation or translation of a Source form, including but
+  not limited to compiled object code, generated documentation,
+  and conversions to other media types.
+
+  "Work" shall mean the work of authorship, whether in Source or
+  Object form, made available under the License, as indicated by a
+  copyright notice that is included in or attached to the work
+  (an example is provided in the Appendix below).
+
+  "Derivative Works" shall mean any work, whether in Source or Object
+  form, that is based on (or derived from) the Work and for which the
+  editorial revisions, annotations, elaborations, or other modifications
+  represent, as a whole, an original work of authorship. For the purposes
+  of this License, Derivative Works shall not include works that remain
+  separable from, or merely link (or bind by name) to the interfaces of,
+  the Work and Derivative Works thereof.
+
+  "Contribution" shall mean any work of authorship, including
+  the original version of the Work and any modifications or additions
+  to that Work or Derivative Works thereof, that is intentionally
+  submitted to Licensor for inclusion in the Work by the copyright owner
+  or by an individual or Legal Entity authorized to submit on behalf of
+  the copyright owner. For the purposes of this definition, "submitted"
+  means any form of electronic, verbal, or written communication sent
+  to the Licensor or its representatives, including but not limited to
+  communication on electronic mailing lists, source code control systems,
+  and issue tracking systems that are managed by, or on behalf of, the
+  Licensor for the purpose of discussing and improving the Work, but
+  excluding communication that is conspicuously marked or otherwise
+  designated in writing by the copyright owner as "Not a Contribution."
+
+  "Contributor" shall mean Licensor and any individual or Legal Entity
+  on behalf of whom a Contribution has been received by Licensor and
+  subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of
+  this License, each Contributor hereby grants to You a perpetual,
+  worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+  copyright license to reproduce, prepare Derivative Works of,
+  publicly display, publicly perform, sublicense, and distribute the
+  Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of
+  this License, each Contributor hereby grants to You a perpetual,
+  worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+  (except as stated in this section) patent license to make, have made,
+  use, offer to sell, sell, import, and otherwise transfer the Work,
+  where such license applies only to those patent claims licensable
+  by such Contributor that are necessarily infringed by their
+  Contribution(s) alone or by combination of their Contribution(s)
+  with the Work to which such Contribution(s) was submitted. If You
+  institute patent litigation against any entity (including a
+  cross-claim or counterclaim in a lawsuit) alleging that the Work
+  or a Contribution incorporated within the Work constitutes direct
+  or contributory patent infringement, then any patent licenses
+  granted to You under this License for that Work shall terminate
+  as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the
+  Work or Derivative Works thereof in any medium, with or without
+  modifications, and in Source or Object form, provided that You
+  meet the following conditions:
+
+  (a) You must give any other recipients of the Work or
+      Derivative Works a copy of this License; and
+
+  (b) You must cause any modified files to carry prominent notices
+      stating that You changed the files; and
+
+  (c) You must retain, in the Source form of any Derivative Works
+      that You distribute, all copyright, patent, trademark, and
+      attribution notices from the Source form of the Work,
+      excluding those notices that do not pertain to any part of
+      the Derivative Works; and
+
+  (d) If the Work includes a "NOTICE" text file as part of its
+      distribution, then any Derivative Works that You distribute must
+      include a readable copy of the attribution notices contained
+      within such NOTICE file, excluding those notices that do not
+      pertain to any part of the Derivative Works, in at least one
+      of the following places: within a NOTICE text file distributed
+      as part of the Derivative Works; within the Source form or
+      documentation, if provided along with the Derivative Works; or,
+      within a display generated by the Derivative Works, if and
+      wherever such third-party notices normally appear. The contents
+      of the NOTICE file are for informational purposes only and
+      do not modify the License. You may add Your own attribution
+      notices within Derivative Works that You distribute, alongside
+      or as an addendum to the NOTICE text from the Work, provided
+      that such additional attribution notices cannot be construed
+      as modifying the License.
+
+  You may add Your own copyright statement to Your modifications and
+  may provide additional or different license terms and conditions
+  for use, reproduction, or distribution of Your modifications, or
+  for any such Derivative Works as a whole, provided Your use,
+  reproduction, and distribution of the Work otherwise complies with
+  the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise,
+  any Contribution intentionally submitted for inclusion in the Work
+  by You to the Licensor shall be under the terms and conditions of
+  this License, without any additional terms or conditions.
+  Notwithstanding the above, nothing herein shall supersede or modify
+  the terms of any separate license agreement you may have executed
+  with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade
+  names, trademarks, service marks, or product names of the Licensor,
+  except as required for reasonable and customary use in describing the
+  origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or
+  agreed to in writing, Licensor provides the Work (and each
+  Contributor provides its Contributions) on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+  implied, including, without limitation, any warranties or conditions
+  of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+  PARTICULAR PURPOSE. You are solely responsible for determining the
+  appropriateness of using or redistributing the Work and assume any
+  risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory,
+  whether in tort (including negligence), contract, or otherwise,
+  unless required by applicable law (such as deliberate and grossly
+  negligent acts) or agreed to in writing, shall any Contributor be
+  liable to You for damages, including any direct, indirect, special,
+  incidental, or consequential damages of any character arising as a
+  result of this License or out of the use or inability to use the
+  Work (including but not limited to damages for loss of goodwill,
+  work stoppage, computer failure or malfunction, or any and all
+  other commercial damages or losses), even if such Contributor
+  has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing
+  the Work or Derivative Works thereof, You may choose to offer,
+  and charge a fee for, acceptance of support, warranty, indemnity,
+  or other liability obligations and/or rights consistent with this
+  License. However, in accepting such obligations, You may act only
+  on Your own behalf and on Your sole responsibility, not on behalf
+  of any other Contributor, and only if You agree to indemnify,
+  defend, and hold each Contributor harmless for any liability
+  incurred by, or claims asserted against, such Contributor by reason
+  of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+  To apply the Apache License to your work, attach the following
+  boilerplate notice, with the fields enclosed by brackets "[]"
+  replaced with your own identifying information. (Don't include
+  the brackets!)  The text should be enclosed in the appropriate
+  comment syntax for the file format. We also recommend that a
+  file or class name and description of purpose be included on the
+  same "printed page" as the copyright notice for easier
+  identification within third-party archives.
+
+Copyright © 2023 Apple Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
diff --git a/macos/third_party/mlx/CMakeLists.txt b/macos/third_party/mlx/CMakeLists.txt
new file mode 100644
index 0000000..041a476
--- /dev/null
+++ b/macos/third_party/mlx/CMakeLists.txt
@@ -0,0 +1,437 @@
+cmake_minimum_required(VERSION 3.25)
+
+if(NOT MLX_VERSION)
+  file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$")
+  string(REGEX MATCH "#define MLX_VERSION_MAJOR ([0-9]+)" _ "${_mlx_h_version}")
+  set(_major ${CMAKE_MATCH_1})
+  string(REGEX MATCH "#define MLX_VERSION_MINOR ([0-9]+)" _ "${_mlx_h_version}")
+  set(_minor ${CMAKE_MATCH_1})
+  string(REGEX MATCH "#define MLX_VERSION_PATCH ([0-9]+)" _ "${_mlx_h_version}")
+  set(_patch ${CMAKE_MATCH_1})
+  set(MLX_PROJECT_VERSION "${_major}.${_minor}.${_patch}")
+  set(MLX_VERSION ${MLX_PROJECT_VERSION})
+else()
+  string(REGEX REPLACE "^([0-9]+\.[0-9]+\.[0-9]+).*" "\\1" MLX_PROJECT_VERSION
+                       ${MLX_VERSION})
+endif()
+
+project(
+  mlx
+  LANGUAGES C CXX
+  VERSION ${MLX_PROJECT_VERSION})
+
+# ----------------------------- Setup -----------------------------
+set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+set(CMAKE_INSTALL_MESSAGE NEVER)
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+
+# ----------------------------- Configuration -----------------------------
+option(MLX_BUILD_TESTS "Build tests for mlx" ON)
+option(MLX_BUILD_EXAMPLES "Build examples for mlx" ON)
+option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF)
+option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF)
+option(MLX_BUILD_METAL "Build metal backend" ON)
+option(MLX_BUILD_CPU "Build cpu backend" ON)
+option(MLX_BUILD_CUDA "Build cuda backend" OFF)
+option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
+option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
+option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
+option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
+option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON)
+option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
+option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON)
+option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
+option(USE_SYSTEM_FMT "Use system's provided fmt library" OFF)
+option(USE_ASAN "Enable AddressSanitizer (ASan)" OFF)
+option(USE_UBSAN "Enable UndefinedBehaviorSanitizer (UBSan)" OFF)
+option(USE_TSAN "Enable ThreadSanitizer (TSan)" OFF)
+
+# --------------------- Processor tests -------------------------
+message(
+  STATUS
+    "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}"
+)
+
+if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+  if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "x86_64")
+    if(NOT MLX_ENABLE_X64_MAC)
+      message(
+        FATAL_ERROR
+          "Building for x86_64 on macOS is not supported."
+          " If you are on an Apple silicon system, check the build"
+          " documentation for possible fixes: "
+          "https://ml-explore.github.io/mlx/build/html/install.html#build-from-source"
+      )
+    else()
+      set(MLX_BUILD_METAL OFF)
+      message(WARNING "Building for x86_64 arch is not officially supported.")
+    endif()
+  endif()
+else()
+  set(MLX_BUILD_METAL OFF)
+endif()
+
+if(MLX_USE_CCACHE)
+  find_program(CCACHE_PROGRAM ccache)
+  if(CCACHE_PROGRAM)
+    message(STATUS "Found CCache: ${CCACHE_PROGRAM}")
+    set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+    set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+    set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+  endif()
+endif()
+
+if(USE_ASAN AND USE_TSAN)
+  message(
+    FATAL_ERROR
+      "AddressSanitizer (ASan) and ThreadSanitizer (TSan) are mutually exclusive and cannot be enabled at the same time."
+  )
+endif()
+
+set(SANITIZER_COMPILE_FLAGS "")
+set(SANITIZER_LINK_FLAGS "")
+
+if(USE_ASAN)
+  if(WIN32 AND MSVC)
+    list(APPEND SANITIZER_COMPILE_FLAGS /fsanitize=address)
+    list(APPEND SANITIZER_LINK_FLAGS /fsanitize=address)
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=address)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=address)
+    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
+    endif()
+  endif()
+endif()
+
+if(USE_UBSAN)
+  if(WIN32 AND MSVC)
+    if(CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
+      list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
+      list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
+    else()
+      message(
+        WARNING
+          "UndefinedBehaviorSanitizer (UBSan) is not directly supported via a simple flag in MSVC."
+      )
+    endif()
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
+  endif()
+endif()
+
+if(USE_TSAN)
+  if(WIN32 AND MSVC)
+    message(
+      FATAL_ERROR
+        "ThreadSanitizer (TSan) is not supported by the MSVC compiler. Please use Clang or GCC."
+    )
+  elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
+    message(FATAL_ERROR "ThreadSanitizer (TSan) is not supported on macOS.")
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=thread)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=thread)
+    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
+    endif()
+  endif()
+endif()
+
+# ----------------------------- Lib -----------------------------
+
+include(FetchContent)
+# Avoid warning about DOWNLOAD_EXTRACT_TIMESTAMP in CMake 3.24:
+cmake_policy(SET CMP0135 NEW)
+
+add_library(mlx)
+
+target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS})
+target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS})
+
+if(MLX_BUILD_CUDA)
+  enable_language(CUDA)
+  find_package(CUDAToolkit REQUIRED)
+  find_package(CUDNN REQUIRED)
+endif()
+
+if(MLX_BUILD_METAL)
+  find_library(METAL_LIB Metal)
+  find_library(FOUNDATION_LIB Foundation)
+  find_library(QUARTZ_LIB QuartzCore)
+  if(METAL_LIB)
+    message(STATUS "Metal found ${METAL_LIB}")
+  else()
+    message(
+      FATAL_ERROR
+        "Metal not found. Set MLX_BUILD_METAL=OFF to build without GPU")
+  endif()
+
+  if(MLX_METAL_DEBUG)
+    add_compile_definitions(MLX_METAL_DEBUG)
+  endif()
+
+  # Throw an error if xcrun not found
+  execute_process(
+    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
+    OUTPUT_VARIABLE MACOS_SDK_VERSION
+    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
+
+  if(${MACOS_SDK_VERSION} LESS 14.0)
+    message(
+      FATAL_ERROR
+        "MLX requires macOS SDK >= 14.0 to be built with MLX_BUILD_METAL=ON")
+  endif()
+  message(STATUS "Building with macOS SDK version ${MACOS_SDK_VERSION}")
+
+  set(METAL_CPP_URL
+      https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip)
+
+  if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
+    if(${CMAKE_OSX_DEPLOYMENT_TARGET} LESS 14.0)
+      message(FATAL_ERROR "MLX requires macOS >= 14.0")
+    endif()
+    set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+  endif()
+  execute_process(
+    COMMAND
+      zsh "-c"
+      "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
+    OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
+  FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
+  FetchContent_MakeAvailable(metal_cpp)
+  target_include_directories(
+    mlx PUBLIC $<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>
+               $<INSTALL_INTERFACE:include/metal_cpp>)
+  target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB})
+endif()
+
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  # With newer clang/gcc versions following libs are implicitly linked, but when
+  # building on old distributions they need to be explicitly listed.
+  target_link_libraries(mlx PRIVATE dl pthread)
+endif()
+
+if(WIN32)
+  if(MSVC)
+    # GGUF does not build with MSVC.
+    set(MLX_BUILD_GGUF OFF)
+  endif()
+  # Generate DLL and EXE in the same dir, otherwise EXE will not be able to run.
+  # This is only done when MLX is built as the top project.
+  if(CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR)
+    set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
+  endif()
+  # Windows implementation of dlfcn.h APIs.
+  FetchContent_Declare(
+    dlfcn-win32
+    GIT_REPOSITORY https://github.com/dlfcn-win32/dlfcn-win32.git
+    GIT_TAG v1.4.2
+    EXCLUDE_FROM_ALL)
+  block()
+  set(BUILD_SHARED_LIBS OFF)
+  FetchContent_MakeAvailable(dlfcn-win32)
+  endblock()
+  target_include_directories(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src")
+  target_link_libraries(mlx PRIVATE dl)
+endif()
+
+if(MLX_BUILD_CPU)
+  find_library(ACCELERATE_LIBRARY Accelerate)
+  if(ACCELERATE_LIBRARY)
+    message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}")
+    set(MLX_BUILD_ACCELERATE ON)
+  else()
+    message(STATUS "Accelerate not found, using default backend.")
+    set(MLX_BUILD_ACCELERATE OFF)
+  endif()
+
+  if(MLX_BUILD_ACCELERATE)
+    target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
+    add_compile_definitions(MLX_USE_ACCELERATE)
+    add_compile_definitions(ACCELERATE_NEW_LAPACK)
+  elseif(WIN32)
+    # Download and link prebuilt binaries of OpenBLAS. Note that we can only
+    # link with the dynamic library, the prebuilt binaries were built with MinGW
+    # so static-linking would require linking with MinGW's runtime.
+    FetchContent_Declare(
+      openblas
+      URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.31/OpenBLAS-0.3.31-x64.zip"
+    )
+    FetchContent_MakeAvailable(openblas)
+    target_link_libraries(mlx
+                          PRIVATE "${openblas_SOURCE_DIR}/lib/libopenblas.lib")
+    target_include_directories(mlx PRIVATE "${openblas_SOURCE_DIR}/include")
+    # Make sure the DLL file is placed in the same dir with executables.
+    set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/libopenblas.dll")
+    add_custom_command(
+      TARGET mlx
+      POST_BUILD
+      COMMAND ${CMAKE_COMMAND} -E copy_if_different ${OPENBLAS_DLL_FILE}
+              ${CMAKE_BINARY_DIR})
+  else()
+    if(${CMAKE_HOST_APPLE})
+      # The blas shipped in macOS SDK is not supported, search homebrew for
+      # openblas instead.
+      set(BLA_VENDOR OpenBLAS)
+      set(LAPACK_ROOT
+          "${LAPACK_ROOT};$ENV{LAPACK_ROOT};/usr/local/opt/openblas")
+    endif()
+    # Search and link with lapack.
+    find_package(LAPACK REQUIRED)
+    if(NOT LAPACK_FOUND)
+      message(FATAL_ERROR "Must have LAPACK installed")
+    endif()
+    find_path(LAPACK_INCLUDE_DIRS lapacke.h /usr/include /usr/local/include
+              /usr/local/opt/openblas/include)
+    message(STATUS "Lapack lib " ${LAPACK_LIBRARIES})
+    message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS})
+    target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS})
+    target_link_libraries(mlx PRIVATE ${LAPACK_LIBRARIES})
+    # List blas after lapack otherwise we may accidentally incldue an old
+    # version of lapack.h from the include dirs of blas.
+    find_package(BLAS REQUIRED)
+    if(NOT BLAS_FOUND)
+      message(FATAL_ERROR "Must have BLAS installed")
+    endif()
+    # TODO find a cleaner way to do this
+    find_path(BLAS_INCLUDE_DIRS cblas.h /usr/include /usr/local/include
+              $ENV{BLAS_HOME}/include)
+    message(STATUS "Blas lib " ${BLAS_LIBRARIES})
+    message(STATUS "Blas include " ${BLAS_INCLUDE_DIRS})
+    target_include_directories(mlx PRIVATE ${BLAS_INCLUDE_DIRS})
+    target_link_libraries(mlx PRIVATE ${BLAS_LIBRARIES})
+  endif()
+else()
+  set(MLX_BUILD_ACCELERATE OFF)
+endif()
+
+message(STATUS "Downloading json")
+FetchContent_Declare(
+  json
+  URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
+FetchContent_MakeAvailable(json)
+target_include_directories(
+  mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)
+
+add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)
+
+target_include_directories(
+  mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
+             $<INSTALL_INTERFACE:include>)
+
+if(USE_SYSTEM_FMT)
+  find_package(fmt REQUIRED)
+else()
+  FetchContent_Declare(
+    fmt
+    GIT_REPOSITORY https://github.com/fmtlib/fmt.git
+    GIT_TAG 12.1.0
+    EXCLUDE_FROM_ALL)
+  FetchContent_MakeAvailable(fmt)
+endif()
+target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:fmt::fmt-header-only>)
+
+if(MLX_BUILD_PYTHON_BINDINGS)
+  message(STATUS "Building Python bindings.")
+  find_package(
+    Python 3.10
+    COMPONENTS Interpreter Development.Module
+    REQUIRED)
+  FetchContent_Declare(
+    nanobind
+    GIT_REPOSITORY https://github.com/wjakob/nanobind.git
+    GIT_TAG v2.10.2
+    GIT_SHALLOW TRUE
+    EXCLUDE_FROM_ALL)
+  FetchContent_MakeAvailable(nanobind)
+  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
+endif()
+
+if(MLX_BUILD_TESTS)
+  include(CTest)
+  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/tests)
+endif()
+
+if(MLX_BUILD_EXAMPLES)
+  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/examples/cpp)
+endif()
+
+if(MLX_BUILD_BENCHMARKS)
+  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/benchmarks/cpp)
+endif()
+
+# ----------------------------- Installation -----------------------------
+include(GNUInstallDirs)
+
+if(WIN32)
+  # Install DLLs to the same dir with extension file (core.pyd) on Windows.
+  set(CMAKE_INSTALL_BINDIR ".")
+  if(MLX_BUILD_CPU)
+    # Install OpenBLAS.
+    install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN)
+  endif()
+endif()
+
+# Install library
+install(
+  TARGETS mlx
+  EXPORT MLXTargets
+  LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
+  INCLUDES
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
+
+# Install headers
+install(
+  DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/mlx
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+  COMPONENT headers
+  FILES_MATCHING
+  PATTERN "*.h"
+  PATTERN "backend/metal/kernels.h" EXCLUDE)
+
+# Install metal dependencies
+if(MLX_BUILD_METAL)
+
+  # Install metal cpp
+  install(
+    DIRECTORY ${metal_cpp_SOURCE_DIR}/
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/metal_cpp
+    COMPONENT metal_cpp_source)
+
+endif()
+
+# Install cmake config
+set(MLX_CMAKE_BUILD_CONFIG ${CMAKE_BINARY_DIR}/MLXConfig.cmake)
+set(MLX_CMAKE_BUILD_VERSION_CONFIG ${CMAKE_BINARY_DIR}/MLXConfigVersion.cmake)
+set(MLX_CMAKE_INSTALL_MODULE_DIR share/cmake/MLX)
+
+install(
+  EXPORT MLXTargets
+  FILE MLXTargets.cmake
+  DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR})
+
+include(CMakePackageConfigHelpers)
+
+write_basic_package_version_file(
+  ${MLX_CMAKE_BUILD_VERSION_CONFIG}
+  COMPATIBILITY SameMajorVersion
+  VERSION ${MLX_VERSION})
+
+configure_package_config_file(
+  ${CMAKE_CURRENT_LIST_DIR}/mlx.pc.in ${MLX_CMAKE_BUILD_CONFIG}
+  INSTALL_DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}
+  NO_CHECK_REQUIRED_COMPONENTS_MACRO
+  PATH_VARS CMAKE_INSTALL_LIBDIR CMAKE_INSTALL_INCLUDEDIR
+            MLX_CMAKE_INSTALL_MODULE_DIR)
+
+install(FILES ${MLX_CMAKE_BUILD_CONFIG} ${MLX_CMAKE_BUILD_VERSION_CONFIG}
+        DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR})
+
+install(DIRECTORY ${CMAKE_MODULE_PATH}/
+        DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR})
diff --git a/macos/third_party/mlx/LICENSE b/macos/third_party/mlx/LICENSE
new file mode 100644
index 0000000..98ff47b
--- /dev/null
+++ b/macos/third_party/mlx/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright © 2023 Apple Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/macos/third_party/mlx/README.md b/macos/third_party/mlx/README.md
new file mode 100644
index 0000000..fb30a24
--- /dev/null
+++ b/macos/third_party/mlx/README.md
@@ -0,0 +1,121 @@
+# MLX
+
+[**Quickstart**](#quickstart) | [**Installation**](#installation) |
+[**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) |
+[**Examples**](#examples)
+
+[![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)
+
+MLX is an array framework for machine learning on Apple silicon,
+brought to you by Apple machine learning research.
+
+Some key features of MLX include:
+
+- **Familiar APIs**: MLX has a Python API that closely follows NumPy. MLX
+   also has fully featured C++, [C](https://github.com/ml-explore/mlx-c), and
+   [Swift](https://github.com/ml-explore/mlx-swift/) APIs, which closely mirror
+   the Python API. MLX has higher-level packages like `mlx.nn` and
+   `mlx.optimizers` with APIs that closely follow PyTorch to simplify building
+   more complex models.
+
+- **Composable function transformations**: MLX supports composable function
+  transformations for automatic differentiation, automatic vectorization,
+  and computation graph optimization.
+
+- **Lazy computation**: Computations in MLX are lazy. Arrays are only
+  materialized when needed.
+
+- **Dynamic graph construction**: Computation graphs in MLX are constructed
+  dynamically. Changing the shapes of function arguments does not trigger
+  slow compilations, and debugging is simple and intuitive.
+
+- **Multi-device**: Operations can run on any of the supported devices
+  (currently the CPU and the GPU).
+
+- **Unified memory**: A notable difference from MLX and other frameworks
+  is the *unified memory model*. Arrays in MLX live in shared memory.
+  Operations on MLX arrays can be performed on any of the supported
+  device types without transferring data.
+
+MLX is designed by machine learning researchers for machine learning
+researchers. The framework is intended to be user-friendly, but still efficient
+to train and deploy models. The design of the framework itself is also
+conceptually simple. We intend to make it easy for researchers to extend and
+improve MLX with the goal of quickly exploring new ideas.
+
+The design of MLX is inspired by frameworks like
+[NumPy](https://numpy.org/doc/stable/index.html),
+[PyTorch](https://pytorch.org/), [Jax](https://github.com/google/jax), and
+[ArrayFire](https://arrayfire.org/).
+
+## Examples
+
+The [MLX examples repo](https://github.com/ml-explore/mlx-examples) has a
+variety of examples, including:
+
+- [Transformer language model](https://github.com/ml-explore/mlx-examples/tree/main/transformer_lm) training.
+- Large-scale text generation with
+  [LLaMA](https://github.com/ml-explore/mlx-examples/tree/main/llms/llama) and
+  finetuning with [LoRA](https://github.com/ml-explore/mlx-examples/tree/main/lora).
+- Generating images with [Stable Diffusion](https://github.com/ml-explore/mlx-examples/tree/main/stable_diffusion).
+- Speech recognition with [OpenAI's Whisper](https://github.com/ml-explore/mlx-examples/tree/main/whisper).
+
+## Quickstart
+
+See the [quick start
+guide](https://ml-explore.github.io/mlx/build/html/usage/quick_start.html)
+in the documentation.
+
+## Installation
+
+MLX is available on [PyPI](https://pypi.org/project/mlx/). To install MLX on
+macOS, run:
+
+```bash
+pip install mlx
+```
+
+To install the CUDA backend on Linux, run:
+
+```bash
+pip install mlx[cuda]
+```
+
+To install a CPU-only Linux package, run:
+
+```bash
+pip install mlx[cpu]
+```
+
+Checkout the
+[documentation](https://ml-explore.github.io/mlx/build/html/install.html#)
+for more information on building the C++ and Python APIs from source.
+
+## Contributing
+
+Check out the [contribution guidelines](https://github.com/ml-explore/mlx/tree/main/CONTRIBUTING.md) for more information
+on contributing to MLX. See the
+[docs](https://ml-explore.github.io/mlx/build/html/install.html) for more
+information on building from source, and running tests.
+
+We are grateful for all of [our
+contributors](https://github.com/ml-explore/mlx/tree/main/ACKNOWLEDGMENTS.md#Individual-Contributors). If you contribute
+to MLX and wish to be acknowledged, please add your name to the list in your
+pull request.
+
+## Citing MLX
+
+The MLX software suite was initially developed with equal contribution by Awni
+Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
+MLX useful in your research and wish to cite it, please use the following
+BibTex entry:
+
+```text
+@software{mlx2023,
+  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
+  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
+  url = {https://github.com/ml-explore},
+  version = {0.0},
+  year = {2023},
+}
+```
diff --git a/macos/third_party/mlx/cmake/FindCUDNN.cmake b/macos/third_party/mlx/cmake/FindCUDNN.cmake
new file mode 100644
index 0000000..e91f2f2
--- /dev/null
+++ b/macos/third_party/mlx/cmake/FindCUDNN.cmake
@@ -0,0 +1,177 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Modified from
+# https://github.com/NVIDIA/cudnn-frontend/blob/main/cmake/cuDNN.cmake
+
+# Return the last file matching the pattern.
+function(find_file_glob VAR PATTERN)
+  file(GLOB _RESULT "${PATTERN}")
+  if(_RESULT)
+    list(LENGTH ${_RESULT} _RESULT_LENGTH)
+    if(_RESULT_LENGTH GREATER 0)
+      list(GET ${_RESULT} -1 _RESULT)
+    endif()
+    set(${VAR}
+        "${_RESULT}"
+        PARENT_SCOPE)
+  endif()
+endfunction()
+
+# Find the dir including the "cudnn.h" file.
+find_path(
+  CUDNN_INCLUDE_DIR cudnn.h
+  HINTS ${CUDNN_INCLUDE_PATH} ${CUDAToolkit_INCLUDE_DIRS}
+  PATH_SUFFIXES include OPTIONAL)
+
+# Glob searching "cudnn.h" for Windows.
+if(WIN32 AND NOT CUDNN_INCLUDE_DIR)
+  find_file_glob(
+    CUDNN_H_PATH
+    "C:/Program Files/NVIDIA/CUDNN/*/include/${CUDAToolkit_VERSION_MAJOR}.*/cudnn.h"
+  )
+  if(CUDNN_H_PATH)
+    get_filename_component(CUDNN_INCLUDE_DIR "${CUDNN_H_PATH}" DIRECTORY)
+  endif()
+endif()
+
+if(NOT CUDNN_INCLUDE_DIR)
+  message(
+    FATAL_ERROR
+      "Unable to find cudnn.h, please make sure cuDNN is installed and pass CUDNN_INCLUDE_PATH to cmake."
+  )
+endif()
+
+# Get cudnn version.
+file(READ "${CUDNN_INCLUDE_DIR}/cudnn_version.h" cudnn_version_header)
+string(REGEX MATCH "#define CUDNN_MAJOR [1-9]+" macrodef
+             "${cudnn_version_header}")
+string(REGEX MATCH "[1-9]+" CUDNN_MAJOR_VERSION "${macrodef}")
+
+# Function for searching library files.
+function(find_cudnn_library NAME)
+  if(NOT "${ARGV1}" STREQUAL "OPTIONAL")
+    set(_CUDNN_REQUIRED TRUE)
+  else()
+    set(_CUDNN_REQUIRED FALSE)
+  endif()
+
+  find_library(
+    ${NAME}_LIBRARY
+    NAMES ${NAME} "lib${NAME}.so.${CUDNN_MAJOR_VERSION}" NAMES_PER_DIR
+    HINTS ${CUDNN_LIBRARY_PATH} ${CUDAToolkit_LIBRARY_DIR}
+    PATH_SUFFIXES lib64 lib/x64 lib OPTIONAL)
+
+  if(WIN32 AND NOT ${NAME}_LIBRARY)
+    find_file_glob(
+      ${NAME}_LIBRARY
+      "C:/Program Files/NVIDIA/CUDNN/*/lib/${CUDAToolkit_VERSION_MAJOR}.*/x64/${NAME}.lib"
+    )
+  endif()
+
+  if(NOT ${NAME}_LIBRARY AND ${_CUDNN_REQUIRED})
+    message(
+      FATAL_ERROR
+        "Unable to find ${NAME}, please make sure cuDNN is installed and pass CUDNN_LIBRARY_PATH to cmake."
+    )
+  endif()
+
+  if(${NAME}_LIBRARY)
+    add_library(CUDNN::${NAME} UNKNOWN IMPORTED)
+    set_target_properties(
+      CUDNN::${NAME}
+      PROPERTIES INTERFACE_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_DIR}
+                 IMPORTED_LOCATION ${${NAME}_LIBRARY})
+    set(${NAME}_LIBRARY
+        "${${NAME}_LIBRARY}"
+        PARENT_SCOPE)
+  else()
+    message(STATUS "${NAME} not found.")
+  endif()
+endfunction()
+
+# Search for the main cudnn library.
+find_cudnn_library(cudnn)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(CUDNN REQUIRED_VARS CUDNN_INCLUDE_DIR
+                                                      cudnn_LIBRARY)
+
+if(CUDNN_INCLUDE_DIR AND cudnn_LIBRARY)
+  set(CUDNN_FOUND
+      ON
+      CACHE INTERNAL "cuDNN Library Found")
+else()
+  set(CUDNN_FOUND
+      OFF
+      CACHE INTERNAL "cuDNN Library Not Found")
+endif()
+
+# Find out all the DLL files for Windows.
+if(WIN32 AND cudnn_LIBRARY)
+  get_filename_component(CUDNN_BIN_DIR "${cudnn_LIBRARY}" DIRECTORY)
+  string(REPLACE "/lib/" "/bin/" CUDNN_BIN_DIR "${CUDNN_BIN_DIR}")
+  file(
+    GLOB CUDNN_DLL_NAMES
+    RELATIVE "${CUDNN_BIN_DIR}"
+    "${CUDNN_BIN_DIR}/*.dll")
+endif()
+
+# Create an interface library that users can link with.
+add_library(CUDNN::cudnn_all INTERFACE IMPORTED)
+target_link_libraries(CUDNN::cudnn_all INTERFACE CUDNN::cudnn)
+target_include_directories(
+  CUDNN::cudnn_all INTERFACE $<INSTALL_INTERFACE:include>
+                             $<BUILD_INTERFACE:${CUDNN_INCLUDE_DIR}>)
+
+# Add other components of cudnn.
+if(CUDNN_MAJOR_VERSION EQUAL 8)
+  find_cudnn_library(cudnn_adv_infer)
+  find_cudnn_library(cudnn_adv_train)
+  find_cudnn_library(cudnn_cnn_infer)
+  find_cudnn_library(cudnn_cnn_train)
+  find_cudnn_library(cudnn_ops_infer)
+  find_cudnn_library(cudnn_ops_train)
+
+  target_link_libraries(
+    CUDNN::cudnn_all
+    INTERFACE CUDNN::cudnn_adv_train CUDNN::cudnn_ops_train
+              CUDNN::cudnn_cnn_train CUDNN::cudnn_adv_infer
+              CUDNN::cudnn_cnn_infer CUDNN::cudnn_ops_infer)
+
+elseif(CUDNN_MAJOR_VERSION EQUAL 9)
+  find_cudnn_library(cudnn_graph)
+  find_cudnn_library(cudnn_engines_runtime_compiled)
+  find_cudnn_library(cudnn_ops OPTIONAL)
+  find_cudnn_library(cudnn_cnn OPTIONAL)
+  find_cudnn_library(cudnn_adv OPTIONAL)
+  find_cudnn_library(cudnn_engines_precompiled OPTIONAL)
+  find_cudnn_library(cudnn_heuristic OPTIONAL)
+
+  target_link_libraries(
+    CUDNN::cudnn_all
+    INTERFACE CUDNN::cudnn_graph
+              CUDNN::cudnn_engines_runtime_compiled
+              CUDNN::cudnn_ops
+              CUDNN::cudnn_cnn
+              CUDNN::cudnn_adv
+              CUDNN::cudnn_engines_precompiled
+              CUDNN::cudnn_heuristic)
+endif()
diff --git a/macos/third_party/mlx/cmake/FindNCCL.cmake b/macos/third_party/mlx/cmake/FindNCCL.cmake
new file mode 100644
index 0000000..b318932
--- /dev/null
+++ b/macos/third_party/mlx/cmake/FindNCCL.cmake
@@ -0,0 +1,54 @@
+# FindNCCL.cmake This module finds the NVIDIA NCCL library and its include
+# directories.
+
+set(NCCL_ROOT_DIR
+    $ENV{NCCL_ROOT_DIR}
+    CACHE PATH "Folder contains NVIDIA NCCL")
+
+find_path(
+  NCCL_INCLUDE_DIRS
+  NAMES nccl.h
+  HINTS ${NCCL_INCLUDE_DIR} ${NCCL_ROOT_DIR} ${NCCL_ROOT_DIR}/include
+        ${CUDA_TOOLKIT_ROOT_DIR}/include)
+
+if($ENV{USE_STATIC_NCCL})
+  message(
+    STATUS "USE_STATIC_NCCL detected. Linking against static NCCL library")
+  set(NCCL_LIBNAME "libnccl_static.a")
+else()
+  set(NCCL_LIBNAME "nccl")
+endif()
+
+find_library(
+  NCCL_LIBRARIES
+  NAMES ${NCCL_LIBNAME}
+  HINTS ${NCCL_LIB_DIR}
+        ${NCCL_ROOT_DIR}
+        ${NCCL_ROOT_DIR}/lib
+        ${NCCL_ROOT_DIR}/lib/x86_64-linux-gnu
+        ${NCCL_ROOT_DIR}/lib64
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib
+        ${CUDA_TOOLKIT_ROOT_DIR}/lib64)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(NCCL DEFAULT_MSG NCCL_INCLUDE_DIRS
+                                  NCCL_LIBRARIES)
+
+if(NCCL_FOUND)
+  set(NCCL_HEADER_FILE "${NCCL_INCLUDE_DIRS}/nccl.h")
+  message(
+    STATUS "Determining NCCL version from the header file: ${NCCL_HEADER_FILE}")
+  file(
+    STRINGS ${NCCL_HEADER_FILE} NCCL_MAJOR_VERSION_DEFINED
+    REGEX "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+[0-9]+.*$"
+    LIMIT_COUNT 1)
+  if(NCCL_MAJOR_VERSION_DEFINED)
+    string(REGEX REPLACE "^[ \t]*#define[ \t]+NCCL_MAJOR[ \t]+" ""
+                         NCCL_MAJOR_VERSION ${NCCL_MAJOR_VERSION_DEFINED})
+    message(STATUS "NCCL_MAJOR_VERSION: ${NCCL_MAJOR_VERSION}")
+  endif()
+  message(
+    STATUS
+      "Found NCCL (include: ${NCCL_INCLUDE_DIRS}, library: ${NCCL_LIBRARIES})")
+  mark_as_advanced(NCCL_ROOT_DIR NCCL_INCLUDE_DIRS NCCL_LIBRARIES)
+endif()
diff --git a/macos/third_party/mlx/cmake/Findnvpl.cmake b/macos/third_party/mlx/cmake/Findnvpl.cmake
new file mode 100644
index 0000000..011d303
--- /dev/null
+++ b/macos/third_party/mlx/cmake/Findnvpl.cmake
@@ -0,0 +1,3 @@
+# This file does nothing but to suppress the cmake warning: "By not providing
+# Findnvpl.cmake in CMAKE_MODULE_PATH...", which is caused by the
+# find_package(nvpl) from cmake's builtin FindLAPACK.cmake module.
diff --git a/macos/third_party/mlx/cmake/extension.cmake b/macos/third_party/mlx/cmake/extension.cmake
new file mode 100644
index 0000000..13db804
--- /dev/null
+++ b/macos/third_party/mlx/cmake/extension.cmake
@@ -0,0 +1,50 @@
+include(CMakeParseArguments)
+
+# clang format off
+#
+# ##############################################################################
+# Build metal library
+#
+# Adds a custom target ${TARGET} to build ${OUTPUT_DIRECTORY}/{TITLE}.metallib
+# from list ${SOURCES}, including list ${INCLUDE_DIRS}, depends on list ${DEPS}
+#
+# Args: TARGET: Custom target to be added for the metal library TITLE: Name of
+# the .metallib OUTPUT_DIRECTORY: Where to place ${TITLE}.metallib SOURCES: List
+# of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency
+# files (like headers) DEBUG: Boolean, if true, enables debug compile options
+# for this specific library. If not provided, uses global MLX_METAL_DEBUG.
+#
+# clang format on
+
+macro(mlx_build_metallib)
+  # Parse args
+  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG)
+  set(multiValueArgs SOURCES INCLUDE_DIRS DEPS)
+  cmake_parse_arguments(MTLLIB "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+
+  # Set output
+  set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
+
+  # Collect compile options
+  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
+  if(MLX_METAL_DEBUG OR MTLLIB_DEBUG)
+    set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only
+                               -frecord-sources)
+  endif()
+
+  # Prepare metallib build command
+  add_custom_command(
+    OUTPUT ${MTLLIB_BUILD_TARGET}
+    COMMAND
+      xcrun -sdk macosx metal
+      "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
+      ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
+    DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
+    COMMAND_EXPAND_LISTS
+    COMMENT "Building ${MTLLIB_TITLE}.metallib"
+    VERBATIM)
+
+  # Add metallib custom target
+  add_custom_target(${MTLLIB_TARGET} DEPENDS ${MTLLIB_BUILD_TARGET})
+
+endmacro(mlx_build_metallib)
diff --git a/macos/third_party/mlx/mlx.pc.in b/macos/third_party/mlx/mlx.pc.in
new file mode 100644
index 0000000..c4e2515
--- /dev/null
+++ b/macos/third_party/mlx/mlx.pc.in
@@ -0,0 +1,52 @@
+# Find MLX
+#
+# Defines the following variables:
+#
+#   MLX_FOUND            : True if MLX is found
+#   MLX_INCLUDE_DIRS     : Include directory
+#   MLX_LIBRARIES        : Libraries to link against
+#   MLX_CXX_FLAGS        : Additional compiler flags
+#   MLX_BUILD_ACCELERATE : True if MLX was built with accelerate 
+#   MLX_BUILD_METAL      : True if MLX was built with metal 
+
+@PACKAGE_INIT@
+
+include(@PACKAGE_MLX_CMAKE_INSTALL_MODULE_DIR@/MLXTargets.cmake)
+include(@PACKAGE_MLX_CMAKE_INSTALL_MODULE_DIR@/extension.cmake)
+
+set_and_check(MLX_LIBRARY_DIRS @PACKAGE_CMAKE_INSTALL_LIBDIR@)
+set_and_check(MLX_INCLUDE_DIRS @PACKAGE_CMAKE_INSTALL_INCLUDEDIR@)
+set(MLX_LIBRARIES mlx)
+
+find_library(MLX_LIBRARY mlx PATHS ${MLX_LIBRARY_DIRS})
+
+if (@MLX_BUILD_ACCELERATE@)
+    set(MLX_BUILD_ACCELERATE @MLX_BUILD_ACCELERATE@)
+    set(MLX_CXX_FLAGS ${MLX_CXX_FLAGS} -DACCELERATE_NEW_LAPACK)
+endif()
+
+if (@MLX_BUILD_METAL@)
+    set(MLX_BUILD_METAL @MLX_BUILD_METAL@)
+    set(MLX_CXX_FLAGS ${MLX_CXX_FLAGS} -D_METAL_)
+    set(MLX_INCLUDE_DIRS 
+        "${MLX_INCLUDE_DIRS};"
+        @PACKAGE_CMAKE_INSTALL_INCLUDEDIR@/metal_cpp
+    )
+    if(@MLX_METAL_VERSION@ GREATER_EQUAL 310)
+      set(MLX_INCLUDE_DIRS
+        "${MLX_INCLUDE_DIRS};"
+        @PACKAGE_CMAKE_INSTALL_INCLUDEDIR@/mlx/backend/metal/kernels/metal_3_1)
+    else()
+      set(MLX_INCLUDE_DIRS
+        "${MLX_INCLUDE_DIRS};"
+        @PACKAGE_CMAKE_INSTALL_INCLUDEDIR@/mlx/backend/metal/kernels/metal_3_0)
+    endif()
+endif()
+
+set_target_properties(mlx PROPERTIES
+    CXX_STANDARD 17
+    INTERFACE_COMPILE_OPTIONS "${MLX_CXX_FLAGS}"
+)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(MLX DEFAULT_MSG MLX_LIBRARY MLX_INCLUDE_DIRS)
diff --git a/macos/third_party/mlx/mlx/3rdparty/.clang-format b/macos/third_party/mlx/mlx/3rdparty/.clang-format
new file mode 100644
index 0000000..47a38a9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/3rdparty/.clang-format
@@ -0,0 +1,2 @@
+DisableFormat: true
+SortIncludes: Never
diff --git a/macos/third_party/mlx/mlx/3rdparty/pocketfft.h b/macos/third_party/mlx/mlx/3rdparty/pocketfft.h
new file mode 100644
index 0000000..03a4589
--- /dev/null
+++ b/macos/third_party/mlx/mlx/3rdparty/pocketfft.h
@@ -0,0 +1,3581 @@
+/*
+This file is part of pocketfft.
+
+Copyright (C) 2010-2022 Max-Planck-Society
+Copyright (C) 2019-2020 Peter Bell
+
+For the odd-sized DCT-IV transforms:
+  Copyright (C) 2003, 2007-14 Matteo Frigo
+  Copyright (C) 2003, 2007-14 Massachusetts Institute of Technology
+
+Authors: Martin Reinecke, Peter Bell
+
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+* Redistributions in binary form must reproduce the above copyright notice, this
+  list of conditions and the following disclaimer in the documentation and/or
+  other materials provided with the distribution.
+* Neither the name of the copyright holder nor the names of its contributors may
+  be used to endorse or promote products derived from this software without
+  specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+
+#ifndef POCKETFFT_HDRONLY_H
+#define POCKETFFT_HDRONLY_H
+
+#ifndef __cplusplus
+#error This file is C++ and requires a C++ compiler.
+#endif
+
+#if !(__cplusplus >= 201103L || _MSVC_LANG+0L >= 201103L)
+#error This file requires at least C++11 support.
+#endif
+
+#ifndef POCKETFFT_CACHE_SIZE
+#define POCKETFFT_CACHE_SIZE 0
+#endif
+
+#include <cmath>
+#include <cstdlib>
+#include <stdexcept>
+#include <memory>
+#include <vector>
+#include <complex>
+#include <algorithm>
+#if POCKETFFT_CACHE_SIZE!=0
+#include <array>
+#include <mutex>
+#endif
+
+#ifndef POCKETFFT_NO_MULTITHREADING
+#include <mutex>
+#include <condition_variable>
+#include <thread>
+#include <queue>
+#include <atomic>
+#include <functional>
+#include <new>
+
+#ifdef POCKETFFT_PTHREADS
+#  include <pthread.h>
+#endif
+#endif
+
+#if defined(__GNUC__)
+#define POCKETFFT_NOINLINE __attribute__((noinline))
+#define POCKETFFT_RESTRICT __restrict__
+#elif defined(_MSC_VER)
+#define POCKETFFT_NOINLINE __declspec(noinline)
+#define POCKETFFT_RESTRICT __restrict
+#else
+#define POCKETFFT_NOINLINE
+#define POCKETFFT_RESTRICT
+#endif
+
+namespace pocketfft {
+
+namespace detail {
+using std::size_t;
+using std::ptrdiff_t;
+
+// Always use std:: for <cmath> functions
+template <typename T> T cos(T) = delete;
+template <typename T> T sin(T) = delete;
+template <typename T> T sqrt(T) = delete;
+
+using shape_t = std::vector<size_t>;
+using stride_t = std::vector<ptrdiff_t>;
+
+constexpr bool FORWARD  = true,
+               BACKWARD = false;
+
+// only enable vector support for gcc>=5.0 and clang>=5.0
+#ifndef POCKETFFT_NO_VECTORS
+#define POCKETFFT_NO_VECTORS
+#if defined(__INTEL_COMPILER)
+// do nothing. This is necessary because this compiler also sets __GNUC__.
+#elif defined(__clang__)
+// AppleClang has their own version numbering
+#ifdef __apple_build_version__
+#  if (__clang_major__ > 9) || (__clang_major__ == 9 && __clang_minor__ >= 1)
+#     undef POCKETFFT_NO_VECTORS
+#  endif
+#elif __clang_major__ >= 5
+#  undef POCKETFFT_NO_VECTORS
+#endif
+#elif defined(__GNUC__)
+#if __GNUC__>=5
+#undef POCKETFFT_NO_VECTORS
+#endif
+#endif
+#endif
+
+template<typename T> struct VLEN { static constexpr size_t val=1; };
+
+#ifndef POCKETFFT_NO_VECTORS
+#if (defined(__AVX512F__))
+template<> struct VLEN<float> { static constexpr size_t val=16; };
+template<> struct VLEN<double> { static constexpr size_t val=8; };
+#elif (defined(__AVX__))
+template<> struct VLEN<float> { static constexpr size_t val=8; };
+template<> struct VLEN<double> { static constexpr size_t val=4; };
+#elif (defined(__SSE2__))
+template<> struct VLEN<float> { static constexpr size_t val=4; };
+template<> struct VLEN<double> { static constexpr size_t val=2; };
+#elif (defined(__VSX__))
+template<> struct VLEN<float> { static constexpr size_t val=4; };
+template<> struct VLEN<double> { static constexpr size_t val=2; };
+#elif (defined(__ARM_NEON__) || defined(__ARM_NEON))
+template<> struct VLEN<float> { static constexpr size_t val=4; };
+template<> struct VLEN<double> { static constexpr size_t val=2; };
+#else
+#define POCKETFFT_NO_VECTORS
+#endif
+#endif
+
+// the __MINGW32__ part in the conditional below works around the problem that
+// the standard C++ library on Windows does not provide aligned_alloc() even
+// though the MinGW compiler and MSVC may advertise C++17 compliance.
+#if (__cplusplus >= 201703L) && (!defined(__MINGW32__)) && (!defined(_MSC_VER))
+inline void *aligned_alloc(size_t align, size_t size)
+  {
+  // aligned_alloc() requires that the requested size is a multiple of "align"
+  void *ptr = ::aligned_alloc(align,(size+align-1)&(~(align-1)));
+  if (!ptr) throw std::bad_alloc();
+  return ptr;
+  }
+inline void aligned_dealloc(void *ptr)
+    { free(ptr); }
+#else // portable emulation
+inline void *aligned_alloc(size_t align, size_t size)
+  {
+  align = std::max(align, alignof(max_align_t));
+  void *ptr = malloc(size+align);
+  if (!ptr) throw std::bad_alloc();
+  void *res = reinterpret_cast<void *>
+    ((reinterpret_cast<uintptr_t>(ptr) & ~(uintptr_t(align-1))) + uintptr_t(align));
+  (reinterpret_cast<void**>(res))[-1] = ptr;
+  return res;
+  }
+inline void aligned_dealloc(void *ptr)
+  { if (ptr) free((reinterpret_cast<void**>(ptr))[-1]); }
+#endif
+
+template<typename T> class arr
+  {
+  private:
+    T *p;
+    size_t sz;
+
+#if defined(POCKETFFT_NO_VECTORS)
+    static T *ralloc(size_t num)
+      {
+      if (num==0) return nullptr;
+      void *res = malloc(num*sizeof(T));
+      if (!res) throw std::bad_alloc();
+      return reinterpret_cast<T *>(res);
+      }
+    static void dealloc(T *ptr)
+      { free(ptr); }
+#else
+    static T *ralloc(size_t num)
+      {
+      if (num==0) return nullptr;
+      void *ptr = aligned_alloc(64, num*sizeof(T));
+      return static_cast<T*>(ptr);
+      }
+    static void dealloc(T *ptr)
+      { aligned_dealloc(ptr); }
+#endif
+
+  public:
+    arr() : p(0), sz(0) {}
+    arr(size_t n) : p(ralloc(n)), sz(n) {}
+    arr(arr &&other)
+      : p(other.p), sz(other.sz)
+      { other.p=nullptr; other.sz=0; }
+    ~arr() { dealloc(p); }
+
+    void resize(size_t n)
+      {
+      if (n==sz) return;
+      dealloc(p);
+      p = ralloc(n);
+      sz = n;
+      }
+
+    T &operator[](size_t idx) { return p[idx]; }
+    const T &operator[](size_t idx) const { return p[idx]; }
+
+    T *data() { return p; }
+    const T *data() const { return p; }
+
+    size_t size() const { return sz; }
+  };
+
+template<typename T> struct cmplx {
+  T r, i;
+  cmplx() {}
+  cmplx(T r_, T i_) : r(r_), i(i_) {}
+  void Set(T r_, T i_) { r=r_; i=i_; }
+  void Set(T r_) { r=r_; i=T(0); }
+  cmplx &operator+= (const cmplx &other)
+    { r+=other.r; i+=other.i; return *this; }
+  template<typename T2>cmplx &operator*= (T2 other)
+    { r*=other; i*=other; return *this; }
+  template<typename T2>cmplx &operator*= (const cmplx<T2> &other)
+    {
+    T tmp = r*other.r - i*other.i;
+    i = r*other.i + i*other.r;
+    r = tmp;
+    return *this;
+    }
+  template<typename T2>cmplx &operator+= (const cmplx<T2> &other)
+    { r+=other.r; i+=other.i; return *this; }
+  template<typename T2>cmplx &operator-= (const cmplx<T2> &other)
+    { r-=other.r; i-=other.i; return *this; }
+  template<typename T2> auto operator* (const T2 &other) const
+    -> cmplx<decltype(r*other)>
+    { return {r*other, i*other}; }
+  template<typename T2> auto operator+ (const cmplx<T2> &other) const
+    -> cmplx<decltype(r+other.r)>
+    { return {r+other.r, i+other.i}; }
+  template<typename T2> auto operator- (const cmplx<T2> &other) const
+    -> cmplx<decltype(r+other.r)>
+    { return {r-other.r, i-other.i}; }
+  template<typename T2> auto operator* (const cmplx<T2> &other) const
+    -> cmplx<decltype(r+other.r)>
+    { return {r*other.r-i*other.i, r*other.i + i*other.r}; }
+  template<bool fwd, typename T2> auto special_mul (const cmplx<T2> &other) const
+    -> cmplx<decltype(r+other.r)>
+    {
+    using Tres = cmplx<decltype(r+other.r)>;
+    return fwd ? Tres(r*other.r+i*other.i, i*other.r-r*other.i)
+               : Tres(r*other.r-i*other.i, r*other.i+i*other.r);
+    }
+};
+template<typename T> inline void PM(T &a, T &b, T c, T d)
+  { a=c+d; b=c-d; }
+template<typename T> inline void PMINPLACE(T &a, T &b)
+  { T t = a; a+=b; b=t-b; }
+template<typename T> inline void MPINPLACE(T &a, T &b)
+  { T t = a; a-=b; b=t+b; }
+template<typename T> cmplx<T> conj(const cmplx<T> &a)
+  { return {a.r, -a.i}; }
+template<bool fwd, typename T, typename T2> void special_mul (const cmplx<T> &v1, const cmplx<T2> &v2, cmplx<T> &res)
+  {
+  res = fwd ? cmplx<T>(v1.r*v2.r+v1.i*v2.i, v1.i*v2.r-v1.r*v2.i)
+            : cmplx<T>(v1.r*v2.r-v1.i*v2.i, v1.r*v2.i+v1.i*v2.r);
+  }
+
+template<typename T> void ROT90(cmplx<T> &a)
+  { auto tmp_=a.r; a.r=-a.i; a.i=tmp_; }
+template<bool fwd, typename T> void ROTX90(cmplx<T> &a)
+  { auto tmp_= fwd ? -a.r : a.r; a.r = fwd ? a.i : -a.i; a.i=tmp_; }
+
+//
+// twiddle factor section
+//
+template<typename T> class sincos_2pibyn
+  {
+  private:
+    using Thigh = typename std::conditional<(sizeof(T)>sizeof(double)), T, double>::type;
+    size_t N, mask, shift;
+    arr<cmplx<Thigh>> v1, v2;
+
+    static cmplx<Thigh> calc(size_t x, size_t n, Thigh ang)
+      {
+      x<<=3;
+      if (x<4*n) // first half
+        {
+        if (x<2*n) // first quadrant
+          {
+          if (x<n) return cmplx<Thigh>(std::cos(Thigh(x)*ang), std::sin(Thigh(x)*ang));
+          return cmplx<Thigh>(std::sin(Thigh(2*n-x)*ang), std::cos(Thigh(2*n-x)*ang));
+          }
+        else // second quadrant
+          {
+          x-=2*n;
+          if (x<n) return cmplx<Thigh>(-std::sin(Thigh(x)*ang), std::cos(Thigh(x)*ang));
+          return cmplx<Thigh>(-std::cos(Thigh(2*n-x)*ang), std::sin(Thigh(2*n-x)*ang));
+          }
+        }
+      else
+        {
+        x=8*n-x;
+        if (x<2*n) // third quadrant
+          {
+          if (x<n) return cmplx<Thigh>(std::cos(Thigh(x)*ang), -std::sin(Thigh(x)*ang));
+          return cmplx<Thigh>(std::sin(Thigh(2*n-x)*ang), -std::cos(Thigh(2*n-x)*ang));
+          }
+        else // fourth quadrant
+          {
+          x-=2*n;
+          if (x<n) return cmplx<Thigh>(-std::sin(Thigh(x)*ang), -std::cos(Thigh(x)*ang));
+          return cmplx<Thigh>(-std::cos(Thigh(2*n-x)*ang), -std::sin(Thigh(2*n-x)*ang));
+          }
+        }
+      }
+
+  public:
+    POCKETFFT_NOINLINE sincos_2pibyn(size_t n)
+      : N(n)
+      {
+      constexpr auto pi = 3.141592653589793238462643383279502884197L;
+      Thigh ang = Thigh(0.25L*pi/n);
+      size_t nval = (n+2)/2;
+      shift = 1;
+      while((size_t(1)<<shift)*(size_t(1)<<shift) < nval) ++shift;
+      mask = (size_t(1)<<shift)-1;
+      v1.resize(mask+1);
+      v1[0].Set(Thigh(1), Thigh(0));
+      for (size_t i=1; i<v1.size(); ++i)
+        v1[i]=calc(i,n,ang);
+      v2.resize((nval+mask)/(mask+1));
+      v2[0].Set(Thigh(1), Thigh(0));
+      for (size_t i=1; i<v2.size(); ++i)
+        v2[i]=calc(i*(mask+1),n,ang);
+      }
+
+    cmplx<T> operator[](size_t idx) const
+      {
+      if (2*idx<=N)
+        {
+        auto x1=v1[idx&mask], x2=v2[idx>>shift];
+        return cmplx<T>(T(x1.r*x2.r-x1.i*x2.i), T(x1.r*x2.i+x1.i*x2.r));
+        }
+      idx = N-idx;
+      auto x1=v1[idx&mask], x2=v2[idx>>shift];
+      return cmplx<T>(T(x1.r*x2.r-x1.i*x2.i), -T(x1.r*x2.i+x1.i*x2.r));
+      }
+  };
+
+struct util // hack to avoid duplicate symbols
+  {
+  static POCKETFFT_NOINLINE size_t largest_prime_factor (size_t n)
+    {
+    size_t res=1;
+    while ((n&1)==0)
+      { res=2; n>>=1; }
+    for (size_t x=3; x*x<=n; x+=2)
+      while ((n%x)==0)
+        { res=x; n/=x; }
+    if (n>1) res=n;
+    return res;
+    }
+
+  static POCKETFFT_NOINLINE double cost_guess (size_t n)
+    {
+    constexpr double lfp=1.1; // penalty for non-hardcoded larger factors
+    size_t ni=n;
+    double result=0.;
+    while ((n&1)==0)
+      { result+=2; n>>=1; }
+    for (size_t x=3; x*x<=n; x+=2)
+      while ((n%x)==0)
+        {
+        result+= (x<=5) ? double(x) : lfp*double(x); // penalize larger prime factors
+        n/=x;
+        }
+    if (n>1) result+=(n<=5) ? double(n) : lfp*double(n);
+    return result*double(ni);
+    }
+
+  /* returns the smallest composite of 2, 3, 5, 7 and 11 which is >= n */
+  static POCKETFFT_NOINLINE size_t good_size_cmplx(size_t n)
+    {
+    if (n<=12) return n;
+
+    size_t bestfac=2*n;
+    for (size_t f11=1; f11<bestfac; f11*=11)
+      for (size_t f117=f11; f117<bestfac; f117*=7)
+        for (size_t f1175=f117; f1175<bestfac; f1175*=5)
+          {
+          size_t x=f1175;
+          while (x<n) x*=2;
+          for (;;)
+            {
+            if (x<n)
+              x*=3;
+            else if (x>n)
+              {
+              if (x<bestfac) bestfac=x;
+              if (x&1) break;
+              x>>=1;
+              }
+            else
+              return n;
+            }
+          }
+    return bestfac;
+    }
+
+  /* returns the smallest composite of 2, 3, 5 which is >= n */
+  static POCKETFFT_NOINLINE size_t good_size_real(size_t n)
+    {
+    if (n<=6) return n;
+
+    size_t bestfac=2*n;
+    for (size_t f5=1; f5<bestfac; f5*=5)
+      {
+      size_t x = f5;
+      while (x<n) x *= 2;
+      for (;;)
+        {
+        if (x<n)
+          x*=3;
+        else if (x>n)
+          {
+          if (x<bestfac) bestfac=x;
+          if (x&1) break;
+          x>>=1;
+          }
+        else
+          return n;
+        }
+      }
+    return bestfac;
+    }
+
+  static size_t prod(const shape_t &shape)
+    {
+    size_t res=1;
+    for (auto sz: shape)
+      res*=sz;
+    return res;
+    }
+
+  static POCKETFFT_NOINLINE void sanity_check(const shape_t &shape,
+    const stride_t &stride_in, const stride_t &stride_out, bool inplace)
+    {
+    auto ndim = shape.size();
+    if (ndim<1) throw std::runtime_error("ndim must be >= 1");
+    if ((stride_in.size()!=ndim) || (stride_out.size()!=ndim))
+      throw std::runtime_error("stride dimension mismatch");
+    if (inplace && (stride_in!=stride_out))
+      throw std::runtime_error("stride mismatch");
+    }
+
+  static POCKETFFT_NOINLINE void sanity_check(const shape_t &shape,
+    const stride_t &stride_in, const stride_t &stride_out, bool inplace,
+    const shape_t &axes)
+    {
+    sanity_check(shape, stride_in, stride_out, inplace);
+    auto ndim = shape.size();
+    shape_t tmp(ndim,0);
+    for (auto ax : axes)
+      {
+      if (ax>=ndim) throw std::invalid_argument("bad axis number");
+      if (++tmp[ax]>1) throw std::invalid_argument("axis specified repeatedly");
+      }
+    }
+
+  static POCKETFFT_NOINLINE void sanity_check(const shape_t &shape,
+    const stride_t &stride_in, const stride_t &stride_out, bool inplace,
+    size_t axis)
+    {
+    sanity_check(shape, stride_in, stride_out, inplace);
+    if (axis>=shape.size()) throw std::invalid_argument("bad axis number");
+    }
+
+#ifdef POCKETFFT_NO_MULTITHREADING
+  static size_t thread_count (size_t /*nthreads*/, const shape_t &/*shape*/,
+    size_t /*axis*/, size_t /*vlen*/)
+    { return 1; }
+#else
+  static size_t thread_count (size_t nthreads, const shape_t &shape,
+    size_t axis, size_t vlen)
+    {
+    if (nthreads==1) return 1;
+    size_t size = prod(shape);
+    size_t parallel = size / (shape[axis] * vlen);
+    if (shape[axis] < 1000)
+      parallel /= 4;
+    size_t max_threads = nthreads == 0 ?
+      std::thread::hardware_concurrency() : nthreads;
+    return std::max(size_t(1), std::min(parallel, max_threads));
+    }
+#endif
+  };
+
+namespace threading {
+
+#ifdef POCKETFFT_NO_MULTITHREADING
+
+constexpr inline size_t thread_id() { return 0; }
+constexpr inline size_t num_threads() { return 1; }
+
+template <typename Func>
+void thread_map(size_t /* nthreads */, Func f)
+  { f(); }
+
+#else
+
+inline size_t &thread_id()
+  {
+  static thread_local size_t thread_id_=0;
+  return thread_id_;
+  }
+inline size_t &num_threads()
+  {
+  static thread_local size_t num_threads_=1;
+  return num_threads_;
+  }
+static const size_t max_threads = std::max(1u, std::thread::hardware_concurrency());
+
+class latch
+  {
+    std::atomic<size_t> num_left_;
+    std::mutex mut_;
+    std::condition_variable completed_;
+    using lock_t = std::unique_lock<std::mutex>;
+
+  public:
+    latch(size_t n): num_left_(n) {}
+
+    void count_down()
+      {
+      lock_t lock(mut_);
+      if (--num_left_)
+        return;
+      completed_.notify_all();
+      }
+
+    void wait()
+      {
+      lock_t lock(mut_);
+      completed_.wait(lock, [this]{ return is_ready(); });
+      }
+    bool is_ready() { return num_left_ == 0; }
+  };
+
+template <typename T> class concurrent_queue
+  {
+    std::queue<T> q_;
+    std::mutex mut_;
+    std::atomic<size_t> size_;
+    using lock_t = std::lock_guard<std::mutex>;
+
+  public:
+
+    void push(T val)
+      {
+      lock_t lock(mut_);
+      ++size_;
+      q_.push(std::move(val));
+      }
+
+    bool try_pop(T &val)
+      {
+      if (size_ == 0) return false;
+      lock_t lock(mut_);
+      // Queue might have been emptied while we acquired the lock
+      if (q_.empty()) return false;
+
+      val = std::move(q_.front());
+      --size_;
+      q_.pop();
+      return true;
+      }
+
+    bool empty() const { return size_==0; }
+  };
+
+// C++ allocator with support for over-aligned types
+template <typename T> struct aligned_allocator
+  {
+  using value_type = T;
+  template <class U>
+  aligned_allocator(const aligned_allocator<U>&) {}
+  aligned_allocator() = default;
+
+  T *allocate(size_t n)
+    {
+    void* mem = aligned_alloc(alignof(T), n*sizeof(T));
+    return static_cast<T*>(mem);
+    }
+
+  void deallocate(T *p, size_t /*n*/)
+    { aligned_dealloc(p); }
+  };
+
+class thread_pool
+  {
+    // A reasonable guess, probably close enough for most hardware
+    static constexpr size_t cache_line_size = 64;
+    struct alignas(cache_line_size) worker
+      {
+      std::thread thread;
+      std::condition_variable work_ready;
+      std::mutex mut;
+      std::atomic_flag busy_flag = ATOMIC_FLAG_INIT;
+      std::function<void()> work;
+
+      void worker_main(
+        std::atomic<bool> &shutdown_flag,
+        std::atomic<size_t> &unscheduled_tasks,
+        concurrent_queue<std::function<void()>> &overflow_work)
+        {
+        using lock_t = std::unique_lock<std::mutex>;
+        bool expect_work = true;
+        while (!shutdown_flag || expect_work)
+          {
+          std::function<void()> local_work;
+          if (expect_work || unscheduled_tasks == 0)
+            {
+            lock_t lock(mut);
+            // Wait until there is work to be executed
+            work_ready.wait(lock, [&]{ return (work || shutdown_flag); });
+            local_work.swap(work);
+            expect_work = false;
+            }
+
+          bool marked_busy = false;
+          if (local_work)
+            {
+            marked_busy = true;
+            local_work();
+            }
+
+          if (!overflow_work.empty())
+            {
+            if (!marked_busy && busy_flag.test_and_set())
+              {
+              expect_work = true;
+              continue;
+              }
+            marked_busy = true;
+
+            while (overflow_work.try_pop(local_work))
+              {
+              --unscheduled_tasks;
+              local_work();
+              }
+            }
+
+          if (marked_busy) busy_flag.clear();
+          }
+        }
+      };
+
+    concurrent_queue<std::function<void()>> overflow_work_;
+    std::mutex mut_;
+    std::vector<worker, aligned_allocator<worker>> workers_;
+    std::atomic<bool> shutdown_;
+    std::atomic<size_t> unscheduled_tasks_;
+    using lock_t = std::lock_guard<std::mutex>;
+
+    void create_threads()
+      {
+      lock_t lock(mut_);
+      size_t nthreads=workers_.size();
+      for (size_t i=0; i<nthreads; ++i)
+        {
+        try
+          {
+          auto *worker = &workers_[i];
+          worker->busy_flag.clear();
+          worker->work = nullptr;
+          worker->thread = std::thread([worker, this]
+            {
+            worker->worker_main(shutdown_, unscheduled_tasks_, overflow_work_);
+            });
+          }
+        catch (...)
+          {
+          shutdown_locked();
+          throw;
+          }
+        }
+      }
+
+    void shutdown_locked()
+      {
+      shutdown_ = true;
+      for (auto &worker : workers_)
+        worker.work_ready.notify_all();
+
+      for (auto &worker : workers_)
+        if (worker.thread.joinable())
+          worker.thread.join();
+      }
+
+  public:
+    explicit thread_pool(size_t nthreads):
+      workers_(nthreads)
+      { create_threads(); }
+
+    thread_pool(): thread_pool(max_threads) {}
+
+    ~thread_pool() { shutdown(); }
+
+    void submit(std::function<void()> work)
+      {
+      lock_t lock(mut_);
+      if (shutdown_)
+        throw std::runtime_error("Work item submitted after shutdown");
+
+      ++unscheduled_tasks_;
+
+      // First check for any idle workers and wake those
+      for (auto &worker : workers_)
+        if (!worker.busy_flag.test_and_set())
+          {
+          --unscheduled_tasks_;
+          {
+          lock_t lock(worker.mut);
+          worker.work = std::move(work);
+          }
+          worker.work_ready.notify_one();
+          return;
+          }
+
+      // If no workers were idle, push onto the overflow queue for later
+      overflow_work_.push(std::move(work));
+      }
+
+    void shutdown()
+      {
+      lock_t lock(mut_);
+      shutdown_locked();
+      }
+
+    void restart()
+      {
+      shutdown_ = false;
+      create_threads();
+      }
+  };
+
+inline thread_pool & get_pool()
+  {
+  static thread_pool pool;
+#ifdef POCKETFFT_PTHREADS
+  static std::once_flag f;
+  std::call_once(f,
+    []{
+    pthread_atfork(
+      +[]{ get_pool().shutdown(); },  // prepare
+      +[]{ get_pool().restart(); },   // parent
+      +[]{ get_pool().restart(); }    // child
+      );
+    });
+#endif
+
+  return pool;
+  }
+
+/** Map a function f over nthreads */
+template <typename Func>
+void thread_map(size_t nthreads, Func f)
+  {
+  if (nthreads == 0)
+    nthreads = max_threads;
+
+  if (nthreads == 1)
+    { f(); return; }
+
+  auto & pool = get_pool();
+  latch counter(nthreads);
+  std::exception_ptr ex;
+  std::mutex ex_mut;
+  for (size_t i=0; i<nthreads; ++i)
+    {
+    pool.submit(
+      [&f, &counter, &ex, &ex_mut, i, nthreads] {
+      thread_id() = i;
+      num_threads() = nthreads;
+      try { f(); }
+      catch (...)
+        {
+        std::lock_guard<std::mutex> lock(ex_mut);
+        ex = std::current_exception();
+        }
+      counter.count_down();
+      });
+    }
+  counter.wait();
+  if (ex)
+    std::rethrow_exception(ex);
+  }
+
+#endif
+
+}
+
+//
+// complex FFTPACK transforms
+//
+
+template<typename T0> class cfftp
+  {
+  private:
+    struct fctdata
+      {
+      size_t fct;
+      cmplx<T0> *tw, *tws;
+      };
+
+    size_t length;
+    arr<cmplx<T0>> mem;
+    std::vector<fctdata> fact;
+
+    void add_factor(size_t factor)
+      { fact.push_back({factor, nullptr, nullptr}); }
+
+template<bool fwd, typename T> void pass2 (size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa) const
+  {
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+2*c)]; };
+  auto WA = [wa, ido](size_t x, size_t i)
+    { return wa[i-1+x*(ido-1)]; };
+
+  if (ido==1)
+    for (size_t k=0; k<l1; ++k)
+      {
+      CH(0,k,0) = CC(0,0,k)+CC(0,1,k);
+      CH(0,k,1) = CC(0,0,k)-CC(0,1,k);
+      }
+  else
+    for (size_t k=0; k<l1; ++k)
+      {
+      CH(0,k,0) = CC(0,0,k)+CC(0,1,k);
+      CH(0,k,1) = CC(0,0,k)-CC(0,1,k);
+      for (size_t i=1; i<ido; ++i)
+        {
+        CH(i,k,0) = CC(i,0,k)+CC(i,1,k);
+        special_mul<fwd>(CC(i,0,k)-CC(i,1,k),WA(0,i),CH(i,k,1));
+        }
+      }
+  }
+
+#define POCKETFFT_PREP3(idx) \
+        T t0 = CC(idx,0,k), t1, t2; \
+        PM (t1,t2,CC(idx,1,k),CC(idx,2,k)); \
+        CH(idx,k,0)=t0+t1;
+#define POCKETFFT_PARTSTEP3a(u1,u2,twr,twi) \
+        { \
+        T ca=t0+t1*twr; \
+        T cb{-t2.i*twi, t2.r*twi}; \
+        PM(CH(0,k,u1),CH(0,k,u2),ca,cb) ;\
+        }
+#define POCKETFFT_PARTSTEP3b(u1,u2,twr,twi) \
+        { \
+        T ca=t0+t1*twr; \
+        T cb{-t2.i*twi, t2.r*twi}; \
+        special_mul<fwd>(ca+cb,WA(u1-1,i),CH(i,k,u1)); \
+        special_mul<fwd>(ca-cb,WA(u2-1,i),CH(i,k,u2)); \
+        }
+template<bool fwd, typename T> void pass3 (size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 tw1r=-0.5,
+               tw1i= (fwd ? -1: 1) * T0(0.8660254037844386467637231707529362L);
+
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+3*c)]; };
+  auto WA = [wa, ido](size_t x, size_t i)
+    { return wa[i-1+x*(ido-1)]; };
+
+  if (ido==1)
+    for (size_t k=0; k<l1; ++k)
+      {
+      POCKETFFT_PREP3(0)
+      POCKETFFT_PARTSTEP3a(1,2,tw1r,tw1i)
+      }
+  else
+    for (size_t k=0; k<l1; ++k)
+      {
+      {
+      POCKETFFT_PREP3(0)
+      POCKETFFT_PARTSTEP3a(1,2,tw1r,tw1i)
+      }
+      for (size_t i=1; i<ido; ++i)
+        {
+        POCKETFFT_PREP3(i)
+        POCKETFFT_PARTSTEP3b(1,2,tw1r,tw1i)
+        }
+      }
+  }
+
+#undef POCKETFFT_PARTSTEP3b
+#undef POCKETFFT_PARTSTEP3a
+#undef POCKETFFT_PREP3
+
+template<bool fwd, typename T> void pass4 (size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa) const
+  {
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+4*c)]; };
+  auto WA = [wa, ido](size_t x, size_t i)
+    { return wa[i-1+x*(ido-1)]; };
+
+  if (ido==1)
+    for (size_t k=0; k<l1; ++k)
+      {
+      T t1, t2, t3, t4;
+      PM(t2,t1,CC(0,0,k),CC(0,2,k));
+      PM(t3,t4,CC(0,1,k),CC(0,3,k));
+      ROTX90<fwd>(t4);
+      PM(CH(0,k,0),CH(0,k,2),t2,t3);
+      PM(CH(0,k,1),CH(0,k,3),t1,t4);
+      }
+  else
+    for (size_t k=0; k<l1; ++k)
+      {
+      {
+      T t1, t2, t3, t4;
+      PM(t2,t1,CC(0,0,k),CC(0,2,k));
+      PM(t3,t4,CC(0,1,k),CC(0,3,k));
+      ROTX90<fwd>(t4);
+      PM(CH(0,k,0),CH(0,k,2),t2,t3);
+      PM(CH(0,k,1),CH(0,k,3),t1,t4);
+      }
+      for (size_t i=1; i<ido; ++i)
+        {
+        T t1, t2, t3, t4;
+        T cc0=CC(i,0,k), cc1=CC(i,1,k),cc2=CC(i,2,k),cc3=CC(i,3,k);
+        PM(t2,t1,cc0,cc2);
+        PM(t3,t4,cc1,cc3);
+        ROTX90<fwd>(t4);
+        CH(i,k,0) = t2+t3;
+        special_mul<fwd>(t1+t4,WA(0,i),CH(i,k,1));
+        special_mul<fwd>(t2-t3,WA(1,i),CH(i,k,2));
+        special_mul<fwd>(t1-t4,WA(2,i),CH(i,k,3));
+        }
+      }
+  }
+
+#define POCKETFFT_PREP5(idx) \
+        T t0 = CC(idx,0,k), t1, t2, t3, t4; \
+        PM (t1,t4,CC(idx,1,k),CC(idx,4,k)); \
+        PM (t2,t3,CC(idx,2,k),CC(idx,3,k)); \
+        CH(idx,k,0).r=t0.r+t1.r+t2.r; \
+        CH(idx,k,0).i=t0.i+t1.i+t2.i;
+
+#define POCKETFFT_PARTSTEP5a(u1,u2,twar,twbr,twai,twbi) \
+        { \
+        T ca,cb; \
+        ca.r=t0.r+twar*t1.r+twbr*t2.r; \
+        ca.i=t0.i+twar*t1.i+twbr*t2.i; \
+        cb.i=twai*t4.r twbi*t3.r; \
+        cb.r=-(twai*t4.i twbi*t3.i); \
+        PM(CH(0,k,u1),CH(0,k,u2),ca,cb); \
+        }
+
+#define POCKETFFT_PARTSTEP5b(u1,u2,twar,twbr,twai,twbi) \
+        { \
+        T ca,cb,da,db; \
+        ca.r=t0.r+twar*t1.r+twbr*t2.r; \
+        ca.i=t0.i+twar*t1.i+twbr*t2.i; \
+        cb.i=twai*t4.r twbi*t3.r; \
+        cb.r=-(twai*t4.i twbi*t3.i); \
+        special_mul<fwd>(ca+cb,WA(u1-1,i),CH(i,k,u1)); \
+        special_mul<fwd>(ca-cb,WA(u2-1,i),CH(i,k,u2)); \
+        }
+template<bool fwd, typename T> void pass5 (size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 tw1r= T0(0.3090169943749474241022934171828191L),
+               tw1i= (fwd ? -1: 1) * T0(0.9510565162951535721164393333793821L),
+               tw2r= T0(-0.8090169943749474241022934171828191L),
+               tw2i= (fwd ? -1: 1) * T0(0.5877852522924731291687059546390728L);
+
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+5*c)]; };
+  auto WA = [wa, ido](size_t x, size_t i)
+    { return wa[i-1+x*(ido-1)]; };
+
+  if (ido==1)
+    for (size_t k=0; k<l1; ++k)
+      {
+      POCKETFFT_PREP5(0)
+      POCKETFFT_PARTSTEP5a(1,4,tw1r,tw2r,+tw1i,+tw2i)
+      POCKETFFT_PARTSTEP5a(2,3,tw2r,tw1r,+tw2i,-tw1i)
+      }
+  else
+    for (size_t k=0; k<l1; ++k)
+      {
+      {
+      POCKETFFT_PREP5(0)
+      POCKETFFT_PARTSTEP5a(1,4,tw1r,tw2r,+tw1i,+tw2i)
+      POCKETFFT_PARTSTEP5a(2,3,tw2r,tw1r,+tw2i,-tw1i)
+      }
+      for (size_t i=1; i<ido; ++i)
+        {
+        POCKETFFT_PREP5(i)
+        POCKETFFT_PARTSTEP5b(1,4,tw1r,tw2r,+tw1i,+tw2i)
+        POCKETFFT_PARTSTEP5b(2,3,tw2r,tw1r,+tw2i,-tw1i)
+        }
+      }
+  }
+
+#undef POCKETFFT_PARTSTEP5b
+#undef POCKETFFT_PARTSTEP5a
+#undef POCKETFFT_PREP5
+
+#define POCKETFFT_PREP7(idx) \
+        T t1 = CC(idx,0,k), t2, t3, t4, t5, t6, t7; \
+        PM (t2,t7,CC(idx,1,k),CC(idx,6,k)); \
+        PM (t3,t6,CC(idx,2,k),CC(idx,5,k)); \
+        PM (t4,t5,CC(idx,3,k),CC(idx,4,k)); \
+        CH(idx,k,0).r=t1.r+t2.r+t3.r+t4.r; \
+        CH(idx,k,0).i=t1.i+t2.i+t3.i+t4.i;
+
+#define POCKETFFT_PARTSTEP7a0(u1,u2,x1,x2,x3,y1,y2,y3,out1,out2) \
+        { \
+        T ca,cb; \
+        ca.r=t1.r+x1*t2.r+x2*t3.r+x3*t4.r; \
+        ca.i=t1.i+x1*t2.i+x2*t3.i+x3*t4.i; \
+        cb.i=y1*t7.r y2*t6.r y3*t5.r; \
+        cb.r=-(y1*t7.i y2*t6.i y3*t5.i); \
+        PM(out1,out2,ca,cb); \
+        }
+#define POCKETFFT_PARTSTEP7a(u1,u2,x1,x2,x3,y1,y2,y3) \
+        POCKETFFT_PARTSTEP7a0(u1,u2,x1,x2,x3,y1,y2,y3,CH(0,k,u1),CH(0,k,u2))
+#define POCKETFFT_PARTSTEP7(u1,u2,x1,x2,x3,y1,y2,y3) \
+        { \
+        T da,db; \
+        POCKETFFT_PARTSTEP7a0(u1,u2,x1,x2,x3,y1,y2,y3,da,db) \
+        special_mul<fwd>(da,WA(u1-1,i),CH(i,k,u1)); \
+        special_mul<fwd>(db,WA(u2-1,i),CH(i,k,u2)); \
+        }
+
+template<bool fwd, typename T> void pass7(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 tw1r= T0(0.6234898018587335305250048840042398L),
+               tw1i= (fwd ? -1 : 1) * T0(0.7818314824680298087084445266740578L),
+               tw2r= T0(-0.2225209339563144042889025644967948L),
+               tw2i= (fwd ? -1 : 1) * T0(0.9749279121818236070181316829939312L),
+               tw3r= T0(-0.9009688679024191262361023195074451L),
+               tw3i= (fwd ? -1 : 1) * T0(0.433883739117558120475768332848359L);
+
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+7*c)]; };
+  auto WA = [wa, ido](size_t x, size_t i)
+    { return wa[i-1+x*(ido-1)]; };
+
+  if (ido==1)
+    for (size_t k=0; k<l1; ++k)
+      {
+      POCKETFFT_PREP7(0)
+      POCKETFFT_PARTSTEP7a(1,6,tw1r,tw2r,tw3r,+tw1i,+tw2i,+tw3i)
+      POCKETFFT_PARTSTEP7a(2,5,tw2r,tw3r,tw1r,+tw2i,-tw3i,-tw1i)
+      POCKETFFT_PARTSTEP7a(3,4,tw3r,tw1r,tw2r,+tw3i,-tw1i,+tw2i)
+      }
+  else
+    for (size_t k=0; k<l1; ++k)
+      {
+      {
+      POCKETFFT_PREP7(0)
+      POCKETFFT_PARTSTEP7a(1,6,tw1r,tw2r,tw3r,+tw1i,+tw2i,+tw3i)
+      POCKETFFT_PARTSTEP7a(2,5,tw2r,tw3r,tw1r,+tw2i,-tw3i,-tw1i)
+      POCKETFFT_PARTSTEP7a(3,4,tw3r,tw1r,tw2r,+tw3i,-tw1i,+tw2i)
+      }
+      for (size_t i=1; i<ido; ++i)
+        {
+        POCKETFFT_PREP7(i)
+        POCKETFFT_PARTSTEP7(1,6,tw1r,tw2r,tw3r,+tw1i,+tw2i,+tw3i)
+        POCKETFFT_PARTSTEP7(2,5,tw2r,tw3r,tw1r,+tw2i,-tw3i,-tw1i)
+        POCKETFFT_PARTSTEP7(3,4,tw3r,tw1r,tw2r,+tw3i,-tw1i,+tw2i)
+        }
+      }
+  }
+
+#undef POCKETFFT_PARTSTEP7
+#undef POCKETFFT_PARTSTEP7a0
+#undef POCKETFFT_PARTSTEP7a
+#undef POCKETFFT_PREP7
+
+template <bool fwd, typename T> void ROTX45(T &a) const
+  {
+  constexpr T0 hsqt2=T0(0.707106781186547524400844362104849L);
+  if (fwd)
+    { auto tmp_=a.r; a.r=hsqt2*(a.r+a.i); a.i=hsqt2*(a.i-tmp_); }
+  else
+    { auto tmp_=a.r; a.r=hsqt2*(a.r-a.i); a.i=hsqt2*(a.i+tmp_); }
+  }
+template <bool fwd, typename T> void ROTX135(T &a) const
+  {
+  constexpr T0 hsqt2=T0(0.707106781186547524400844362104849L);
+  if (fwd)
+    { auto tmp_=a.r; a.r=hsqt2*(a.i-a.r); a.i=hsqt2*(-tmp_-a.i); }
+  else
+    { auto tmp_=a.r; a.r=hsqt2*(-a.r-a.i); a.i=hsqt2*(tmp_-a.i); }
+  }
+
+template<bool fwd, typename T> void pass8 (size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa) const
+  {
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+8*c)]; };
+  auto WA = [wa, ido](size_t x, size_t i)
+    { return wa[i-1+x*(ido-1)]; };
+
+  if (ido==1)
+    for (size_t k=0; k<l1; ++k)
+      {
+      T a0, a1, a2, a3, a4, a5, a6, a7;
+      PM(a1,a5,CC(0,1,k),CC(0,5,k));
+      PM(a3,a7,CC(0,3,k),CC(0,7,k));
+      PMINPLACE(a1,a3);
+      ROTX90<fwd>(a3);
+
+      ROTX90<fwd>(a7);
+      PMINPLACE(a5,a7);
+      ROTX45<fwd>(a5);
+      ROTX135<fwd>(a7);
+
+      PM(a0,a4,CC(0,0,k),CC(0,4,k));
+      PM(a2,a6,CC(0,2,k),CC(0,6,k));
+      PM(CH(0,k,0),CH(0,k,4),a0+a2,a1);
+      PM(CH(0,k,2),CH(0,k,6),a0-a2,a3);
+      ROTX90<fwd>(a6);
+      PM(CH(0,k,1),CH(0,k,5),a4+a6,a5);
+      PM(CH(0,k,3),CH(0,k,7),a4-a6,a7);
+      }
+  else
+    for (size_t k=0; k<l1; ++k)
+      {
+      {
+      T a0, a1, a2, a3, a4, a5, a6, a7;
+      PM(a1,a5,CC(0,1,k),CC(0,5,k));
+      PM(a3,a7,CC(0,3,k),CC(0,7,k));
+      PMINPLACE(a1,a3);
+      ROTX90<fwd>(a3);
+
+      ROTX90<fwd>(a7);
+      PMINPLACE(a5,a7);
+      ROTX45<fwd>(a5);
+      ROTX135<fwd>(a7);
+
+      PM(a0,a4,CC(0,0,k),CC(0,4,k));
+      PM(a2,a6,CC(0,2,k),CC(0,6,k));
+      PM(CH(0,k,0),CH(0,k,4),a0+a2,a1);
+      PM(CH(0,k,2),CH(0,k,6),a0-a2,a3);
+      ROTX90<fwd>(a6);
+      PM(CH(0,k,1),CH(0,k,5),a4+a6,a5);
+      PM(CH(0,k,3),CH(0,k,7),a4-a6,a7);
+      }
+      for (size_t i=1; i<ido; ++i)
+        {
+        T a0, a1, a2, a3, a4, a5, a6, a7;
+        PM(a1,a5,CC(i,1,k),CC(i,5,k));
+        PM(a3,a7,CC(i,3,k),CC(i,7,k));
+        ROTX90<fwd>(a7);
+        PMINPLACE(a1,a3);
+        ROTX90<fwd>(a3);
+        PMINPLACE(a5,a7);
+        ROTX45<fwd>(a5);
+        ROTX135<fwd>(a7);
+        PM(a0,a4,CC(i,0,k),CC(i,4,k));
+        PM(a2,a6,CC(i,2,k),CC(i,6,k));
+        PMINPLACE(a0,a2);
+        CH(i,k,0) = a0+a1;
+        special_mul<fwd>(a0-a1,WA(3,i),CH(i,k,4));
+        special_mul<fwd>(a2+a3,WA(1,i),CH(i,k,2));
+        special_mul<fwd>(a2-a3,WA(5,i),CH(i,k,6));
+        ROTX90<fwd>(a6);
+        PMINPLACE(a4,a6);
+        special_mul<fwd>(a4+a5,WA(0,i),CH(i,k,1));
+        special_mul<fwd>(a4-a5,WA(4,i),CH(i,k,5));
+        special_mul<fwd>(a6+a7,WA(2,i),CH(i,k,3));
+        special_mul<fwd>(a6-a7,WA(6,i),CH(i,k,7));
+        }
+      }
+   }
+
+
+#define POCKETFFT_PREP11(idx) \
+        T t1 = CC(idx,0,k), t2, t3, t4, t5, t6, t7, t8, t9, t10, t11; \
+        PM (t2,t11,CC(idx,1,k),CC(idx,10,k)); \
+        PM (t3,t10,CC(idx,2,k),CC(idx, 9,k)); \
+        PM (t4,t9 ,CC(idx,3,k),CC(idx, 8,k)); \
+        PM (t5,t8 ,CC(idx,4,k),CC(idx, 7,k)); \
+        PM (t6,t7 ,CC(idx,5,k),CC(idx, 6,k)); \
+        CH(idx,k,0).r=t1.r+t2.r+t3.r+t4.r+t5.r+t6.r; \
+        CH(idx,k,0).i=t1.i+t2.i+t3.i+t4.i+t5.i+t6.i;
+
+#define POCKETFFT_PARTSTEP11a0(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5,out1,out2) \
+        { \
+        T ca = t1 + t2*x1 + t3*x2 + t4*x3 + t5*x4 +t6*x5, \
+          cb; \
+        cb.i=y1*t11.r y2*t10.r y3*t9.r y4*t8.r y5*t7.r; \
+        cb.r=-(y1*t11.i y2*t10.i y3*t9.i y4*t8.i y5*t7.i ); \
+        PM(out1,out2,ca,cb); \
+        }
+#define POCKETFFT_PARTSTEP11a(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5) \
+        POCKETFFT_PARTSTEP11a0(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5,CH(0,k,u1),CH(0,k,u2))
+#define POCKETFFT_PARTSTEP11(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5) \
+        { \
+        T da,db; \
+        POCKETFFT_PARTSTEP11a0(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5,da,db) \
+        special_mul<fwd>(da,WA(u1-1,i),CH(i,k,u1)); \
+        special_mul<fwd>(db,WA(u2-1,i),CH(i,k,u2)); \
+        }
+
+template<bool fwd, typename T> void pass11 (size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 tw1r= T0(0.8412535328311811688618116489193677L),
+               tw1i= (fwd ? -1 : 1) * T0(0.5406408174555975821076359543186917L),
+               tw2r= T0(0.4154150130018864255292741492296232L),
+               tw2i= (fwd ? -1 : 1) * T0(0.9096319953545183714117153830790285L),
+               tw3r= T0(-0.1423148382732851404437926686163697L),
+               tw3i= (fwd ? -1 : 1) * T0(0.9898214418809327323760920377767188L),
+               tw4r= T0(-0.6548607339452850640569250724662936L),
+               tw4i= (fwd ? -1 : 1) * T0(0.7557495743542582837740358439723444L),
+               tw5r= T0(-0.9594929736144973898903680570663277L),
+               tw5i= (fwd ? -1 : 1) * T0(0.2817325568414296977114179153466169L);
+
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+11*c)]; };
+  auto WA = [wa, ido](size_t x, size_t i)
+    { return wa[i-1+x*(ido-1)]; };
+
+  if (ido==1)
+    for (size_t k=0; k<l1; ++k)
+      {
+      POCKETFFT_PREP11(0)
+      POCKETFFT_PARTSTEP11a(1,10,tw1r,tw2r,tw3r,tw4r,tw5r,+tw1i,+tw2i,+tw3i,+tw4i,+tw5i)
+      POCKETFFT_PARTSTEP11a(2, 9,tw2r,tw4r,tw5r,tw3r,tw1r,+tw2i,+tw4i,-tw5i,-tw3i,-tw1i)
+      POCKETFFT_PARTSTEP11a(3, 8,tw3r,tw5r,tw2r,tw1r,tw4r,+tw3i,-tw5i,-tw2i,+tw1i,+tw4i)
+      POCKETFFT_PARTSTEP11a(4, 7,tw4r,tw3r,tw1r,tw5r,tw2r,+tw4i,-tw3i,+tw1i,+tw5i,-tw2i)
+      POCKETFFT_PARTSTEP11a(5, 6,tw5r,tw1r,tw4r,tw2r,tw3r,+tw5i,-tw1i,+tw4i,-tw2i,+tw3i)
+      }
+  else
+    for (size_t k=0; k<l1; ++k)
+      {
+      {
+      POCKETFFT_PREP11(0)
+      POCKETFFT_PARTSTEP11a(1,10,tw1r,tw2r,tw3r,tw4r,tw5r,+tw1i,+tw2i,+tw3i,+tw4i,+tw5i)
+      POCKETFFT_PARTSTEP11a(2, 9,tw2r,tw4r,tw5r,tw3r,tw1r,+tw2i,+tw4i,-tw5i,-tw3i,-tw1i)
+      POCKETFFT_PARTSTEP11a(3, 8,tw3r,tw5r,tw2r,tw1r,tw4r,+tw3i,-tw5i,-tw2i,+tw1i,+tw4i)
+      POCKETFFT_PARTSTEP11a(4, 7,tw4r,tw3r,tw1r,tw5r,tw2r,+tw4i,-tw3i,+tw1i,+tw5i,-tw2i)
+      POCKETFFT_PARTSTEP11a(5, 6,tw5r,tw1r,tw4r,tw2r,tw3r,+tw5i,-tw1i,+tw4i,-tw2i,+tw3i)
+      }
+      for (size_t i=1; i<ido; ++i)
+        {
+        POCKETFFT_PREP11(i)
+        POCKETFFT_PARTSTEP11(1,10,tw1r,tw2r,tw3r,tw4r,tw5r,+tw1i,+tw2i,+tw3i,+tw4i,+tw5i)
+        POCKETFFT_PARTSTEP11(2, 9,tw2r,tw4r,tw5r,tw3r,tw1r,+tw2i,+tw4i,-tw5i,-tw3i,-tw1i)
+        POCKETFFT_PARTSTEP11(3, 8,tw3r,tw5r,tw2r,tw1r,tw4r,+tw3i,-tw5i,-tw2i,+tw1i,+tw4i)
+        POCKETFFT_PARTSTEP11(4, 7,tw4r,tw3r,tw1r,tw5r,tw2r,+tw4i,-tw3i,+tw1i,+tw5i,-tw2i)
+        POCKETFFT_PARTSTEP11(5, 6,tw5r,tw1r,tw4r,tw2r,tw3r,+tw5i,-tw1i,+tw4i,-tw2i,+tw3i)
+        }
+      }
+  }
+
+#undef POCKETFFT_PARTSTEP11
+#undef POCKETFFT_PARTSTEP11a0
+#undef POCKETFFT_PARTSTEP11a
+#undef POCKETFFT_PREP11
+
+template<bool fwd, typename T> void passg (size_t ido, size_t ip,
+  size_t l1, T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const cmplx<T0> * POCKETFFT_RESTRICT wa,
+  const cmplx<T0> * POCKETFFT_RESTRICT csarr) const
+  {
+  const size_t cdim=ip;
+  size_t ipph = (ip+1)/2;
+  size_t idl1 = ido*l1;
+
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto CC = [cc,ido,cdim](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+cdim*c)]; };
+  auto CX = [cc, ido, l1](size_t a, size_t b, size_t c) -> T&
+    { return cc[a+ido*(b+l1*c)]; };
+  auto CX2 = [cc, idl1](size_t a, size_t b) -> T&
+    { return cc[a+idl1*b]; };
+  auto CH2 = [ch, idl1](size_t a, size_t b) -> const T&
+    { return ch[a+idl1*b]; };
+
+  arr<cmplx<T0>> wal(ip);
+  wal[0] = cmplx<T0>(1., 0.);
+  for (size_t i=1; i<ip; ++i)
+    wal[i]=cmplx<T0>(csarr[i].r,fwd ? -csarr[i].i : csarr[i].i);
+
+  for (size_t k=0; k<l1; ++k)
+    for (size_t i=0; i<ido; ++i)
+      CH(i,k,0) = CC(i,0,k);
+  for (size_t j=1, jc=ip-1; j<ipph; ++j, --jc)
+    for (size_t k=0; k<l1; ++k)
+      for (size_t i=0; i<ido; ++i)
+        PM(CH(i,k,j),CH(i,k,jc),CC(i,j,k),CC(i,jc,k));
+  for (size_t k=0; k<l1; ++k)
+    for (size_t i=0; i<ido; ++i)
+      {
+      T tmp = CH(i,k,0);
+      for (size_t j=1; j<ipph; ++j)
+        tmp+=CH(i,k,j);
+      CX(i,k,0) = tmp;
+      }
+  for (size_t l=1, lc=ip-1; l<ipph; ++l, --lc)
+    {
+    // j=0
+    for (size_t ik=0; ik<idl1; ++ik)
+      {
+      CX2(ik,l).r = CH2(ik,0).r+wal[l].r*CH2(ik,1).r+wal[2*l].r*CH2(ik,2).r;
+      CX2(ik,l).i = CH2(ik,0).i+wal[l].r*CH2(ik,1).i+wal[2*l].r*CH2(ik,2).i;
+      CX2(ik,lc).r=-wal[l].i*CH2(ik,ip-1).i-wal[2*l].i*CH2(ik,ip-2).i;
+      CX2(ik,lc).i=wal[l].i*CH2(ik,ip-1).r+wal[2*l].i*CH2(ik,ip-2).r;
+      }
+
+    size_t iwal=2*l;
+    size_t j=3, jc=ip-3;
+    for (; j<ipph-1; j+=2, jc-=2)
+      {
+      iwal+=l; if (iwal>ip) iwal-=ip;
+      cmplx<T0> xwal=wal[iwal];
+      iwal+=l; if (iwal>ip) iwal-=ip;
+      cmplx<T0> xwal2=wal[iwal];
+      for (size_t ik=0; ik<idl1; ++ik)
+        {
+        CX2(ik,l).r += CH2(ik,j).r*xwal.r+CH2(ik,j+1).r*xwal2.r;
+        CX2(ik,l).i += CH2(ik,j).i*xwal.r+CH2(ik,j+1).i*xwal2.r;
+        CX2(ik,lc).r -= CH2(ik,jc).i*xwal.i+CH2(ik,jc-1).i*xwal2.i;
+        CX2(ik,lc).i += CH2(ik,jc).r*xwal.i+CH2(ik,jc-1).r*xwal2.i;
+        }
+      }
+    for (; j<ipph; ++j, --jc)
+      {
+      iwal+=l; if (iwal>ip) iwal-=ip;
+      cmplx<T0> xwal=wal[iwal];
+      for (size_t ik=0; ik<idl1; ++ik)
+        {
+        CX2(ik,l).r += CH2(ik,j).r*xwal.r;
+        CX2(ik,l).i += CH2(ik,j).i*xwal.r;
+        CX2(ik,lc).r -= CH2(ik,jc).i*xwal.i;
+        CX2(ik,lc).i += CH2(ik,jc).r*xwal.i;
+        }
+      }
+    }
+
+  // shuffling and twiddling
+  if (ido==1)
+    for (size_t j=1, jc=ip-1; j<ipph; ++j, --jc)
+      for (size_t ik=0; ik<idl1; ++ik)
+        {
+        T t1=CX2(ik,j), t2=CX2(ik,jc);
+        PM(CX2(ik,j),CX2(ik,jc),t1,t2);
+        }
+  else
+    {
+    for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)
+      for (size_t k=0; k<l1; ++k)
+        {
+        T t1=CX(0,k,j), t2=CX(0,k,jc);
+        PM(CX(0,k,j),CX(0,k,jc),t1,t2);
+        for (size_t i=1; i<ido; ++i)
+          {
+          T x1, x2;
+          PM(x1,x2,CX(i,k,j),CX(i,k,jc));
+          size_t idij=(j-1)*(ido-1)+i-1;
+          special_mul<fwd>(x1,wa[idij],CX(i,k,j));
+          idij=(jc-1)*(ido-1)+i-1;
+          special_mul<fwd>(x2,wa[idij],CX(i,k,jc));
+          }
+        }
+    }
+  }
+
+template<bool fwd, typename T> void pass_all(T c[], T0 fct) const
+  {
+  if (length==1) { c[0]*=fct; return; }
+  size_t l1=1;
+  arr<T> ch(length);
+  T *p1=c, *p2=ch.data();
+
+  for(size_t k1=0; k1<fact.size(); k1++)
+    {
+    size_t ip=fact[k1].fct;
+    size_t l2=ip*l1;
+    size_t ido = length/l2;
+    if     (ip==4)
+      pass4<fwd> (ido, l1, p1, p2, fact[k1].tw);
+    else if(ip==8)
+      pass8<fwd>(ido, l1, p1, p2, fact[k1].tw);
+    else if(ip==2)
+      pass2<fwd>(ido, l1, p1, p2, fact[k1].tw);
+    else if(ip==3)
+      pass3<fwd> (ido, l1, p1, p2, fact[k1].tw);
+    else if(ip==5)
+      pass5<fwd> (ido, l1, p1, p2, fact[k1].tw);
+    else if(ip==7)
+      pass7<fwd> (ido, l1, p1, p2, fact[k1].tw);
+    else if(ip==11)
+      pass11<fwd> (ido, l1, p1, p2, fact[k1].tw);
+    else
+      {
+      passg<fwd>(ido, ip, l1, p1, p2, fact[k1].tw, fact[k1].tws);
+      std::swap(p1,p2);
+      }
+    std::swap(p1,p2);
+    l1=l2;
+    }
+  if (p1!=c)
+    {
+    if (fct!=1.)
+      for (size_t i=0; i<length; ++i)
+        c[i] = ch[i]*fct;
+    else
+      std::copy_n (p1, length, c);
+    }
+  else
+    if (fct!=1.)
+      for (size_t i=0; i<length; ++i)
+        c[i] *= fct;
+  }
+
+  public:
+    template<typename T> void exec(T c[], T0 fct, bool fwd) const
+      { fwd ? pass_all<true>(c, fct) : pass_all<false>(c, fct); }
+
+  private:
+    POCKETFFT_NOINLINE void factorize()
+      {
+      size_t len=length;
+      while ((len&7)==0)
+        { add_factor(8); len>>=3; }
+      while ((len&3)==0)
+        { add_factor(4); len>>=2; }
+      if ((len&1)==0)
+        {
+        len>>=1;
+        // factor 2 should be at the front of the factor list
+        add_factor(2);
+        std::swap(fact[0].fct, fact.back().fct);
+        }
+      for (size_t divisor=3; divisor*divisor<=len; divisor+=2)
+        while ((len%divisor)==0)
+          {
+          add_factor(divisor);
+          len/=divisor;
+          }
+      if (len>1) add_factor(len);
+      }
+
+    size_t twsize() const
+      {
+      size_t twsize=0, l1=1;
+      for (size_t k=0; k<fact.size(); ++k)
+        {
+        size_t ip=fact[k].fct, ido= length/(l1*ip);
+        twsize+=(ip-1)*(ido-1);
+        if (ip>11)
+          twsize+=ip;
+        l1*=ip;
+        }
+      return twsize;
+      }
+
+    void comp_twiddle()
+      {
+      sincos_2pibyn<T0> twiddle(length);
+      size_t l1=1;
+      size_t memofs=0;
+      for (size_t k=0; k<fact.size(); ++k)
+        {
+        size_t ip=fact[k].fct, ido=length/(l1*ip);
+        fact[k].tw=mem.data()+memofs;
+        memofs+=(ip-1)*(ido-1);
+        for (size_t j=1; j<ip; ++j)
+          for (size_t i=1; i<ido; ++i)
+            fact[k].tw[(j-1)*(ido-1)+i-1] = twiddle[j*l1*i];
+        if (ip>11)
+          {
+          fact[k].tws=mem.data()+memofs;
+          memofs+=ip;
+          for (size_t j=0; j<ip; ++j)
+            fact[k].tws[j] = twiddle[j*l1*ido];
+          }
+        l1*=ip;
+        }
+      }
+
+  public:
+    POCKETFFT_NOINLINE cfftp(size_t length_)
+      : length(length_)
+      {
+      if (length==0) throw std::runtime_error("zero-length FFT requested");
+      if (length==1) return;
+      factorize();
+      mem.resize(twsize());
+      comp_twiddle();
+      }
+  };
+
+//
+// real-valued FFTPACK transforms
+//
+
+template<typename T0> class rfftp
+  {
+  private:
+    struct fctdata
+      {
+      size_t fct;
+      T0 *tw, *tws;
+      };
+
+    size_t length;
+    arr<T0> mem;
+    std::vector<fctdata> fact;
+
+    void add_factor(size_t factor)
+      { fact.push_back({factor, nullptr, nullptr}); }
+
+/* (a+ib) = conj(c+id) * (e+if) */
+template<typename T1, typename T2, typename T3> inline void MULPM
+  (T1 &a, T1 &b, T2 c, T2 d, T3 e, T3 f) const
+  {  a=c*e+d*f; b=c*f-d*e; }
+
+template<typename T> void radf2 (size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido,l1](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+l1*c)]; };
+  auto CH = [ch,ido](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+2*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    PM (CH(0,0,k),CH(ido-1,1,k),CC(0,k,0),CC(0,k,1));
+  if ((ido&1)==0)
+    for (size_t k=0; k<l1; k++)
+      {
+      CH(    0,1,k) = -CC(ido-1,k,1);
+      CH(ido-1,0,k) =  CC(ido-1,k,0);
+      }
+  if (ido<=2) return;
+  for (size_t k=0; k<l1; k++)
+    for (size_t i=2; i<ido; i+=2)
+      {
+      size_t ic=ido-i;
+      T tr2, ti2;
+      MULPM (tr2,ti2,WA(0,i-2),WA(0,i-1),CC(i-1,k,1),CC(i,k,1));
+      PM (CH(i-1,0,k),CH(ic-1,1,k),CC(i-1,k,0),tr2);
+      PM (CH(i  ,0,k),CH(ic  ,1,k),ti2,CC(i  ,k,0));
+      }
+  }
+
+// a2=a+b; b2=i*(b-a);
+#define POCKETFFT_REARRANGE(rx, ix, ry, iy) \
+  {\
+  auto t1=rx+ry, t2=ry-rx, t3=ix+iy, t4=ix-iy; \
+  rx=t1; ix=t3; ry=t4; iy=t2; \
+  }
+
+template<typename T> void radf3(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 taur=-0.5, taui=T0(0.8660254037844386467637231707529362L);
+
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido,l1](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+l1*c)]; };
+  auto CH = [ch,ido](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+3*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    {
+    T cr2=CC(0,k,1)+CC(0,k,2);
+    CH(0,0,k) = CC(0,k,0)+cr2;
+    CH(0,2,k) = taui*(CC(0,k,2)-CC(0,k,1));
+    CH(ido-1,1,k) = CC(0,k,0)+taur*cr2;
+    }
+  if (ido==1) return;
+  for (size_t k=0; k<l1; k++)
+    for (size_t i=2; i<ido; i+=2)
+      {
+      size_t ic=ido-i;
+      T di2, di3, dr2, dr3;
+      MULPM (dr2,di2,WA(0,i-2),WA(0,i-1),CC(i-1,k,1),CC(i,k,1)); // d2=conj(WA0)*CC1
+      MULPM (dr3,di3,WA(1,i-2),WA(1,i-1),CC(i-1,k,2),CC(i,k,2)); // d3=conj(WA1)*CC2
+      POCKETFFT_REARRANGE(dr2, di2, dr3, di3);
+      CH(i-1,0,k) = CC(i-1,k,0)+dr2; // c add
+      CH(i  ,0,k) = CC(i  ,k,0)+di2;
+      T tr2 = CC(i-1,k,0)+taur*dr2; // c add
+      T ti2 = CC(i  ,k,0)+taur*di2;
+      T tr3 = taui*dr3;  // t3 = taui*i*(d3-d2)?
+      T ti3 = taui*di3;
+      PM(CH(i-1,2,k),CH(ic-1,1,k),tr2,tr3); // PM(i) = t2+t3
+      PM(CH(i  ,2,k),CH(ic  ,1,k),ti3,ti2); // PM(ic) = conj(t2-t3)
+      }
+  }
+
+template<typename T> void radf4(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 hsqt2=T0(0.707106781186547524400844362104849L);
+
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido,l1](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+l1*c)]; };
+  auto CH = [ch,ido](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+4*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    {
+    T tr1,tr2;
+    PM (tr1,CH(0,2,k),CC(0,k,3),CC(0,k,1));
+    PM (tr2,CH(ido-1,1,k),CC(0,k,0),CC(0,k,2));
+    PM (CH(0,0,k),CH(ido-1,3,k),tr2,tr1);
+    }
+  if ((ido&1)==0)
+    for (size_t k=0; k<l1; k++)
+      {
+      T ti1=-hsqt2*(CC(ido-1,k,1)+CC(ido-1,k,3));
+      T tr1= hsqt2*(CC(ido-1,k,1)-CC(ido-1,k,3));
+      PM (CH(ido-1,0,k),CH(ido-1,2,k),CC(ido-1,k,0),tr1);
+      PM (CH(    0,3,k),CH(    0,1,k),ti1,CC(ido-1,k,2));
+      }
+  if (ido<=2) return;
+  for (size_t k=0; k<l1; k++)
+    for (size_t i=2; i<ido; i+=2)
+      {
+      size_t ic=ido-i;
+      T ci2, ci3, ci4, cr2, cr3, cr4, ti1, ti2, ti3, ti4, tr1, tr2, tr3, tr4;
+      MULPM(cr2,ci2,WA(0,i-2),WA(0,i-1),CC(i-1,k,1),CC(i,k,1));
+      MULPM(cr3,ci3,WA(1,i-2),WA(1,i-1),CC(i-1,k,2),CC(i,k,2));
+      MULPM(cr4,ci4,WA(2,i-2),WA(2,i-1),CC(i-1,k,3),CC(i,k,3));
+      PM(tr1,tr4,cr4,cr2);
+      PM(ti1,ti4,ci2,ci4);
+      PM(tr2,tr3,CC(i-1,k,0),cr3);
+      PM(ti2,ti3,CC(i  ,k,0),ci3);
+      PM(CH(i-1,0,k),CH(ic-1,3,k),tr2,tr1);
+      PM(CH(i  ,0,k),CH(ic  ,3,k),ti1,ti2);
+      PM(CH(i-1,2,k),CH(ic-1,1,k),tr3,ti4);
+      PM(CH(i  ,2,k),CH(ic  ,1,k),tr4,ti3);
+      }
+  }
+
+template<typename T> void radf5(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 tr11= T0(0.3090169943749474241022934171828191L),
+               ti11= T0(0.9510565162951535721164393333793821L),
+               tr12= T0(-0.8090169943749474241022934171828191L),
+               ti12= T0(0.5877852522924731291687059546390728L);
+
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido,l1](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+l1*c)]; };
+  auto CH = [ch,ido](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+5*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    {
+    T cr2, cr3, ci4, ci5;
+    PM (cr2,ci5,CC(0,k,4),CC(0,k,1));
+    PM (cr3,ci4,CC(0,k,3),CC(0,k,2));
+    CH(0,0,k)=CC(0,k,0)+cr2+cr3;
+    CH(ido-1,1,k)=CC(0,k,0)+tr11*cr2+tr12*cr3;
+    CH(0,2,k)=ti11*ci5+ti12*ci4;
+    CH(ido-1,3,k)=CC(0,k,0)+tr12*cr2+tr11*cr3;
+    CH(0,4,k)=ti12*ci5-ti11*ci4;
+    }
+  if (ido==1) return;
+  for (size_t k=0; k<l1;++k)
+    for (size_t i=2, ic=ido-2; i<ido; i+=2, ic-=2)
+      {
+      T di2, di3, di4, di5, dr2, dr3, dr4, dr5;
+      MULPM (dr2,di2,WA(0,i-2),WA(0,i-1),CC(i-1,k,1),CC(i,k,1));
+      MULPM (dr3,di3,WA(1,i-2),WA(1,i-1),CC(i-1,k,2),CC(i,k,2));
+      MULPM (dr4,di4,WA(2,i-2),WA(2,i-1),CC(i-1,k,3),CC(i,k,3));
+      MULPM (dr5,di5,WA(3,i-2),WA(3,i-1),CC(i-1,k,4),CC(i,k,4));
+      POCKETFFT_REARRANGE(dr2, di2, dr5, di5);
+      POCKETFFT_REARRANGE(dr3, di3, dr4, di4);
+      CH(i-1,0,k)=CC(i-1,k,0)+dr2+dr3;
+      CH(i  ,0,k)=CC(i  ,k,0)+di2+di3;
+      T tr2=CC(i-1,k,0)+tr11*dr2+tr12*dr3;
+      T ti2=CC(i  ,k,0)+tr11*di2+tr12*di3;
+      T tr3=CC(i-1,k,0)+tr12*dr2+tr11*dr3;
+      T ti3=CC(i  ,k,0)+tr12*di2+tr11*di3;
+      T tr5 = ti11*dr5 + ti12*dr4;
+      T ti5 = ti11*di5 + ti12*di4;
+      T tr4 = ti12*dr5 - ti11*dr4;
+      T ti4 = ti12*di5 - ti11*di4;
+      PM(CH(i-1,2,k),CH(ic-1,1,k),tr2,tr5);
+      PM(CH(i  ,2,k),CH(ic  ,1,k),ti5,ti2);
+      PM(CH(i-1,4,k),CH(ic-1,3,k),tr3,tr4);
+      PM(CH(i  ,4,k),CH(ic  ,3,k),ti4,ti3);
+      }
+  }
+
+#undef POCKETFFT_REARRANGE
+
+template<typename T> void radfg(size_t ido, size_t ip, size_t l1,
+  T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa, const T0 * POCKETFFT_RESTRICT csarr) const
+  {
+  const size_t cdim=ip;
+  size_t ipph=(ip+1)/2;
+  size_t idl1 = ido*l1;
+
+  auto CC = [cc,ido,cdim](size_t a, size_t b, size_t c) -> T&
+    { return cc[a+ido*(b+cdim*c)]; };
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> const T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto C1 = [cc,ido,l1] (size_t a, size_t b, size_t c) -> T&
+    { return cc[a+ido*(b+l1*c)]; };
+  auto C2 = [cc,idl1] (size_t a, size_t b) -> T&
+    { return cc[a+idl1*b]; };
+  auto CH2 = [ch,idl1] (size_t a, size_t b) -> T&
+    { return ch[a+idl1*b]; };
+
+  if (ido>1)
+    {
+    for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)              // 114
+      {
+      size_t is=(j-1)*(ido-1),
+             is2=(jc-1)*(ido-1);
+      for (size_t k=0; k<l1; ++k)                            // 113
+        {
+        size_t idij=is;
+        size_t idij2=is2;
+        for (size_t i=1; i<=ido-2; i+=2)                      // 112
+          {
+          T t1=C1(i,k,j ), t2=C1(i+1,k,j ),
+            t3=C1(i,k,jc), t4=C1(i+1,k,jc);
+          T x1=wa[idij]*t1 + wa[idij+1]*t2,
+            x2=wa[idij]*t2 - wa[idij+1]*t1,
+            x3=wa[idij2]*t3 + wa[idij2+1]*t4,
+            x4=wa[idij2]*t4 - wa[idij2+1]*t3;
+          PM(C1(i,k,j),C1(i+1,k,jc),x3,x1);
+          PM(C1(i+1,k,j),C1(i,k,jc),x2,x4);
+          idij+=2;
+          idij2+=2;
+          }
+        }
+      }
+    }
+
+  for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)                // 123
+    for (size_t k=0; k<l1; ++k)                              // 122
+      MPINPLACE(C1(0,k,jc), C1(0,k,j));
+
+//everything in C
+//memset(ch,0,ip*l1*ido*sizeof(double));
+
+  for (size_t l=1,lc=ip-1; l<ipph; ++l,--lc)                 // 127
+    {
+    for (size_t ik=0; ik<idl1; ++ik)                         // 124
+      {
+      CH2(ik,l ) = C2(ik,0)+csarr[2*l]*C2(ik,1)+csarr[4*l]*C2(ik,2);
+      CH2(ik,lc) = csarr[2*l+1]*C2(ik,ip-1)+csarr[4*l+1]*C2(ik,ip-2);
+      }
+    size_t iang = 2*l;
+    size_t j=3, jc=ip-3;
+    for (; j<ipph-3; j+=4,jc-=4)              // 126
+      {
+      iang+=l; if (iang>=ip) iang-=ip;
+      T0 ar1=csarr[2*iang], ai1=csarr[2*iang+1];
+      iang+=l; if (iang>=ip) iang-=ip;
+      T0 ar2=csarr[2*iang], ai2=csarr[2*iang+1];
+      iang+=l; if (iang>=ip) iang-=ip;
+      T0 ar3=csarr[2*iang], ai3=csarr[2*iang+1];
+      iang+=l; if (iang>=ip) iang-=ip;
+      T0 ar4=csarr[2*iang], ai4=csarr[2*iang+1];
+      for (size_t ik=0; ik<idl1; ++ik)                       // 125
+        {
+        CH2(ik,l ) += ar1*C2(ik,j )+ar2*C2(ik,j +1)
+                     +ar3*C2(ik,j +2)+ar4*C2(ik,j +3);
+        CH2(ik,lc) += ai1*C2(ik,jc)+ai2*C2(ik,jc-1)
+                     +ai3*C2(ik,jc-2)+ai4*C2(ik,jc-3);
+        }
+      }
+    for (; j<ipph-1; j+=2,jc-=2)              // 126
+      {
+      iang+=l; if (iang>=ip) iang-=ip;
+      T0 ar1=csarr[2*iang], ai1=csarr[2*iang+1];
+      iang+=l; if (iang>=ip) iang-=ip;
+      T0 ar2=csarr[2*iang], ai2=csarr[2*iang+1];
+      for (size_t ik=0; ik<idl1; ++ik)                       // 125
+        {
+        CH2(ik,l ) += ar1*C2(ik,j )+ar2*C2(ik,j +1);
+        CH2(ik,lc) += ai1*C2(ik,jc)+ai2*C2(ik,jc-1);
+        }
+      }
+    for (; j<ipph; ++j,--jc)              // 126
+      {
+      iang+=l; if (iang>=ip) iang-=ip;
+      T0 ar=csarr[2*iang], ai=csarr[2*iang+1];
+      for (size_t ik=0; ik<idl1; ++ik)                       // 125
+        {
+        CH2(ik,l ) += ar*C2(ik,j );
+        CH2(ik,lc) += ai*C2(ik,jc);
+        }
+      }
+    }
+  for (size_t ik=0; ik<idl1; ++ik)                         // 101
+    CH2(ik,0) = C2(ik,0);
+  for (size_t j=1; j<ipph; ++j)                              // 129
+    for (size_t ik=0; ik<idl1; ++ik)                         // 128
+      CH2(ik,0) += C2(ik,j);
+
+// everything in CH at this point!
+//memset(cc,0,ip*l1*ido*sizeof(double));
+
+  for (size_t k=0; k<l1; ++k)                                // 131
+    for (size_t i=0; i<ido; ++i)                             // 130
+      CC(i,0,k) = CH(i,k,0);
+
+  for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)                // 137
+    {
+    size_t j2=2*j-1;
+    for (size_t k=0; k<l1; ++k)                              // 136
+      {
+      CC(ido-1,j2,k) = CH(0,k,j);
+      CC(0,j2+1,k) = CH(0,k,jc);
+      }
+    }
+
+  if (ido==1) return;
+
+  for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)                // 140
+    {
+    size_t j2=2*j-1;
+    for(size_t k=0; k<l1; ++k)                               // 139
+      for(size_t i=1, ic=ido-i-2; i<=ido-2; i+=2, ic-=2)      // 138
+        {
+        CC(i   ,j2+1,k) = CH(i  ,k,j )+CH(i  ,k,jc);
+        CC(ic  ,j2  ,k) = CH(i  ,k,j )-CH(i  ,k,jc);
+        CC(i+1 ,j2+1,k) = CH(i+1,k,j )+CH(i+1,k,jc);
+        CC(ic+1,j2  ,k) = CH(i+1,k,jc)-CH(i+1,k,j );
+        }
+    }
+  }
+
+template<typename T> void radb2(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+2*c)]; };
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    PM (CH(0,k,0),CH(0,k,1),CC(0,0,k),CC(ido-1,1,k));
+  if ((ido&1)==0)
+    for (size_t k=0; k<l1; k++)
+      {
+      CH(ido-1,k,0) = 2*CC(ido-1,0,k);
+      CH(ido-1,k,1) =-2*CC(0    ,1,k);
+      }
+  if (ido<=2) return;
+  for (size_t k=0; k<l1;++k)
+    for (size_t i=2; i<ido; i+=2)
+      {
+      size_t ic=ido-i;
+      T ti2, tr2;
+      PM (CH(i-1,k,0),tr2,CC(i-1,0,k),CC(ic-1,1,k));
+      PM (ti2,CH(i  ,k,0),CC(i  ,0,k),CC(ic  ,1,k));
+      MULPM (CH(i,k,1),CH(i-1,k,1),WA(0,i-2),WA(0,i-1),ti2,tr2);
+      }
+  }
+
+template<typename T> void radb3(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 taur=-0.5, taui=T0(0.8660254037844386467637231707529362L);
+
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+3*c)]; };
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    {
+    T tr2=2*CC(ido-1,1,k);
+    T cr2=CC(0,0,k)+taur*tr2;
+    CH(0,k,0)=CC(0,0,k)+tr2;
+    T ci3=2*taui*CC(0,2,k);
+    PM (CH(0,k,2),CH(0,k,1),cr2,ci3);
+    }
+  if (ido==1) return;
+  for (size_t k=0; k<l1; k++)
+    for (size_t i=2, ic=ido-2; i<ido; i+=2, ic-=2)
+      {
+      T tr2=CC(i-1,2,k)+CC(ic-1,1,k); // t2=CC(I) + conj(CC(ic))
+      T ti2=CC(i  ,2,k)-CC(ic  ,1,k);
+      T cr2=CC(i-1,0,k)+taur*tr2;     // c2=CC +taur*t2
+      T ci2=CC(i  ,0,k)+taur*ti2;
+      CH(i-1,k,0)=CC(i-1,0,k)+tr2;         // CH=CC+t2
+      CH(i  ,k,0)=CC(i  ,0,k)+ti2;
+      T cr3=taui*(CC(i-1,2,k)-CC(ic-1,1,k));// c3=taui*(CC(i)-conj(CC(ic)))
+      T ci3=taui*(CC(i  ,2,k)+CC(ic  ,1,k));
+      T di2, di3, dr2, dr3;
+      PM(dr3,dr2,cr2,ci3); // d2= (cr2-ci3, ci2+cr3) = c2+i*c3
+      PM(di2,di3,ci2,cr3); // d3= (cr2+ci3, ci2-cr3) = c2-i*c3
+      MULPM(CH(i,k,1),CH(i-1,k,1),WA(0,i-2),WA(0,i-1),di2,dr2); // ch = WA*d2
+      MULPM(CH(i,k,2),CH(i-1,k,2),WA(1,i-2),WA(1,i-1),di3,dr3);
+      }
+  }
+
+template<typename T> void radb4(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 sqrt2=T0(1.414213562373095048801688724209698L);
+
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+4*c)]; };
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    {
+    T tr1, tr2;
+    PM (tr2,tr1,CC(0,0,k),CC(ido-1,3,k));
+    T tr3=2*CC(ido-1,1,k);
+    T tr4=2*CC(0,2,k);
+    PM (CH(0,k,0),CH(0,k,2),tr2,tr3);
+    PM (CH(0,k,3),CH(0,k,1),tr1,tr4);
+    }
+  if ((ido&1)==0)
+    for (size_t k=0; k<l1; k++)
+      {
+      T tr1,tr2,ti1,ti2;
+      PM (ti1,ti2,CC(0    ,3,k),CC(0    ,1,k));
+      PM (tr2,tr1,CC(ido-1,0,k),CC(ido-1,2,k));
+      CH(ido-1,k,0)=tr2+tr2;
+      CH(ido-1,k,1)=sqrt2*(tr1-ti1);
+      CH(ido-1,k,2)=ti2+ti2;
+      CH(ido-1,k,3)=-sqrt2*(tr1+ti1);
+      }
+  if (ido<=2) return;
+  for (size_t k=0; k<l1;++k)
+    for (size_t i=2; i<ido; i+=2)
+      {
+      T ci2, ci3, ci4, cr2, cr3, cr4, ti1, ti2, ti3, ti4, tr1, tr2, tr3, tr4;
+      size_t ic=ido-i;
+      PM (tr2,tr1,CC(i-1,0,k),CC(ic-1,3,k));
+      PM (ti1,ti2,CC(i  ,0,k),CC(ic  ,3,k));
+      PM (tr4,ti3,CC(i  ,2,k),CC(ic  ,1,k));
+      PM (tr3,ti4,CC(i-1,2,k),CC(ic-1,1,k));
+      PM (CH(i-1,k,0),cr3,tr2,tr3);
+      PM (CH(i  ,k,0),ci3,ti2,ti3);
+      PM (cr4,cr2,tr1,tr4);
+      PM (ci2,ci4,ti1,ti4);
+      MULPM (CH(i,k,1),CH(i-1,k,1),WA(0,i-2),WA(0,i-1),ci2,cr2);
+      MULPM (CH(i,k,2),CH(i-1,k,2),WA(1,i-2),WA(1,i-1),ci3,cr3);
+      MULPM (CH(i,k,3),CH(i-1,k,3),WA(2,i-2),WA(2,i-1),ci4,cr4);
+      }
+  }
+
+template<typename T> void radb5(size_t ido, size_t l1,
+  const T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa) const
+  {
+  constexpr T0 tr11= T0(0.3090169943749474241022934171828191L),
+               ti11= T0(0.9510565162951535721164393333793821L),
+               tr12= T0(-0.8090169943749474241022934171828191L),
+               ti12= T0(0.5877852522924731291687059546390728L);
+
+  auto WA = [wa,ido](size_t x, size_t i) { return wa[i+x*(ido-1)]; };
+  auto CC = [cc,ido](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+5*c)]; };
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+
+  for (size_t k=0; k<l1; k++)
+    {
+    T ti5=CC(0,2,k)+CC(0,2,k);
+    T ti4=CC(0,4,k)+CC(0,4,k);
+    T tr2=CC(ido-1,1,k)+CC(ido-1,1,k);
+    T tr3=CC(ido-1,3,k)+CC(ido-1,3,k);
+    CH(0,k,0)=CC(0,0,k)+tr2+tr3;
+    T cr2=CC(0,0,k)+tr11*tr2+tr12*tr3;
+    T cr3=CC(0,0,k)+tr12*tr2+tr11*tr3;
+    T ci4, ci5;
+    MULPM(ci5,ci4,ti5,ti4,ti11,ti12);
+    PM(CH(0,k,4),CH(0,k,1),cr2,ci5);
+    PM(CH(0,k,3),CH(0,k,2),cr3,ci4);
+    }
+  if (ido==1) return;
+  for (size_t k=0; k<l1;++k)
+    for (size_t i=2, ic=ido-2; i<ido; i+=2, ic-=2)
+      {
+      T tr2, tr3, tr4, tr5, ti2, ti3, ti4, ti5;
+      PM(tr2,tr5,CC(i-1,2,k),CC(ic-1,1,k));
+      PM(ti5,ti2,CC(i  ,2,k),CC(ic  ,1,k));
+      PM(tr3,tr4,CC(i-1,4,k),CC(ic-1,3,k));
+      PM(ti4,ti3,CC(i  ,4,k),CC(ic  ,3,k));
+      CH(i-1,k,0)=CC(i-1,0,k)+tr2+tr3;
+      CH(i  ,k,0)=CC(i  ,0,k)+ti2+ti3;
+      T cr2=CC(i-1,0,k)+tr11*tr2+tr12*tr3;
+      T ci2=CC(i  ,0,k)+tr11*ti2+tr12*ti3;
+      T cr3=CC(i-1,0,k)+tr12*tr2+tr11*tr3;
+      T ci3=CC(i  ,0,k)+tr12*ti2+tr11*ti3;
+      T ci4, ci5, cr5, cr4;
+      MULPM(cr5,cr4,tr5,tr4,ti11,ti12);
+      MULPM(ci5,ci4,ti5,ti4,ti11,ti12);
+      T dr2, dr3, dr4, dr5, di2, di3, di4, di5;
+      PM(dr4,dr3,cr3,ci4);
+      PM(di3,di4,ci3,cr4);
+      PM(dr5,dr2,cr2,ci5);
+      PM(di2,di5,ci2,cr5);
+      MULPM(CH(i,k,1),CH(i-1,k,1),WA(0,i-2),WA(0,i-1),di2,dr2);
+      MULPM(CH(i,k,2),CH(i-1,k,2),WA(1,i-2),WA(1,i-1),di3,dr3);
+      MULPM(CH(i,k,3),CH(i-1,k,3),WA(2,i-2),WA(2,i-1),di4,dr4);
+      MULPM(CH(i,k,4),CH(i-1,k,4),WA(3,i-2),WA(3,i-1),di5,dr5);
+      }
+  }
+
+template<typename T> void radbg(size_t ido, size_t ip, size_t l1,
+  T * POCKETFFT_RESTRICT cc, T * POCKETFFT_RESTRICT ch,
+  const T0 * POCKETFFT_RESTRICT wa, const T0 * POCKETFFT_RESTRICT csarr) const
+  {
+  const size_t cdim=ip;
+  size_t ipph=(ip+1)/ 2;
+  size_t idl1 = ido*l1;
+
+  auto CC = [cc,ido,cdim](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+cdim*c)]; };
+  auto CH = [ch,ido,l1](size_t a, size_t b, size_t c) -> T&
+    { return ch[a+ido*(b+l1*c)]; };
+  auto C1 = [cc,ido,l1](size_t a, size_t b, size_t c) -> const T&
+    { return cc[a+ido*(b+l1*c)]; };
+  auto C2 = [cc,idl1](size_t a, size_t b) -> T&
+    { return cc[a+idl1*b]; };
+  auto CH2 = [ch,idl1](size_t a, size_t b) -> T&
+    { return ch[a+idl1*b]; };
+
+  for (size_t k=0; k<l1; ++k)        // 102
+    for (size_t i=0; i<ido; ++i)     // 101
+      CH(i,k,0) = CC(i,0,k);
+  for (size_t j=1, jc=ip-1; j<ipph; ++j, --jc)   // 108
+    {
+    size_t j2=2*j-1;
+    for (size_t k=0; k<l1; ++k)
+      {
+      CH(0,k,j ) = 2*CC(ido-1,j2,k);
+      CH(0,k,jc) = 2*CC(0,j2+1,k);
+      }
+    }
+
+  if (ido!=1)
+    {
+    for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)   // 111
+      {
+      size_t j2=2*j-1;
+      for (size_t k=0; k<l1; ++k)
+        for (size_t i=1, ic=ido-i-2; i<=ido-2; i+=2, ic-=2)      // 109
+          {
+          CH(i  ,k,j ) = CC(i  ,j2+1,k)+CC(ic  ,j2,k);
+          CH(i  ,k,jc) = CC(i  ,j2+1,k)-CC(ic  ,j2,k);
+          CH(i+1,k,j ) = CC(i+1,j2+1,k)-CC(ic+1,j2,k);
+          CH(i+1,k,jc) = CC(i+1,j2+1,k)+CC(ic+1,j2,k);
+          }
+      }
+    }
+  for (size_t l=1,lc=ip-1; l<ipph; ++l,--lc)
+    {
+    for (size_t ik=0; ik<idl1; ++ik)
+      {
+      C2(ik,l ) = CH2(ik,0)+csarr[2*l]*CH2(ik,1)+csarr[4*l]*CH2(ik,2);
+      C2(ik,lc) = csarr[2*l+1]*CH2(ik,ip-1)+csarr[4*l+1]*CH2(ik,ip-2);
+      }
+    size_t iang=2*l;
+    size_t j=3,jc=ip-3;
+    for(; j<ipph-3; j+=4,jc-=4)
+      {
+      iang+=l; if(iang>ip) iang-=ip;
+      T0 ar1=csarr[2*iang], ai1=csarr[2*iang+1];
+      iang+=l; if(iang>ip) iang-=ip;
+      T0 ar2=csarr[2*iang], ai2=csarr[2*iang+1];
+      iang+=l; if(iang>ip) iang-=ip;
+      T0 ar3=csarr[2*iang], ai3=csarr[2*iang+1];
+      iang+=l; if(iang>ip) iang-=ip;
+      T0 ar4=csarr[2*iang], ai4=csarr[2*iang+1];
+      for (size_t ik=0; ik<idl1; ++ik)
+        {
+        C2(ik,l ) += ar1*CH2(ik,j )+ar2*CH2(ik,j +1)
+                    +ar3*CH2(ik,j +2)+ar4*CH2(ik,j +3);
+        C2(ik,lc) += ai1*CH2(ik,jc)+ai2*CH2(ik,jc-1)
+                    +ai3*CH2(ik,jc-2)+ai4*CH2(ik,jc-3);
+        }
+      }
+    for(; j<ipph-1; j+=2,jc-=2)
+      {
+      iang+=l; if(iang>ip) iang-=ip;
+      T0 ar1=csarr[2*iang], ai1=csarr[2*iang+1];
+      iang+=l; if(iang>ip) iang-=ip;
+      T0 ar2=csarr[2*iang], ai2=csarr[2*iang+1];
+      for (size_t ik=0; ik<idl1; ++ik)
+        {
+        C2(ik,l ) += ar1*CH2(ik,j )+ar2*CH2(ik,j +1);
+        C2(ik,lc) += ai1*CH2(ik,jc)+ai2*CH2(ik,jc-1);
+        }
+      }
+    for(; j<ipph; ++j,--jc)
+      {
+      iang+=l; if(iang>ip) iang-=ip;
+      T0 war=csarr[2*iang], wai=csarr[2*iang+1];
+      for (size_t ik=0; ik<idl1; ++ik)
+        {
+        C2(ik,l ) += war*CH2(ik,j );
+        C2(ik,lc) += wai*CH2(ik,jc);
+        }
+      }
+    }
+  for (size_t j=1; j<ipph; ++j)
+    for (size_t ik=0; ik<idl1; ++ik)
+      CH2(ik,0) += CH2(ik,j);
+  for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)   // 124
+    for (size_t k=0; k<l1; ++k)
+      PM(CH(0,k,jc),CH(0,k,j),C1(0,k,j),C1(0,k,jc));
+
+  if (ido==1) return;
+
+  for (size_t j=1, jc=ip-1; j<ipph; ++j, --jc)  // 127
+    for (size_t k=0; k<l1; ++k)
+      for (size_t i=1; i<=ido-2; i+=2)
+        {
+        CH(i  ,k,j ) = C1(i  ,k,j)-C1(i+1,k,jc);
+        CH(i  ,k,jc) = C1(i  ,k,j)+C1(i+1,k,jc);
+        CH(i+1,k,j ) = C1(i+1,k,j)+C1(i  ,k,jc);
+        CH(i+1,k,jc) = C1(i+1,k,j)-C1(i  ,k,jc);
+        }
+
+// All in CH
+
+  for (size_t j=1; j<ip; ++j)
+    {
+    size_t is = (j-1)*(ido-1);
+    for (size_t k=0; k<l1; ++k)
+      {
+      size_t idij = is;
+      for (size_t i=1; i<=ido-2; i+=2)
+        {
+        T t1=CH(i,k,j), t2=CH(i+1,k,j);
+        CH(i  ,k,j) = wa[idij]*t1-wa[idij+1]*t2;
+        CH(i+1,k,j) = wa[idij]*t2+wa[idij+1]*t1;
+        idij+=2;
+        }
+      }
+    }
+  }
+
+    template<typename T> void copy_and_norm(T *c, T *p1, T0 fct) const
+      {
+      if (p1!=c)
+        {
+        if (fct!=1.)
+          for (size_t i=0; i<length; ++i)
+            c[i] = fct*p1[i];
+        else
+          std::copy_n (p1, length, c);
+        }
+      else
+        if (fct!=1.)
+          for (size_t i=0; i<length; ++i)
+            c[i] *= fct;
+      }
+
+  public:
+    template<typename T> void exec(T c[], T0 fct, bool r2hc) const
+      {
+      if (length==1) { c[0]*=fct; return; }
+      size_t nf=fact.size();
+      arr<T> ch(length);
+      T *p1=c, *p2=ch.data();
+
+      if (r2hc)
+        for(size_t k1=0, l1=length; k1<nf;++k1)
+          {
+          size_t k=nf-k1-1;
+          size_t ip=fact[k].fct;
+          size_t ido=length / l1;
+          l1 /= ip;
+          if(ip==4)
+            radf4(ido, l1, p1, p2, fact[k].tw);
+          else if(ip==2)
+            radf2(ido, l1, p1, p2, fact[k].tw);
+          else if(ip==3)
+            radf3(ido, l1, p1, p2, fact[k].tw);
+          else if(ip==5)
+            radf5(ido, l1, p1, p2, fact[k].tw);
+          else
+            { radfg(ido, ip, l1, p1, p2, fact[k].tw, fact[k].tws); std::swap (p1,p2); }
+          std::swap (p1,p2);
+          }
+      else
+        for(size_t k=0, l1=1; k<nf; k++)
+          {
+          size_t ip = fact[k].fct,
+                 ido= length/(ip*l1);
+          if(ip==4)
+            radb4(ido, l1, p1, p2, fact[k].tw);
+          else if(ip==2)
+            radb2(ido, l1, p1, p2, fact[k].tw);
+          else if(ip==3)
+            radb3(ido, l1, p1, p2, fact[k].tw);
+          else if(ip==5)
+            radb5(ido, l1, p1, p2, fact[k].tw);
+          else
+            radbg(ido, ip, l1, p1, p2, fact[k].tw, fact[k].tws);
+          std::swap (p1,p2);
+          l1*=ip;
+          }
+
+      copy_and_norm(c,p1,fct);
+      }
+
+  private:
+    void factorize()
+      {
+      size_t len=length;
+      while ((len%4)==0)
+        { add_factor(4); len>>=2; }
+      if ((len%2)==0)
+        {
+        len>>=1;
+        // factor 2 should be at the front of the factor list
+        add_factor(2);
+        std::swap(fact[0].fct, fact.back().fct);
+        }
+      for (size_t divisor=3; divisor*divisor<=len; divisor+=2)
+        while ((len%divisor)==0)
+          {
+          add_factor(divisor);
+          len/=divisor;
+          }
+      if (len>1) add_factor(len);
+      }
+
+    size_t twsize() const
+      {
+      size_t twsz=0, l1=1;
+      for (size_t k=0; k<fact.size(); ++k)
+        {
+        size_t ip=fact[k].fct, ido=length/(l1*ip);
+        twsz+=(ip-1)*(ido-1);
+        if (ip>5) twsz+=2*ip;
+        l1*=ip;
+        }
+      return twsz;
+      }
+
+    void comp_twiddle()
+      {
+      sincos_2pibyn<T0> twid(length);
+      size_t l1=1;
+      T0 *ptr=mem.data();
+      for (size_t k=0; k<fact.size(); ++k)
+        {
+        size_t ip=fact[k].fct, ido=length/(l1*ip);
+        if (k<fact.size()-1) // last factor doesn't need twiddles
+          {
+          fact[k].tw=ptr; ptr+=(ip-1)*(ido-1);
+          for (size_t j=1; j<ip; ++j)
+            for (size_t i=1; i<=(ido-1)/2; ++i)
+              {
+              fact[k].tw[(j-1)*(ido-1)+2*i-2] = twid[j*l1*i].r;
+              fact[k].tw[(j-1)*(ido-1)+2*i-1] = twid[j*l1*i].i;
+              }
+          }
+        if (ip>5) // special factors required by *g functions
+          {
+          fact[k].tws=ptr; ptr+=2*ip;
+          fact[k].tws[0] = 1.;
+          fact[k].tws[1] = 0.;
+          for (size_t i=2, ic=2*ip-2; i<=ic; i+=2, ic-=2)
+            {
+            fact[k].tws[i  ] = twid[i/2*(length/ip)].r;
+            fact[k].tws[i+1] = twid[i/2*(length/ip)].i;
+            fact[k].tws[ic]   = twid[i/2*(length/ip)].r;
+            fact[k].tws[ic+1] = -twid[i/2*(length/ip)].i;
+            }
+          }
+        l1*=ip;
+        }
+      }
+
+  public:
+    POCKETFFT_NOINLINE rfftp(size_t length_)
+      : length(length_)
+      {
+      if (length==0) throw std::runtime_error("zero-length FFT requested");
+      if (length==1) return;
+      factorize();
+      mem.resize(twsize());
+      comp_twiddle();
+      }
+};
+
+//
+// complex Bluestein transforms
+//
+
+template<typename T0> class fftblue
+  {
+  private:
+    size_t n, n2;
+    cfftp<T0> plan;
+    arr<cmplx<T0>> mem;
+    cmplx<T0> *bk, *bkf;
+
+    template<bool fwd, typename T> void fft(cmplx<T> c[], T0 fct) const
+      {
+      arr<cmplx<T>> akf(n2);
+
+      /* initialize a_k and FFT it */
+      for (size_t m=0; m<n; ++m)
+        special_mul<fwd>(c[m],bk[m],akf[m]);
+      auto zero = akf[0]*T0(0);
+      for (size_t m=n; m<n2; ++m)
+        akf[m]=zero;
+
+      plan.exec (akf.data(),1.,true);
+
+      /* do the convolution */
+      akf[0] = akf[0].template special_mul<!fwd>(bkf[0]);
+      for (size_t m=1; m<(n2+1)/2; ++m)
+        {
+        akf[m] = akf[m].template special_mul<!fwd>(bkf[m]);
+        akf[n2-m] = akf[n2-m].template special_mul<!fwd>(bkf[m]);
+        }
+      if ((n2&1)==0)
+        akf[n2/2] = akf[n2/2].template special_mul<!fwd>(bkf[n2/2]);
+
+      /* inverse FFT */
+      plan.exec (akf.data(),1.,false);
+
+      /* multiply by b_k */
+      for (size_t m=0; m<n; ++m)
+        c[m] = akf[m].template special_mul<fwd>(bk[m])*fct;
+      }
+
+  public:
+    POCKETFFT_NOINLINE fftblue(size_t length)
+      : n(length), n2(util::good_size_cmplx(n*2-1)), plan(n2), mem(n+n2/2+1),
+        bk(mem.data()), bkf(mem.data()+n)
+      {
+      /* initialize b_k */
+      sincos_2pibyn<T0> tmp(2*n);
+      bk[0].Set(1, 0);
+
+      size_t coeff=0;
+      for (size_t m=1; m<n; ++m)
+        {
+        coeff+=2*m-1;
+        if (coeff>=2*n) coeff-=2*n;
+        bk[m] = tmp[coeff];
+        }
+
+      /* initialize the zero-padded, Fourier transformed b_k. Add normalisation. */
+      arr<cmplx<T0>> tbkf(n2);
+      T0 xn2 = T0(1)/T0(n2);
+      tbkf[0] = bk[0]*xn2;
+      for (size_t m=1; m<n; ++m)
+        tbkf[m] = tbkf[n2-m] = bk[m]*xn2;
+      for (size_t m=n;m<=(n2-n);++m)
+        tbkf[m].Set(0.,0.);
+      plan.exec(tbkf.data(),1.,true);
+      for (size_t i=0; i<n2/2+1; ++i)
+        bkf[i] = tbkf[i];
+      }
+
+    template<typename T> void exec(cmplx<T> c[], T0 fct, bool fwd) const
+      { fwd ? fft<true>(c,fct) : fft<false>(c,fct); }
+
+    template<typename T> void exec_r(T c[], T0 fct, bool fwd)
+      {
+      arr<cmplx<T>> tmp(n);
+      if (fwd)
+        {
+        auto zero = T0(0)*c[0];
+        for (size_t m=0; m<n; ++m)
+          tmp[m].Set(c[m], zero);
+        fft<true>(tmp.data(),fct);
+        c[0] = tmp[0].r;
+        std::copy_n (&tmp[1].r, n-1, &c[1]);
+        }
+      else
+        {
+        tmp[0].Set(c[0],c[0]*0);
+        std::copy_n (c+1, n-1, &tmp[1].r);
+        if ((n&1)==0) tmp[n/2].i=T0(0)*c[0];
+        for (size_t m=1; 2*m<n; ++m)
+          tmp[n-m].Set(tmp[m].r, -tmp[m].i);
+        fft<false>(tmp.data(),fct);
+        for (size_t m=0; m<n; ++m)
+          c[m] = tmp[m].r;
+        }
+      }
+  };
+
+//
+// flexible (FFTPACK/Bluestein) complex 1D transform
+//
+
+template<typename T0> class pocketfft_c
+  {
+  private:
+    std::unique_ptr<cfftp<T0>> packplan;
+    std::unique_ptr<fftblue<T0>> blueplan;
+    size_t len;
+
+  public:
+    POCKETFFT_NOINLINE pocketfft_c(size_t length)
+      : len(length)
+      {
+      if (length==0) throw std::runtime_error("zero-length FFT requested");
+      size_t tmp = (length<50) ? 0 : util::largest_prime_factor(length);
+      if (tmp*tmp <= length)
+        {
+        packplan=std::unique_ptr<cfftp<T0>>(new cfftp<T0>(length));
+        return;
+        }
+      double comp1 = util::cost_guess(length);
+      double comp2 = 2*util::cost_guess(util::good_size_cmplx(2*length-1));
+      comp2*=1.5; /* fudge factor that appears to give good overall performance */
+      if (comp2<comp1) // use Bluestein
+        blueplan=std::unique_ptr<fftblue<T0>>(new fftblue<T0>(length));
+      else
+        packplan=std::unique_ptr<cfftp<T0>>(new cfftp<T0>(length));
+      }
+
+    template<typename T> POCKETFFT_NOINLINE void exec(cmplx<T> c[], T0 fct, bool fwd) const
+      { packplan ? packplan->exec(c,fct,fwd) : blueplan->exec(c,fct,fwd); }
+
+    size_t length() const { return len; }
+  };
+
+//
+// flexible (FFTPACK/Bluestein) real-valued 1D transform
+//
+
+template<typename T0> class pocketfft_r
+  {
+  private:
+    std::unique_ptr<rfftp<T0>> packplan;
+    std::unique_ptr<fftblue<T0>> blueplan;
+    size_t len;
+
+  public:
+    POCKETFFT_NOINLINE pocketfft_r(size_t length)
+      : len(length)
+      {
+      if (length==0) throw std::runtime_error("zero-length FFT requested");
+      size_t tmp = (length<50) ? 0 : util::largest_prime_factor(length);
+      if (tmp*tmp <= length)
+        {
+        packplan=std::unique_ptr<rfftp<T0>>(new rfftp<T0>(length));
+        return;
+        }
+      double comp1 = 0.5*util::cost_guess(length);
+      double comp2 = 2*util::cost_guess(util::good_size_cmplx(2*length-1));
+      comp2*=1.5; /* fudge factor that appears to give good overall performance */
+      if (comp2<comp1) // use Bluestein
+        blueplan=std::unique_ptr<fftblue<T0>>(new fftblue<T0>(length));
+      else
+        packplan=std::unique_ptr<rfftp<T0>>(new rfftp<T0>(length));
+      }
+
+    template<typename T> POCKETFFT_NOINLINE void exec(T c[], T0 fct, bool fwd) const
+      { packplan ? packplan->exec(c,fct,fwd) : blueplan->exec_r(c,fct,fwd); }
+
+    size_t length() const { return len; }
+  };
+
+
+//
+// sine/cosine transforms
+//
+
+template<typename T0> class T_dct1
+  {
+  private:
+    pocketfft_r<T0> fftplan;
+
+  public:
+    POCKETFFT_NOINLINE T_dct1(size_t length)
+      : fftplan(2*(length-1)) {}
+
+    template<typename T> POCKETFFT_NOINLINE void exec(T c[], T0 fct, bool ortho,
+      int /*type*/, bool /*cosine*/) const
+      {
+      constexpr T0 sqrt2=T0(1.414213562373095048801688724209698L);
+      size_t N=fftplan.length(), n=N/2+1;
+      if (ortho)
+        { c[0]*=sqrt2; c[n-1]*=sqrt2; }
+      arr<T> tmp(N);
+      tmp[0] = c[0];
+      for (size_t i=1; i<n; ++i)
+        tmp[i] = tmp[N-i] = c[i];
+      fftplan.exec(tmp.data(), fct, true);
+      c[0] = tmp[0];
+      for (size_t i=1; i<n; ++i)
+        c[i] = tmp[2*i-1];
+      if (ortho)
+        { c[0]*=sqrt2*T0(0.5); c[n-1]*=sqrt2*T0(0.5); }
+      }
+
+    size_t length() const { return fftplan.length()/2+1; }
+  };
+
+template<typename T0> class T_dst1
+  {
+  private:
+    pocketfft_r<T0> fftplan;
+
+  public:
+    POCKETFFT_NOINLINE T_dst1(size_t length)
+      : fftplan(2*(length+1)) {}
+
+    template<typename T> POCKETFFT_NOINLINE void exec(T c[], T0 fct,
+      bool /*ortho*/, int /*type*/, bool /*cosine*/) const
+      {
+      size_t N=fftplan.length(), n=N/2-1;
+      arr<T> tmp(N);
+      tmp[0] = tmp[n+1] = c[0]*0;
+      for (size_t i=0; i<n; ++i)
+        { tmp[i+1]=c[i]; tmp[N-1-i]=-c[i]; }
+      fftplan.exec(tmp.data(), fct, true);
+      for (size_t i=0; i<n; ++i)
+        c[i] = -tmp[2*i+2];
+      }
+
+    size_t length() const { return fftplan.length()/2-1; }
+  };
+
+template<typename T0> class T_dcst23
+  {
+  private:
+    pocketfft_r<T0> fftplan;
+    std::vector<T0> twiddle;
+
+  public:
+    POCKETFFT_NOINLINE T_dcst23(size_t length)
+      : fftplan(length), twiddle(length)
+      {
+      sincos_2pibyn<T0> tw(4*length);
+      for (size_t i=0; i<length; ++i)
+        twiddle[i] = tw[i+1].r;
+      }
+
+    template<typename T> POCKETFFT_NOINLINE void exec(T c[], T0 fct, bool ortho,
+      int type, bool cosine) const
+      {
+      constexpr T0 sqrt2=T0(1.414213562373095048801688724209698L);
+      size_t N=length();
+      size_t NS2 = (N+1)/2;
+      if (type==2)
+        {
+        if (!cosine)
+          for (size_t k=1; k<N; k+=2)
+            c[k] = -c[k];
+        c[0] *= 2;
+        if ((N&1)==0) c[N-1]*=2;
+        for (size_t k=1; k<N-1; k+=2)
+          MPINPLACE(c[k+1], c[k]);
+        fftplan.exec(c, fct, false);
+        for (size_t k=1, kc=N-1; k<NS2; ++k, --kc)
+          {
+          T t1 = twiddle[k-1]*c[kc]+twiddle[kc-1]*c[k];
+          T t2 = twiddle[k-1]*c[k]-twiddle[kc-1]*c[kc];
+          c[k] = T0(0.5)*(t1+t2); c[kc]=T0(0.5)*(t1-t2);
+          }
+        if ((N&1)==0)
+          c[NS2] *= twiddle[NS2-1];
+        if (!cosine)
+          for (size_t k=0, kc=N-1; k<kc; ++k, --kc)
+            std::swap(c[k], c[kc]);
+        if (ortho) c[0]*=sqrt2*T0(0.5);
+        }
+      else
+        {
+        if (ortho) c[0]*=sqrt2;
+        if (!cosine)
+          for (size_t k=0, kc=N-1; k<NS2; ++k, --kc)
+            std::swap(c[k], c[kc]);
+        for (size_t k=1, kc=N-1; k<NS2; ++k, --kc)
+          {
+          T t1=c[k]+c[kc], t2=c[k]-c[kc];
+          c[k] = twiddle[k-1]*t2+twiddle[kc-1]*t1;
+          c[kc]= twiddle[k-1]*t1-twiddle[kc-1]*t2;
+          }
+        if ((N&1)==0)
+          c[NS2] *= 2*twiddle[NS2-1];
+        fftplan.exec(c, fct, true);
+        for (size_t k=1; k<N-1; k+=2)
+          MPINPLACE(c[k], c[k+1]);
+        if (!cosine)
+          for (size_t k=1; k<N; k+=2)
+            c[k] = -c[k];
+        }
+      }
+
+    size_t length() const { return fftplan.length(); }
+  };
+
+template<typename T0> class T_dcst4
+  {
+  private:
+    size_t N;
+    std::unique_ptr<pocketfft_c<T0>> fft;
+    std::unique_ptr<pocketfft_r<T0>> rfft;
+    arr<cmplx<T0>> C2;
+
+  public:
+    POCKETFFT_NOINLINE T_dcst4(size_t length)
+      : N(length),
+        fft((N&1) ? nullptr : new pocketfft_c<T0>(N/2)),
+        rfft((N&1)? new pocketfft_r<T0>(N) : nullptr),
+        C2((N&1) ? 0 : N/2)
+      {
+      if ((N&1)==0)
+        {
+        sincos_2pibyn<T0> tw(16*N);
+        for (size_t i=0; i<N/2; ++i)
+          C2[i] = conj(tw[8*i+1]);
+        }
+      }
+
+    template<typename T> POCKETFFT_NOINLINE void exec(T c[], T0 fct,
+      bool /*ortho*/, int /*type*/, bool cosine) const
+      {
+      size_t n2 = N/2;
+      if (!cosine)
+        for (size_t k=0, kc=N-1; k<n2; ++k, --kc)
+          std::swap(c[k], c[kc]);
+      if (N&1)
+        {
+        // The following code is derived from the FFTW3 function apply_re11()
+        // and is released under the 3-clause BSD license with friendly
+        // permission of Matteo Frigo and Steven G. Johnson.
+
+        arr<T> y(N);
+        {
+        size_t i=0, m=n2;
+        for (; m<N; ++i, m+=4)
+          y[i] = c[m];
+        for (; m<2*N; ++i, m+=4)
+          y[i] = -c[2*N-m-1];
+        for (; m<3*N; ++i, m+=4)
+          y[i] = -c[m-2*N];
+        for (; m<4*N; ++i, m+=4)
+          y[i] = c[4*N-m-1];
+        for (; i<N; ++i, m+=4)
+          y[i] = c[m-4*N];
+        }
+        rfft->exec(y.data(), fct, true);
+        {
+        auto SGN = [](size_t i)
+           {
+           constexpr T0 sqrt2=T0(1.414213562373095048801688724209698L);
+           return (i&2) ? -sqrt2 : sqrt2;
+           };
+        c[n2] = y[0]*SGN(n2+1);
+        size_t i=0, i1=1, k=1;
+        for (; k<n2; ++i, ++i1, k+=2)
+          {
+          c[i    ] = y[2*k-1]*SGN(i1)     + y[2*k  ]*SGN(i);
+          c[N -i1] = y[2*k-1]*SGN(N -i)   - y[2*k  ]*SGN(N -i1);
+          c[n2-i1] = y[2*k+1]*SGN(n2-i)   - y[2*k+2]*SGN(n2-i1);
+          c[n2+i1] = y[2*k+1]*SGN(n2+i+2) + y[2*k+2]*SGN(n2+i1);
+          }
+        if (k == n2)
+          {
+          c[i   ] = y[2*k-1]*SGN(i+1) + y[2*k]*SGN(i);
+          c[N-i1] = y[2*k-1]*SGN(i+2) + y[2*k]*SGN(i1);
+          }
+        }
+
+        // FFTW-derived code ends here
+        }
+      else
+        {
+        // even length algorithm from
+        // https://www.appletonaudio.com/blog/2013/derivation-of-fast-dct-4-algorithm-based-on-dft/
+        arr<cmplx<T>> y(n2);
+        for(size_t i=0; i<n2; ++i)
+          {
+          y[i].Set(c[2*i],c[N-1-2*i]);
+          y[i] *= C2[i];
+          }
+        fft->exec(y.data(), fct, true);
+        for(size_t i=0, ic=n2-1; i<n2; ++i, --ic)
+          {
+          c[2*i  ] =  2*(y[i ].r*C2[i ].r-y[i ].i*C2[i ].i);
+          c[2*i+1] = -2*(y[ic].i*C2[ic].r+y[ic].r*C2[ic].i);
+          }
+        }
+      if (!cosine)
+        for (size_t k=1; k<N; k+=2)
+          c[k] = -c[k];
+      }
+
+    size_t length() const { return N; }
+  };
+
+
+//
+// multi-D infrastructure
+//
+
+template<typename T> std::shared_ptr<T> get_plan(size_t length)
+  {
+#if POCKETFFT_CACHE_SIZE==0
+  return std::make_shared<T>(length);
+#else
+  constexpr size_t nmax=POCKETFFT_CACHE_SIZE;
+  static std::array<std::shared_ptr<T>, nmax> cache;
+  static std::array<size_t, nmax> last_access{{0}};
+  static size_t access_counter = 0;
+  static std::mutex mut;
+
+  auto find_in_cache = [&]() -> std::shared_ptr<T>
+    {
+    for (size_t i=0; i<nmax; ++i)
+      if (cache[i] && (cache[i]->length()==length))
+        {
+        // no need to update if this is already the most recent entry
+        if (last_access[i]!=access_counter)
+          {
+          last_access[i] = ++access_counter;
+          // Guard against overflow
+          if (access_counter == 0)
+            last_access.fill(0);
+          }
+        return cache[i];
+        }
+
+    return nullptr;
+    };
+
+  {
+  std::lock_guard<std::mutex> lock(mut);
+  auto p = find_in_cache();
+  if (p) return p;
+  }
+  auto plan = std::make_shared<T>(length);
+  {
+  std::lock_guard<std::mutex> lock(mut);
+  auto p = find_in_cache();
+  if (p) return p;
+
+  size_t lru = 0;
+  for (size_t i=1; i<nmax; ++i)
+    if (last_access[i] < last_access[lru])
+      lru = i;
+
+  cache[lru] = plan;
+  last_access[lru] = ++access_counter;
+  }
+  return plan;
+#endif
+  }
+
+class arr_info
+  {
+  protected:
+    shape_t shp;
+    stride_t str;
+
+  public:
+    arr_info(const shape_t &shape_, const stride_t &stride_)
+      : shp(shape_), str(stride_) {}
+    size_t ndim() const { return shp.size(); }
+    size_t size() const { return util::prod(shp); }
+    const shape_t &shape() const { return shp; }
+    size_t shape(size_t i) const { return shp[i]; }
+    const stride_t &stride() const { return str; }
+    const ptrdiff_t &stride(size_t i) const { return str[i]; }
+  };
+
+template<typename T> class cndarr: public arr_info
+  {
+  protected:
+    const char *d;
+
+  public:
+    cndarr(const void *data_, const shape_t &shape_, const stride_t &stride_)
+      : arr_info(shape_, stride_),
+        d(reinterpret_cast<const char *>(data_)) {}
+    const T &operator[](ptrdiff_t ofs) const
+      { return *reinterpret_cast<const T *>(d+ofs); }
+  };
+
+template<typename T> class ndarr: public cndarr<T>
+  {
+  public:
+    ndarr(void *data_, const shape_t &shape_, const stride_t &stride_)
+      : cndarr<T>::cndarr(const_cast<const void *>(data_), shape_, stride_)
+      {}
+    T &operator[](ptrdiff_t ofs)
+      { return *reinterpret_cast<T *>(const_cast<char *>(cndarr<T>::d+ofs)); }
+  };
+
+template<size_t N> class multi_iter
+  {
+  private:
+    shape_t pos;
+    const arr_info &iarr, &oarr;
+    ptrdiff_t p_ii, p_i[N], str_i, p_oi, p_o[N], str_o;
+    size_t idim, rem;
+
+    void advance_i()
+      {
+      for (int i_=int(pos.size())-1; i_>=0; --i_)
+        {
+        auto i = size_t(i_);
+        if (i==idim) continue;
+        p_ii += iarr.stride(i);
+        p_oi += oarr.stride(i);
+        if (++pos[i] < iarr.shape(i))
+          return;
+        pos[i] = 0;
+        p_ii -= ptrdiff_t(iarr.shape(i))*iarr.stride(i);
+        p_oi -= ptrdiff_t(oarr.shape(i))*oarr.stride(i);
+        }
+      }
+
+  public:
+    multi_iter(const arr_info &iarr_, const arr_info &oarr_, size_t idim_)
+      : pos(iarr_.ndim(), 0), iarr(iarr_), oarr(oarr_), p_ii(0),
+        str_i(iarr.stride(idim_)), p_oi(0), str_o(oarr.stride(idim_)),
+        idim(idim_), rem(iarr.size()/iarr.shape(idim))
+      {
+      auto nshares = threading::num_threads();
+      if (nshares==1) return;
+      if (nshares==0) throw std::runtime_error("can't run with zero threads");
+      auto myshare = threading::thread_id();
+      if (myshare>=nshares) throw std::runtime_error("impossible share requested");
+      size_t nbase = rem/nshares;
+      size_t additional = rem%nshares;
+      size_t lo = myshare*nbase + ((myshare<additional) ? myshare : additional);
+      size_t hi = lo+nbase+(myshare<additional);
+      size_t todo = hi-lo;
+
+      size_t chunk = rem;
+      for (size_t i=0; i<pos.size(); ++i)
+        {
+        if (i==idim) continue;
+        chunk /= iarr.shape(i);
+        size_t n_advance = lo/chunk;
+        pos[i] += n_advance;
+        p_ii += ptrdiff_t(n_advance)*iarr.stride(i);
+        p_oi += ptrdiff_t(n_advance)*oarr.stride(i);
+        lo -= n_advance*chunk;
+        }
+      rem = todo;
+      }
+    void advance(size_t n)
+      {
+      if (rem<n) throw std::runtime_error("underrun");
+      for (size_t i=0; i<n; ++i)
+        {
+        p_i[i] = p_ii;
+        p_o[i] = p_oi;
+        advance_i();
+        }
+      rem -= n;
+      }
+    ptrdiff_t iofs(size_t i) const { return p_i[0] + ptrdiff_t(i)*str_i; }
+    ptrdiff_t iofs(size_t j, size_t i) const { return p_i[j] + ptrdiff_t(i)*str_i; }
+    ptrdiff_t oofs(size_t i) const { return p_o[0] + ptrdiff_t(i)*str_o; }
+    ptrdiff_t oofs(size_t j, size_t i) const { return p_o[j] + ptrdiff_t(i)*str_o; }
+    size_t length_in() const { return iarr.shape(idim); }
+    size_t length_out() const { return oarr.shape(idim); }
+    ptrdiff_t stride_in() const { return str_i; }
+    ptrdiff_t stride_out() const { return str_o; }
+    size_t remaining() const { return rem; }
+  };
+
+class simple_iter
+  {
+  private:
+    shape_t pos;
+    const arr_info &arr;
+    ptrdiff_t p;
+    size_t rem;
+
+  public:
+    simple_iter(const arr_info &arr_)
+      : pos(arr_.ndim(), 0), arr(arr_), p(0), rem(arr_.size()) {}
+    void advance()
+      {
+      --rem;
+      for (int i_=int(pos.size())-1; i_>=0; --i_)
+        {
+        auto i = size_t(i_);
+        p += arr.stride(i);
+        if (++pos[i] < arr.shape(i))
+          return;
+        pos[i] = 0;
+        p -= ptrdiff_t(arr.shape(i))*arr.stride(i);
+        }
+      }
+    ptrdiff_t ofs() const { return p; }
+    size_t remaining() const { return rem; }
+  };
+
+class rev_iter
+  {
+  private:
+    shape_t pos;
+    const arr_info &arr;
+    std::vector<char> rev_axis;
+    std::vector<char> rev_jump;
+    size_t last_axis, last_size;
+    shape_t shp;
+    ptrdiff_t p, rp;
+    size_t rem;
+
+  public:
+    rev_iter(const arr_info &arr_, const shape_t &axes)
+      : pos(arr_.ndim(), 0), arr(arr_), rev_axis(arr_.ndim(), 0),
+        rev_jump(arr_.ndim(), 1), p(0), rp(0)
+      {
+      for (auto ax: axes)
+        rev_axis[ax]=1;
+      last_axis = axes.back();
+      last_size = arr.shape(last_axis)/2 + 1;
+      shp = arr.shape();
+      shp[last_axis] = last_size;
+      rem=1;
+      for (auto i: shp)
+        rem *= i;
+      }
+    void advance()
+      {
+      --rem;
+      for (int i_=int(pos.size())-1; i_>=0; --i_)
+        {
+        auto i = size_t(i_);
+        p += arr.stride(i);
+        if (!rev_axis[i])
+          rp += arr.stride(i);
+        else
+          {
+          rp -= arr.stride(i);
+          if (rev_jump[i])
+            {
+            rp += ptrdiff_t(arr.shape(i))*arr.stride(i);
+            rev_jump[i] = 0;
+            }
+          }
+        if (++pos[i] < shp[i])
+          return;
+        pos[i] = 0;
+        p -= ptrdiff_t(shp[i])*arr.stride(i);
+        if (rev_axis[i])
+          {
+          rp -= ptrdiff_t(arr.shape(i)-shp[i])*arr.stride(i);
+          rev_jump[i] = 1;
+          }
+        else
+          rp -= ptrdiff_t(shp[i])*arr.stride(i);
+        }
+      }
+    ptrdiff_t ofs() const { return p; }
+    ptrdiff_t rev_ofs() const { return rp; }
+    size_t remaining() const { return rem; }
+  };
+
+template<typename T> struct VTYPE {};
+template <typename T> using vtype_t = typename VTYPE<T>::type;
+
+#ifndef POCKETFFT_NO_VECTORS
+template<> struct VTYPE<float>
+  {
+  using type = float __attribute__ ((vector_size (VLEN<float>::val*sizeof(float))));
+  };
+template<> struct VTYPE<double>
+  {
+  using type = double __attribute__ ((vector_size (VLEN<double>::val*sizeof(double))));
+  };
+template<> struct VTYPE<long double>
+  {
+  using type = long double __attribute__ ((vector_size (VLEN<long double>::val*sizeof(long double))));
+  };
+#endif
+
+template<typename T> arr<char> alloc_tmp(const shape_t &shape,
+  size_t axsize, size_t elemsize)
+  {
+  auto othersize = util::prod(shape)/axsize;
+  auto tmpsize = axsize*((othersize>=VLEN<T>::val) ? VLEN<T>::val : 1);
+  return arr<char>(tmpsize*elemsize);
+  }
+template<typename T> arr<char> alloc_tmp(const shape_t &shape,
+  const shape_t &axes, size_t elemsize)
+  {
+  size_t fullsize=util::prod(shape);
+  size_t tmpsize=0;
+  for (size_t i=0; i<axes.size(); ++i)
+    {
+    auto axsize = shape[axes[i]];
+    auto othersize = fullsize/axsize;
+    auto sz = axsize*((othersize>=VLEN<T>::val) ? VLEN<T>::val : 1);
+    if (sz>tmpsize) tmpsize=sz;
+    }
+  return arr<char>(tmpsize*elemsize);
+  }
+
+template <typename T, size_t vlen> void copy_input(const multi_iter<vlen> &it,
+  const cndarr<cmplx<T>> &src, cmplx<vtype_t<T>> *POCKETFFT_RESTRICT dst)
+  {
+  for (size_t i=0; i<it.length_in(); ++i)
+    for (size_t j=0; j<vlen; ++j)
+      {
+      dst[i].r[j] = src[it.iofs(j,i)].r;
+      dst[i].i[j] = src[it.iofs(j,i)].i;
+      }
+  }
+
+template <typename T, size_t vlen> void copy_input(const multi_iter<vlen> &it,
+  const cndarr<T> &src, vtype_t<T> *POCKETFFT_RESTRICT dst)
+  {
+  for (size_t i=0; i<it.length_in(); ++i)
+    for (size_t j=0; j<vlen; ++j)
+      dst[i][j] = src[it.iofs(j,i)];
+  }
+
+template <typename T, size_t vlen> void copy_input(const multi_iter<vlen> &it,
+  const cndarr<T> &src, T *POCKETFFT_RESTRICT dst)
+  {
+  if (dst == &src[it.iofs(0)]) return;  // in-place
+  for (size_t i=0; i<it.length_in(); ++i)
+    dst[i] = src[it.iofs(i)];
+  }
+
+template<typename T, size_t vlen> void copy_output(const multi_iter<vlen> &it,
+  const cmplx<vtype_t<T>> *POCKETFFT_RESTRICT src, ndarr<cmplx<T>> &dst)
+  {
+  for (size_t i=0; i<it.length_out(); ++i)
+    for (size_t j=0; j<vlen; ++j)
+      dst[it.oofs(j,i)].Set(src[i].r[j],src[i].i[j]);
+  }
+
+template<typename T, size_t vlen> void copy_output(const multi_iter<vlen> &it,
+  const vtype_t<T> *POCKETFFT_RESTRICT src, ndarr<T> &dst)
+  {
+  for (size_t i=0; i<it.length_out(); ++i)
+    for (size_t j=0; j<vlen; ++j)
+      dst[it.oofs(j,i)] = src[i][j];
+  }
+
+template<typename T, size_t vlen> void copy_output(const multi_iter<vlen> &it,
+  const T *POCKETFFT_RESTRICT src, ndarr<T> &dst)
+  {
+  if (src == &dst[it.oofs(0)]) return;  // in-place
+  for (size_t i=0; i<it.length_out(); ++i)
+    dst[it.oofs(i)] = src[i];
+  }
+
+template <typename T> struct add_vec { using type = vtype_t<T>; };
+template <typename T> struct add_vec<cmplx<T>>
+  { using type = cmplx<vtype_t<T>>; };
+template <typename T> using add_vec_t = typename add_vec<T>::type;
+
+template<typename Tplan, typename T, typename T0, typename Exec>
+POCKETFFT_NOINLINE void general_nd(const cndarr<T> &in, ndarr<T> &out,
+  const shape_t &axes, T0 fct, size_t nthreads, const Exec & exec,
+  const bool allow_inplace=true)
+  {
+  std::shared_ptr<Tplan> plan;
+
+  for (size_t iax=0; iax<axes.size(); ++iax)
+    {
+    size_t len=in.shape(axes[iax]);
+    if ((!plan) || (len!=plan->length()))
+      plan = get_plan<Tplan>(len);
+
+    threading::thread_map(
+      util::thread_count(nthreads, in.shape(), axes[iax], VLEN<T>::val),
+      [&] {
+        constexpr auto vlen = VLEN<T0>::val;
+        auto storage = alloc_tmp<T0>(in.shape(), len, sizeof(T));
+        const auto &tin(iax==0? in : out);
+        multi_iter<vlen> it(tin, out, axes[iax]);
+#ifndef POCKETFFT_NO_VECTORS
+        if (vlen>1)
+          while (it.remaining()>=vlen)
+            {
+            it.advance(vlen);
+            auto tdatav = reinterpret_cast<add_vec_t<T> *>(storage.data());
+            exec(it, tin, out, tdatav, *plan, fct);
+            }
+#endif
+        while (it.remaining()>0)
+          {
+          it.advance(1);
+          auto buf = allow_inplace && it.stride_out() == sizeof(T) ?
+            &out[it.oofs(0)] : reinterpret_cast<T *>(storage.data());
+          exec(it, tin, out, buf, *plan, fct);
+          }
+      });  // end of parallel region
+    fct = T0(1); // factor has been applied, use 1 for remaining axes
+    }
+  }
+
+struct ExecC2C
+  {
+  bool forward;
+
+  template <typename T0, typename T, size_t vlen> void operator () (
+    const multi_iter<vlen> &it, const cndarr<cmplx<T0>> &in,
+    ndarr<cmplx<T0>> &out, T * buf, const pocketfft_c<T0> &plan, T0 fct) const
+    {
+    copy_input(it, in, buf);
+    plan.exec(buf, fct, forward);
+    copy_output(it, buf, out);
+    }
+  };
+
+template <typename T, size_t vlen> void copy_hartley(const multi_iter<vlen> &it,
+  const vtype_t<T> *POCKETFFT_RESTRICT src, ndarr<T> &dst)
+  {
+  for (size_t j=0; j<vlen; ++j)
+    dst[it.oofs(j,0)] = src[0][j];
+  size_t i=1, i1=1, i2=it.length_out()-1;
+  for (i=1; i<it.length_out()-1; i+=2, ++i1, --i2)
+    for (size_t j=0; j<vlen; ++j)
+      {
+        dst[it.oofs(j,i1)] = src[i][j]+src[i+1][j];
+        dst[it.oofs(j,i2)] = src[i][j]-src[i+1][j];
+      }
+  if (i<it.length_out())
+    for (size_t j=0; j<vlen; ++j)
+      dst[it.oofs(j,i1)] = src[i][j];
+  }
+
+template <typename T, size_t vlen> void copy_hartley(const multi_iter<vlen> &it,
+  const T *POCKETFFT_RESTRICT src, ndarr<T> &dst)
+  {
+  dst[it.oofs(0)] = src[0];
+  size_t i=1, i1=1, i2=it.length_out()-1;
+  for (i=1; i<it.length_out()-1; i+=2, ++i1, --i2)
+    {
+    dst[it.oofs(i1)] = src[i]+src[i+1];
+    dst[it.oofs(i2)] = src[i]-src[i+1];
+    }
+  if (i<it.length_out())
+    dst[it.oofs(i1)] = src[i];
+  }
+
+struct ExecHartley
+  {
+  template <typename T0, typename T, size_t vlen> void operator () (
+    const multi_iter<vlen> &it, const cndarr<T0> &in, ndarr<T0> &out,
+    T * buf, const pocketfft_r<T0> &plan, T0 fct) const
+    {
+    copy_input(it, in, buf);
+    plan.exec(buf, fct, true);
+    copy_hartley(it, buf, out);
+    }
+  };
+
+struct ExecDcst
+  {
+  bool ortho;
+  int type;
+  bool cosine;
+
+  template <typename T0, typename T, typename Tplan, size_t vlen>
+  void operator () (const multi_iter<vlen> &it, const cndarr<T0> &in,
+    ndarr<T0> &out, T * buf, const Tplan &plan, T0 fct) const
+    {
+    copy_input(it, in, buf);
+    plan.exec(buf, fct, ortho, type, cosine);
+    copy_output(it, buf, out);
+    }
+  };
+
+template<typename T> POCKETFFT_NOINLINE void general_r2c(
+  const cndarr<T> &in, ndarr<cmplx<T>> &out, size_t axis, bool forward, T fct,
+  size_t nthreads)
+  {
+  auto plan = get_plan<pocketfft_r<T>>(in.shape(axis));
+  size_t len=in.shape(axis);
+  threading::thread_map(
+    util::thread_count(nthreads, in.shape(), axis, VLEN<T>::val),
+    [&] {
+    constexpr auto vlen = VLEN<T>::val;
+    auto storage = alloc_tmp<T>(in.shape(), len, sizeof(T));
+    multi_iter<vlen> it(in, out, axis);
+#ifndef POCKETFFT_NO_VECTORS
+    if (vlen>1)
+      while (it.remaining()>=vlen)
+        {
+        it.advance(vlen);
+        auto tdatav = reinterpret_cast<vtype_t<T> *>(storage.data());
+        copy_input(it, in, tdatav);
+        plan->exec(tdatav, fct, true);
+        for (size_t j=0; j<vlen; ++j)
+          out[it.oofs(j,0)].Set(tdatav[0][j]);
+        size_t i=1, ii=1;
+        if (forward)
+          for (; i<len-1; i+=2, ++ii)
+            for (size_t j=0; j<vlen; ++j)
+              out[it.oofs(j,ii)].Set(tdatav[i][j], tdatav[i+1][j]);
+        else
+          for (; i<len-1; i+=2, ++ii)
+            for (size_t j=0; j<vlen; ++j)
+              out[it.oofs(j,ii)].Set(tdatav[i][j], -tdatav[i+1][j]);
+        if (i<len)
+          for (size_t j=0; j<vlen; ++j)
+            out[it.oofs(j,ii)].Set(tdatav[i][j]);
+        }
+#endif
+    while (it.remaining()>0)
+      {
+      it.advance(1);
+      auto tdata = reinterpret_cast<T *>(storage.data());
+      copy_input(it, in, tdata);
+      plan->exec(tdata, fct, true);
+      out[it.oofs(0)].Set(tdata[0]);
+      size_t i=1, ii=1;
+      if (forward)
+        for (; i<len-1; i+=2, ++ii)
+          out[it.oofs(ii)].Set(tdata[i], tdata[i+1]);
+      else
+        for (; i<len-1; i+=2, ++ii)
+          out[it.oofs(ii)].Set(tdata[i], -tdata[i+1]);
+      if (i<len)
+        out[it.oofs(ii)].Set(tdata[i]);
+      }
+    });  // end of parallel region
+  }
+template<typename T> POCKETFFT_NOINLINE void general_c2r(
+  const cndarr<cmplx<T>> &in, ndarr<T> &out, size_t axis, bool forward, T fct,
+  size_t nthreads)
+  {
+  auto plan = get_plan<pocketfft_r<T>>(out.shape(axis));
+  size_t len=out.shape(axis);
+  threading::thread_map(
+    util::thread_count(nthreads, in.shape(), axis, VLEN<T>::val),
+    [&] {
+      constexpr auto vlen = VLEN<T>::val;
+      auto storage = alloc_tmp<T>(out.shape(), len, sizeof(T));
+      multi_iter<vlen> it(in, out, axis);
+#ifndef POCKETFFT_NO_VECTORS
+      if (vlen>1)
+        while (it.remaining()>=vlen)
+          {
+          it.advance(vlen);
+          auto tdatav = reinterpret_cast<vtype_t<T> *>(storage.data());
+          for (size_t j=0; j<vlen; ++j)
+            tdatav[0][j]=in[it.iofs(j,0)].r;
+          {
+          size_t i=1, ii=1;
+          if (forward)
+            for (; i<len-1; i+=2, ++ii)
+              for (size_t j=0; j<vlen; ++j)
+                {
+                tdatav[i  ][j] =  in[it.iofs(j,ii)].r;
+                tdatav[i+1][j] = -in[it.iofs(j,ii)].i;
+                }
+          else
+            for (; i<len-1; i+=2, ++ii)
+              for (size_t j=0; j<vlen; ++j)
+                {
+                tdatav[i  ][j] = in[it.iofs(j,ii)].r;
+                tdatav[i+1][j] = in[it.iofs(j,ii)].i;
+                }
+          if (i<len)
+            for (size_t j=0; j<vlen; ++j)
+              tdatav[i][j] = in[it.iofs(j,ii)].r;
+          }
+          plan->exec(tdatav, fct, false);
+          copy_output(it, tdatav, out);
+          }
+#endif
+      while (it.remaining()>0)
+        {
+        it.advance(1);
+        auto tdata = reinterpret_cast<T *>(storage.data());
+        tdata[0]=in[it.iofs(0)].r;
+        {
+        size_t i=1, ii=1;
+        if (forward)
+          for (; i<len-1; i+=2, ++ii)
+            {
+            tdata[i  ] =  in[it.iofs(ii)].r;
+            tdata[i+1] = -in[it.iofs(ii)].i;
+            }
+        else
+          for (; i<len-1; i+=2, ++ii)
+            {
+            tdata[i  ] = in[it.iofs(ii)].r;
+            tdata[i+1] = in[it.iofs(ii)].i;
+            }
+        if (i<len)
+          tdata[i] = in[it.iofs(ii)].r;
+        }
+        plan->exec(tdata, fct, false);
+        copy_output(it, tdata, out);
+        }
+    });  // end of parallel region
+  }
+
+struct ExecR2R
+  {
+  bool r2h, forward;
+
+  template <typename T0, typename T, size_t vlen> void operator () (
+    const multi_iter<vlen> &it, const cndarr<T0> &in, ndarr<T0> &out, T * buf,
+    const pocketfft_r<T0> &plan, T0 fct) const
+    {
+    copy_input(it, in, buf);
+    if ((!r2h) && forward)
+      for (size_t i=2; i<it.length_out(); i+=2)
+        buf[i] = -buf[i];
+    plan.exec(buf, fct, r2h);
+    if (r2h && (!forward))
+      for (size_t i=2; i<it.length_out(); i+=2)
+        buf[i] = -buf[i];
+    copy_output(it, buf, out);
+    }
+  };
+
+template<typename T> void c2c(const shape_t &shape, const stride_t &stride_in,
+  const stride_t &stride_out, const shape_t &axes, bool forward,
+  const std::complex<T> *data_in, std::complex<T> *data_out, T fct,
+  size_t nthreads=1)
+  {
+  if (util::prod(shape)==0) return;
+  util::sanity_check(shape, stride_in, stride_out, data_in==data_out, axes);
+  cndarr<cmplx<T>> ain(data_in, shape, stride_in);
+  ndarr<cmplx<T>> aout(data_out, shape, stride_out);
+  general_nd<pocketfft_c<T>>(ain, aout, axes, fct, nthreads, ExecC2C{forward});
+  }
+
+template<typename T> void dct(const shape_t &shape,
+  const stride_t &stride_in, const stride_t &stride_out, const shape_t &axes,
+  int type, const T *data_in, T *data_out, T fct, bool ortho, size_t nthreads=1)
+  {
+  if ((type<1) || (type>4)) throw std::invalid_argument("invalid DCT type");
+  if (util::prod(shape)==0) return;
+  util::sanity_check(shape, stride_in, stride_out, data_in==data_out, axes);
+  cndarr<T> ain(data_in, shape, stride_in);
+  ndarr<T> aout(data_out, shape, stride_out);
+  const ExecDcst exec{ortho, type, true};
+  if (type==1)
+    general_nd<T_dct1<T>>(ain, aout, axes, fct, nthreads, exec);
+  else if (type==4)
+    general_nd<T_dcst4<T>>(ain, aout, axes, fct, nthreads, exec);
+  else
+    general_nd<T_dcst23<T>>(ain, aout, axes, fct, nthreads, exec);
+  }
+
+template<typename T> void dst(const shape_t &shape,
+  const stride_t &stride_in, const stride_t &stride_out, const shape_t &axes,
+  int type, const T *data_in, T *data_out, T fct, bool ortho, size_t nthreads=1)
+  {
+  if ((type<1) || (type>4)) throw std::invalid_argument("invalid DST type");
+  if (util::prod(shape)==0) return;
+  util::sanity_check(shape, stride_in, stride_out, data_in==data_out, axes);
+  cndarr<T> ain(data_in, shape, stride_in);
+  ndarr<T> aout(data_out, shape, stride_out);
+  const ExecDcst exec{ortho, type, false};
+  if (type==1)
+    general_nd<T_dst1<T>>(ain, aout, axes, fct, nthreads, exec);
+  else if (type==4)
+    general_nd<T_dcst4<T>>(ain, aout, axes, fct, nthreads, exec);
+  else
+    general_nd<T_dcst23<T>>(ain, aout, axes, fct, nthreads, exec);
+  }
+
+template<typename T> void r2c(const shape_t &shape_in,
+  const stride_t &stride_in, const stride_t &stride_out, size_t axis,
+  bool forward, const T *data_in, std::complex<T> *data_out, T fct,
+  size_t nthreads=1)
+  {
+  if (util::prod(shape_in)==0) return;
+  util::sanity_check(shape_in, stride_in, stride_out, false, axis);
+  cndarr<T> ain(data_in, shape_in, stride_in);
+  shape_t shape_out(shape_in);
+  shape_out[axis] = shape_in[axis]/2 + 1;
+  ndarr<cmplx<T>> aout(data_out, shape_out, stride_out);
+  general_r2c(ain, aout, axis, forward, fct, nthreads);
+  }
+
+template<typename T> void r2c(const shape_t &shape_in,
+  const stride_t &stride_in, const stride_t &stride_out, const shape_t &axes,
+  bool forward, const T *data_in, std::complex<T> *data_out, T fct,
+  size_t nthreads=1)
+  {
+  if (util::prod(shape_in)==0) return;
+  util::sanity_check(shape_in, stride_in, stride_out, false, axes);
+  r2c(shape_in, stride_in, stride_out, axes.back(), forward, data_in, data_out,
+    fct, nthreads);
+  if (axes.size()==1) return;
+
+  shape_t shape_out(shape_in);
+  shape_out[axes.back()] = shape_in[axes.back()]/2 + 1;
+  auto newaxes = shape_t{axes.begin(), --axes.end()};
+  c2c(shape_out, stride_out, stride_out, newaxes, forward, data_out, data_out,
+    T(1), nthreads);
+  }
+
+template<typename T> void c2r(const shape_t &shape_out,
+  const stride_t &stride_in, const stride_t &stride_out, size_t axis,
+  bool forward, const std::complex<T> *data_in, T *data_out, T fct,
+  size_t nthreads=1)
+  {
+  if (util::prod(shape_out)==0) return;
+  util::sanity_check(shape_out, stride_in, stride_out, false, axis);
+  shape_t shape_in(shape_out);
+  shape_in[axis] = shape_out[axis]/2 + 1;
+  cndarr<cmplx<T>> ain(data_in, shape_in, stride_in);
+  ndarr<T> aout(data_out, shape_out, stride_out);
+  general_c2r(ain, aout, axis, forward, fct, nthreads);
+  }
+
+template<typename T> void c2r(const shape_t &shape_out,
+  const stride_t &stride_in, const stride_t &stride_out, const shape_t &axes,
+  bool forward, const std::complex<T> *data_in, T *data_out, T fct,
+  size_t nthreads=1)
+  {
+  if (util::prod(shape_out)==0) return;
+  if (axes.size()==1)
+    return c2r(shape_out, stride_in, stride_out, axes[0], forward,
+      data_in, data_out, fct, nthreads);
+  util::sanity_check(shape_out, stride_in, stride_out, false, axes);
+  auto shape_in = shape_out;
+  shape_in[axes.back()] = shape_out[axes.back()]/2 + 1;
+  auto nval = util::prod(shape_in);
+  stride_t stride_inter(shape_in.size());
+  stride_inter.back() = sizeof(cmplx<T>);
+  for (int i=int(shape_in.size())-2; i>=0; --i)
+    stride_inter[size_t(i)] =
+      stride_inter[size_t(i+1)]*ptrdiff_t(shape_in[size_t(i+1)]);
+  arr<std::complex<T>> tmp(nval);
+  auto newaxes = shape_t{axes.begin(), --axes.end()};
+  c2c(shape_in, stride_in, stride_inter, newaxes, forward, data_in, tmp.data(),
+    T(1), nthreads);
+  c2r(shape_out, stride_inter, stride_out, axes.back(), forward,
+    tmp.data(), data_out, fct, nthreads);
+  }
+
+template<typename T> void r2r_fftpack(const shape_t &shape,
+  const stride_t &stride_in, const stride_t &stride_out, const shape_t &axes,
+  bool real2hermitian, bool forward, const T *data_in, T *data_out, T fct,
+  size_t nthreads=1)
+  {
+  if (util::prod(shape)==0) return;
+  util::sanity_check(shape, stride_in, stride_out, data_in==data_out, axes);
+  cndarr<T> ain(data_in, shape, stride_in);
+  ndarr<T> aout(data_out, shape, stride_out);
+  general_nd<pocketfft_r<T>>(ain, aout, axes, fct, nthreads,
+    ExecR2R{real2hermitian, forward});
+  }
+
+template<typename T> void r2r_separable_hartley(const shape_t &shape,
+  const stride_t &stride_in, const stride_t &stride_out, const shape_t &axes,
+  const T *data_in, T *data_out, T fct, size_t nthreads=1)
+  {
+  if (util::prod(shape)==0) return;
+  util::sanity_check(shape, stride_in, stride_out, data_in==data_out, axes);
+  cndarr<T> ain(data_in, shape, stride_in);
+  ndarr<T> aout(data_out, shape, stride_out);
+  general_nd<pocketfft_r<T>>(ain, aout, axes, fct, nthreads, ExecHartley{},
+    false);
+  }
+
+template<typename T> void r2r_genuine_hartley(const shape_t &shape,
+  const stride_t &stride_in, const stride_t &stride_out, const shape_t &axes,
+  const T *data_in, T *data_out, T fct, size_t nthreads=1)
+  {
+  if (util::prod(shape)==0) return;
+  if (axes.size()==1)
+    return r2r_separable_hartley(shape, stride_in, stride_out, axes, data_in,
+      data_out, fct, nthreads);
+  util::sanity_check(shape, stride_in, stride_out, data_in==data_out, axes);
+  shape_t tshp(shape);
+  tshp[axes.back()] = tshp[axes.back()]/2+1;
+  arr<std::complex<T>> tdata(util::prod(tshp));
+  stride_t tstride(shape.size());
+  tstride.back()=sizeof(std::complex<T>);
+  for (size_t i=tstride.size()-1; i>0; --i)
+    tstride[i-1]=tstride[i]*ptrdiff_t(tshp[i]);
+  r2c(shape, stride_in, tstride, axes, true, data_in, tdata.data(), fct, nthreads);
+  cndarr<cmplx<T>> atmp(tdata.data(), tshp, tstride);
+  ndarr<T> aout(data_out, shape, stride_out);
+  simple_iter iin(atmp);
+  rev_iter iout(aout, axes);
+  while(iin.remaining()>0)
+    {
+    auto v = atmp[iin.ofs()];
+    aout[iout.ofs()] = v.r+v.i;
+    aout[iout.rev_ofs()] = v.r-v.i;
+    iin.advance(); iout.advance();
+    }
+  }
+
+} // namespace detail
+
+using detail::FORWARD;
+using detail::BACKWARD;
+using detail::shape_t;
+using detail::stride_t;
+using detail::c2c;
+using detail::c2r;
+using detail::r2c;
+using detail::r2r_fftpack;
+using detail::r2r_separable_hartley;
+using detail::r2r_genuine_hartley;
+using detail::dct;
+using detail::dst;
+
+} // namespace pocketfft
+
+#undef POCKETFFT_NOINLINE
+#undef POCKETFFT_RESTRICT
+
+#endif // POCKETFFT_HDRONLY_H
diff --git a/macos/third_party/mlx/mlx/CMakeLists.txt b/macos/third_party/mlx/mlx/CMakeLists.txt
new file mode 100644
index 0000000..9c0fd38
--- /dev/null
+++ b/macos/third_party/mlx/mlx/CMakeLists.txt
@@ -0,0 +1,107 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/dtype_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/export.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/einsum.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fast.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/graph_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/random.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scheduler.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/metal.h)
+
+# Define MLX_VERSION only in the version.cpp file.
+add_library(mlx_version OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp)
+target_compile_definitions(mlx_version PRIVATE MLX_VERSION="${MLX_VERSION}")
+target_include_directories(mlx_version PRIVATE ${PROJECT_SOURCE_DIR})
+target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:mlx_version>)
+
+# Do not export symbols by default.
+set_target_properties(
+  mlx mlx_version
+  PROPERTIES VISIBILITY_INLINES_HIDDEN ON
+             CXX_VISIBILITY_PRESET hidden
+             CUDA_VISIBILITY_PRESET hidden)
+
+# Define MLX_EXPORT for shared libraries, MLX_STATIC for static libraries.
+set_target_properties(mlx PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
+if(BUILD_SHARED_LIBS)
+  target_compile_definitions(mlx_version PUBLIC MLX_EXPORT)
+else()
+  target_compile_definitions(mlx PUBLIC MLX_STATIC)
+  target_compile_definitions(mlx_version PUBLIC MLX_STATIC)
+endif()
+
+if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+  # Supress warnings: note: parameter passing for argument of type
+  # 'std::pair<float, float>' when C++17 is enabled changed to match C++14 in
+  # GCC 10.1
+  target_compile_options(mlx PRIVATE -Wno-psabi)
+endif()
+
+if(MSVC)
+  # Some of CUDA's headers include windows.h, which defines min/max macros.
+  target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN)
+  # Unicode support in fmt does not compile in .cu files.
+  target_compile_definitions(mlx PRIVATE FMT_UNICODE=0)
+  # Disable some MSVC warnings to speed up compilation.
+  target_compile_options(
+    mlx
+    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4244 /wd4267>
+    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/wd4068
+            /wd4146
+            /wd4700
+            /wd4804
+            /wd4805>
+            $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4244
+            -Xcompiler=/wd4267>)
+  # Enable /bigobj for heavily templated code (e.g., binary.cpp) that exceeds
+  # the default 65,535 section limit in COFF object files.
+  target_compile_options(
+    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>
+                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/bigobj>)
+  # Use modern preprocessor, otherwise CCCL would complain.
+  target_compile_options(
+    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>
+                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/Zc:preprocessor>)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
+
+if(MLX_BUILD_CPU)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cpu)
+else()
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_cpu)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/distributed)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
+
+if(MLX_BUILD_METAL)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/metal)
+else()
+  target_sources(mlx
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/no_metal.cpp)
+endif()
+
+if(MLX_BUILD_CUDA)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda)
+else()
+  target_sources(mlx
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda/no_cuda.cpp)
+endif()
+
+if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/gpu)
+else()
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_gpu)
+endif()
diff --git a/macos/third_party/mlx/mlx/allocator.h b/macos/third_party/mlx/mlx/allocator.h
new file mode 100644
index 0000000..824deac
--- /dev/null
+++ b/macos/third_party/mlx/mlx/allocator.h
@@ -0,0 +1,75 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "mlx/api.h"
+
+namespace mlx::core::allocator {
+
+// Simple wrapper around buffer pointers
+// WARNING: Only Buffer objects constructed from and those that wrap
+//          raw pointers from mlx::allocator are supported.
+class MLX_API Buffer {
+ private:
+  void* ptr_;
+
+ public:
+  explicit Buffer(void* ptr) : ptr_(ptr) {};
+
+  // Get the raw data pointer from the buffer
+  void* raw_ptr();
+
+  // Get the buffer pointer from the buffer
+  const void* ptr() const {
+    return ptr_;
+  };
+  void* ptr() {
+    return ptr_;
+  };
+};
+
+class MLX_API Allocator {
+  /** Abstract base class for a memory allocator. */
+ public:
+  virtual Buffer malloc(size_t size) = 0;
+  virtual void free(Buffer buffer) = 0;
+  virtual size_t size(Buffer buffer) const = 0;
+  virtual Buffer make_buffer(void* ptr, size_t size) {
+    return Buffer{nullptr};
+  };
+  virtual void release(Buffer buffer) {}
+
+  Allocator() = default;
+  Allocator(const Allocator& other) = delete;
+  Allocator(Allocator&& other) = delete;
+  Allocator& operator=(const Allocator& other) = delete;
+  Allocator& operator=(Allocator&& other) = delete;
+  virtual ~Allocator() = default;
+};
+
+MLX_API Allocator& allocator();
+
+inline Buffer malloc(size_t size) {
+  return allocator().malloc(size);
+}
+
+inline void free(Buffer buffer) {
+  allocator().free(buffer);
+}
+
+// Make a Buffer from a raw pointer of the given size without a copy.  If a
+// no-copy conversion is not possible then the returned buffer.ptr() will be
+// nullptr. Any buffer created with this function must be released with
+// release(buffer)
+inline Buffer make_buffer(void* ptr, size_t size) {
+  return allocator().make_buffer(ptr, size);
+};
+
+// Release a buffer from the allocator made with make_buffer
+inline void release(Buffer buffer) {
+  allocator().release(buffer);
+}
+
+} // namespace mlx::core::allocator
diff --git a/macos/third_party/mlx/mlx/api.h b/macos/third_party/mlx/mlx/api.h
new file mode 100644
index 0000000..8aed091
--- /dev/null
+++ b/macos/third_party/mlx/mlx/api.h
@@ -0,0 +1,29 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+// MLX_API macro for controlling symbol visibility, must add for public APIs.
+//
+// Usage:
+//   MLX_API void some_function(...);
+//   class MLX_API SomeClass { ... };
+
+#if defined(MLX_STATIC)
+
+// Static library build - no import/export decorations needed
+#define MLX_API
+
+#else
+
+// Shared library build.
+#if defined(_WIN32)
+#if defined(MLX_EXPORT)
+#define MLX_API __declspec(dllexport)
+#else
+#define MLX_API __declspec(dllimport)
+#endif // defined(MLX_EXPORT)
+#else
+#define MLX_API __attribute__((visibility("default")))
+#endif // defined(_WIN32)
+
+#endif // defined(MLX_STATIC)
diff --git a/macos/third_party/mlx/mlx/array.cpp b/macos/third_party/mlx/mlx/array.cpp
new file mode 100644
index 0000000..1769f0f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/array.cpp
@@ -0,0 +1,354 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <functional>
+#include <unordered_map>
+
+#include "mlx/array.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/transforms_impl.h"
+
+namespace mlx::core {
+
+array::array(const std::complex<float>& val, Dtype dtype /* = complex64 */)
+    : array_desc_(std::make_shared<ArrayDesc>(Shape{}, dtype)) {
+  auto cval = static_cast<complex64_t>(val);
+  init(&cval);
+}
+
+array::array(
+    Shape shape,
+    Dtype dtype,
+    std::shared_ptr<Primitive> primitive,
+    std::vector<array> inputs)
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              std::move(shape),
+              dtype,
+              std::move(primitive),
+              std::move(inputs))) {
+  if (has_primitive() && this->primitive().stream().device == Device::gpu) {
+    for (auto& in : this->inputs()) {
+      if (in.dtype() == float64) {
+        throw std::invalid_argument("float64 is not supported on the GPU");
+      }
+    }
+    if (this->dtype() == float64) {
+      throw std::invalid_argument("float64 is not supported on the GPU");
+    }
+  }
+}
+
+std::vector<array> array::make_arrays(
+    std::vector<Shape> shapes,
+    const std::vector<Dtype>& dtypes,
+    const std::shared_ptr<Primitive>& primitive,
+    const std::vector<array>& inputs) {
+  std::vector<array> outputs;
+  for (size_t i = 0; i < shapes.size(); ++i) {
+    outputs.emplace_back(std::move(shapes[i]), dtypes[i], primitive, inputs);
+  }
+  // For each node in |outputs|, its siblings are the other nodes.
+  for (size_t i = 0; i < outputs.size(); ++i) {
+    auto siblings = outputs;
+    siblings.erase(siblings.begin() + i);
+    outputs[i].set_siblings(std::move(siblings), i);
+  }
+  return outputs;
+}
+
+array array::unsafe_weak_copy(const array& other) {
+  auto cpy = array(other.shape(), other.dtype(), nullptr, {});
+  cpy.set_data(
+      other.buffer(),
+      other.data_size(),
+      other.strides(),
+      other.flags(),
+      [](auto) {});
+  cpy.array_desc_->offset = other.array_desc_->offset;
+  return cpy;
+}
+
+array::array(std::initializer_list<float> data)
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              float32)) {
+  init(data.begin());
+}
+
+array::array(std::initializer_list<int> data, Dtype dtype)
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              dtype)) {
+  init(data.begin());
+}
+
+array::array(
+    void* data,
+    Shape shape,
+    Dtype dtype,
+    const std::function<void(void*)>& deleter)
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+  auto buffer = allocator::make_buffer(data, nbytes());
+  if (buffer.ptr() == nullptr) {
+    set_data(allocator::malloc(nbytes()));
+    auto ptr = static_cast<char*>(data);
+    std::copy(ptr, ptr + nbytes(), this->data<char>());
+    deleter(data);
+  } else {
+    auto wrapped_deleter = [deleter](allocator::Buffer buffer) {
+      auto ptr = buffer.raw_ptr();
+      allocator::release(buffer);
+      return deleter(ptr);
+    };
+    set_data(buffer, std::move(wrapped_deleter));
+  }
+}
+
+/* Build an array from a shared buffer */
+array::array(allocator::Buffer data, Shape shape, Dtype dtype, Deleter deleter)
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+  set_data(data, deleter);
+}
+
+void array::detach() {
+  array_desc_->primitive = nullptr;
+  for (auto& s : array_desc_->siblings) {
+    s.array_desc_->primitive = nullptr;
+  }
+  for (auto& s : array_desc_->siblings) {
+    s.array_desc_->inputs.clear();
+    s.array_desc_->siblings.clear();
+    s.array_desc_->position = 0;
+  }
+  array_desc_->inputs.clear();
+  array_desc_->siblings.clear();
+  array_desc_->position = 0;
+}
+
+bool array::is_available() const {
+  if (status() == Status::available) {
+    return true;
+  } else if (
+      status() == Status::evaluated &&
+      (!event().valid() || event().is_signaled())) {
+    set_status(Status::available);
+    return true;
+  }
+  return false;
+}
+
+void array::wait() {
+  if (!is_available()) {
+    if (event().valid()) {
+      event().wait();
+      detach_event();
+    }
+    set_status(Status::available);
+  }
+}
+
+void array::eval() {
+  // Ensure the array is ready to be read
+  if (status() == Status::unscheduled) {
+    mlx::core::eval({*this});
+  } else {
+    wait();
+  }
+}
+
+bool array::is_tracer() const {
+  return (array_desc_->is_tracer && detail::in_tracing()) ||
+      detail::retain_graph();
+}
+
+void array::set_data(allocator::Buffer buffer, Deleter d) {
+  array_desc_->data = std::make_shared<Data>(buffer, d);
+  array_desc_->offset = 0;
+  array_desc_->data_size = size();
+  array_desc_->flags.contiguous = true;
+  array_desc_->flags.row_contiguous = true;
+  auto max_dim = std::max_element(shape().begin(), shape().end());
+  array_desc_->flags.col_contiguous = size() <= 1 || size() == *max_dim;
+}
+
+void array::set_data(
+    allocator::Buffer buffer,
+    size_t data_size,
+    Strides strides,
+    Flags flags,
+    Deleter d) {
+  array_desc_->data = std::make_shared<Data>(buffer, d);
+  array_desc_->offset = 0;
+  array_desc_->data_size = data_size;
+  array_desc_->strides = std::move(strides);
+  array_desc_->flags = flags;
+}
+
+void array::copy_shared_buffer(
+    const array& other,
+    const Strides& strides,
+    Flags flags,
+    size_t data_size,
+    int64_t offset /* = 0 */) {
+  array_desc_->data = other.array_desc_->data;
+  array_desc_->strides = strides;
+  array_desc_->flags = flags;
+  array_desc_->data_size = data_size;
+  array_desc_->offset =
+      sizeof(char) * itemsize() * offset + other.array_desc_->offset;
+}
+
+void array::copy_shared_buffer(const array& other) {
+  copy_shared_buffer(other, other.strides(), other.flags(), other.data_size());
+}
+
+array::~array() {
+  if (array_desc_ == nullptr) {
+    return;
+  }
+
+  // Detached/detaching
+  if (array_desc_->primitive == nullptr) {
+    return;
+  }
+
+  // Break circular reference for non-detached arrays with siblings
+  if (auto n = siblings().size(); n > 0) {
+    bool do_detach = true;
+    // If all siblings have siblings.size() references except
+    // the one we are currently destroying (which has siblings.size() + 1)
+    // then there are no more external references
+    do_detach &= (array_desc_.use_count() == (n + 1));
+    for (auto& s : siblings()) {
+      do_detach &= (s.array_desc_.use_count() == n);
+      if (!do_detach) {
+        break;
+      }
+    }
+    if (do_detach) {
+      for (auto& s : siblings()) {
+        for (auto& ss : s.siblings()) {
+          // Set to null here to avoid descending into array destructor
+          // for siblings
+          ss.array_desc_ = nullptr;
+        }
+        s.array_desc_->siblings.clear();
+      }
+    }
+  }
+}
+
+void array::ArrayDesc::init() {
+  strides.resize(shape.size());
+  size = 1;
+  for (int i = shape.size() - 1; i >= 0; --i) {
+    strides[i] = size;
+    size *= shape[i];
+  }
+  for (const auto& in : inputs) {
+    is_tracer |= in.is_tracer();
+  }
+}
+
+array::ArrayDesc::ArrayDesc(Shape shape, Dtype dtype)
+    : shape(std::move(shape)), dtype(dtype), status(Status::available) {
+  init();
+}
+
+array::ArrayDesc::ArrayDesc(
+    Shape shape,
+    Dtype dtype,
+    std::shared_ptr<Primitive> primitive,
+    std::vector<array> inputs)
+    : shape(std::move(shape)),
+      dtype(dtype),
+      primitive(std::move(primitive)),
+      status(Status::unscheduled),
+      inputs(std::move(inputs)) {
+  init();
+}
+
+array::ArrayDesc::~ArrayDesc() {
+  // When an array description is destroyed it will delete a bunch of arrays
+  // that may also destroy their corresponding descriptions and so on and so
+  // forth.
+  //
+  // This calls recursively the destructor and can result in stack overflow, we
+  // instead put them in a vector and destroy them one at a time resulting in a
+  // max stack depth of 2.
+  if (inputs.empty()) {
+    return;
+  }
+
+  std::vector<std::shared_ptr<ArrayDesc>> for_deletion;
+
+  auto append_deletable_inputs = [&for_deletion](ArrayDesc& ad) {
+    std::unordered_map<std::uintptr_t, array> input_map;
+    for (array& a : ad.inputs) {
+      if (a.array_desc_) {
+        input_map.insert({a.id(), a});
+        for (auto& s : a.siblings()) {
+          input_map.insert({s.id(), s});
+        }
+      }
+    }
+    ad.inputs.clear();
+    for (auto& [_, a] : input_map) {
+      bool is_deletable =
+          (a.array_desc_.use_count() <= a.siblings().size() + 1);
+      // An array with siblings is deletable only if all of its siblings
+      // are deletable
+      for (auto& s : a.siblings()) {
+        if (!is_deletable) {
+          break;
+        }
+        int is_input = (input_map.find(s.id()) != input_map.end());
+        is_deletable &=
+            s.array_desc_.use_count() <= a.siblings().size() + is_input;
+      }
+      if (is_deletable) {
+        for_deletion.push_back(std::move(a.array_desc_));
+      }
+    }
+  };
+
+  append_deletable_inputs(*this);
+
+  while (!for_deletion.empty()) {
+    // top is going to be deleted at the end of the block *after* the arrays
+    // with inputs have been moved into the vector
+    auto top = std::move(for_deletion.back());
+    for_deletion.pop_back();
+    append_deletable_inputs(*top);
+
+    // Clear out possible siblings to break circular references
+    for (auto& s : top->siblings) {
+      // Set to null here to avoid descending into top-level
+      // array destructor for siblings
+      s.array_desc_ = nullptr;
+    }
+    top->siblings.clear();
+  }
+}
+
+array::ArrayIterator::ArrayIterator(const array& arr, int idx)
+    : arr(arr), idx(idx) {
+  if (arr.ndim() == 0) {
+    throw std::invalid_argument("Cannot iterate over 0-d array.");
+  }
+}
+
+array::ArrayIterator::reference array::ArrayIterator::operator*() const {
+  auto start = Shape(arr.ndim(), 0);
+  auto end = arr.shape();
+  auto shape = arr.shape();
+  shape.erase(shape.begin());
+  start[0] = idx;
+  end[0] = idx + 1;
+  return reshape(slice(arr, start, end), shape);
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/array.h b/macos/third_party/mlx/mlx/array.h
new file mode 100644
index 0000000..60d5e50
--- /dev/null
+++ b/macos/third_party/mlx/mlx/array.h
@@ -0,0 +1,647 @@
+// Copyright © 2023 Apple Inc.
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <vector>
+
+#include "mlx/allocator.h"
+#include "mlx/api.h"
+#include "mlx/dtype.h"
+#include "mlx/event.h"
+#include "mlx/small_vector.h"
+
+namespace mlx::core {
+
+// Forward declaration
+class Primitive;
+
+using Deleter = std::function<void(allocator::Buffer)>;
+using ShapeElem = int32_t;
+using Shape = SmallVector<ShapeElem>;
+using Strides = SmallVector<int64_t>;
+
+class MLX_API array {
+  /* An array is really a node in a graph. It contains a shared ArrayDesc
+   * object */
+
+ public:
+  /** Construct a scalar array with zero dimensions. */
+  template <typename T>
+  explicit array(T val, Dtype dtype = TypeToDtype<T>());
+
+  /* Special case since std::complex can't be implicitly converted to other
+   * types. */
+  explicit array(const std::complex<float>& val, Dtype dtype = complex64);
+
+  template <typename It>
+  explicit array(
+      It data,
+      Shape shape,
+      Dtype dtype =
+          TypeToDtype<typename std::iterator_traits<It>::value_type>());
+
+  template <typename T>
+  explicit array(std::initializer_list<T> data, Dtype dtype = TypeToDtype<T>());
+
+  /* Special case so empty lists default to float32. */
+  explicit array(std::initializer_list<float> data);
+
+  /* Special case so array({}, type) is an empty array. */
+  explicit array(std::initializer_list<int> data, Dtype dtype);
+
+  template <typename T>
+  explicit array(
+      std::initializer_list<T> data,
+      Shape shape,
+      Dtype dtype = TypeToDtype<T>());
+
+  /* Build an array from a raw pointer. The constructor will attempt to use the
+   * input data without a copy. The deleter will be called when the array no
+   * longer needs the underlying memory - after the array is destroyed in the
+   * no-copy case and after the copy otherwise. */
+  explicit array(
+      void* data,
+      Shape shape,
+      Dtype dtype,
+      const std::function<void(void*)>& deleter);
+
+  /* Build an array from a buffer */
+  explicit array(
+      allocator::Buffer data,
+      Shape shape,
+      Dtype dtype,
+      Deleter deleter = allocator::free);
+
+  /** Assignment to rvalue does not compile. */
+  array& operator=(const array& other) && = delete;
+  array& operator=(array&& other) && = delete;
+
+  /** Default copy and move constructors otherwise. */
+  array& operator=(array&& other) & = default;
+  array(const array& other) = default;
+  array(array&& other) = default;
+
+  array& operator=(const array& other) & {
+    if (this->id() != other.id()) {
+      this->array_desc_ = other.array_desc_;
+    }
+    return *this;
+  }
+
+  /** The size of the array's datatype in bytes. */
+  size_t itemsize() const {
+    return size_of(dtype());
+  }
+
+  /** The number of elements in the array. */
+  size_t size() const {
+    return array_desc_->size;
+  }
+
+  /** The number of bytes in the array. */
+  size_t nbytes() const {
+    return size() * itemsize();
+  }
+
+  /** The number of dimensions of the array. */
+  size_t ndim() const {
+    return array_desc_->shape.size();
+  }
+
+  /** The shape of the array as a vector of integers. */
+  const Shape& shape() const {
+    return array_desc_->shape;
+  }
+
+  /**
+   *  Get the size of the corresponding dimension.
+   *
+   *  This function supports negative indexing and provides
+   *  bounds checking. */
+  auto shape(int dim) const {
+    return shape().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
+  }
+
+  /** The strides of the array. */
+  const Strides& strides() const {
+    return array_desc_->strides;
+  }
+
+  /**
+   *  Get the stride of the corresponding dimension.
+   *
+   *  This function supports negative indexing and provides
+   *  bounds checking. */
+  auto strides(int dim) const {
+    return strides().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
+  }
+
+  /** Get the arrays data type. */
+  Dtype dtype() const {
+    return array_desc_->dtype;
+  }
+
+  /** Evaluate the array. */
+  void eval();
+
+  /** Get the value from a scalar array. */
+  template <typename T>
+  T item();
+
+  template <typename T>
+  T item() const;
+
+  struct MLX_API ArrayIterator {
+    using iterator_category = std::random_access_iterator_tag;
+    using difference_type = size_t;
+    using value_type = const array;
+    using reference = value_type;
+
+    explicit ArrayIterator(const array& arr, int idx = 0);
+
+    reference operator*() const;
+
+    ArrayIterator& operator+(difference_type diff) {
+      idx += diff;
+      return *this;
+    }
+
+    ArrayIterator& operator++() {
+      idx++;
+      return *this;
+    }
+
+    friend bool operator==(const ArrayIterator& a, const ArrayIterator& b) {
+      return a.arr.id() == b.arr.id() && a.idx == b.idx;
+    }
+    friend bool operator!=(const ArrayIterator& a, const ArrayIterator& b) {
+      return !(a == b);
+    }
+
+   private:
+    const array& arr;
+    int idx;
+  };
+
+  ArrayIterator begin() const {
+    return ArrayIterator(*this);
+  }
+  ArrayIterator end() const {
+    return ArrayIterator(*this, shape(0));
+  }
+
+  /**
+   * The following methods should be used with caution.
+   * They are intended for use by the backend implementation and the
+   * API may change.
+   */
+
+  array(
+      Shape shape,
+      Dtype dtype,
+      std::shared_ptr<Primitive> primitive,
+      std::vector<array> inputs);
+
+  static std::vector<array> make_arrays(
+      std::vector<Shape> shapes,
+      const std::vector<Dtype>& dtypes,
+      const std::shared_ptr<Primitive>& primitive,
+      const std::vector<array>& inputs);
+
+  /**
+   * Get a new array that refers to the same data as the input but with a
+   * non-owning pointer to it. Note the array is detached from the graph and has
+   * no inputs, siblings or primitive.
+   */
+  static array unsafe_weak_copy(const array& other);
+
+  /** A unique identifier for an array. */
+  std::uintptr_t id() const {
+    return reinterpret_cast<std::uintptr_t>(array_desc_.get());
+  }
+
+  /** A unique identifier for an arrays primitive. */
+  std::uintptr_t primitive_id() const {
+    return reinterpret_cast<std::uintptr_t>(array_desc_->primitive.get());
+  }
+
+  struct Data {
+    allocator::Buffer buffer;
+    Deleter d;
+    Data(allocator::Buffer buffer, Deleter d = allocator::free)
+        : buffer(buffer), d(d) {}
+    // Not copyable
+    Data(const Data& d) = delete;
+    Data& operator=(const Data& d) = delete;
+    Data(Data&& o) : buffer(o.buffer), d(o.d) {
+      o.buffer = allocator::Buffer(nullptr);
+      o.d = [](allocator::Buffer) {};
+    }
+    ~Data() {
+      d(buffer);
+    }
+  };
+
+  struct Flags {
+    // True iff there are no gaps in the underlying data. Each item
+    // in the underlying data buffer belongs to at least one index.
+    //
+    // True iff:
+    // prod(shape[i] for i in range(ndim) if strides[i] > 0) == data_size()
+    bool contiguous : 1;
+
+    // True iff:
+    // strides[-1] == 1 and
+    // all(strides[i] == (shape[i+1]*strides[i+1]) or shape[i] == 1 for i in
+    // range(ndim - 1))
+    bool row_contiguous : 1;
+
+    // True iff:
+    // strides[0] == 1 and
+    // all(strides[i] == (shape[i-1]*strides[i-1]) or shape[i] == 1 for i in
+    // range(1, ndim))
+    bool col_contiguous : 1;
+  };
+
+  /** The array's primitive. */
+  Primitive& primitive() const {
+    return *(array_desc_->primitive);
+  }
+
+  /** A shared pointer to the array's primitive. */
+  std::shared_ptr<Primitive>& primitive_ptr() const {
+    return array_desc_->primitive;
+  }
+
+  /** Check if the array has an attached primitive or is a leaf node. */
+  bool has_primitive() const {
+    return array_desc_->primitive != nullptr;
+  }
+
+  /** The array's inputs. */
+  const std::vector<array>& inputs() const {
+    return array_desc_->inputs;
+  }
+
+  std::vector<array>& inputs() {
+    return array_desc_->inputs;
+  }
+
+  /** True indicates the arrays buffer is safe to reuse */
+  bool is_donatable() const {
+    return array_desc_.use_count() == 1 && (array_desc_->data.use_count() == 1);
+  }
+
+  /** The array's siblings. */
+  const std::vector<array>& siblings() const {
+    return array_desc_->siblings;
+  }
+
+  /** The array's siblings. */
+  std::vector<array>& siblings() {
+    return array_desc_->siblings;
+  }
+
+  /** The array's position in the sibling list. */
+  int sibling_position() const {
+    return array_desc_->position;
+  }
+
+  void set_siblings(std::vector<array> siblings, uint16_t position) {
+    array_desc_->siblings = std::move(siblings);
+    array_desc_->position = position;
+  }
+
+  /** The outputs of the array's primitive (i.e. this array and
+   * its siblings) in the order the primitive expects. */
+  std::vector<array> outputs() const {
+    auto idx = array_desc_->position;
+    std::vector<array> outputs;
+    outputs.reserve(siblings().size() + 1);
+    outputs.insert(outputs.end(), siblings().begin(), siblings().begin() + idx);
+    outputs.push_back(*this);
+    outputs.insert(outputs.end(), siblings().begin() + idx, siblings().end());
+    return outputs;
+  }
+
+  /** Detach the array from the graph. */
+  void detach();
+
+  /** Get the Flags bit-field. */
+  const Flags& flags() const {
+    return array_desc_->flags;
+  }
+
+  /** The size (in elements) of the underlying buffer the array points to.
+   *
+   * This can be different than the actual size of the array if the array has
+   * been broadcast or irregularly strided.  If ``first`` is the offset into
+   * the data buffer of the first element of the array (i.e. the offset
+   * corresponding to ``arr[0, 0, ...]``) and last is the offset into the
+   * data buffer of the last element of the array (i.e. the offset
+   * corresponding to ``arr[-1, -1, ...]``) then ``data_size = last - first``.
+   * Note, ``data_size`` is in units of ``item_size`` (not bytes).
+   **/
+  size_t data_size() const {
+    return array_desc_->data_size;
+  }
+
+  allocator::Buffer& buffer() {
+    return array_desc_->data->buffer;
+  }
+  const allocator::Buffer& buffer() const {
+    return array_desc_->data->buffer;
+  }
+
+  size_t buffer_size() const {
+    return allocator::allocator().size(buffer());
+  }
+
+  // Return the shared pointer to the array::Data struct
+  const std::shared_ptr<Data>& data_shared_ptr() const {
+    return array_desc_->data;
+  }
+
+  // Return a raw pointer to the arrays data. This function may do a copy if
+  // the underlying buffer is not accessible on the CPU. When accessing the
+  // data for GPU kernels, be sure to use the correct method / function for the
+  // given backend to access the GPU pointer.
+  template <typename T>
+  T* data() {
+    return reinterpret_cast<T*>(
+        (static_cast<char*>(buffer().raw_ptr()) + array_desc_->offset));
+  }
+
+  template <typename T>
+  const T* data() const {
+    return const_cast<array&>(*this).data<T>();
+  }
+
+  int64_t offset() const {
+    return array_desc_->offset;
+  }
+
+  enum Status {
+    // The output of a computation which has not been scheduled.
+    // For example, the status of `x` in `auto x = a + b`.
+    unscheduled,
+
+    // The array's `eval_*` function has been run, but the computation is not
+    // necessarily complete. The array will have memory allocated and if it is
+    // not a tracer then it will be detached from the graph.
+    evaluated,
+
+    // If the array is the output of a computation then the computation
+    // is complete. Constant arrays are always available (e.g. `array({1, 2,
+    // 3})`)
+    available
+  };
+
+  // Check if the array is safe to read.
+  bool is_available() const;
+
+  // Wait on the array to be available. After this `is_available` returns
+  // `true`.
+  void wait();
+
+  Status status() const {
+    return array_desc_->status;
+  }
+
+  void set_status(Status s) const {
+    array_desc_->status = s;
+  }
+
+  // Get the array's shared event
+  Event& event() const {
+    return array_desc_->event;
+  }
+
+  // Attach an event to a not yet evaluated array
+  void attach_event(Event e) const {
+    array_desc_->event = std::move(e);
+  }
+
+  void detach_event() const {
+    array_desc_->event = Event{};
+  }
+
+  // Mark the array as a tracer array (true) or not.
+  void set_tracer(bool is_tracer) {
+    array_desc_->is_tracer = is_tracer;
+  }
+  // Check if the array is a tracer array
+  bool is_tracer() const;
+
+  void set_data(allocator::Buffer buffer, Deleter d = allocator::free);
+
+  void set_data(
+      allocator::Buffer buffer,
+      size_t data_size,
+      Strides strides,
+      Flags flags,
+      Deleter d = allocator::free);
+
+  void copy_shared_buffer(
+      const array& other,
+      const Strides& strides,
+      Flags flags,
+      size_t data_size,
+      int64_t offset = 0);
+
+  void copy_shared_buffer(const array& other);
+
+  void overwrite_descriptor(const array& other) {
+    array_desc_ = other.array_desc_;
+  }
+
+  ~array();
+
+ private:
+  // Initialize the arrays data
+  template <typename It>
+  void init(const It src);
+
+  struct MLX_API ArrayDesc {
+    Shape shape;
+    Strides strides;
+    size_t size;
+    Dtype dtype;
+    std::shared_ptr<Primitive> primitive;
+
+    Status status;
+
+    // An event on the array used for synchronization
+    Event event;
+
+    // Indicates an array is being used in a graph transform
+    // and should not be detached from the graph
+    bool is_tracer{false};
+
+    // This is a shared pointer so that *different* arrays
+    // can share the underlying data buffer.
+    std::shared_ptr<Data> data;
+
+    // Offset from beginning of data pointer
+    int64_t offset{0};
+
+    // The size in elements of the data buffer the array accesses
+    size_t data_size{0};
+
+    // Contains useful meta data about the array
+    Flags flags{true, true, true};
+
+    std::vector<array> inputs;
+    // An array to keep track of the siblings from a multi-output
+    // primitive.
+    std::vector<array> siblings;
+    // The arrays position in the output list
+    uint32_t position{0};
+
+    explicit ArrayDesc(Shape shape, Dtype dtype);
+
+    explicit ArrayDesc(
+        Shape shape,
+        Dtype dtype,
+        std::shared_ptr<Primitive> primitive,
+        std::vector<array> inputs);
+
+    ~ArrayDesc();
+
+   private:
+    // Initialize size, strides, and other metadata
+    void init();
+  };
+
+  // The ArrayDesc contains the details of the materialized array including the
+  // shape, strides, the data type. It also includes
+  // the primitive which knows how to compute the array's data from its inputs
+  // and the list of array's inputs for the primitive.
+  std::shared_ptr<ArrayDesc> array_desc_;
+};
+
+template <typename T>
+array::array(T val, Dtype dtype /* = TypeToDtype<T>() */)
+    : array_desc_(std::make_shared<ArrayDesc>(Shape{}, dtype)) {
+  init(&val);
+}
+
+template <typename It>
+array::array(
+  It data,
+  Shape shape,
+  Dtype dtype /* = TypeToDtype<typename std::iterator_traits<It>::value_type>() */) :
+    array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+  init(data);
+}
+
+template <typename T>
+array::array(
+    std::initializer_list<T> data,
+    Dtype dtype /* = TypeToDtype<T>() */)
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              dtype)) {
+  init(data.begin());
+}
+
+template <typename T>
+array::array(
+    std::initializer_list<T> data,
+    Shape shape,
+    Dtype dtype /* = TypeToDtype<T>() */)
+    : array_desc_(std::make_shared<ArrayDesc>(std::move(shape), dtype)) {
+  if (data.size() != size()) {
+    throw std::invalid_argument(
+        "Data size and provided shape mismatch in array construction.");
+  }
+  init(data.begin());
+}
+
+template <typename T>
+T array::item() {
+  if (size() != 1) {
+    throw std::invalid_argument("item can only be called on arrays of size 1.");
+  }
+  eval();
+  return *data<T>();
+}
+
+template <typename T>
+T array::item() const {
+  if (size() != 1) {
+    throw std::invalid_argument("item can only be called on arrays of size 1.");
+  }
+  if (status() == Status::unscheduled) {
+    throw std::invalid_argument(
+        "item() const can only be called on evaled arrays");
+  }
+  const_cast<array*>(this)->eval();
+  return *data<T>();
+}
+
+template <typename It>
+void array::init(It src) {
+  set_data(allocator::malloc(size() * size_of(dtype())));
+  switch (dtype()) {
+    case bool_:
+      std::copy(src, src + size(), data<bool>());
+      break;
+    case uint8:
+      std::copy(src, src + size(), data<uint8_t>());
+      break;
+    case uint16:
+      std::copy(src, src + size(), data<uint16_t>());
+      break;
+    case uint32:
+      std::copy(src, src + size(), data<uint32_t>());
+      break;
+    case uint64:
+      std::copy(src, src + size(), data<uint64_t>());
+      break;
+    case int8:
+      std::copy(src, src + size(), data<int8_t>());
+      break;
+    case int16:
+      std::copy(src, src + size(), data<int16_t>());
+      break;
+    case int32:
+      std::copy(src, src + size(), data<int32_t>());
+      break;
+    case int64:
+      std::copy(src, src + size(), data<int64_t>());
+      break;
+    case float16:
+      std::copy(src, src + size(), data<float16_t>());
+      break;
+    case float32:
+      std::copy(src, src + size(), data<float>());
+      break;
+    case float64:
+      std::copy(src, src + size(), data<double>());
+      break;
+    case bfloat16:
+      std::copy(src, src + size(), data<bfloat16_t>());
+      break;
+    case complex64:
+      std::copy(src, src + size(), data<complex64_t>());
+      break;
+  }
+}
+
+/* Utilities for determining whether a template parameter is array. */
+template <typename T>
+inline constexpr bool is_array_v =
+    std::is_same_v<std::remove_cv_t<std::remove_reference_t<T>>, array>;
+
+template <typename... T>
+inline constexpr bool is_arrays_v = (is_array_v<T> && ...);
+
+template <typename... T>
+using enable_for_arrays_t = typename std::enable_if_t<is_arrays_v<T...>>;
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/common/CMakeLists.txt
new file mode 100644
index 0000000..6c4e250
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/CMakeLists.txt
@@ -0,0 +1,9 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/broadcasting.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
diff --git a/macos/third_party/mlx/mlx/backend/common/binary.h b/macos/third_party/mlx/mlx/backend/common/binary.h
new file mode 100644
index 0000000..78607ef
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/binary.h
@@ -0,0 +1,97 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/allocator.h"
+#include "mlx/array.h"
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+enum class BinaryOpType {
+  ScalarScalar,
+  ScalarVector,
+  VectorScalar,
+  VectorVector,
+  General,
+};
+
+inline BinaryOpType get_binary_op_type(const array& a, const array& b) {
+  BinaryOpType bopt;
+  if (a.data_size() == 1 && b.data_size() == 1) {
+    bopt = BinaryOpType::ScalarScalar;
+  } else if (a.data_size() == 1 && b.flags().contiguous) {
+    bopt = BinaryOpType::ScalarVector;
+  } else if (b.data_size() == 1 && a.flags().contiguous) {
+    bopt = BinaryOpType::VectorScalar;
+  } else if (
+      (a.flags().row_contiguous && b.flags().row_contiguous) ||
+      (a.flags().col_contiguous && b.flags().col_contiguous)) {
+    bopt = BinaryOpType::VectorVector;
+  } else {
+    bopt = BinaryOpType::General;
+  }
+  return bopt;
+}
+
+inline void set_binary_op_output_data(
+    const array& a,
+    const array& b,
+    array& out,
+    BinaryOpType bopt,
+    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+  bool b_donatable = is_donatable(b, out);
+  bool a_donatable = is_donatable(a, out);
+  switch (bopt) {
+    case BinaryOpType::ScalarScalar:
+      out.set_data(mallocfn(out.itemsize()), 1, a.strides(), a.flags());
+      break;
+    case BinaryOpType::ScalarVector:
+      if (b_donatable) {
+        out.copy_shared_buffer(b);
+      } else {
+        out.set_data(
+            mallocfn(b.data_size() * out.itemsize()),
+            b.data_size(),
+            b.strides(),
+            b.flags());
+      }
+      break;
+    case BinaryOpType::VectorScalar:
+      if (a_donatable) {
+        out.copy_shared_buffer(a);
+      } else {
+        out.set_data(
+            mallocfn(a.data_size() * out.itemsize()),
+            a.data_size(),
+            a.strides(),
+            a.flags());
+      }
+      break;
+    case BinaryOpType::VectorVector:
+      if (a_donatable) {
+        out.copy_shared_buffer(a);
+      } else if (b_donatable) {
+        out.copy_shared_buffer(b);
+      } else {
+        out.set_data(
+            mallocfn(a.data_size() * out.itemsize()),
+            a.data_size(),
+            a.strides(),
+            a.flags());
+      }
+      break;
+    case BinaryOpType::General:
+      if (a_donatable && a.flags().row_contiguous && a.size() == out.size()) {
+        out.copy_shared_buffer(a);
+      } else if (
+          b_donatable && b.flags().row_contiguous && b.size() == out.size()) {
+        out.copy_shared_buffer(b);
+      } else {
+        out.set_data(mallocfn(out.nbytes()));
+      }
+      break;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/broadcasting.cpp b/macos/third_party/mlx/mlx/backend/common/broadcasting.cpp
new file mode 100644
index 0000000..0bb52e0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/broadcasting.cpp
@@ -0,0 +1,24 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+void broadcast(const array& in, array& out) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+  Strides strides(out.ndim(), 0);
+  int diff = out.ndim() - in.ndim();
+  for (int i = in.ndim() - 1; i >= 0; --i) {
+    strides[i + diff] = (in.shape()[i] == 1) ? 0 : in.strides()[i];
+  }
+  auto flags = in.flags();
+  if (out.size() > in.size()) {
+    flags.row_contiguous = flags.col_contiguous = false;
+  }
+  out.copy_shared_buffer(in, strides, flags, in.data_size());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/broadcasting.h b/macos/third_party/mlx/mlx/backend/common/broadcasting.h
new file mode 100644
index 0000000..29651e9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/broadcasting.h
@@ -0,0 +1,11 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+void broadcast(const array& in, array& out);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/buffer_cache.h b/macos/third_party/mlx/mlx/backend/common/buffer_cache.h
new file mode 100644
index 0000000..27648b2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/buffer_cache.h
@@ -0,0 +1,158 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <functional>
+#include <map>
+
+namespace mlx::core {
+
+template <typename T>
+class BufferCache {
+ public:
+  BufferCache(
+      size_t page_size,
+      std::function<size_t(T*)> get_size,
+      std::function<void(T*)> free)
+      : page_size_(page_size),
+        get_size_(std::move(get_size)),
+        free_(std::move(free)) {}
+
+  ~BufferCache() {
+    clear();
+  }
+
+  BufferCache(const BufferCache&) = delete;
+  BufferCache& operator=(const BufferCache&) = delete;
+
+  T* reuse_from_cache(size_t size) {
+    // Find the closest buffer in pool.
+    auto it = buffer_pool_.lower_bound(size);
+    if (it == buffer_pool_.end() ||
+        it->first >= std::min(2 * size, size + 2 * page_size_)) {
+      return nullptr;
+    }
+
+    // Collect from the cache.
+    T* buf = it->second->buf;
+    pool_size_ -= it->first;
+
+    // Remove from record.
+    remove_from_list(it->second);
+    buffer_pool_.erase(it);
+    return buf;
+  }
+
+  void recycle_to_cache(T* buf) {
+    assert(buf);
+    // Add to cache.
+    BufferHolder* bh = new BufferHolder(buf);
+    add_at_head(bh);
+    size_t size = get_size_(buf);
+    pool_size_ += size;
+    buffer_pool_.emplace(size, bh);
+  }
+
+  int release_cached_buffers(size_t min_bytes_to_free) {
+    if (min_bytes_to_free >= 0.9 * pool_size_) {
+      return clear();
+    } else {
+      int n_release = 0;
+      size_t total_bytes_freed = 0;
+
+      while (tail_ && (total_bytes_freed < min_bytes_to_free)) {
+        // Release buffer.
+        size_t size = get_size_(tail_->buf);
+        total_bytes_freed += size;
+        free_(tail_->buf);
+        n_release++;
+
+        // Remove from record.
+        auto its = buffer_pool_.equal_range(size);
+        auto it = std::find_if(its.first, its.second, [this](const auto& el) {
+          return el.second == tail_;
+        });
+        assert(it != buffer_pool_.end());
+        buffer_pool_.erase(it);
+        remove_from_list(tail_);
+      }
+
+      pool_size_ -= total_bytes_freed;
+      return n_release;
+    }
+  }
+
+  int clear() {
+    int n_release = 0;
+    for (auto& [size, holder] : buffer_pool_) {
+      free_(holder->buf);
+      n_release++;
+      delete holder;
+    }
+    buffer_pool_.clear();
+    pool_size_ = 0;
+    head_ = nullptr;
+    tail_ = nullptr;
+    return n_release;
+  }
+
+  size_t cache_size() const {
+    return pool_size_;
+  }
+
+  size_t page_size() const {
+    return page_size_;
+  }
+
+ private:
+  struct BufferHolder {
+   public:
+    explicit BufferHolder(T* buf_) : buf(buf_) {}
+
+    BufferHolder* prev{nullptr};
+    BufferHolder* next{nullptr};
+    T* buf;
+  };
+
+  void add_at_head(BufferHolder* to_add) {
+    if (!head_) {
+      head_ = to_add;
+      tail_ = to_add;
+    } else {
+      head_->prev = to_add;
+      to_add->next = head_;
+      head_ = to_add;
+    }
+  }
+
+  void remove_from_list(BufferHolder* to_remove) {
+    if (to_remove->prev && to_remove->next) { // if middle
+      to_remove->prev->next = to_remove->next;
+      to_remove->next->prev = to_remove->prev;
+    } else if (to_remove->prev && to_remove == tail_) { // if tail
+      tail_ = to_remove->prev;
+      tail_->next = nullptr;
+    } else if (to_remove == head_ && to_remove->next) { // if head
+      head_ = to_remove->next;
+      head_->prev = nullptr;
+    } else if (to_remove == head_ && to_remove == tail_) { // if only element
+      head_ = nullptr;
+      tail_ = nullptr;
+    }
+
+    delete to_remove;
+  }
+
+  std::multimap<size_t, BufferHolder*> buffer_pool_;
+  BufferHolder* head_{nullptr};
+  BufferHolder* tail_{nullptr};
+  size_t pool_size_{0};
+
+  const size_t page_size_;
+  std::function<size_t(T*)> get_size_;
+  std::function<void(T*)> free_;
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/common.cpp b/macos/third_party/mlx/mlx/backend/common/common.cpp
new file mode 100644
index 0000000..2cda88a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/common.cpp
@@ -0,0 +1,305 @@
+// Copyright © 2024 Apple Inc.
+#include <cassert>
+
+#include "mlx/backend/common/broadcasting.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void AsStrided::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  auto& in = inputs[0];
+
+  if (!in.flags().row_contiguous) {
+    // Just ensuring that inputs[0] came from the ops which would ensure the
+    // input is row contiguous.
+    throw std::runtime_error(
+        "AsStrided must be used with row contiguous arrays only.");
+  }
+
+  // Compute the flags given the shape and strides
+  bool row_contiguous = true, col_contiguous = true;
+  size_t r = 1, c = 1;
+  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
+    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
+    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
+    r *= shape_[i];
+    c *= shape_[j];
+  }
+  auto flags = in.flags();
+  // TODO: Compute the contiguous flag in a better way cause now we are
+  //       unnecessarily strict.
+  flags.contiguous = row_contiguous || col_contiguous;
+  flags.row_contiguous = row_contiguous;
+  flags.col_contiguous = col_contiguous;
+
+  // There is no easy way to compute the actual data size so we use out.size().
+  // The contiguous flag will almost certainly not be set so no code should
+  // rely on data_size anyway.
+  size_t data_size = out.size();
+
+  return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
+}
+
+void Broadcast::eval(const std::vector<array>& inputs, array& out) {
+  broadcast(inputs[0], out);
+}
+
+void BroadcastAxes::eval(const std::vector<array>& inputs, array& out) {
+  broadcast(inputs[0], out);
+}
+
+void Copy::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  out.copy_shared_buffer(inputs[0]);
+}
+
+void CustomTransforms::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() > outputs.size());
+  for (int i = 0, j = inputs.size() - outputs.size(); i < outputs.size();
+       i++, j++) {
+    outputs[i].copy_shared_buffer(inputs[j]);
+  }
+}
+
+void Depends::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() > outputs.size());
+  for (int i = 0; i < outputs.size(); i++) {
+    outputs[i].copy_shared_buffer(inputs[i]);
+  }
+}
+
+void ExpandDims::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  auto strides = in.strides();
+  for (auto ax : axes_) {
+    strides.insert(strides.begin() + ax, 1);
+  }
+  out.copy_shared_buffer(in, strides, in.flags(), in.data_size());
+}
+
+void NumberOfElements::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  double numel = 1;
+  for (auto ax : axes_) {
+    numel *= inputs[0].shape(ax);
+  }
+
+  if (inverted_) {
+    numel = 1.0 / numel;
+  }
+
+  switch (out.dtype()) {
+    case bool_:
+      *out.data<bool>() = static_cast<bool>(numel);
+      break;
+    case uint8:
+      *out.data<uint8_t>() = static_cast<uint8_t>(numel);
+      break;
+    case uint16:
+      *out.data<uint16_t>() = static_cast<uint16_t>(numel);
+      break;
+    case uint32:
+      *out.data<uint32_t>() = static_cast<uint32_t>(numel);
+      break;
+    case uint64:
+      *out.data<uint64_t>() = static_cast<uint64_t>(numel);
+      break;
+    case int8:
+      *out.data<int8_t>() = static_cast<int8_t>(numel);
+      break;
+    case int16:
+      *out.data<int16_t>() = static_cast<int16_t>(numel);
+      break;
+    case int32:
+      *out.data<int32_t>() = static_cast<int32_t>(numel);
+      break;
+    case int64:
+      *out.data<int64_t>() = static_cast<int64_t>(numel);
+      break;
+    case float16:
+      *out.data<float16_t>() = static_cast<float16_t>(numel);
+      break;
+    case float32:
+      *out.data<float>() = static_cast<float>(numel);
+      break;
+    case bfloat16:
+      *out.data<bfloat16_t>() = static_cast<bfloat16_t>(numel);
+      break;
+    case float64:
+      *out.data<double>() = static_cast<double>(numel);
+      break;
+    case complex64:
+      *out.data<complex64_t>() = static_cast<complex64_t>(numel);
+      break;
+  }
+}
+
+std::pair<bool, Strides> prepare_reshape(const array& in, const array& out) {
+  // Special case for empty arrays or row contiguous arrays
+  if (in.size() == 0 || in.flags().row_contiguous) {
+    return {false, out.strides()};
+  }
+
+  // Special case for scalars
+  if (in.ndim() == 0) {
+    return {false, Strides(out.ndim(), 0)};
+  }
+
+  // Firstly let's collapse all the contiguous dimensions of the input
+  auto [shape, strides] = collapse_contiguous_dims(in);
+
+  // If shapes fit exactly in the contiguous dims then no copy is necessary so
+  // let's check.
+  Strides out_strides;
+  bool copy_necessary = false;
+  int j = 0;
+  for (int i = 0; i < out.ndim(); i++) {
+    int N = out.shape(i);
+    if (j < shape.size() && shape[j] % N == 0) {
+      shape[j] /= N;
+      out_strides.push_back(shape[j] * strides[j]);
+      j += (shape[j] == 1);
+    } else if (N == 1) {
+      // i > 0 because otherwise j < shape.size() && shape[j] % 1 == 0
+      out_strides.push_back(out_strides.back());
+    } else {
+      copy_necessary = true;
+      break;
+    }
+  }
+
+  return {copy_necessary, out_strides};
+}
+
+void shared_buffer_reshape(
+    const array& in,
+    const Strides& out_strides,
+    array& out) {
+  auto flags = in.flags();
+  if (flags.row_contiguous) {
+    // For row contiguous reshapes:
+    // - Shallow copy the buffer
+    // - If reshaping into a vector (all singleton dimensions except one) it
+    //    becomes col contiguous again.
+    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
+    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
+  }
+  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
+}
+
+void Split::eval(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+
+  auto& in = inputs[0];
+
+  auto compute_new_flags = [](const auto& shape,
+                              const auto& strides,
+                              size_t in_data_size,
+                              auto flags) {
+    size_t data_size = 1;
+    size_t f_stride = 1;
+    size_t b_stride = 1;
+    flags.row_contiguous = true;
+    flags.col_contiguous = true;
+    for (int i = 0, ri = shape.size() - 1; ri >= 0; i++, ri--) {
+      flags.col_contiguous &= strides[i] == f_stride || shape[i] == 1;
+      flags.row_contiguous &= strides[ri] == b_stride || shape[ri] == 1;
+      f_stride *= shape[i];
+      b_stride *= shape[ri];
+      if (strides[i] > 0) {
+        data_size *= shape[i];
+      }
+    }
+
+    if (data_size == 1) {
+      // Broadcasted scalar array is contiguous.
+      flags.contiguous = true;
+    } else if (data_size == in_data_size) {
+      // Means we sliced a broadcasted dimension so leave the "no holes" flag
+      // alone.
+    } else {
+      // We sliced something. So either we are row or col contiguous or we
+      // punched a hole.
+      flags.contiguous &= flags.row_contiguous || flags.col_contiguous;
+    }
+
+    return std::pair<decltype(flags), size_t>{flags, data_size};
+  };
+
+  std::vector<int> indices(1, 0);
+  indices.insert(indices.end(), indices_.begin(), indices_.end());
+  for (int i = 0; i < indices.size(); i++) {
+    size_t offset = indices[i] * in.strides()[axis_];
+    auto [new_flags, data_size] = compute_new_flags(
+        outputs[i].shape(), in.strides(), in.data_size(), in.flags());
+    outputs[i].copy_shared_buffer(
+        in, in.strides(), new_flags, data_size, offset);
+  }
+}
+
+void Squeeze::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  Strides strides;
+  for (int i = 0, j = 0; i < in.ndim(); ++i) {
+    if (j < axes_.size() && i == axes_[j]) {
+      j++;
+    } else {
+      strides.push_back(in.strides(i));
+    }
+  }
+  out.copy_shared_buffer(in, strides, in.flags(), in.data_size());
+}
+
+void StopGradient::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  out.copy_shared_buffer(inputs[0]);
+}
+
+void Transpose::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  Strides out_strides(out.ndim());
+  auto& in = inputs[0];
+  for (int ax = 0; ax < axes_.size(); ++ax) {
+    out_strides[ax] = in.strides()[axes_[ax]];
+  }
+
+  // Conditions for {row/col}_contiguous
+  // - array must be contiguous (no gaps)
+  // - underlying buffer size should have the same size as the array
+  // - cumulative product of shapes is equal to the strides (we can ignore axes
+  //   with size == 1)
+  //   - in the forward direction (column contiguous)
+  //   - in the reverse direction (row contiguous)
+  // - vectors are both row and col contiguous (hence if both row/col are
+  //   true, they stay true)
+  auto flags = in.flags();
+  if (flags.contiguous && in.data_size() == in.size()) {
+    int64_t f_stride = 1;
+    int64_t b_stride = 1;
+    flags.col_contiguous = true;
+    flags.row_contiguous = true;
+    for (int i = 0, ri = out.ndim() - 1; i < out.ndim(); ++i, --ri) {
+      flags.col_contiguous &= (out_strides[i] == f_stride || out.shape(i) == 1);
+      f_stride *= out.shape(i);
+      flags.row_contiguous &=
+          (out_strides[ri] == b_stride || out.shape(ri) == 1);
+      b_stride *= out.shape(ri);
+    }
+  }
+  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/compiled.cpp b/macos/third_party/mlx/mlx/backend/common/compiled.cpp
new file mode 100644
index 0000000..aceeb1f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/compiled.cpp
@@ -0,0 +1,243 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+void print_constant(std::ostream& os, const array& x) {
+  switch (x.dtype()) {
+    case float32:
+      return print_float_constant<float>(os, x);
+    case float16:
+      return print_float_constant<float16_t>(os, x);
+    case bfloat16:
+      return print_float_constant<bfloat16_t>(os, x);
+    case float64:
+      return print_float_constant<double>(os, x);
+    case complex64:
+      return print_complex_constant<complex64_t>(os, x);
+    case int8:
+      os << static_cast<int32_t>(x.item<int8_t>());
+      return;
+    case int16:
+      return print_int_constant<int16_t>(os, x);
+    case int32:
+      return print_int_constant<int32_t>(os, x);
+    case int64:
+      return print_int_constant<int64_t>(os, x);
+    case uint8:
+      os << static_cast<uint32_t>(x.item<uint8_t>());
+      return;
+    case uint16:
+      return print_int_constant<uint16_t>(os, x);
+    case uint32:
+      return print_int_constant<uint32_t>(os, x);
+    case uint64:
+      return print_int_constant<uint64_t>(os, x);
+    case bool_:
+      os << std::boolalpha << x.item<bool>();
+      return;
+    default:
+      throw std::runtime_error("Unsupported constant type");
+  }
+}
+
+std::string get_type_string(Dtype d) {
+  switch (d) {
+    case float32:
+      return "float";
+    case float16:
+      return "float16_t";
+    case bfloat16:
+      return "bfloat16_t";
+    case float64:
+      return "double";
+    case complex64:
+      return "complex64_t";
+    case bool_:
+      return "bool";
+    case int8:
+      return "int8_t";
+    case int16:
+      return "int16_t";
+    case int32:
+      return "int32_t";
+    case int64:
+      return "int64_t";
+    case uint8:
+      return "uint8_t";
+    case uint16:
+      return "uint16_t";
+    case uint32:
+      return "uint32_t";
+    case uint64:
+      return "uint64_t";
+    default: {
+      std::ostringstream msg;
+      msg << "Unsupported compilation type " << d;
+      throw std::runtime_error(msg.str());
+    }
+  }
+}
+
+bool compiled_check_contiguity(
+    const std::vector<array>& inputs,
+    const Shape& shape) {
+  bool contiguous = true;
+  bool all_contig = true;
+  bool all_row_contig = true;
+  bool all_col_contig = true;
+  int non_scalar_inputs = 0;
+  for (const auto& x : inputs) {
+    if (is_scalar(x)) {
+      continue;
+    }
+    non_scalar_inputs++;
+    bool shape_eq = x.shape() == shape;
+    all_contig &= (x.flags().contiguous && shape_eq);
+    all_row_contig &= (x.flags().row_contiguous && shape_eq);
+    all_col_contig &= (x.flags().col_contiguous && shape_eq);
+  }
+  if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) {
+    contiguous = false;
+  } else if (non_scalar_inputs == 1 && !all_contig) {
+    contiguous = false;
+  } else if (non_scalar_inputs == 0 && !shape.empty()) {
+    contiguous = false;
+  }
+  return contiguous;
+}
+
+void compiled_allocate_outputs(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const std::function<bool(size_t)>& is_constant,
+    bool contiguous,
+    const std::function<allocator::Buffer(size_t)>&
+        mallocfn /* = allocator::malloc */) {
+  if (contiguous) {
+    int o = 0;
+    Strides strides;
+    size_t data_size;
+    array::Flags flags;
+    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+      auto& in = inputs[i];
+      // Conditions for donation
+      // - Correct size
+      // - Not a scalar
+      // - Donatable
+      // - Not a constant
+      if (in.itemsize() == outputs[o].itemsize() && !is_scalar(in) &&
+          in.is_donatable() && !is_constant(i)) {
+        outputs[o++].copy_shared_buffer(in);
+      }
+      // Get representative input flags to properly set non-donated outputs
+      if (strides.empty() && in.size() == outputs[0].size()) {
+        strides = in.strides();
+        flags = in.flags();
+        data_size = in.data_size();
+      }
+    }
+    for (; o < outputs.size(); ++o) {
+      outputs[o].set_data(
+          mallocfn(data_size * outputs[o].itemsize()),
+          data_size,
+          strides,
+          flags);
+    }
+  } else {
+    int o = 0;
+    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+      auto& in = inputs[i];
+      // Conditions for donation
+      // - Row contiguous
+      // - Donatable
+      // - Correct size
+      // - Not a constant
+      if (in.flags().row_contiguous && in.size() == outputs[o].size() &&
+          in.itemsize() == outputs[o].itemsize() && in.is_donatable() &&
+          !is_constant(i)) {
+        outputs[o].copy_shared_buffer(
+            in, outputs[o].strides(), in.flags(), in.data_size());
+        o++;
+      }
+    }
+    for (; o < outputs.size(); ++o) {
+      outputs[o].set_data(mallocfn(outputs[o].nbytes()));
+    }
+  }
+}
+
+std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(
+    const std::vector<array>& inputs,
+    const array& out,
+    const std::function<bool(size_t)>& is_constant) {
+  const Shape& shape = out.shape();
+  bool contiguous = compiled_check_contiguity(inputs, shape);
+  if (contiguous) {
+    return {true, shape, {}};
+  }
+
+  std::vector<Strides> strides_vec{out.strides()};
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    // Skip constants.
+    if (is_constant(i)) {
+      continue;
+    }
+
+    // Skip scalar inputs.
+    const auto& x = inputs[i];
+    if (is_scalar(x)) {
+      continue;
+    }
+
+    // Broadcast the inputs to the output shape.
+    Strides xstrides;
+    size_t j = 0;
+    for (; j < shape.size() - x.ndim(); ++j) {
+      if (shape[j] == 1) {
+        xstrides.push_back(out.strides()[j]);
+      } else {
+        xstrides.push_back(0);
+      }
+    }
+    for (size_t i = 0; i < x.ndim(); ++i, ++j) {
+      if (x.shape(i) == 1) {
+        if (shape[j] == 1) {
+          xstrides.push_back(out.strides()[j]);
+        } else {
+          xstrides.push_back(0);
+        }
+      } else {
+        xstrides.push_back(x.strides()[i]);
+      }
+    }
+    strides_vec.push_back(std::move(xstrides));
+  }
+
+  auto tup = collapse_contiguous_dims(shape, strides_vec, INT32_MAX);
+  return {false, std::move(std::get<0>(tup)), std::move(std::get<1>(tup))};
+}
+
+bool compiled_use_large_index(
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    bool contiguous) {
+  if (contiguous) {
+    size_t max_size = 0;
+    for (const auto& in : inputs) {
+      max_size = std::max(max_size, in.data_size());
+    }
+    return max_size > UINT32_MAX;
+  } else {
+    size_t max_size = 0;
+    for (const auto& o : outputs) {
+      max_size = std::max(max_size, o.size());
+    }
+    return max_size > UINT32_MAX;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/compiled.h b/macos/third_party/mlx/mlx/backend/common/compiled.h
new file mode 100644
index 0000000..3be3713
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/compiled.h
@@ -0,0 +1,77 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+
+#include <functional>
+#include <iomanip>
+
+#include "mlx/array.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+inline bool is_static_cast(const Primitive& p) {
+  return (typeid(p) == typeid(Broadcast) || typeid(p) == typeid(AsType));
+}
+
+std::string get_type_string(Dtype d);
+
+template <typename T>
+void print_float_constant(std::ostream& os, const array& x) {
+  auto old_precision = os.precision();
+  if constexpr (std::is_same_v<T, double>) {
+    os << std::setprecision(std::numeric_limits<double>::digits10 + 1);
+  } else {
+    os << std::setprecision(std::numeric_limits<float>::digits10 + 1);
+  }
+  os << x.item<T>() << std::setprecision(old_precision);
+}
+
+template <typename T>
+void print_int_constant(std::ostream& os, const array& x) {
+  os << x.item<T>();
+}
+
+template <typename T>
+void print_complex_constant(std::ostream& os, const array& x) {
+  auto old_precision = os.precision();
+  T constant = x.item<T>();
+
+  os << get_type_string(x.dtype()) << "("
+     << std::setprecision(std::numeric_limits<float>::digits10 + 1)
+     << constant.real() << ", " << constant.imag() << ")"
+     << std::setprecision(old_precision);
+}
+
+void print_constant(std::ostream& os, const array& x);
+
+inline bool is_scalar(const array& x) {
+  return x.ndim() == 0;
+}
+
+// Check if we can use a contiguous operation given inputs and the output shape
+bool compiled_check_contiguity(
+    const std::vector<array>& inputs,
+    const Shape& shape);
+
+// Allocate space for the outputs possibly with input donation
+void compiled_allocate_outputs(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const std::function<bool(size_t)>& is_constant,
+    bool contiguous,
+    const std::function<allocator::Buffer(size_t)>& mallocfn =
+        allocator::malloc);
+
+// Collapse contiguous dims ignoring scalars and constants.
+std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims(
+    const std::vector<array>& inputs,
+    const array& out,
+    const std::function<bool(size_t)>& is_constant);
+
+// Return whether the kernel should use large index.
+bool compiled_use_large_index(
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    bool contiguous);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/copy.h b/macos/third_party/mlx/mlx/backend/common/copy.h
new file mode 100644
index 0000000..859ce04
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/copy.h
@@ -0,0 +1,50 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+enum class CopyType {
+  // Copy a raw scalar input into the full contiguous output
+  Scalar,
+
+  // Copy the raw input buffer contiguously into a raw output buffer of the same
+  // size
+  Vector,
+
+  // Copy the full virtual input to the full contiguous output
+  General,
+
+  // Copy the full virtual input to the full virtual output. We assume the
+  // input and output have the same shape.
+  GeneralGeneral
+};
+
+inline bool set_copy_output_data(
+    const array& in,
+    array& out,
+    CopyType ctype,
+    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+  if (ctype == CopyType::Vector) {
+    // If the input is donateable, we are doing a vector copy and the types
+    // have the same size, then the input buffer can hold the output.
+    if (is_donatable(in, out)) {
+      out.copy_shared_buffer(in);
+      return true;
+    } else {
+      out.set_data(
+          mallocfn(in.data_size() * out.itemsize()),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+      return false;
+    }
+  } else {
+    out.set_data(mallocfn(out.nbytes()));
+    return false;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/hadamard.h b/macos/third_party/mlx/mlx/backend/common/hadamard.h
new file mode 100644
index 0000000..ba5c4e4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/hadamard.h
@@ -0,0 +1,109 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <map>
+
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+// From http://neilsloane.com/hadamard/
+constexpr std::string_view h12 = R"(
++-++++++++++
+--+-+-+-+-+-
++++-++----++
++---+--+-++-
++++++-++----
++-+---+--+-+
+++--+++-++--
++--++---+--+
+++----+++-++
++--+-++---+-
+++++----+++-
++-+--+-++---
+)";
+
+constexpr std::string_view h20 = R"(
++----+----++--++-++-
+-+----+---+++---+-++
+--+----+---+++-+-+-+
+---+----+---+++++-+-
+----+----++--++-++-+
+-+++++-----+--+++--+
++-+++-+---+-+--+++--
+++-++--+---+-+--+++-
++++-+---+---+-+--+++
+++++-----++--+-+--++
+--++-+-++-+-----++++
+---++-+-++-+---+-+++
++---++-+-+--+--++-++
+++---++-+----+-+++-+
+-++---++-+----+++++-
+-+--+--++-+----+----
++-+-----++-+----+---
+-+-+-+---+--+----+--
+--+-+++------+----+-
++--+--++------+----+
+)";
+
+constexpr std::string_view h28 = R"(
++------++----++-+--+-+--++--
+-+-----+++-----+-+--+-+--++-
+--+-----+++---+-+-+----+--++
+---+-----+++---+-+-+-+--+--+
+----+-----+++---+-+-+++--+--
+-----+-----++++--+-+--++--+-
+------++----++-+--+-+--++--+
+--++++-+-------++--+++-+--+-
+---++++-+-----+-++--+-+-+--+
++---+++--+----++-++--+-+-+--
+++---++---+----++-++--+-+-+-
++++---+----+----++-++--+-+-+
+++++--------+-+--++-++--+-+-
+-++++--------+++--++--+--+-+
+-+-++-++--++--+--------++++-
++-+-++--+--++--+--------++++
+-+-+-++--+--++--+----+---+++
++-+-+-++--+--+---+---++---++
+++-+-+-++--+------+--+++---+
+-++-+-+-++--+------+-++++---
++-++-+---++--+------+-++++--
+-++--++-+-++-+++----++------
++-++--++-+-++-+++-----+-----
+++-++---+-+-++-+++-----+----
+-++-++-+-+-+-+--+++-----+---
+--++-++++-+-+----+++-----+--
++--++-+-++-+-+----+++-----+-
+++--++-+-++-+-+----++------+
+)";
+
+inline const std::map<int, std::string_view> hadamard_matrices() {
+  return {{12, h12}, {20, h20}, {28, h28}};
+}
+
+inline std::pair<int, int> decompose_hadamard(int n) {
+  // n = m*2^k
+  int m = 1;
+  if (!is_power_of_2(n)) {
+    auto h_matrices = hadamard_matrices();
+    for (auto [factor, _] : h_matrices) {
+      if (n % factor == 0) {
+        m = factor;
+        n /= factor;
+        break;
+      }
+    }
+    if (m == 1) {
+      throw std::invalid_argument(
+          "[hadamard] Only supports n = m*2^k where m in (1, 12, 20, 28).");
+    }
+  }
+  if (n > (1 << 26)) {
+    throw std::invalid_argument(
+        "[hadamard] Only supports n = m*2^k where k <= 26");
+  }
+  return {n, m};
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/load.cpp b/macos/third_party/mlx/mlx/backend/common/load.cpp
new file mode 100644
index 0000000..ce41963
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/load.cpp
@@ -0,0 +1,57 @@
+// Copyright © 2023 Apple Inc.
+
+#include <algorithm>
+#include <utility>
+
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+
+namespace {
+
+template <const uint8_t scalar_size>
+void swap_endianness(uint8_t* data_bytes, size_t N) {
+  struct Elem {
+    uint8_t bytes[scalar_size];
+  };
+
+  Elem* data = reinterpret_cast<Elem*>(data_bytes);
+
+  for (size_t i = 0; i < N; i++) {
+    for (size_t j = 0; j < (scalar_size / 2); j++) {
+      std::swap(data[i].bytes[j], data[i].bytes[scalar_size - j - 1]);
+    }
+  }
+}
+
+} // namespace
+
+namespace mlx::core {
+
+void Load::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto read_task = [out_ptr = out.data<char>(),
+                    size = out.size(),
+                    itemsize = out.itemsize(),
+                    offset = offset_,
+                    reader = reader_,
+                    swap_endianness_ = swap_endianness_]() mutable {
+    reader->read(out_ptr, size * itemsize, offset);
+    if (swap_endianness_) {
+      switch (itemsize) {
+        case 2:
+          swap_endianness<2>(reinterpret_cast<uint8_t*>(out_ptr), size);
+          break;
+        case 4:
+          swap_endianness<4>(reinterpret_cast<uint8_t*>(out_ptr), size);
+          break;
+        case 8:
+          swap_endianness<8>(reinterpret_cast<uint8_t*>(out_ptr), size);
+          break;
+      }
+    }
+  };
+  auto fut = io::thread_pool().enqueue(std::move(read_task)).share();
+  scheduler::enqueue(stream(), [fut = std::move(fut)]() { fut.wait(); });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/matmul.h b/macos/third_party/mlx/mlx/backend/common/matmul.h
new file mode 100644
index 0000000..2545c4f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/matmul.h
@@ -0,0 +1,67 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/utils.h"
+
+#include <sstream>
+
+namespace mlx::core {
+
+inline std::tuple<Shape, Strides, Strides> collapse_batches(
+    const array& a,
+    const array& b) {
+  if (a.ndim() == 2) {
+    return {Shape{1}, Strides{0}, Strides{0}};
+  }
+
+  Shape A_bshape{a.shape().begin(), a.shape().end() - 2};
+  Strides A_bstride{a.strides().begin(), a.strides().end() - 2};
+  Strides B_bstride{b.strides().begin(), b.strides().end() - 2};
+
+  auto [batch_shape, batch_strides] =
+      collapse_contiguous_dims(A_bshape, std::vector{A_bstride, B_bstride});
+
+  auto a_batch_strides = batch_strides[0];
+  auto b_batch_strides = batch_strides[1];
+
+  if (batch_shape.empty()) {
+    batch_shape.push_back(1);
+    a_batch_strides.push_back(0);
+    b_batch_strides.push_back(0);
+  }
+
+  return std::make_tuple(batch_shape, a_batch_strides, b_batch_strides);
+}
+
+inline std::tuple<Shape, Strides, Strides, Strides>
+collapse_batches(const array& a, const array& b, const array& c) {
+  if (a.ndim() == 2) {
+    return {Shape{1}, Strides{0}, Strides{0}, Strides{0}};
+  }
+
+  Shape A_bshape{a.shape().begin(), a.shape().end() - 2};
+  Strides A_bstride{a.strides().begin(), a.strides().end() - 2};
+  Strides B_bstride{b.strides().begin(), b.strides().end() - 2};
+  Strides C_bstride{c.strides().begin(), c.strides().end() - 2};
+
+  auto [batch_shape, batch_strides] = collapse_contiguous_dims(
+      A_bshape, std::vector{A_bstride, B_bstride, C_bstride});
+
+  auto A_batch_stride = batch_strides[0];
+  auto B_batch_stride = batch_strides[1];
+  auto C_batch_stride = batch_strides[2];
+
+  if (batch_shape.empty()) {
+    batch_shape.push_back(1);
+    A_batch_stride.push_back(0);
+    B_batch_stride.push_back(0);
+    C_batch_stride.push_back(0);
+  }
+
+  return std::make_tuple(
+      batch_shape, A_batch_stride, B_batch_stride, C_batch_stride);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/reduce.cpp b/macos/third_party/mlx/mlx/backend/common/reduce.cpp
new file mode 100644
index 0000000..ceef464
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/reduce.cpp
@@ -0,0 +1,154 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/reduce.h"
+
+namespace mlx::core {
+
+std::pair<Shape, Strides> shapes_without_reduction_axes(
+    Shape shape,
+    Strides strides,
+    const std::vector<int>& axes) {
+  for (int i = axes.size() - 1; i >= 0; i--) {
+    int a = axes[i];
+    shape.erase(shape.begin() + a);
+    strides.erase(strides.begin() + a);
+  }
+
+  return std::make_pair(shape, strides);
+}
+
+std::pair<Shape, Strides> shapes_without_reduction_axes(
+    const array& x,
+    const std::vector<int>& axes) {
+  auto shape = x.shape();
+  auto strides = x.strides();
+  return shapes_without_reduction_axes(
+      std::move(shape), std::move(strides), axes);
+}
+
+ReductionPlan get_reduction_plan(const array& x, const std::vector<int>& axes) {
+  // The data is all there and we are reducing over everything
+  if (x.size() == x.data_size() && axes.size() == x.ndim() &&
+      x.flags().contiguous) {
+    return ContiguousAllReduce;
+  }
+
+  // Row contiguous input so the output is row contiguous
+  if (x.flags().row_contiguous) {
+    // Merge consecutive axes
+    Shape shape = {x.shape(axes[0])};
+    Strides strides = {x.strides()[axes[0]]};
+    for (int i = 1; i < axes.size(); i++) {
+      if (axes[i] - 1 == axes[i - 1] && x.shape(axes[i]) > 1) {
+        shape.back() *= x.shape(axes[i]);
+        strides.back() = x.strides()[axes[i]];
+      } else {
+        shape.push_back(x.shape(axes[i]));
+        strides.push_back(x.strides()[axes[i]]);
+      }
+    }
+
+    // Remove singleton axes from the plan
+    for (int i = shape.size() - 1; i >= 0; i--) {
+      if (shape[i] == 1) {
+        shape.erase(shape.begin() + i);
+        strides.erase(strides.begin() + i);
+      }
+    }
+
+    if (strides.back() == 1) {
+      return ReductionPlan(ContiguousReduce, shape, strides);
+    } else if (strides.back() > 1) {
+      return ReductionPlan(ContiguousStridedReduce, shape, strides);
+    }
+  }
+
+  // Let's check if we can optimize our access patterns
+  //
+  // 1. We have a reduction axis with stride 1. Simply call
+  //    GeneralContiguousReduce and be done with it.
+  // 2. We have transpositions and we are not reducing over the axis with
+  //    stride 1. However, we are reducing over an axis where everything is
+  //    contiguous in memory to the right of that axis. We can call strided
+  //    reduce and be done with it.
+  // 2. We have weird transpositions and expands. Copy the strides to the
+  //    output, then call strided reduce.
+
+  // Sort reduction axes by stride in order to merge them and figure out if we
+  // have a contiguous reduction.
+  std::vector<std::pair<int, int64_t>> reductions;
+  for (auto a : axes) {
+    if (x.shape(a) > 1) {
+      reductions.push_back(std::make_pair(x.shape(a), x.strides()[a]));
+    }
+  }
+  std::sort(reductions.begin(), reductions.end(), [](auto a, auto b) {
+    bool a_is_zero = a.second == 0;
+    bool b_is_zero = b.second == 0;
+    return (a_is_zero != b_is_zero) ? a.second < b.second : a.second > b.second;
+  });
+  // Extract the two smallest and try to merge them in case the contiguous
+  // reduction can be bigger than just the last axis.
+  for (int i = reductions.size() - 1; i >= 1; i--) {
+    auto a = reductions[i];
+    auto b = reductions[i - 1];
+
+    // b.stride = a.shape * a.stride then a and b are contiguous
+    if (b.second == a.first * a.second) {
+      reductions.erase(reductions.begin() + i);
+      reductions[i - 1] = std::make_pair(a.first * b.first, a.second);
+    }
+  }
+
+  Shape shape;
+  Strides strides;
+  for (auto r : reductions) {
+    shape.push_back(r.first);
+    strides.push_back(r.second);
+  }
+
+  // We can call the contiguous reduction op for every weird way the input is
+  // structured in the rest of the axes.
+  if (strides.back() == 1) {
+    return ReductionPlan(GeneralContiguousReduce, shape, strides);
+  }
+
+  // Delegate to the general strided reduction op if the axes after
+  // strides.back() are contiguous.
+  if (strides.back() > 1) {
+    int64_t size = 1;
+    bool have_expand = false;
+    for (int i = x.ndim() - 1; i >= 0; i--) {
+      if (axes.back() == i) {
+        continue;
+      }
+
+      auto stride_i = x.strides()[i];
+      auto shape_i = x.shape(i);
+      if (stride_i == 0) {
+        if (shape_i == 1) {
+          continue;
+        }
+
+        have_expand = true;
+        break;
+      }
+
+      if (stride_i != size && shape_i != 1) {
+        break;
+      }
+      size *= shape_i;
+    }
+    // In the case of an expanded dimension we are being conservative and
+    // require the smallest reduction stride to be smaller than the maximum row
+    // contiguous size. The reason is that we can't easily know if the reduced
+    // axis is before or after an expanded dimension.
+    if (size > strides.back() || (size == strides.back() && !have_expand)) {
+      return ReductionPlan(GeneralStridedReduce, shape, strides);
+    }
+  }
+
+  return ReductionPlan(GeneralReduce, shape, strides);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/reduce.h b/macos/third_party/mlx/mlx/backend/common/reduce.h
new file mode 100644
index 0000000..8b24f4f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/reduce.h
@@ -0,0 +1,59 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+enum ReductionOpType {
+  // Self-explanatory. Read everything and produce 1 output.
+  ContiguousAllReduce,
+
+  // The input is contiguous and the last axis is reduced
+  // N1xR1xN2xR2x...xNnxRn
+  ContiguousReduce,
+
+  // The input is contiguous and the last axis is not reduced
+  // R1xN1xR2xN2x...xRnxNn
+  ContiguousStridedReduce,
+
+  // The input is not contiguous but the last axis is and it is reduced so we
+  // need to figure out the offsets but we can call the contiguous reduce after
+  // that.
+  // N3xR1xN1xR4x...xRn
+  GeneralContiguousReduce,
+
+  // The input is not contiguous but the last reduction axis and the last axis
+  // are so we need to figure out the offset but we can call the strided reduce
+  // after that.
+  GeneralStridedReduce,
+
+  // The input is not contiguous after the reduction axis and it may contain
+  // 0-stride axes or transpositions. We could copy the strides and produce a
+  // transposed outcome or we can read the input out of order and write the
+  // output in order.
+  GeneralReduce
+};
+
+struct ReductionPlan {
+  ReductionOpType type;
+  Shape shape;
+  Strides strides;
+
+  ReductionPlan(ReductionOpType type_, Shape shape_, Strides strides_)
+      : type(type_), shape(std::move(shape_)), strides(std::move(strides_)) {}
+  ReductionPlan(ReductionOpType type_) : type(type_) {}
+};
+
+ReductionPlan get_reduction_plan(const array& x, const std::vector<int>& axes);
+
+std::pair<Shape, Strides> shapes_without_reduction_axes(
+    const array& x,
+    const std::vector<int>& axes);
+std::pair<Shape, Strides> shapes_without_reduction_axes(
+    Shape shape,
+    Strides strides,
+    const std::vector<int>& axes);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/slicing.cpp b/macos/third_party/mlx/mlx/backend/common/slicing.cpp
new file mode 100644
index 0000000..cad2fd6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/slicing.cpp
@@ -0,0 +1,71 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+std::tuple<int64_t, Strides> prepare_slice(
+    const array& in,
+    const Shape& start_indices,
+    const Shape& strides) {
+  int64_t data_offset = 0;
+  Strides inp_strides(in.ndim(), 0);
+  for (int i = 0; i < in.ndim(); ++i) {
+    data_offset += start_indices[i] * in.strides()[i];
+    inp_strides[i] = in.strides()[i] * strides[i];
+  }
+  return std::make_tuple(data_offset, inp_strides);
+}
+
+void shared_buffer_slice(
+    const array& in,
+    const Strides& out_strides,
+    int64_t data_offset,
+    size_t data_size,
+    array& out) {
+  // Compute row/col contiguity
+  auto [no_bsx_size, is_row_contiguous, is_col_contiguous] =
+      check_contiguity(out.shape(), out_strides);
+
+  auto flags = in.flags();
+  flags.row_contiguous = is_row_contiguous;
+  flags.col_contiguous = is_col_contiguous;
+  flags.contiguous = (no_bsx_size == data_size);
+
+  out.copy_shared_buffer(in, out_strides, flags, data_size, data_offset);
+}
+
+void slice(
+    const array& in,
+    array& out,
+    const Shape& start_indices,
+    const Shape& strides) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  // Calculate out strides, initial offset
+  auto [data_offset, inp_strides] = prepare_slice(in, start_indices, strides);
+
+  // Get the location of the end based on the inp strides and out.shape()
+  int64_t low_idx = 0;
+  int64_t high_idx = 0;
+  for (int i = 0; i < inp_strides.size(); ++i) {
+    auto delta = inp_strides[i] * (out.shape()[i] - 1);
+    if (inp_strides[i] > 0) {
+      high_idx += delta;
+    } else {
+      low_idx += delta;
+    }
+  }
+  int64_t data_size = (high_idx - low_idx) + 1;
+  if (data_size < 0) {
+    std::ostringstream msg;
+    msg << "[slice] Computed invalid data size: " << data_size << ".";
+    throw std::runtime_error(msg.str());
+  }
+  shared_buffer_slice(in, inp_strides, data_offset, data_size, out);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/slicing.h b/macos/third_party/mlx/mlx/backend/common/slicing.h
new file mode 100644
index 0000000..b667d26
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/slicing.h
@@ -0,0 +1,20 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+std::tuple<int64_t, Strides> prepare_slice(
+    const array& in,
+    const Shape& start_indices,
+    const Shape& strides);
+
+void slice(
+    const array& in,
+    array& out,
+    const Shape& start_indices,
+    const Shape& strides);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/ternary.h b/macos/third_party/mlx/mlx/backend/common/ternary.h
new file mode 100644
index 0000000..c63a572
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/ternary.h
@@ -0,0 +1,85 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+#include "mlx/allocator.h"
+#include "mlx/array.h"
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+// TODO: Add support for more combinations of input types.
+enum class TernaryOpType {
+  ScalarScalarScalar,
+  VectorVectorVector,
+  VectorVectorScalar,
+  VectorScalarVector,
+  General,
+};
+
+inline TernaryOpType
+get_ternary_op_type(const array& a, const array& b, const array& c) {
+  TernaryOpType topt;
+  if (a.data_size() == 1 && b.data_size() == 1 && c.data_size() == 1) {
+    topt = TernaryOpType::ScalarScalarScalar;
+  } else if (
+      (a.flags().row_contiguous && b.flags().row_contiguous &&
+       c.flags().row_contiguous) ||
+      (a.flags().col_contiguous && b.flags().col_contiguous &&
+       c.flags().col_contiguous)) {
+    topt = TernaryOpType::VectorVectorVector;
+  } else if (
+      b.data_size() == 1 && a.flags().row_contiguous &&
+      c.flags().row_contiguous) {
+    topt = TernaryOpType::VectorScalarVector;
+  } else if (
+      c.data_size() == 1 && a.flags().row_contiguous &&
+      b.flags().row_contiguous) {
+    topt = TernaryOpType::VectorVectorScalar;
+  } else {
+    topt = TernaryOpType::General;
+  }
+  return topt;
+}
+
+inline void set_ternary_op_output_data(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    TernaryOpType topt,
+    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+  auto maybe_donate = [&out](const array& x) {
+    if (is_donatable(x, out)) {
+      out.copy_shared_buffer(x);
+      return true;
+    }
+    return false;
+  };
+
+  switch (topt) {
+    case TernaryOpType::ScalarScalarScalar:
+      out.set_data(mallocfn(out.itemsize()), 1, b.strides(), b.flags());
+      break;
+    case TernaryOpType::VectorVectorVector:
+      if (!(maybe_donate(a) || maybe_donate(b) || maybe_donate(c))) {
+        out.set_data(
+            mallocfn(out.itemsize() * b.data_size()),
+            b.data_size(),
+            b.strides(),
+            b.flags());
+      }
+      break;
+    case TernaryOpType::VectorVectorScalar:
+    case TernaryOpType::VectorScalarVector:
+    case TernaryOpType::General:
+      // Try to donate an input which is row_contiguous
+      if (!((a.flags().row_contiguous && maybe_donate(a)) ||
+            (b.flags().row_contiguous && maybe_donate(b)) ||
+            (c.flags().row_contiguous && maybe_donate(c)))) {
+        out.set_data(mallocfn(out.nbytes()));
+      }
+      break;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/unary.h b/macos/third_party/mlx/mlx/backend/common/unary.h
new file mode 100644
index 0000000..b19fc98
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/unary.h
@@ -0,0 +1,29 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+inline void set_unary_output_data(
+    const array& in,
+    array& out,
+    std::function<allocator::Buffer(size_t)> mallocfn = allocator::malloc) {
+  if (in.flags().contiguous) {
+    if (is_donatable(in, out)) {
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(
+          mallocfn(in.data_size() * out.itemsize()),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+    }
+  } else {
+    out.set_data(mallocfn(out.nbytes()));
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/utils.cpp b/macos/third_party/mlx/mlx/backend/common/utils.cpp
new file mode 100644
index 0000000..ae169e3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/utils.cpp
@@ -0,0 +1,231 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <dlfcn.h>
+
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+std::filesystem::path current_binary_dir() {
+  static std::filesystem::path binary_dir = []() {
+    Dl_info info;
+    if (!dladdr(reinterpret_cast<void*>(&current_binary_dir), &info)) {
+      throw std::runtime_error("Unable to get current binary dir.");
+    }
+    return std::filesystem::path(info.dli_fname).parent_path();
+  }();
+  return binary_dir;
+}
+
+std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
+    const Shape& shape,
+    const std::vector<Strides>& strides,
+    int64_t size_cap) {
+  // Make a vector that has axes separated with -1. Collapse all axes between
+  // -1.
+  Shape to_collapse;
+  if (shape.size() > 0) {
+    if (shape[0] != 1) {
+      to_collapse.push_back(0);
+    }
+    size_t size = shape[0];
+    for (int i = 1; i < shape.size(); i++) {
+      bool contiguous = true;
+      size *= shape[i];
+      for (const auto& st : strides) {
+        if (st[i] * shape[i] != st[i - 1] || size > size_cap) {
+          contiguous = false;
+          size = shape[i];
+          break;
+        }
+      }
+      if (!contiguous) {
+        to_collapse.push_back(-1);
+      }
+      if (shape[i] != 1) {
+        to_collapse.push_back(i);
+      }
+    }
+    to_collapse.push_back(-1);
+  }
+
+  Shape out_shape;
+  std::vector<Strides> out_strides(strides.size());
+  for (int i = 0;;) {
+    while (i < to_collapse.size() && to_collapse[i] == -1) {
+      ++i;
+    };
+    if (i == to_collapse.size()) {
+      break;
+    }
+    int current_shape = shape[to_collapse[i]];
+    int k = i;
+    while (to_collapse[++k] != -1) {
+      current_shape *= shape[to_collapse[k]];
+    }
+    out_shape.push_back(current_shape);
+    for (int j = 0; j < strides.size(); j++) {
+      const auto& st = strides[j];
+      out_strides[j].push_back(st[to_collapse[k - 1]]);
+    }
+    i = k + 1;
+  }
+
+  if (!shape.empty() && out_shape.empty()) {
+    out_shape.push_back(1);
+    for (auto& out_stride : out_strides) {
+      out_stride.push_back(0);
+    }
+  }
+  return std::make_tuple(out_shape, out_strides);
+}
+
+std::pair<Shape, Strides> collapse_contiguous_dims(
+    const Shape& shape,
+    const Strides& strides,
+    int64_t size_cap) {
+  Shape collapsed_shape;
+  Strides collapsed_strides;
+
+  if (shape.size() > 0) {
+    collapsed_shape.push_back(shape[0]);
+    collapsed_strides.push_back(strides[0]);
+    for (int i = 1; i < shape.size(); i++) {
+      if (shape[i] == 1) {
+        continue;
+      } else if (
+          strides[i] * shape[i] != collapsed_strides.back() ||
+          collapsed_shape.back() * static_cast<int64_t>(shape[i]) > size_cap) {
+        collapsed_shape.push_back(shape[i]);
+        collapsed_strides.push_back(strides[i]);
+      } else {
+        collapsed_shape.back() *= shape[i];
+        collapsed_strides.back() = strides[i];
+      }
+    }
+  }
+
+  return std::make_pair(collapsed_shape, collapsed_strides);
+}
+
+std::pair<Shape, Strides> collapse_contiguous_dims(
+    const array& a,
+    int64_t size_cap /* = std::numeric_limits<int32_t>::max()*/) {
+  return collapse_contiguous_dims(a.shape(), a.strides(), size_cap);
+}
+
+Dims get_block_dims_common(int dim0, int dim1, int dim2, int pow2 /* = 10 */) {
+  int pows[3] = {0, 0, 0};
+  int sum = 0;
+  while (true) {
+    int presum = sum;
+    // Check all the pows
+    if (dim0 >= (1 << (pows[0] + 1))) {
+      pows[0]++;
+      sum++;
+    }
+    if (sum == 10) {
+      break;
+    }
+    if (dim1 >= (1 << (pows[1] + 1))) {
+      pows[1]++;
+      sum++;
+    }
+    if (sum == 10) {
+      break;
+    }
+    if (dim2 >= (1 << (pows[2] + 1))) {
+      pows[2]++;
+      sum++;
+    }
+    if (sum == presum || sum == pow2) {
+      break;
+    }
+  }
+  return std::make_tuple(1ul << pows[0], 1ul << pows[1], 1ul << pows[2]);
+}
+
+Dims get_2d_grid_dims_common(const Shape& shape, const Strides& strides) {
+  // Dims with strides of 0 are ignored as they
+  // correspond to broadcasted dimensions
+  size_t grid_x = 1;
+  size_t grid_y = 1;
+  for (int i = 0; i < shape.size(); ++i) {
+    if (strides[i] == 0) {
+      continue;
+    }
+    if (grid_x * shape[i] < UINT32_MAX) {
+      grid_x *= shape[i];
+    } else {
+      grid_y *= shape[i];
+    }
+  }
+  if (grid_y > UINT32_MAX || grid_x > UINT32_MAX) {
+    throw std::runtime_error("Unable to safely factor shape.");
+  }
+  if (grid_y > grid_x) {
+    std::swap(grid_x, grid_y);
+  }
+  return std::make_tuple(
+      static_cast<uint32_t>(grid_x), static_cast<uint32_t>(grid_y), 1);
+}
+
+Dims get_2d_grid_dims_common(
+    const Shape& shape,
+    const Strides& strides,
+    size_t divisor) {
+  // Compute the 2d grid dimensions such that the total size of the grid is
+  // divided by divisor.
+  size_t grid_x = 1;
+  size_t grid_y = 1;
+  for (int i = 0; i < shape.size(); ++i) {
+    if (strides[i] == 0) {
+      continue;
+    }
+
+    // No need to add this shape we can just remove it from the divisor.
+    if (divisor % shape[i] == 0) {
+      divisor /= shape[i];
+      continue;
+    }
+
+    if (grid_x * shape[i] < UINT32_MAX) {
+      grid_x *= shape[i];
+    } else {
+      grid_y *= shape[i];
+    }
+
+    if (divisor > 1) {
+      if (grid_x % divisor == 0) {
+        grid_x /= divisor;
+        divisor = 1;
+      } else if (grid_y % divisor == 0) {
+        grid_y /= divisor;
+        divisor = 1;
+      }
+    }
+  }
+  if (grid_y > UINT32_MAX || grid_x > UINT32_MAX) {
+    throw std::runtime_error("Unable to safely factor shape.");
+  }
+  if (grid_y > grid_x) {
+    std::swap(grid_x, grid_y);
+  }
+  if (divisor > 1) {
+    grid_x = ((grid_x + divisor - 1) / divisor) * divisor;
+  }
+  return std::make_tuple(
+      static_cast<uint32_t>(grid_x), static_cast<uint32_t>(grid_y), 1);
+}
+
+std::pair<Dims, Dims> get_grid_and_block_common(int dim0, int dim1, int dim2) {
+  auto [bx, by, bz] = get_block_dims_common(dim0, dim1, dim2);
+  auto gx = (dim0 + bx - 1) / bx;
+  auto gy = (dim1 + by - 1) / by;
+  auto gz = (dim2 + bz - 1) / bz;
+
+  return std::make_pair(
+      std::make_tuple(gx, gy, gz), std::make_tuple(bx, by, bz));
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/common/utils.h b/macos/third_party/mlx/mlx/backend/common/utils.h
new file mode 100644
index 0000000..1b6902f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/common/utils.h
@@ -0,0 +1,205 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <filesystem>
+#include <tuple>
+#include <vector>
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+// Return the directory that contains current shared library.
+std::filesystem::path current_binary_dir();
+
+inline int64_t
+elem_to_loc(int elem, const Shape& shape, const Strides& strides) {
+  int64_t loc = 0;
+  for (int i = shape.size() - 1; i >= 0; --i) {
+    auto q_and_r = ldiv(elem, shape[i]);
+    loc += q_and_r.rem * strides[i];
+    elem = q_and_r.quot;
+  }
+  return loc;
+}
+
+inline int64_t elem_to_loc(int elem, const array& a) {
+  if (a.flags().row_contiguous) {
+    return elem;
+  }
+  return elem_to_loc(elem, a.shape(), a.strides());
+}
+
+inline Strides make_contiguous_strides(const Shape& shape) {
+  Strides strides(shape.size(), 1);
+  for (int i = shape.size() - 1; i > 0; i--) {
+    strides[i - 1] = strides[i] * shape[i];
+  }
+  return strides;
+}
+
+// Collapse dims that are contiguous to possibly route to a better kernel
+// e.g. for x = transpose(array({0, 1, 2, 3, 4, 5, 6, 7}, {2, 2, 2}), {2, 0, 1})
+// should return {{2, 4}, {{1, 2}}}.
+//
+// When multiple arrays are passed they should all have the same shape. The
+// collapsed axes are also the same so one shape is returned.
+std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
+    const Shape& shape,
+    const std::vector<Strides>& strides,
+    int64_t size_cap = std::numeric_limits<int32_t>::max());
+
+inline std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
+    const std::vector<array>& xs,
+    size_t size_cap = std::numeric_limits<int32_t>::max()) {
+  std::vector<Strides> strides;
+  for (auto& x : xs) {
+    strides.emplace_back(x.strides());
+  }
+  return collapse_contiguous_dims(xs[0].shape(), strides, size_cap);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline auto collapse_contiguous_dims(Arrays&&... xs) {
+  return collapse_contiguous_dims(
+      std::vector<array>{std::forward<Arrays>(xs)...});
+}
+
+// The single array version of the above.
+std::pair<Shape, Strides> collapse_contiguous_dims(
+    const Shape& shape,
+    const Strides& strides,
+    int64_t size_cap = std::numeric_limits<int32_t>::max());
+std::pair<Shape, Strides> collapse_contiguous_dims(
+    const array& a,
+    int64_t size_cap = std::numeric_limits<int32_t>::max());
+
+// Compute the thread block dimensions which fit the given
+// input dimensions.
+// - The thread block dimensions will be powers of two
+// - The thread block size will be less than 2^pow2
+using Dims = std::tuple<uint32_t, uint32_t, uint32_t>;
+Dims get_block_dims_common(int dim0, int dim1, int dim2, int pow2 = 10);
+
+// Computes a 2D grid where each element is < UINT_MAX
+// Assumes:
+// - overall size (product of non-broadcasted dimensions) is < UINT_MAX^2
+// - shape and strides correspond to a contiguous (no holes) but
+//   possibly broadcasted array
+Dims get_2d_grid_dims_common(const Shape& shape, const Strides& strides);
+
+// Same as above but we do an implicit division with divisor.
+// Basically, equivalent to factorizing
+//    Prod(s \forall s in shape if strides[s] > 0) / divisor.
+Dims get_2d_grid_dims_common(
+    const Shape& shape,
+    const Strides& strides,
+    size_t divisor);
+
+// Get both the block and a grid of blocks that covers dim0, dim1 and dim2.
+std::pair<Dims, Dims> get_grid_and_block_common(int dim0, int dim1, int dim2);
+
+struct ContiguousIterator {
+  inline void step() {
+    int dims = shape_.size();
+    if (dims == 0) {
+      return;
+    }
+    int i = dims - 1;
+    while (pos_[i] == (shape_[i] - 1) && i > 0) {
+      pos_[i] = 0;
+      loc -= (shape_[i] - 1) * strides_[i];
+      i--;
+    }
+    pos_[i]++;
+    loc += strides_[i];
+  }
+
+  void seek(int64_t n) {
+    loc = 0;
+    for (int i = shape_.size() - 1; i >= 0; --i) {
+      auto q_and_r = ldiv(n, shape_[i]);
+      loc += q_and_r.rem * strides_[i];
+      pos_[i] = q_and_r.rem;
+      n = q_and_r.quot;
+    }
+  }
+
+  void reset() {
+    loc = 0;
+    std::fill(pos_.begin(), pos_.end(), 0);
+  }
+
+  ContiguousIterator() {};
+
+  explicit ContiguousIterator(const array& a)
+      : shape_(a.shape()), strides_(a.strides()) {
+    if (!shape_.empty()) {
+      std::tie(shape_, strides_) = collapse_contiguous_dims(shape_, strides_);
+      pos_ = Shape(shape_.size(), 0);
+    }
+  }
+
+  explicit ContiguousIterator(
+      const Shape& shape,
+      const Strides& strides,
+      int dims)
+      : shape_(shape.begin(), shape.begin() + dims),
+        strides_(strides.begin(), strides.begin() + dims) {
+    if (!shape_.empty()) {
+      std::tie(shape_, strides_) = collapse_contiguous_dims(shape_, strides_);
+      pos_ = Shape(shape_.size(), 0);
+    }
+  }
+
+  int64_t loc{0};
+
+ private:
+  Shape shape_;
+  Strides strides_;
+  Shape pos_;
+};
+
+inline auto check_contiguity(const Shape& shape, const Strides& strides) {
+  size_t no_broadcast_data_size = 1;
+  int64_t f_stride = 1;
+  int64_t b_stride = 1;
+  bool is_row_contiguous = true;
+  bool is_col_contiguous = true;
+
+  for (int i = 0, ri = shape.size() - 1; ri >= 0; i++, ri--) {
+    is_col_contiguous &= strides[i] == f_stride || shape[i] == 1;
+    is_row_contiguous &= strides[ri] == b_stride || shape[ri] == 1;
+    f_stride *= shape[i];
+    b_stride *= shape[ri];
+    if (strides[i] > 0) {
+      no_broadcast_data_size *= shape[i];
+    }
+  }
+
+  return std::make_tuple(
+      no_broadcast_data_size, is_row_contiguous, is_col_contiguous);
+}
+
+inline bool is_donatable(const array& in, const array& out) {
+  constexpr size_t donation_extra = 16384;
+
+  return in.is_donatable() && in.itemsize() == out.itemsize() &&
+      in.buffer_size() <= out.nbytes() + donation_extra;
+}
+
+std::pair<bool, Strides> prepare_reshape(const array& in, const array& out);
+
+void shared_buffer_reshape(
+    const array& in,
+    const Strides& out_strides,
+    array& out);
+
+template <typename T>
+inline SmallVector<T> remove_index(SmallVector<T> vec, size_t index) {
+  vec.erase(std::next(vec.begin(), index));
+  return vec;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/cpu/CMakeLists.txt
new file mode 100644
index 0000000..f2c9694
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/CMakeLists.txt
@@ -0,0 +1,88 @@
+if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+  set(COMPILER ${CMAKE_C_COMPILER})
+  set(CLANG TRUE)
+else()
+  set(COMPILER ${CMAKE_CXX_COMPILER})
+endif()
+
+set(COMPILE_DEPS
+    ${PROJECT_SOURCE_DIR}/mlx/types/half_types.h
+    ${PROJECT_SOURCE_DIR}/mlx/types/fp16.h
+    ${PROJECT_SOURCE_DIR}/mlx/types/bf16.h
+    ${PROJECT_SOURCE_DIR}/mlx/types/complex.h
+    simd/simd.h
+    simd/base_simd.h
+    simd/math.h
+    simd/type.h
+    unary_ops.h
+    binary_ops.h)
+
+if(MSVC)
+  set(SHELL_EXT ps1)
+  set(SHELL_CMD powershell -ExecutionPolicy Bypass -File)
+else()
+  set(SHELL_EXT sh)
+  set(SHELL_CMD bash)
+endif()
+
+add_custom_command(
+  OUTPUT compiled_preamble.cpp
+  COMMAND
+    ${SHELL_CMD} ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.${SHELL_EXT}
+    ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp ${COMPILER}
+    ${PROJECT_SOURCE_DIR} ${CLANG} ${CMAKE_SYSTEM_PROCESSOR}
+  DEPENDS make_compiled_preamble.${SHELL_EXT} compiled_preamble.h
+          ${COMPILE_DEPS})
+
+add_custom_target(cpu_compiled_preamble DEPENDS compiled_preamble.cpp)
+
+add_dependencies(mlx cpu_compiled_preamble)
+
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eig.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eigh.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/encoder.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cblas.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/masked_mm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/select.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/luf.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/qrf.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/svd.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/inverse.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/cholesky.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/unary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
+          ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp)
+
+if(MLX_BUILD_ACCELERATE)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/bnns.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/simd_fp16.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/gemms/simd_bf16.cpp)
+endif()
+
+if(IOS)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../no_cpu/compiled.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/jit_compiler.cpp)
+endif()
diff --git a/macos/third_party/mlx/mlx/backend/cpu/arange.h b/macos/third_party/mlx/mlx/backend/cpu/arange.h
new file mode 100644
index 0000000..9e9b03b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/arange.h
@@ -0,0 +1,28 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cpu/encoder.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T>
+void arange(T start, T next, array& out, size_t size, Stream stream) {
+  auto ptr = out.data<T>();
+  auto step_size = next - start;
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([ptr, start, step_size, size]() mutable {
+    for (int i = 0; i < size; ++i) {
+      ptr[i] = start;
+      start += step_size;
+    }
+  });
+}
+
+} // namespace
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/arg_reduce.cpp b/macos/third_party/mlx/mlx/backend/cpu/arg_reduce.cpp
new file mode 100644
index 0000000..6646891
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/arg_reduce.cpp
@@ -0,0 +1,124 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename InT, typename OpT>
+void arg_reduce(const array& in, array& out, const OpT& op, int axis) {
+  auto axis_size = in.shape()[axis];
+  auto axis_stride = in.strides()[axis];
+  Strides strides = remove_index(in.strides(), axis);
+  Shape shape = remove_index(in.shape(), axis);
+  auto in_ptr = in.data<InT>();
+  auto out_ptr = out.data<uint32_t>();
+
+  for (uint32_t i = 0; i < out.size(); ++i) {
+    auto loc = elem_to_loc(i, shape, strides);
+    auto local_in_ptr = in_ptr + loc;
+    uint32_t ind_v = 0;
+    InT v = (*local_in_ptr);
+    for (uint32_t j = 0; j < axis_size; ++j, local_in_ptr += axis_stride) {
+      op(j, (*local_in_ptr), &ind_v, &v);
+    }
+    out_ptr[i] = ind_v;
+  }
+}
+
+template <typename InT>
+void arg_reduce_dispatch(
+    const array& in,
+    array& out,
+    ArgReduce::ReduceType rtype,
+    int axis) {
+  switch (rtype) {
+    case ArgReduce::ArgMin: {
+      auto op = [](auto ind_x, auto x, auto ind_y, auto y) {
+        if (x < (*y)) {
+          (*y) = x;
+          (*ind_y) = ind_x;
+        }
+      };
+      arg_reduce<InT>(in, out, op, axis);
+      break;
+    }
+    case ArgReduce::ArgMax: {
+      auto op = [](auto ind_x, auto x, auto ind_y, auto y) {
+        if (x > (*y)) {
+          (*y) = x;
+          (*ind_y) = ind_x;
+        }
+      };
+      arg_reduce<InT>(in, out, op, axis);
+      break;
+    }
+  }
+}
+
+} // namespace
+
+void ArgReduce::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    reduce_type_ = reduce_type_,
+                    axis_ = axis_]() mutable {
+    switch (in.dtype()) {
+      case bool_:
+        arg_reduce_dispatch<bool>(in, out, reduce_type_, axis_);
+        break;
+      case uint8:
+        arg_reduce_dispatch<uint8_t>(in, out, reduce_type_, axis_);
+        break;
+      case uint16:
+        arg_reduce_dispatch<uint16_t>(in, out, reduce_type_, axis_);
+        break;
+      case uint32:
+        arg_reduce_dispatch<uint32_t>(in, out, reduce_type_, axis_);
+        break;
+      case uint64:
+        arg_reduce_dispatch<uint64_t>(in, out, reduce_type_, axis_);
+        break;
+      case int8:
+        arg_reduce_dispatch<int8_t>(in, out, reduce_type_, axis_);
+        break;
+      case int16:
+        arg_reduce_dispatch<int16_t>(in, out, reduce_type_, axis_);
+        break;
+      case int32:
+        arg_reduce_dispatch<int32_t>(in, out, reduce_type_, axis_);
+        break;
+      case int64:
+        arg_reduce_dispatch<int64_t>(in, out, reduce_type_, axis_);
+        break;
+      case float16:
+        arg_reduce_dispatch<float16_t>(in, out, reduce_type_, axis_);
+        break;
+      case float32:
+        arg_reduce_dispatch<float>(in, out, reduce_type_, axis_);
+        break;
+      case bfloat16:
+        arg_reduce_dispatch<bfloat16_t>(in, out, reduce_type_, axis_);
+        break;
+      case float64:
+        arg_reduce_dispatch<double>(in, out, reduce_type_, axis_);
+        break;
+      case complex64:
+        arg_reduce_dispatch<complex64_t>(in, out, reduce_type_, axis_);
+        break;
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/binary.cpp b/macos/third_party/mlx/mlx/backend/cpu/binary.cpp
new file mode 100644
index 0000000..c1db2e1
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/binary.cpp
@@ -0,0 +1,269 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+#include <cmath>
+#include <sstream>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/binary.h"
+#include "mlx/backend/cpu/binary_ops.h"
+#include "mlx/backend/cpu/binary_two.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Add(), stream());
+}
+
+void DivMod::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto bopt = get_binary_op_type(a, b);
+  auto& out_a = outputs[0];
+  auto& out_b = outputs[1];
+  set_binary_op_output_data(a, b, out_a, bopt);
+  set_binary_op_output_data(a, b, out_b, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out_a);
+  encoder.set_output_array(out_b);
+
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out_a = array::unsafe_weak_copy(out_a),
+                    out_b = array::unsafe_weak_copy(out_b),
+                    bopt]() mutable {
+    auto integral_op = [](auto x, auto y) {
+      return std::make_pair(x / y, x % y);
+    };
+    auto float_op = [](auto x, auto y) {
+      return std::make_pair(std::trunc(x / y), std::fmod(x, y));
+    };
+
+    switch (out_a.dtype()) {
+      case bool_:
+        binary_op<bool>(a, b, out_a, out_b, integral_op, bopt);
+      case uint8:
+        binary_op<uint8_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case uint16:
+        binary_op<uint16_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case uint32:
+        binary_op<uint32_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case uint64:
+        binary_op<uint64_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case int8:
+        binary_op<int8_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case int16:
+        binary_op<int16_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case int32:
+        binary_op<int32_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case int64:
+        binary_op<int64_t>(a, b, out_a, out_b, integral_op, bopt);
+        break;
+      case float16:
+        binary_op<float16_t>(a, b, out_a, out_b, float_op, bopt);
+        break;
+      case float32:
+        binary_op<float>(a, b, out_a, out_b, float_op, bopt);
+        break;
+      case float64:
+        binary_op<double>(a, b, out_a, out_b, float_op, bopt);
+        break;
+      case bfloat16:
+        binary_op<bfloat16_t>(a, b, out_a, out_b, float_op, bopt);
+        break;
+      case complex64:
+        // Should never get here
+        throw std::runtime_error("[DivMod] Complex type not supported");
+        break;
+    }
+  });
+}
+
+void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Divide(), stream());
+}
+
+void Remainder::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Remainder(), stream());
+}
+
+void Equal::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  if (equal_nan_) {
+    auto bopt = get_binary_op_type(a, b);
+    set_binary_op_output_data(a, b, out, bopt);
+
+    auto& encoder = cpu::get_command_encoder(stream());
+    encoder.set_input_array(a);
+    encoder.set_input_array(b);
+    encoder.set_output_array(out);
+    encoder.dispatch([a = array::unsafe_weak_copy(a),
+                      b = array::unsafe_weak_copy(b),
+                      out = array::unsafe_weak_copy(out),
+                      bopt]() mutable {
+      switch (a.dtype()) {
+        case float16:
+          binary_op<float16_t, bool, detail::NaNEqual>(a, b, out, bopt);
+          break;
+        case float32:
+          binary_op<float, bool, detail::NaNEqual>(a, b, out, bopt);
+          break;
+        case float64:
+          binary_op<double, bool, detail::NaNEqual>(a, b, out, bopt);
+          break;
+        case bfloat16:
+          binary_op<bfloat16_t, bool, detail::NaNEqual>(a, b, out, bopt);
+          break;
+        case complex64:
+          binary_op<complex64_t, bool, detail::NaNEqual>(a, b, out, bopt);
+          break;
+        default:
+          throw std::runtime_error(
+              "[NanEqual::eval_cpu] Only for floating point types.");
+      }
+    });
+  } else {
+    comparison_op_cpu(a, b, out, detail::Equal(), stream());
+  }
+}
+
+void Greater::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  comparison_op_cpu(inputs[0], inputs[1], out, detail::Greater(), stream());
+}
+
+void GreaterEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  comparison_op_cpu(
+      inputs[0], inputs[1], out, detail::GreaterEqual(), stream());
+}
+
+void Less::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  comparison_op_cpu(inputs[0], inputs[1], out, detail::Less(), stream());
+}
+
+void LessEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  comparison_op_cpu(inputs[0], inputs[1], out, detail::LessEqual(), stream());
+}
+
+void LogAddExp::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_float_op_cpu(a, b, out, detail::LogAddExp(), stream());
+}
+
+void LogicalAnd::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2); // LogicalAnd requires two input arrays
+  auto& in1 = inputs[0];
+  auto& in2 = inputs[1];
+  binary_op_cpu(in1, in2, out, detail::LogicalAnd(), stream());
+}
+
+void LogicalOr::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2); // LogicalOr requires two input arrays
+  auto& in1 = inputs[0];
+  auto& in2 = inputs[1];
+  binary_op_cpu(in1, in2, out, detail::LogicalOr(), stream());
+}
+
+void Maximum::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Maximum(), stream());
+}
+
+void Minimum::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Minimum(), stream());
+}
+
+void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Multiply(), stream());
+}
+
+void NotEqual::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  comparison_op_cpu(inputs[0], inputs[1], out, detail::NotEqual(), stream());
+}
+
+void Power::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Power(), stream());
+}
+
+void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  binary_op_cpu(a, b, out, detail::Subtract(), stream());
+}
+
+void BitwiseBinary::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  switch (op_) {
+    case BitwiseBinary::And:
+      binary_int_op_cpu(a, b, out, detail::BitwiseAnd(), stream());
+      break;
+    case BitwiseBinary::Or:
+      binary_int_op_cpu(a, b, out, detail::BitwiseOr(), stream());
+      break;
+    case BitwiseBinary::Xor:
+      binary_int_op_cpu(a, b, out, detail::BitwiseXor(), stream());
+      break;
+    case BitwiseBinary::LeftShift:
+      binary_int_op_cpu(a, b, out, detail::LeftShift(), stream());
+      break;
+    case BitwiseBinary::RightShift:
+      binary_int_op_cpu(a, b, out, detail::RightShift(), stream());
+      break;
+  }
+}
+
+void ArcTan2::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  const auto& a = inputs[0];
+  const auto& b = inputs[1];
+  binary_float_op_cpu(a, b, out, detail::ArcTan2(), stream());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/binary.h b/macos/third_party/mlx/mlx/backend/cpu/binary.h
new file mode 100644
index 0000000..acaca50
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/binary.h
@@ -0,0 +1,517 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+#include <cassert>
+
+#include "mlx/array.h"
+#include "mlx/backend/common/binary.h"
+#include "mlx/backend/common/utils.h"
+
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+
+namespace mlx::core {
+
+template <typename Op>
+struct VectorScalar {
+  template <typename T, typename U>
+  void operator()(const T* a, const T* b, U* dst, int size) {
+    T scalar = *b;
+    constexpr int N = simd::max_size<T>;
+    while (size >= N) {
+      simd::store(dst, Op{}(simd::load<T, N>(a), simd::Simd<T, N>(scalar)));
+      dst += N;
+      a += N;
+      size -= N;
+    }
+    while (size-- > 0) {
+      *dst = Op{}(*a, scalar);
+      dst++;
+      a++;
+    }
+  }
+};
+
+template <typename Op>
+struct ScalarVector {
+  template <typename T, typename U>
+  void operator()(const T* a, const T* b, U* dst, int size) {
+    T scalar = *a;
+    constexpr int N = simd::max_size<T>;
+    while (size >= N) {
+      simd::store(dst, Op{}(simd::Simd<T, N>(scalar), simd::load<T, N>(b)));
+      dst += N;
+      b += N;
+      size -= N;
+    }
+    while (size-- > 0) {
+      *dst = Op{}(scalar, *b);
+      dst++;
+      b++;
+    }
+  }
+};
+
+template <typename Op>
+struct VectorVector {
+  template <typename T, typename U>
+  void operator()(const T* a, const T* b, U* dst, int size) {
+    constexpr int N = simd::max_size<T>;
+    while (size >= N) {
+      simd::store(dst, Op{}(simd::load<T, N>(a), simd::load<T, N>(b)));
+      dst += N;
+      a += N;
+      b += N;
+      size -= N;
+    }
+    while (size-- > 0) {
+      *dst = Op{}(*a, *b);
+      dst++;
+      a++;
+      b++;
+    }
+  }
+};
+
+template <typename T, typename U, typename Op, int D, bool Strided>
+void binary_op_dims(
+    const T* a,
+    const T* b,
+    U* out,
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& out_strides,
+    int axis) {
+  auto stride_a = a_strides[axis];
+  auto stride_b = b_strides[axis];
+  auto stride_out = out_strides[axis];
+  auto N = shape[axis];
+
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      binary_op_dims<T, U, Op, D - 1, Strided>(
+          a, b, out, shape, a_strides, b_strides, out_strides, axis + 1);
+    } else {
+      if constexpr (Strided) {
+        Op{}(a, b, out, stride_out);
+      } else {
+        *out = Op{}(*a, *b);
+      }
+    }
+    out += stride_out;
+    a += stride_a;
+    b += stride_b;
+  }
+}
+
+template <typename T, typename U, bool Strided, typename Op>
+void binary_op_dispatch_dims(
+    const T* a,
+    const T* b,
+    U* out,
+    int dim,
+    int size,
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& out_strides) {
+  switch (dim) {
+    case 1:
+      binary_op_dims<T, U, Op, 1, Strided>(
+          a, b, out, shape, a_strides, b_strides, out_strides, 0);
+      return;
+    case 2:
+      binary_op_dims<T, U, Op, 2, Strided>(
+          a, b, out, shape, a_strides, b_strides, out_strides, 0);
+      return;
+    case 3:
+      binary_op_dims<T, U, Op, 3, Strided>(
+          a, b, out, shape, a_strides, b_strides, out_strides, 0);
+      return;
+  }
+
+  ContiguousIterator a_it(shape, a_strides, dim - 3);
+  ContiguousIterator b_it(shape, b_strides, dim - 3);
+  auto stride = out_strides[dim - 4];
+  for (int64_t elem = 0; elem < size; elem += stride) {
+    binary_op_dims<T, U, Op, 3, Strided>(
+        a + a_it.loc,
+        b + b_it.loc,
+        out + elem,
+        shape,
+        a_strides,
+        b_strides,
+        out_strides,
+        dim - 3);
+    a_it.step();
+    b_it.step();
+  }
+}
+
+template <typename T, typename U, typename Op>
+void binary_op(const array& a, const array& b, array& out, BinaryOpType bopt) {
+  // The full computation is scalar scalar so call the base op once
+  auto a_ptr = a.data<T>();
+  auto b_ptr = b.data<T>();
+
+  auto out_ptr = out.data<U>();
+  if (bopt == BinaryOpType::ScalarScalar) {
+    *out_ptr = Op{}(*a_ptr, *b_ptr);
+    return;
+  }
+
+  // The full computation is scalar vector so delegate to the op
+  if (bopt == BinaryOpType::ScalarVector) {
+    ScalarVector<Op>{}(a_ptr, b_ptr, out_ptr, b.data_size());
+    return;
+  }
+
+  // The full computation is vector scalar so delegate to the op
+  if (bopt == BinaryOpType::VectorScalar) {
+    VectorScalar<Op>{}(a_ptr, b_ptr, out_ptr, a.data_size());
+    return;
+  }
+
+  // The full computation is vector vector so delegate to the op
+  if (bopt == BinaryOpType::VectorVector) {
+    VectorVector<Op>{}(a_ptr, b_ptr, out_ptr, a.size());
+    return;
+  }
+
+  // General computation so let's try to optimize
+  auto [new_shape, new_strides] = collapse_contiguous_dims(
+      a.shape(), {a.strides(), b.strides(), out.strides()});
+  auto& a_strides = new_strides[0];
+  auto& b_strides = new_strides[1];
+  auto& strides = new_strides[2];
+
+  // Get the left-most dim such that the array is row contiguous after
+  auto leftmost_rc_dim = [&strides](const auto& arr_strides) {
+    int d = arr_strides.size() - 1;
+    for (; d >= 0 && arr_strides[d] == strides[d]; d--) {
+    }
+    return d + 1;
+  };
+  auto a_rc_dim = leftmost_rc_dim(a_strides);
+  auto b_rc_dim = leftmost_rc_dim(b_strides);
+
+  // Get the left-most dim such that the array is a broadcasted "scalar" after
+  auto leftmost_s_dim = [](const auto& arr_strides) {
+    int d = arr_strides.size() - 1;
+    for (; d >= 0 && arr_strides[d] == 0; d--) {
+    }
+    return d + 1;
+  };
+  auto a_s_dim = leftmost_s_dim(a_strides);
+  auto b_s_dim = leftmost_s_dim(b_strides);
+
+  auto ndim = new_shape.size();
+
+  // Case 1: LxM and FxM where L and F are broadcastable and M is row
+  // contiguous
+  int dim = ndim;
+  if (int d = std::max(a_rc_dim, b_rc_dim); d < ndim) {
+    bopt = BinaryOpType::VectorVector;
+    dim = d;
+    // Case 2: LxM and Fx1 where L and F are broadcastable and M is row
+    // contiguous
+  } else if (int d = std::max(a_rc_dim, b_s_dim); d < ndim) {
+    bopt = BinaryOpType::VectorScalar;
+    dim = d;
+    // Case 3: Lx1 and FxM where L and F are broadcastable and M is row
+    // contiguous
+  } else if (int d = std::max(a_s_dim, b_rc_dim); d < ndim) {
+    bopt = BinaryOpType::ScalarVector;
+    dim = d;
+  }
+
+  // Can be sure dim > 0 since otherwise we would have used one of the fully
+  // contiguous methods above. Except for the case that the flags do not
+  // correspond to the underlying contiguity.
+  if (dim == 0 || strides[dim - 1] < 16) {
+    bopt = BinaryOpType::General;
+    dim = ndim;
+  }
+
+  switch (bopt) {
+    case BinaryOpType::VectorVector:
+      binary_op_dispatch_dims<T, U, true, VectorVector<Op>>(
+          a_ptr,
+          b_ptr,
+          out_ptr,
+          dim,
+          a.size(),
+          new_shape,
+          a_strides,
+          b_strides,
+          strides);
+      break;
+    case BinaryOpType::VectorScalar:
+      binary_op_dispatch_dims<T, U, true, VectorScalar<Op>>(
+          a_ptr,
+          b_ptr,
+          out_ptr,
+          dim,
+          a.size(),
+          new_shape,
+          a_strides,
+          b_strides,
+          strides);
+      break;
+    case BinaryOpType::ScalarVector:
+      binary_op_dispatch_dims<T, U, true, ScalarVector<Op>>(
+          a_ptr,
+          b_ptr,
+          out_ptr,
+          dim,
+          a.size(),
+          new_shape,
+          a_strides,
+          b_strides,
+          strides);
+      break;
+    default:
+      binary_op_dispatch_dims<T, U, false, Op>(
+          a_ptr,
+          b_ptr,
+          out_ptr,
+          dim,
+          a.size(),
+          new_shape,
+          a_strides,
+          b_strides,
+          strides);
+      break;
+  }
+}
+
+template <typename T, typename Op>
+void binary_op(const array& a, const array& b, array& out, BinaryOpType bopt) {
+  binary_op<T, T, Op>(a, b, out, bopt);
+}
+
+template <typename Op>
+void binary_op_cpu(
+    const array& a,
+    const array& b,
+    array& out,
+    Op op,
+    Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        binary_op<bool, Op>(a, b, out, bopt);
+        break;
+      case uint8:
+        binary_op<uint8_t, Op>(a, b, out, bopt);
+        break;
+      case uint16:
+        binary_op<uint16_t, Op>(a, b, out, bopt);
+        break;
+      case uint32:
+        binary_op<uint32_t, Op>(a, b, out, bopt);
+        break;
+      case uint64:
+        binary_op<uint64_t, Op>(a, b, out, bopt);
+        break;
+      case int8:
+        binary_op<int8_t, Op>(a, b, out, bopt);
+        break;
+      case int16:
+        binary_op<int16_t, Op>(a, b, out, bopt);
+        break;
+      case int32:
+        binary_op<int32_t, Op>(a, b, out, bopt);
+        break;
+      case int64:
+        binary_op<int64_t, Op>(a, b, out, bopt);
+        break;
+      case float16:
+        binary_op<float16_t, Op>(a, b, out, bopt);
+        break;
+      case float32:
+        binary_op<float, Op>(a, b, out, bopt);
+        break;
+      case float64:
+        binary_op<double, Op>(a, b, out, bopt);
+        break;
+      case bfloat16:
+        binary_op<bfloat16_t, Op>(a, b, out, bopt);
+        break;
+      case complex64:
+        binary_op<complex64_t, Op>(a, b, out, bopt);
+        break;
+    }
+  });
+}
+
+template <typename Op>
+void comparison_op_cpu(
+    const array& a,
+    const array& b,
+    array& out,
+    Op op,
+    Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (a.dtype()) {
+      case bool_:
+        binary_op<bool, bool, Op>(a, b, out, bopt);
+        break;
+      case uint8:
+        binary_op<uint8_t, bool, Op>(a, b, out, bopt);
+        break;
+      case uint16:
+        binary_op<uint16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case uint32:
+        binary_op<uint32_t, bool, Op>(a, b, out, bopt);
+        break;
+      case uint64:
+        binary_op<uint64_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int8:
+        binary_op<int8_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int16:
+        binary_op<int16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int32:
+        binary_op<int32_t, bool, Op>(a, b, out, bopt);
+        break;
+      case int64:
+        binary_op<int64_t, bool, Op>(a, b, out, bopt);
+        break;
+      case float16:
+        binary_op<float16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case float32:
+        binary_op<float, bool, Op>(a, b, out, bopt);
+        break;
+      case float64:
+        binary_op<double, bool, Op>(a, b, out, bopt);
+        break;
+      case bfloat16:
+        binary_op<bfloat16_t, bool, Op>(a, b, out, bopt);
+        break;
+      case complex64:
+        binary_op<complex64_t, bool, Op>(a, b, out, bopt);
+        break;
+    }
+  });
+}
+
+template <typename Op>
+void binary_float_op_cpu(
+    const array& a,
+    const array& b,
+    array& out,
+    Op op,
+    Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (out.dtype()) {
+      case float16:
+        binary_op<float16_t, Op>(a, b, out, bopt);
+        break;
+      case float32:
+        binary_op<float, Op>(a, b, out, bopt);
+        break;
+      case float64:
+        binary_op<double, Op>(a, b, out, bopt);
+        break;
+      case bfloat16:
+        binary_op<bfloat16_t, Op>(a, b, out, bopt);
+        break;
+      case complex64:
+        binary_op<complex64_t, Op>(a, b, out, bopt);
+        break;
+      default:
+        throw std::runtime_error(
+            "[binary_float] Only supports floating point types.");
+    }
+  });
+}
+
+template <typename Op>
+void binary_int_op_cpu(
+    const array& a,
+    const array& b,
+    array& out,
+    Op op,
+    Stream stream) {
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    out = array::unsafe_weak_copy(out),
+                    bopt]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        binary_op<bool, Op>(a, b, out, bopt);
+      case uint8:
+        binary_op<uint8_t, Op>(a, b, out, bopt);
+        break;
+      case uint16:
+        binary_op<uint16_t, Op>(a, b, out, bopt);
+        break;
+      case uint32:
+        binary_op<uint32_t, Op>(a, b, out, bopt);
+        break;
+      case uint64:
+        binary_op<uint64_t, Op>(a, b, out, bopt);
+        break;
+      case int8:
+        binary_op<int8_t, Op>(a, b, out, bopt);
+        break;
+      case int16:
+        binary_op<int16_t, Op>(a, b, out, bopt);
+        break;
+      case int32:
+        binary_op<int32_t, Op>(a, b, out, bopt);
+        break;
+      case int64:
+        binary_op<int64_t, Op>(a, b, out, bopt);
+        break;
+      default:
+        throw std::runtime_error("[binary_int] Type not supported");
+        break;
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/binary_ops.h b/macos/third_party/mlx/mlx/backend/cpu/binary_ops.h
new file mode 100644
index 0000000..d50751c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/binary_ops.h
@@ -0,0 +1,98 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cpu/simd/simd.h"
+
+namespace mlx::core::detail {
+
+using namespace mlx::core::simd;
+
+#define BINARY_SINGLE()                                 \
+  template <typename T>                                 \
+  T operator()(T x, T y) {                              \
+    return (*this)(Simd<T, 1>(x), Simd<T, 1>(y)).value; \
+  }
+
+#define DEFAULT_BINARY_OP(Op, op)                       \
+  struct Op {                                           \
+    template <int N, typename T>                        \
+    Simd<T, N> operator()(Simd<T, N> x, Simd<T, N> y) { \
+      return op(x, y);                                  \
+    }                                                   \
+    BINARY_SINGLE()                                     \
+  };
+
+DEFAULT_BINARY_OP(Add, operator+)
+DEFAULT_BINARY_OP(ArcTan2, atan2)
+DEFAULT_BINARY_OP(Divide, operator/)
+DEFAULT_BINARY_OP(Multiply, operator*)
+DEFAULT_BINARY_OP(Subtract, operator-)
+DEFAULT_BINARY_OP(LogicalAnd, operator&&)
+DEFAULT_BINARY_OP(LogicalOr, operator||)
+DEFAULT_BINARY_OP(BitwiseAnd, operator&)
+DEFAULT_BINARY_OP(BitwiseOr, operator|)
+DEFAULT_BINARY_OP(BitwiseXor, operator^)
+DEFAULT_BINARY_OP(LeftShift, operator<<)
+DEFAULT_BINARY_OP(RightShift, operator>>)
+DEFAULT_BINARY_OP(Remainder, remainder)
+DEFAULT_BINARY_OP(Maximum, maximum)
+DEFAULT_BINARY_OP(Minimum, minimum)
+DEFAULT_BINARY_OP(Power, pow)
+
+#define DEFAULT_BOOL_OP(Op, op)                            \
+  struct Op {                                              \
+    template <int N, typename T>                           \
+    Simd<bool, N> operator()(Simd<T, N> x, Simd<T, N> y) { \
+      return op(x, y);                                     \
+    }                                                      \
+    template <typename T>                                  \
+    bool operator()(T x, T y) {                            \
+      return (*this)(Simd<T, 1>(x), Simd<T, 1>(y)).value;  \
+    }                                                      \
+  };
+
+DEFAULT_BOOL_OP(Equal, operator==)
+DEFAULT_BOOL_OP(Greater, operator>)
+DEFAULT_BOOL_OP(GreaterEqual, operator>=)
+DEFAULT_BOOL_OP(Less, operator<)
+DEFAULT_BOOL_OP(LessEqual, operator<=)
+DEFAULT_BOOL_OP(NotEqual, operator!=)
+
+struct NaNEqual {
+  template <int N, typename T>
+  Simd<bool, N> operator()(Simd<T, N> x, Simd<T, N> y) {
+    return x == y || (isnan(x) && isnan(y));
+  }
+  template <typename T>
+  bool operator()(T x, T y) {
+    return (*this)(Simd<T, 1>(x), Simd<T, 1>(y)).value;
+  }
+};
+
+struct LogAddExp {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x, Simd<T, N> y) {
+    auto maxval = maximum(x, y);
+    auto minval = minimum(x, y);
+    auto mask = minval == -inf || maxval == inf;
+    auto out = maxval + log1p(exp(minval - maxval));
+    return select(mask, Simd<T, N>(maxval), Simd<T, N>(out));
+  }
+  BINARY_SINGLE()
+};
+
+struct Select {
+  template <typename T>
+  T operator()(bool condition, T x, T y) {
+    return (*this)(Simd<bool, 1>(condition), Simd<T, 1>(x), Simd<T, 1>(y))
+        .value;
+  }
+
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<bool, N> condition, Simd<T, N> x, Simd<T, N> y) {
+    return select(condition, x, y);
+  }
+};
+
+} // namespace mlx::core::detail
diff --git a/macos/third_party/mlx/mlx/backend/cpu/binary_two.h b/macos/third_party/mlx/mlx/backend/cpu/binary_two.h
new file mode 100644
index 0000000..fa0ca79
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/binary_two.h
@@ -0,0 +1,166 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/binary.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T, typename U, typename Op, int D>
+void binary_op_dims(
+    const T* a,
+    const T* b,
+    U* out_a,
+    U* out_b,
+    Op op,
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& out_strides,
+    int axis) {
+  auto stride_a = a_strides[axis];
+  auto stride_b = b_strides[axis];
+  auto stride_out = out_strides[axis];
+  auto N = shape[axis];
+
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      binary_op_dims<T, U, Op, D - 1>(
+          a,
+          b,
+          out_a,
+          out_b,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          axis + 1);
+    } else {
+      std::tie(*out_a, *out_b) = op(*a, *b);
+    }
+    a += stride_a;
+    b += stride_b;
+    out_a += stride_out;
+    out_b += stride_out;
+  }
+}
+
+template <typename T, typename U, typename Op>
+void binary_op_dispatch_dims(
+    const array& a,
+    const array& b,
+    array& out_a,
+    array& out_b,
+    Op op) {
+  auto [shape, strides] = collapse_contiguous_dims(
+      a.shape(), {a.strides(), b.strides(), out_a.strides()});
+  const T* a_ptr = a.data<T>();
+  const T* b_ptr = b.data<T>();
+  U* out_a_ptr = out_a.data<U>();
+  U* out_b_ptr = out_b.data<U>();
+
+  const auto& a_strides = strides[0];
+  const auto& b_strides = strides[1];
+  const auto& out_strides = strides[2];
+  int ndim = shape.size();
+  switch (ndim) {
+    case 1:
+      binary_op_dims<T, U, Op, 1>(
+          a_ptr,
+          b_ptr,
+          out_a_ptr,
+          out_b_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
+      return;
+    case 2:
+      binary_op_dims<T, U, Op, 2>(
+          a_ptr,
+          b_ptr,
+          out_a_ptr,
+          out_b_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          out_strides,
+          0);
+      return;
+  }
+
+  ContiguousIterator a_it(shape, a_strides, ndim - 2);
+  ContiguousIterator b_it(shape, b_strides, ndim - 2);
+  auto stride = out_strides[ndim - 3];
+  for (size_t elem = 0; elem < a.size(); elem += stride) {
+    binary_op_dims<T, U, Op, 2>(
+        a_ptr + a_it.loc,
+        b_ptr + b_it.loc,
+        out_a_ptr + elem,
+        out_b_ptr + elem,
+        op,
+        shape,
+        a_strides,
+        b_strides,
+        out_strides,
+        ndim - 2);
+    a_it.step();
+    b_it.step();
+  }
+}
+
+template <typename T, typename U = T, typename Op>
+void binary_op(
+    const array& a,
+    const array& b,
+    array& out_a,
+    array& out_b,
+    Op op,
+    BinaryOpType bopt) {
+  // The full computation is scalar scalar so call the base op once
+  if (bopt == BinaryOpType::General) {
+    binary_op_dispatch_dims<T, U, Op>(a, b, out_a, out_b, op);
+    return;
+  }
+
+  auto a_ptr = a.data<T>();
+  auto b_ptr = b.data<T>();
+  auto out_a_ptr = out_a.data<U>();
+  auto out_b_ptr = out_b.data<U>();
+  if (bopt == BinaryOpType::ScalarScalar) {
+    std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+  } else if (bopt == BinaryOpType::ScalarVector) {
+    for (size_t i = 0; i < b.data_size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      b_ptr++;
+    }
+  } else if (bopt == BinaryOpType::VectorScalar) {
+    for (size_t i = 0; i < a.data_size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      a_ptr++;
+    }
+  } else { // VectorVector
+    for (size_t i = 0; i < a.size(); ++i) {
+      std::tie(*out_a_ptr, *out_b_ptr) = op(*a_ptr, *b_ptr);
+      out_a_ptr++;
+      out_b_ptr++;
+      a_ptr++;
+      b_ptr++;
+    }
+  }
+}
+
+} // namespace
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/cholesky.cpp b/macos/third_party/mlx/mlx/backend/cpu/cholesky.cpp
new file mode 100644
index 0000000..3c5bbbc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/cholesky.cpp
@@ -0,0 +1,85 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/linalg.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T>
+void cholesky_impl(const array& a, array& factor, bool upper, Stream stream) {
+  // Lapack uses the column-major convention. We take advantage of the fact that
+  // the matrix should be symmetric:
+  //   (A)ᵀ = A
+  // and that a column-major lower triangular matrix is a row-major upper
+  // triangular matrix, so uplo is the opposite of what we would expect from
+  // upper
+
+  // The decomposition is computed in place, so just copy the input to the
+  // output.
+  copy_cpu(
+      a,
+      factor,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(factor);
+  encoder.dispatch([matrix = factor.data<T>(),
+                    upper,
+                    N = a.shape(-1),
+                    size = a.size()]() mutable {
+    char uplo = (upper) ? 'L' : 'U';
+    size_t num_matrices = size / (N * N);
+    for (int i = 0; i < num_matrices; i++) {
+      // Compute Cholesky factorization.
+      int info;
+      potrf<T>(
+          /* uplo = */ &uplo,
+          /* n = */ &N,
+          /* a = */ matrix,
+          /* lda = */ &N,
+          /* info = */ &info);
+
+      // TODO: We do nothing when the matrix is not positive semi-definite
+      // because throwing an error would result in a crash. If we figure out how
+      // to catch errors from the implementation we should throw.
+      if (info < 0) {
+        std::stringstream msg;
+        msg << "[Cholesky::eval_cpu] Cholesky decomposition failed with error code "
+            << info;
+        throw std::runtime_error(msg.str());
+      }
+
+      // Zero out the upper/lower triangle while advancing the pointer to the
+      // next matrix at the same time.
+      for (int row = 0; row < N; row++) {
+        if (upper) {
+          std::fill(matrix, matrix + row, 0);
+        } else {
+          std::fill(matrix + row + 1, matrix + N, 0);
+        }
+        matrix += N;
+      }
+    }
+  });
+}
+
+void Cholesky::eval_cpu(const std::vector<array>& inputs, array& output) {
+  switch (inputs[0].dtype()) {
+    case float32:
+      cholesky_impl<float>(inputs[0], output, upper_, stream());
+      break;
+    case float64:
+      cholesky_impl<double>(inputs[0], output, upper_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Cholesky::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/compiled.cpp b/macos/third_party/mlx/mlx/backend/cpu/compiled.cpp
new file mode 100644
index 0000000..0806368
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/compiled.cpp
@@ -0,0 +1,357 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <dlfcn.h>
+#include <filesystem>
+#include <fstream>
+#include <list>
+#include <mutex>
+#include <shared_mutex>
+
+#include <fmt/format.h>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cpu/compiled_preamble.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/jit_compiler.h"
+#include "mlx/device.h"
+#include "mlx/graph_utils.h"
+#include "mlx/version.h"
+
+namespace mlx::core {
+
+struct CompilerCache {
+  struct DLib {
+    DLib(const std::string& libname) {
+      lib = dlopen(libname.c_str(), RTLD_NOW);
+      if (!lib) {
+        std::ostringstream msg;
+        msg << "Could not load C++ shared library " << dlerror();
+        throw std::runtime_error(msg.str());
+      }
+    }
+
+    ~DLib() {
+      dlclose(lib);
+    }
+    void* lib;
+  };
+  // Statics to cache compiled libraries and functions
+  std::list<DLib> libs;
+  std::unordered_map<std::string, void*> kernels;
+  std::shared_mutex mtx;
+};
+
+static CompilerCache& cache() {
+  static CompilerCache cache_;
+  return cache_;
+};
+
+// GPU compile is always available if the GPU is available and since we are in
+// this file CPU compile is also available.
+namespace detail {
+bool compile_available_for_device(const Device& device) {
+  return true;
+}
+
+} // namespace detail
+
+// Return a pointer to a compiled function
+void* compile(
+    const std::string& kernel_name,
+    const std::function<std::string(void)>& source_builder) {
+  {
+    std::shared_lock lock(cache().mtx);
+    if (auto it = cache().kernels.find(kernel_name);
+        it != cache().kernels.end()) {
+      return it->second;
+    }
+  }
+
+  std::unique_lock lock(cache().mtx);
+  if (auto it = cache().kernels.find(kernel_name);
+      it != cache().kernels.end()) {
+    return it->second;
+  }
+  std::string source_code = source_builder();
+  std::string kernel_file_name;
+
+  // Deal with long kernel names. Maximum length for filename on macOS is 255
+  // characters, and on Windows the maximum length for whole path is 260. Clip
+  // file name with a little extra room and append a 16 character hash.
+#ifdef _WIN32
+  constexpr int max_file_name_length = 140;
+#else
+  constexpr int max_file_name_length = 245;
+#endif
+  if (kernel_name.size() > max_file_name_length) {
+    std::ostringstream file_name;
+    file_name
+        << std::string_view(kernel_name).substr(0, max_file_name_length - 16);
+    auto file_id =
+        std::hash<std::string>{}(kernel_name.substr(max_file_name_length - 16));
+    file_name << "_" << std::hex << std::setw(16) << file_id << std::dec;
+    kernel_file_name = file_name.str();
+  } else {
+    kernel_file_name = kernel_name;
+  }
+
+  auto output_dir =
+      std::filesystem::temp_directory_path() / "mlx" / version() / "cpu";
+  if (!std::filesystem::exists(output_dir)) {
+    std::filesystem::create_directories(output_dir);
+  }
+
+  std::string shared_lib_name = "lib" + kernel_file_name + ".so";
+  auto shared_lib_path = (output_dir / shared_lib_name).string();
+  bool lib_exists = false;
+  {
+    std::ifstream f(shared_lib_path.c_str());
+    lib_exists = f.good();
+  }
+
+  if (!lib_exists) {
+    // Open source file and write source code to it
+    std::string source_file_name = kernel_file_name + ".cpp";
+    auto source_file_path = (output_dir / source_file_name).string();
+
+    std::ofstream source_file(source_file_path);
+    source_file << source_code;
+    source_file.close();
+
+    try {
+      JitCompiler::exec(
+          JitCompiler::build_command(
+              output_dir, source_file_name, shared_lib_name));
+    } catch (const std::exception& error) {
+      throw std::runtime_error(
+          fmt::format(
+              "[Compile::eval_cpu] Failed to compile function {0}: {1}",
+              kernel_name,
+              error.what()));
+    }
+  }
+
+  // load library
+  cache().libs.emplace_back(shared_lib_path);
+
+  // Load function
+  void* fun = dlsym(cache().libs.back().lib, kernel_name.c_str());
+  if (!fun) {
+    std::ostringstream msg;
+    msg << "[Compile::eval_cpu] Failed to load compiled function "
+        << kernel_name << std::endl
+        << dlerror();
+    throw std::runtime_error(msg.str());
+  }
+  cache().kernels.insert({kernel_name, fun});
+  return fun;
+}
+
+inline void build_kernel(
+    std::ostream& os,
+    const std::string& kernel_name,
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    const std::vector<array>& tape,
+    const std::function<bool(size_t)>& is_constant,
+    bool contiguous,
+    int ndim) {
+  NodeNamer namer;
+
+#ifdef _MSC_VER
+  // Export the symbol
+  os << "__declspec(dllexport) ";
+#endif
+
+  // Start the kernel
+  os << "void " << kernel_name
+     << "(int* shape, int64_t** strides, void** args) {" << std::endl;
+
+  // Add the input arguments
+  int cnt = 0;
+  int strides_index = 1;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    // Skip constants from the input list
+    if (is_constant(i)) {
+      continue;
+    }
+
+    const auto& x = inputs[i];
+    auto& xname = namer.get_name(x);
+
+    auto tstr = get_type_string(x.dtype());
+    os << "  " << tstr << "* " << xname << " = (" << tstr << "*)args[" << cnt++
+       << "];" << std::endl;
+    // Scalars and contiguous need no strides
+    if (!is_scalar(x) && !contiguous) {
+      os << "  const int64_t* " << xname << "_strides = strides["
+         << strides_index++ << "];" << std::endl;
+    }
+  }
+
+  // Add the output arguments
+  for (auto& x : outputs) {
+    auto tstr = get_type_string(x.dtype());
+    os << "  " << tstr << "* " << namer.get_name(x) << " = (" << tstr
+       << "*)args[" << cnt++ << "];" << std::endl;
+  }
+  // Add output size
+  if (contiguous) {
+    os << "  const size_t size = (size_t)args[" << cnt++ << "];" << std::endl;
+  }
+
+  if (contiguous) {
+    os << "  for (size_t i = 0; i < size; ++i) {" << std::endl;
+  } else {
+    for (int d = 0; d < ndim; ++d) {
+      os << "  for (int i" << d << " = 0; i" << d << " < shape[" << d
+         << "]; ++i" << d << ") {" << std::endl;
+    }
+  }
+
+  // Read the inputs in tmps
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    const auto& x = inputs[i];
+    auto& xname = namer.get_name(x);
+
+    if (is_constant(i)) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = ";
+      print_constant(os, x);
+      os << ";" << std::endl;
+    } else if (is_scalar(x)) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = "
+         << xname << "[0];" << std::endl;
+    } else if (contiguous) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = "
+         << xname << "[i];" << std::endl;
+    } else {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = *"
+         << xname << ";" << std::endl;
+    }
+  }
+
+  // Actually write the computation
+  for (auto& x : tape) {
+    os << "  " << get_type_string(x.dtype()) << " tmp_" << namer.get_name(x)
+       << " = ";
+    if (is_static_cast(x.primitive())) {
+      os << "static_cast<" << get_type_string(x.dtype()) << ">(tmp_"
+         << namer.get_name(x.inputs()[0]) << ");" << std::endl;
+    } else {
+      os << x.primitive().name();
+      os << "()(";
+      for (int i = 0; i < x.inputs().size() - 1; i++) {
+        os << "tmp_" << namer.get_name(x.inputs()[i]) << ", ";
+      }
+      os << "tmp_" << namer.get_name(x.inputs().back()) << ");" << std::endl;
+    }
+  }
+
+  // Write the outputs from tmps
+  for (auto& x : outputs) {
+    if (contiguous) {
+      os << "  " << namer.get_name(x) << "[i] = tmp_" << namer.get_name(x)
+         << ";" << std::endl;
+    } else {
+      os << "  *" << namer.get_name(x) << "++ = tmp_" << namer.get_name(x)
+         << ";" << std::endl;
+    }
+  }
+
+  // Close loops
+  if (contiguous) {
+    os << "  }" << std::endl;
+  } else {
+    for (int d = ndim - 1; d >= 0; --d) {
+      // Update pointers
+      for (size_t i = 0; i < inputs.size(); ++i) {
+        const auto& x = inputs[i];
+        if (is_constant(i) || is_scalar(x)) {
+          continue;
+        }
+        auto& xname = namer.get_name(x);
+        os << "  " << xname << " += " << xname << "_strides[" << d << "];"
+           << std::endl;
+        if (d < ndim - 1) {
+          os << "  " << xname << " -= " << xname << "_strides[" << d + 1 << "]"
+             << " * shape[" << d + 1 << "];" << std::endl;
+        }
+      }
+      os << "  }" << std::endl;
+    }
+  }
+
+  // Finish the kernel
+  os << "}" << std::endl;
+}
+
+void Compiled::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& encoder = cpu::get_command_encoder(stream());
+
+  // Collapse contiguous dims to route to a faster kernel if possible. Also
+  // handle all broadcasting.
+  auto [contiguous, shape, strides] =
+      compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
+
+  // Collect function input arguments.
+  std::vector<void*> args;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    if (is_constant_(i)) {
+      continue;
+    }
+    const auto& x = inputs[i];
+    encoder.set_input_array(x);
+    args.push_back((void*)x.data<void>());
+  }
+
+  // Get the kernel name from the lib
+  int ndim = shape.size();
+  auto kernel_name = kernel_lib_ + (contiguous ? "_contiguous" : "_strided_");
+  if (!contiguous) {
+    kernel_name += std::to_string(ndim);
+  }
+
+  // Get the function
+  auto fn_ptr = compile(kernel_name, [&, contiguous = contiguous]() {
+    std::ostringstream kernel;
+    kernel << get_kernel_preamble() << std::endl;
+    kernel << "extern \"C\"  {" << std::endl;
+    build_kernel(
+        kernel,
+        kernel_name,
+        inputs_,
+        outputs_,
+        tape_,
+        is_constant_,
+        contiguous,
+        ndim);
+    // Close extern "C"
+    kernel << "}" << std::endl;
+    return kernel.str();
+  });
+
+  compiled_allocate_outputs(inputs, outputs, is_constant_, contiguous);
+
+  for (auto& x : outputs) {
+    args.push_back(x.data<void>());
+    encoder.set_output_array(x);
+  }
+  if (contiguous) {
+    args.push_back((void*)outputs[0].data_size());
+  }
+  auto fun = reinterpret_cast<void (*)(int*, int64_t**, void**)>(fn_ptr);
+  encoder.dispatch([fun,
+                    args = std::move(args),
+                    strides = std::move(strides),
+                    shape = std::move(shape)]() mutable {
+    SmallVector<int64_t*> strides_ptrs;
+    for (auto& s : strides) {
+      strides_ptrs.push_back(s.data());
+    }
+    fun(shape.data(), strides_ptrs.data(), args.data());
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/compiled_preamble.h b/macos/third_party/mlx/mlx/backend/cpu/compiled_preamble.h
new file mode 100644
index 0000000..31ca1b4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/compiled_preamble.h
@@ -0,0 +1,12 @@
+// Copyright © 2023-24 Apple Inc.
+
+#pragma once
+
+// clang-format off
+#include "mlx/types/half_types.h"
+#include "mlx/types/complex.h"
+#include "mlx/backend/cpu/unary_ops.h"
+#include "mlx/backend/cpu/binary_ops.h"
+// clang-format on
+
+const char* get_kernel_preamble();
diff --git a/macos/third_party/mlx/mlx/backend/cpu/conv.cpp b/macos/third_party/mlx/mlx/backend/cpu/conv.cpp
new file mode 100644
index 0000000..5d4638a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/conv.cpp
@@ -0,0 +1,1351 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cassert>
+#include <numeric>
+
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+///////////////////////////////////////////////////////////////////////////////
+// Naive reference conv
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+void slow_conv_1D(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_input_array(wt);
+  encoder.set_output_array(out);
+
+  encoder.dispatch([start_wt_ptr = wt.data<T>(),
+                    in_ptr = in.data<T>(),
+                    out_ptr = out.data<T>(),
+
+                    N = in.shape(
+                        0), // Batch size, should be the same as out.shape(0)
+                    iH = 1 +
+                        in_dilation[0] * (in.shape(1) - 1), // Input spatial dim
+                    oH = out.shape(1), // Output spatial dim
+                    wH = wt.shape(1), // Weight spatial dim
+                    groups = in.shape(2) / wt.shape(2),
+                    O = wt.shape(0), // Out channels
+                    C_per_group = wt.shape(2),
+
+                    in_stride_N = in.strides()[0],
+                    in_stride_H = in.strides()[1],
+                    in_stride_C = in.strides()[2],
+
+                    wt_stride_O = wt.strides()[0],
+                    wt_stride_H = wt.strides()[1],
+                    wt_stride_C = wt.strides()[2],
+
+                    out_stride_N = out.strides()[0],
+                    out_stride_H = out.strides()[1],
+                    out_stride_O = out.strides()[2],
+
+                    flip,
+                    padding_lo = padding_lo[0],
+                    padding_hi = padding_hi[0],
+                    wt_stride = wt_strides[0],
+                    wt_dilation = wt_dilation[0],
+                    in_dilation = in_dilation[0]]() mutable {
+    auto O_per_group = O / groups;
+
+    for (int n = 0; n < N; ++n) {
+      for (int oh = 0; oh < oH; ++oh) {
+        for (int g = 0; g < groups; ++g) {
+          for (int o = g * O_per_group; o < (g + 1) * O_per_group; ++o) {
+            const T* filter_wt_ptr = start_wt_ptr + o * wt_stride_O;
+            float r = 0.;
+
+            for (int wh = 0; wh < wH; ++wh) {
+              const T* wt_ptr = filter_wt_ptr + wh * wt_stride_H;
+
+              int wh_flip = flip ? (wH - wh - 1) : wh;
+              int ih = oh * wt_stride - padding_lo + wh_flip * wt_dilation;
+
+              auto ih_div = std::div(ih, in_dilation);
+
+              if (ih >= 0 && ih < iH && ih_div.rem == 0) {
+                for (int c = g * C_per_group; c < (g + 1) * C_per_group; ++c) {
+                  r +=
+                      static_cast<float>(
+                          in_ptr[ih_div.quot * in_stride_H + c * in_stride_C]) *
+                      static_cast<float>(
+                          wt_ptr[(c % C_per_group) * wt_stride_C]);
+                } // c
+
+              } // ih check
+            } // wh
+
+            out_ptr[oh * out_stride_H + o * out_stride_O] = static_cast<T>(r);
+          } // o
+        } // g
+      } // oh
+
+      in_ptr += in_stride_N;
+      out_ptr += out_stride_N;
+    } // n
+  });
+}
+
+template <typename T>
+void slow_conv_2D(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_input_array(wt);
+  encoder.set_output_array(out);
+
+  encoder.dispatch(
+      [st_wt_ptr = wt.data<T>(),
+       st_in_ptr = in.data<T>(),
+       st_out_ptr = out.data<T>(),
+
+       N = in.shape(0), // Batch size, should be the same as out.shape(0)
+       iH = 1 + in_dilation[0] * (in.shape(1) - 1), // Input spatial dim
+       iW = 1 + in_dilation[1] * (in.shape(2) - 1), // Input spatial dim
+       C = in.shape(3), // In channels
+       oH = out.shape(1), // Output spatial dim
+       oW = out.shape(2), // Output spatial dim
+       O = wt.shape(0), // Out channels
+       wH = wt.shape(1), // Weight spatial dim
+       wW = wt.shape(2), // Weight spatial dim
+
+       groups = in.shape(3) / wt.shape(3),
+       C_per_group = wt.shape(3),
+
+       in_stride_N = in.strides()[0],
+       in_stride_H = in.strides()[1],
+       in_stride_W = in.strides()[2],
+       in_stride_C = in.strides()[3],
+
+       wt_stride_O = wt.strides()[0],
+       wt_stride_H = wt.strides()[1],
+       wt_stride_W = wt.strides()[2],
+       wt_stride_C = wt.strides()[3],
+
+       out_stride_N = out.strides()[0],
+       out_stride_H = out.strides()[1],
+       out_stride_W = out.strides()[2],
+       out_stride_O = out.strides()[3],
+
+       padding_lo,
+       padding_hi,
+       wt_strides,
+       wt_dilation,
+       in_dilation,
+       flip]() mutable {
+        bool is_idil_one = in_dilation[0] == 1 && in_dilation[1] == 1;
+
+        const int O_per_group = O / groups;
+        auto pt_conv_no_checks =
+            [&](const T* in_ptr, const T* wt_ptr, T* out_ptr, int oh, int ow) {
+              out_ptr += oh * out_stride_H + ow * out_stride_W;
+              int ih_base = oh * wt_strides[0] - padding_lo[0];
+              int iw_base = ow * wt_strides[1] - padding_lo[1];
+
+              for (int g = 0; g < groups; ++g) {
+                for (int o = g * O_per_group; o < (g + 1) * O_per_group; ++o) {
+                  float r = 0.;
+
+                  for (int wh = 0; wh < wH; ++wh) {
+                    for (int ww = 0; ww < wW; ++ww) {
+                      int wh_flip = flip ? wH - wh - 1 : wh;
+                      int ww_flip = flip ? wW - ww - 1 : ww;
+                      int ih = ih_base + wh_flip * wt_dilation[0];
+                      int iw = iw_base + ww_flip * wt_dilation[1];
+
+                      const T* wt_ptr_pt =
+                          wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
+                      const T* in_ptr_pt =
+                          in_ptr + ih * in_stride_H + iw * in_stride_W;
+
+                      for (int c = g * C_per_group; c < (g + 1) * C_per_group;
+                           ++c) {
+                        r += static_cast<float>(in_ptr_pt[c * in_stride_C]) *
+                            static_cast<float>(
+                                 wt_ptr_pt[(c % C_per_group) * wt_stride_C]);
+                      } // c
+                    } // ww
+                  } // wh
+
+                  out_ptr[0] = static_cast<T>(r);
+                  out_ptr += out_stride_O;
+                  wt_ptr += wt_stride_O;
+                } // o
+              } // g
+            };
+
+        int jump_h = flip ? -wt_dilation[0] : wt_dilation[0];
+        int jump_w = flip ? -wt_dilation[1] : wt_dilation[1];
+
+        int init_h = (flip ? (wH - 1) * wt_dilation[0] : 0);
+        int init_w = (flip ? (wW - 1) * wt_dilation[1] : 0);
+
+        int f_wgt_jump_h =
+            std::lcm(in_dilation[0], wt_dilation[0]) / wt_dilation[0];
+        int f_wgt_jump_w =
+            std::lcm(in_dilation[1], wt_dilation[1]) / wt_dilation[1];
+
+        int f_out_jump_h =
+            std::lcm(in_dilation[0], wt_strides[0]) / wt_strides[0];
+        int f_out_jump_w =
+            std::lcm(in_dilation[1], wt_strides[1]) / wt_strides[1];
+
+        std::vector<int> base_h(f_out_jump_h);
+        std::vector<int> base_w(f_out_jump_w);
+
+        for (int i = 0; i < f_out_jump_h; ++i) {
+          int ih_loop = i * wt_strides[0] - padding_lo[0] + init_h;
+
+          int wh_base = 0;
+          while (wh_base < wH && ih_loop % in_dilation[0] != 0) {
+            wh_base++;
+            ih_loop += jump_h;
+          }
+
+          base_h[i] = wh_base;
+        }
+
+        for (int j = 0; j < f_out_jump_w; ++j) {
+          int iw_loop = j * wt_strides[1] - padding_lo[1] + init_w;
+
+          int ww_base = 0;
+          while (ww_base < wW && iw_loop % in_dilation[1] != 0) {
+            ww_base++;
+            iw_loop += jump_w;
+          }
+
+          base_w[j] = ww_base;
+        }
+
+        auto pt_conv_all_checks =
+            [&](const T* in_ptr, const T* wt_ptr, T* out_ptr, int oh, int ow) {
+              out_ptr += oh * out_stride_H + ow * out_stride_W;
+
+              int ih_base = oh * wt_strides[0] - padding_lo[0];
+              int iw_base = ow * wt_strides[1] - padding_lo[1];
+
+              int wh_base = base_h[oh % f_out_jump_h];
+              int ww_base = base_w[ow % f_out_jump_w];
+
+              for (int g = 0; g < groups; ++g) {
+                for (int o = g * O_per_group; o < (g + 1) * O_per_group; ++o) {
+                  float r = 0.;
+
+                  for (int wh = wh_base; wh < wH; wh += f_wgt_jump_h) {
+                    for (int ww = ww_base; ww < wW; ww += f_wgt_jump_w) {
+                      int wh_flip = flip ? wH - wh - 1 : wh;
+                      int ww_flip = flip ? wW - ww - 1 : ww;
+                      int ih = ih_base + wh_flip * wt_dilation[0];
+                      int iw = iw_base + ww_flip * wt_dilation[1];
+
+                      if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) {
+                        const T* wt_ptr_pt =
+                            wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
+
+                        int ih_dil = !is_idil_one ? (ih / in_dilation[0]) : ih;
+                        int iw_dil = !is_idil_one ? (iw / in_dilation[1]) : iw;
+
+                        const T* in_ptr_pt = in_ptr + ih_dil * in_stride_H +
+                            iw_dil * in_stride_W;
+
+                        for (int c = g * C_per_group; c < (g + 1) * C_per_group;
+                             ++c) {
+                          r += static_cast<float>(in_ptr_pt[c * in_stride_C]) *
+                              static_cast<float>(
+                                   wt_ptr_pt[(c % C_per_group) * wt_stride_C]);
+                        } // c
+
+                      } // ih, iw check
+                    } // ww
+                  } // wh
+
+                  out_ptr[0] = static_cast<T>(r);
+                  out_ptr += out_stride_O;
+                  wt_ptr += wt_stride_O;
+                } // o
+              } // g
+            };
+
+        int oH_border_0 = 0;
+        int oH_border_1 = is_idil_one
+            ? ((padding_lo[0] + wt_strides[0] - 1) / wt_strides[0])
+            : oH;
+        int oH_border_2 = std::max(
+            oH_border_1,
+            (iH + padding_lo[0] - wH * wt_dilation[0]) / wt_strides[0]);
+        int oH_border_3 = oH;
+
+        int oW_border_0 = 0;
+        int oW_border_1 = is_idil_one
+            ? ((padding_lo[1] + wt_strides[1] - 1) / wt_strides[1])
+            : oW;
+        int oW_border_2 = std::max(
+            oW_border_1,
+            (iW + padding_lo[1] - wW * wt_dilation[1]) / wt_strides[1]);
+        int oW_border_3 = oW;
+
+        for (int n = 0; n < N; ++n) {
+          // Case 1: oh might put us out of bounds
+          for (int oh = oH_border_0; oh < oH_border_1; ++oh) {
+            for (int ow = 0; ow < oW; ++ow) {
+              pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, oh, ow);
+            } // ow
+          } // oh
+
+          // Case 2: oh in bounds
+          for (int oh = oH_border_1; oh < oH_border_2; ++oh) {
+            // Case a: ow might put us out of bounds
+            for (int ow = oW_border_0; ow < oW_border_1; ++ow) {
+              pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, oh, ow);
+            } // ow
+
+            // Case b: ow in bounds
+            for (int ow = oW_border_1; ow < oW_border_2; ++ow) {
+              pt_conv_no_checks(st_in_ptr, st_wt_ptr, st_out_ptr, oh, ow);
+            } // ow
+
+            // Case c: ow might put us out of bounds
+            for (int ow = oW_border_2; ow < oW_border_3; ++ow) {
+              pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, oh, ow);
+            } // ow
+
+          } // oh
+
+          // Case 3: oh might put us out of bounds
+          for (int oh = oH_border_2; oh < oH_border_3; ++oh) {
+            for (int ow = 0; ow < oW; ++ow) {
+              pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, oh, ow);
+            } // ow
+          } // oh
+
+          st_in_ptr += in_stride_N;
+          st_out_ptr += out_stride_N;
+
+        } // n
+      });
+}
+
+template <typename T>
+void slow_conv_3D(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_input_array(wt);
+  encoder.set_output_array(out);
+
+  encoder.dispatch([st_wt_ptr = wt.data<T>(),
+                    st_in_ptr = in.data<T>(),
+                    st_out_ptr = out.data<T>(),
+
+                    N = in.shape(
+                        0), // Batch size, should be the same as out.shape(0)
+                    iD = 1 +
+                        in_dilation[0] * (in.shape(1) - 1), // Input spatial dim
+                    iH = 1 +
+                        in_dilation[1] * (in.shape(2) - 1), // Input spatial dim
+                    iW = 1 +
+                        in_dilation[2] * (in.shape(3) - 1), // Input spatial dim
+                    oD = out.shape(1), // Output spatial dim
+                    oH = out.shape(2), // Output spatial dim
+                    oW = out.shape(3), // Output spatial dim
+                    O = wt.shape(0), // Out channels
+                    C = wt.shape(4), // In channels
+                    wD = wt.shape(1), // Weight spatial dim
+                    wH = wt.shape(2), // Weight spatial dim
+                    wW = wt.shape(3), // Weight spatial dim
+
+                    in_stride_N = in.strides()[0],
+                    in_stride_D = in.strides()[1],
+                    in_stride_H = in.strides()[2],
+                    in_stride_W = in.strides()[3],
+                    in_stride_C = in.strides()[4],
+
+                    wt_stride_O = wt.strides()[0],
+                    wt_stride_D = wt.strides()[1],
+                    wt_stride_H = wt.strides()[2],
+                    wt_stride_W = wt.strides()[3],
+                    wt_stride_C = wt.strides()[4],
+
+                    out_stride_N = out.strides()[0],
+                    out_stride_D = out.strides()[1],
+                    out_stride_H = out.strides()[2],
+                    out_stride_W = out.strides()[3],
+                    out_stride_O = out.strides()[4],
+                    padding_lo,
+                    padding_hi,
+                    wt_strides,
+                    wt_dilation,
+                    in_dilation,
+                    flip]() mutable {
+    bool is_idil_one =
+        in_dilation[0] == 1 && in_dilation[1] == 1 && in_dilation[2] == 1;
+
+    auto pt_conv_no_checks = [&](const T* in_ptr,
+                                 const T* wt_ptr,
+                                 T* out_ptr,
+                                 int od,
+                                 int oh,
+                                 int ow) {
+      out_ptr += od * out_stride_D + oh * out_stride_H + ow * out_stride_W;
+      int id_base = od * wt_strides[0] - padding_lo[0];
+      int ih_base = oh * wt_strides[1] - padding_lo[1];
+      int iw_base = ow * wt_strides[2] - padding_lo[2];
+
+      for (int o = 0; o < O; ++o) {
+        float r = 0.;
+
+        for (int wd = 0; wd < wD; ++wd) {
+          for (int wh = 0; wh < wH; ++wh) {
+            for (int ww = 0; ww < wW; ++ww) {
+              int wd_flip = flip ? wD - wd - 1 : wd;
+              int wh_flip = flip ? wH - wh - 1 : wh;
+              int ww_flip = flip ? wW - ww - 1 : ww;
+              int id = id_base + wd_flip * wt_dilation[0];
+              int ih = ih_base + wh_flip * wt_dilation[1];
+              int iw = iw_base + ww_flip * wt_dilation[2];
+
+              const T* wt_ptr_pt = wt_ptr + wd * wt_stride_D +
+                  wh * wt_stride_H + ww * wt_stride_W;
+              const T* in_ptr_pt = in_ptr + id * in_stride_D +
+                  ih * in_stride_H + iw * in_stride_W;
+
+              for (int c = 0; c < C; ++c) {
+                r += static_cast<float>(in_ptr_pt[0]) *
+                    static_cast<float>(wt_ptr_pt[0]);
+                in_ptr_pt += in_stride_C;
+                wt_ptr_pt += wt_stride_C;
+              } // c
+
+            } // ww
+          } // wh
+        } // wd
+
+        out_ptr[0] = static_cast<T>(r);
+        out_ptr += out_stride_O;
+        wt_ptr += wt_stride_O;
+      } // o
+    };
+
+    int jump_d = flip ? -wt_dilation[0] : wt_dilation[0];
+    int jump_h = flip ? -wt_dilation[1] : wt_dilation[1];
+    int jump_w = flip ? -wt_dilation[2] : wt_dilation[2];
+
+    int init_d = (flip ? (wD - 1) * wt_dilation[0] : 0);
+    int init_h = (flip ? (wH - 1) * wt_dilation[1] : 0);
+    int init_w = (flip ? (wW - 1) * wt_dilation[2] : 0);
+
+    int f_wgt_jump_d =
+        std::lcm(in_dilation[0], wt_dilation[0]) / wt_dilation[0];
+    int f_wgt_jump_h =
+        std::lcm(in_dilation[1], wt_dilation[1]) / wt_dilation[1];
+    int f_wgt_jump_w =
+        std::lcm(in_dilation[2], wt_dilation[2]) / wt_dilation[2];
+
+    int f_out_jump_d = std::lcm(in_dilation[0], wt_strides[0]) / wt_strides[0];
+    int f_out_jump_h = std::lcm(in_dilation[1], wt_strides[1]) / wt_strides[1];
+    int f_out_jump_w = std::lcm(in_dilation[2], wt_strides[2]) / wt_strides[2];
+
+    std::vector<int> base_d(f_out_jump_d);
+    std::vector<int> base_h(f_out_jump_h);
+    std::vector<int> base_w(f_out_jump_w);
+
+    for (int i = 0; i < f_out_jump_d; ++i) {
+      int id_loop = i * wt_strides[0] - padding_lo[0] + init_d;
+
+      int wd_base = 0;
+      while (wd_base < wD && id_loop % in_dilation[0] != 0) {
+        wd_base++;
+        id_loop += jump_d;
+      }
+
+      base_d[i] = wd_base;
+    }
+
+    for (int i = 0; i < f_out_jump_h; ++i) {
+      int ih_loop = i * wt_strides[1] - padding_lo[1] + init_h;
+
+      int wh_base = 0;
+      while (wh_base < wH && ih_loop % in_dilation[1] != 0) {
+        wh_base++;
+        ih_loop += jump_h;
+      }
+
+      base_h[i] = wh_base;
+    }
+
+    for (int j = 0; j < f_out_jump_w; ++j) {
+      int iw_loop = j * wt_strides[2] - padding_lo[2] + init_w;
+
+      int ww_base = 0;
+      while (ww_base < wW && iw_loop % in_dilation[2] != 0) {
+        ww_base++;
+        iw_loop += jump_w;
+      }
+
+      base_w[j] = ww_base;
+    }
+
+    auto pt_conv_all_checks = [&](const T* in_ptr,
+                                  const T* wt_ptr,
+                                  T* out_ptr,
+                                  int od,
+                                  int oh,
+                                  int ow) {
+      out_ptr += od * out_stride_D + oh * out_stride_H + ow * out_stride_W;
+
+      int id_base = od * wt_strides[0] - padding_lo[0];
+      int ih_base = oh * wt_strides[1] - padding_lo[1];
+      int iw_base = ow * wt_strides[2] - padding_lo[2];
+
+      int wd_base = base_d[od % f_out_jump_d];
+      int wh_base = base_h[oh % f_out_jump_h];
+      int ww_base = base_w[ow % f_out_jump_w];
+
+      for (int o = 0; o < O; ++o) {
+        float r = 0.;
+
+        for (int wd = wd_base; wd < wD; wd += f_wgt_jump_d) {
+          for (int wh = wh_base; wh < wH; wh += f_wgt_jump_h) {
+            for (int ww = ww_base; ww < wW; ww += f_wgt_jump_w) {
+              int wd_flip = flip ? wD - wd - 1 : wd;
+              int wh_flip = flip ? wH - wh - 1 : wh;
+              int ww_flip = flip ? wW - ww - 1 : ww;
+              int id = id_base + wd_flip * wt_dilation[0];
+              int ih = ih_base + wh_flip * wt_dilation[1];
+              int iw = iw_base + ww_flip * wt_dilation[2];
+
+              if (id >= 0 && id < iD && ih >= 0 && ih < iH && iw >= 0 &&
+                  iw < iW) {
+                const T* wt_ptr_pt = wt_ptr + wd * wt_stride_D +
+                    wh * wt_stride_H + ww * wt_stride_W;
+
+                int id_dil = !is_idil_one ? (id / in_dilation[0]) : id;
+                int ih_dil = !is_idil_one ? (ih / in_dilation[1]) : ih;
+                int iw_dil = !is_idil_one ? (iw / in_dilation[2]) : iw;
+
+                const T* in_ptr_pt = in_ptr + id_dil * in_stride_D +
+                    ih_dil * in_stride_H + iw_dil * in_stride_W;
+
+                for (int c = 0; c < C; ++c) {
+                  r += static_cast<float>(in_ptr_pt[0]) *
+                      static_cast<float>(wt_ptr_pt[0]);
+                  in_ptr_pt += in_stride_C;
+                  wt_ptr_pt += wt_stride_C;
+                } // c
+
+              } // iD, ih, iw check
+            } // ww
+          } // wh
+        } // wd
+
+        out_ptr[0] = static_cast<T>(r);
+        out_ptr += out_stride_O;
+        wt_ptr += wt_stride_O;
+      } // o
+    };
+
+    int oD_border_0 = 0;
+    int oD_border_1 = is_idil_one
+        ? ((padding_lo[0] + wt_strides[0] - 1) / wt_strides[0])
+        : oD;
+    int oD_border_2 = std::max(
+        oD_border_1,
+        (iD + padding_lo[0] - wD * wt_dilation[0]) / wt_strides[0]);
+    int oD_border_3 = oD;
+
+    int oH_border_0 = 0;
+    int oH_border_1 = is_idil_one
+        ? ((padding_lo[1] + wt_strides[1] - 1) / wt_strides[1])
+        : oH;
+    int oH_border_2 = std::max(
+        oH_border_1,
+        (iH + padding_lo[1] - wH * wt_dilation[1]) / wt_strides[1]);
+    int oH_border_3 = oH;
+
+    int oW_border_0 = 0;
+    int oW_border_1 = is_idil_one
+        ? ((padding_lo[2] + wt_strides[2] - 1) / wt_strides[2])
+        : oW;
+    int oW_border_2 = std::max(
+        oW_border_1,
+        (iW + padding_lo[2] - wW * wt_dilation[2]) / wt_strides[2]);
+    int oW_border_3 = oW;
+
+    for (int n = 0; n < N; ++n) {
+      // Case 1: od might put us out of bounds
+      for (int od = oD_border_0; od < oD_border_1; ++od) {
+        for (int oh = 0; oh < oH; ++oh) {
+          for (int ow = 0; ow < oW; ++ow) {
+            pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
+          } // ow
+        } // oh
+      } // od
+
+      // Case 2: od in bounds
+      for (int od = oD_border_1; od < oD_border_2; ++od) {
+        // Case 2.1: oh might put us out of bounds
+        for (int oh = oH_border_0; oh < oH_border_1; ++oh) {
+          for (int ow = 0; ow < oW; ++ow) {
+            pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
+          } // ow
+        } // oh
+
+        // Case 2.2: oh in bounds
+        for (int oh = oH_border_1; oh < oH_border_2; ++oh) {
+          // Case 2.2.1: ow might put us out of bounds
+          for (int ow = oW_border_0; ow < oW_border_1; ++ow) {
+            pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
+          } // ow
+
+          // Case 2.2.2: ow in bounds
+          for (int ow = oW_border_1; ow < oW_border_2; ++ow) {
+            pt_conv_no_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
+          } // ow
+
+          // Case 2.2.3: ow might put us out of bounds
+          for (int ow = oW_border_2; ow < oW_border_3; ++ow) {
+            pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
+          } // ow
+        } // oh
+
+        // Case 2.3: oh might put us out of bounds
+        for (int oh = oH_border_2; oh < oH_border_3; ++oh) {
+          for (int ow = 0; ow < oW; ++ow) {
+            pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
+          } // ow
+        } // oh
+      } // od
+
+      // Case 3: od might put us out of bounds
+      for (int od = oD_border_2; od < oD_border_3; ++od) {
+        for (int oh = 0; oh < oH; ++oh) {
+          for (int ow = 0; ow < oW; ++ow) {
+            pt_conv_all_checks(st_in_ptr, st_wt_ptr, st_out_ptr, od, oh, ow);
+          } // ow
+        } // oh
+      } // od
+
+      st_in_ptr += in_stride_N;
+      st_out_ptr += out_stride_N;
+
+    } // n
+  });
+}
+
+void dispatch_slow_conv_1D(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  if (in.dtype() == float32) {
+    return slow_conv_1D<float>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else if (in.dtype() == float16) {
+    return slow_conv_1D<float16_t>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else if (in.dtype() == bfloat16) {
+    return slow_conv_1D<bfloat16_t>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else {
+    throw std::invalid_argument(
+        "[Convolution::eval] got unsupported data type.");
+  }
+}
+
+void dispatch_slow_conv_2D(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  if (in.dtype() == float32) {
+    return slow_conv_2D<float>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else if (in.dtype() == float16) {
+    return slow_conv_2D<float16_t>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else if (in.dtype() == bfloat16) {
+    return slow_conv_2D<bfloat16_t>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else {
+    throw std::invalid_argument(
+        "[Convolution::eval] got unsupported data type.");
+  }
+}
+
+void dispatch_slow_conv_3D(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  if (in.dtype() == float32) {
+    return slow_conv_3D<float>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else if (in.dtype() == float16) {
+    return slow_conv_3D<float16_t>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else if (in.dtype() == bfloat16) {
+    return slow_conv_3D<bfloat16_t>(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        in_dilation,
+        flip,
+        stream);
+  } else {
+    throw std::invalid_argument(
+        "[Convolution::eval] got unsupported data type.");
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Explicit gemm conv
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+void flip_spatial_dims_inplace(
+    T* x,
+    size_t in_channels,
+    size_t out_channels,
+    size_t spatial_size) {
+  for (size_t i = 0; i < out_channels; i++) {
+    T* top = x + i * spatial_size * in_channels;
+    T* bottom =
+        x + i * spatial_size * in_channels + (spatial_size - 1) * in_channels;
+    for (size_t j = 0; j < spatial_size / 2; j++) {
+      for (size_t k = 0; k < in_channels; k++) {
+        std::swap(top[k], bottom[k]);
+      }
+      top += in_channels;
+      bottom -= in_channels;
+    }
+  }
+}
+
+void explicit_gemm_conv_1D_cpu(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    Stream stream) {
+  const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
+  const int iH = in.shape(1); // Input spatial dim
+  const int C = in.shape(2); // Input channels
+  const int oH = out.shape(1); // Output spatial dim
+  const int O = wt.shape(0); // Out channels
+  const int wH = wt.shape(1); // Weight spatial dim
+
+  const int groups = C / wt.shape(2);
+  const int C_per_group = wt.shape(2);
+  const int O_per_group = O / groups;
+
+  auto conv_dtype = float32;
+  auto& encoder = cpu::get_command_encoder(stream);
+
+  // Pad input
+  Shape padded_shape = {N, iH + padding_lo[0] + padding_hi[0], C};
+  array in_padded(padded_shape, conv_dtype, nullptr, {});
+
+  // Fill with zeros
+  std::vector<array> temps;
+  temps.push_back(array(0, conv_dtype));
+  copy_cpu(temps.back(), in_padded, CopyType::Scalar, stream);
+
+  // Pick input slice from padded
+  size_t data_offset = padding_lo[0] * in_padded.strides()[1];
+  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
+  in_padded_slice.copy_shared_buffer(
+      in_padded,
+      in_padded.strides(),
+      in_padded.flags(),
+      in_padded_slice.size(),
+      data_offset);
+  // Copy input values into the slice
+  copy_cpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, stream);
+  temps.push_back(in_padded_slice);
+
+  // Make strided view
+  Shape strided_shape = {N, oH, wH, C};
+
+  Strides strided_strides = {
+      in_padded.strides()[0],
+      in_padded.strides()[1] * wt_strides[0],
+      in_padded.strides()[1],
+      in_padded.strides()[2]};
+  auto flags = in_padded.flags();
+  if (groups > 1) {
+    // Transpose the last two dimensions for grouped convolutions
+    std::swap(strided_shape[2], strided_shape[3]);
+    std::swap(strided_strides[2], strided_strides[3]);
+  }
+
+  array in_strided_view(strided_shape, in_padded.dtype(), nullptr, {});
+  in_strided_view.copy_shared_buffer(
+      in_padded, strided_strides, flags, in_strided_view.size(), 0);
+
+  // Materialize strided view
+  Shape strided_reshape = {N * oH, wH * C};
+  array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
+  copy_cpu(in_strided_view, in_strided, CopyType::General, stream);
+  temps.push_back(in_strided);
+
+  // Check wt dtype and prepare
+  auto gemm_wt = wt;
+  auto gemm_out = out;
+
+  if (groups > 1) {
+    // Transpose the last two dimensions for grouped convolutions
+    array wt_transpose(
+        {wt.shape(0), wt.shape(2), wt.shape(1)}, wt.dtype(), nullptr, {});
+    wt_transpose.copy_shared_buffer(
+        wt,
+        {wt.strides(0), wt.strides(2), wt.strides(1)},
+        wt.flags(),
+        wt.size(),
+        0);
+    gemm_wt = array(wt_transpose.shape(), float32, nullptr, {});
+    copy_cpu(wt_transpose, gemm_wt, CopyType::General, stream);
+    temps.push_back(gemm_wt);
+  } else if (wt.dtype() != float32 || !wt.flags().row_contiguous) {
+    auto ctype =
+        wt.flags().row_contiguous ? CopyType::Vector : CopyType::General;
+    gemm_wt = array(wt.shape(), float32, nullptr, {});
+    copy_cpu(wt, gemm_wt, ctype, stream);
+    temps.push_back(gemm_wt);
+  }
+
+  if (out.dtype() != float32) {
+    gemm_out = array(out.shape(), float32, nullptr, {});
+    gemm_out.set_data(allocator::malloc(gemm_out.nbytes()));
+    temps.push_back(gemm_out);
+  }
+
+  encoder.set_input_array(in_strided);
+  encoder.set_input_array(gemm_wt);
+  encoder.set_output_array(gemm_out);
+
+  encoder.dispatch([in_strided_ptr = in_strided.data<float>(),
+                    gemm_wt_ptr = gemm_wt.data<float>(),
+                    gemm_out_ptr = gemm_out.data<float>(),
+                    groups,
+                    strided_reshape = strided_reshape[0],
+                    O,
+                    C,
+                    wH,
+                    O_per_group,
+                    C_per_group]() {
+    for (int g = 0; g < groups; ++g) {
+      // Perform gemm
+      cblas_sgemm(
+          CblasRowMajor,
+          CblasNoTrans, // no trans A
+          CblasTrans, // transB
+          strided_reshape, // M
+          O_per_group, // N
+          C_per_group * wH, // K
+          1.0f, // alpha
+          in_strided_ptr + g * C_per_group * wH, // A
+          wH * C, // lda
+          gemm_wt_ptr + g * O_per_group * C_per_group * wH, // B
+          wH * C_per_group, // ldb
+          0.0f, // beta
+          gemm_out_ptr + g * O_per_group, // C
+          O // ldc
+      );
+    }
+  });
+
+  // Copy results if needed
+  if (out.dtype() != float32) {
+    copy_cpu_inplace(gemm_out, out, CopyType::Vector, stream);
+  }
+  encoder.add_temporaries(std::move(temps));
+}
+
+void explicit_gemm_conv_ND_cpu(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const bool flip,
+    Stream stream) {
+  const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
+  const auto iDim =
+      Shape(in.shape().begin() + 1, in.shape().end() - 1); // Input spatial dim
+  const auto oDim = Shape(
+      out.shape().begin() + 1, out.shape().end() - 1); // Output spatial dim
+  const int O = wt.shape(0); // Out channels
+  const int C = wt.shape(-1); // In channels
+  const auto wDim =
+      Shape(wt.shape().begin() + 1, wt.shape().end() - 1); // Weight spatial dim
+
+  auto conv_dtype = float32;
+
+  auto& encoder = cpu::get_command_encoder(stream);
+
+  // Pad input
+  Shape padded_shape(in.shape().size());
+  padded_shape.front() = N;
+  for (size_t i = 0; i < iDim.size(); i++) {
+    padded_shape[i + 1] = iDim[i] + padding_lo[i] + padding_hi[i];
+  }
+  padded_shape.back() = C;
+  array in_padded(padded_shape, conv_dtype, nullptr, {});
+
+  // Fill with zeros
+  std::vector<array> temps = {array(0, conv_dtype)};
+  copy_cpu(temps.back(), in_padded, CopyType::Scalar, stream);
+
+  // Pick input slice from padded
+  size_t data_offset = 0;
+  for (size_t i = 0; i < padding_lo.size(); i++) {
+    data_offset += padding_lo[i] * in_padded.strides()[i + 1];
+  }
+
+  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
+  in_padded_slice.copy_shared_buffer(
+      in_padded,
+      in_padded.strides(),
+      in_padded.flags(),
+      in_padded_slice.size(),
+      data_offset);
+
+  // Copy input values into the slice
+  copy_cpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, stream);
+  temps.push_back(in_padded_slice);
+
+  // Make strided view
+  Shape strided_shape(oDim.size() + wDim.size() + 2);
+  strided_shape.front() = N;
+  for (size_t i = 0; i < oDim.size(); i++) {
+    strided_shape[i + 1] = oDim[i];
+  }
+  for (size_t i = 0; i < wDim.size(); i++) {
+    strided_shape[i + 1 + oDim.size()] = wDim[i];
+  }
+  strided_shape.back() = C;
+
+  Strides strided_strides(in.shape().size() * 2 - 2);
+  strided_strides[0] = in_padded.strides()[0];
+  for (size_t i = 0; i < wt_strides.size(); i++) {
+    strided_strides[i + 1] = in_padded.strides()[i + 1] * wt_strides[i];
+  }
+  for (size_t i = 1; i < in_padded.strides().size(); i++) {
+    strided_strides[i + wt_strides.size()] = in_padded.strides()[i];
+  }
+
+  auto flags = in_padded.flags();
+
+  array in_strided_view(strided_shape, in_padded.dtype(), nullptr, {});
+  in_strided_view.copy_shared_buffer(
+      in_padded, strided_strides, flags, in_strided_view.size(), 0);
+
+  // Materialize strided view
+  Shape strided_reshape = {N, C};
+  for (const auto& o : oDim) {
+    strided_reshape[0] *= o;
+  }
+  for (const auto& w : wDim) {
+    strided_reshape[1] *= w;
+  }
+
+  array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
+  copy_cpu(in_strided_view, in_strided, CopyType::General, stream);
+  temps.push_back(in_strided);
+
+  // Check wt dtype and prepare
+  auto gemm_wt = wt;
+  auto gemm_out = out;
+
+  if (wt.dtype() != float32 || !wt.flags().row_contiguous) {
+    auto ctype =
+        wt.flags().row_contiguous ? CopyType::Vector : CopyType::General;
+    gemm_wt = array(wt.shape(), float32, nullptr, {});
+    copy_cpu(wt, gemm_wt, ctype, stream);
+    temps.push_back(gemm_wt);
+  }
+
+  if (flip) {
+    auto gemm_wt_ = array(gemm_wt.shape(), float32, nullptr, {});
+    copy_cpu(gemm_wt, gemm_wt_, CopyType::Vector, stream);
+    temps.push_back(gemm_wt_);
+
+    // Calculate the total size of the spatial dimensions
+    int spatial_size = 1;
+    for (int d = 1; d < gemm_wt.ndim() - 1; ++d) {
+      spatial_size *= gemm_wt.shape(d);
+    }
+    encoder.set_output_array(gemm_wt_);
+    encoder.dispatch([gemm_wt_ptr = gemm_wt_.data<float>(),
+                      out_channels = gemm_wt.shape(0),
+                      in_channels = gemm_wt.shape(-1),
+                      spatial_size]() {
+      flip_spatial_dims_inplace<float>(
+          gemm_wt_ptr, in_channels, out_channels, spatial_size);
+    });
+    gemm_wt = gemm_wt_;
+  }
+
+  if (out.dtype() != float32) {
+    gemm_out = array(out.shape(), float32, nullptr, {});
+    gemm_out.set_data(allocator::malloc(gemm_out.nbytes()));
+    temps.push_back(gemm_out);
+  }
+
+  encoder.set_input_array(in_strided);
+  encoder.set_input_array(gemm_wt);
+  encoder.set_output_array(gemm_out);
+
+  encoder.dispatch([in_strided_ptr = in_strided.data<float>(),
+                    gemm_wt_ptr = gemm_wt.data<float>(),
+                    gemm_out_ptr = gemm_out.data<float>(),
+                    strided_reshape = std::move(strided_reshape),
+                    O]() {
+    // Perform gemm
+    cblas_sgemm(
+        CblasRowMajor,
+        CblasNoTrans, // no trans A
+        CblasTrans, // transB
+        strided_reshape[0], // M
+        O, // N
+        strided_reshape[1], // K
+        1.0f, // alpha
+        in_strided_ptr,
+        strided_reshape[1], // lda
+        gemm_wt_ptr,
+        strided_reshape[1], // ldb
+        0.0f, // beta
+        gemm_out_ptr,
+        O // ldc
+    );
+  });
+
+  // Copy results if needed
+  if (out.dtype() != float32) {
+    copy_cpu_inplace(gemm_out, out, CopyType::Vector, stream);
+  }
+  encoder.add_temporaries(std::move(temps));
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Conv routing
+///////////////////////////////////////////////////////////////////////////////
+
+void conv_1D_cpu(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  const int groups = in.shape().back() / wt.shape().back();
+  if (wt_dilation[0] == 1 && in_dilation[0] == 1 && !flip) {
+    return explicit_gemm_conv_1D_cpu(
+        in, wt, out, padding_lo, padding_hi, wt_strides, wt_dilation, stream);
+  }
+  if (wt_dilation[0] == 1 && in_dilation[0] == 1 && groups == 1) {
+    return explicit_gemm_conv_ND_cpu(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        flip,
+        stream);
+  }
+
+  return dispatch_slow_conv_1D(
+      in,
+      wt,
+      out,
+      padding_lo,
+      padding_hi,
+      wt_strides,
+      wt_dilation,
+      in_dilation,
+      flip,
+      stream);
+}
+
+void conv_2D_cpu(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  const int groups = in.shape().back() / wt.shape().back();
+  if (wt_dilation[0] == 1 && wt_dilation[1] == 1 && in_dilation[0] == 1 &&
+      in_dilation[1] == 1 && groups == 1) {
+    return explicit_gemm_conv_ND_cpu(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        flip,
+        stream);
+  }
+  return dispatch_slow_conv_2D(
+      in,
+      wt,
+      out,
+      padding_lo,
+      padding_hi,
+      wt_strides,
+      wt_dilation,
+      in_dilation,
+      flip,
+      stream);
+}
+
+void conv_3D_cpu(
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    Stream stream) {
+  const int groups = in.shape().back() / wt.shape().back();
+  if (wt_dilation[0] == 1 && wt_dilation[1] == 1 && wt_dilation[2] == 1 &&
+      in_dilation[0] == 1 && in_dilation[1] == 1 && in_dilation[2] == 1 &&
+      groups == 1) {
+    return explicit_gemm_conv_ND_cpu(
+        in,
+        wt,
+        out,
+        padding_lo,
+        padding_hi,
+        wt_strides,
+        wt_dilation,
+        flip,
+        stream);
+  }
+
+  return dispatch_slow_conv_3D(
+      in,
+      wt,
+      out,
+      padding_lo,
+      padding_hi,
+      wt_strides,
+      wt_dilation,
+      in_dilation,
+      flip,
+      stream);
+}
+
+} // namespace
+
+void Convolution::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& in = inputs[0];
+  auto& wt = inputs[1];
+
+  // 3D convolution
+  if (in.ndim() == (3 + 2)) {
+    return conv_3D_cpu(
+        in,
+        wt,
+        out,
+        padding_lo_,
+        padding_hi_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_,
+        stream());
+  }
+  // 2D convolution
+  else if (in.ndim() == (2 + 2)) {
+    return conv_2D_cpu(
+        in,
+        wt,
+        out,
+        padding_lo_,
+        padding_hi_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_,
+        stream());
+  }
+  // 1D convolution
+  else if (in.ndim() == (1 + 2)) {
+    return conv_1D_cpu(
+        in,
+        wt,
+        out,
+        padding_lo_,
+        padding_hi_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_,
+        stream());
+  }
+  // Throw error
+  else {
+    std::ostringstream msg;
+    msg << "[Convolution::eval] Convolution currently only supports"
+        << " 1D, 2D and 3D convolutions. Got inputs with " << in.ndim() - 2
+        << " spatial dimensions";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/copy.cpp b/macos/third_party/mlx/mlx/backend/cpu/copy.cpp
new file mode 100644
index 0000000..f9ff226
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/copy.cpp
@@ -0,0 +1,386 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <numeric>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename SrcT, typename DstT>
+void copy_single(const array& src, array& dst) {
+  auto src_ptr = src.data<SrcT>();
+  auto dst_ptr = dst.data<DstT>();
+  auto size = dst.size();
+  auto val = static_cast<DstT>(src_ptr[0]);
+  std::fill_n(dst_ptr, size, val);
+}
+
+template <typename SrcT, typename DstT>
+void copy_vector(const array& src, array& dst) {
+  auto src_ptr = src.data<SrcT>();
+  auto dst_ptr = dst.data<DstT>();
+  auto size = src.data_size();
+  std::copy(src_ptr, src_ptr + size, dst_ptr);
+}
+
+template <typename SrcT, typename DstT, int D>
+inline void copy_dims(
+    const SrcT* src,
+    DstT* dst,
+    const Shape& shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
+    int axis) {
+  auto stride_src = i_strides[axis];
+  auto stride_dst = o_strides[axis];
+  auto N = shape[axis];
+
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      copy_dims<SrcT, DstT, D - 1>(
+          src, dst, shape, i_strides, o_strides, axis + 1);
+    } else {
+      *dst = static_cast<DstT>(*src);
+    }
+    src += stride_src;
+    dst += stride_dst;
+  }
+}
+
+template <typename SrcT, typename DstT>
+void copy_general_general(
+    const array& src,
+    array& dst,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    const std::optional<array>& dynamic_i_offset,
+    const std::optional<array>& dynamic_o_offset) {
+  auto src_ptr = src.data<SrcT>() + i_offset;
+  auto dst_ptr = dst.data<DstT>() + o_offset;
+  auto i_offset_ptr =
+      dynamic_i_offset ? dynamic_i_offset->data<int64_t>() : nullptr;
+  auto o_offset_ptr =
+      dynamic_o_offset ? dynamic_o_offset->data<int64_t>() : nullptr;
+  auto size = src.size();
+  if (data_shape.empty()) {
+    auto val = static_cast<DstT>(*src_ptr);
+    *dst_ptr = val;
+    return;
+  }
+  auto [shape, strides] =
+      collapse_contiguous_dims(data_shape, {i_strides, o_strides});
+
+  int ndim = shape.size();
+  if (ndim < 3) {
+    if (i_offset_ptr) {
+      src_ptr += i_offset_ptr[0];
+    }
+    if (o_offset_ptr) {
+      dst_ptr += o_offset_ptr[0];
+    }
+
+    if (ndim == 1) {
+      copy_dims<SrcT, DstT, 1>(
+          src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+    } else if (ndim == 2) {
+      copy_dims<SrcT, DstT, 2>(
+          src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+    } else if (ndim == 3) {
+      copy_dims<SrcT, DstT, 3>(
+          src_ptr, dst_ptr, shape, strides[0], strides[1], 0);
+    }
+    return;
+  }
+  if (i_offset_ptr) {
+    src_ptr += i_offset_ptr[0];
+  }
+  if (o_offset_ptr) {
+    dst_ptr += o_offset_ptr[0];
+  }
+
+  ContiguousIterator in(shape, strides[0], ndim - 3);
+  ContiguousIterator out(shape, strides[1], ndim - 3);
+  auto stride = std::accumulate(
+      shape.end() - 3, shape.end(), 1, std::multiplies<int64_t>());
+  for (int64_t elem = 0; elem < size; elem += stride) {
+    copy_dims<SrcT, DstT, 3>(
+        src_ptr + in.loc,
+        dst_ptr + out.loc,
+        shape,
+        strides[0],
+        strides[1],
+        ndim - 3);
+    in.step();
+    out.step();
+  }
+}
+
+template <typename SrcT, typename DstT>
+inline void copy_general_general(const array& src, array& dst) {
+  copy_general_general<SrcT, DstT>(
+      src,
+      dst,
+      src.shape(),
+      src.strides(),
+      dst.strides(),
+      0,
+      0,
+      std::nullopt,
+      std::nullopt);
+}
+
+template <typename SrcT, typename DstT>
+void copy_general(
+    const array& src,
+    array& dst,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides&,
+    int64_t i_offset,
+    int64_t o_offset,
+    const std::optional<array>& dynamic_i_offset,
+    const std::optional<array>& dynamic_o_offset) {
+  copy_general_general<SrcT, DstT>(
+      src,
+      dst,
+      data_shape,
+      i_strides,
+      make_contiguous_strides(data_shape),
+      i_offset,
+      o_offset,
+      dynamic_i_offset,
+      dynamic_o_offset);
+}
+
+template <typename SrcT, typename DstT>
+inline void copy_general(const array& src, array& dst) {
+  copy_general_general<SrcT, DstT>(
+      src,
+      dst,
+      src.shape(),
+      src.strides(),
+      make_contiguous_strides(src.shape()),
+      0,
+      0,
+      std::nullopt,
+      std::nullopt);
+}
+
+template <typename SrcT, typename DstT, typename... Args>
+void copy(const array& src, array& dst, CopyType ctype, Args&&... args) {
+  switch (ctype) {
+    case CopyType::Scalar:
+      copy_single<SrcT, DstT>(src, dst);
+      return;
+    case CopyType::Vector:
+      copy_vector<SrcT, DstT>(src, dst);
+      return;
+    case CopyType::General:
+      copy_general<SrcT, DstT>(src, dst, std::forward<Args>(args)...);
+      return;
+    case CopyType::GeneralGeneral:
+      copy_general_general<SrcT, DstT>(src, dst, std::forward<Args>(args)...);
+      return;
+  }
+}
+
+template <typename SrcT, typename... Args>
+void copy(const array& src, array& dst, CopyType ctype, Args&&... args) {
+  switch (dst.dtype()) {
+    case bool_:
+      copy<SrcT, bool>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint8:
+      copy<SrcT, uint8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint16:
+      copy<SrcT, uint16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint32:
+      copy<SrcT, uint32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint64:
+      copy<SrcT, uint64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int8:
+      copy<SrcT, int8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int16:
+      copy<SrcT, int16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int32:
+      copy<SrcT, int32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int64:
+      copy<SrcT, int64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case float16:
+      copy<SrcT, float16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case float32:
+      copy<SrcT, float>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case float64:
+      copy<SrcT, double>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case bfloat16:
+      copy<SrcT, bfloat16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case complex64:
+      copy<SrcT, complex64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+  }
+}
+
+template <typename... Args>
+inline void copy_inplace_dispatch(
+    const array& src,
+    array& dst,
+    CopyType ctype,
+    Args&&... args) {
+  switch (src.dtype()) {
+    case bool_:
+      copy<bool>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint8:
+      copy<uint8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint16:
+      copy<uint16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint32:
+      copy<uint32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case uint64:
+      copy<uint64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int8:
+      copy<int8_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int16:
+      copy<int16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int32:
+      copy<int32_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case int64:
+      copy<int64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case float16:
+      copy<float16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case float32:
+      copy<float>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case float64:
+      copy<double>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case bfloat16:
+      copy<bfloat16_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+    case complex64:
+      copy<complex64_t>(src, dst, ctype, std::forward<Args>(args)...);
+      break;
+  }
+}
+
+} // namespace
+
+void copy_cpu_inplace(
+    const array& src,
+    array& dst,
+    CopyType ctype,
+    Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(src);
+  encoder.set_output_array(dst);
+  encoder.dispatch(
+      [src = array::unsafe_weak_copy(src),
+       dst = array::unsafe_weak_copy(dst),
+       ctype]() mutable { copy_inplace_dispatch(src, dst, ctype); });
+}
+
+void copy_cpu(const array& src, array& dst, CopyType ctype, Stream stream) {
+  bool donated = set_copy_output_data(src, dst, ctype);
+  if (donated && src.dtype() == dst.dtype()) {
+    // If the output has the same type as the input then there is nothing to
+    // copy, just use the buffer.
+    return;
+  }
+  if (ctype == CopyType::GeneralGeneral) {
+    ctype = CopyType::General;
+  }
+  copy_cpu_inplace(src, dst, ctype, stream);
+}
+
+void copy_cpu_inplace(
+    const array& src,
+    array& dst,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype,
+    Stream stream,
+    const std::optional<array>& dynamic_i_offset, /* = std::nullopt */
+    const std::optional<array>& dynamic_o_offset /* = std::nullopt */) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(src);
+  encoder.set_output_array(dst);
+  auto weak_copy_if_set = [](auto x) -> std::optional<array> {
+    if (x) {
+      return array::unsafe_weak_copy(*x);
+    } else {
+      return std::nullopt;
+    }
+  };
+  encoder.dispatch(
+      [src = array::unsafe_weak_copy(src),
+       dst = array::unsafe_weak_copy(dst),
+       data_shape,
+       i_strides,
+       o_strides,
+       i_offset,
+       o_offset,
+       ctype,
+       dynamic_i_offset = weak_copy_if_set(dynamic_i_offset),
+       dynamic_o_offset = weak_copy_if_set(dynamic_o_offset)]() mutable {
+        switch (ctype) {
+          case CopyType::General:
+          case CopyType::GeneralGeneral:
+            copy_inplace_dispatch(
+                src,
+                dst,
+                ctype,
+                data_shape,
+                i_strides,
+                o_strides,
+                i_offset,
+                o_offset,
+                dynamic_i_offset,
+                dynamic_o_offset);
+            break;
+          case CopyType::Scalar:
+          case CopyType::Vector:
+            copy_inplace_dispatch(src, dst, ctype);
+        }
+      });
+}
+
+array contiguous_copy_cpu(const array& arr, Stream stream) {
+  array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
+  copy_cpu(arr, arr_copy, CopyType::General, stream);
+  return arr_copy;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/copy.h b/macos/third_party/mlx/mlx/backend/cpu/copy.h
new file mode 100644
index 0000000..0072913
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/copy.h
@@ -0,0 +1,36 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <optional>
+
+#include "mlx/array.h"
+#include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+void copy_cpu(const array& src, array& dst, CopyType ctype, Stream stream);
+void copy_cpu_inplace(
+    const array& src,
+    array& dst,
+    CopyType ctype,
+    Stream stream);
+
+void copy_cpu_inplace(
+    const array& src,
+    array& dst,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype,
+    Stream stream,
+    const std::optional<array>& dynamic_i_offset = std::nullopt,
+    const std::optional<array>& dynamic_o_offset = std::nullopt);
+
+// Return a contiguous array with same shape that copies the data of |arr|.
+array contiguous_copy_cpu(const array& arr, Stream stream);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/device_info.cpp b/macos/third_party/mlx/mlx/backend/cpu/device_info.cpp
new file mode 100644
index 0000000..b709b2a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/device_info.cpp
@@ -0,0 +1,113 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cpu/device_info.h"
+
+#ifdef __APPLE__
+#include <sys/sysctl.h>
+#include <sys/utsname.h>
+#elif defined(_WIN32)
+#include <windows.h>
+#else
+#include <sys/utsname.h>
+#include <fstream>
+#endif
+
+namespace mlx::core::cpu {
+
+namespace {
+
+// Get CPU architecture string at runtime
+std::string get_cpu_architecture() {
+#ifdef _WIN32
+  // Use GetNativeSystemInfo to get the actual hardware architecture,
+  // even when running under WoW64 emulation
+  SYSTEM_INFO sysInfo;
+  GetNativeSystemInfo(&sysInfo);
+  switch (sysInfo.wProcessorArchitecture) {
+    case PROCESSOR_ARCHITECTURE_AMD64:
+      return "x86_64";
+    case PROCESSOR_ARCHITECTURE_ARM64:
+      return "arm64";
+    case PROCESSOR_ARCHITECTURE_INTEL:
+      return "x86";
+    case PROCESSOR_ARCHITECTURE_ARM:
+      return "arm";
+    default:
+      return "unknown";
+  }
+#else
+  // Use uname() for runtime detection on Unix-like systems.
+  // This returns the actual hardware architecture (e.g., "arm64" on Apple
+  // Silicon even when running x86_64 binaries via Rosetta 2)
+  struct utsname info;
+  if (uname(&info) == 0) {
+    return std::string(info.machine);
+  }
+  return "unknown";
+#endif
+}
+
+// Get CPU device name (brand string)
+std::string get_cpu_name() {
+#ifdef __APPLE__
+  char model[256];
+  size_t len = sizeof(model);
+  if (sysctlbyname("machdep.cpu.brand_string", &model, &len, NULL, 0) == 0) {
+    return std::string(model);
+  }
+#elif defined(_WIN32)
+  // Read CPU brand string from registry
+  HKEY hKey;
+  if (RegOpenKeyExA(
+          HKEY_LOCAL_MACHINE,
+          "HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0",
+          0,
+          KEY_READ,
+          &hKey) == ERROR_SUCCESS) {
+    char brand[256];
+    DWORD size = sizeof(brand);
+    if (RegQueryValueExA(
+            hKey, "ProcessorNameString", NULL, NULL, (LPBYTE)brand, &size) ==
+        ERROR_SUCCESS) {
+      RegCloseKey(hKey);
+      return std::string(brand);
+    }
+    RegCloseKey(hKey);
+  }
+#else
+  // Try reading from /proc/cpuinfo on Linux
+  std::ifstream cpuinfo("/proc/cpuinfo");
+  if (cpuinfo.is_open()) {
+    std::string line;
+    while (std::getline(cpuinfo, line)) {
+      if (line.starts_with("model name")) {
+        if (auto n = line.find(": "); n != std::string::npos) {
+          return line.substr(n + 2);
+        }
+      }
+    }
+  }
+#endif
+  return get_cpu_architecture();
+}
+
+} // anonymous namespace
+
+bool is_available() {
+  return true;
+}
+
+int device_count() {
+  return 1;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int /* device_index */) {
+  static auto info =
+      std::unordered_map<std::string, std::variant<std::string, size_t>>{
+          {"device_name", get_cpu_name()},
+          {"architecture", get_cpu_architecture()}};
+  return info;
+}
+
+} // namespace mlx::core::cpu
diff --git a/macos/third_party/mlx/mlx/backend/cpu/device_info.h b/macos/third_party/mlx/mlx/backend/cpu/device_info.h
new file mode 100644
index 0000000..1e23233
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/device_info.h
@@ -0,0 +1,28 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+namespace mlx::core::cpu {
+
+bool is_available();
+
+/**
+ * Get the number of available CPU devices.
+ *
+ * For CPU, always returns 1.
+ */
+int device_count();
+
+/**
+ * Get CPU device information.
+ *
+ * Returns a map with basic CPU device properties.
+ */
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index = 0);
+
+} // namespace mlx::core::cpu
diff --git a/macos/third_party/mlx/mlx/backend/cpu/distributed.cpp b/macos/third_party/mlx/mlx/backend/cpu/distributed.cpp
new file mode 100644
index 0000000..22dc4b4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/distributed.cpp
@@ -0,0 +1,103 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/distributed/primitives.h"
+
+namespace mlx::core::distributed {
+
+std::pair<array, bool> ensure_row_contiguous(const array& arr, Stream stream) {
+  if (arr.flags().row_contiguous) {
+    return {arr, false};
+  } else {
+    return {contiguous_copy_cpu(arr, stream), true};
+  }
+};
+
+void AllReduce::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto donate_or_copy = [s = stream()](const array& in, array& out) {
+    if (in.flags().row_contiguous) {
+      if (in.is_donatable()) {
+        out.copy_shared_buffer(in);
+      } else {
+        out.set_data(allocator::malloc(out.nbytes()));
+      }
+      return in;
+    } else {
+      array arr_copy = contiguous_copy_cpu(in, s);
+      out.copy_shared_buffer(arr_copy);
+      return arr_copy;
+    }
+  };
+
+  auto in = donate_or_copy(inputs[0], outputs[0]);
+  switch (reduce_type_) {
+    case Sum:
+      distributed::detail::all_sum(group(), in, outputs[0], stream());
+      break;
+    case Max:
+      distributed::detail::all_max(group(), in, outputs[0], stream());
+      break;
+    case Min:
+      distributed::detail::all_min(group(), in, outputs[0], stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "Only all reduce sum, min and max are supported for now");
+  }
+}
+
+void AllGather::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto [in, copied] = ensure_row_contiguous(inputs[0], stream());
+  outputs[0].set_data(allocator::malloc(outputs[0].nbytes()));
+  distributed::detail::all_gather(group(), in, outputs[0], stream());
+  if (copied) {
+    auto& enc = cpu::get_command_encoder(stream());
+    enc.add_temporary(in);
+  }
+}
+
+void Send::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto [in, copied] = ensure_row_contiguous(inputs[0], stream());
+  distributed::detail::send(group(), in, dst_, stream());
+  outputs[0].copy_shared_buffer(inputs[0]);
+  if (copied) {
+    auto& enc = cpu::get_command_encoder(stream());
+    enc.add_temporary(in);
+  }
+}
+
+void Recv::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 0);
+  assert(outputs.size() == 1);
+
+  outputs[0].set_data(allocator::malloc(outputs[0].nbytes()));
+  distributed::detail::recv(group(), outputs[0], src_, stream());
+}
+
+void ReduceScatter::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("[ReduceScatter] Not implemented yet.");
+}
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/backend/cpu/eig.cpp b/macos/third_party/mlx/mlx/backend/cpu/eig.cpp
new file mode 100644
index 0000000..fb63eef
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/eig.cpp
@@ -0,0 +1,281 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/array.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/linalg.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T>
+complex64_t to_complex(T r, T i) {
+  return {static_cast<float>(r), static_cast<float>(i)};
+}
+
+template <typename T, class Enable = void>
+struct EigWork {};
+
+template <typename T>
+struct EigWork<
+    T,
+    typename std::enable_if<std::is_floating_point<T>::value>::type> {
+  using O = complex64_t;
+
+  char jobl;
+  char jobr;
+  int N;
+  int lwork;
+  int info;
+  std::vector<array::Data> buffers;
+
+  EigWork(char jobl_, char jobr_, int N_, bool compute_eigenvectors)
+      : jobl(jobl_), jobr(jobr_), N(N_), lwork(-1) {
+    T work;
+    int n_vecs_l = compute_eigenvectors ? N_ : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        nullptr,
+        &N,
+        nullptr,
+        nullptr,
+        nullptr,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        &work,
+        &lwork,
+        &info);
+    lwork = static_cast<int>(work);
+
+    buffers.emplace_back(allocator::malloc(sizeof(T) * N * 2));
+    if (compute_eigenvectors) {
+      buffers.emplace_back(allocator::malloc(sizeof(T) * N * N * 2));
+    }
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, O* values, O* vectors) {
+    auto eig_tmp = static_cast<T*>(buffers[0].buffer.raw_ptr());
+    T* vec_tmp = nullptr;
+    if (vectors) {
+      vec_tmp = static_cast<T*>(buffers[1].buffer.raw_ptr());
+    }
+    auto work = static_cast<T*>(buffers.back().buffer.raw_ptr());
+
+    int n_vecs_l = vectors ? N : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        a,
+        &N,
+        eig_tmp,
+        eig_tmp + N,
+        vectors ? vec_tmp : nullptr,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        work,
+        &lwork,
+        &info);
+
+    for (int i = 0; i < N; ++i) {
+      values[i] = to_complex(eig_tmp[i], eig_tmp[N + i]);
+    }
+
+    if (vectors) {
+      for (int i = 0; i < N; ++i) {
+        if (values[i].imag() != 0) {
+          for (int j = 0; j < N; ++j) {
+            vectors[i * N + j] =
+                to_complex(vec_tmp[i * N + j], -vec_tmp[(i + 1) * N + j]);
+            vectors[(i + 1) * N + j] =
+                to_complex(vec_tmp[i * N + j], vec_tmp[(i + 1) * N + j]);
+          }
+          i += 1;
+        } else {
+          for (int j = 0; j < N; ++j) {
+            vectors[i * N + j] = to_complex(vec_tmp[i * N + j], T(0.0));
+          }
+        }
+      }
+    }
+  }
+};
+
+template <>
+struct EigWork<std::complex<float>> {
+  using T = std::complex<float>;
+  using R = float;
+  using O = T;
+
+  char jobl;
+  char jobr;
+  int N;
+  int lwork;
+  int lrwork;
+  int info;
+  std::vector<array::Data> buffers;
+
+  EigWork(char jobl_, char jobr_, int N_, bool compute_eigenvectors)
+      : jobl(jobl_), jobr(jobr_), N(N_), lwork(-1), lrwork(2 * N_) {
+    T work;
+    R rwork;
+    int n_vecs_l = compute_eigenvectors ? N_ : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        nullptr,
+        &N,
+        nullptr,
+        nullptr,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        &work,
+        &lwork,
+        &rwork,
+        &info);
+    lwork = static_cast<int>(work.real());
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+    buffers.emplace_back(allocator::malloc(sizeof(R) * lrwork));
+  }
+
+  void run(T* a, T* values, T* vectors) {
+    int n_vecs_l = vectors ? N : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        a,
+        &N,
+        values,
+        vectors,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        static_cast<T*>(buffers[0].buffer.raw_ptr()),
+        &lwork,
+        static_cast<R*>(buffers[1].buffer.raw_ptr()),
+        &info);
+  }
+};
+
+template <typename T>
+void eig_impl(
+    array& a,
+    array& vectors,
+    array& values,
+    bool compute_eigenvectors,
+    Stream stream) {
+  auto a_ptr = a.data<T>();
+  auto val_ptr = values.data<complex64_t>();
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(values);
+  complex64_t* vec_ptr = nullptr;
+  if (compute_eigenvectors) {
+    encoder.set_output_array(vectors);
+    vec_ptr = vectors.data<complex64_t>();
+  }
+  encoder.dispatch([a_ptr,
+                    val_ptr,
+                    vec_ptr,
+                    compute_eigenvectors,
+                    N = vectors.shape(-1),
+                    size = vectors.size()]() mutable {
+    char jobr = 'N';
+    char jobl = compute_eigenvectors ? 'V' : 'N';
+
+    EigWork<T> work(jobl, jobr, N, compute_eigenvectors);
+
+    for (size_t i = 0; i < size / (N * N); ++i) {
+      work.run(a_ptr, val_ptr, vec_ptr);
+      a_ptr += N * N;
+      val_ptr += N;
+      if (vec_ptr) {
+        vec_ptr += N * N;
+      }
+      if (work.info != 0) {
+        std::stringstream msg;
+        msg << "[Eig::eval_cpu] Eigenvalue decomposition failed with error code "
+            << work.info;
+        throw std::runtime_error(msg.str());
+      }
+    }
+  });
+  encoder.add_temporary(a);
+}
+
+} // namespace
+
+void Eig::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  const auto& a = inputs[0];
+  auto& values = outputs[0];
+
+  auto vectors = compute_eigenvectors_
+      ? outputs[1]
+      : array(a.shape(), complex64, nullptr, {});
+
+  auto a_copy = array(a.shape(), a.dtype(), nullptr, {});
+  copy_cpu(
+      a,
+      a_copy,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream());
+
+  values.set_data(allocator::malloc(values.nbytes()));
+
+  if (compute_eigenvectors_) {
+    // Set the strides and flags so the eigenvectors
+    // are in the columns of the output
+    auto flags = vectors.flags();
+    auto strides = vectors.strides();
+    auto ndim = a.ndim();
+    std::swap(strides[ndim - 1], strides[ndim - 2]);
+
+    if (a.size() > 1) {
+      flags.row_contiguous = false;
+      if (ndim > 2) {
+        flags.col_contiguous = false;
+      } else {
+        flags.col_contiguous = true;
+      }
+    }
+    vectors.set_data(
+        allocator::malloc(vectors.nbytes()), vectors.size(), strides, flags);
+  }
+  switch (a.dtype()) {
+    case float32:
+      eig_impl<float>(a_copy, vectors, values, compute_eigenvectors_, stream());
+      break;
+    case float64:
+      eig_impl<double>(
+          a_copy, vectors, values, compute_eigenvectors_, stream());
+      break;
+    case complex64:
+      eig_impl<std::complex<float>>(
+          a_copy, vectors, values, compute_eigenvectors_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Eig::eval_cpu] only supports float32, float64, or complex64.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/eigh.cpp b/macos/third_party/mlx/mlx/backend/cpu/eigh.cpp
new file mode 100644
index 0000000..d457c1f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/eigh.cpp
@@ -0,0 +1,241 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/array.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/linalg.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T, class Enable = void>
+struct EighWork {};
+
+template <typename T>
+struct EighWork<
+    T,
+    typename std::enable_if<std::is_floating_point<T>::value>::type> {
+  using R = T;
+
+  char jobz;
+  char uplo;
+  int N;
+  int lwork;
+  int liwork;
+  int info;
+  std::vector<array::Data> buffers;
+
+  EighWork(char jobz_, char uplo_, int N_)
+      : jobz(jobz_), uplo(uplo_), N(N_), lwork(-1), liwork(-1) {
+    T work;
+    int iwork;
+    syevd<T>(
+        &jobz,
+        &uplo,
+        &N,
+        nullptr,
+        &N,
+        nullptr,
+        &work,
+        &lwork,
+        &iwork,
+        &liwork,
+        &info);
+    lwork = static_cast<int>(work);
+    liwork = iwork;
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+    buffers.emplace_back(allocator::malloc(sizeof(int) * liwork));
+  }
+
+  void run(T* vectors, T* values) {
+    syevd<T>(
+        &jobz,
+        &uplo,
+        &N,
+        vectors,
+        &N,
+        values,
+        static_cast<T*>(buffers[0].buffer.raw_ptr()),
+        &lwork,
+        static_cast<int*>(buffers[1].buffer.raw_ptr()),
+        &liwork,
+        &info);
+  }
+};
+
+template <>
+struct EighWork<std::complex<float>> {
+  using T = std::complex<float>;
+  using R = float;
+
+  char jobz;
+  char uplo;
+  int N;
+  int lwork;
+  int lrwork;
+  int liwork;
+  int info;
+  std::vector<array::Data> buffers;
+
+  EighWork(char jobz_, char uplo_, int N_)
+      : jobz(jobz_), uplo(uplo_), N(N_), lwork(-1), lrwork(-1), liwork(-1) {
+    T work;
+    R rwork;
+    int iwork;
+    heevd<T>(
+        &jobz,
+        &uplo,
+        &N,
+        nullptr,
+        &N,
+        nullptr,
+        &work,
+        &lwork,
+        &rwork,
+        &lrwork,
+        &iwork,
+        &liwork,
+        &info);
+    lwork = static_cast<int>(work.real());
+    lrwork = static_cast<int>(rwork);
+    liwork = iwork;
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+    buffers.emplace_back(allocator::malloc(sizeof(R) * lrwork));
+    buffers.emplace_back(allocator::malloc(sizeof(int) * liwork));
+  }
+
+  void run(T* vectors, R* values) {
+    heevd<T>(
+        &jobz,
+        &uplo,
+        &N,
+        vectors,
+        &N,
+        values,
+        static_cast<T*>(buffers[0].buffer.raw_ptr()),
+        &lwork,
+        static_cast<R*>(buffers[1].buffer.raw_ptr()),
+        &lrwork,
+        static_cast<int*>(buffers[2].buffer.raw_ptr()),
+        &liwork,
+        &info);
+    if (jobz == 'V') {
+      // We have pre-transposed the vectors but we also must conjugate them
+      // when they are complex.
+      //
+      // We could vectorize this but it is so fast in comparison to heevd that
+      // it doesn't really matter.
+      for (int i = 0; i < N; i++) {
+        for (int j = 0; j < N; j++) {
+          *vectors = std::conj(*vectors);
+          vectors++;
+        }
+      }
+    }
+  }
+};
+
+template <typename T>
+void eigh_impl(
+    array& vectors,
+    array& values,
+    const std::string& uplo,
+    bool compute_eigenvectors,
+    Stream stream) {
+  using R = typename EighWork<T>::R;
+
+  auto vec_ptr = vectors.data<T>();
+  auto eig_ptr = values.data<R>();
+  char jobz = compute_eigenvectors ? 'V' : 'N';
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(vectors);
+  encoder.set_output_array(values);
+  encoder.dispatch([vec_ptr,
+                    eig_ptr,
+                    jobz,
+                    uplo = uplo[0],
+                    N = vectors.shape(-1),
+                    size = vectors.size()]() mutable {
+    // Work query
+    EighWork<T> work(jobz, uplo, N);
+
+    // Work loop
+    for (size_t i = 0; i < size / (N * N); ++i) {
+      work.run(vec_ptr, eig_ptr);
+      vec_ptr += N * N;
+      eig_ptr += N;
+      if (work.info != 0) {
+        std::stringstream msg;
+        msg << "[Eigh::eval_cpu] Eigenvalue decomposition failed with error code "
+            << work.info;
+        throw std::runtime_error(msg.str());
+      }
+    }
+  });
+  if (!compute_eigenvectors) {
+    encoder.add_temporary(vectors);
+  }
+}
+
+} // namespace
+
+void Eigh::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  const auto& a = inputs[0];
+  auto& values = outputs[0];
+
+  auto vectors = compute_eigenvectors_
+      ? outputs[1]
+      : array(a.shape(), a.dtype(), nullptr, {});
+
+  values.set_data(allocator::malloc(values.nbytes()));
+
+  copy_cpu(
+      a,
+      vectors,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream());
+
+  if (compute_eigenvectors_) {
+    // Set the strides and flags so the eigenvectors
+    // are in the columns of the output
+    auto flags = vectors.flags();
+    auto strides = vectors.strides();
+    auto ndim = a.ndim();
+    std::swap(strides[ndim - 1], strides[ndim - 2]);
+
+    if (a.size() > 1) {
+      flags.row_contiguous = false;
+      if (ndim > 2) {
+        flags.col_contiguous = false;
+      } else {
+        flags.col_contiguous = true;
+      }
+    }
+    vectors.copy_shared_buffer(vectors, strides, flags, vectors.data_size());
+  }
+  switch (a.dtype()) {
+    case float32:
+      eigh_impl<float>(vectors, values, uplo_, compute_eigenvectors_, stream());
+      break;
+    case float64:
+      eigh_impl<double>(
+          vectors, values, uplo_, compute_eigenvectors_, stream());
+      break;
+    case complex64:
+      eigh_impl<std::complex<float>>(
+          vectors, values, uplo_, compute_eigenvectors_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Eigh::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/encoder.cpp b/macos/third_party/mlx/mlx/backend/cpu/encoder.cpp
new file mode 100644
index 0000000..f3fedef
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/encoder.cpp
@@ -0,0 +1,16 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cpu/encoder.h"
+
+namespace mlx::core::cpu {
+
+CommandEncoder& get_command_encoder(Stream stream) {
+  static std::unordered_map<int, CommandEncoder> encoder_map;
+  auto it = encoder_map.find(stream.index);
+  if (it == encoder_map.end()) {
+    it = encoder_map.emplace(stream.index, stream).first;
+  }
+  return it->second;
+}
+
+} // namespace mlx::core::cpu
diff --git a/macos/third_party/mlx/mlx/backend/cpu/encoder.h b/macos/third_party/mlx/mlx/backend/cpu/encoder.h
new file mode 100644
index 0000000..e04179e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/encoder.h
@@ -0,0 +1,67 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+
+#include "mlx/array.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core::cpu {
+
+// Number of dispatches per scheduler task
+constexpr int DISPATCHES_PER_TASK = 10;
+
+struct MLX_API CommandEncoder {
+  CommandEncoder(Stream stream) : stream_(stream) {}
+
+  CommandEncoder(const CommandEncoder&) = delete;
+  CommandEncoder& operator=(const CommandEncoder&) = delete;
+  CommandEncoder(CommandEncoder&&) = delete;
+  CommandEncoder& operator=(CommandEncoder&&) = delete;
+
+  void set_input_array(const array& a) {}
+  void set_output_array(array& a) {}
+
+  // Hold onto a temporary until any already scheduled tasks which use it as
+  // an input are complete.
+  void add_temporary(array arr) {
+    temporaries_.push_back(std::move(arr));
+  }
+
+  void add_temporaries(std::vector<array> arrays) {
+    temporaries_.insert(
+        temporaries_.end(),
+        std::make_move_iterator(arrays.begin()),
+        std::make_move_iterator(arrays.end()));
+  }
+
+  std::vector<array>& temporaries() {
+    return temporaries_;
+  }
+
+  template <class F, class... Args>
+  void dispatch(F&& f, Args&&... args) {
+    num_ops_ = (num_ops_ + 1) % DISPATCHES_PER_TASK;
+    auto task = std::bind(std::forward<F>(f), std::forward<Args>(args)...);
+    if (num_ops_ == 0) {
+      scheduler::notify_new_task(stream_);
+      auto task_wrap = [s = stream_, task = std::move(task)]() mutable {
+        task();
+        scheduler::notify_task_completion(s);
+      };
+      scheduler::enqueue(stream_, std::move(task_wrap));
+    } else {
+      scheduler::enqueue(stream_, std::move(task));
+    }
+  }
+
+ private:
+  Stream stream_;
+  std::vector<array> temporaries_;
+  int num_ops_{0};
+};
+
+MLX_API CommandEncoder& get_command_encoder(Stream stream);
+
+} // namespace mlx::core::cpu
diff --git a/macos/third_party/mlx/mlx/backend/cpu/eval.cpp b/macos/third_party/mlx/mlx/backend/cpu/eval.cpp
new file mode 100644
index 0000000..b23c8d5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/eval.cpp
@@ -0,0 +1,40 @@
+// Copyright © 2025 Apple Inc.
+#include "mlx/backend/cpu/eval.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::cpu {
+
+void eval(array& arr) {
+  auto s = arr.primitive().stream();
+
+  auto outputs = arr.outputs();
+  {
+    // If the array is a tracer hold a reference
+    // to its inputs so they don't get donated
+    std::vector<array> inputs;
+    if (arr.is_tracer()) {
+      inputs = arr.inputs();
+    }
+    arr.primitive().eval_cpu(arr.inputs(), outputs);
+  }
+
+  std::unordered_set<std::shared_ptr<array::Data>> buffers;
+  for (auto& in : arr.inputs()) {
+    buffers.insert(in.data_shared_ptr());
+  }
+  for (auto& s : arr.siblings()) {
+    buffers.insert(s.data_shared_ptr());
+  }
+  // Remove the output if it was donated to by an input
+  if (auto it = buffers.find(arr.data_shared_ptr()); it != buffers.end()) {
+    buffers.erase(it);
+  }
+  auto& encoder = cpu::get_command_encoder(s);
+  encoder.dispatch([buffers = std::move(buffers),
+                    temps = std::move(encoder.temporaries())]() {});
+}
+
+} // namespace mlx::core::cpu
diff --git a/macos/third_party/mlx/mlx/backend/cpu/eval.h b/macos/third_party/mlx/mlx/backend/cpu/eval.h
new file mode 100644
index 0000000..20156d6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/eval.h
@@ -0,0 +1,12 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/stream.h"
+
+namespace mlx::core::cpu {
+
+void eval(array& arr);
+
+} // namespace mlx::core::cpu
diff --git a/macos/third_party/mlx/mlx/backend/cpu/fft.cpp b/macos/third_party/mlx/mlx/backend/cpu/fft.cpp
new file mode 100644
index 0000000..d9e5f80
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/fft.cpp
@@ -0,0 +1,120 @@
+// Copyright © 2023 Apple Inc.
+
+#include <numeric>
+
+#include "mlx/3rdparty/pocketfft.h"
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void FFT::eval_cpu(const std::vector<array>& inputs, array& out) {
+  auto& in = inputs[0];
+  std::vector<std::ptrdiff_t> strides_in(
+      in.strides().begin(), in.strides().end());
+  for (auto& s : strides_in) {
+    s *= in.itemsize();
+  }
+  std::vector<std::ptrdiff_t> strides_out(
+      out.strides().begin(), out.strides().end());
+  for (auto& s : strides_out) {
+    s *= out.itemsize();
+  }
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  std::vector<size_t> shape;
+  if (out.dtype() == float32) {
+    shape.insert(shape.end(), out.shape().begin(), out.shape().end());
+  } else {
+    shape.insert(shape.end(), in.shape().begin(), in.shape().end());
+  }
+
+  float scale = 1.0f;
+  if (inverse_) {
+    size_t nelem = std::accumulate(
+        axes_.begin(), axes_.end(), 1, [&shape](auto x, auto y) {
+          return x * shape[y];
+        });
+    scale /= nelem;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+
+  if (in.dtype() == complex64 && out.dtype() == complex64) {
+    auto in_ptr =
+        reinterpret_cast<const std::complex<float>*>(in.data<complex64_t>());
+    auto out_ptr =
+        reinterpret_cast<std::complex<float>*>(out.data<complex64_t>());
+    encoder.dispatch([shape = std::move(shape),
+                      strides_in = std::move(strides_in),
+                      strides_out = std::move(strides_out),
+                      axes = axes_,
+                      inverse = inverse_,
+                      in_ptr,
+                      out_ptr,
+                      scale]() {
+      pocketfft::c2c(
+          shape,
+          strides_in,
+          strides_out,
+          axes,
+          !inverse,
+          in_ptr,
+          out_ptr,
+          scale);
+    });
+  } else if (in.dtype() == float32 && out.dtype() == complex64) {
+    auto in_ptr = in.data<float>();
+    auto out_ptr =
+        reinterpret_cast<std::complex<float>*>(out.data<complex64_t>());
+    encoder.dispatch([shape = std::move(shape),
+                      strides_in = std::move(strides_in),
+                      strides_out = std::move(strides_out),
+                      axes = axes_,
+                      inverse = inverse_,
+                      in_ptr,
+                      out_ptr,
+                      scale]() {
+      pocketfft::r2c(
+          shape,
+          strides_in,
+          strides_out,
+          axes,
+          !inverse,
+          in_ptr,
+          out_ptr,
+          scale);
+    });
+  } else if (in.dtype() == complex64 && out.dtype() == float32) {
+    auto in_ptr =
+        reinterpret_cast<const std::complex<float>*>(in.data<complex64_t>());
+    auto out_ptr = out.data<float>();
+    encoder.dispatch([shape = std::move(shape),
+                      strides_in = std::move(strides_in),
+                      strides_out = std::move(strides_out),
+                      axes = axes_,
+                      inverse = inverse_,
+                      in_ptr,
+                      out_ptr,
+                      scale]() {
+      pocketfft::c2r(
+          shape,
+          strides_in,
+          strides_out,
+          axes,
+          !inverse,
+          in_ptr,
+          out_ptr,
+          scale);
+    });
+  } else {
+    throw std::runtime_error(
+        "[FFT] Received unexpected input and output type combination.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/gemm.h b/macos/third_party/mlx/mlx/backend/cpu/gemm.h
new file mode 100644
index 0000000..d665cb9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/gemm.h
@@ -0,0 +1,26 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+template <typename T>
+void matmul(
+    const T* a,
+    const T* b,
+    T* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/gemms/bnns.cpp b/macos/third_party/mlx/mlx/backend/cpu/gemms/bnns.cpp
new file mode 100644
index 0000000..2ec0fd4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/gemms/bnns.cpp
@@ -0,0 +1,214 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <Accelerate/Accelerate.h>
+
+#include "mlx/array.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/gemm.h"
+#include "mlx/dtype.h"
+
+namespace mlx::core {
+
+template <typename T>
+constexpr BNNSDataType to_bnns_dtype();
+
+template <>
+constexpr BNNSDataType to_bnns_dtype<float>() {
+  return BNNSDataType(BNNSDataTypeFloatBit | 32);
+}
+template <>
+constexpr BNNSDataType to_bnns_dtype<float16_t>() {
+  return BNNSDataType(BNNSDataTypeFloatBit | 16);
+}
+
+template <>
+constexpr BNNSDataType to_bnns_dtype<bfloat16_t>() {
+  return BNNSDataTypeBFloat16;
+}
+
+template <typename T>
+void matmul_bnns(
+    const T* a,
+    const T* b,
+    T* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
+
+  BNNSDataType bnns_dtype = to_bnns_dtype<T>();
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
+  if (beta != 1.0 && beta != 0.0) {
+    // scale the output
+    for (auto i = 0; i < batch_size * M * N; ++i) {
+      out[i] *= beta;
+    }
+    beta = 1.0;
+  }
+  const BNNSLayerParametersBroadcastMatMul gemm_params{
+      /* float alpha = */ alpha,
+      /* float beta = */ beta,
+      /* bool transA = */ a_transposed,
+      /* bool transB = */ b_transposed,
+      /* bool quadratic = */ false,
+      /* bool a_is_weights = */ false,
+      /* bool b_is_weights = */ false,
+      /* BNNSNDArrayDescriptor iA_desc = */
+      BNNSNDArrayDescriptor{
+          /* BNNSNDArrayFlags flags = */ BNNSNDArrayFlagBackpropSet,
+          /* BNNSDataLayout layout = */ BNNSDataLayoutRowMajorMatrix,
+
+          /* size_t size[BNNS_MAX_TENSOR_DIMENSION] = */
+          {lda, (M * K) / lda, 0, 0, 0, 0, 0, 0},
+          /* size_t stride[BNNS_MAX_TENSOR_DIMENSION] = */
+          {1, lda, 0, 0, 0, 0, 0, 0},
+
+          /* void * _Nullable data = */ nullptr,
+          /* BNNSDataType data_type = */ bnns_dtype,
+
+          /* void * _Nullable table_data = */ nullptr,
+          /* BNNSDataType table_data_type = */ bnns_dtype,
+
+          /* float data_scale = */ 1.0,
+          /* float data_bias = */ 0.0,
+      },
+      /* BNNSNDArrayDescriptor iB_desc = */
+      BNNSNDArrayDescriptor{
+          /* BNNSNDArrayFlags flags = */ BNNSNDArrayFlagBackpropSet,
+          /* BNNSDataLayout layout = */ BNNSDataLayoutRowMajorMatrix,
+
+          /* size_t size[BNNS_MAX_TENSOR_DIMENSION] = */
+          {ldb, (K * N) / ldb, 0, 0, 0, 0, 0, 0},
+          /* size_t stride[BNNS_MAX_TENSOR_DIMENSION] = */
+          {1, ldb, 0, 0, 0, 0, 0, 0},
+
+          /* void * _Nullable data = */ nullptr,
+          /* BNNSDataType data_type = */ bnns_dtype,
+
+          /* void * _Nullable table_data = */ nullptr,
+          /* BNNSDataType table_data_type = */ bnns_dtype,
+
+          /* float data_scale = */ 1.0,
+          /* float data_bias = */ 0.0,
+      },
+      /* BNNSNDArrayDescriptor o_desc = */
+      BNNSNDArrayDescriptor{
+          /* BNNSNDArrayFlags flags = */ BNNSNDArrayFlagBackpropSet,
+          /* BNNSDataLayout layout = */ BNNSDataLayoutRowMajorMatrix,
+
+          /* size_t size[BNNS_MAX_TENSOR_DIMENSION] = */
+          {N, M, 0, 0, 0, 0, 0, 0},
+          /* size_t stride[BNNS_MAX_TENSOR_DIMENSION] = */
+          {1, N, 0, 0, 0, 0, 0, 0},
+
+          /* void * _Nullable data = */ nullptr,
+          /* BNNSDataType data_type = */ bnns_dtype,
+
+          /* void * _Nullable table_data = */ nullptr,
+          /* BNNSDataType table_data_type = */ bnns_dtype,
+
+          /* float data_scale = */ 1.0,
+          /* float data_bias = */ 0.0,
+      },
+  };
+
+  auto bnns_filter =
+      BNNSFilterCreateLayerBroadcastMatMul(&gemm_params, nullptr);
+
+  for (int i = 0; i < batch_size; ++i) {
+    BNNSFilterApplyTwoInput(
+        bnns_filter,
+        reinterpret_cast<const uint8_t*>(
+            a + elem_to_loc(M * K * i, a_shape, a_strides)),
+        reinterpret_cast<const uint8_t*>(
+            b + elem_to_loc(K * N * i, b_shape, b_strides)),
+        reinterpret_cast<uint8_t*>(out + M * N * i));
+  }
+
+  BNNSFilterDestroy(bnns_filter);
+#pragma GCC diagnostic pop
+}
+
+template <>
+void matmul<float16_t>(
+    const float16_t* a,
+    const float16_t* b,
+    float16_t* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  matmul_bnns(
+      a,
+      b,
+      out,
+      a_transposed,
+      b_transposed,
+      lda,
+      ldb,
+      ldc,
+      alpha,
+      beta,
+      batch_size,
+      a_shape,
+      a_strides,
+      b_shape,
+      b_strides);
+}
+
+template <>
+void matmul<bfloat16_t>(
+    const bfloat16_t* a,
+    const bfloat16_t* b,
+    bfloat16_t* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  matmul_bnns(
+      a,
+      b,
+      out,
+      a_transposed,
+      b_transposed,
+      lda,
+      ldb,
+      ldc,
+      alpha,
+      beta,
+      batch_size,
+      a_shape,
+      a_strides,
+      b_shape,
+      b_strides);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/gemms/cblas.cpp b/macos/third_party/mlx/mlx/backend/cpu/gemms/cblas.cpp
new file mode 100644
index 0000000..765e9f5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/gemms/cblas.cpp
@@ -0,0 +1,134 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/gemm.h"
+#include "mlx/backend/cpu/lapack.h"
+
+namespace mlx::core {
+
+template <>
+void matmul<float>(
+    const float* a,
+    const float* b,
+    float* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
+
+  for (int i = 0; i < batch_size; ++i) {
+    cblas_sgemm(
+        CblasRowMajor,
+        a_transposed ? CblasTrans : CblasNoTrans, // transA
+        b_transposed ? CblasTrans : CblasNoTrans, // transB
+        M,
+        N,
+        K,
+        alpha,
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
+        lda,
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
+        ldb,
+        beta,
+        out + M * N * i,
+        ldc);
+  }
+}
+
+template <>
+void matmul<double>(
+    const double* a,
+    const double* b,
+    double* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
+
+  for (int i = 0; i < batch_size; ++i) {
+    cblas_dgemm(
+        CblasRowMajor,
+        a_transposed ? CblasTrans : CblasNoTrans, // transA
+        b_transposed ? CblasTrans : CblasNoTrans, // transB
+        M,
+        N,
+        K,
+        alpha,
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
+        lda,
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
+        ldb,
+        beta,
+        out + M * N * i,
+        ldc);
+  }
+}
+
+template <>
+void matmul<complex64_t>(
+    const complex64_t* a,
+    const complex64_t* b,
+    complex64_t* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
+  auto calpha = static_cast<complex64_t>(alpha);
+  auto cbeta = static_cast<complex64_t>(beta);
+
+  for (int i = 0; i < batch_size; ++i) {
+    cblas_cgemm(
+        CblasRowMajor,
+        a_transposed ? CblasTrans : CblasNoTrans, // transA
+        b_transposed ? CblasTrans : CblasNoTrans, // transB
+        M,
+        N,
+        K,
+        &calpha,
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
+        lda,
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
+        ldb,
+        &cbeta,
+        out + M * N * i,
+        ldc);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_bf16.cpp b/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_bf16.cpp
new file mode 100644
index 0000000..58f5964
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_bf16.cpp
@@ -0,0 +1,45 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/gemm.h"
+#include "mlx/backend/cpu/gemms/simd_gemm.h"
+
+namespace mlx::core {
+
+template <>
+void matmul<bfloat16_t>(
+    const bfloat16_t* a,
+    const bfloat16_t* b,
+    bfloat16_t* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
+  for (int i = 0; i < batch_size; ++i) {
+    simd_gemm<bfloat16_t, float>(
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
+        out + M * N * i,
+        a_transposed,
+        b_transposed,
+        M,
+        N,
+        K,
+        alpha,
+        beta);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_fp16.cpp b/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_fp16.cpp
new file mode 100644
index 0000000..93467da
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_fp16.cpp
@@ -0,0 +1,45 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/gemm.h"
+#include "mlx/backend/cpu/gemms/simd_gemm.h"
+
+namespace mlx::core {
+
+template <>
+void matmul<float16_t>(
+    const float16_t* a,
+    const float16_t* b,
+    float16_t* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    size_t ldc,
+    float alpha,
+    float beta,
+    size_t batch_size,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides) {
+  auto ndim = a_shape.size();
+  size_t M = a_shape[ndim - 2];
+  size_t N = b_shape[ndim - 1];
+  size_t K = a_shape[ndim - 1];
+  for (int i = 0; i < batch_size; ++i) {
+    simd_gemm<float16_t, float>(
+        a + elem_to_loc(M * K * i, a_shape, a_strides),
+        b + elem_to_loc(K * N * i, b_shape, b_strides),
+        out + M * N * i,
+        a_transposed,
+        b_transposed,
+        M,
+        N,
+        K,
+        alpha,
+        beta);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_gemm.h b/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_gemm.h
new file mode 100644
index 0000000..a23c7de
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/gemms/simd_gemm.h
@@ -0,0 +1,139 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+#include "mlx/backend/cpu/simd/simd.h"
+
+namespace mlx::core {
+
+inline int ceildiv(int a, int b) {
+  return (a + b - 1) / b;
+}
+
+template <int block_size, typename T, typename AccT>
+void load_block(
+    const T* in,
+    AccT* out,
+    int M,
+    int N,
+    int i,
+    int j,
+    bool transpose) {
+  if (transpose) {
+    for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
+      for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
+        out[jj * block_size + ii] =
+            in[(i * block_size + ii) * N + j * block_size + jj];
+      }
+    }
+  } else {
+    for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
+      for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
+        out[ii * block_size + jj] =
+            in[(i * block_size + ii) * N + j * block_size + jj];
+      }
+    }
+  }
+}
+
+template <typename T, typename AccT>
+void simd_gemm(
+    const T* a,
+    const T* b,
+    T* c,
+    bool a_trans,
+    bool b_trans,
+    int M,
+    int N,
+    int K,
+    float alpha,
+    float beta) {
+  constexpr int block_size = 16;
+  constexpr int simd_size = simd::max_size<AccT>;
+  static_assert(
+      (block_size % simd_size) == 0,
+      "Block size must be divisible by SIMD size");
+
+  int last_k_block_size = K - block_size * (K / block_size);
+  int last_k_simd_block = (last_k_block_size / simd_size) * simd_size;
+  for (int i = 0; i < ceildiv(M, block_size); i++) {
+    for (int j = 0; j < ceildiv(N, block_size); j++) {
+      AccT c_block[block_size * block_size] = {0.0};
+      AccT a_block[block_size * block_size];
+      AccT b_block[block_size * block_size];
+
+      int k = 0;
+      for (; k < K / block_size; k++) {
+        // Load a and b blocks
+        if (a_trans) {
+          load_block<block_size>(a, a_block, K, M, k, i, true);
+        } else {
+          load_block<block_size>(a, a_block, M, K, i, k, false);
+        }
+        if (b_trans) {
+          load_block<block_size>(b, b_block, N, K, j, k, false);
+        } else {
+          load_block<block_size>(b, b_block, K, N, k, j, true);
+        }
+
+        // Multiply and accumulate
+        for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
+          for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
+            for (int kk = 0; kk < block_size; kk += simd_size) {
+              auto av =
+                  simd::load<AccT, simd_size>(a_block + ii * block_size + kk);
+              auto bv =
+                  simd::load<AccT, simd_size>(b_block + jj * block_size + kk);
+              c_block[ii * block_size + jj] += simd::sum(av * bv);
+            }
+          }
+        }
+      }
+      if (last_k_block_size) {
+        // Load a and b blocks
+        if (a_trans) {
+          load_block<block_size>(a, a_block, K, M, k, i, true);
+        } else {
+          load_block<block_size>(a, a_block, M, K, i, k, false);
+        }
+        if (b_trans) {
+          load_block<block_size>(b, b_block, N, K, j, k, false);
+        } else {
+          load_block<block_size>(b, b_block, K, N, k, j, true);
+        }
+
+        // Multiply and accumulate
+        for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
+          for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
+            int kk = 0;
+            for (; kk < last_k_simd_block; kk += simd_size) {
+              auto av =
+                  simd::load<AccT, simd_size>(a_block + ii * block_size + kk);
+              auto bv =
+                  simd::load<AccT, simd_size>(b_block + jj * block_size + kk);
+              c_block[ii * block_size + jj] += simd::sum(av * bv);
+            }
+            for (; kk < last_k_block_size; ++kk) {
+              c_block[ii * block_size + jj] +=
+                  a_block[ii * block_size + kk] * b_block[jj * block_size + kk];
+            }
+          }
+        }
+      }
+
+      // Store
+      for (int ii = 0; ii < block_size && i * block_size + ii < M; ++ii) {
+        for (int jj = 0; jj < block_size && j * block_size + jj < N; ++jj) {
+          auto c_idx = (i * block_size + ii) * N + j * block_size + jj;
+          if (beta != 0) {
+            c[c_idx] = static_cast<T>(
+                alpha * c_block[ii * block_size + jj] + beta * c[c_idx]);
+          } else {
+            c[c_idx] = static_cast<T>(alpha * c_block[ii * block_size + jj]);
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/hadamard.cpp b/macos/third_party/mlx/mlx/backend/cpu/hadamard.cpp
new file mode 100644
index 0000000..bf7e1dc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/hadamard.cpp
@@ -0,0 +1,121 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/common/hadamard.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+// n = 2^k component
+template <typename T>
+void hadamard_n(T* out, int n, int m, float scale, size_t size) {
+  for (int b = 0; b < size / n; b++) {
+    size_t loc = b * n;
+    T* data_ptr = out + loc;
+    int h = 1;
+    int n_over_2 = n / 2;
+    while (h < n) {
+      for (int i = 0; i < n / 2; i++) {
+        int k = i & (h - 1);
+        int j = ((i - k) << 1) + k;
+        float x = *(data_ptr + j);
+        float y = *(data_ptr + j + h);
+        *(data_ptr + j) = x + y;
+        *(data_ptr + j + h) = x - y;
+        if (h == n_over_2) {
+          *(data_ptr + j) *= scale;
+          *(data_ptr + j + h) *= scale;
+        }
+      }
+      h <<= 1;
+    }
+  }
+}
+
+// m component
+template <typename T>
+void hadamard_m(T* out, int n, int m, float scale, size_t size) {
+  auto h_matrices = hadamard_matrices();
+  auto& matrix = h_matrices[m];
+  auto start = 1;
+  auto end = matrix.find('\n', start);
+  std::vector<bool> hmat_vec;
+  while (end != std::string_view::npos) {
+    auto row = matrix.substr(start, end - start);
+    for (int i = 0; i < row.length(); i++) {
+      hmat_vec.push_back(row[i] == '+');
+    }
+    start = end + 1;
+    end = matrix.find('\n', start);
+  }
+
+  for (int b = 0; b < size / m / n; b++) {
+    size_t loc = b * n * m;
+    T* data_ptr = out + loc;
+    for (int i = 0; i < n; i++) {
+      std::vector<float> out(m);
+      for (int j = 0; j < m; j++) {
+        for (int k = 0; k < m; k++) {
+          float x = *(data_ptr + i + k * n);
+          if (hmat_vec[k + j * m]) {
+            out[j] += x;
+          } else {
+            out[j] -= x;
+          }
+        }
+      }
+      for (int j = 0; j < m; j++) {
+        *(data_ptr + i + j * n) = out[j] * scale;
+      }
+    }
+  }
+}
+
+template <typename T>
+void hadamard(array& out, int n, int m, float scale, Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  auto out_ptr = out.data<T>();
+  encoder.dispatch([out_ptr, size = out.size(), n, m, scale]() {
+    float n_scale = m > 1 ? 1.0 : scale;
+    hadamard_n<T>(out_ptr, n, m, n_scale, size);
+    if (m > 1) {
+      hadamard_m<T>(out_ptr, n, m, scale, size);
+    }
+  });
+}
+
+void Hadamard::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+
+  // Copy input to output
+  if (in.flags().row_contiguous && in.is_donatable()) {
+    out.copy_shared_buffer(in);
+  } else {
+    copy_cpu(
+        in,
+        out,
+        in.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+        stream());
+  }
+
+  int axis = out.ndim() - 1;
+  auto [n, m] = decompose_hadamard(out.shape(axis));
+
+  switch (in.dtype()) {
+    case float32:
+      return hadamard<float>(out, n, m, scale_, stream());
+    case float16:
+      return hadamard<float16_t>(out, n, m, scale_, stream());
+    case bfloat16:
+      return hadamard<bfloat16_t>(out, n, m, scale_, stream());
+    default:
+      throw std::invalid_argument("[hadamard] Unsupported type.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/indexing.cpp b/macos/third_party/mlx/mlx/backend/cpu/indexing.cpp
new file mode 100644
index 0000000..a1b412c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/indexing.cpp
@@ -0,0 +1,854 @@
+// Copyright © 2023 Apple Inc.
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+
+#include "mlx/allocator.h"
+#include "mlx/primitives.h"
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+
+namespace mlx::core {
+
+template <typename IdxT>
+inline size_t offset_neg_idx(IdxT idx, size_t size) {
+  return (idx < 0) ? idx + size : idx;
+}
+
+template <>
+inline size_t offset_neg_idx(uint32_t idx, size_t) {
+  return idx;
+}
+
+struct None {
+  template <typename T>
+  void operator()(T x, T* y) {
+    (*y) = x;
+  }
+};
+struct Sum {
+  template <typename T>
+  void operator()(T x, T* y) {
+    (*y) += x;
+  }
+};
+
+struct Prod {
+  template <typename T>
+  void operator()(T x, T* y) {
+    (*y) *= x;
+  }
+};
+
+struct Max {
+  template <typename T>
+  void operator()(T x, T* y) {
+    (*y) = (*y > x) ? *y : x;
+  }
+};
+
+struct Min {
+  template <typename T>
+  void operator()(T x, T* y) {
+    (*y) = (*y < x) ? *y : x;
+  }
+};
+
+template <typename T, typename IdxT>
+void gather(
+    const array& src,
+    const std::vector<array>& inds,
+    array& out,
+    const std::vector<int>& axes,
+    const Shape& slice_sizes) {
+  // If the array is row contiguous then we can do a contiguous copy given
+  // two conditions on the slice size:
+  // - Any number of leading ones in the slice sizes are allowed
+  // - All other slice sizes match the corresponding dimension except the
+  //   first non-singleton slice size
+  // If the array is col contiguous then the reverse is the case:
+  // - Any number of trailing ones in the slice sizes are allowed
+  // - All other slice sizes match the corresponding dimension except the
+  //   first non-singleton slice size from the end
+
+  bool can_copy = false;
+  if (src.flags().row_contiguous) {
+    can_copy = true;
+
+    // Ignore leading 1s
+    int i = 0;
+    for (; i < slice_sizes.size() && slice_sizes[i] == 1; ++i)
+      ;
+
+    // Check the remaining
+    i++;
+    for (; i < src.ndim() && can_copy; ++i) {
+      can_copy = (src.shape(i) == slice_sizes[i]);
+    }
+  } else if (src.flags().col_contiguous) {
+    can_copy = true;
+
+    // Ignore trailing 1s
+    int i = slice_sizes.size() - 1;
+    for (; i >= 0 && slice_sizes[i] == 1; --i)
+      ;
+
+    // Skip the next slice size and check the remaining
+    i--;
+    for (; i >= 0 && can_copy; --i) {
+      can_copy = (src.shape(i) == slice_sizes[i]);
+    }
+  }
+  size_t slice_size = 1;
+  for (auto s : slice_sizes) {
+    slice_size *= s;
+  }
+  size_t ind_size = slice_size == 0 ? 0 : out.size() / slice_size;
+  const T* src_ptr = src.data<T>();
+  T* dst_ptr = out.data<T>();
+
+  std::vector<ContiguousIterator> its(inds.begin(), inds.end());
+  ContiguousIterator src_it;
+  if (!can_copy && src.ndim() > 0) {
+    src_it = ContiguousIterator(slice_sizes, src.strides(), src.ndim());
+  }
+
+  size_t out_idx = 0;
+  for (int idx = 0; idx < ind_size; idx++) {
+    size_t src_idx = 0;
+    for (int ii = 0; ii < inds.size(); ++ii) {
+      auto ax = axes[ii];
+      auto idx_loc = its[ii].loc;
+      its[ii].step();
+      auto idx_val =
+          offset_neg_idx(inds[ii].data<IdxT>()[idx_loc], src.shape(ax));
+      src_idx += (idx_val * src.strides()[ax]);
+    }
+
+    if (slice_size == 1) {
+      dst_ptr[out_idx++] = src_ptr[src_idx];
+    } else if (can_copy) {
+      std::copy(
+          src_ptr + src_idx, src_ptr + src_idx + slice_size, dst_ptr + out_idx);
+      out_idx += slice_size;
+    } else {
+      for (int jj = 0; jj < slice_size; jj++) {
+        dst_ptr[out_idx++] = src_ptr[src_idx + src_it.loc];
+        src_it.step();
+      }
+      src_it.reset();
+    }
+  }
+}
+
+template <typename IdxT>
+void dispatch_gather(
+    const array& src,
+    const std::vector<array>& inds,
+    array& out,
+    const std::vector<int>& axes,
+    const Shape& size) {
+  switch (out.dtype()) {
+    case bool_:
+      gather<bool, IdxT>(src, inds, out, axes, size);
+      break;
+    case uint8:
+      gather<uint8_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case uint16:
+      gather<uint16_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case uint32:
+      gather<uint32_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case uint64:
+      gather<uint64_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case int8:
+      gather<int8_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case int16:
+      gather<int16_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case int32:
+      gather<int32_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case int64:
+      gather<int64_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case float16:
+      gather<float16_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case float32:
+      gather<float, IdxT>(src, inds, out, axes, size);
+      break;
+    case float64:
+      gather<double, IdxT>(src, inds, out, axes, size);
+      break;
+    case bfloat16:
+      gather<bfloat16_t, IdxT>(src, inds, out, axes, size);
+      break;
+    case complex64:
+      gather<complex64_t, IdxT>(src, inds, out, axes, size);
+      break;
+  }
+}
+
+void Gather::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& src = inputs[0];
+  std::vector<array> inds;
+  for (auto it = inputs.begin() + 1; it < inputs.end(); ++it) {
+    inds.push_back(array::unsafe_weak_copy(*it));
+  }
+  auto& encoder = cpu::get_command_encoder(stream());
+  for (auto& in : inputs) {
+    encoder.set_input_array(in);
+  }
+  encoder.set_output_array(out);
+  encoder.dispatch([axes_ = axes_,
+                    slice_sizes_ = slice_sizes_,
+                    src = array::unsafe_weak_copy(src),
+                    inds = std::move(inds),
+                    out = array::unsafe_weak_copy(out)]() mutable {
+    if (inds.empty()) {
+      dispatch_gather<uint8_t>(src, inds, out, axes_, slice_sizes_);
+      return;
+    }
+
+    switch (inds[0].dtype()) {
+      case uint8:
+        dispatch_gather<uint8_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      case uint16:
+        dispatch_gather<uint16_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      case uint32:
+        dispatch_gather<uint32_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      case uint64:
+        dispatch_gather<uint64_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      case int8:
+        dispatch_gather<int8_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      case int16:
+        dispatch_gather<int16_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      case int32:
+        dispatch_gather<int32_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      case int64:
+        dispatch_gather<int64_t>(src, inds, out, axes_, slice_sizes_);
+        break;
+      default:
+        throw std::runtime_error(
+            "[Gather::eval_cpu] Cannot gather with indices type.");
+        break;
+    }
+  });
+}
+template <typename T, typename IdxT>
+void gather_axis(
+    const array& src,
+    const array& ind,
+    array& out,
+    const int axis) {
+  auto shape = remove_index(ind.shape(), axis);
+  ContiguousIterator ind_it(
+      shape, remove_index(ind.strides(), axis), src.ndim() - 1);
+  ContiguousIterator src_it(
+      shape, remove_index(src.strides(), axis), src.ndim() - 1);
+
+  auto ind_ptr = ind.data<IdxT>();
+  auto src_ptr = src.data<T>();
+  auto dst_ptr = out.data<T>();
+  auto ind_ax_stride = ind.strides(axis);
+  auto src_ax_stride = src.strides(axis);
+  auto dst_ax_stride = out.strides(axis);
+  auto ind_ax_size = ind.shape(axis);
+  auto src_ax_size = src.shape(axis);
+
+  size_t size_pre = 1;
+  size_t size_post = 1;
+  for (int i = 0; i < axis; ++i) {
+    size_pre *= ind.shape(i);
+  }
+  for (int i = axis + 1; i < ind.ndim(); ++i) {
+    size_post *= ind.shape(i);
+  }
+
+  size_t stride_pre = size_post * ind_ax_size;
+  for (size_t i = 0; i < size_pre; i++) {
+    for (size_t k = 0; k < size_post; k++) {
+      for (int j = 0; j < ind_ax_size; ++j) {
+        auto ind_val = offset_neg_idx(
+            ind_ptr[ind_it.loc + j * ind_ax_stride], src_ax_size);
+        dst_ptr[k + j * dst_ax_stride] =
+            src_ptr[src_it.loc + ind_val * src_ax_stride];
+      }
+      ind_it.step();
+      src_it.step();
+    }
+    dst_ptr += stride_pre;
+  }
+}
+
+template <typename IdxT>
+void dispatch_gather_axis(
+    const array& src,
+    const array& inds,
+    array& out,
+    const int axis) {
+  switch (out.dtype()) {
+    case bool_:
+      gather_axis<bool, IdxT>(src, inds, out, axis);
+      break;
+    case uint8:
+      gather_axis<uint8_t, IdxT>(src, inds, out, axis);
+      break;
+    case uint16:
+      gather_axis<uint16_t, IdxT>(src, inds, out, axis);
+      break;
+    case uint32:
+      gather_axis<uint32_t, IdxT>(src, inds, out, axis);
+      break;
+    case uint64:
+      gather_axis<uint64_t, IdxT>(src, inds, out, axis);
+      break;
+    case int8:
+      gather_axis<int8_t, IdxT>(src, inds, out, axis);
+      break;
+    case int16:
+      gather_axis<int16_t, IdxT>(src, inds, out, axis);
+      break;
+    case int32:
+      gather_axis<int32_t, IdxT>(src, inds, out, axis);
+      break;
+    case int64:
+      gather_axis<int64_t, IdxT>(src, inds, out, axis);
+      break;
+    case float16:
+      gather_axis<float16_t, IdxT>(src, inds, out, axis);
+      break;
+    case float32:
+      gather_axis<float, IdxT>(src, inds, out, axis);
+      break;
+    case float64:
+      gather_axis<double, IdxT>(src, inds, out, axis);
+      break;
+    case bfloat16:
+      gather_axis<bfloat16_t, IdxT>(src, inds, out, axis);
+      break;
+    case complex64:
+      gather_axis<complex64_t, IdxT>(src, inds, out, axis);
+      break;
+  }
+}
+
+void GatherAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& src = inputs[0];
+  auto& inds = inputs[1];
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(src);
+  encoder.set_input_array(inds);
+  encoder.set_output_array(out);
+  encoder.dispatch([axis_ = axis_,
+                    src = array::unsafe_weak_copy(src),
+                    inds = array::unsafe_weak_copy(inds),
+                    out = array::unsafe_weak_copy(out)]() mutable {
+    switch (inds.dtype()) {
+      case uint8:
+        dispatch_gather_axis<uint8_t>(src, inds, out, axis_);
+        break;
+      case uint16:
+        dispatch_gather_axis<uint16_t>(src, inds, out, axis_);
+        break;
+      case uint32:
+        dispatch_gather_axis<uint32_t>(src, inds, out, axis_);
+        break;
+      case uint64:
+        dispatch_gather_axis<uint64_t>(src, inds, out, axis_);
+        break;
+      case int8:
+        dispatch_gather_axis<int8_t>(src, inds, out, axis_);
+        break;
+      case int16:
+        dispatch_gather_axis<int16_t>(src, inds, out, axis_);
+        break;
+      case int32:
+        dispatch_gather_axis<int32_t>(src, inds, out, axis_);
+        break;
+      case int64:
+        dispatch_gather_axis<int64_t>(src, inds, out, axis_);
+        break;
+      default:
+        throw std::runtime_error(
+            "[GatherAxis::eval_cpu] Cannot gather with indices type.");
+        break;
+    }
+  });
+}
+
+template <typename InT, typename IdxT, typename OpT>
+void scatter(
+    const array& updates,
+    array& out,
+    const std::vector<array>& inds,
+    const std::vector<int>& axes) {
+  int nind = inds.size();
+  auto inds_ndim = updates.ndim() - out.ndim();
+  size_t n_updates = nind ? inds[0].size() : 1;
+
+  Shape update_shape(
+      updates.shape().begin() + inds_ndim, updates.shape().end());
+  size_t update_size = 1;
+  for (auto us : update_shape) {
+    update_size *= us;
+  }
+
+  std::vector<ContiguousIterator> its(inds.begin(), inds.end());
+  ContiguousIterator update_it(updates);
+  ContiguousIterator out_it(update_shape, out.strides(), out.ndim());
+
+  auto out_ptr = out.data<InT>();
+  auto upd_ptr = updates.data<InT>();
+  for (int i = 0; i < n_updates; ++i) {
+    size_t out_offset = 0;
+    for (int j = 0; j < inds.size(); ++j) {
+      auto ax = axes[j];
+      auto idx_loc = its[j].loc;
+      its[j].step();
+      auto idx_val =
+          offset_neg_idx(inds[j].data<IdxT>()[idx_loc], out.shape(ax));
+      out_offset += (idx_val * out.strides()[ax]);
+    }
+    update_it.seek(i * update_size);
+    for (int j = 0; j < update_size; ++j) {
+      OpT{}(upd_ptr[update_it.loc], out_ptr + out_offset + out_it.loc);
+      update_it.step();
+      out_it.step();
+    }
+    out_it.reset();
+    update_it.reset();
+  }
+}
+
+template <typename InT, typename IdxT>
+void dispatch_scatter_inds(
+    array& out,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    Scatter::ReduceType rtype) {
+  switch (rtype) {
+    case Scatter::None:
+      scatter<InT, IdxT, None>(updates, out, indices, axes);
+      break;
+    case Scatter::Sum:
+      scatter<InT, IdxT, Sum>(updates, out, indices, axes);
+      break;
+    case Scatter::Prod:
+      scatter<InT, IdxT, Prod>(updates, out, indices, axes);
+      break;
+    case Scatter::Max:
+      scatter<InT, IdxT, Max>(updates, out, indices, axes);
+      break;
+    case Scatter::Min:
+      scatter<InT, IdxT, Min>(updates, out, indices, axes);
+      break;
+  }
+}
+
+template <typename InT>
+void dispatch_scatter(
+    array& out,
+    const std::vector<array>& inds,
+    const array& updates,
+    const std::vector<int>& axes,
+    Scatter::ReduceType rtype) {
+  if (inds.empty()) {
+    dispatch_scatter_inds<InT, uint8_t>(out, inds, updates, axes, rtype);
+    return;
+  }
+
+  switch (inds[0].dtype()) {
+    case uint8:
+      dispatch_scatter_inds<InT, uint8_t>(out, inds, updates, axes, rtype);
+      break;
+    case uint16:
+      dispatch_scatter_inds<InT, uint16_t>(out, inds, updates, axes, rtype);
+      break;
+    case uint32:
+      dispatch_scatter_inds<InT, uint32_t>(out, inds, updates, axes, rtype);
+      break;
+    case uint64:
+      dispatch_scatter_inds<InT, uint64_t>(out, inds, updates, axes, rtype);
+      break;
+    case int8:
+      dispatch_scatter_inds<InT, int8_t>(out, inds, updates, axes, rtype);
+      break;
+    case int16:
+      dispatch_scatter_inds<InT, int16_t>(out, inds, updates, axes, rtype);
+      break;
+    case int32:
+      dispatch_scatter_inds<InT, int32_t>(out, inds, updates, axes, rtype);
+      break;
+    case int64:
+      dispatch_scatter_inds<InT, int64_t>(out, inds, updates, axes, rtype);
+      break;
+    default:
+      throw std::runtime_error(
+          "[Scatter::eval_cpu] Cannot scatter with indices type.");
+  }
+}
+
+void Scatter::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() >= 2);
+
+  auto& src = inputs[0];
+  auto& updates = inputs.back();
+
+  // Copy src into out (copy allocates memory for out)
+  auto ctype =
+      src.flags().row_contiguous ? CopyType::Vector : CopyType::General;
+  copy_cpu(src, out, ctype, stream());
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  std::vector<array> inds;
+  for (auto it = inputs.begin() + 1; it < inputs.end() - 1; ++it) {
+    encoder.set_input_array(*it);
+    inds.push_back(array::unsafe_weak_copy(*it));
+  }
+  encoder.set_input_array(updates);
+  encoder.set_output_array(out);
+  encoder.dispatch([axes_ = axes_,
+                    reduce_type_ = reduce_type_,
+                    updates = array::unsafe_weak_copy(updates),
+                    inds = std::move(inds),
+                    out = array::unsafe_weak_copy(out)]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        dispatch_scatter<bool>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case uint8:
+        dispatch_scatter<uint8_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case uint16:
+        dispatch_scatter<uint16_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case uint32:
+        dispatch_scatter<uint32_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case uint64:
+        dispatch_scatter<uint64_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case int8:
+        dispatch_scatter<int8_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case int16:
+        dispatch_scatter<int16_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case int32:
+        dispatch_scatter<int32_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case int64:
+        dispatch_scatter<int64_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case float16:
+        dispatch_scatter<float16_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case float32:
+        dispatch_scatter<float>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case float64:
+        dispatch_scatter<double>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case bfloat16:
+        dispatch_scatter<bfloat16_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+      case complex64:
+        dispatch_scatter<complex64_t>(out, inds, updates, axes_, reduce_type_);
+        break;
+    }
+  });
+}
+
+template <typename T, typename IdxT, typename OpT>
+void scatter_axis(array& out, const array idx, const array& upd, int axis) {
+  auto shape = remove_index(idx.shape(), axis);
+  ContiguousIterator idx_it(
+      shape, remove_index(idx.strides(), axis), upd.ndim() - 1);
+  ContiguousIterator upd_it(
+      shape, remove_index(upd.strides(), axis), upd.ndim() - 1);
+
+  auto idx_ptr = idx.data<IdxT>();
+  auto upd_ptr = upd.data<T>();
+  auto dst_ptr = out.data<T>();
+  auto idx_ax_stride = idx.strides(axis);
+  auto upd_ax_stride = upd.strides(axis);
+  auto dst_ax_stride = out.strides(axis);
+  auto idx_ax_size = idx.shape(axis);
+  auto dst_ax_size = out.shape(axis);
+
+  size_t size_pre = 1;
+  size_t size_post = 1;
+  for (int i = 0; i < axis; ++i) {
+    size_pre *= idx.shape(i);
+  }
+  for (int i = axis + 1; i < idx.ndim(); ++i) {
+    size_post *= idx.shape(i);
+  }
+  size_t stride_pre = size_post * dst_ax_size;
+  for (size_t i = 0; i < size_pre; i++) {
+    for (size_t k = 0; k < size_post; k++) {
+      for (int j = 0; j < idx_ax_size; ++j) {
+        auto ind_val = offset_neg_idx(
+            idx_ptr[idx_it.loc + j * idx_ax_stride], dst_ax_size);
+        OpT{}(
+            upd_ptr[upd_it.loc + j * upd_ax_stride],
+            dst_ptr + k + ind_val * dst_ax_stride);
+      }
+      idx_it.step();
+      upd_it.step();
+    }
+    dst_ptr += stride_pre;
+  }
+}
+
+template <typename InT, typename IdxT>
+void dispatch_scatter_axis_op(
+    array& out,
+    const array& idx,
+    const array& updates,
+    int axis,
+    ScatterAxis::ReduceType rtype) {
+  switch (rtype) {
+    case ScatterAxis::None:
+      scatter_axis<InT, IdxT, None>(out, idx, updates, axis);
+      break;
+    case ScatterAxis::Sum:
+      scatter_axis<InT, IdxT, Sum>(out, idx, updates, axis);
+      break;
+  }
+}
+
+template <typename InT>
+void dispatch_scatter_axis(
+    array& out,
+    const array& idx,
+    const array& updates,
+    int axis,
+    ScatterAxis::ReduceType rtype) {
+  switch (idx.dtype()) {
+    case uint8:
+      dispatch_scatter_axis_op<InT, uint8_t>(out, idx, updates, axis, rtype);
+      break;
+    case uint16:
+      dispatch_scatter_axis_op<InT, uint16_t>(out, idx, updates, axis, rtype);
+      break;
+    case uint32:
+      dispatch_scatter_axis_op<InT, uint32_t>(out, idx, updates, axis, rtype);
+      break;
+    case uint64:
+      dispatch_scatter_axis_op<InT, uint64_t>(out, idx, updates, axis, rtype);
+      break;
+    case int8:
+      dispatch_scatter_axis_op<InT, int8_t>(out, idx, updates, axis, rtype);
+      break;
+    case int16:
+      dispatch_scatter_axis_op<InT, int16_t>(out, idx, updates, axis, rtype);
+      break;
+    case int32:
+      dispatch_scatter_axis_op<InT, int32_t>(out, idx, updates, axis, rtype);
+      break;
+    case int64:
+      dispatch_scatter_axis_op<InT, int64_t>(out, idx, updates, axis, rtype);
+      break;
+    default:
+      throw std::runtime_error(
+          "[ScatterAxis::eval_cpu] Cannot scatter with indices type.");
+  }
+}
+
+void ScatterAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() >= 2);
+
+  auto& src = inputs[0];
+  auto& idx = inputs[1];
+  auto& updates = inputs[2];
+
+  // Copy src into out (copy allocates memory for out)
+  auto ctype =
+      src.flags().row_contiguous ? CopyType::Vector : CopyType::General;
+  copy_cpu(src, out, ctype, stream());
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(idx);
+  encoder.set_input_array(updates);
+  encoder.set_output_array(out);
+  encoder.dispatch([axis_ = axis_,
+                    reduce_type_ = reduce_type_,
+                    idx = array::unsafe_weak_copy(idx),
+                    updates = array::unsafe_weak_copy(updates),
+                    out = array::unsafe_weak_copy(out)]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        dispatch_scatter_axis<bool>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case uint8:
+        dispatch_scatter_axis<uint8_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case uint16:
+        dispatch_scatter_axis<uint16_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case uint32:
+        dispatch_scatter_axis<uint32_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case uint64:
+        dispatch_scatter_axis<uint64_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case int8:
+        dispatch_scatter_axis<int8_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case int16:
+        dispatch_scatter_axis<int16_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case int32:
+        dispatch_scatter_axis<int32_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case int64:
+        dispatch_scatter_axis<int64_t>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case float16:
+        dispatch_scatter_axis<float16_t>(
+            out, idx, updates, axis_, reduce_type_);
+        break;
+      case float32:
+        dispatch_scatter_axis<float>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case float64:
+        dispatch_scatter_axis<double>(out, idx, updates, axis_, reduce_type_);
+        break;
+      case bfloat16:
+        dispatch_scatter_axis<bfloat16_t>(
+            out, idx, updates, axis_, reduce_type_);
+        break;
+      case complex64:
+        dispatch_scatter_axis<complex64_t>(
+            out, idx, updates, axis_, reduce_type_);
+        break;
+    }
+  });
+}
+
+template <typename T>
+void masked_scatter_impl(const array& mask, const array& src, array& out) {
+  ContiguousIterator mask_it(mask);
+  ContiguousIterator src_it(src);
+  ContiguousIterator out_it(out);
+
+  const bool* mask_ptr = mask.data<bool>();
+  const T* src_ptr = src.data<T>();
+  T* dst_ptr = out.data<T>();
+
+  const size_t batch_count = mask.shape(0);
+  const size_t mask_batch_size = mask.size() / batch_count;
+  const size_t src_batch_size = src.size() / batch_count;
+
+  for (size_t b = 0; b < batch_count; ++b) {
+    size_t src_consumed = 0;
+    src_it.seek(b * src_batch_size);
+
+    for (size_t i = 0; i < mask_batch_size; ++i) {
+      if (mask_ptr[mask_it.loc]) {
+        if (src_consumed >= src_batch_size) {
+          throw std::runtime_error(
+              "[MaskedScatter::eval_cpu] Source does not have enough elements for mask.");
+        }
+        dst_ptr[out_it.loc] = src_ptr[src_it.loc];
+        src_it.step();
+        ++src_consumed;
+      }
+      mask_it.step();
+      out_it.step();
+    }
+  }
+}
+
+void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 3);
+
+  auto& dst = inputs[0];
+  auto& mask = inputs[1];
+  auto& src = inputs[2];
+
+  // Copy src into out (copy allocates memory for out)
+  auto ctype =
+      dst.flags().row_contiguous ? CopyType::Vector : CopyType::General;
+  copy_cpu(dst, out, ctype, stream());
+
+  if (mask.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(mask);
+  encoder.set_input_array(src);
+  encoder.set_output_array(out);
+  encoder.dispatch([mask = array::unsafe_weak_copy(mask),
+                    src = array::unsafe_weak_copy(src),
+                    out = array::unsafe_weak_copy(out)]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        masked_scatter_impl<bool>(mask, src, out);
+        break;
+      case uint8:
+        masked_scatter_impl<uint8_t>(mask, src, out);
+        break;
+      case uint16:
+        masked_scatter_impl<uint16_t>(mask, src, out);
+        break;
+      case uint32:
+        masked_scatter_impl<uint32_t>(mask, src, out);
+        break;
+      case uint64:
+        masked_scatter_impl<uint64_t>(mask, src, out);
+        break;
+      case int8:
+        masked_scatter_impl<int8_t>(mask, src, out);
+        break;
+      case int16:
+        masked_scatter_impl<int16_t>(mask, src, out);
+        break;
+      case int32:
+        masked_scatter_impl<int32_t>(mask, src, out);
+        break;
+      case int64:
+        masked_scatter_impl<int64_t>(mask, src, out);
+        break;
+      case float16:
+        masked_scatter_impl<float16_t>(mask, src, out);
+        break;
+      case float32:
+        masked_scatter_impl<float>(mask, src, out);
+        break;
+      case float64:
+        masked_scatter_impl<double>(mask, src, out);
+        break;
+      case bfloat16:
+        masked_scatter_impl<bfloat16_t>(mask, src, out);
+        break;
+      case complex64:
+        masked_scatter_impl<complex64_t>(mask, src, out);
+        break;
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/inverse.cpp b/macos/third_party/mlx/mlx/backend/cpu/inverse.cpp
new file mode 100644
index 0000000..ddc979d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/inverse.cpp
@@ -0,0 +1,160 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T>
+void general_inv(T* inv, int N) {
+  int info;
+  auto ipiv = array::Data{allocator::malloc(sizeof(int) * N)};
+  // Compute LU factorization.
+  getrf<T>(
+      /* m = */ &N,
+      /* n = */ &N,
+      /* a = */ inv,
+      /* lda = */ &N,
+      /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] LU factorization failed with error code "
+       << info;
+    throw std::runtime_error(ss.str());
+  }
+
+  static const int lwork_query = -1;
+  T workspace_size = 0;
+
+  // Compute workspace size.
+  getri<T>(
+      /* m = */ &N,
+      /* a = */ nullptr,
+      /* lda = */ &N,
+      /* ipiv = */ nullptr,
+      /* work = */ &workspace_size,
+      /* lwork = */ &lwork_query,
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] LU workspace calculation failed with error code "
+       << info;
+    throw std::runtime_error(ss.str());
+  }
+
+  const int lwork = workspace_size;
+  auto scratch = array::Data{allocator::malloc(sizeof(T) * lwork)};
+
+  // Compute inverse.
+  getri<T>(
+      /* m = */ &N,
+      /* a = */ inv,
+      /* lda = */ &N,
+      /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+      /* work = */ static_cast<T*>(scratch.buffer.raw_ptr()),
+      /* lwork = */ &lwork,
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] inversion failed with error code " << info;
+    throw std::runtime_error(ss.str());
+  }
+}
+
+template <typename T>
+void tri_inv(T* inv, int N, bool upper) {
+  const char uplo = upper ? 'L' : 'U';
+  const char diag = 'N';
+  int info;
+  trtri<T>(
+      /* uplo = */ &uplo,
+      /* diag = */ &diag,
+      /* N = */ &N,
+      /* a = */ inv,
+      /* lda = */ &N,
+      /* info = */ &info);
+
+  // zero out the other triangle
+  if (upper) {
+    for (int i = 0; i < N; i++) {
+      std::fill(inv, inv + i, 0.0f);
+      inv += N;
+    }
+  } else {
+    for (int i = 0; i < N; i++) {
+      std::fill(inv + i + 1, inv + N, 0.0f);
+      inv += N;
+    }
+  }
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "[Inverse::eval_cpu] triangular inversion failed with error code "
+       << info;
+    throw std::runtime_error(ss.str());
+  }
+}
+
+template <typename T>
+void inverse_impl(
+    const array& a,
+    array& inv,
+    bool tri,
+    bool upper,
+    Stream stream) {
+  // Lapack uses the column-major convention. We take advantage of the following
+  // identity to avoid transposing (see
+  // https://math.stackexchange.com/a/340234):
+  //   (A⁻¹)ᵀ = (Aᵀ)⁻¹
+
+  // The inverse is computed in place, so just copy the input to the output.
+  copy_cpu(
+      a,
+      inv,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
+
+  const int N = a.shape(-1);
+  const size_t num_matrices = a.size() / (N * N);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(inv);
+
+  auto inv_ptr = inv.data<T>();
+  if (tri) {
+    encoder.dispatch([inv_ptr, N, num_matrices, upper]() {
+      for (int i = 0; i < num_matrices; i++) {
+        tri_inv<T>(inv_ptr + N * N * i, N, upper);
+      }
+    });
+  } else {
+    encoder.dispatch([inv_ptr, N, num_matrices]() {
+      for (int i = 0; i < num_matrices; i++) {
+        general_inv<T>(inv_ptr + N * N * i, N);
+      }
+    });
+  }
+}
+
+void Inverse::eval_cpu(const std::vector<array>& inputs, array& output) {
+  switch (inputs[0].dtype()) {
+    case float32:
+      inverse_impl<float>(inputs[0], output, tri_, upper_, stream());
+      break;
+    case float64:
+      inverse_impl<double>(inputs[0], output, tri_, upper_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[Inverse::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/jit_compiler.cpp b/macos/third_party/mlx/mlx/backend/cpu/jit_compiler.cpp
new file mode 100644
index 0000000..267e8fd
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/jit_compiler.cpp
@@ -0,0 +1,166 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/cpu/jit_compiler.h"
+
+#include <algorithm>
+#include <sstream>
+#include <vector>
+
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+#ifdef _MSC_VER
+
+namespace {
+
+// Split string into array.
+std::vector<std::string> str_split(const std::string& str, char delimiter) {
+  std::vector<std::string> tokens;
+  std::string token;
+  std::istringstream tokenStream(str);
+  while (std::getline(tokenStream, token, delimiter)) {
+    tokens.push_back(token);
+  }
+  return tokens;
+}
+
+// Get path information about MSVC.
+struct VisualStudioInfo {
+  VisualStudioInfo() {
+#ifdef _M_ARM64
+    arch = "arm64";
+#else
+    arch = "x64";
+#endif
+    // Get path of Visual Studio.
+    // Use -latest to get only the most recent installation when multiple
+    // versions are installed, avoiding path concatenation issues.
+    std::string vs_path = JitCompiler::exec(
+        fmt::format(
+            "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
+            " -latest -property installationPath",
+            std::getenv("ProgramFiles(x86)")));
+    if (vs_path.empty()) {
+      throw std::runtime_error("Can not find Visual Studio.");
+    }
+    // Trim any trailing whitespace/newlines from the path
+    vs_path.erase(
+        std::find_if(
+            vs_path.rbegin(),
+            vs_path.rend(),
+            [](unsigned char ch) { return !std::isspace(ch); })
+            .base(),
+        vs_path.end());
+    // Read the envs from vcvarsall.
+    std::string envs = JitCompiler::exec(
+        fmt::format(
+            "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
+            vs_path,
+            arch));
+    for (const std::string& line : str_split(envs, '\n')) {
+      // Each line is in the format "ENV_NAME=values".
+      auto pos = line.find_first_of('=');
+      if (pos == std::string::npos || pos == 0 || pos == line.size() - 1)
+        continue;
+      std::string name = line.substr(0, pos);
+      std::string value = line.substr(pos + 1);
+      if (name == "LIB") {
+        libpaths = str_split(value, ';');
+      } else if (name == "VCToolsInstallDir" || name == "VCTOOLSINSTALLDIR") {
+        cl_exe = fmt::format("{0}\\bin\\Host{1}\\{1}\\cl.exe", value, arch);
+      }
+    }
+  }
+  std::string arch;
+  std::string cl_exe;
+  std::vector<std::string> libpaths;
+};
+
+const VisualStudioInfo& GetVisualStudioInfo() {
+  static VisualStudioInfo info;
+  return info;
+}
+
+} // namespace
+
+#endif // _MSC_VER
+
+std::string JitCompiler::build_command(
+    const std::filesystem::path& dir,
+    const std::string& source_file_name,
+    const std::string& shared_lib_name) {
+#ifdef _MSC_VER
+  const VisualStudioInfo& info = GetVisualStudioInfo();
+  std::string libpaths;
+  for (const std::string& lib : info.libpaths) {
+    libpaths += fmt::format(" /libpath:\"{0}\"", lib);
+  }
+  return fmt::format(
+      "\""
+      "cd /D \"{0}\" && "
+      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" "
+      "/link /out:\"{3}\" {4} 2>&1"
+      "\"",
+      dir.string(),
+      info.cl_exe,
+      source_file_name,
+      shared_lib_name,
+      libpaths);
+#else
+  return fmt::format(
+      "g++ -std=c++17 -O3 -Wall -fPIC -shared \"{0}\" -o \"{1}\" 2>&1",
+      (dir / source_file_name).string(),
+      (dir / shared_lib_name).string());
+#endif
+}
+
+std::string JitCompiler::exec(const std::string& cmd) {
+#ifdef _MSC_VER
+  FILE* pipe = _popen(cmd.c_str(), "r");
+#else
+  FILE* pipe = popen(cmd.c_str(), "r");
+#endif
+  if (!pipe) {
+    throw std::runtime_error("popen() failed.");
+  }
+  char buffer[128];
+  std::string ret;
+  while (fgets(buffer, sizeof(buffer), pipe)) {
+    ret += buffer;
+  }
+  // Trim trailing spaces.
+  ret.erase(
+      std::find_if(
+          ret.rbegin(),
+          ret.rend(),
+          [](unsigned char ch) { return !std::isspace(ch); })
+          .base(),
+      ret.end());
+
+#ifdef _MSC_VER
+  int status = _pclose(pipe);
+#else
+  int status = pclose(pipe);
+#endif
+  if (status == -1) {
+    throw std::runtime_error("pclose() failed.");
+  }
+#if defined(_WIN32) || defined(__FreeBSD__)
+  int code = status;
+#else
+  int code = WEXITSTATUS(status);
+#endif
+  if (code != 0) {
+    throw std::runtime_error(
+        fmt::format(
+            "Failed to execute command with return code {0}: \"{1}\", "
+            "the output is: {2}",
+            code,
+            cmd,
+            ret));
+  }
+  return ret;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/jit_compiler.h b/macos/third_party/mlx/mlx/backend/cpu/jit_compiler.h
new file mode 100644
index 0000000..3a9e988
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/jit_compiler.h
@@ -0,0 +1,20 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+
+#include <filesystem>
+
+namespace mlx::core {
+
+class JitCompiler {
+ public:
+  // Build a shell command that compiles a source code file to a shared library.
+  static std::string build_command(
+      const std::filesystem::path& dir,
+      const std::string& source_file_name,
+      const std::string& shared_lib_name);
+
+  // Run a command and get its output.
+  static std::string exec(const std::string& cmd);
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/lapack.h b/macos/third_party/mlx/mlx/backend/cpu/lapack.h
new file mode 100644
index 0000000..1c3ba1a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/lapack.h
@@ -0,0 +1,80 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <complex>
+#define LAPACK_COMPLEX_CUSTOM
+#define lapack_complex_float std::complex<float>
+#define lapack_complex_double std::complex<double>
+#define lapack_complex_float_real(z) ((z).real())
+#define lapack_complex_float_imag(z) ((z).imag())
+#define lapack_complex_double_real(z) ((z).real())
+#define lapack_complex_double_imag(z) ((z).imag())
+
+#ifdef MLX_USE_ACCELERATE
+#include <Accelerate/Accelerate.h>
+#else
+#include <cblas.h>
+#include <lapack.h>
+#endif
+
+#if defined(LAPACK_GLOBAL) || defined(LAPACK_NAME)
+
+// This is to work around a change in the function signatures of lapack >= 3.9.1
+// where functions taking char* also include a strlen argument, see a similar
+// change in OpenCV:
+// https://github.com/opencv/opencv/blob/1eb061f89de0fb85c4c75a2deeb0f61a961a63ad/cmake/OpenCVFindLAPACK.cmake#L57
+#define MLX_LAPACK_FUNC(f) LAPACK_##f
+
+#else
+
+#define MLX_LAPACK_FUNC(f) f##_
+
+#endif
+
+#define INSTANTIATE_LAPACK_REAL(FUNC)                        \
+  template <typename T, typename... Args>                    \
+  void FUNC(Args... args) {                                  \
+    if constexpr (std::is_same_v<T, float>) {                \
+      MLX_LAPACK_FUNC(s##FUNC)(std::forward<Args>(args)...); \
+    } else if constexpr (std::is_same_v<T, double>) {        \
+      MLX_LAPACK_FUNC(d##FUNC)(std::forward<Args>(args)...); \
+    }                                                        \
+  }
+
+INSTANTIATE_LAPACK_REAL(geqrf)
+INSTANTIATE_LAPACK_REAL(orgqr)
+INSTANTIATE_LAPACK_REAL(syevd)
+INSTANTIATE_LAPACK_REAL(potrf)
+INSTANTIATE_LAPACK_REAL(getrf)
+INSTANTIATE_LAPACK_REAL(getri)
+INSTANTIATE_LAPACK_REAL(trtri)
+
+#define INSTANTIATE_LAPACK_COMPLEX(FUNC)                            \
+  template <typename T, typename... Args>                           \
+  void FUNC(Args... args) {                                         \
+    if constexpr (std::is_same_v<T, std::complex<float>>) {         \
+      MLX_LAPACK_FUNC(c##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<double>>) { \
+      MLX_LAPACK_FUNC(z##FUNC)(std::forward<Args>(args)...);        \
+    }                                                               \
+  }
+
+INSTANTIATE_LAPACK_COMPLEX(heevd)
+
+#define INSTANTIATE_LAPACK_ALL(FUNC)                                \
+  template <typename T, typename... Args>                           \
+  void FUNC(Args... args) {                                         \
+    if constexpr (std::is_same_v<T, float>) {                       \
+      MLX_LAPACK_FUNC(s##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, double>) {               \
+      MLX_LAPACK_FUNC(d##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<float>>) {  \
+      MLX_LAPACK_FUNC(c##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<double>>) { \
+      MLX_LAPACK_FUNC(z##FUNC)(std::forward<Args>(args)...);        \
+    }                                                               \
+  }
+
+INSTANTIATE_LAPACK_ALL(geev)
+INSTANTIATE_LAPACK_ALL(gesdd)
diff --git a/macos/third_party/mlx/mlx/backend/cpu/logsumexp.cpp b/macos/third_party/mlx/mlx/backend/cpu/logsumexp.cpp
new file mode 100644
index 0000000..d061907
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/logsumexp.cpp
@@ -0,0 +1,139 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cassert>
+#include <cmath>
+
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/primitives.h"
+#include "mlx/types/limits.h"
+
+namespace mlx::core {
+
+namespace {
+
+using namespace mlx::core::simd;
+
+template <typename T, typename AccT>
+void logsumexp(const array& in, array& out, Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+
+  const T* in_ptr = in.data<T>();
+  T* out_ptr = out.data<T>();
+
+  int M = in.shape().back();
+  int L = in.data_size() / M;
+
+  encoder.dispatch([in_ptr, out_ptr, M, L]() mutable {
+    constexpr int N = std::min(max_size<AccT>, max_size<T>);
+
+    const T* current_in_ptr;
+
+    for (int i = 0; i < L; i++, in_ptr += M, out_ptr += 1) {
+      // Find the maximum
+      current_in_ptr = in_ptr;
+      Simd<AccT, N> vmaximum(-numeric_limits<AccT>::infinity());
+      size_t s = M;
+      while (s >= N) {
+        Simd<AccT, N> vals = load<T, N>(current_in_ptr);
+        vmaximum = maximum(vals, vmaximum);
+        current_in_ptr += N;
+        s -= N;
+      }
+
+      AccT maximum = max(vmaximum);
+      while (s-- > 0) {
+        maximum = std::max(maximum, static_cast<AccT>(*current_in_ptr));
+        current_in_ptr++;
+      }
+
+      // Compute the normalizer and the exponentials
+      Simd<AccT, N> vnormalizer(0.0);
+      current_in_ptr = in_ptr;
+      s = M;
+      while (s >= N) {
+        Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
+        vexp = exp(vexp - maximum);
+        vnormalizer = vnormalizer + vexp;
+        current_in_ptr += N;
+        s -= N;
+      }
+      AccT normalizer = sum(vnormalizer);
+      while (s-- > 0) {
+        AccT _exp = std::exp(*current_in_ptr - maximum);
+        normalizer += _exp;
+        current_in_ptr++;
+      }
+      // Normalize
+      *out_ptr = std::isinf(maximum)
+          ? static_cast<T>(maximum)
+          : static_cast<T>(std::log(normalizer) + maximum);
+    }
+  });
+}
+
+} // namespace
+
+void LogSumExp::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  // Make sure that the last dimension is contiguous
+  auto s = stream();
+  auto& encoder = cpu::get_command_encoder(s);
+  auto ensure_contiguous = [&s, &encoder](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      return x;
+    } else {
+      array x_copy = contiguous_copy_cpu(x, s);
+      encoder.add_temporary(x_copy);
+      return x_copy;
+    }
+  };
+
+  auto in = ensure_contiguous(inputs[0]);
+  if (in.flags().row_contiguous) {
+    out.set_data(allocator::malloc(out.nbytes()));
+  } else {
+    auto n = in.shape(-1);
+    auto flags = in.flags();
+    auto strides = in.strides();
+    for (auto& s : strides) {
+      s /= n;
+    }
+    bool col_contig = strides[0] == 1;
+    for (int i = 1; col_contig && i < strides.size(); ++i) {
+      col_contig &=
+          (out.shape(i) == 1 || strides[i - 1] == out.shape(i) * strides[i]);
+    }
+    flags.col_contiguous = col_contig;
+    out.set_data(
+        allocator::malloc(in.nbytes() / n),
+        in.data_size() / n,
+        std::move(strides),
+        flags);
+  }
+
+  switch (in.dtype()) {
+    case float32:
+      logsumexp<float, float>(in, out, stream());
+      break;
+    case float16:
+      logsumexp<float16_t, float>(in, out, stream());
+      break;
+    case bfloat16:
+      logsumexp<bfloat16_t, float>(in, out, stream());
+      break;
+    case float64:
+      logsumexp<double, double>(in, out, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[logsumexp] only supports floating point types");
+      break;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/luf.cpp b/macos/third_party/mlx/mlx/backend/cpu/luf.cpp
new file mode 100644
index 0000000..5f1507e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/luf.cpp
@@ -0,0 +1,120 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T>
+void luf_impl(
+    const array& a,
+    array& lu,
+    array& pivots,
+    array& row_indices,
+    Stream stream) {
+  int M = a.shape(-2);
+  int N = a.shape(-1);
+  int K = std::min(M, N);
+
+  // Copy a into lu and make it col contiguous
+  auto ndim = lu.ndim();
+  auto flags = lu.flags();
+  flags.col_contiguous = ndim == 2;
+  flags.row_contiguous = false;
+  flags.contiguous = true;
+  auto strides = lu.strides();
+  strides[ndim - 1] = M;
+  strides[ndim - 2] = 1;
+  lu.set_data(allocator::malloc(lu.nbytes()), lu.nbytes(), strides, flags);
+  copy_cpu_inplace(
+      a,
+      lu,
+      a.shape(),
+      a.strides(),
+      strides,
+      0,
+      0,
+      CopyType::GeneralGeneral,
+      stream);
+
+  auto a_ptr = lu.data<T>();
+  pivots.set_data(allocator::malloc(pivots.nbytes()));
+  row_indices.set_data(allocator::malloc(row_indices.nbytes()));
+  auto pivots_ptr = pivots.data<uint32_t>();
+  auto row_indices_ptr = row_indices.data<uint32_t>();
+  size_t num_matrices = a.size() / (M * N);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(lu);
+  encoder.set_output_array(pivots);
+  encoder.set_output_array(row_indices);
+
+  encoder.dispatch(
+      [a_ptr, pivots_ptr, row_indices_ptr, num_matrices, M, N, K]() mutable {
+        int info;
+        for (size_t i = 0; i < num_matrices; ++i) {
+          // Compute LU factorization of A
+          getrf<T>(
+              /* m */ &M,
+              /* n */ &N,
+              /* a */ a_ptr,
+              /* lda */ &M,
+              /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
+              /* info */ &info);
+
+          if (info != 0) {
+            std::stringstream ss;
+            ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
+               << ((info > 0) ? " because matrix is singular"
+                              : " because argument had an illegal value");
+            throw std::runtime_error(ss.str());
+          }
+
+          // Subtract 1 to get 0-based index
+          int j = 0;
+          for (; j < K; ++j) {
+            pivots_ptr[j]--;
+            row_indices_ptr[j] = j;
+          }
+          for (; j < M; ++j) {
+            row_indices_ptr[j] = j;
+          }
+          for (int j = K - 1; j >= 0; --j) {
+            auto piv = pivots_ptr[j];
+            auto t1 = row_indices_ptr[piv];
+            auto t2 = row_indices_ptr[j];
+            row_indices_ptr[j] = t1;
+            row_indices_ptr[piv] = t2;
+          }
+
+          // Advance pointers to the next matrix
+          a_ptr += M * N;
+          pivots_ptr += K;
+          row_indices_ptr += M;
+        }
+      });
+}
+
+void LUF::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  switch (inputs[0].dtype()) {
+    case float32:
+      luf_impl<float>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      break;
+    case float64:
+      luf_impl<double>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[LUF::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/make_compiled_preamble.ps1 b/macos/third_party/mlx/mlx/backend/cpu/make_compiled_preamble.ps1
new file mode 100644
index 0000000..0cd2d1f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/make_compiled_preamble.ps1
@@ -0,0 +1,38 @@
+# This script generates a C++ function that provides the CPU
+# code for use with kernel generation.
+#
+# Copyright © 2024 Apple Inc.
+
+$OUTPUT_FILE = $args[0]
+$CL = $args[1]
+$SRCDIR = $args[2]
+
+# Get command result as array.
+$CONTENT = & $CL /std:c++17 /EP "/I$SRCDIR" /Tp "$SRCDIR/mlx/backend/cpu/compiled_preamble.h"
+# Remove empty lines.
+# Otherwise there will be too much empty lines making the result unreadable.
+$CONTENT = $CONTENT | Where-Object { $_.Trim() -ne '' }
+# Concatenate to string.
+$CONTENT = $CONTENT -join "`n"
+
+# Append extra content.
+$CONTENT = @"
+$($CONTENT)
+using namespace mlx::core;
+using namespace mlx::core::detail;
+"@
+
+# Convert each char to ASCII code.
+# Unlike the unix script that outputs string literal directly, the output from
+# MSVC is way too large to be embedded as string and compilation will fail, so
+# we store it as static array instead.
+$CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0'
+
+$OUTPUT = @"
+const char* get_kernel_preamble() {
+  static char preamble[] = { $CHARCODES };
+  return preamble;
+}
+"@
+
+Set-Content -Path $OUTPUT_FILE -Value $OUTPUT
diff --git a/macos/third_party/mlx/mlx/backend/cpu/make_compiled_preamble.sh b/macos/third_party/mlx/mlx/backend/cpu/make_compiled_preamble.sh
new file mode 100644
index 0000000..88b4c46
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/make_compiled_preamble.sh
@@ -0,0 +1,41 @@
+#!/bin/bash
+#
+# This script generates a C++ function that provides the CPU
+# code for use with kernel generation.
+#
+# Copyright © 2023-24 Apple Inc.
+
+
+OUTPUT_FILE=$1
+GCC=$2
+SRCDIR=$3
+CLANG=$4
+ARCH=$5
+
+if [ "$CLANG" = "TRUE" ]; then
+  read -r -d '' INCLUDES <<- EOM
+#include <cmath>
+#include <complex>
+#include <cstdint>
+#include <vector>
+#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+#include <arm_fp16.h>
+#endif
+EOM
+CC_FLAGS="-arch ${ARCH} -nobuiltininc -nostdinc"
+else
+CC_FLAGS="-std=c++17"
+fi
+
+CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
+
+cat << EOF > "$OUTPUT_FILE"
+const char* get_kernel_preamble() {
+return R"preamble(
+$INCLUDES
+$CONTENT
+using namespace mlx::core;
+using namespace mlx::core::detail;
+)preamble";
+}
+EOF
diff --git a/macos/third_party/mlx/mlx/backend/cpu/masked_mm.cpp b/macos/third_party/mlx/mlx/backend/cpu/masked_mm.cpp
new file mode 100644
index 0000000..688479c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/masked_mm.cpp
@@ -0,0 +1,608 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cstring>
+
+#include "mlx/array.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/gemm.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T, typename mask_t>
+inline void mask_matrix(
+    T* data,
+    const mask_t* mask,
+    int block_size,
+    const int X,
+    const int Y,
+    const int64_t X_data_str,
+    const int64_t Y_data_str,
+    const int64_t X_mask_str,
+    const int64_t Y_mask_str,
+    const size_t mask_offset) {
+  int tX = (X + block_size - 1) / block_size;
+  int tY = (Y + block_size - 1) / block_size;
+
+  for (int i = 0; i < tX; i++) {
+    for (int j = 0; j < tY; j++) {
+      mask_t do_mask = mask[mask_offset + i * X_mask_str + j * Y_mask_str];
+      if (do_mask != 1) {
+        int loc_x = i * block_size;
+        int loc_y = j * block_size;
+        T* data_block = data + loc_x * X_data_str + loc_y * Y_data_str;
+
+        int size_x = std::min(block_size, X - loc_x);
+        int size_y = std::min(block_size, Y - loc_y);
+        for (int ii = 0; ii < size_x; ii++) {
+          for (int jj = 0; jj < size_y; jj++) {
+            if constexpr (std::is_same_v<mask_t, bool>) {
+              data_block[ii * X_data_str + jj * Y_data_str] = T(0.);
+            } else {
+              data_block[ii * X_data_str + jj * Y_data_str] *= do_mask;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+inline void segmented_mm(
+    const T* a,
+    const T* b,
+    const uint32_t* segments,
+    T* out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    const Shape& a_shape,
+    const Strides& a_strides,
+    const Shape& b_shape,
+    const Strides& b_strides,
+    size_t num_segments,
+    const Shape& segments_shape,
+    const Strides& segments_strides) {
+  int ndim = a_shape.size();
+  Shape a_copy = a_shape;
+  Shape b_copy = b_shape;
+  int32_t M = a_copy[ndim - 2];
+  int32_t N = b_copy[ndim - 1];
+  for (int i = 0; i < num_segments; i++) {
+    uint32_t k_start =
+        segments[elem_to_loc(2 * i, segments_shape, segments_strides)];
+    uint32_t k_end =
+        segments[elem_to_loc(2 * i + 1, segments_shape, segments_strides)];
+    if (k_end <= k_start) {
+      std::fill_n(out + i * M * N, M * N, T(0));
+      continue;
+    }
+    a_copy[ndim - 1] = k_end - k_start;
+    b_copy[ndim - 2] = k_end - k_start;
+    matmul<T>(
+        a + k_start * a_strides[ndim - 1],
+        b + k_start * b_strides[ndim - 2],
+        out + i * M * N,
+        a_transposed,
+        b_transposed,
+        lda,
+        ldb,
+        N,
+        1.0,
+        0.0,
+        1,
+        a_copy,
+        a_strides,
+        b_copy,
+        b_strides);
+  }
+}
+
+} // namespace
+
+void BlockMaskedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
+  if (out.dtype() != float32) {
+    throw std::runtime_error(
+        "[BlockMaskedMM::eval] Currently only supports float32.");
+  }
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+
+  auto check_transpose =
+      [s = stream()](const array& arr, bool do_copy, bool expand_all = false) {
+        auto stx = arr.strides()[arr.ndim() - 2];
+        auto sty = arr.strides()[arr.ndim() - 1];
+        if (!expand_all && stx == arr.shape(-1) && sty == 1) {
+          if (do_copy) {
+            array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
+            copy_cpu(arr, arr_copy, CopyType::Vector, s);
+            return std::make_tuple(false, stx, arr_copy, true);
+          }
+          return std::make_tuple(false, stx, arr, false);
+        } else if (!expand_all && stx == 1 && sty == arr.shape(-2)) {
+          if (do_copy) {
+            array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
+            copy_cpu(arr, arr_copy, CopyType::Vector, s);
+            return std::make_tuple(true, sty, arr_copy, true);
+          }
+          return std::make_tuple(true, sty, arr, false);
+        } else {
+          int64_t stx = arr.shape(-1);
+          array arr_copy = contiguous_copy_cpu(arr, s);
+          return std::make_tuple(false, stx, arr_copy, true);
+        }
+      };
+
+  bool has_op_mask = inputs.size() > 3;
+  bool has_out_mask = inputs.size() == 3 || inputs.size() == 5;
+  auto [a_transposed, lda, a, a_copied] =
+      check_transpose(a_pre, has_op_mask, inputs.back().dtype() != bool_);
+  auto [b_transposed, ldb, b, b_copied] =
+      check_transpose(b_pre, has_op_mask, inputs.back().dtype() != bool_);
+
+  size_t M = a.shape(-2);
+  size_t N = b.shape(-1);
+  size_t K = a.shape(-1);
+
+  if (M == 0 || N == 0) {
+    return;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  if (K == 0) {
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<void>(), nbytes = out.nbytes()]() {
+      std::memset(out_ptr, 0, nbytes);
+    });
+    return;
+  }
+
+  auto mask_array = [](const void* mask,
+                       float* data,
+                       int block_size,
+                       int batch_idx,
+                       int X,
+                       int Y,
+                       size_t X_data_str,
+                       size_t Y_data_str,
+                       const Shape& mask_shape,
+                       const Strides& mask_strides,
+                       bool is_bool) {
+    auto ndim = mask_shape.size();
+    auto mask_offset = elem_to_loc(
+        mask_shape[ndim - 1] * mask_shape[ndim - 2] * batch_idx,
+        mask_shape,
+        mask_strides);
+
+    auto X_mask_str = mask_strides[ndim - 2];
+    auto Y_mask_str = mask_strides[ndim - 1];
+
+    if (is_bool) {
+      return mask_matrix(
+          data,
+          static_cast<const bool*>(mask),
+          block_size,
+          X,
+          Y,
+          X_data_str,
+          Y_data_str,
+          X_mask_str,
+          Y_mask_str,
+          mask_offset);
+    } else {
+      return mask_matrix(
+          data,
+          static_cast<const float*>(mask),
+          block_size,
+          X,
+          Y,
+          X_data_str,
+          Y_data_str,
+          X_mask_str,
+          Y_mask_str,
+          mask_offset);
+    }
+  };
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  const void* a_mask_ptr = nullptr;
+  const void* b_mask_ptr = nullptr;
+  const void* out_mask_ptr = nullptr;
+  Shape a_mask_shape;
+  Shape b_mask_shape;
+  Shape out_mask_shape;
+  Strides a_mask_strides;
+  Strides b_mask_strides;
+  Strides out_mask_strides;
+  bool a_mask_bool = false;
+  bool b_mask_bool = false;
+  bool out_mask_bool = false;
+  if (has_op_mask) {
+    auto& a_mask = inputs[inputs.size() - 2];
+    auto& b_mask = inputs[inputs.size() - 1];
+    a_mask_ptr = a_mask.data<void>();
+    b_mask_ptr = b_mask.data<void>();
+    a_mask_shape = a_mask.shape();
+    b_mask_shape = b_mask.shape();
+    a_mask_strides = a_mask.strides();
+    b_mask_strides = b_mask.strides();
+    a_mask_bool = (a_mask.dtype() == bool_);
+    b_mask_bool = (b_mask.dtype() == bool_);
+    encoder.set_input_array(a_mask);
+    encoder.set_input_array(b_mask);
+  }
+  if (has_out_mask) {
+    auto& out_mask = inputs[2];
+    out_mask_ptr = out_mask.data<void>();
+    out_mask_bool = (out_mask.dtype() == bool_);
+    encoder.set_input_array(out_mask);
+    out_mask_shape = out_mask.shape();
+    out_mask_strides = out_mask.strides();
+  }
+  encoder.set_output_array(out);
+  auto a_ptr = a.data<float>();
+  auto b_ptr = b.data<float>();
+  auto out_ptr = out.data<float>();
+  size_t num_matrices = out.size() / (M * size_t(N));
+  auto ldc = out.shape(-1);
+
+  encoder.dispatch([a_ptr,
+                    b_ptr,
+                    out_ptr,
+                    a_mask_ptr,
+                    b_mask_ptr,
+                    out_mask_ptr,
+                    has_op_mask,
+                    has_out_mask,
+                    block_size = block_size_,
+                    num_matrices,
+                    M,
+                    N,
+                    K,
+                    a_transposed = a_transposed,
+                    b_transposed = b_transposed,
+                    lda = lda,
+                    ldb = ldb,
+                    ldc,
+                    a_shape = a.shape(),
+                    a_strides = a.strides(),
+                    b_shape = b.shape(),
+                    b_strides = b.strides(),
+                    a_mask_shape = std::move(a_mask_shape),
+                    b_mask_shape = std::move(b_mask_shape),
+                    out_mask_shape = std::move(out_mask_shape),
+                    a_mask_strides = std::move(a_mask_strides),
+                    b_mask_strides = std::move(b_mask_strides),
+                    out_mask_strides = std::move(out_mask_strides),
+                    mask_array,
+                    a_mask_bool,
+                    b_mask_bool,
+                    out_mask_bool]() {
+    for (int i = 0; i < num_matrices; ++i) {
+      // Adjust pointer
+      float* ai = a_ptr + elem_to_loc(M * K * i, a_shape, a_strides);
+      float* bi = b_ptr + elem_to_loc(K * N * i, b_shape, b_strides);
+      float* ci = out_ptr + M * N * i;
+
+      // Zero out blocks in a and b if needed
+      if (has_op_mask) {
+        mask_array(
+            a_mask_ptr,
+            ai,
+            block_size,
+            i,
+            M,
+            K,
+            a_transposed ? 1 : lda,
+            a_transposed ? lda : 1,
+            a_mask_shape,
+            a_mask_strides,
+            a_mask_bool);
+
+        mask_array(
+            b_mask_ptr,
+            bi,
+            block_size,
+            i,
+            K,
+            N,
+            b_transposed ? 1 : ldb,
+            b_transposed ? ldb : 1,
+            b_mask_shape,
+            b_mask_strides,
+            b_mask_bool);
+      }
+
+      // Do matmul
+      cblas_sgemm(
+          CblasRowMajor,
+          a_transposed ? CblasTrans : CblasNoTrans, // transA
+          b_transposed ? CblasTrans : CblasNoTrans, // transB
+          M,
+          N,
+          K,
+          1.0, // alpha
+          ai,
+          lda,
+          bi,
+          ldb,
+          0.0, // beta
+          ci,
+          ldc);
+
+      // Zero out blocks in out
+      if (has_out_mask) {
+        mask_array(
+            out_mask_ptr,
+            ci,
+            block_size,
+            i,
+            M,
+            N,
+            N,
+            1,
+            out_mask_shape,
+            out_mask_strides,
+            out_mask_bool);
+      }
+    }
+  });
+  if (a_copied) {
+    encoder.add_temporary(a);
+  }
+  if (b_copied) {
+    encoder.add_temporary(b);
+  }
+}
+
+void GatherMM::eval_cpu(const std::vector<array>& inputs, array& out) {
+  if (out.dtype() != float32) {
+    throw std::runtime_error(
+        "[GatherMM::eval] Currently only supports float32.");
+  }
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+
+  std::vector<array> temps;
+  auto check_transpose = [s = stream(), &temps](const array& arr) {
+    auto stx = arr.strides()[arr.ndim() - 2];
+    auto sty = arr.strides()[arr.ndim() - 1];
+    if (stx == arr.shape(-1) && sty == 1) {
+      return std::make_tuple(false, stx, arr);
+    } else if (stx == 1 && sty == arr.shape(-2)) {
+      return std::make_tuple(true, sty, arr);
+    } else {
+      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      copy_cpu(arr, temps.back(), CopyType::General, s);
+      int64_t stx = arr.shape(-1);
+      return std::make_tuple(false, stx, temps.back());
+    }
+  };
+
+  auto [a_transposed, lda, a] = check_transpose(a_pre);
+  auto [b_transposed, ldb, b] = check_transpose(b_pre);
+
+  size_t M = a.shape(-2);
+  size_t N = b.shape(-1);
+  size_t K = a.shape(-1);
+
+  if (M == 0 || N == 0) {
+    return;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  if (K == 0) {
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<float>(), size = out.size()]() {
+      std::fill_n(out_ptr, size, 0);
+    });
+    return;
+  }
+
+  // Get batch dims
+  auto batch_size_out = out.size() / (M * N);
+  size_t matrix_stride_out = M * N;
+
+  auto get_batch_dims = [](const auto& v) {
+    return decltype(v){v.begin(), v.end() - 2};
+  };
+
+  auto& lhs_indices = inputs[2];
+  auto& rhs_indices = inputs[3];
+
+  auto batch_shape = get_batch_dims(out.shape());
+
+  auto batch_shape_A = get_batch_dims(a.shape());
+  auto batch_strides_A = get_batch_dims(a.strides());
+  auto batch_shape_B = get_batch_dims(b.shape());
+  auto batch_strides_B = get_batch_dims(b.strides());
+
+  const uint32_t* lhs_indices_ptr = lhs_indices.data<uint32_t>();
+  const uint32_t* rhs_indices_ptr = rhs_indices.data<uint32_t>();
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(lhs_indices);
+  encoder.set_input_array(rhs_indices);
+  encoder.set_output_array(out);
+  auto ldc = out.shape(-1);
+
+  encoder.dispatch([a_ptr = a.data<float>(),
+                    b_ptr = b.data<float>(),
+                    out_ptr = out.data<float>(),
+                    M,
+                    N,
+                    K,
+                    lda = lda,
+                    ldb = ldb,
+                    a_transposed = a_transposed,
+                    b_transposed = b_transposed,
+                    ldc,
+                    lhs_indices_ptr,
+                    rhs_indices_ptr,
+                    lhs_indices_shape = lhs_indices.shape(),
+                    lhs_indices_strides = lhs_indices.strides(),
+                    rhs_indices_shape = rhs_indices.shape(),
+                    rhs_indices_strides = rhs_indices.strides(),
+                    batch_size_out,
+                    matrix_stride_out,
+                    batch_shape_A = std::move(batch_shape_A),
+                    batch_shape_B = std::move(batch_shape_B),
+                    batch_strides_A = std::move(batch_strides_A),
+                    batch_strides_B = std::move(batch_strides_B)]() {
+    for (int i = 0; i < batch_size_out; i++) {
+      // Get index
+      uint32_t indx_A = lhs_indices_ptr[elem_to_loc(
+          i, lhs_indices_shape, lhs_indices_strides)];
+      uint32_t indx_B = rhs_indices_ptr[elem_to_loc(
+          i, rhs_indices_shape, rhs_indices_strides)];
+
+      cblas_sgemm(
+          CblasRowMajor,
+          a_transposed ? CblasTrans : CblasNoTrans, // transA
+          b_transposed ? CblasTrans : CblasNoTrans, // transB
+          M,
+          N,
+          K,
+          1.0f, // alpha
+          a_ptr + elem_to_loc(indx_A, batch_shape_A, batch_strides_A),
+          lda,
+          b_ptr + elem_to_loc(indx_B, batch_shape_B, batch_strides_B),
+          ldb,
+          0.0f, // beta
+          out_ptr + matrix_stride_out * i,
+          ldc);
+    }
+  });
+  encoder.add_temporaries(std::move(temps));
+}
+
+void SegmentedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& encoder = cpu::get_command_encoder(stream());
+  auto check_transpose = [&s, &encoder](const array& x) {
+    auto stx = x.strides()[x.ndim() - 2];
+    auto sty = x.strides()[x.ndim() - 1];
+    if (stx == x.shape(-1) && sty == 1) {
+      return std::make_tuple(false, stx, x);
+    } else if (stx == 1 && sty == x.shape(-2)) {
+      return std::make_tuple(true, sty, x);
+    } else {
+      array xc(x.shape(), x.dtype(), nullptr, {});
+      copy_cpu(x, xc, CopyType::General, s);
+      encoder.add_temporary(xc);
+      int64_t stx = x.shape(-1);
+      return std::make_tuple(false, stx, xc);
+    }
+  };
+
+  auto [a_transposed, lda, a] = check_transpose(inputs[0]);
+  auto [b_transposed, ldb, b] = check_transpose(inputs[1]);
+  auto& segments = inputs[2];
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(segments);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    segments = array::unsafe_weak_copy(segments),
+                    out_ptr = out.data<void>(),
+                    a_transposed = a_transposed,
+                    b_transposed = b_transposed,
+                    lda = lda,
+                    ldb = ldb]() {
+    switch (a.dtype()) {
+      case float64:
+        segmented_mm<double>(
+            a.data<double>(),
+            b.data<double>(),
+            segments.data<uint32_t>(),
+            static_cast<double*>(out_ptr),
+            a_transposed,
+            b_transposed,
+            lda,
+            ldb,
+            a.shape(),
+            a.strides(),
+            b.shape(),
+            b.strides(),
+            segments.size() / 2,
+            segments.shape(),
+            segments.strides());
+        break;
+      case float32:
+        segmented_mm<float>(
+            a.data<float>(),
+            b.data<float>(),
+            segments.data<uint32_t>(),
+            static_cast<float*>(out_ptr),
+            a_transposed,
+            b_transposed,
+            lda,
+            ldb,
+            a.shape(),
+            a.strides(),
+            b.shape(),
+            b.strides(),
+            segments.size() / 2,
+            segments.shape(),
+            segments.strides());
+        break;
+      case float16:
+        segmented_mm<float16_t>(
+            a.data<float16_t>(),
+            b.data<float16_t>(),
+            segments.data<uint32_t>(),
+            static_cast<float16_t*>(out_ptr),
+            a_transposed,
+            b_transposed,
+            lda,
+            ldb,
+            a.shape(),
+            a.strides(),
+            b.shape(),
+            b.strides(),
+            segments.size() / 2,
+            segments.shape(),
+            segments.strides());
+        break;
+      case bfloat16:
+        segmented_mm<bfloat16_t>(
+            a.data<bfloat16_t>(),
+            b.data<bfloat16_t>(),
+            segments.data<uint32_t>(),
+            static_cast<bfloat16_t*>(out_ptr),
+            a_transposed,
+            b_transposed,
+            lda,
+            ldb,
+            a.shape(),
+            a.strides(),
+            b.shape(),
+            b.strides(),
+            segments.size() / 2,
+            segments.shape(),
+            segments.strides());
+        break;
+      default:
+        throw std::invalid_argument(
+            "Segmented mm supports only real float types.");
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/matmul.cpp b/macos/third_party/mlx/mlx/backend/cpu/matmul.cpp
new file mode 100644
index 0000000..7df3316
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/matmul.cpp
@@ -0,0 +1,166 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstring>
+#include "mlx/array.h"
+#include "mlx/backend/cpu/binary.h"
+#include "mlx/backend/cpu/binary_ops.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/gemm.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T>
+void matmul_dispatch(
+    const array& a,
+    const array& b,
+    array& out,
+    bool a_transposed,
+    bool b_transposed,
+    size_t lda,
+    size_t ldb,
+    float alpha,
+    float beta,
+    Stream stream) {
+  const T* a_ptr = a.data<T>();
+  const T* b_ptr = b.data<T>();
+  T* out_ptr = out.data<T>();
+  size_t ldc = out.shape(-1);
+  size_t batch_size = a.size() / (a.shape(-2) * a.shape(-1));
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  encoder.dispatch([a_ptr,
+                    b_ptr,
+                    out_ptr,
+                    a_transposed,
+                    b_transposed,
+                    lda,
+                    ldb,
+                    ldc,
+                    alpha,
+                    beta,
+                    batch_size,
+                    a_shape = a.shape(),
+                    a_strides = a.strides(),
+                    b_shape = b.shape(),
+                    b_strides = b.strides()]() {
+    matmul<T>(
+        a_ptr,
+        b_ptr,
+        out_ptr,
+        a_transposed,
+        b_transposed,
+        lda,
+        ldb,
+        ldc,
+        alpha,
+        beta,
+        batch_size,
+        a_shape,
+        a_strides,
+        b_shape,
+        b_strides);
+  });
+}
+
+void matmul_general(
+    const array& a_pre,
+    const array& b_pre,
+    array& out,
+    Stream stream,
+    float alpha = 1.0f,
+    float beta = 0.0f) {
+  std::vector<array> temps;
+  auto check_transpose = [stream, &temps](const array& arr) {
+    auto stx = arr.strides()[arr.ndim() - 2];
+    auto sty = arr.strides()[arr.ndim() - 1];
+    if (stx == arr.shape(-1) && sty == 1) {
+      return std::make_tuple(false, stx, arr);
+    } else if (stx == 1 && sty == arr.shape(-2)) {
+      return std::make_tuple(true, sty, arr);
+    } else {
+      temps.push_back(array(arr.shape(), arr.dtype(), nullptr, {}));
+      copy_cpu(arr, temps.back(), CopyType::General, stream);
+      stx = arr.shape(-1);
+      return std::make_tuple(false, stx, temps.back());
+    }
+  };
+
+  auto [a_transposed, lda, a] = check_transpose(a_pre);
+  auto [b_transposed, ldb, b] = check_transpose(b_pre);
+  size_t M = a.shape(-2);
+  size_t N = b.shape(-1);
+  if (M == 0 || N == 0) {
+    return;
+  }
+
+  if (out.dtype() == float32) {
+    matmul_dispatch<float>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == float16) {
+    matmul_dispatch<float16_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == bfloat16) {
+    matmul_dispatch<bfloat16_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == float64) {
+    matmul_dispatch<double>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else if (out.dtype() == complex64) {
+    matmul_dispatch<complex64_t>(
+        a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream);
+  } else {
+    throw std::runtime_error("[Matmul::eval_cpu] Invalid type.");
+  }
+  cpu::get_command_encoder(stream).add_temporaries(std::move(temps));
+}
+
+void Matmul::eval_cpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+  if (inputs[0].shape(-1) == 0) {
+    auto& encoder = cpu::get_command_encoder(stream());
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<void>(), nbytes = out.nbytes()]() {
+      std::memset(out_ptr, 0, nbytes);
+    });
+    return;
+  }
+  matmul_general(inputs[0], inputs[1], out, stream());
+}
+
+void AddMM::eval_cpu(const std::vector<array>& inputs, array& out) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(out.nbytes()));
+    return;
+  }
+
+  // Handle empty matrix case (K=0)
+  if (inputs[0].shape(-1) == 0) {
+    auto& c = inputs[2];
+    if (beta_ == 1.0f) {
+      CopyType ctype = c.data_size() == 1
+          ? CopyType::Scalar
+          : (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+      copy_cpu(c, out, ctype, stream());
+    } else {
+      array beta_scalar = array(beta_, c.dtype());
+      auto& encoder = cpu::get_command_encoder(stream());
+      binary_float_op_cpu(c, beta_scalar, out, detail::Multiply(), stream());
+      encoder.add_temporary(std::move(beta_scalar));
+    }
+    return;
+  }
+
+  // Fill output with C
+  auto& c = inputs[2];
+  CopyType ctype = c.data_size() == 1
+      ? CopyType::Scalar
+      : (c.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+  copy_cpu(c, out, ctype, stream());
+  matmul_general(inputs[0], inputs[1], out, stream(), alpha_, beta_);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/primitives.cpp b/macos/third_party/mlx/mlx/backend/cpu/primitives.cpp
new file mode 100644
index 0000000..4e59b1e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/primitives.cpp
@@ -0,0 +1,478 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <numeric>
+#include <sstream>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/slicing.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/arange.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/threefry.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+void reshape(const array& in, array& out) {
+  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
+  if (copy_necessary) {
+    out.set_data(allocator::malloc(out.nbytes()));
+    copy_cpu_inplace(in, out, CopyType::General, out.primitive().stream());
+  } else {
+    shared_buffer_reshape(in, out_strides, out);
+  }
+}
+
+static std::pair<array, bool> compute_dynamic_offset(
+    const array& indices,
+    const Strides& strides,
+    const std::vector<int>& axes,
+    Stream stream) {
+  array offset({1}, int64, nullptr, {});
+  bool donate = indices.is_donatable() &&
+      (indices.data_size() * indices.itemsize()) >= offset.itemsize();
+  if (donate) {
+    offset.copy_shared_buffer(indices);
+  } else {
+    offset.set_data(allocator::malloc(offset.itemsize()));
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(indices);
+  encoder.set_output_array(offset);
+  auto compute_offset =
+      [strides, axes, offset = offset.data<int64_t>()](const auto* indices) {
+        int64_t offset_ = 0;
+        for (int i = 0; i < axes.size(); ++i) {
+          offset_ += indices[i] * strides[axes[i]];
+        }
+        offset[0] = offset_;
+      };
+  switch (indices.dtype()) {
+    case int8:
+    case uint8:
+      encoder.dispatch(compute_offset, indices.data<uint8_t>());
+      break;
+    case int16:
+    case uint16:
+      encoder.dispatch(compute_offset, indices.data<uint16_t>());
+      break;
+    case int32:
+    case uint32:
+      encoder.dispatch(compute_offset, indices.data<uint32_t>());
+      break;
+    case int64:
+    case uint64:
+      encoder.dispatch(compute_offset, indices.data<uint64_t>());
+      break;
+    default:
+      throw std::runtime_error("Invalid indices type.");
+  }
+  return {offset, donate};
+}
+
+void AsStrided::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Broadcast::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void BroadcastAxes::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Copy::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void CustomTransforms::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+void Depends::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+void ExpandDims::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void NumberOfElements::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Slice::eval_cpu(const std::vector<array>& inputs, array& out) {
+  slice(inputs[0], out, start_indices_, strides_);
+}
+void Split::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  eval(inputs, outputs);
+}
+void Squeeze::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void StopGradient::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+void Transpose::eval_cpu(const std::vector<array>& inputs, array& out) {
+  eval(inputs, out);
+}
+
+void Arange::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 0);
+  out.set_data(allocator::malloc(out.nbytes()));
+  switch (out.dtype()) {
+    case bool_:
+      throw std::runtime_error("Bool type unsupported for arange.");
+      break;
+    case uint8:
+      arange<uint8_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint16:
+      arange<uint16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint32:
+      arange<uint32_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case uint64:
+      arange<uint64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int8:
+      arange<int8_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int16:
+      arange<int16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int32:
+      arange<int32_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case int64:
+      arange<int64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float16:
+      arange<float16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float32:
+      arange<float>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case float64:
+      arange<double>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case bfloat16:
+      arange<bfloat16_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+    case complex64:
+      arange<complex64_t>(start_, start_ + step_, out, out.size(), stream());
+      break;
+  }
+}
+
+void AsType::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General;
+  copy_cpu(in, out, ctype, stream());
+}
+
+void Concatenate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  std::vector<int> sizes;
+  sizes.push_back(0);
+  for (auto& p : inputs) {
+    sizes.push_back(p.shape(axis_));
+  }
+  std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto strides = out.strides();
+  auto flags = out.flags();
+  flags.row_contiguous = false;
+  flags.col_contiguous = false;
+  flags.contiguous = false;
+  for (int i = 0; i < inputs.size(); i++) {
+    array out_slice(inputs[i].shape(), out.dtype(), nullptr, {});
+    size_t data_offset = strides[axis_] * sizes[i];
+    out_slice.copy_shared_buffer(
+        out, strides, flags, out_slice.size(), data_offset);
+    copy_cpu_inplace(inputs[i], out_slice, CopyType::GeneralGeneral, stream());
+  }
+}
+
+void Contiguous::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  constexpr size_t extra_bytes = 16384;
+  if (in.buffer_size() <= out.nbytes() + extra_bytes &&
+      (in.flags().row_contiguous ||
+       (allow_col_major_ && in.flags().col_contiguous))) {
+    out.copy_shared_buffer(in);
+  } else {
+    copy_cpu(in, out, CopyType::General, stream());
+  }
+}
+
+void Flatten::eval_cpu(const std::vector<array>& inputs, array& out) {
+  reshape(inputs[0], out);
+}
+
+void Unflatten::eval_cpu(const std::vector<array>& inputs, array& out) {
+  reshape(inputs[0], out);
+}
+
+void Full::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  assert(in.dtype() == out.dtype());
+  CopyType ctype;
+  if (in.data_size() == 1) {
+    ctype = CopyType::Scalar;
+  } else if (in.flags().contiguous) {
+    ctype = CopyType::Vector;
+  } else {
+    ctype = CopyType::General;
+  }
+  copy_cpu(in, out, ctype, stream());
+}
+
+void Pad::eval_cpu(const std::vector<array>& inputs, array& out) {
+  // Inputs must be base input array and scalar val array
+  assert(inputs.size() == 2);
+  auto& in = inputs[0];
+  auto& val = inputs[1];
+
+  // Padding value must be a scalar
+  assert(val.size() == 1);
+
+  // Padding value, input and output must be of the same type
+  assert(val.dtype() == in.dtype() && in.dtype() == out.dtype());
+
+  // Fill output with val
+  copy_cpu(val, out, CopyType::Scalar, stream());
+
+  // Find offset for start of input values
+  size_t data_offset = 0;
+  for (int i = 0; i < axes_.size(); i++) {
+    auto ax = axes_[i] < 0 ? out.ndim() + axes_[i] : axes_[i];
+    data_offset += out.strides()[ax] * low_pad_size_[i];
+  }
+
+  // Extract slice from output where input will be pasted
+  array out_slice(in.shape(), out.dtype(), nullptr, {});
+  out_slice.copy_shared_buffer(
+      out, out.strides(), out.flags(), out_slice.size(), data_offset);
+
+  // Copy input values into the slice
+  copy_cpu_inplace(in, out_slice, CopyType::GeneralGeneral, stream());
+}
+
+void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  // keys has shape (N1, ..., NK, 2)
+  // out has shape (N1, ..., NK, M1, M2, ...)
+  auto& keys = inputs[0];
+  size_t num_keys = keys.size() / 2;
+
+  size_t elems_per_key = out.size() / num_keys;
+  size_t bytes_per_key = out.itemsize() * elems_per_key;
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto kptr = inputs[0].data<uint32_t>();
+  auto cptr = out.data<char>();
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(inputs[0]);
+  encoder.set_output_array(out);
+  encoder.dispatch([kptr,
+                    cptr,
+                    bytes_per_key,
+                    num_keys,
+                    kshape = keys.shape(),
+                    kstrides = keys.strides()]() mutable {
+    auto copy_remaining = [&](char* cptr, size_t loc, uint32_t v) {
+      if (4 * loc + 4 <= bytes_per_key) {
+        reinterpret_cast<uint32_t*>(cptr)[loc] = v;
+      } else {
+        std::copy(
+            reinterpret_cast<char*>(&v),
+            reinterpret_cast<char*>(&v) + bytes_per_key - 4 * loc,
+            cptr + 4 * loc);
+      }
+    };
+
+    size_t out_skip = (bytes_per_key + 4 - 1) / 4;
+    auto half_size = out_skip / 2;
+    bool even = out_skip % 2 == 0;
+    for (int i = 0; i < num_keys; ++i, cptr += bytes_per_key) {
+      auto ptr = reinterpret_cast<uint32_t*>(cptr);
+      // Get ith key
+      auto kidx = 2 * i;
+      auto k1_elem = elem_to_loc(kidx, kshape, kstrides);
+      auto k2_elem = elem_to_loc(kidx + 1, kshape, kstrides);
+      auto key = std::make_pair(kptr[k1_elem], kptr[k2_elem]);
+
+      std::pair<uintptr_t, uintptr_t> count{0, half_size + !even};
+      for (; count.first + 1 < half_size; count.first++, count.second++) {
+        std::tie(ptr[count.first], ptr[count.second]) =
+            random::threefry2x32_hash(key, count);
+      }
+      if (count.first < half_size) {
+        auto rb = random::threefry2x32_hash(key, count);
+        ptr[count.first++] = rb.first;
+        copy_remaining(cptr, count.second, rb.second);
+      }
+      if (!even) {
+        count.second = 0;
+        copy_remaining(
+            cptr, half_size, random::threefry2x32_hash(key, count).first);
+      }
+    }
+  });
+}
+
+void Reshape::eval_cpu(const std::vector<array>& inputs, array& out) {
+  reshape(inputs[0], out);
+}
+
+void DynamicSlice::eval_cpu(const std::vector<array>& inputs, array& out) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+  auto& in = inputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto [in_offset, donated] =
+      compute_dynamic_offset(inputs[1], in.strides(), axes_, stream());
+  copy_cpu_inplace(
+      /* const array& src = */ in,
+      /* array& dst = */ out,
+      /* const Shape& data_shape = */ out.shape(),
+      /* const Strides& i_strides = */ in.strides(),
+      /* const Strides& o_strides = */ out.strides(),
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ 0,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream(),
+      /* const std::optional<array>& dynamic_i_offset = */ in_offset,
+      /* const std::optional<array>& dynamic_o_offset = */ std::nullopt);
+  if (!donated) {
+    cpu::get_command_encoder(stream()).add_temporary(std::move(in_offset));
+  }
+}
+
+void DynamicSliceUpdate::eval_cpu(
+    const std::vector<array>& inputs,
+    array& out) {
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  // Copy or move src to dst
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+
+  auto [out_offset, donated] =
+      compute_dynamic_offset(inputs[2], out.strides(), axes_, stream());
+  copy_cpu_inplace(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const std::vector<int>& data_shape = */ upd.shape(),
+      /* const std::vector<stride_t>& i_strides = */ upd.strides(),
+      /* const std::vector<stride_t>& o_strides = */ out.strides(),
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ 0,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream(),
+      /* const std::optional<array>& dynamic_i_offset = */ std::nullopt,
+      /* const std::optional<array>& dynamic_o_offset = */ out_offset);
+  if (!donated) {
+    cpu::get_command_encoder(stream()).add_temporary(std::move(out_offset));
+  }
+}
+
+void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  // Check if materialization is needed
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [data_offset, out_strides] =
+      prepare_slice(out, start_indices_, strides_);
+
+  // Do copy
+  copy_cpu_inplace(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const std::vector<int>& data_shape = */ upd.shape(),
+      /* const std::vector<stride_t>& i_strides = */ upd.strides(),
+      /* const std::vector<stride_t>& o_strides = */ out_strides,
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ data_offset,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      stream());
+}
+
+void View::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  auto ibytes = size_of(in.dtype());
+  auto obytes = size_of(out.dtype());
+  // Conditions for buffer copying (disjunction):
+  // - type size is the same
+  // - type size is smaller and the last axis is contiguous
+  // - the entire array is row contiguous
+  if (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) ||
+      in.flags().row_contiguous) {
+    auto strides = in.strides();
+    for (int i = 0; i < static_cast<int>(strides.size()) - 1; ++i) {
+      strides[i] *= ibytes;
+      strides[i] /= obytes;
+    }
+    out.copy_shared_buffer(
+        in, strides, in.flags(), in.data_size() * ibytes / obytes);
+  } else {
+    auto tmp = array(
+        in.shape(), in.dtype() == bool_ ? uint8 : in.dtype(), nullptr, {});
+    tmp.set_data(allocator::malloc(tmp.nbytes()));
+    if (in.dtype() == bool_) {
+      auto in_tmp = array(in.shape(), uint8, nullptr, {});
+      in_tmp.copy_shared_buffer(in);
+      copy_cpu_inplace(in_tmp, tmp, CopyType::General, stream());
+    } else {
+      copy_cpu_inplace(in, tmp, CopyType::General, stream());
+    }
+
+    auto flags = out.flags();
+    flags.contiguous = true;
+    flags.row_contiguous = true;
+    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
+    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
+    out.copy_shared_buffer(tmp, out.strides(), flags, out.size());
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/qrf.cpp b/macos/third_party/mlx/mlx/backend/cpu/qrf.cpp
new file mode 100644
index 0000000..13c7e11
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/qrf.cpp
@@ -0,0 +1,147 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T>
+void qrf_impl(const array& a, array& q, array& r, Stream stream) {
+  const int M = a.shape(-2);
+  const int N = a.shape(-1);
+  const int lda = M;
+  size_t num_matrices = a.size() / (M * N);
+
+  // Copy A to inplace input and make it col-contiguous
+  array in(a.shape(), a.dtype(), nullptr, {});
+  auto flags = in.flags();
+
+  // Copy the input to be column contiguous
+  flags.col_contiguous = num_matrices == 1;
+  flags.row_contiguous = false;
+  auto strides = in.strides();
+  strides[in.ndim() - 2] = 1;
+  strides[in.ndim() - 1] = M;
+  in.set_data(allocator::malloc(in.nbytes()), in.nbytes(), strides, flags);
+  copy_cpu_inplace(a, in, CopyType::GeneralGeneral, stream);
+  auto& encoder = cpu::get_command_encoder(stream);
+  q.set_data(allocator::malloc(q.nbytes()));
+  r.set_data(allocator::malloc(r.nbytes()));
+
+  auto in_ptr = in.data<T>();
+  auto r_ptr = r.data<T>();
+  auto q_ptr = q.data<T>();
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(q);
+  encoder.set_output_array(r);
+  encoder.dispatch([in_ptr, q_ptr, r_ptr, M, N, lda, num_matrices]() {
+    int num_reflectors = std::min(M, N);
+    auto tau = allocator::malloc(sizeof(T) * num_matrices * num_reflectors);
+
+    T optimal_work;
+    int lwork = -1;
+    int info;
+
+    // Compute workspace size
+    geqrf<T>(&M, &N, nullptr, &lda, nullptr, &optimal_work, &lwork, &info);
+
+    // Update workspace size
+    lwork = optimal_work;
+    auto work = allocator::malloc(sizeof(T) * lwork);
+
+    // Loop over matrices
+    for (int i = 0; i < num_matrices; ++i) {
+      // Solve
+      geqrf<T>(
+          &M,
+          &N,
+          in_ptr + M * N * i,
+          &lda,
+          static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
+          static_cast<T*>(work.raw_ptr()),
+          &lwork,
+          &info);
+    }
+    allocator::free(work);
+
+    for (int i = 0; i < num_matrices; ++i) {
+      /// num_reflectors x N
+      for (int j = 0; j < num_reflectors; ++j) {
+        for (int k = 0; k < j; ++k) {
+          r_ptr[i * N * num_reflectors + j * N + k] = 0;
+        }
+        for (int k = j; k < N; ++k) {
+          r_ptr[i * N * num_reflectors + j * N + k] =
+              in_ptr[i * N * M + j + k * M];
+        }
+      }
+    }
+
+    // Get work size
+    lwork = -1;
+    orgqr<T>(
+        &M,
+        &num_reflectors,
+        &num_reflectors,
+        nullptr,
+        &lda,
+        nullptr,
+        &optimal_work,
+        &lwork,
+        &info);
+    lwork = optimal_work;
+    work = allocator::malloc(sizeof(T) * lwork);
+
+    // Loop over matrices
+    for (int i = 0; i < num_matrices; ++i) {
+      // Compute Q
+      orgqr<T>(
+          &M,
+          &num_reflectors,
+          &num_reflectors,
+          in_ptr + M * N * i,
+          &lda,
+          static_cast<T*>(tau.raw_ptr()) + num_reflectors * i,
+          static_cast<T*>(work.raw_ptr()),
+          &lwork,
+          &info);
+    }
+
+    for (int i = 0; i < num_matrices; ++i) {
+      // M x num_reflectors
+      for (int j = 0; j < M; ++j) {
+        for (int k = 0; k < num_reflectors; ++k) {
+          q_ptr[i * M * num_reflectors + j * num_reflectors + k] =
+              in_ptr[i * N * M + j + k * M];
+        }
+      }
+    }
+
+    // Cleanup
+    allocator::free(work);
+    allocator::free(tau);
+  });
+  encoder.add_temporary(in);
+}
+
+void QRF::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  switch (inputs[0].dtype()) {
+    case float32:
+      qrf_impl<float>(inputs[0], outputs[0], outputs[1], stream());
+      break;
+    case float64:
+      qrf_impl<double>(inputs[0], outputs[0], outputs[1], stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[QRF::eval_cpu] only supports float32 or float64.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/quantized.cpp b/macos/third_party/mlx/mlx/backend/cpu/quantized.cpp
new file mode 100644
index 0000000..3c3643a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/quantized.cpp
@@ -0,0 +1,1370 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/backend/common/unary.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/backend/cpu/unary.h"
+#include "mlx/backend/cpu/unary_ops.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+array ensure_row_contiguous(
+    const array& arr,
+    cpu::CommandEncoder& encoder,
+    Stream s) {
+  if (arr.flags().row_contiguous) {
+    return arr;
+  } else {
+    auto arr_cpy = contiguous_copy_cpu(arr, s);
+    encoder.add_temporary(arr_cpy);
+    return arr_cpy;
+  }
+};
+
+const static float FP4_LUT[16] = {
+    +0.0f,
+    +0.5f,
+    +1.0f,
+    +1.5f,
+    +2.0f,
+    +3.0f,
+    +4.0f,
+    +6.0f,
+    -0.0f,
+    -0.5f,
+    -1.0f,
+    -1.5f,
+    -2.0f,
+    -3.0f,
+    -4.0f,
+    -6.0f};
+
+template <typename T, int group_size>
+static inline T dequantize_scale(uint8_t s) {
+  if constexpr (group_size == 16) {
+    return static_cast<T>(detail::FromFP8{}(s));
+  } else {
+    using FOrI = union {
+      bfloat16_t f;
+      uint16_t i;
+    };
+    FOrI out;
+    out.i = (s == 0 ? 0x40 : (static_cast<uint16_t>(s) << 7));
+    return static_cast<T>(out.f);
+  }
+}
+
+inline constexpr short get_pack_factor(int bits, int wsize = 8) {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
+  auto power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
+template <typename T, int bits>
+void extract_bits(const uint8_t* w_in, T* w_out) {
+  static_assert(bits == 3 || bits == 5 || bits == 6);
+  if (bits == 3) {
+    w_out[0] = static_cast<T>(w_in[0] & 0x7);
+    w_out[1] = static_cast<T>((w_in[0] & 0x38) >> 3);
+    w_out[2] = static_cast<T>(((w_in[0] & 0xc0) >> 6) + ((w_in[1] & 0x1) << 2));
+    w_out[3] = static_cast<T>((w_in[1] & 0xe) >> 1);
+    w_out[4] = static_cast<T>((w_in[1] & 0x70) >> 4);
+    w_out[5] = static_cast<T>(((w_in[1] & 0x80) >> 7) + ((w_in[2] & 0x3) << 1));
+    w_out[6] = static_cast<T>((w_in[2] & 0x1c) >> 2);
+    w_out[7] = static_cast<T>((w_in[2] & 0xe0) >> 5);
+  } else if (bits == 5) {
+    w_out[0] = static_cast<T>(w_in[0] & 0x1f);
+    w_out[1] = static_cast<T>(((w_in[0] & 0xe0) >> 5) + ((w_in[1] & 0x3) << 3));
+    w_out[2] = static_cast<T>((w_in[1] & 0x7c) >> 2);
+    w_out[3] = static_cast<T>(((w_in[1] & 0x80) >> 7) + ((w_in[2] & 0xf) << 1));
+    w_out[4] = static_cast<T>(((w_in[2] & 0xf0) >> 4) + ((w_in[3] & 0x1) << 4));
+    w_out[5] = static_cast<T>((w_in[3] & 0x3e) >> 1);
+    w_out[6] = static_cast<T>(((w_in[3] & 0xc0) >> 6) + ((w_in[4] & 0x7) << 2));
+    w_out[7] = static_cast<T>((w_in[4] & 0xf8) >> 3);
+
+  } else if (bits == 6) {
+    w_out[0] = static_cast<T>(w_in[0] & 0x3f);
+    w_out[1] =
+        static_cast<T>(((w_in[0] >> 6) & 0x03) + ((w_in[1] & 0x0f) << 2));
+    w_out[2] =
+        static_cast<T>(((w_in[1] >> 4) & 0x0f) + ((w_in[2] & 0x03) << 4));
+    w_out[3] = static_cast<T>((w_in[2] >> 2) & 0x3f);
+  }
+}
+
+template <typename T, int bits, int group_size>
+void _qmm(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K) {
+  constexpr int bitmask = (1 << bits) - 1;
+  constexpr int pack_factor = get_pack_factor(bits, 8);
+  constexpr int bytes_per_pack = get_bytes_per_pack(bits);
+  constexpr int packs_in_group = group_size / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint8_t* w_local = (const uint8_t*)w;
+    const T* scales_local = scales;
+    const T* biases_local = biases;
+
+    std::fill(result, result + N, 0);
+
+    for (int k = 0; k < K; k++) {
+      T* result_local = result;
+      T xi = *x++;
+
+      for (int n = 0; n < N; n += group_size) {
+        T scale = *scales_local++;
+        T bias = *biases_local++;
+        for (int ng = 0; ng < packs_in_group; ng++) {
+          if constexpr (bits == 3 || bits == 5 || bits == 6) {
+            T wl[pack_factor];
+            extract_bits<T, bits>(w_local, wl);
+#pragma clang loop unroll(full)
+            for (int p = 0; p < pack_factor; p++) {
+              (*result_local++) += xi * (scale * wl[p] + bias);
+            }
+            w_local += bytes_per_pack;
+
+          } else {
+            uint8_t wi = *w_local++;
+#pragma clang loop unroll(full)
+            for (int p = 0; p < pack_factor; p++) {
+              (*result_local++) +=
+                  xi * (scale * static_cast<T>(wi & bitmask) + bias);
+              if (bits != 8) {
+                wi >>= bits;
+              }
+            }
+          }
+        }
+      }
+    }
+
+    result += N;
+  }
+}
+
+template <typename T, int bits, int group_size>
+void _qmm_t(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K) {
+  constexpr int bitmask = (1 << bits) - 1;
+
+  constexpr int pack_factor = get_pack_factor(bits, 8);
+  constexpr int bytes_per_pack = get_bytes_per_pack(bits);
+  constexpr int packs_in_group = group_size / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint8_t* w_local = (const uint8_t*)w;
+    const T* scales_local = scales;
+    const T* biases_local = biases;
+
+    for (int n = 0; n < N; n++) {
+      const T* x_local = x;
+      T sum = 0;
+      for (int k = 0; k < K; k += group_size) {
+        T scale = *scales_local++;
+        T bias = *biases_local++;
+
+        for (int kw = 0; kw < packs_in_group; kw++) {
+          if constexpr (bits == 3 || bits == 5 || bits == 6) {
+            T wl[pack_factor];
+            extract_bits<T, bits>(w_local, wl);
+#pragma clang loop unroll(full)
+            for (int p = 0; p < pack_factor; p++) {
+              sum += x_local[p] * (scale * wl[p] + bias);
+            }
+            w_local += bytes_per_pack;
+            x_local += pack_factor;
+
+          } else {
+            uint8_t wi = *w_local++;
+#pragma clang loop unroll(full)
+            for (int p = 0; p < pack_factor; p++) {
+              sum +=
+                  (*x_local++) * (scale * static_cast<T>(wi & bitmask) + bias);
+              if (bits != 8) {
+                wi >>= bits;
+              }
+            }
+          }
+        }
+      }
+      *result = sum;
+      result++;
+    }
+
+    x += K;
+  }
+}
+
+template <int bits, int S>
+simd::Simd<uint32_t, S> extract_bits_simd(const uint32_t* w) {
+  constexpr int bitmask = (1 << bits) - 1;
+  simd::Simd<uint32_t, S> wi;
+  if constexpr (bits == 4 && S == 8) {
+    constexpr std::array<uint32_t, 8> shifts_ = {{0, 4, 8, 12, 16, 20, 24, 28}};
+    auto shifts(*(simd::Simd<uint32_t, S>*)&shifts_);
+    wi = simd::Simd<uint32_t, S>(*w);
+    wi = wi >> shifts;
+    wi = wi & bitmask;
+  } else if constexpr (bits == 8 && S == 8) {
+    constexpr std::array<uint32_t, 8> shifts_ = {{0, 8, 16, 24, 0, 8, 16, 24}};
+    auto shifts(*(simd::Simd<uint32_t, S>*)&shifts_);
+    auto l = simd::Simd<uint32_t, S / 2>(*w++);
+    auto r = simd::Simd<uint32_t, S / 2>(*w);
+    wi = simd::Simd<uint32_t, S>(l, r);
+    wi = wi >> shifts;
+    wi = wi & bitmask;
+  } else {
+    // Appease compiler.. but should never get here
+    throw std::runtime_error("Unsupported combination for simd qmm.");
+  }
+  return wi;
+}
+
+template <typename T, int bits, int group_size>
+void _qmm_t_simd(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K) {
+  constexpr int pack_factor = 32 / bits;
+  constexpr int packs_in_group = group_size / pack_factor;
+  constexpr int S = simd::max_size<T>;
+  static_assert(
+      S % pack_factor == 0, "SIMD size must be divisible by pack factor");
+  constexpr int packs_per_simd = S / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint32_t* w_local = w;
+    const T* scales_local = scales;
+    const T* biases_local = biases;
+
+    for (int n = 0; n < N; n++) {
+      simd::Simd<float, S> acc(0);
+      auto x_local = x;
+      for (int k = 0; k < K; k += group_size) {
+        T scale = *scales_local++;
+        T bias = *biases_local++;
+
+        for (int kw = 0; kw < packs_in_group; kw += packs_per_simd) {
+          auto wf = simd::Simd<float, S>(extract_bits_simd<bits, S>(w_local));
+          w_local += packs_per_simd;
+          wf = wf * scale;
+          wf = wf + bias;
+          simd::Simd<float, S> x_simd = simd::load<T, S>(x_local);
+          acc = acc + x_simd * wf;
+          x_local += S;
+        }
+      }
+
+      *result = T(simd::sum(acc));
+      result++;
+    }
+    x += K;
+  }
+}
+
+template <typename T, int bits, int group_size>
+void _qmm_dispatch_transpose(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K,
+    bool transposed_w) {
+  if (transposed_w) {
+    // the simd size must be a multiple of the number of elements per word
+    if constexpr (32 % bits == 0 && simd::max_size<T> % (32 / bits) == 0) {
+      _qmm_t_simd<T, bits, group_size>(result, x, w, scales, biases, M, N, K);
+    } else {
+      _qmm_t<T, bits, group_size>(result, x, w, scales, biases, M, N, K);
+    }
+  } else {
+    _qmm<T, bits, group_size>(result, x, w, scales, biases, M, N, K);
+  }
+}
+
+template <typename T, int bits>
+void _qmm_dispatch_group(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K,
+    int group_size,
+    bool transposed_w) {
+  switch (group_size) {
+    case 32:
+      _qmm_dispatch_transpose<T, bits, 32>(
+          result, x, w, scales, biases, M, N, K, transposed_w);
+      break;
+    case 64:
+      _qmm_dispatch_transpose<T, bits, 64>(
+          result, x, w, scales, biases, M, N, K, transposed_w);
+      break;
+    case 128:
+      _qmm_dispatch_transpose<T, bits, 128>(
+          result, x, w, scales, biases, M, N, K, transposed_w);
+      break;
+    default:
+      throw std::invalid_argument(
+          "Quantization group size must be 32, 64 or 128.");
+  }
+}
+
+template <typename T>
+void _qmm_dispatch_typed(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const T* scales,
+    const T* biases,
+    int M,
+    int N,
+    int K,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  switch (bits) {
+    case 2:
+      _qmm_dispatch_group<T, 2>(
+          result, x, w, scales, biases, M, N, K, group_size, transposed_w);
+      break;
+    case 3:
+      _qmm_dispatch_group<T, 3>(
+          result, x, w, scales, biases, M, N, K, group_size, transposed_w);
+      break;
+    case 4:
+      _qmm_dispatch_group<T, 4>(
+          result, x, w, scales, biases, M, N, K, group_size, transposed_w);
+      break;
+    case 5:
+      _qmm_dispatch_group<T, 5>(
+          result, x, w, scales, biases, M, N, K, group_size, transposed_w);
+      break;
+    case 6:
+      _qmm_dispatch_group<T, 6>(
+          result, x, w, scales, biases, M, N, K, group_size, transposed_w);
+      break;
+    case 8:
+      _qmm_dispatch_group<T, 8>(
+          result, x, w, scales, biases, M, N, K, group_size, transposed_w);
+      break;
+    default:
+      throw std::invalid_argument("Quantization bits must be 2, 3, 4, 6 or 8.");
+  }
+}
+
+template <typename T>
+void _qmm_dispatch_typed(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int bits,
+    int group_size,
+    bool transposed_w) {
+  int K = x.shape(-1);
+  int M = x.ndim() > 1 ? x.shape(-2) : 1;
+  int N = out.shape(-1);
+  int w_els = w.ndim() > 2 ? w.shape(-1) * w.shape(-2) : 0;
+  int g_els = w.ndim() > 2 ? scales.shape(-1) * scales.shape(-2) : 0;
+  int batch_size = x.size() / (K * M);
+
+  auto out_ptr = out.data<T>();
+  auto x_ptr = x.data<T>();
+  auto w_ptr = w.data<uint32_t>();
+  auto scales_ptr = scales.data<T>();
+  auto biases_ptr = biases.data<T>();
+  for (int i = 0; i < batch_size; i++) {
+    _qmm_dispatch_typed<T>(
+        out_ptr + i * M * N,
+        x_ptr + elem_to_loc(i * M * K, x.shape(), x.strides()),
+        w_ptr + elem_to_loc(i * w_els, w.shape(), w.strides()),
+        scales_ptr + elem_to_loc(i * g_els, scales.shape(), scales.strides()),
+        biases_ptr + elem_to_loc(i * g_els, biases.shape(), biases.strides()),
+        M,
+        N,
+        K,
+        bits,
+        group_size,
+        transposed_w);
+  }
+}
+
+void _qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int bits,
+    int group_size,
+    bool transposed_w) {
+  switch (x.dtype()) {
+    case float32:
+      _qmm_dispatch_typed<float>(
+          out, x, w, scales, biases, bits, group_size, transposed_w);
+      break;
+    case float16:
+      _qmm_dispatch_typed<float16_t>(
+          out, x, w, scales, biases, bits, group_size, transposed_w);
+      break;
+    case bfloat16:
+      _qmm_dispatch_typed<bfloat16_t>(
+          out, x, w, scales, biases, bits, group_size, transposed_w);
+      break;
+    default:
+      throw std::invalid_argument(
+          "[quantized_matmul] only floating types are supported");
+  }
+}
+
+template <typename T, int group_size, int bits>
+void fp_qmm(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const uint8_t* scales,
+    int M,
+    int N,
+    int K) {
+  constexpr int pack_factor = get_pack_factor(bits, 8);
+  constexpr int packs_in_group = group_size / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint8_t* w_local = (const uint8_t*)w;
+    const uint8_t* scales_local = scales;
+
+    std::fill(result, result + N, 0);
+
+    for (int k = 0; k < K; k++) {
+      T* result_local = result;
+      T xi = *x++;
+
+      for (int n = 0; n < N; n += group_size) {
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
+        for (int ng = 0; ng < packs_in_group; ng++) {
+          if constexpr (bits == 4) {
+            (*result_local++) +=
+                xi * scale * static_cast<T>(FP4_LUT[w_local[0] & 0xf]);
+            (*result_local++) +=
+                xi * scale * static_cast<T>(FP4_LUT[(w_local[0] >> 4) & 0xf]);
+          } else {
+            (*result_local++) +=
+                xi * scale * static_cast<T>(detail::FromFP8{}(w_local[0]));
+          }
+          w_local++;
+        }
+      }
+    }
+    result += N;
+  }
+}
+
+template <typename T, int group_size, int bits>
+void fp_qmm_t(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const uint8_t* scales,
+    int M,
+    int N,
+    int K) {
+  constexpr int pack_factor = get_pack_factor(bits, 8);
+  constexpr int packs_in_group = group_size / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint8_t* w_local = (const uint8_t*)w;
+    const uint8_t* scales_local = scales;
+
+    for (int n = 0; n < N; n++) {
+      const T* x_local = x;
+      T sum = 0;
+      for (int k = 0; k < K; k += group_size) {
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
+
+        T gsum = 0;
+        for (int kw = 0; kw < packs_in_group; kw++) {
+          if constexpr (bits == 4) {
+            gsum += (*x_local++) * static_cast<T>(FP4_LUT[w_local[0] & 0xf]);
+            gsum +=
+                (*x_local++) * static_cast<T>(FP4_LUT[(w_local[0] >> 4) & 0xf]);
+          } else {
+            gsum +=
+                (*x_local++) * static_cast<T>(detail::FromFP8{}(w_local[0]));
+          }
+          w_local++;
+        }
+        sum += scale * gsum;
+      }
+      *result = sum;
+      result++;
+    }
+
+    x += K;
+  }
+}
+
+template <int S, int bits>
+simd::Simd<float, S> fp_extract_bits_simd(const uint32_t* w) {
+  if constexpr (S == 8 && bits == 4) {
+    constexpr std::array<uint32_t, 8> shifts_ = {{0, 4, 8, 12, 16, 20, 24, 28}};
+    auto shifts(*(simd::Simd<uint32_t, S>*)&shifts_);
+    auto wi = simd::Simd<uint32_t, S>(*w);
+    wi = wi >> shifts;
+    wi = wi & 0xf;
+    simd::Simd<float, S> w_out;
+    for (int i = 0; i < S; ++i) {
+      w_out[i] = FP4_LUT[wi[i]];
+    }
+    return w_out;
+  } else if constexpr (S == 8 && bits == 8) {
+    auto w_out = simd::load<uint8_t, S>(reinterpret_cast<const uint8_t*>(w));
+    return detail::FromFP8{}(w_out);
+  } else {
+    // Appease compiler.. but should never get here
+    throw std::runtime_error("Unsupported combination for simd qmm.");
+  }
+}
+
+template <typename T, int group_size, int bits>
+void fp_qmm_t_simd(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const uint8_t* scales,
+    int M,
+    int N,
+    int K) {
+  constexpr int pack_factor = get_pack_factor(bits, 32);
+  constexpr int packs_in_group = group_size / pack_factor;
+  constexpr int S = simd::max_size<T>;
+  static_assert(
+      S % pack_factor == 0, "SIMD size must be divisible by pack factor");
+  constexpr int packs_per_simd = S / pack_factor;
+
+  for (int m = 0; m < M; m++) {
+    const uint32_t* w_local = w;
+    const uint8_t* scales_local = scales;
+
+    for (int n = 0; n < N; n++) {
+      simd::Simd<float, S> acc(0);
+      auto x_local = x;
+      for (int k = 0; k < K; k += group_size) {
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
+
+        simd::Simd<float, S> g_acc(0);
+        for (int kw = 0; kw < packs_in_group; kw += packs_per_simd) {
+          // Extract bits
+          auto wf = fp_extract_bits_simd<S, bits>(w_local);
+          w_local += packs_per_simd;
+          simd::Simd<float, S> x_simd = simd::load<T, S>(x_local);
+          g_acc = g_acc + x_simd * wf;
+          x_local += S;
+        }
+        acc = acc + scale * g_acc;
+      }
+
+      *result = T(simd::sum(acc));
+      result++;
+    }
+    x += K;
+  }
+}
+
+template <typename T, int group_size, int bits>
+void fp_qmm_dispatch_transpose(
+    T* result,
+    const T* x,
+    const uint32_t* w,
+    const uint8_t* scales,
+    int M,
+    int N,
+    int K,
+    bool transposed_w) {
+  if (transposed_w) {
+    // the simd size must be a multiple of the number of elements per word
+    if constexpr (simd::max_size<T> % 8 == 0) {
+      fp_qmm_t_simd<T, group_size, bits>(result, x, w, scales, M, N, K);
+    } else {
+      fp_qmm_t<T, group_size, bits>(result, x, w, scales, M, N, K);
+    }
+  } else {
+    fp_qmm<T, group_size, bits>(result, x, w, scales, M, N, K);
+  }
+}
+
+template <typename T, int group_size, int bits>
+void fp_qmm_dispatch_mode(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    bool transposed_w) {
+  int K = x.shape(-1);
+  int M = x.ndim() > 1 ? x.shape(-2) : 1;
+  int N = out.shape(-1);
+  int w_els = w.ndim() > 2 ? w.shape(-1) * w.shape(-2) : 0;
+  int g_els = w.ndim() > 2 ? scales.shape(-1) * scales.shape(-2) : 0;
+  int batch_size = x.size() / (K * M);
+
+  auto out_ptr = out.data<T>();
+  auto x_ptr = x.data<T>();
+  auto w_ptr = w.data<uint32_t>();
+  auto scales_ptr = scales.data<uint8_t>();
+  for (int i = 0; i < batch_size; i++) {
+    fp_qmm_dispatch_transpose<T, group_size, bits>(
+        out_ptr + i * M * N,
+        x_ptr + elem_to_loc(i * M * K, x.shape(), x.strides()),
+        w_ptr + elem_to_loc(i * w_els, w.shape(), w.strides()),
+        scales_ptr + elem_to_loc(i * g_els, scales.shape(), scales.strides()),
+        M,
+        N,
+        K,
+        transposed_w);
+  }
+}
+
+template <typename T>
+void fp_qmm_dispatch_typed(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  if (bits == 8) {
+    fp_qmm_dispatch_mode<T, 32, 8>(out, x, w, scales, transposed_w);
+  } else if (group_size == 32) {
+    fp_qmm_dispatch_mode<T, 32, 4>(out, x, w, scales, transposed_w);
+  } else {
+    fp_qmm_dispatch_mode<T, 16, 4>(out, x, w, scales, transposed_w);
+  }
+}
+
+void fp_qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  switch (x.dtype()) {
+    case bfloat16:
+      fp_qmm_dispatch_typed<bfloat16_t>(
+          out, x, w, scales, group_size, bits, transposed_w);
+      break;
+    case float16:
+      fp_qmm_dispatch_typed<float16_t>(
+          out, x, w, scales, group_size, bits, transposed_w);
+      break;
+    case float32:
+      fp_qmm_dispatch_typed<float>(
+          out, x, w, scales, group_size, bits, transposed_w);
+      break;
+    default:
+      throw std::invalid_argument(
+          "[quantized_matmul] only floating types are supported");
+  }
+}
+
+template <typename T>
+void _bs_qmm_dispatch_typed(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int bits,
+    int group_size,
+    bool transposed_w) {
+  int K = x.shape(-1);
+  int M = x.shape(-2);
+  int N = out.shape(-1);
+
+  int w_els = w.shape(-1) * w.shape(-2);
+  int g_els = scales.shape(-1) * scales.shape(-2);
+
+  auto out_ptr = out.data<T>();
+  auto x_ptr = x.data<T>();
+  auto w_ptr = w.data<uint32_t>();
+  auto scales_ptr = scales.data<T>();
+  auto biases_ptr = biases.data<T>();
+  auto lhs_indices_ptr = lhs_indices.data<uint32_t>();
+  auto rhs_indices_ptr = rhs_indices.data<uint32_t>();
+
+  for (int i = 0; i < lhs_indices.size(); i++) {
+    int x_idx = lhs_indices_ptr[elem_to_loc(
+        i, lhs_indices.shape(), lhs_indices.strides())];
+    int w_idx = rhs_indices_ptr[elem_to_loc(
+        i, rhs_indices.shape(), rhs_indices.strides())];
+    _qmm_dispatch_typed<T>(
+        out_ptr + i * M * N,
+        x_ptr + elem_to_loc(x_idx * M * K, x.shape(), x.strides()),
+        w_ptr + elem_to_loc(w_idx * w_els, w.shape(), w.strides()),
+        scales_ptr +
+            elem_to_loc(w_idx * g_els, scales.shape(), scales.strides()),
+        biases_ptr +
+            elem_to_loc(w_idx * g_els, biases.shape(), biases.strides()),
+        M,
+        N,
+        K,
+        bits,
+        group_size,
+        transposed_w);
+  }
+}
+
+void _bs_qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int bits,
+    int group_size,
+    bool transposed_w) {
+  switch (x.dtype()) {
+    case float32:
+      _bs_qmm_dispatch_typed<float>(
+          out,
+          x,
+          w,
+          scales,
+          biases,
+          lhs_indices,
+          rhs_indices,
+          bits,
+          group_size,
+          transposed_w);
+      break;
+    case float16:
+      _bs_qmm_dispatch_typed<float16_t>(
+          out,
+          x,
+          w,
+          scales,
+          biases,
+          lhs_indices,
+          rhs_indices,
+          bits,
+          group_size,
+          transposed_w);
+      break;
+    case bfloat16:
+      _bs_qmm_dispatch_typed<bfloat16_t>(
+          out,
+          x,
+          w,
+          scales,
+          biases,
+          lhs_indices,
+          rhs_indices,
+          bits,
+          group_size,
+          transposed_w);
+      break;
+    default:
+      throw std::invalid_argument(
+          "[quantized_matmul] only floating types are supported");
+  }
+}
+template <typename T, int group_size, int bits>
+void fp_bs_qmm_dispatch_mode(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    bool transposed_w) {
+  int K = x.shape(-1);
+  int M = x.shape(-2);
+  int N = out.shape(-1);
+
+  int w_els = w.shape(-1) * w.shape(-2);
+  int g_els = scales.shape(-1) * scales.shape(-2);
+
+  auto out_ptr = out.data<T>();
+  auto x_ptr = x.data<T>();
+  auto w_ptr = w.data<uint32_t>();
+  auto scales_ptr = scales.data<uint8_t>();
+  auto lhs_indices_ptr = lhs_indices.data<uint32_t>();
+  auto rhs_indices_ptr = rhs_indices.data<uint32_t>();
+
+  for (int i = 0; i < lhs_indices.size(); i++) {
+    int x_idx = lhs_indices_ptr[elem_to_loc(
+        i, lhs_indices.shape(), lhs_indices.strides())];
+    int w_idx = rhs_indices_ptr[elem_to_loc(
+        i, rhs_indices.shape(), rhs_indices.strides())];
+    fp_qmm_dispatch_transpose<T, group_size, bits>(
+        out_ptr + i * M * N,
+        x_ptr + elem_to_loc(x_idx * M * K, x.shape(), x.strides()),
+        w_ptr + elem_to_loc(w_idx * w_els, w.shape(), w.strides()),
+        scales_ptr +
+            elem_to_loc(w_idx * g_els, scales.shape(), scales.strides()),
+        M,
+        N,
+        K,
+        transposed_w);
+  }
+}
+
+template <typename T>
+void fp_bs_qmm_dispatch_typed(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  if (bits == 8) {
+    fp_bs_qmm_dispatch_mode<T, 32, 8>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  } else if (group_size == 32) {
+    fp_bs_qmm_dispatch_mode<T, 32, 4>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  } else {
+    fp_bs_qmm_dispatch_mode<T, 16, 4>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  }
+}
+
+void fp_bs_qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  switch (x.dtype()) {
+    case float32:
+      fp_bs_qmm_dispatch_typed<float>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
+      break;
+    case float16:
+      fp_bs_qmm_dispatch_typed<float16_t>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
+      break;
+    case bfloat16:
+      fp_bs_qmm_dispatch_typed<bfloat16_t>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
+      break;
+    default:
+      throw std::invalid_argument(
+          "[quantized_matmul] only floating types are supported");
+  }
+}
+
+} // namespace
+
+void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
+  auto& x_pre = inputs[0];
+  auto& w_pre = inputs[1];
+  auto& scales_pre = inputs[2];
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  auto x = ensure_row_contiguous(x_pre, encoder, stream());
+  auto w = ensure_row_contiguous(w_pre, encoder, stream());
+  auto scales = ensure_row_contiguous(scales_pre, encoder, stream());
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(scales);
+  encoder.set_output_array(out);
+  if (mode_ == QuantizationMode::Affine) {
+    auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
+    encoder.set_input_array(biases);
+    encoder.dispatch([out = array::unsafe_weak_copy(out),
+                      x = array::unsafe_weak_copy(x),
+                      w = array::unsafe_weak_copy(w),
+                      scales = array::unsafe_weak_copy(scales),
+                      biases = array::unsafe_weak_copy(biases),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
+                      transpose_ = transpose_]() mutable {
+      _qmm_dispatch(out, x, w, scales, biases, group_size_, bits_, transpose_);
+    });
+  } else {
+    encoder.dispatch([out = array::unsafe_weak_copy(out),
+                      x = array::unsafe_weak_copy(x),
+                      w = array::unsafe_weak_copy(w),
+                      scales = array::unsafe_weak_copy(scales),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
+                      transpose_ = transpose_]() mutable {
+      fp_qmm_dispatch(out, x, w, scales, group_size_, bits_, transpose_);
+    });
+  }
+}
+
+void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
+  auto& x_pre = inputs[0];
+  auto& w_pre = inputs[1];
+  auto& scales_pre = inputs[2];
+  auto& lhs_indices = inputs[inputs.size() - 2];
+  auto& rhs_indices = inputs[inputs.size() - 1];
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  auto ensure_row_contiguous_last_dims = [s = stream(),
+                                          &encoder](const array& arr) {
+    auto stride_0 = arr.strides()[arr.ndim() - 2];
+    auto stride_1 = arr.strides()[arr.ndim() - 1];
+    if (stride_0 == arr.shape(-1) && stride_1 == 1) {
+      return arr;
+    } else {
+      auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {});
+      copy_cpu(arr, arr_cpy, CopyType::General, s);
+      encoder.add_temporary(arr_cpy);
+      return arr_cpy;
+    }
+  };
+
+  auto x = ensure_row_contiguous_last_dims(x_pre);
+  auto w = ensure_row_contiguous_last_dims(w_pre);
+  auto scales = ensure_row_contiguous_last_dims(scales_pre);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(lhs_indices);
+  encoder.set_input_array(rhs_indices);
+  encoder.set_output_array(out);
+  if (mode_ == QuantizationMode::Affine) {
+    auto biases = ensure_row_contiguous_last_dims(inputs[3]);
+    encoder.set_input_array(biases);
+    encoder.dispatch([out = array::unsafe_weak_copy(out),
+                      x = array::unsafe_weak_copy(x),
+                      w = array::unsafe_weak_copy(w),
+                      scales = array::unsafe_weak_copy(scales),
+                      biases = array::unsafe_weak_copy(biases),
+                      lhs_indices = array::unsafe_weak_copy(lhs_indices),
+                      rhs_indices = array::unsafe_weak_copy(rhs_indices),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
+                      transpose_ = transpose_]() mutable {
+      _bs_qmm_dispatch(
+          out,
+          x,
+          w,
+          scales,
+          biases,
+          lhs_indices,
+          rhs_indices,
+          group_size_,
+          bits_,
+          transpose_);
+    });
+  } else {
+    encoder.dispatch([out = array::unsafe_weak_copy(out),
+                      x = array::unsafe_weak_copy(x),
+                      w = array::unsafe_weak_copy(w),
+                      scales = array::unsafe_weak_copy(scales),
+                      lhs_indices = array::unsafe_weak_copy(lhs_indices),
+                      rhs_indices = array::unsafe_weak_copy(rhs_indices),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
+                      transpose_ = transpose_]() mutable {
+      fp_bs_qmm_dispatch(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size_,
+          bits_,
+          transpose_);
+    });
+  }
+}
+
+uint8_t to_fp8_e8m0(float x) {
+  if (!std::isfinite(x)) {
+    return 0xFF;
+  }
+  if (x < 0.0f) {
+    return 0x00;
+  }
+  float le = std::log2(x);
+  int n = int(std::round(le));
+
+  n = n < -127 ? -127 : n;
+  n = n > 127 ? 127 : n;
+  return static_cast<uint8_t>(n + 127);
+}
+
+uint8_t to_fp4_e2m1(float x) {
+  if (std::isnan(x)) {
+    return 0x7;
+  }
+
+  const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0;
+  x = std::abs(x);
+
+  uint8_t bits;
+  if (x > 5.0f) {
+    bits = 0x7;
+  } else if (x >= 3.5f) {
+    bits = 0x6;
+  } else if (x > 2.5f) {
+    bits = 0x5;
+  } else if (x >= 1.75f) {
+    bits = 0x4;
+  } else if (x > 1.25f) {
+    bits = 0x3;
+  } else if (x >= 0.75f) {
+    bits = 0x2;
+  } else if (x > 0.25f) {
+    bits = 0x1;
+  } else {
+    bits = 0x0;
+  }
+  return bits | sign_bit;
+}
+
+template <typename T>
+void fp_quantize_dequantize(
+    const array& w_arr,
+    array& out_arr,
+    int bits,
+    int group_size,
+    size_t w_size) {
+  auto w = w_arr.data<T>();
+  auto out = out_arr.data<T>();
+
+  size_t n_groups = w_size / group_size;
+
+  for (size_t i = 0; i < n_groups; ++i) {
+    size_t idx = i * group_size;
+    float scale = -std::numeric_limits<float>::infinity();
+    for (int j = 0; j < group_size; ++j) {
+      scale = std::max(scale, std::abs(w[idx + j]));
+    }
+    scale /= bits == 4 ? 6.0f : 448.0f;
+    if (group_size == 16) {
+      scale = dequantize_scale<float, 16>(detail::ToFP8()(scale));
+    } else {
+      scale = dequantize_scale<float, 32>(to_fp8_e8m0(scale));
+    }
+
+    for (int j = 0; j < group_size; ++j) {
+      float w_el = scale == 0 ? 0.0f : w[idx + j] / scale;
+      float output;
+      if (bits == 8) {
+        output = detail::FromFP8()(detail::ToFP8()(w_el));
+      } else {
+        output = FP4_LUT[to_fp4_e2m1(w_el)];
+      }
+      out[idx + j] = static_cast<T>(scale * output);
+    }
+  }
+}
+
+void dispatch_quantize_dequantize(
+    const array& w,
+    array& out,
+    int bits,
+    int group_size) {
+  if (w.dtype() == float16) {
+    fp_quantize_dequantize<float16_t>(w, out, bits, group_size, w.size());
+  } else if (w.dtype() == bfloat16) {
+    fp_quantize_dequantize<bfloat16_t>(w, out, bits, group_size, w.size());
+  } else if (w.dtype() == float32) {
+    fp_quantize_dequantize<float>(w, out, bits, group_size, w.size());
+  } else {
+    throw std::runtime_error(
+        "[quantize_dequantize] Only supports floating point inputs");
+  }
+}
+
+template <typename T, typename U>
+void quantize(
+    const T* w,
+    U* out,
+    T* scales,
+    T* biases,
+    int bits,
+    int group_size,
+    size_t w_size) {
+  float n_bins = (1 << bits) - 1;
+  float eps = 1e-7;
+
+  bool power_of_2_bits = is_power_of_2(bits);
+  int el_per_int = get_pack_factor(bits, 32);
+  int bytes_per_pack = get_bytes_per_pack(bits);
+  int int_per_group = group_size * bytes_per_pack / el_per_int;
+  size_t n_groups = w_size / group_size;
+
+  for (size_t i = 0; i < n_groups; ++i) {
+    size_t w_idx = i * group_size;
+    float w_min = std::numeric_limits<float>::infinity();
+    float w_max = -w_min;
+    for (int j = 0; j < group_size; ++j) {
+      w_max = std::max(w_max, (float)w[w_idx + j]);
+      w_min = std::min(w_min, (float)w[w_idx + j]);
+    }
+    bool mask = std::abs(w_min) > std::abs(w_max);
+    float scale = std::max((w_max - w_min) / n_bins, eps);
+    scale = mask ? scale : -scale;
+
+    float edge = mask ? w_min : w_max;
+    float q0 = std::rint(edge / scale);
+    float bias = 0;
+    if (q0 != 0) {
+      scale = edge / q0;
+      bias = edge;
+    }
+    size_t out_idx = i * int_per_group;
+    for (int j = 0; j < int_per_group / bytes_per_pack; ++j) {
+      uint64_t out_el = 0;
+      for (int k = 0; k < el_per_int; ++k) {
+        float w_el = w[w_idx + j * el_per_int + k];
+        w_el = std::rint((w_el - bias) / scale);
+        w_el = std::min(std::max(w_el, 0.0f), n_bins);
+        out_el |= static_cast<uint64_t>(w_el) << (k * bits);
+      }
+      if (power_of_2_bits) {
+        out[out_idx + j] = out_el;
+      } else if (bits == 5) {
+        out[out_idx + bytes_per_pack * j] = out_el & 0xff;
+        out[out_idx + bytes_per_pack * j + 1] = (out_el & 0xff00) >> 8;
+        out[out_idx + bytes_per_pack * j + 2] = (out_el & 0xff0000) >> 16;
+        out[out_idx + bytes_per_pack * j + 3] = (out_el & 0xff000000) >> 24;
+        out[out_idx + bytes_per_pack * j + 4] = (out_el & 0xff00000000) >> 32;
+      } else {
+        out[out_idx + bytes_per_pack * j] = out_el & 0xff;
+        out[out_idx + bytes_per_pack * j + 1] = (out_el & 0xff00) >> 8;
+        out[out_idx + bytes_per_pack * j + 2] = (out_el & 0xff0000) >> 16;
+      }
+    }
+    scales[i] = static_cast<T>(scale);
+    biases[i] = static_cast<T>(bias);
+  }
+}
+
+template <typename T, typename U>
+void dispatch_quantize(
+    const array& w,
+    array& out,
+    array& scales,
+    array& biases,
+    int bits,
+    int group_size) {
+  auto w_ptr = w.data<T>();
+  auto out_ptr = out.data<U>();
+  auto scales_ptr = scales.data<T>();
+  auto biases_ptr = biases.data<T>();
+  quantize<T, U>(
+      w_ptr, out_ptr, scales_ptr, biases_ptr, bits, group_size, w.size());
+}
+
+void fast::Quantize::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& encoder = cpu::get_command_encoder(stream());
+  auto w = ensure_row_contiguous(inputs[0], encoder, stream());
+  auto& out = outputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& scales = outputs[1];
+  auto& biases = outputs[2];
+  scales.set_data(allocator::malloc(scales.nbytes()));
+  biases.set_data(allocator::malloc(biases.nbytes()));
+  encoder.set_input_array(w);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(biases);
+  encoder.set_output_array(out);
+  encoder.dispatch([w = array::unsafe_weak_copy(w),
+                    out = array::unsafe_weak_copy(out),
+                    scales = array::unsafe_weak_copy(scales),
+                    biases = array::unsafe_weak_copy(biases),
+                    group_size_ = group_size_,
+                    bits_ = bits_]() mutable {
+    if (w.dtype() == float16) {
+      if (is_power_of_2(bits_)) {
+        dispatch_quantize<float16_t, uint32_t>(
+            w, out, scales, biases, bits_, group_size_);
+      } else {
+        dispatch_quantize<float16_t, uint8_t>(
+            w, out, scales, biases, bits_, group_size_);
+      }
+    } else if (w.dtype() == bfloat16) {
+      if (is_power_of_2(bits_)) {
+        dispatch_quantize<bfloat16_t, uint32_t>(
+            w, out, scales, biases, bits_, group_size_);
+      } else {
+        dispatch_quantize<bfloat16_t, uint8_t>(
+            w, out, scales, biases, bits_, group_size_);
+      }
+    } else if (w.dtype() == float32) {
+      if (is_power_of_2(bits_)) {
+        dispatch_quantize<float, uint32_t>(
+            w, out, scales, biases, bits_, group_size_);
+      } else {
+        dispatch_quantize<float, uint8_t>(
+            w, out, scales, biases, bits_, group_size_);
+      }
+    } else {
+      throw std::runtime_error(
+          "[fast::Quantize::eval_cpu] Only supports floating point inputs");
+    }
+  });
+}
+
+void fast::ConvertFP8::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& in = inputs[0];
+  auto& out = outputs[0];
+  set_unary_output_data(in, out);
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    to_fp8 = to_fp8_]() mutable {
+    if (to_fp8) {
+      switch (in.dtype()) {
+        case float16:
+          unary_op<float16_t, uint8_t>(in, out, detail::ToFP8());
+          break;
+        case bfloat16:
+          unary_op<bfloat16_t, uint8_t>(in, out, detail::ToFP8());
+          break;
+        default:
+          unary_op<float, uint8_t>(in, out, detail::ToFP8());
+          break;
+      }
+    } else {
+      switch (out.dtype()) {
+        case float16:
+          unary_op<uint8_t, float16_t>(in, out, detail::FromFP8());
+          break;
+        case bfloat16:
+          unary_op<uint8_t, bfloat16_t>(in, out, detail::FromFP8());
+          break;
+        default:
+          unary_op<uint8_t, float>(in, out, detail::FromFP8());
+          break;
+      }
+    }
+  });
+}
+
+void QQMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
+  auto& encoder = cpu::get_command_encoder(stream());
+
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    bool donate_x = inputs[0].is_donatable();
+    auto x = ensure_row_contiguous(inputs[0], encoder, stream());
+    auto w = ensure_row_contiguous(inputs[1], encoder, stream());
+    auto scales = ensure_row_contiguous(inputs[2], encoder, stream());
+
+    out.set_data(allocator::malloc(out.nbytes()));
+
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype());
+    if (!donate_x) {
+      encoder.add_temporary(xhat);
+    }
+    encoder.set_input_array(x);
+    encoder.set_input_array(w);
+    encoder.set_input_array(scales);
+    encoder.set_output_array(out);
+    encoder.dispatch([out = array::unsafe_weak_copy(out),
+                      x = array::unsafe_weak_copy(x),
+                      xhat = array::unsafe_weak_copy(xhat),
+                      w = array::unsafe_weak_copy(w),
+                      scales = array::unsafe_weak_copy(scales),
+                      group_size_ = group_size_,
+                      bits_ = bits_]() mutable {
+      dispatch_quantize_dequantize(x, xhat, bits_, group_size_);
+      fp_qmm_dispatch(out, xhat, w, scales, group_size_, bits_, true);
+    });
+    return;
+  } else {
+    throw std::runtime_error("[QQMatmul] NYI for the general case");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/reduce.cpp b/macos/third_party/mlx/mlx/backend/cpu/reduce.cpp
new file mode 100644
index 0000000..41764f4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/reduce.cpp
@@ -0,0 +1,587 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+#include <functional>
+#include <limits>
+
+#include "mlx/backend/common/reduce.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename U>
+struct Limits {
+  static const U max;
+  static const U min;
+};
+
+#define instantiate_default_limit(type)                           \
+  template <>                                                     \
+  struct Limits<type> {                                           \
+    static constexpr type max = std::numeric_limits<type>::max(); \
+    static constexpr type min = std::numeric_limits<type>::min(); \
+  };
+
+instantiate_default_limit(uint8_t);
+instantiate_default_limit(uint16_t);
+instantiate_default_limit(uint32_t);
+instantiate_default_limit(uint64_t);
+instantiate_default_limit(int8_t);
+instantiate_default_limit(int16_t);
+instantiate_default_limit(int32_t);
+instantiate_default_limit(int64_t);
+
+#define instantiate_float_limit(type) \
+  template <>                         \
+  struct Limits<type> {               \
+    static const type max;            \
+    static const type min;            \
+  };
+
+instantiate_float_limit(float16_t);
+instantiate_float_limit(bfloat16_t);
+instantiate_float_limit(float);
+instantiate_float_limit(double);
+instantiate_float_limit(complex64_t);
+
+template <>
+struct Limits<bool> {
+  static constexpr bool max = true;
+  static constexpr bool min = false;
+};
+
+const float Limits<float>::max = std::numeric_limits<float>::infinity();
+const float Limits<float>::min = -std::numeric_limits<float>::infinity();
+const bfloat16_t Limits<bfloat16_t>::max =
+    std::numeric_limits<float>::infinity();
+const bfloat16_t Limits<bfloat16_t>::min =
+    -std::numeric_limits<float>::infinity();
+const float16_t Limits<float16_t>::max = std::numeric_limits<float>::infinity();
+const float16_t Limits<float16_t>::min =
+    -std::numeric_limits<float>::infinity();
+const double Limits<double>::max = std::numeric_limits<double>::infinity();
+const double Limits<double>::min = -std::numeric_limits<double>::infinity();
+const complex64_t Limits<complex64_t>::max =
+    std::numeric_limits<float>::infinity();
+const complex64_t Limits<complex64_t>::min =
+    -std::numeric_limits<float>::infinity();
+
+template <typename T, typename U, typename Op>
+void strided_reduce(
+    const T* x,
+    U* accumulator,
+    int size,
+    size_t stride,
+    Op op) {
+  constexpr int N = std::min(simd::max_size<T>, simd::max_size<U>);
+  for (int i = 0; i < size; i++) {
+    U* moving_accumulator = accumulator;
+    auto s = stride;
+    while (s >= N) {
+      auto acc = simd::load<U, N>(moving_accumulator);
+      auto v = simd::Simd<U, N>(simd::load<T, N>(x));
+      simd::store<U, N>(moving_accumulator, op(acc, v));
+      moving_accumulator += N;
+      x += N;
+      s -= N;
+    }
+    while (s-- > 0) {
+      *moving_accumulator = op(*moving_accumulator, *x);
+      moving_accumulator++;
+      x++;
+    }
+  }
+};
+
+template <typename T, typename U, typename Op>
+void contiguous_reduce(const T* x, U* accumulator, int size, Op op, U init) {
+  constexpr int N = std::min(simd::max_size<T>, simd::max_size<U>);
+  simd::Simd<U, N> accumulator_v(init);
+  while (size >= N) {
+    accumulator_v = op(accumulator_v, simd::Simd<U, N>(simd::load<T, N>(x)));
+    x += N;
+    size -= N;
+  }
+  *accumulator = op(*accumulator, op(accumulator_v));
+  while (size-- > 0) {
+    *accumulator = op(*accumulator, *x);
+    x++;
+  }
+}
+
+// Helper for the ndimensional strided loop
+void nd_loop(
+    std::function<void(int)> callback,
+    const Shape& shape,
+    const Strides& strides) {
+  std::function<void(int, int)> loop_inner;
+  loop_inner = [&](int dim, int offset) {
+    if (dim < shape.size() - 1) {
+      auto size = shape[dim];
+      auto stride = strides[dim];
+      for (int i = 0; i < size; i++) {
+        loop_inner(dim + 1, offset + i * stride);
+      }
+    } else {
+      auto size = shape[dim];
+      auto stride = strides[dim];
+      for (int i = 0; i < size; i++) {
+        callback(offset + i * stride);
+      }
+    }
+  };
+  loop_inner(0, 0);
+}
+
+template <typename T, typename U, typename Op>
+void reduction_op(
+    const array& x,
+    array& out,
+    const std::vector<int>& axes,
+    U init) {
+  ReductionPlan plan = get_reduction_plan(x, axes);
+
+  auto in_ptr = x.data<T>();
+  auto out_ptr = out.data<U>();
+  if (plan.type == ContiguousAllReduce) {
+    *out_ptr = init;
+    contiguous_reduce(in_ptr, out_ptr, x.size(), Op{}, init);
+    return;
+  }
+
+  if (plan.type == ContiguousReduce && plan.shape.size() == 1) {
+    int reduction_size = plan.shape[0];
+    for (int i = 0; i < out.size(); i++, out_ptr++, in_ptr += reduction_size) {
+      *out_ptr = init;
+      contiguous_reduce(in_ptr, out_ptr, reduction_size, Op{}, init);
+    }
+    return;
+  }
+
+  if (plan.type == GeneralContiguousReduce || plan.type == ContiguousReduce) {
+    int reduction_size = plan.shape.back();
+    plan.shape.pop_back();
+    plan.strides.pop_back();
+    // Unrolling the following loop (and implementing it in order for
+    // ContiguousReduce) should hold extra performance boost.
+    auto [shape, strides] = shapes_without_reduction_axes(x, axes);
+    if (plan.shape.size() == 0) {
+      for (int i = 0; i < out.size(); i++, out_ptr++) {
+        int offset = elem_to_loc(i, shape, strides);
+        *out_ptr = init;
+        contiguous_reduce(in_ptr + offset, out_ptr, reduction_size, Op{}, init);
+      }
+    } else {
+      for (int i = 0; i < out.size(); i++, out_ptr++) {
+        int offset = elem_to_loc(i, shape, strides);
+        *out_ptr = init;
+        nd_loop(
+            [&](int extra_offset) {
+              contiguous_reduce(
+                  in_ptr + offset + extra_offset,
+                  out_ptr,
+                  reduction_size,
+                  Op{},
+                  init);
+            },
+            plan.shape,
+            plan.strides);
+      }
+    }
+    return;
+  }
+
+  if (plan.type == ContiguousStridedReduce && plan.shape.size() == 1) {
+    int reduction_size = plan.shape.back();
+    size_t reduction_stride = plan.strides.back();
+    plan.shape.pop_back();
+    plan.strides.pop_back();
+    for (int i = 0; i < out.size(); i += reduction_stride) {
+      std::fill_n(out_ptr, reduction_stride, init);
+      strided_reduce(in_ptr, out_ptr, reduction_size, reduction_stride, Op{});
+      in_ptr += reduction_stride * reduction_size;
+      out_ptr += reduction_stride;
+    }
+    return;
+  }
+
+  if (plan.type == GeneralStridedReduce ||
+      plan.type == ContiguousStridedReduce) {
+    int reduction_size = plan.shape.back();
+    size_t reduction_stride = plan.strides.back();
+    plan.shape.pop_back();
+    plan.strides.pop_back();
+    auto [shape, strides] = shapes_without_reduction_axes(x, axes);
+
+    if (plan.shape.size() == 0) {
+      for (int i = 0; i < out.size(); i += reduction_stride) {
+        int offset = elem_to_loc(i, shape, strides);
+        std::fill_n(out_ptr, reduction_stride, init);
+        strided_reduce(
+            in_ptr + offset, out_ptr, reduction_size, reduction_stride, Op{});
+        out_ptr += reduction_stride;
+      }
+    } else {
+      for (int i = 0; i < out.size(); i += reduction_stride) {
+        int offset = elem_to_loc(i, shape, strides);
+        std::fill_n(out_ptr, reduction_stride, init);
+        nd_loop(
+            [&](int extra_offset) {
+              strided_reduce(
+                  in_ptr + offset + extra_offset,
+                  out_ptr,
+                  reduction_size,
+                  reduction_stride,
+                  Op{});
+            },
+            plan.shape,
+            plan.strides);
+        out_ptr += reduction_stride;
+      }
+    }
+    return;
+  }
+
+  if (plan.type == GeneralReduce) {
+    auto [shape, strides] = shapes_without_reduction_axes(x, axes);
+
+    for (int i = 0; i < out.size(); i++, out_ptr++) {
+      int offset = elem_to_loc(i, shape, strides);
+      U val = init;
+      nd_loop(
+          [&](int extra_offset) {
+            val = Op{}(val, *(in_ptr + offset + extra_offset));
+          },
+          plan.shape,
+          plan.strides);
+      *out_ptr = val;
+    }
+  }
+}
+
+struct AndReduce {
+  template <typename T>
+  bool operator()(bool x, T y) {
+    return x & (y != 0);
+  }
+
+  bool operator()(bool x, bool y) {
+    return x & y;
+  }
+
+  template <int N, typename T>
+  simd::Simd<bool, N> operator()(simd::Simd<bool, N> y, simd::Simd<T, N> x) {
+    return x & (y != 0);
+  };
+
+  template <int N>
+  simd::Simd<bool, N> operator()(simd::Simd<bool, N> y, simd::Simd<bool, N> x) {
+    return x & y;
+  };
+
+  template <int N, typename T>
+  bool operator()(simd::Simd<T, N> x) {
+    return simd::all(x);
+  };
+};
+
+struct OrReduce {
+  template <typename T>
+  bool operator()(bool x, T y) {
+    return x | (y != 0);
+  }
+
+  bool operator()(bool x, bool y) {
+    return x | y;
+  }
+
+  template <int N, typename T>
+  simd::Simd<bool, N> operator()(simd::Simd<bool, N> y, simd::Simd<T, N> x) {
+    return x | (y != 0);
+  };
+
+  template <int N>
+  simd::Simd<bool, N> operator()(simd::Simd<bool, N> y, simd::Simd<bool, N> x) {
+    return x | y;
+  };
+
+  template <int N, typename T>
+  bool operator()(simd::Simd<T, N> x) {
+    return simd::any(x);
+  };
+};
+
+struct MaxReduce {
+  template <typename T>
+  T operator()(T y, T x) {
+    return (*this)(simd::Simd<T, 1>(x), simd::Simd<T, 1>(y)).value;
+  };
+
+  template <int N, typename T>
+  simd::Simd<T, N> operator()(simd::Simd<T, N> y, simd::Simd<T, N> x) {
+    return simd::maximum(x, y);
+  };
+
+  template <int N, typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(simd::Simd<T, N> x) {
+    return simd::max(x);
+  };
+
+  template <int N, typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(simd::Simd<T, N> x) {
+    if (simd::any(x != x)) {
+      return static_cast<T>(NAN);
+    }
+    return simd::max(x);
+  };
+};
+
+struct MinReduce {
+  template <typename T>
+  T operator()(T y, T x) {
+    return (*this)(simd::Simd<T, 1>(x), simd::Simd<T, 1>(y)).value;
+  };
+
+  template <int N, typename T>
+  simd::Simd<T, N> operator()(simd::Simd<T, N> y, simd::Simd<T, N> x) {
+    return simd::minimum(x, y);
+  };
+
+  template <int N, typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(simd::Simd<T, N> x) {
+    return simd::min(x);
+  };
+
+  template <int N, typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(simd::Simd<T, N> x) {
+    if (simd::any(x != x)) {
+      return static_cast<T>(NAN);
+    }
+    return simd::min(x);
+  };
+};
+
+struct SumReduce {
+  template <typename T, typename U>
+  U operator()(U y, T x) {
+    return x + y;
+  };
+
+  template <int N, typename T, typename U>
+  simd::Simd<U, N> operator()(simd::Simd<U, N> y, simd::Simd<T, N> x) {
+    return y + x;
+  };
+
+  template <int N, typename T>
+  T operator()(simd::Simd<T, N> x) {
+    return simd::sum(x);
+  };
+};
+
+struct ProdReduce {
+  template <typename T, typename U>
+  U operator()(U y, T x) {
+    return x * y;
+  };
+
+  template <int N, typename T, typename U>
+  simd::Simd<U, N> operator()(simd::Simd<U, N> y, simd::Simd<T, N> x) {
+    return x * y;
+  };
+
+  template <int N, typename T>
+  T operator()(simd::Simd<T, N> x) {
+    return simd::prod(x);
+  };
+};
+
+template <typename InT>
+void reduce_dispatch_and_or(
+    const array& in,
+    array& out,
+    Reduce::ReduceType rtype,
+    const std::vector<int>& axes) {
+  if (rtype == Reduce::And) {
+    reduction_op<InT, bool, AndReduce>(in, out, axes, true);
+  } else {
+    reduction_op<InT, bool, OrReduce>(in, out, axes, false);
+  }
+}
+
+template <typename InT>
+void reduce_dispatch_sum_prod(
+    const array& in,
+    array& out,
+    Reduce::ReduceType rtype,
+    const std::vector<int>& axes) {
+  if (rtype == Reduce::Sum) {
+    if constexpr (std::is_integral_v<InT> && sizeof(InT) <= 4) {
+      reduction_op<InT, int32_t, SumReduce>(in, out, axes, 0);
+    } else {
+      reduction_op<InT, InT, SumReduce>(in, out, axes, 0);
+    }
+  } else {
+    if constexpr (std::is_integral_v<InT> && sizeof(InT) <= 4) {
+      reduction_op<InT, int32_t, ProdReduce>(in, out, axes, 1);
+    } else {
+      reduction_op<InT, InT, ProdReduce>(in, out, axes, 1);
+    }
+  }
+}
+
+template <typename InT>
+void reduce_dispatch_min_max(
+    const array& in,
+    array& out,
+    Reduce::ReduceType rtype,
+    const std::vector<int>& axes) {
+  if (rtype == Reduce::Max) {
+    auto init = Limits<InT>::min;
+    reduction_op<InT, InT, MaxReduce>(in, out, axes, init);
+  } else {
+    auto init = Limits<InT>::max;
+    reduction_op<InT, InT, MinReduce>(in, out, axes, init);
+  }
+}
+
+void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    reduce_type_ = reduce_type_,
+                    axes_ = axes_]() mutable {
+    switch (reduce_type_) {
+      case Reduce::And:
+      case Reduce::Or: {
+        switch (in.dtype()) {
+          case bool_:
+          case uint8:
+          case int8:
+            reduce_dispatch_and_or<int8_t>(in, out, reduce_type_, axes_);
+            break;
+          case int16:
+          case uint16:
+          case float16:
+          case bfloat16:
+            reduce_dispatch_and_or<int16_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint32:
+          case int32:
+          case float32:
+            reduce_dispatch_and_or<int32_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint64:
+          case int64:
+          case float64:
+          case complex64:
+            reduce_dispatch_and_or<int64_t>(in, out, reduce_type_, axes_);
+            break;
+        }
+        break;
+      }
+      case Reduce::Sum:
+      case Reduce::Prod: {
+        switch (in.dtype()) {
+          case bool_:
+          case uint8:
+            reduce_dispatch_sum_prod<uint8_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint16:
+            reduce_dispatch_sum_prod<uint16_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint32:
+            reduce_dispatch_sum_prod<uint32_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint64:
+            reduce_dispatch_sum_prod<uint64_t>(in, out, reduce_type_, axes_);
+            break;
+          case int8:
+            reduce_dispatch_sum_prod<int8_t>(in, out, reduce_type_, axes_);
+            break;
+          case int16:
+            reduce_dispatch_sum_prod<int16_t>(in, out, reduce_type_, axes_);
+            break;
+          case int32:
+            reduce_dispatch_sum_prod<int32_t>(in, out, reduce_type_, axes_);
+            break;
+          case int64:
+            reduce_dispatch_sum_prod<int64_t>(in, out, reduce_type_, axes_);
+            break;
+          case float16:
+            reduce_dispatch_sum_prod<float16_t>(in, out, reduce_type_, axes_);
+            break;
+          case bfloat16:
+            reduce_dispatch_sum_prod<bfloat16_t>(in, out, reduce_type_, axes_);
+            break;
+          case float32:
+            reduce_dispatch_sum_prod<float>(in, out, reduce_type_, axes_);
+            break;
+          case float64:
+            reduce_dispatch_sum_prod<double>(in, out, reduce_type_, axes_);
+            break;
+          case complex64:
+            reduce_dispatch_sum_prod<complex64_t>(in, out, reduce_type_, axes_);
+            break;
+        }
+        break;
+      }
+      case Reduce::Max:
+      case Reduce::Min: {
+        switch (in.dtype()) {
+          case bool_:
+            reduce_dispatch_min_max<bool>(in, out, reduce_type_, axes_);
+            break;
+          case uint8:
+            reduce_dispatch_min_max<uint8_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint16:
+            reduce_dispatch_min_max<uint16_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint32:
+            reduce_dispatch_min_max<uint32_t>(in, out, reduce_type_, axes_);
+            break;
+          case uint64:
+            reduce_dispatch_min_max<uint64_t>(in, out, reduce_type_, axes_);
+            break;
+          case int8:
+            reduce_dispatch_min_max<int8_t>(in, out, reduce_type_, axes_);
+            break;
+          case int16:
+            reduce_dispatch_min_max<int16_t>(in, out, reduce_type_, axes_);
+            break;
+          case int32:
+            reduce_dispatch_min_max<int32_t>(in, out, reduce_type_, axes_);
+            break;
+          case int64:
+            reduce_dispatch_min_max<int64_t>(in, out, reduce_type_, axes_);
+            break;
+          case float16:
+            reduce_dispatch_min_max<float16_t>(in, out, reduce_type_, axes_);
+            break;
+          case float32:
+            reduce_dispatch_min_max<float>(in, out, reduce_type_, axes_);
+            break;
+          case float64:
+            reduce_dispatch_min_max<double>(in, out, reduce_type_, axes_);
+            break;
+          case bfloat16:
+            reduce_dispatch_min_max<bfloat16_t>(in, out, reduce_type_, axes_);
+            break;
+          case complex64:
+            reduce_dispatch_min_max<complex64_t>(in, out, reduce_type_, axes_);
+            break;
+        }
+        break;
+      }
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/scan.cpp b/macos/third_party/mlx/mlx/backend/cpu/scan.cpp
new file mode 100644
index 0000000..4dc2f50
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/scan.cpp
@@ -0,0 +1,338 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/binary_ops.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename T, typename U, typename Op>
+void contiguous_scan(
+    const T* input,
+    U* output,
+    int count,
+    int stride,
+    bool reverse,
+    bool inclusive,
+    const Op& op,
+    U init) {
+  if (!reverse) {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        *output = *input;
+        for (int j = 1; j < stride; j++) {
+          input++;
+          output++;
+          *output = op(*(output - 1), *input);
+        }
+        output++;
+        input++;
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        *output = init;
+        for (int j = 1; j < stride; j++) {
+          *(output + 1) = op(*output, *input);
+          input++;
+          output++;
+        }
+        output++;
+        input++;
+      }
+    }
+  } else {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        output += stride - 1;
+        input += stride - 1;
+        *output = *input;
+        for (int j = 1; j < stride; j++) {
+          input--;
+          output--;
+          *output = op(*(output + 1), *input);
+        }
+        output += stride;
+        input += stride;
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        output += stride - 1;
+        input += stride - 1;
+        *output = init;
+        for (int j = 1; j < stride; j++) {
+          *(output - 1) = op(*output, *input);
+          input--;
+          output--;
+        }
+        output += stride;
+        input += stride;
+      }
+    }
+  }
+};
+
+template <typename T, typename U, typename Op>
+void strided_scan(
+    const T* input,
+    U* output,
+    int count,
+    int size,
+    int stride,
+    bool reverse,
+    bool inclusive,
+    const Op& op,
+    U init) {
+  // TODO: Vectorize the following naive implementation
+  if (!reverse) {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        std::copy(input, input + stride, output);
+        output += stride;
+        input += stride;
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            *output = op(*(output - stride), *input);
+            output++;
+            input++;
+          }
+        }
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        std::fill(output, output + stride, init);
+        output += stride;
+        input += stride;
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            *output = op(*(output - stride), *(input - stride));
+            output++;
+            input++;
+          }
+        }
+      }
+    }
+  } else {
+    if (inclusive) {
+      for (int i = 0; i < count; i++) {
+        output += (size - 1) * stride;
+        input += (size - 1) * stride;
+        std::copy(input, input + stride, output);
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            output--;
+            input--;
+            *output = op(*(output + stride), *input);
+          }
+        }
+        output += size * stride;
+        input += size * stride;
+      }
+    } else {
+      for (int i = 0; i < count; i++) {
+        output += (size - 1) * stride;
+        input += (size - 1) * stride;
+        std::fill(output, output + stride, init);
+        for (int j = 1; j < size; j++) {
+          for (int k = 0; k < stride; k++) {
+            output--;
+            input--;
+            *output = op(*(output + stride), *(input + stride));
+          }
+        }
+        output += size * stride;
+        input += size * stride;
+      }
+    }
+  }
+};
+
+template <typename T, typename U, typename Op>
+void scan_op(
+    const array& in,
+    array& out,
+    int axis,
+    bool reverse,
+    bool inclusive,
+    const Op& op,
+    U init) {
+  if (in.flags().row_contiguous) {
+    if (in.strides()[axis] == 1) {
+      contiguous_scan(
+          in.data<T>(),
+          out.data<U>(),
+          in.size() / in.shape(axis),
+          in.shape(axis),
+          reverse,
+          inclusive,
+          op,
+          init);
+    } else {
+      strided_scan(
+          in.data<T>(),
+          out.data<U>(),
+          in.size() / in.shape(axis) / in.strides()[axis],
+          in.shape(axis),
+          in.strides()[axis],
+          reverse,
+          inclusive,
+          op,
+          init);
+    }
+  } else {
+    throw std::runtime_error("Scan op supports only contiguous inputs");
+  }
+}
+
+template <typename T, typename U>
+void scan_dispatch(
+    Scan::ReduceType rtype,
+    const array& in,
+    array& out,
+    int axis,
+    bool reverse,
+    bool inclusive) {
+  switch (rtype) {
+    case Scan::Sum: {
+      auto op = [](U y, T x) { return y + x; };
+      auto init = static_cast<U>(0);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::Prod: {
+      auto op = [](U y, T x) { return y * x; };
+      auto init = static_cast<U>(1);
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::Min: {
+      auto op = [](U y, T x) { return x < y ? x : y; };
+      auto init = (issubdtype(in.dtype(), floating))
+          ? static_cast<U>(std::numeric_limits<float>::infinity())
+          : std::numeric_limits<U>::max();
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::Max: {
+      auto op = [](U y, T x) { return x < y ? y : x; };
+      auto init = (issubdtype(in.dtype(), floating))
+          ? static_cast<U>(-std::numeric_limits<float>::infinity())
+          : std::numeric_limits<U>::min();
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+    case Scan::LogAddExp: {
+      auto op = [](U a, T b) {
+        return detail::LogAddExp{}(a, static_cast<U>(b));
+      };
+      auto init = (issubdtype(in.dtype(), floating))
+          ? static_cast<U>(-std::numeric_limits<float>::infinity())
+          : std::numeric_limits<U>::min();
+      scan_op<T, U>(in, out, axis, reverse, inclusive, op, init);
+      break;
+    }
+  }
+}
+
+} // namespace
+
+void Scan::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  auto& encoder = cpu::get_command_encoder(stream());
+
+  // Ensure contiguity
+  auto in = inputs[0];
+  if (!in.flags().row_contiguous) {
+    in = contiguous_copy_cpu(in, stream());
+    encoder.add_temporary(in);
+  }
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    axis_ = axis_,
+                    reduce_type_ = reduce_type_,
+                    reverse_ = reverse_,
+                    inclusive_ = inclusive_]() mutable {
+    switch (in.dtype()) {
+      case bool_: {
+        // We could do a full dtype x dtype switch but this is the only case
+        // where we accumulate in a different type, for now.
+        //
+        // TODO: If we add the option to accumulate floats in higher precision
+        //       floats perhaps we should add the full all-to-all dispatch.
+        if (reduce_type_ == Scan::Sum && out.dtype() == int32) {
+          scan_dispatch<bool, int32_t>(
+              reduce_type_, in, out, axis_, reverse_, inclusive_);
+        } else {
+          scan_dispatch<bool, bool>(
+              reduce_type_, in, out, axis_, reverse_, inclusive_);
+        }
+        break;
+      }
+      case uint8:
+        scan_dispatch<uint8_t, uint8_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case uint16:
+        scan_dispatch<uint16_t, uint16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case uint32:
+        scan_dispatch<uint32_t, uint32_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case uint64:
+        scan_dispatch<uint64_t, uint64_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int8:
+        scan_dispatch<int8_t, int8_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int16:
+        scan_dispatch<int16_t, int16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int32:
+        scan_dispatch<int32_t, int32_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case int64:
+        scan_dispatch<int64_t, int64_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case float16:
+        scan_dispatch<float16_t, float16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case float32:
+        scan_dispatch<float, float>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case float64:
+        scan_dispatch<double, double>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case bfloat16:
+        scan_dispatch<bfloat16_t, bfloat16_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+      case complex64:
+        scan_dispatch<complex64_t, complex64_t>(
+            reduce_type_, in, out, axis_, reverse_, inclusive_);
+        break;
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/select.cpp b/macos/third_party/mlx/mlx/backend/cpu/select.cpp
new file mode 100644
index 0000000..bf6a9b8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/select.cpp
@@ -0,0 +1,95 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/cpu/binary_ops.h"
+#include "mlx/backend/cpu/ternary.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename Op>
+void select_op(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op,
+    Stream stream) {
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+  set_ternary_op_output_data(a, b, c, out, topt);
+
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(c);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    b = array::unsafe_weak_copy(b),
+                    c = array::unsafe_weak_copy(c),
+                    out = array::unsafe_weak_copy(out),
+                    op,
+                    topt]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        ternary_op<bool, bool, bool, bool>(a, b, c, out, op, topt);
+        break;
+      case uint8:
+        ternary_op<bool, uint8_t, uint8_t, uint8_t>(a, b, c, out, op, topt);
+        break;
+      case uint16:
+        ternary_op<bool, uint16_t, uint16_t, uint16_t>(a, b, c, out, op, topt);
+        break;
+      case uint32:
+        ternary_op<bool, uint32_t, uint32_t, uint32_t>(a, b, c, out, op, topt);
+        break;
+      case uint64:
+        ternary_op<bool, uint64_t, uint64_t, uint64_t>(a, b, c, out, op, topt);
+        break;
+      case int8:
+        ternary_op<bool, int8_t, int8_t, int8_t>(a, b, c, out, op, topt);
+        break;
+      case int16:
+        ternary_op<bool, int16_t, int16_t, int16_t>(a, b, c, out, op, topt);
+        break;
+      case int32:
+        ternary_op<bool, int32_t, int32_t, int32_t>(a, b, c, out, op, topt);
+        break;
+      case int64:
+        ternary_op<bool, int64_t, int64_t, int64_t>(a, b, c, out, op, topt);
+        break;
+      case float16:
+        ternary_op<bool, float16_t, float16_t, float16_t>(
+            a, b, c, out, op, topt);
+        break;
+      case float32:
+        ternary_op<bool, float, float, float>(a, b, c, out, op, topt);
+        break;
+      case float64:
+        ternary_op<bool, double, double, double>(a, b, c, out, op, topt);
+        break;
+      case bfloat16:
+        ternary_op<bool, bfloat16_t, bfloat16_t, bfloat16_t>(
+            a, b, c, out, op, topt);
+        break;
+      case complex64:
+        ternary_op<bool, complex64_t, complex64_t, complex64_t>(
+            a, b, c, out, op, topt);
+        break;
+    }
+  });
+}
+
+} // namespace
+
+void Select::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 3);
+  const auto& condition = inputs[0];
+  const auto& a = inputs[1];
+  const auto& b = inputs[2];
+  select_op(condition, a, b, out, detail::Select(), stream());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/simd/accelerate_fp16_simd.h b/macos/third_party/mlx/mlx/backend/cpu/simd/accelerate_fp16_simd.h
new file mode 100644
index 0000000..9505448
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/simd/accelerate_fp16_simd.h
@@ -0,0 +1,56 @@
+#pragma once
+
+#include "mlx/backend/cpu/simd/base_simd.h"
+
+#if MLX_SIMD_LIBRARY_VERSION < 6
+#include "mlx/backend/cpu/simd/neon_fp16_simd.h"
+#endif
+
+namespace mlx::core::simd {
+
+#if MLX_SIMD_LIBRARY_VERSION >= 6
+constexpr int N = 8;
+template <int N>
+struct ScalarT<float16_t, N> {
+  using v = _Float16;
+};
+#endif
+
+template <>
+inline constexpr int max_size<float16_t> = N;
+
+#define SIMD_FP16_DEFAULT_UNARY(op)                    \
+  template <>                                          \
+  inline Simd<float16_t, N> op(Simd<float16_t, N> v) { \
+    Simd<float, N> in = v;                             \
+    return op(in);                                     \
+  }
+
+SIMD_FP16_DEFAULT_UNARY(acos)
+SIMD_FP16_DEFAULT_UNARY(acosh)
+SIMD_FP16_DEFAULT_UNARY(asin)
+SIMD_FP16_DEFAULT_UNARY(asinh)
+SIMD_FP16_DEFAULT_UNARY(atan)
+SIMD_FP16_DEFAULT_UNARY(atanh)
+SIMD_FP16_DEFAULT_UNARY(cosh)
+SIMD_FP16_DEFAULT_UNARY(expm1)
+SIMD_FP16_DEFAULT_UNARY(log)
+SIMD_FP16_DEFAULT_UNARY(log2)
+SIMD_FP16_DEFAULT_UNARY(log10)
+SIMD_FP16_DEFAULT_UNARY(log1p)
+SIMD_FP16_DEFAULT_UNARY(sinh)
+SIMD_FP16_DEFAULT_UNARY(tan)
+SIMD_FP16_DEFAULT_UNARY(tanh)
+
+#define SIMD_FP16_DEFAULT_BINARY(op)                                         \
+  template <>                                                                \
+  inline Simd<float16_t, N> op(Simd<float16_t, N> x, Simd<float16_t, N> y) { \
+    Simd<float, N> a = x;                                                    \
+    Simd<float, N> b = y;                                                    \
+    return op(a, b);                                                         \
+  }
+SIMD_FP16_DEFAULT_BINARY(atan2)
+SIMD_FP16_DEFAULT_BINARY(remainder)
+SIMD_FP16_DEFAULT_BINARY(pow)
+
+} // namespace mlx::core::simd
diff --git a/macos/third_party/mlx/mlx/backend/cpu/simd/accelerate_simd.h b/macos/third_party/mlx/mlx/backend/cpu/simd/accelerate_simd.h
new file mode 100644
index 0000000..f62c67d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/simd/accelerate_simd.h
@@ -0,0 +1,329 @@
+#pragma once
+
+#include <arm_neon.h>
+#include <simd/math.h>
+#include <simd/vector.h>
+
+#include <stdint.h>
+#include <cmath>
+#include <complex>
+
+#include "mlx/backend/cpu/simd/base_simd.h"
+
+// There seems to be a bug in simd/base_simd.h
+// __XROS_2_0 is not defined, the expression evaluates
+// to true instead of false setting the SIMD library
+// higher than it should be even on macOS < 15
+#if __MAC_OS_X_VERSION_MIN_REQUIRED >= 150000 ||  \
+    __IPHONE_OS_VERSION_MIN_REQUIRED >= 180000 || \
+    __WATCH_OS_VERSION_MIN_REQUIRED >= 110000 ||  \
+    __WATCH_OS_VERSION_MIN_REQUIRED >= 110000 ||  \
+    __TV_OS_VERSION_MIN_REQUIRED >= 180000
+#define MLX_SIMD_LIBRARY_VERSION 6
+#else
+#define MLX_SIMD_LIBRARY_VERSION 5
+#endif
+
+namespace mlx::core::simd {
+
+// Apple simd namespace
+namespace asd = ::simd;
+
+// This indirection is needed to remap certain types to ones that accelerate
+// SIMD can handle
+template <typename T, int N>
+struct ScalarT {
+  using v = T;
+};
+template <int N>
+struct ScalarT<bool, N> {
+  using v = char;
+};
+template <int N>
+struct ScalarT<int8_t, N> {
+  using v = char;
+};
+template <int N>
+struct ScalarT<uint64_t, N> {
+  using v = unsigned long;
+};
+template <int N>
+struct ScalarT<int64_t, N> {
+  using v = long;
+};
+
+template <typename T, int N>
+struct Simd {
+  static constexpr int size = N;
+  using scalar_t = typename ScalarT<T, N>::v;
+
+  Simd<T, N>() {}
+
+  template <typename U>
+  Simd<T, N>(Simd<U, N> other) : value(asd::convert<scalar_t>(other.value)) {}
+
+  template <typename U>
+  Simd<T, N>(U v) : value(v){};
+
+  Simd<T, N>(Simd<T, N / 2> x, Simd<T, N / 2> y) {
+    value = asd::make<typename asd::Vector<scalar_t, N>::packed_t>(
+        x.value, y.value);
+  };
+
+  T operator[](int idx) const {
+    return reinterpret_cast<const T*>(&value)[idx];
+  }
+
+  T& operator[](int idx) {
+    return reinterpret_cast<T*>(&value)[idx];
+  }
+
+  typename asd::Vector<scalar_t, N>::packed_t value;
+};
+
+// Values chosen based on benchmarks on M3 Max
+// TODO: consider choosing these more optimally
+template <>
+inline constexpr int max_size<int8_t> = 16;
+template <>
+inline constexpr int max_size<int16_t> = 16;
+template <>
+inline constexpr int max_size<int> = 8;
+template <>
+inline constexpr int max_size<int64_t> = 4;
+template <>
+inline constexpr int max_size<uint8_t> = 16;
+template <>
+inline constexpr int max_size<uint16_t> = 16;
+template <>
+inline constexpr int max_size<uint32_t> = 8;
+template <>
+inline constexpr int max_size<uint64_t> = 4;
+template <>
+inline constexpr int max_size<float> = 8;
+template <>
+inline constexpr int max_size<double> = 4;
+
+#define SIMD_DEFAULT_UNARY(name, op) \
+  template <typename T, int N>       \
+  Simd<T, N> name(Simd<T, N> v) {    \
+    return op(v.value);              \
+  }
+
+SIMD_DEFAULT_UNARY(abs, asd::abs)
+SIMD_DEFAULT_UNARY(floor, asd::floor)
+SIMD_DEFAULT_UNARY(acos, asd::acos)
+SIMD_DEFAULT_UNARY(acosh, asd::acosh)
+SIMD_DEFAULT_UNARY(asin, asd::asin)
+SIMD_DEFAULT_UNARY(asinh, asd::asinh)
+SIMD_DEFAULT_UNARY(atan, asd::atan)
+SIMD_DEFAULT_UNARY(atanh, asd::atanh)
+SIMD_DEFAULT_UNARY(ceil, asd::ceil)
+SIMD_DEFAULT_UNARY(cosh, asd::cosh)
+SIMD_DEFAULT_UNARY(expm1, asd::expm1)
+SIMD_DEFAULT_UNARY(log, asd::log)
+SIMD_DEFAULT_UNARY(log2, asd::log2)
+SIMD_DEFAULT_UNARY(log10, asd::log10)
+SIMD_DEFAULT_UNARY(log1p, asd::log1p)
+SIMD_DEFAULT_UNARY(rint, asd::rint)
+SIMD_DEFAULT_UNARY(sinh, asd::sinh)
+SIMD_DEFAULT_UNARY(sqrt, asd::sqrt)
+SIMD_DEFAULT_UNARY(rsqrt, asd::rsqrt)
+SIMD_DEFAULT_UNARY(recip, asd::recip)
+SIMD_DEFAULT_UNARY(tan, asd::tan)
+SIMD_DEFAULT_UNARY(tanh, asd::tanh)
+
+template <typename T, int N>
+Simd<T, N> operator-(Simd<T, N> v) {
+  return -v.value;
+}
+
+template <typename T, int N>
+Simd<T, N> operator~(Simd<T, N> v) {
+  return ~v.value;
+}
+
+template <typename T, int N>
+Simd<bool, N> isnan(Simd<T, N> v) {
+  return asd::convert<char>(v.value != v.value);
+}
+
+// No simd_boolN in accelerate, use int8_t instead
+template <typename T, int N>
+Simd<bool, N> operator!(Simd<T, N> v) {
+  return asd::convert<char>(!v.value);
+}
+
+#define SIMD_DEFAULT_BINARY(OP)                                              \
+  template <typename T, typename U, int N>                                   \
+  Simd<T, N> operator OP(Simd<T, N> x, U y) {                                \
+    return asd::convert<typename Simd<T, N>::scalar_t>(x.value OP y);        \
+  }                                                                          \
+  template <typename T1, typename T2, int N>                                 \
+  Simd<T2, N> operator OP(T1 x, Simd<T2, N> y) {                             \
+    return asd::convert<typename Simd<T2, N>::scalar_t>(x OP y.value);       \
+  }                                                                          \
+  template <typename T1, typename T2, int N>                                 \
+  Simd<T1, N> operator OP(Simd<T1, N> x, Simd<T2, N> y) {                    \
+    return asd::convert<typename Simd<T1, N>::scalar_t>(x.value OP y.value); \
+  }
+
+SIMD_DEFAULT_BINARY(+)
+SIMD_DEFAULT_BINARY(-)
+SIMD_DEFAULT_BINARY(/)
+SIMD_DEFAULT_BINARY(*)
+SIMD_DEFAULT_BINARY(<<)
+SIMD_DEFAULT_BINARY(>>)
+SIMD_DEFAULT_BINARY(|)
+SIMD_DEFAULT_BINARY(^)
+SIMD_DEFAULT_BINARY(&)
+SIMD_DEFAULT_BINARY(&&)
+SIMD_DEFAULT_BINARY(||)
+
+#define SIMD_DEFAULT_COMPARISONS(OP)                        \
+  template <int N, typename T, typename U>                  \
+  Simd<bool, N> operator OP(Simd<T, N> a, U b) {            \
+    return asd::convert<char>(a.value OP b);                \
+  }                                                         \
+  template <int N, typename T, typename U>                  \
+  Simd<bool, N> operator OP(T a, Simd<U, N> b) {            \
+    return asd::convert<char>(a OP b.value);                \
+  }                                                         \
+  template <int N, typename T1, typename T2>                \
+  Simd<bool, N> operator OP(Simd<T1, N> a, Simd<T2, N> b) { \
+    return asd::convert<char>(a.value OP b.value);          \
+  }
+
+SIMD_DEFAULT_COMPARISONS(>)
+SIMD_DEFAULT_COMPARISONS(<)
+SIMD_DEFAULT_COMPARISONS(>=)
+SIMD_DEFAULT_COMPARISONS(<=)
+SIMD_DEFAULT_COMPARISONS(==)
+SIMD_DEFAULT_COMPARISONS(!=)
+
+template <typename T, int N>
+Simd<T, N> clz(Simd<T, N> x) {
+  auto a = *(uint32x4_t*)(&x);
+  auto b = *((uint32x4_t*)(&x) + 1);
+  a = vclzq_u32(a);
+  b = vclzq_u32(b);
+  return asd::make_uint8(a, b);
+}
+
+template <typename T, int N>
+Simd<T, N> atan2(Simd<T, N> a, Simd<T, N> b) {
+  return asd::atan2(a.value, b.value);
+}
+
+template <typename T, int N>
+Simd<T, N> maximum(Simd<T, N> a, Simd<T, N> b) {
+  auto out = Simd<T, N>(asd::max(a.value, b.value));
+  if constexpr (!std::is_integral_v<T>) {
+    out = select(isnan(b), b, select(isnan(a), a, out));
+  }
+  return out;
+}
+
+template <typename T, int N>
+Simd<T, N> minimum(Simd<T, N> a, Simd<T, N> b) {
+  auto out = Simd<T, N>(asd::min(a.value, b.value));
+  if constexpr (!std::is_integral_v<T>) {
+    out = select(isnan(b), b, select(isnan(a), a, out));
+  }
+  return out;
+}
+
+template <typename T, int N>
+Simd<T, N> remainder(Simd<T, N> a, Simd<T, N> b) {
+  Simd<T, N> r;
+  if constexpr (!std::is_integral_v<T>) {
+    r = asd::remainder(a.value, b.value);
+  } else {
+    r = a - b * (a / b);
+  }
+  if constexpr (std::is_signed_v<T>) {
+    auto mask = r != 0 && (r < 0 != b < 0);
+    r = select(mask, r + b, r);
+  }
+  return r;
+}
+
+template <typename MaskT, typename T1, typename T2, int N>
+Simd<T1, N> select(Simd<MaskT, N> mask, Simd<T1, N> x, Simd<T2, N> y) {
+  static_assert(std::is_same_v<MaskT, bool>);
+  if constexpr (sizeof(T1) == 1) {
+    return asd::bitselect(y.value, x.value, asd::convert<char>(mask.value));
+  } else if constexpr (sizeof(T1) == 2) {
+    return asd::bitselect(y.value, x.value, asd::convert<short>(mask.value));
+  } else if constexpr (sizeof(T1) == 4) {
+    return asd::bitselect(y.value, x.value, asd::convert<int>(mask.value));
+  } else {
+    return asd::bitselect(y.value, x.value, asd::convert<long>(mask.value));
+  }
+}
+
+template <typename T, int N>
+Simd<T, N> pow(Simd<T, N> base, Simd<T, N> exp) {
+  if constexpr (!std::is_integral_v<T>) {
+    return asd::pow(base.value, exp.value);
+  } else {
+    Simd<T, N> res = 1;
+    // Raising an integer to a negative power is undefined
+    if (any(exp < 0)) {
+      return 0;
+    }
+    while (any(exp > 0)) {
+      res = select((exp & 1) != 0, res * base, res);
+      base = select(exp > 0, base * base, base);
+      exp = exp >> 1;
+    }
+    return res;
+  }
+}
+
+template <typename T, int N>
+Simd<T, N> clamp(Simd<T, N> v, Simd<T, N> min, Simd<T, N> max) {
+  return asd::clamp(v.value, min.value, max.value);
+}
+
+template <typename T, typename U, int N>
+Simd<T, N> fma(Simd<T, N> x, Simd<T, N> y, U z) {
+  return asd::muladd(x.value, y.value, Simd<T, N>(z).value);
+}
+
+// Reductions
+
+template <typename T, int N>
+bool all(Simd<T, N> x) {
+  return asd::all(x.value);
+}
+template <typename T, int N>
+bool any(Simd<T, N> x) {
+  return asd::any(x.value);
+}
+template <typename T, int N>
+T sum(Simd<T, N> x) {
+  return asd::reduce_add(x.value);
+}
+template <typename T, int N>
+T max(Simd<T, N> x) {
+  return asd::reduce_max(x.value);
+}
+template <typename T, int N>
+T min(Simd<T, N> x) {
+  return asd::reduce_min(x.value);
+}
+
+template <typename T, int N>
+T prod(Simd<T, N> x) {
+  auto ptr = (T*)&x;
+  auto lhs = load<T, N / 2>(ptr);
+  auto rhs = load<T, N / 2>(ptr + N / 2);
+  return prod(lhs * rhs);
+}
+
+} // namespace mlx::core::simd
+
+#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+#include "mlx/backend/cpu/simd/accelerate_fp16_simd.h"
+#endif
diff --git a/macos/third_party/mlx/mlx/backend/cpu/simd/base_simd.h b/macos/third_party/mlx/mlx/backend/cpu/simd/base_simd.h
new file mode 100644
index 0000000..775f5df
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/simd/base_simd.h
@@ -0,0 +1,319 @@
+#pragma once
+
+// Required for using M_LN2 in MSVC.
+#define _USE_MATH_DEFINES
+
+#include <math.h>
+#include <stdint.h>
+#include <algorithm>
+#include <complex>
+#include <functional>
+
+#ifdef _MSC_VER
+#include <intrin.h> // For _BitScanReverse
+#endif
+
+namespace mlx::core::simd {
+template <typename T, int N>
+struct Simd;
+
+template <typename T>
+static constexpr int max_size = 1;
+
+template <typename T>
+struct Simd<T, 1> {
+  static constexpr int size = 1;
+  T value;
+  Simd() {}
+  template <typename U>
+  Simd(Simd<U, 1> v) : value(v.value) {}
+  template <typename U>
+  Simd(U v) : value(v) {}
+
+  T operator[](int) const {
+    return value;
+  }
+
+  T& operator[](int) {
+    return value;
+  }
+};
+
+template <typename T, int N>
+Simd<T, N> load(const T* x) {
+  return *(Simd<T, N>*)x;
+}
+
+template <typename T, int N>
+void store(T* dst, Simd<T, N> x) {
+  // Maintain invariant that bool is either 0 or 1 as
+  // simd comparison ops set all bits in the result to 1
+  if constexpr (std::is_same_v<T, bool> && N > 1) {
+    x = x & 1;
+  }
+  *(Simd<T, N>*)dst = x;
+}
+
+template <typename, typename = void>
+constexpr bool is_complex = false;
+
+template <typename T>
+constexpr bool is_complex<T, std::void_t<decltype(std::declval<T>().real())>> =
+    true;
+
+template <typename T>
+Simd<T, 1> rint(Simd<T, 1> in) {
+  if constexpr (is_complex<T>) {
+    return Simd<T, 1>{
+        T{std::rint(in.value.real()), std::rint(in.value.imag())}};
+  } else {
+    return Simd<T, 1>{std::rint(in.value)};
+  }
+}
+
+template <typename T>
+Simd<T, 1> rsqrt(Simd<T, 1> in) {
+  return T(1.0) / sqrt(in);
+}
+
+template <typename T>
+Simd<T, 1> recip(Simd<T, 1> in) {
+  return T(1.0) / in;
+}
+
+#define DEFAULT_UNARY(name, op)    \
+  template <typename T>            \
+  Simd<T, 1> name(Simd<T, 1> in) { \
+    return op(in.value);           \
+  }
+
+DEFAULT_UNARY(operator-, std::negate{})
+DEFAULT_UNARY(operator!, std::logical_not{})
+DEFAULT_UNARY(abs, std::abs)
+DEFAULT_UNARY(acos, std::acos)
+DEFAULT_UNARY(acosh, std::acosh)
+DEFAULT_UNARY(asin, std::asin)
+DEFAULT_UNARY(asinh, std::asinh)
+DEFAULT_UNARY(atan, std::atan)
+DEFAULT_UNARY(atanh, std::atanh)
+DEFAULT_UNARY(ceil, std::ceil)
+DEFAULT_UNARY(conj, std::conj)
+DEFAULT_UNARY(cosh, std::cosh)
+DEFAULT_UNARY(expm1, std::expm1)
+DEFAULT_UNARY(floor, std::floor)
+DEFAULT_UNARY(log, std::log)
+DEFAULT_UNARY(log10, std::log10)
+DEFAULT_UNARY(sinh, std::sinh)
+DEFAULT_UNARY(sqrt, std::sqrt)
+DEFAULT_UNARY(tan, std::tan)
+DEFAULT_UNARY(tanh, std::tanh)
+
+template <typename T>
+Simd<T, 1> log1p(Simd<T, 1> in) {
+  if constexpr (is_complex<T>) {
+    auto x = in.value.real();
+    auto y = in.value.imag();
+    auto zabs = std::abs(in.value);
+    auto theta = std::atan2(y, x + 1);
+    if (zabs < 0.5) {
+      auto r = x * (2 + x) + y * y;
+      if (r == 0) { // handle underflow
+        return Simd<T, 1>{T{x, theta}};
+      }
+      return Simd<T, 1>{T{((decltype(x))(0.5)) * std::log1p(r), theta}};
+    } else {
+      auto z0 = std::hypot(x + 1, y);
+      return Simd<T, 1>{T{std::log(z0), theta}};
+    }
+  } else {
+    return Simd<T, 1>{std::log1p(in.value)};
+  }
+}
+
+template <typename T>
+Simd<T, 1> log2(Simd<T, 1> in) {
+  if constexpr (is_complex<T>) {
+    auto out = std::log(in.value);
+    auto scale = decltype(out.real())(M_LN2);
+    return Simd<T, 1>{T{out.real() / scale, out.imag() / scale}};
+  } else {
+    return Simd<T, 1>{std::log2(in.value)};
+  }
+}
+
+template <typename T>
+Simd<T, 1> operator~(Simd<T, 1> in) {
+  return ~in.value;
+}
+
+template <typename T>
+auto real(Simd<T, 1> in) -> Simd<decltype(std::real(in.value)), 1> {
+  return std::real(in.value);
+}
+template <typename T>
+auto imag(Simd<T, 1> in) -> Simd<decltype(std::imag(in.value)), 1> {
+  return std::imag(in.value);
+}
+template <typename T>
+Simd<bool, 1> isnan(Simd<T, 1> in) {
+  return std::isnan(in.value);
+}
+
+#define DEFAULT_BINARY(OP)                                                 \
+  template <typename T1, typename T2>                                      \
+  auto operator OP(Simd<T1, 1> a, Simd<T2, 1> b)                           \
+      ->Simd<decltype(a.value OP b.value), 1> {                            \
+    return a.value OP b.value;                                             \
+  }                                                                        \
+  template <typename T1, typename T2>                                      \
+  auto operator OP(T1 a, Simd<T2, 1> b)->Simd<decltype(a OP b.value), 1> { \
+    return a OP b.value;                                                   \
+  }                                                                        \
+  template <typename T1, typename T2>                                      \
+  auto operator OP(Simd<T1, 1> a, T2 b)->Simd<decltype(a.value OP b), 1> { \
+    return a.value OP b;                                                   \
+  }
+
+DEFAULT_BINARY(+)
+DEFAULT_BINARY(-)
+DEFAULT_BINARY(*)
+DEFAULT_BINARY(/)
+DEFAULT_BINARY(<<)
+DEFAULT_BINARY(>>)
+DEFAULT_BINARY(|)
+DEFAULT_BINARY(^)
+DEFAULT_BINARY(&)
+DEFAULT_BINARY(&&)
+DEFAULT_BINARY(||)
+
+template <typename T>
+Simd<T, 1> clz(Simd<T, 1> x_) {
+#ifdef _MSC_VER
+  // MSVC doesn't have __builtin_clz, use _BitScanReverse instead
+  unsigned long index;
+  if (_BitScanReverse(&index, static_cast<unsigned long>(x_.value))) {
+    return static_cast<T>(31 - index);
+  }
+  return static_cast<T>(32); // All zeros case
+#else
+  return __builtin_clz(x_.value);
+#endif
+}
+
+template <typename T>
+Simd<T, 1> remainder(Simd<T, 1> a_, Simd<T, 1> b_) {
+  T a = a_.value;
+  T b = b_.value;
+  T r;
+  if constexpr (std::is_integral_v<T>) {
+    r = a % b;
+  } else {
+    r = std::remainder(a, b);
+  }
+  if constexpr (std::is_signed_v<T>) {
+    if (r != 0 && (r < 0 != b < 0)) {
+      r += b;
+    }
+  }
+  return r;
+}
+
+template <typename T>
+Simd<T, 1> maximum(Simd<T, 1> a_, Simd<T, 1> b_) {
+  T a = a_.value;
+  T b = b_.value;
+  if constexpr (!std::is_integral_v<T>) {
+    if (std::isnan(a)) {
+      return a;
+    }
+  }
+  return (a > b) ? a : b;
+}
+
+template <typename T>
+Simd<T, 1> minimum(Simd<T, 1> a_, Simd<T, 1> b_) {
+  T a = a_.value;
+  T b = b_.value;
+  if constexpr (!std::is_integral_v<T>) {
+    if (std::isnan(a)) {
+      return a;
+    }
+  }
+  return (a < b) ? a : b;
+}
+
+template <typename T>
+Simd<T, 1> pow(Simd<T, 1> a, Simd<T, 1> b) {
+  T base = a.value;
+  T exp = b.value;
+  if constexpr (!std::is_integral_v<T>) {
+    return std::pow(base, exp);
+  } else {
+    T res = 1;
+    while (exp) {
+      if (exp & 1) {
+        res *= base;
+      }
+      exp >>= 1;
+      base *= base;
+    }
+    return res;
+  }
+}
+
+template <typename T>
+Simd<T, 1> atan2(Simd<T, 1> a, Simd<T, 1> b) {
+  return std::atan2(a.value, b.value);
+}
+
+#define DEFAULT_COMPARISONS(OP)                             \
+  template <typename T1, typename T2>                       \
+  Simd<bool, 1> operator OP(Simd<T1, 1> a, Simd<T2, 1> b) { \
+    return a.value OP b.value;                              \
+  }                                                         \
+  template <typename T1, typename T2>                       \
+  Simd<bool, 1> operator OP(T1 a, Simd<T2, 1> b) {          \
+    return a OP b.value;                                    \
+  }                                                         \
+  template <typename T1, typename T2>                       \
+  Simd<bool, 1> operator OP(Simd<T1, 1> a, T2 b) {          \
+    return a.value OP b;                                    \
+  }
+
+DEFAULT_COMPARISONS(>)
+DEFAULT_COMPARISONS(<)
+DEFAULT_COMPARISONS(>=)
+DEFAULT_COMPARISONS(<=)
+DEFAULT_COMPARISONS(==)
+DEFAULT_COMPARISONS(!=)
+
+template <typename MaskT, typename T>
+Simd<T, 1> select(Simd<MaskT, 1> mask, Simd<T, 1> x, Simd<T, 1> y) {
+  return mask.value ? x.value : y.value;
+}
+
+template <typename T>
+Simd<T, 1> clamp(Simd<T, 1> v, Simd<T, 1> min, Simd<T, 1> max) {
+  return std::clamp(v.value, min.value, max.value);
+}
+
+template <typename T, typename U>
+Simd<T, 1> fma(Simd<T, 1> x, Simd<T, 1> y, U z) {
+  return std::fma(x.value, y.value, Simd<T, 1>(z).value);
+}
+
+// Reductions
+#define DEFAULT_REDUCTION(name, type) \
+  template <typename T>               \
+  type name(Simd<T, 1> x) {           \
+    return x.value;                   \
+  }
+
+DEFAULT_REDUCTION(max, T)
+DEFAULT_REDUCTION(min, T)
+DEFAULT_REDUCTION(sum, T)
+DEFAULT_REDUCTION(prod, T)
+DEFAULT_REDUCTION(any, bool)
+DEFAULT_REDUCTION(all, bool)
+
+} // namespace mlx::core::simd
diff --git a/macos/third_party/mlx/mlx/backend/cpu/simd/math.h b/macos/third_party/mlx/mlx/backend/cpu/simd/math.h
new file mode 100644
index 0000000..f9fc831
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/simd/math.h
@@ -0,0 +1,193 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cpu/simd/type.h"
+
+namespace mlx::core::simd {
+
+constexpr float inf = std::numeric_limits<float>::infinity();
+
+/**
+ * Compute exp(x) in an optimizer friendly way as follows:
+ *
+ * First change the problem to computing 2**y where y = x / ln(2).
+ *
+ * Now we will compute 2**y as 2**y1 * 2**y2 where y1 is the integer part
+ * `ipart` and y2 is fractional part. For the integer part we perform bit
+ * shifting and for the fractional part we use a polynomial approximation.
+ *
+ * The algorithm and constants of the polynomial taken from
+ * https://github.com/akohlmey/fastermath/blob/master/src/exp.c which took them
+ * from Cephes math library.
+ *
+ * Note: The implementation below is a general fast exp. There could be faster
+ *       implementations for numbers strictly < 0.
+ */
+template <typename T, int N>
+Simd<T, N> exp(Simd<T, N> in) {
+  if constexpr (is_complex<T>) {
+    return Simd<T, 1>{std::exp(in.value)};
+  } else {
+    Simd<float, N> x_init = in;
+    auto x = x_init * 1.442695f; // multiply with log_2(e)
+    Simd<float, N> ipart, fpart;
+    ipart = floor(x + 0.5);
+    fpart = x - ipart;
+
+    x = 1.535336188319500e-4f;
+    x = fma(x, fpart, 1.339887440266574e-3f);
+    x = fma(x, fpart, 9.618437357674640e-3f);
+    x = fma(x, fpart, 5.550332471162809e-2f);
+    x = fma(x, fpart, 2.402264791363012e-1f);
+    x = fma(x, fpart, 6.931472028550421e-1f);
+    x = fma(x, fpart, 1.000000000000000f);
+
+    // generate 2**ipart in the floating point representation using integer
+    // bitshifting
+    Simd<int, N> epart = (Simd<int, N>(ipart) + 127) << 23;
+
+    // Deal with NaN and Inf
+    auto result = select(isnan(x_init), x_init, (*(Simd<float, N>*)&epart) * x);
+    result = select(x_init > 88.0f, Simd<float, N>(inf), result);
+    result = select(x_init < -88.0f, Simd<float, N>(0), result);
+    return Simd<T, N>(result);
+  }
+}
+
+/* Implementation from:
+ * https://github.com/JishinMaster/simd_utils/blob/3c1433a86fb38edcc9b02039f3c9a65b16640976/neon_mathfun.h#L357
+ * which originally came from the Cephes math library.
+ */
+template <bool Sine, typename T, int N>
+Simd<T, N> sincos(Simd<T, N> in) {
+  auto sign_mask_sin = in < 0;
+  in = abs(in);
+  Simd<float, N> x = in;
+
+  // scale by 4/Pi
+  auto y = x * 1.27323954473516f;
+
+  // store the integer part of y in mm0
+  Simd<uint32_t, N> emm2 = y;
+
+  // j=(j+1) & (~1) (see the cephes sources)
+  emm2 = emm2 + 1;
+  emm2 = emm2 & ~1;
+
+  y = emm2;
+
+  // Get the polynom selection mask. There is one polynom for 0 <= x <= Pi/4
+  // and another one for Pi/4<x<=Pi/2. Both branches will be computed.
+  auto poly_mask = (emm2 & 2) != 0;
+
+  // The magic pass: "Extended precision modular arithmetic"
+  // x = ((x - y * DP1) - y * DP2) - y * DP3
+  x = fma(y, Simd<float, N>(-0.78515625f), x);
+  x = fma(y, Simd<float, N>(-2.4187564849853515625e-4f), x);
+  x = fma(y, Simd<float, N>(-3.77489497744594108e-8f), x);
+
+  sign_mask_sin = sign_mask_sin ^ ((emm2 & 4) != 0);
+  auto sign_mask_cos = ((emm2 - 2) & 4) != 0;
+
+  // Evaluate the first polynom  (0 <= x <= Pi/4) in y1,
+  // and the second polynom      (Pi/4 <= x <= 0) in y2
+  auto z = x * x;
+
+  auto y1 =
+      fma(z, Simd<float, N>(2.443315711809948e-5f), -1.388731625493765e-3f);
+  auto y2 = fma(z, Simd<float, N>(-1.9515295891e-4f), 8.3321608736e-3f);
+  y1 = fma(y1, z, 4.166664568298827e-2f);
+  y2 = fma(y2, z, -1.6666654611e-1f);
+  y1 = y1 * z;
+  y2 = y2 * z;
+  y1 = y1 * z;
+  y2 = fma(x, y2, x);
+  y1 = fma(z, Simd<float, N>(-0.5f), y1);
+  y1 = y1 + 1.0f;
+
+  if constexpr (Sine) {
+    auto ys = select(poly_mask, y1, y2);
+    return select(sign_mask_sin, -ys, ys);
+  } else {
+    auto yc = select(poly_mask, y2, y1);
+    return select(sign_mask_cos, yc, -yc);
+  }
+}
+
+template <typename T, int N>
+Simd<T, N> sin(Simd<T, N> x) {
+  if constexpr (is_complex<T>) {
+    return std::sin(x.value);
+  } else {
+    return sincos<true>(x);
+  }
+}
+
+template <typename T, int N>
+Simd<T, N> cos(Simd<T, N> x) {
+  if constexpr (is_complex<T>) {
+    return std::cos(x.value);
+  } else {
+    return sincos<false>(x);
+  }
+}
+
+template <typename T, int N>
+Simd<T, N> erf(Simd<T, N> x) {
+  // https://github.com/pytorch/pytorch/blob/abf28982a8cb43342e7669d859de9543fd804cc9/aten/src/ATen/cpu/vec/vec256/vec256_float.h#L175
+  Simd<float, N> v = x;
+  auto t = recip(fma(Simd<float, N>(0.3275911f), abs(v), 1.0f));
+  auto r = fma(Simd<float, N>(1.061405429f), t, -1.453152027f);
+  r = fma(r, t, 1.421413741f);
+  r = fma(r, t, -0.284496736f);
+  r = fma(r, t, 0.254829592f);
+  auto e = -exp(-v * v);
+  auto result = Simd<T, N>(fma(e * t, r, 1.0f));
+  return select(x > 0, result, -result);
+}
+
+template <typename T, int N>
+Simd<T, N> erfinv(Simd<T, N> a_) {
+  Simd<float, N> a = a_;
+  auto t = fma(a, 0.0f - a, 1.0f);
+  t = log(t);
+  auto lhs = [](auto t) {
+    Simd<float, N> p;
+    p = 3.03697567e-10f; //  0x1.4deb44p-32
+    p = fma(p, t, 2.93243101e-8f); //  0x1.f7c9aep-26
+    p = fma(p, t, 1.22150334e-6f); //  0x1.47e512p-20
+    p = fma(p, t, 2.84108955e-5f); //  0x1.dca7dep-16
+    p = fma(p, t, 3.93552968e-4f); //  0x1.9cab92p-12
+    p = fma(p, t, 3.02698812e-3f); //  0x1.8cc0dep-9
+    p = fma(p, t, 4.83185798e-3f); //  0x1.3ca920p-8
+    p = fma(p, t, -2.64646143e-1f); // -0x1.0eff66p-2
+    return fma(p, t, 8.40016484e-1f); //  0x1.ae16a4p-1
+  };
+  auto rhs = [](auto t) {
+    Simd<float, N> p;
+    p = 5.43877832e-9f; //  0x1.75c000p-28
+    p = fma(p, t, 1.43285448e-7f); //  0x1.33b402p-23
+    p = fma(p, t, 1.22774793e-6f); //  0x1.499232p-20
+    p = fma(p, t, 1.12963626e-7f); //  0x1.e52cd2p-24
+    p = fma(p, t, -5.61530760e-5f); // -0x1.d70bd0p-15
+    p = fma(p, t, -1.47697632e-4f); // -0x1.35be90p-13
+    p = fma(p, t, 2.31468678e-3f); //  0x1.2f6400p-9
+    p = fma(p, t, 1.15392581e-2f); //  0x1.7a1e50p-7
+    p = fma(p, t, -2.32015476e-1f); // -0x1.db2aeep-3
+    return fma(p, t, 8.86226892e-1f); //  0x1.c5bf88p-1
+  };
+  auto thresh = 6.125f;
+  // Compute both branches and select if N > 1
+  if constexpr (N == 1) {
+    if ((abs(t) > thresh).value) { // maximum ulp error = 2.35793
+      return a * lhs(t);
+    } else { // maximum ulp error = 2.35002
+      return a * rhs(t);
+    }
+  } else {
+    return a * select(abs(t) > thresh, lhs(t), rhs(t));
+  }
+}
+
+} // namespace mlx::core::simd
diff --git a/macos/third_party/mlx/mlx/backend/cpu/simd/neon_fp16_simd.h b/macos/third_party/mlx/mlx/backend/cpu/simd/neon_fp16_simd.h
new file mode 100644
index 0000000..5d32042
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/simd/neon_fp16_simd.h
@@ -0,0 +1,212 @@
+#pragma once
+
+#include <arm_neon.h>
+
+#include "mlx/backend/cpu/simd/base_simd.h"
+
+namespace mlx::core::simd {
+
+constexpr int N = 8;
+
+template <>
+struct Simd<float16_t, N> {
+  static constexpr int size = N;
+  using scalar_t = float16_t;
+
+  Simd<float16_t, N>() {}
+
+  template <typename U>
+  Simd<float16_t, N>(U v) : value(vdupq_n_f16(v)){};
+
+  Simd<float16_t, N>(float16x8_t v) : value(v){};
+
+  Simd<float16_t, N>(Simd<float, N> other) {
+    auto f32x4_a = *(float32x4_t*)(&other);
+    auto f32x4_b = *((float32x4_t*)(&other) + 1);
+    value = vcvt_high_f16_f32(vcvt_f16_f32(f32x4_a), f32x4_b);
+  };
+
+  Simd<float16_t, N>(Simd<uint16_t, N> other) {
+    value = vcvtq_f16_u16(*(uint16x8_t*)(&other.value));
+  };
+
+  operator Simd<int16_t, N>() {
+    auto v = vcvtq_s16_f16(value);
+    return load<int16_t, N>((int16_t*)&v);
+  };
+
+  operator Simd<float, N>() {
+    float32x4x2_t v;
+    v.val[0] = vcvt_f32_f16(*(float16x4_t*)(&value));
+    v.val[1] = vcvt_high_f32_f16(value);
+    return load<float, N>((float*)&v);
+  }
+  float16_t operator[](int idx) const {
+    return reinterpret_cast<const float16_t*>(&value)[idx];
+  }
+
+  float16_t& operator[](int idx) {
+    return reinterpret_cast<float16_t*>(&value)[idx];
+  }
+
+  float16x8_t value;
+};
+
+#define DEFINE_NEON_UNARY_OP(name, op)                   \
+  inline Simd<float16_t, N> name(Simd<float16_t, N> a) { \
+    return Simd<float16_t, N>{op(a.value)};              \
+  }
+
+DEFINE_NEON_UNARY_OP(abs, vabsq_f16)
+DEFINE_NEON_UNARY_OP(ceil, vrndpq_f16)
+DEFINE_NEON_UNARY_OP(floor, vrndmq_f16)
+DEFINE_NEON_UNARY_OP(sqrt, vsqrtq_f16)
+DEFINE_NEON_UNARY_OP(rsqrt, vrsqrteq_f16)
+DEFINE_NEON_UNARY_OP(recip, vrecpeq_f16)
+DEFINE_NEON_UNARY_OP(rint, vrndnq_f16)
+
+#define DEFINE_NEON_BINARY_OP(name, op)                                        \
+  inline Simd<float16_t, N> name(Simd<float16_t, N> a, Simd<float16_t, N> b) { \
+    return op(a.value, b.value);                                               \
+  }                                                                            \
+  template <typename T>                                                        \
+  Simd<float16_t, N> name(Simd<float16_t, N> a, T b) {                         \
+    return op(a.value, Simd<float16_t, N>(b).value);                           \
+  }                                                                            \
+  template <typename T>                                                        \
+  Simd<float16_t, N> name(T a, Simd<float16_t, N> b) {                         \
+    return op(Simd<float16_t, N>(a).value, b.value);                           \
+  }
+
+inline Simd<float16_t, N> operator!(Simd<float16_t, N> v) {
+  auto out = vceqzq_f16(v.value);
+  return Simd<uint16_t, N>(*(uint16_t*)&out);
+}
+
+inline Simd<float16_t, N> operator-(Simd<float16_t, N> v) {
+  return vnegq_f16(v.value);
+}
+
+DEFINE_NEON_BINARY_OP(maximum, vmaxq_f16)
+DEFINE_NEON_BINARY_OP(minimum, vminq_f16)
+DEFINE_NEON_BINARY_OP(operator+, vaddq_f16)
+DEFINE_NEON_BINARY_OP(operator-, vsubq_f16)
+DEFINE_NEON_BINARY_OP(operator*, vmulq_f16)
+DEFINE_NEON_BINARY_OP(operator/, vdivq_f16)
+
+#define DEFINE_NEON_COMPARISON(Op, op)                   \
+  template <typename T>                                  \
+  Simd<bool, N> operator Op(Simd<float16_t, N> a, T b) { \
+    auto out = op(a.value, Simd<float16_t, N>(b).value); \
+    return Simd<uint16_t, N>(*(uint16_t*)(&out));        \
+  }                                                      \
+  template <typename T>                                  \
+  Simd<bool, N> operator Op(T a, Simd<float16_t, N> b) { \
+    auto out = op(Simd<float16_t, N>(a).value, b.value); \
+    return Simd<uint16_t, N>(*(uint16_t*)(&out));        \
+  }                                                      \
+  inline Simd<bool, N> operator Op(                      \
+      Simd<float16_t, N> a, Simd<float16_t, N> b) {      \
+    auto out = op(a.value, b.value);                     \
+    return Simd<uint16_t, N>(*(uint16_t*)(&out));        \
+  }
+
+DEFINE_NEON_COMPARISON(==, vceqq_f16)
+DEFINE_NEON_COMPARISON(>=, vcgeq_f16)
+DEFINE_NEON_COMPARISON(<=, vcleq_f16)
+DEFINE_NEON_COMPARISON(>, vcgtq_f16)
+DEFINE_NEON_COMPARISON(<, vcltq_f16)
+
+template <typename T>
+Simd<bool, N> operator!=(Simd<float16_t, N> a, T b) {
+  return !(a == b);
+}
+template <typename T>
+Simd<bool, N> operator!=(T a, Simd<float16_t, N> b) {
+  return !(a == b);
+}
+inline Simd<bool, N> operator!=(Simd<float16_t, N> a, Simd<float16_t, N> b) {
+  return !(a == b);
+}
+
+inline Simd<float16_t, N> operator||(
+    Simd<float16_t, N> a,
+    Simd<float16_t, N> b) {
+  return Simd<uint16_t, N>((a != 0) || (b != 0));
+}
+template <typename T>
+Simd<float16_t, N> operator||(Simd<float16_t, N> a, T b) {
+  return Simd<uint16_t, N>((a != 0) || (b != 0));
+}
+template <typename T>
+Simd<float16_t, N> operator||(T a, Simd<float16_t, N> b) {
+  return Simd<uint16_t, N>((a != 0) || (b != 0));
+}
+inline Simd<float16_t, N> operator&&(
+    Simd<float16_t, N> a,
+    Simd<float16_t, N> b) {
+  return Simd<uint16_t, N>((a != 0) && (b != 0));
+}
+template <typename T>
+Simd<float16_t, N> operator&&(Simd<float16_t, N> a, T b) {
+  return Simd<uint16_t, N>((a != 0) && (b != 0));
+}
+template <typename T>
+Simd<float16_t, N> operator&&(T a, Simd<float16_t, N> b) {
+  return Simd<uint16_t, N>((a != 0) && (b != 0));
+}
+
+template <>
+inline Simd<bool, N> isnan(Simd<float16_t, N> v) {
+  return v != v;
+}
+
+template <>
+inline Simd<float16_t, N>
+clamp(Simd<float16_t, N> v, Simd<float16_t, N> min, Simd<float16_t, N> max) {
+  return minimum(maximum(v, min), max);
+}
+
+template <typename T>
+Simd<float16_t, N> fma(Simd<float16_t, N> x, Simd<float16_t, N> y, T z) {
+  return vfmaq_f16(x.value, y.value, Simd<float16_t, N>(z).value);
+}
+
+template <typename MaskT>
+Simd<float16_t, N>
+select(Simd<MaskT, N> mask, Simd<float16_t, N> x, Simd<float16_t, N> y) {
+  return vbslq_f16(Simd<uint16_t, N>(mask).value, x.value, y.value);
+}
+
+// Reductions
+inline float16_t max(Simd<float16_t, N> x) {
+  float16x4_t y;
+  y = vpmax_f16(vget_low_f16(x.value), vget_high_f16(x.value));
+  y = vpmax_f16(y, y);
+  y = vpmax_f16(y, y);
+  return vget_lane_f16(y, 0);
+}
+inline float16_t min(Simd<float16_t, N> x) {
+  float16x4_t y;
+  y = vpmin_f16(vget_low_f16(x.value), vget_high_f16(x.value));
+  y = vpmin_f16(y, y);
+  y = vpmin_f16(y, y);
+  return vget_lane_f16(y, 0);
+}
+inline float16_t sum(Simd<float16_t, N> x) {
+  float16x4_t y;
+  y = vpadd_f16(vget_low_f16(x.value), vget_high_f16(x.value));
+  y = vpadd_f16(y, y);
+  y = vpadd_f16(y, y);
+  return vget_lane_f16(y, 0);
+}
+inline float16_t prod(Simd<float16_t, N> x) {
+  auto hx = vmul_f16(vget_low_f16(x.value), vget_high_f16(x.value));
+  auto out = hx[0];
+  hx[0] *= hx[1];
+  hx[0] *= hx[2];
+  hx[0] *= hx[3];
+  return hx[0];
+}
+
+} // namespace mlx::core::simd
diff --git a/macos/third_party/mlx/mlx/backend/cpu/simd/simd.h b/macos/third_party/mlx/mlx/backend/cpu/simd/simd.h
new file mode 100644
index 0000000..8700f24
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/simd/simd.h
@@ -0,0 +1,4 @@
+#pragma once
+
+#include "mlx/backend/cpu/simd/math.h"
+#include "mlx/backend/cpu/simd/type.h"
diff --git a/macos/third_party/mlx/mlx/backend/cpu/simd/type.h b/macos/third_party/mlx/mlx/backend/cpu/simd/type.h
new file mode 100644
index 0000000..59b6ecc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/simd/type.h
@@ -0,0 +1,11 @@
+#pragma once
+
+#include "mlx/backend/cpu/simd/base_simd.h"
+
+#ifdef MLX_USE_ACCELERATE
+#if defined(__x86_64__)
+// the accelerate_simd implementation require neon -- use base implementation
+#else
+#include "mlx/backend/cpu/simd/accelerate_simd.h"
+#endif
+#endif
diff --git a/macos/third_party/mlx/mlx/backend/cpu/slicing.h b/macos/third_party/mlx/mlx/backend/cpu/slicing.h
new file mode 100644
index 0000000..eda3732
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/slicing.h
@@ -0,0 +1,21 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+std::tuple<int64_t, Strides> prepare_slice(
+    const array& in,
+    const Shape& start_indices,
+    const Shape& strides);
+
+void shared_buffer_slice(
+    const array& in,
+    const Strides& out_strides,
+    size_t data_offset,
+    size_t data_size,
+    array& out);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/softmax.cpp b/macos/third_party/mlx/mlx/backend/cpu/softmax.cpp
new file mode 100644
index 0000000..4c2941e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/softmax.cpp
@@ -0,0 +1,170 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cassert>
+#include <cmath>
+
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/primitives.h"
+#include "mlx/types/limits.h"
+
+namespace mlx::core {
+
+namespace {
+
+using namespace mlx::core::simd;
+
+template <typename T, typename AccT>
+void softmax(const array& in, array& out, Stream stream) {
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+
+  const T* in_ptr = in.data<T>();
+  T* out_ptr = out.data<T>();
+
+  int M = in.shape().back();
+  int L = in.data_size() / M;
+
+  encoder.dispatch([in_ptr, out_ptr, M, L]() mutable {
+    constexpr bool same_t = std::is_same_v<T, AccT>;
+    constexpr int N = std::min(max_size<AccT>, max_size<T>);
+
+    const T* current_in_ptr;
+    T* current_out_ptr;
+
+    for (int i = 0; i < L; i++, in_ptr += M, out_ptr += M) {
+      // Find the maximum
+      current_in_ptr = in_ptr;
+      Simd<AccT, N> vmaximum(-numeric_limits<AccT>::infinity());
+      size_t s = M;
+      while (s >= N) {
+        Simd<AccT, N> vals = load<T, N>(current_in_ptr);
+        vmaximum = maximum(vals, vmaximum);
+        current_in_ptr += N;
+        s -= N;
+      }
+
+      AccT maximum = max(vmaximum);
+      while (s-- > 0) {
+        maximum = std::max(maximum, static_cast<AccT>(*current_in_ptr));
+        current_in_ptr++;
+      }
+
+      // Compute the normalizer and the exponentials
+      Simd<AccT, N> vnormalizer(0.0);
+      current_out_ptr = out_ptr;
+      current_in_ptr = in_ptr;
+      s = M;
+      while (s >= N) {
+        Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
+        vexp = exp(vexp - maximum);
+        if constexpr (same_t) {
+          store(current_out_ptr, vexp);
+        }
+        vnormalizer = vnormalizer + vexp;
+        current_in_ptr += N;
+        current_out_ptr += N;
+        s -= N;
+      }
+      AccT normalizer = sum(vnormalizer);
+      while (s-- > 0) {
+        AccT _exp = std::exp(*current_in_ptr - maximum);
+        if constexpr (same_t) {
+          *current_out_ptr = _exp;
+        }
+        normalizer += _exp;
+        current_in_ptr++;
+        current_out_ptr++;
+      }
+      normalizer = 1 / normalizer;
+
+      // Normalize
+      current_out_ptr = out_ptr;
+      current_in_ptr = in_ptr;
+      s = M;
+      while (s >= N) {
+        if constexpr (same_t) {
+          store(
+              current_out_ptr,
+              Simd<T, N>(load<T, N>(current_out_ptr) * normalizer));
+        } else {
+          Simd<AccT, N> vexp = load<T, N>(current_in_ptr);
+          vexp = exp(vexp - maximum) * normalizer;
+          store(current_out_ptr, Simd<T, N>(vexp));
+          current_in_ptr += N;
+        }
+        current_out_ptr += N;
+        s -= N;
+      }
+      while (s-- > 0) {
+        if constexpr (same_t) {
+          *current_out_ptr *= normalizer;
+        } else {
+          AccT _exp = std::exp(*current_in_ptr - maximum);
+          *current_out_ptr = static_cast<T>(_exp * normalizer);
+          current_in_ptr++;
+        }
+        current_out_ptr++;
+      }
+    }
+  });
+}
+
+} // namespace
+
+void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  // Make sure that the last dimension is contiguous
+  auto set_output = [s = stream(), &out](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_cpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  auto in = set_output(inputs[0]);
+
+  switch (in.dtype()) {
+    case float32:
+      softmax<float, float>(in, out, stream());
+      break;
+    case float16:
+      if (precise_) {
+        softmax<float16_t, float>(in, out, stream());
+      } else {
+        softmax<float16_t, float16_t>(in, out, stream());
+      }
+      break;
+    case bfloat16:
+      if (precise_) {
+        softmax<bfloat16_t, float>(in, out, stream());
+      } else {
+        softmax<bfloat16_t, bfloat16_t>(in, out, stream());
+      }
+      break;
+    case float64:
+      softmax<double, double>(in, out, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[softmax] Only defined for floating point types.");
+      break;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/sort.cpp b/macos/third_party/mlx/mlx/backend/cpu/sort.cpp
new file mode 100644
index 0000000..fcf12d7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/sort.cpp
@@ -0,0 +1,481 @@
+// Copyright © 2023 Apple Inc.
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <numeric>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+// NaN-aware comparator that places NaNs at the end
+template <typename T>
+bool nan_aware_less(T a, T b) {
+  if constexpr (std::is_floating_point_v<T> || std::is_same_v<T, complex64_t>) {
+    if (std::isnan(a))
+      return false;
+    if (std::isnan(b))
+      return true;
+  }
+  return a < b;
+}
+
+template <typename T>
+struct StridedIterator {
+  using iterator_category = std::random_access_iterator_tag;
+  using difference_type = int32_t;
+  using value_type = T;
+  using reference = value_type&;
+  using pointer = value_type*;
+
+  // Constructors
+  StridedIterator() = default;
+
+  explicit StridedIterator(T* ptr, int64_t stride, difference_type offset = 0)
+      : stride_(stride), ptr_(ptr + offset * stride) {}
+
+  explicit StridedIterator(array& arr, int axis, difference_type offset = 0)
+      : StridedIterator(arr.data<T>(), arr.strides()[axis], offset) {}
+
+  // Accessors
+  reference operator*() const {
+    return ptr_[0];
+  }
+
+  reference operator[](difference_type idx) const {
+    return ptr_[idx * stride_];
+  }
+
+  // Comparisons
+  bool operator==(const StridedIterator& other) const {
+    return ptr_ == other.ptr_ && stride_ == other.stride_;
+  }
+
+  bool operator!=(const StridedIterator& other) const {
+    return ptr_ != other.ptr_;
+  }
+
+  bool operator<(const StridedIterator& other) const {
+    return ptr_ < other.ptr_;
+  }
+
+  bool operator>(const StridedIterator& other) const {
+    return ptr_ > other.ptr_;
+  }
+
+  bool operator<=(const StridedIterator& other) const {
+    return ptr_ <= other.ptr_;
+  }
+
+  bool operator>=(const StridedIterator& other) const {
+    return ptr_ >= other.ptr_;
+  }
+
+  difference_type operator-(const StridedIterator& other) const {
+    return (ptr_ - other.ptr_) / stride_;
+  }
+
+  // Moving
+  StridedIterator& operator++() {
+    ptr_ += stride_;
+    return *this;
+  }
+
+  StridedIterator& operator--() {
+    ptr_ -= stride_;
+    return *this;
+  }
+
+  StridedIterator& operator+=(difference_type diff) {
+    ptr_ += diff * stride_;
+    return *this;
+  }
+
+  StridedIterator& operator-=(difference_type diff) {
+    ptr_ -= diff * stride_;
+    return *this;
+  }
+
+  StridedIterator operator+(difference_type diff) {
+    return StridedIterator(ptr_, stride_, diff);
+  }
+
+  StridedIterator operator-(difference_type diff) {
+    return StridedIterator(ptr_, stride_, -diff);
+  }
+
+ private:
+  int64_t stride_;
+  T* ptr_;
+};
+
+template <typename T>
+void sort(array& out, int axis) {
+  // Get axis, shape and stride info
+  axis = axis < 0 ? axis + out.ndim() : axis;
+  size_t in_size = out.size();
+  size_t n_rows = in_size / out.shape(axis);
+
+  auto remaining_shape = out.shape();
+  remaining_shape.erase(remaining_shape.begin() + axis);
+
+  auto remaining_strides = out.strides();
+  remaining_strides.erase(remaining_strides.begin() + axis);
+
+  auto axis_stride = out.strides()[axis];
+  auto axis_size = out.shape(axis);
+
+  // Perform sorting in place
+  ContiguousIterator src_it(
+      remaining_shape, remaining_strides, remaining_shape.size());
+  auto out_ptr = out.data<T>();
+  for (int i = 0; i < n_rows; i++) {
+    T* data_ptr = out_ptr + src_it.loc;
+
+    StridedIterator st(data_ptr, axis_stride, 0);
+    StridedIterator ed(data_ptr, axis_stride, axis_size);
+
+    std::stable_sort(st, ed, nan_aware_less<T>);
+    src_it.step();
+  }
+}
+
+template <typename T, typename IdxT = uint32_t>
+void argsort(const array& in, array& out, int axis) {
+  // Get axis, shape and stride info
+  axis = axis < 0 ? axis + in.ndim() : axis;
+  size_t n_rows = in.size() / in.shape(axis);
+
+  auto in_remaining_shape = in.shape();
+  in_remaining_shape.erase(in_remaining_shape.begin() + axis);
+
+  auto in_remaining_strides = in.strides();
+  in_remaining_strides.erase(in_remaining_strides.begin() + axis);
+
+  auto out_remaining_shape = out.shape();
+  out_remaining_shape.erase(out_remaining_shape.begin() + axis);
+
+  auto out_remaining_strides = out.strides();
+  out_remaining_strides.erase(out_remaining_strides.begin() + axis);
+
+  auto in_stride = in.strides()[axis];
+  auto out_stride = out.strides()[axis];
+  auto axis_size = in.shape(axis);
+
+  // Perform sorting
+  ContiguousIterator in_it(
+      in_remaining_shape, in_remaining_strides, in_remaining_shape.size());
+  ContiguousIterator out_it(
+      out_remaining_shape, out_remaining_strides, out_remaining_shape.size());
+  auto in_ptr = in.data<T>();
+  auto out_ptr = out.data<IdxT>();
+  for (int i = 0; i < n_rows; i++) {
+    const T* data_ptr = in_ptr + in_it.loc;
+    IdxT* idx_ptr = out_ptr + out_it.loc;
+
+    in_it.step();
+    out_it.step();
+
+    StridedIterator st_(idx_ptr, out_stride, 0);
+    StridedIterator ed_(idx_ptr, out_stride, axis_size);
+
+    // Initialize with iota
+    std::iota(st_, ed_, IdxT(0));
+
+    // Sort according to vals
+    StridedIterator st(idx_ptr, out_stride, 0);
+    StridedIterator ed(idx_ptr, out_stride, axis_size);
+
+    std::stable_sort(st, ed, [data_ptr, in_stride](IdxT a, IdxT b) {
+      auto v1 = data_ptr[a * in_stride];
+      auto v2 = data_ptr[b * in_stride];
+
+      // Handle NaNs (place them at the end)
+      if (std::is_floating_point<T>::value) {
+        if (std::isnan(v1))
+          return false;
+        if (std::isnan(v2))
+          return true;
+      }
+
+      return v1 < v2 || (v1 == v2 && a < b);
+    });
+  }
+}
+
+template <typename T>
+void partition(array& out, int axis, int kth) {
+  // Get axis, shape and stride info
+  axis = axis < 0 ? axis + out.ndim() : axis;
+  size_t in_size = out.size();
+  size_t n_rows = in_size / out.shape(axis);
+
+  auto remaining_shape = out.shape();
+  remaining_shape.erase(remaining_shape.begin() + axis);
+
+  auto remaining_strides = out.strides();
+  remaining_strides.erase(remaining_strides.begin() + axis);
+
+  auto axis_stride = out.strides()[axis];
+  int axis_size = out.shape(axis);
+
+  kth = kth < 0 ? kth + axis_size : kth;
+
+  // Perform partition in place
+  ContiguousIterator src_it(
+      remaining_shape, remaining_strides, remaining_shape.size());
+  auto out_ptr = out.data<T>();
+  for (int i = 0; i < n_rows; i++) {
+    T* data_ptr = out_ptr + src_it.loc;
+    src_it.step();
+
+    StridedIterator st(data_ptr, axis_stride, 0);
+    StridedIterator md(data_ptr, axis_stride, kth);
+    StridedIterator ed(data_ptr, axis_stride, axis_size);
+
+    std::nth_element(st, md, ed, nan_aware_less<T>);
+  }
+}
+
+template <typename T, typename IdxT = uint32_t>
+void argpartition(const array& in, array& out, int axis, int kth) {
+  // Get axis, shape and stride info
+  axis = axis < 0 ? axis + in.ndim() : axis;
+  size_t n_rows = in.size() / in.shape(axis);
+
+  auto in_remaining_shape = in.shape();
+  in_remaining_shape.erase(in_remaining_shape.begin() + axis);
+
+  auto in_remaining_strides = in.strides();
+  in_remaining_strides.erase(in_remaining_strides.begin() + axis);
+
+  auto out_remaining_shape = out.shape();
+  out_remaining_shape.erase(out_remaining_shape.begin() + axis);
+
+  auto out_remaining_strides = out.strides();
+  out_remaining_strides.erase(out_remaining_strides.begin() + axis);
+
+  auto in_stride = in.strides()[axis];
+  auto out_stride = out.strides()[axis];
+  auto axis_size = in.shape(axis);
+
+  kth = kth < 0 ? kth + axis_size : kth;
+
+  // Perform partition
+  ContiguousIterator in_it(
+      in_remaining_shape, in_remaining_strides, in_remaining_shape.size());
+  ContiguousIterator out_it(
+      out_remaining_shape, out_remaining_strides, out_remaining_shape.size());
+
+  auto in_ptr = in.data<T>();
+  auto out_ptr = out.data<IdxT>();
+
+  for (int i = 0; i < n_rows; i++) {
+    const T* data_ptr = in_ptr + in_it.loc;
+    IdxT* idx_ptr = out_ptr + out_it.loc;
+    in_it.step();
+    out_it.step();
+
+    StridedIterator st_(idx_ptr, out_stride, 0);
+    StridedIterator ed_(idx_ptr, out_stride, axis_size);
+
+    // Initialize with iota
+    std::iota(st_, ed_, IdxT(0));
+
+    // Sort according to vals
+    StridedIterator st(idx_ptr, out_stride, 0);
+    StridedIterator md(idx_ptr, out_stride, kth);
+    StridedIterator ed(idx_ptr, out_stride, axis_size);
+
+    std::nth_element(st, md, ed, [data_ptr, in_stride](IdxT a, IdxT b) {
+      auto v1 = data_ptr[a * in_stride];
+      auto v2 = data_ptr[b * in_stride];
+
+      // Handle NaNs (place them at the end)
+      if (std::is_floating_point<T>::value) {
+        if (std::isnan(v1))
+          return false;
+        if (std::isnan(v2))
+          return true;
+      }
+
+      return v1 < v2 || (v1 == v2 && a < b);
+    });
+  }
+}
+
+} // namespace
+
+void ArgSort::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+
+  // Allocate output
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_input_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    axis_ = axis_]() mutable {
+    switch (in.dtype()) {
+      case bool_:
+        return argsort<bool>(in, out, axis_);
+      case uint8:
+        return argsort<uint8_t>(in, out, axis_);
+      case uint16:
+        return argsort<uint16_t>(in, out, axis_);
+      case uint32:
+        return argsort<uint32_t>(in, out, axis_);
+      case uint64:
+        return argsort<uint64_t>(in, out, axis_);
+      case int8:
+        return argsort<int8_t>(in, out, axis_);
+      case int16:
+        return argsort<int16_t>(in, out, axis_);
+      case int32:
+        return argsort<int32_t>(in, out, axis_);
+      case int64:
+        return argsort<int64_t>(in, out, axis_);
+      case float32:
+        return argsort<float>(in, out, axis_);
+      case float64:
+        return argsort<double>(in, out, axis_);
+      case float16:
+        return argsort<float16_t>(in, out, axis_);
+      case bfloat16:
+        return argsort<bfloat16_t>(in, out, axis_);
+      case complex64:
+        return argsort<complex64_t>(in, out, axis_);
+    }
+  });
+}
+
+void Sort::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+
+  int axis = axis_;
+  if (axis < 0) {
+    axis += in.ndim();
+  }
+
+  // Copy input to output
+  CopyType ctype = (in.flags().contiguous && in.strides()[axis] != 0)
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, ctype, stream());
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_output_array(out);
+  encoder.dispatch([out = array::unsafe_weak_copy(out), axis]() mutable {
+    dispatch_all_types(out.dtype(), [&](auto type_tag) {
+      sort<MLX_GET_TYPE(type_tag)>(out, axis);
+    });
+  });
+}
+
+void ArgPartition::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+
+  // Allocate output
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(in);
+  encoder.set_input_array(out);
+  encoder.dispatch([in = array::unsafe_weak_copy(in),
+                    out = array::unsafe_weak_copy(out),
+                    axis_ = axis_,
+                    kth_ = kth_]() mutable {
+    switch (in.dtype()) {
+      case bool_:
+        return argpartition<bool>(in, out, axis_, kth_);
+      case uint8:
+        return argpartition<uint8_t>(in, out, axis_, kth_);
+      case uint16:
+        return argpartition<uint16_t>(in, out, axis_, kth_);
+      case uint32:
+        return argpartition<uint32_t>(in, out, axis_, kth_);
+      case uint64:
+        return argpartition<uint64_t>(in, out, axis_, kth_);
+      case int8:
+        return argpartition<int8_t>(in, out, axis_, kth_);
+      case int16:
+        return argpartition<int16_t>(in, out, axis_, kth_);
+      case int32:
+        return argpartition<int32_t>(in, out, axis_, kth_);
+      case int64:
+        return argpartition<int64_t>(in, out, axis_, kth_);
+      case float32:
+        return argpartition<float>(in, out, axis_, kth_);
+      case float64:
+        return argpartition<double>(in, out, axis_, kth_);
+      case float16:
+        return argpartition<float16_t>(in, out, axis_, kth_);
+      case bfloat16:
+        return argpartition<bfloat16_t>(in, out, axis_, kth_);
+      case complex64:
+        return argpartition<complex64_t>(in, out, axis_, kth_);
+    }
+  });
+}
+
+void Partition::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+
+  // Copy input to output
+  CopyType ctype = (in.flags().contiguous && in.strides()[axis_] != 0)
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, ctype, stream());
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_output_array(out);
+  encoder.dispatch([out = array::unsafe_weak_copy(out),
+                    axis_ = axis_,
+                    kth_ = kth_]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        return partition<bool>(out, axis_, kth_);
+      case uint8:
+        return partition<uint8_t>(out, axis_, kth_);
+      case uint16:
+        return partition<uint16_t>(out, axis_, kth_);
+      case uint32:
+        return partition<uint32_t>(out, axis_, kth_);
+      case uint64:
+        return partition<uint64_t>(out, axis_, kth_);
+      case int8:
+        return partition<int8_t>(out, axis_, kth_);
+      case int16:
+        return partition<int16_t>(out, axis_, kth_);
+      case int32:
+        return partition<int32_t>(out, axis_, kth_);
+      case int64:
+        return partition<int64_t>(out, axis_, kth_);
+      case float32:
+        return partition<float>(out, axis_, kth_);
+      case float64:
+        return partition<double>(out, axis_, kth_);
+      case float16:
+        return partition<float16_t>(out, axis_, kth_);
+      case bfloat16:
+        return partition<bfloat16_t>(out, axis_, kth_);
+      case complex64:
+        return partition<complex64_t>(out, axis_, kth_);
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/svd.cpp b/macos/third_party/mlx/mlx/backend/cpu/svd.cpp
new file mode 100644
index 0000000..ca01a0a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/svd.cpp
@@ -0,0 +1,289 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/copy.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/lapack.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename T, class Enable = void>
+struct SVDWork {};
+
+template <typename T>
+struct SVDWork<
+    T,
+    typename std::enable_if<std::is_floating_point<T>::value>::type> {
+  using R = T;
+
+  int N;
+  int M;
+  int K;
+  int lda;
+  int ldu;
+  int ldvt;
+  char jobz;
+  std::vector<array::Data> buffers;
+  int lwork;
+
+  SVDWork(int N, int M, int K, char jobz)
+      : N(N), M(M), K(K), lda(N), ldu(N), ldvt(M), jobz(jobz) {
+    T workspace_dimension = 0;
+
+    // Will contain the indices of eigenvectors that failed to converge (not
+    // used here but required by lapack).
+    buffers.emplace_back(allocator::malloc(sizeof(int) * 8 * K));
+
+    int lwork_query = -1;
+    int info;
+
+    // Compute workspace size.
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ nullptr,
+        /* lda = */ &lda,
+        /* s = */ nullptr,
+        /* u = */ nullptr,
+        /* ldu = */ &ldu,
+        /* vt = */ nullptr,
+        /* ldvt = */ &ldvt,
+        /* work = */ &workspace_dimension,
+        /* lwork = */ &lwork_query,
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    lwork = workspace_dimension;
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, R* s, T* u, T* vt) {
+    int info;
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ a,
+        /* lda = */ &lda,
+        /* s = */ s,
+        // According to the identity above, lapack will write Vᵀᵀ as U.
+        /* u = */ u,
+        /* ldu = */ &ldu,
+        // According to the identity above, lapack will write Uᵀ as Vᵀ.
+        /* vt = */ vt,
+        /* ldvt = */ &ldvt,
+        /* work = */ static_cast<T*>(buffers[1].buffer.raw_ptr()),
+        /* lwork = */ &lwork,
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+  }
+};
+
+template <>
+struct SVDWork<std::complex<float>> {
+  using T = std::complex<float>;
+  using R = float;
+
+  int N;
+  int M;
+  int K;
+  int lda;
+  int ldu;
+  int ldvt;
+  char jobz;
+  std::vector<array::Data> buffers;
+  int lwork;
+
+  SVDWork(int N, int M, int K, char jobz)
+      : N(N), M(M), K(K), lda(N), ldu(N), ldvt(M), jobz(jobz) {
+    T workspace_dimension = 0;
+
+    // Will contain the indices of eigenvectors that failed to converge (not
+    // used here but required by lapack).
+    buffers.emplace_back(allocator::malloc(sizeof(int) * 8 * K));
+
+    const int lrwork =
+        jobz == 'A' ? std::max(1, 5 * K * K + 5 * K) : std::max(1, 7 * K);
+    buffers.emplace_back(allocator::malloc(sizeof(float) * lrwork));
+
+    int lwork_query = -1;
+    int work_query = -1;
+    int info;
+
+    // Compute workspace size.
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ nullptr,
+        /* lda = */ &lda,
+        /* s = */ nullptr,
+        /* u = */ nullptr,
+        /* ldu = */ &ldu,
+        /* vt = */ nullptr,
+        /* ldvt = */ &ldvt,
+        /* work = */ &workspace_dimension,
+        /* lwork = */ &lwork_query,
+        /* rwork = */ static_cast<float*>(buffers[1].buffer.raw_ptr()),
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    lwork = workspace_dimension.real();
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, R* s, T* u, T* vt) {
+    int info;
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ a,
+        /* lda = */ &lda,
+        /* s = */ s,
+        // According to the identity above, lapack will write Vᵀᵀ as U.
+        /* u = */ u,
+        /* ldu = */ &ldu,
+        // According to the identity above, lapack will write Uᵀ as Vᵀ.
+        /* vt = */ vt,
+        /* ldvt = */ &ldvt,
+        /* work = */ static_cast<T*>(buffers[2].buffer.raw_ptr()),
+        /* lwork = */ &lwork,
+        /* rwork = */ static_cast<float*>(buffers[1].buffer.raw_ptr()),
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+  }
+};
+
+template <typename T>
+void svd_impl(
+    const array& a,
+    std::vector<array>& outputs,
+    bool compute_uv,
+    Stream stream) {
+  // Lapack uses the column-major convention. To avoid having to transpose
+  // the input and then transpose the outputs, we swap the indices/sizes of the
+  // matrices and take advantage of the following identity (see
+  // https://math.stackexchange.com/a/30077)
+  //    A = UΣVᵀ
+  //    Aᵀ = VΣUᵀ
+  // As a result some of the indices/sizes are swapped as noted above.
+
+  // Rows and cols of the original matrix in row-major order.
+  const int M = a.shape(-2);
+  const int N = a.shape(-1);
+  const int K = std::min(M, N);
+
+  using R = typename SVDWork<T>::R;
+
+  size_t num_matrices = a.size() / (M * N);
+
+  // lapack clobbers the input, so we have to make a copy.
+  array in(a.shape(), a.dtype(), nullptr, {});
+  copy_cpu(
+      a,
+      in,
+      a.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+      stream);
+
+  // Allocate outputs.
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  auto in_ptr = in.data<T>();
+  T* u_ptr;
+  R* s_ptr;
+  T* vt_ptr;
+
+  if (compute_uv) {
+    array& u = outputs[0];
+    array& s = outputs[1];
+    array& vt = outputs[2];
+
+    u.set_data(allocator::malloc(u.nbytes()));
+    s.set_data(allocator::malloc(s.nbytes()));
+    vt.set_data(allocator::malloc(vt.nbytes()));
+
+    encoder.set_output_array(u);
+    encoder.set_output_array(s);
+    encoder.set_output_array(vt);
+
+    s_ptr = s.data<R>();
+    u_ptr = u.data<T>();
+    vt_ptr = vt.data<T>();
+  } else {
+    array& s = outputs[0];
+
+    s.set_data(allocator::malloc(s.nbytes()));
+
+    encoder.set_output_array(s);
+
+    s_ptr = s.data<R>();
+    u_ptr = nullptr;
+    vt_ptr = nullptr;
+  }
+
+  encoder.dispatch([in_ptr, u_ptr, s_ptr, vt_ptr, M, N, K, num_matrices]() {
+    auto jobz = (u_ptr) ? 'A' : 'N';
+    SVDWork<T> svd_work(N, M, K, jobz);
+    // Loop over matrices.
+    for (int i = 0; i < num_matrices; i++) {
+      svd_work.run(
+          in_ptr + M * N * i,
+          s_ptr + K * i,
+          vt_ptr ? vt_ptr + N * N * i : nullptr,
+          u_ptr ? u_ptr + M * M * i : nullptr);
+    }
+  });
+  encoder.add_temporary(in);
+}
+
+void SVD::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  switch (inputs[0].dtype()) {
+    case float32:
+      svd_impl<float>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    case float64:
+      svd_impl<double>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    case complex64:
+      svd_impl<std::complex<float>>(inputs[0], outputs, compute_uv_, stream());
+      break;
+    default:
+      throw std::runtime_error(
+          "[SVD::eval_cpu] only supports float32, float64, or complex64.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/ternary.h b/macos/third_party/mlx/mlx/backend/cpu/ternary.h
new file mode 100644
index 0000000..a27a7f2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/ternary.h
@@ -0,0 +1,154 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+#include "mlx/array.h"
+#include "mlx/backend/common/ternary.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/encoder.h"
+
+namespace mlx::core {
+
+template <typename T1, typename T2, typename T3, typename U, typename Op, int D>
+void ternary_op_dims(
+    const T1* a,
+    const T2* b,
+    const T3* c,
+    U* out,
+    Op op,
+    const Shape& shape,
+    const Strides& a_strides,
+    const Strides& b_strides,
+    const Strides& c_strides,
+    const Strides& out_strides,
+    int axis) {
+  auto stride_a = a_strides[axis];
+  auto stride_b = b_strides[axis];
+  auto stride_c = c_strides[axis];
+  auto stride_out = out_strides[axis];
+  auto N = shape[axis];
+
+  for (int i = 0; i < N; i++) {
+    if constexpr (D > 1) {
+      ternary_op_dims<T1, T2, T3, U, Op, D - 1>(
+          a,
+          b,
+          c,
+          out,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          c_strides,
+          out_strides,
+          axis + 1);
+    } else {
+      *out = op(*a, *b, *c);
+    }
+    a += stride_a;
+    b += stride_b;
+    c += stride_c;
+    out += stride_out;
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op_dispatch_dims(
+    const T1* a_ptr,
+    const T2* b_ptr,
+    const T3* c_ptr,
+    U* out_ptr,
+    Op op,
+    size_t size,
+    Shape& shape,
+    std::vector<Strides>& strides) {
+  const auto& a_strides = strides[0];
+  const auto& b_strides = strides[1];
+  const auto& c_strides = strides[2];
+  const auto& out_strides = strides[3];
+  int ndim = shape.size();
+  switch (ndim) {
+    case 1:
+      ternary_op_dims<T1, T2, T3, U, Op, 1>(
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          out_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          c_strides,
+          out_strides,
+          0);
+      return;
+    case 2:
+      ternary_op_dims<T1, T2, T3, U, Op, 2>(
+          a_ptr,
+          b_ptr,
+          c_ptr,
+          out_ptr,
+          op,
+          shape,
+          a_strides,
+          b_strides,
+          c_strides,
+          out_strides,
+          0);
+      return;
+  }
+
+  ContiguousIterator a_it(shape, a_strides, ndim - 2);
+  ContiguousIterator b_it(shape, b_strides, ndim - 2);
+  ContiguousIterator c_it(shape, c_strides, ndim - 2);
+  auto stride = out_strides[ndim - 3];
+  for (size_t elem = 0; elem < size; elem += stride) {
+    ternary_op_dims<T1, T2, T3, U, Op, 2>(
+        a_ptr + a_it.loc,
+        b_ptr + b_it.loc,
+        c_ptr + c_it.loc,
+        out_ptr + elem,
+        op,
+        shape,
+        a_strides,
+        b_strides,
+        c_strides,
+        out_strides,
+        ndim - 2);
+    a_it.step();
+    b_it.step();
+    c_it.step();
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op,
+    TernaryOpType topt) {
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+  U* out_ptr = out.data<U>();
+
+  if (topt == TernaryOpType::ScalarScalarScalar) {
+    *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
+  } else if (topt == TernaryOpType::VectorVectorVector) {
+    for (size_t i = 0; i < out.size(); ++i) {
+      *out_ptr = op(*a_ptr, *b_ptr, *c_ptr);
+      a_ptr++;
+      b_ptr++;
+      c_ptr++;
+      out_ptr++;
+    }
+  } else {
+    auto [shape, strides] = collapse_contiguous_dims(
+        a.shape(), {a.strides(), b.strides(), c.strides(), out.strides()});
+    ternary_op_dispatch_dims<T1, T2, T3, U>(
+        a_ptr, b_ptr, c_ptr, out_ptr, op, out.size(), shape, strides);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/threefry.cpp b/macos/third_party/mlx/mlx/backend/cpu/threefry.cpp
new file mode 100644
index 0000000..8056b84
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/threefry.cpp
@@ -0,0 +1,31 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/backend/cpu/threefry.h"
+
+namespace mlx::core::random {
+
+std::pair<uint32_t, uint32_t> threefry2x32_hash(
+    const std::pair<uint32_t, uint32_t>& key,
+    std::pair<uint32_t, uint32_t> count) {
+  constexpr static uint32_t rotations[2][4] = {
+      {13, 15, 26, 6}, {17, 29, 16, 24}};
+
+  uint32_t ks[3] = {key.first, key.second, key.first ^ key.second ^ 0x1BD11BDA};
+
+  count.first += ks[0];
+  count.second += ks[1];
+
+  for (int i = 0; i < 5; ++i) {
+    for (auto r : rotations[i % 2]) {
+      count.first += count.second;
+      count.second = (count.second << r) | (count.second >> (32 - r));
+      count.second ^= count.first;
+    }
+    count.first += ks[(i + 1) % 3];
+    count.second += ks[(i + 2) % 3] + i + 1;
+  }
+
+  return count;
+}
+
+} // namespace mlx::core::random
diff --git a/macos/third_party/mlx/mlx/backend/cpu/threefry.h b/macos/third_party/mlx/mlx/backend/cpu/threefry.h
new file mode 100644
index 0000000..0fc485f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/threefry.h
@@ -0,0 +1,21 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <cstdint>
+#include <utility>
+
+namespace mlx::core::random {
+
+/** Applies the Threefry 2x32 hash function.
+ * This code is based on the Jax counter-based and splittable PRNG
+ * https://github.com/google/jax/blob/main/docs/jep/263-prng.md
+ *
+ * Original Threefry reference:
+ * http://www.thesalmons.org/john/random123/papers/random123sc11.pdf
+ */
+std::pair<uint32_t, uint32_t> threefry2x32_hash(
+    const std::pair<uint32_t, uint32_t>& key,
+    std::pair<uint32_t, uint32_t> count);
+
+} // namespace mlx::core::random
diff --git a/macos/third_party/mlx/mlx/backend/cpu/unary.cpp b/macos/third_party/mlx/mlx/backend/cpu/unary.cpp
new file mode 100644
index 0000000..eafe988
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/unary.cpp
@@ -0,0 +1,238 @@
+// Copyright © 2024 Apple Inc.
+
+// Required for using M_LN2 in MSVC.
+#define _USE_MATH_DEFINES
+
+#include <cassert>
+
+#include "mlx/backend/cpu/unary.h"
+#include "mlx/backend/cpu/unary_ops.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), unsignedinteger) || in.dtype() == bool_) {
+    // No-op for unsigned types
+    out.copy_shared_buffer(in);
+  } else {
+    unary_signed(in, out, detail::Abs(), stream());
+  }
+}
+
+void ArcCos::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcCos(), stream());
+}
+
+void ArcCosh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcCosh(), stream());
+}
+
+void ArcSin::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcSin(), stream());
+}
+
+void ArcSinh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcSinh(), stream());
+}
+
+void ArcTan::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcTan(), stream());
+}
+
+void ArcTanh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::ArcTanh(), stream());
+}
+
+void BitwiseInvert::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_int(in, out, detail::BitwiseInvert(), stream());
+}
+
+void Ceil::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_fp(in, out, detail::Ceil(), stream());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+void Conjugate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  unary_complex(inputs[0], out, detail::Conjugate(), stream());
+}
+
+void Cos::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Cos(), stream());
+}
+
+void Cosh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Cosh(), stream());
+}
+
+void Erf::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_real_fp(in, out, detail::Erf(), stream());
+}
+
+void ErfInv::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_real_fp(in, out, detail::ErfInv(), stream());
+}
+
+void Exp::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Exp(), stream());
+}
+
+void Expm1::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Expm1(), stream());
+}
+
+void Floor::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_fp(in, out, detail::Floor(), stream());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+void Imag::eval_cpu(const std::vector<array>& inputs, array& out) {
+  unary_complex_to_float(inputs[0], out, detail::Imag(), stream());
+}
+
+void Log::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  switch (base_) {
+    case Base::e:
+      unary_fp(in, out, detail::Log(), stream());
+      break;
+    case Base::two:
+      unary_fp(in, out, detail::Log2(), stream());
+      break;
+    case Base::ten:
+      unary_fp(in, out, detail::Log10(), stream());
+      break;
+  }
+}
+
+void Log1p::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Log1p(), stream());
+}
+
+void LogicalNot::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  unary(in, out, detail::LogicalNot(), stream());
+}
+
+void Negative::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  unary(in, out, detail::Negative(), stream());
+}
+
+void Real::eval_cpu(const std::vector<array>& inputs, array& out) {
+  unary_complex_to_float(inputs[0], out, detail::Real(), stream());
+}
+
+void Round::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_fp(in, out, detail::Round(), stream());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+void Sigmoid::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Sigmoid(), stream());
+}
+
+void Sign::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (in.dtype() == bool_) {
+    out.copy_shared_buffer(in);
+  } else {
+    unary(in, out, detail::Sign(), stream());
+  }
+}
+
+void Sin::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Sin(), stream());
+}
+
+void Sinh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Sinh(), stream());
+}
+
+void Square::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  unary(in, out, detail::Square(), stream());
+}
+
+void Sqrt::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  if (recip_) {
+    unary_fp(in, out, detail::Rsqrt(), stream());
+  } else {
+    unary_fp(in, out, detail::Sqrt(), stream());
+  }
+}
+
+void Tan::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Tan(), stream());
+}
+
+void Tanh::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  unary_fp(in, out, detail::Tanh(), stream());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/unary.h b/macos/third_party/mlx/mlx/backend/cpu/unary.h
new file mode 100644
index 0000000..4fab6a7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/unary.h
@@ -0,0 +1,281 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/unary.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/simd/simd.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+template <typename T, typename U = T, typename Op>
+void unary_op(const T* a, U* out, size_t shape, size_t stride) {
+  for (size_t i = 0; i < shape; i += 1) {
+    out[i] = Op{}(*a);
+    a += stride;
+  }
+}
+
+template <typename T, typename U = T, typename Op>
+void unary_op(const array& a, array& out, Op) {
+  const T* src = a.data<T>();
+  U* dst = out.data<U>();
+  auto ndim = a.ndim();
+  if (a.flags().contiguous) {
+    auto size = a.data_size();
+    constexpr int N = std::min(simd::max_size<T>, simd::max_size<U>);
+    while (size >= N) {
+      simd::store(dst, simd::Simd<U, N>(Op{}(simd::load<T, N>(src))));
+      size -= N;
+      src += N;
+      dst += N;
+    }
+    while (size > 0) {
+      *dst = Op{}(*src);
+      size--;
+      dst++;
+      src++;
+    }
+  } else {
+    size_t shape = ndim > 0 ? a.shape().back() : 1;
+    size_t stride = ndim > 0 ? a.strides().back() : 1;
+    if (ndim <= 1) {
+      unary_op<T, U, Op>(src, dst, shape, stride);
+      return;
+    }
+    auto it = ContiguousIterator(a.shape(), a.strides(), ndim - 1);
+    for (size_t elem = 0; elem < a.size(); elem += shape) {
+      unary_op<T, U, Op>(src + it.loc, dst + elem, shape, stride);
+      it.step();
+    }
+  }
+}
+
+template <typename Op>
+void unary(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case bool_:
+        unary_op<bool>(a, out, op);
+        break;
+      case uint8:
+        unary_op<uint8_t>(a, out, op);
+        break;
+      case uint16:
+        unary_op<uint16_t>(a, out, op);
+        break;
+      case uint32:
+        unary_op<uint32_t>(a, out, op);
+        break;
+      case uint64:
+        unary_op<uint64_t>(a, out, op);
+        break;
+      case int8:
+        unary_op<int8_t>(a, out, op);
+        break;
+      case int16:
+        unary_op<int16_t>(a, out, op);
+        break;
+      case int32:
+        unary_op<int32_t>(a, out, op);
+        break;
+      case int64:
+        unary_op<int64_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case complex64:
+        unary_op<complex64_t>(a, out, op);
+        break;
+    }
+  });
+}
+
+template <typename Op>
+void unary_real_fp(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      default:
+        std::ostringstream err;
+        err << "[unary_real] Does not support " << out.dtype();
+        throw std::runtime_error(err.str());
+    }
+  });
+}
+template <typename Op>
+void unary_fp(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      case complex64:
+        unary_op<complex64_t>(a, out, op);
+        break;
+      default:
+        std::ostringstream err;
+        err << "[unary_fp] Does not support " << out.dtype();
+        throw std::runtime_error(err.str());
+    }
+  });
+}
+
+template <typename Op>
+void unary_signed(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case int8:
+        unary_op<int8_t>(a, out, op);
+        break;
+      case int16:
+        unary_op<int16_t>(a, out, op);
+        break;
+      case int32:
+        unary_op<int32_t>(a, out, op);
+        break;
+      case int64:
+        unary_op<int64_t>(a, out, op);
+        break;
+      case float16:
+        unary_op<float16_t>(a, out, op);
+        break;
+      case float32:
+        unary_op<float>(a, out, op);
+        break;
+      case float64:
+        unary_op<double>(a, out, op);
+        break;
+      case bfloat16:
+        unary_op<bfloat16_t>(a, out, op);
+        break;
+      case complex64:
+        unary_op<complex64_t>(a, out, op);
+        break;
+      default:
+        throw std::runtime_error("[Abs] Called on unsigned type");
+    }
+  });
+}
+
+template <typename Op>
+void unary_complex(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable { unary_op<complex64_t>(a, out, op); });
+}
+
+template <typename Op>
+void unary_complex_to_float(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch(
+      [a = array::unsafe_weak_copy(a),
+       out = array::unsafe_weak_copy(out),
+       op = op]() mutable { unary_op<complex64_t, float>(a, out, op); });
+}
+
+template <typename Op>
+void unary_int(const array& a, array& out, Op op, Stream stream) {
+  set_unary_output_data(a, out);
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(a);
+  encoder.set_output_array(out);
+  encoder.dispatch([a = array::unsafe_weak_copy(a),
+                    out = array::unsafe_weak_copy(out),
+                    op = op]() mutable {
+    switch (out.dtype()) {
+      case uint8:
+        unary_op<uint8_t>(a, out, op);
+        break;
+      case uint16:
+        unary_op<uint16_t>(a, out, op);
+        break;
+      case uint32:
+        unary_op<uint32_t>(a, out, op);
+        break;
+      case uint64:
+        unary_op<uint64_t>(a, out, op);
+        break;
+      case int8:
+        unary_op<int8_t>(a, out, op);
+        break;
+      case int16:
+        unary_op<int16_t>(a, out, op);
+        break;
+      case int32:
+        unary_op<int32_t>(a, out, op);
+        break;
+      case int64:
+        unary_op<int64_t>(a, out, op);
+        break;
+      default:
+        std::ostringstream err;
+        err << "[unary_int] Does not support " << out.dtype();
+        throw std::runtime_error(err.str());
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cpu/unary_ops.h b/macos/third_party/mlx/mlx/backend/cpu/unary_ops.h
new file mode 100644
index 0000000..f441e88
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cpu/unary_ops.h
@@ -0,0 +1,175 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <stdint.h>
+#include <cmath>
+#include <complex>
+
+#include "mlx/backend/cpu/simd/simd.h"
+
+namespace mlx::core::detail {
+
+using namespace mlx::core::simd;
+
+#define SINGLE()                         \
+  template <typename T>                  \
+  T operator()(T x) {                    \
+    return (*this)(Simd<T, 1>(x)).value; \
+  }
+
+#define DEFAULT_OP(Op, op)                \
+  struct Op {                             \
+    template <int N, typename T>          \
+    Simd<T, N> operator()(Simd<T, N> x) { \
+      return simd::op(x);                 \
+    }                                     \
+    SINGLE()                              \
+  };
+
+DEFAULT_OP(Abs, abs)
+DEFAULT_OP(ArcCos, acos)
+DEFAULT_OP(ArcCosh, acosh)
+DEFAULT_OP(ArcSin, asin)
+DEFAULT_OP(ArcSinh, asinh)
+DEFAULT_OP(ArcTan, atan)
+DEFAULT_OP(ArcTanh, atanh)
+DEFAULT_OP(BitwiseInvert, operator~)
+DEFAULT_OP(Ceil, ceil)
+DEFAULT_OP(Conjugate, conj)
+DEFAULT_OP(Cos, cos)
+DEFAULT_OP(Cosh, cosh)
+DEFAULT_OP(Erf, erf)
+DEFAULT_OP(ErfInv, erfinv)
+DEFAULT_OP(Exp, exp)
+DEFAULT_OP(Expm1, expm1)
+DEFAULT_OP(Floor, floor);
+DEFAULT_OP(Log, log);
+DEFAULT_OP(Log2, log2);
+DEFAULT_OP(Log10, log10);
+DEFAULT_OP(Log1p, log1p);
+DEFAULT_OP(LogicalNot, operator!)
+DEFAULT_OP(Negative, operator-)
+DEFAULT_OP(Round, rint);
+DEFAULT_OP(Sin, sin)
+DEFAULT_OP(Sinh, sinh)
+DEFAULT_OP(Sqrt, sqrt)
+DEFAULT_OP(Rsqrt, rsqrt)
+DEFAULT_OP(Tan, tan)
+DEFAULT_OP(Tanh, tanh)
+
+struct Imag {
+  template <int N>
+  Simd<float, N> operator()(Simd<complex64_t, N> x) {
+    return simd::imag(x);
+  }
+  SINGLE()
+};
+
+struct Real {
+  template <int N>
+  Simd<float, N> operator()(Simd<complex64_t, N> x) {
+    return simd::real(x);
+  }
+  SINGLE()
+};
+
+struct Sigmoid {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x) {
+    auto y = 1.0f / (1.0f + simd::exp(simd::abs(x)));
+    return simd::select(x < Simd<T, N>{0}, y, Simd<T, N>{1} - y);
+  }
+  SINGLE()
+};
+
+struct Sign {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x) {
+    auto z = Simd<T, N>{0};
+    auto o = Simd<T, N>{1};
+    auto m = Simd<T, N>{-1};
+    if constexpr (std::is_unsigned_v<T>) {
+      return simd::select(x == z, z, o);
+    } else if constexpr (std::is_same_v<T, complex64_t>) {
+      return simd::select(x == z, x, Simd<T, N>(x / simd::abs(x)));
+    } else {
+      return simd::select(x < z, m, simd::select(x > z, o, z));
+    }
+  }
+  SINGLE()
+};
+
+struct Square {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x) {
+    return x * x;
+  }
+  SINGLE()
+};
+
+template <int N>
+Simd<float, N> fp32_from_bits(Simd<uint32_t, N> x) {
+  return *(Simd<float, N>*)(&x);
+}
+template <int N>
+Simd<uint32_t, N> fp32_to_bits(Simd<float, N> x) {
+  return *(Simd<uint32_t, N>*)(&x);
+}
+
+struct ToFP8 {
+  template <typename T, int N>
+  Simd<uint8_t, N> operator()(Simd<T, N> f) {
+    uint32_t fp8_max = 543 << 21;
+    auto denorm_mask = Simd<uint32_t, N>(141 << 23);
+    Simd<uint32_t, N> f_bits;
+    Simd<float, N> f32 = f;
+    f_bits = fp32_to_bits(f32);
+    Simd<uint8_t, N> result = 0u;
+    auto sign = f_bits & 0x80000000;
+    f_bits = f_bits ^ sign;
+
+    auto f_bits_low =
+        fp32_to_bits(fp32_from_bits(f_bits) + fp32_from_bits(denorm_mask));
+    auto result_low = Simd<uint8_t, N>(f_bits_low - denorm_mask);
+
+    auto mant_odd = Simd<uint8_t, N>((f_bits >> 20) & 1);
+    auto f_bits_high = f_bits + (((uint32_t)(7 - 127) << 23) + 0x7FFFF);
+    f_bits_high = f_bits_high + Simd<uint32_t, N>(mant_odd);
+
+    auto result_high = Simd<uint8_t, N>(f_bits_high >> 20);
+    result = select(f_bits < (121 << 23), result_low, result_high);
+
+    auto result_sat = Simd<uint8_t, N>(0x7E);
+    result = select(f_bits >= fp8_max, result_sat, result);
+    return result | Simd<uint8_t, N>(sign >> 24);
+  }
+
+  template <typename T>
+  uint8_t operator()(T x) {
+    return (*this)(Simd<T, 1>(x)).value;
+  }
+};
+
+struct FromFP8 {
+  template <int N>
+  Simd<float, N> operator()(Simd<uint8_t, N> x) {
+    auto v = Simd<uint16_t, N>(x & 127) << 7;
+    Simd<float, N> out;
+    if constexpr (simd::max_size<float16_t> >= N) {
+      auto converted = *(Simd<float16_t, N>*)(&v);
+      out = converted * 256.0;
+    } else {
+      for (int i = 0; i < N; ++i) {
+        auto converted = *(float16_t*)(&v[i]);
+        out[i] = converted * 256.0;
+      }
+    }
+    auto sign = Simd<bool, N>(x & 128);
+    return select(sign, -out, out);
+  }
+  float operator()(uint8_t x) {
+    return (*this)(Simd<uint8_t, 1>(x)).value;
+  }
+};
+} // namespace mlx::core::detail
diff --git a/macos/third_party/mlx/mlx/backend/cuda/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/cuda/CMakeLists.txt
new file mode 100644
index 0000000..013b24b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/CMakeLists.txt
@@ -0,0 +1,265 @@
+# Filename rules in cuda backend:
+#
+# * Use .cu/.cuh if code contains device code, and .cpp/.h if not.
+# * Device-only code should be put in device/ subdir.
+# * Files in device/ subdir should not include files outside.
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/arange.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/binary_two.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy/copy_contiguous.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy/copy_general.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy/copy_general_dynamic.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy/copy_general_input.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/event.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/layer_norm.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/random.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce/all_reduce.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce/col_reduce.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce/init_reduce.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce/row_reduce.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/rms_norm.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/rope.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/scan.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/sort.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/ternary.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/fp_quantize.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/convert_fp8.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/binary)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary)
+
+# fp4 is not available on < 12.8
+if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0)
+  target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/)
+  target_sources(mlx
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/no_qqmm_impl.cpp)
+else()
+  target_sources(
+    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_impl.cpp
+                ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp)
+endif()
+
+if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)
+  target_sources(
+    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm_batched_12_9.cu)
+else()
+  target_sources(
+    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm_batched_12_0.cpp)
+endif()
+
+# Embed kernel sources in binary for JIT compilation.
+file(
+  GLOB MLX_JIT_SOURCES
+  RELATIVE ${CMAKE_CURRENT_SOURCE_DIR}
+  "${CMAKE_CURRENT_SOURCE_DIR}/device/*.h"
+  "${CMAKE_CURRENT_SOURCE_DIR}/device/*.cuh")
+string(JOIN ":" MLX_JIT_SOURCES_ARG ${MLX_JIT_SOURCES})
+add_custom_command(
+  OUTPUT gen/cuda_jit_sources.h
+  COMMAND
+    ${CMAKE_COMMAND} -DMLX_SOURCE_ROOT=${CMAKE_CURRENT_SOURCE_DIR}
+    -DMLX_JIT_SOURCES=${MLX_JIT_SOURCES_ARG} -P
+    "${CMAKE_CURRENT_SOURCE_DIR}/bin2h.cmake"
+  DEPENDS bin2h.cmake ${MLX_JIT_SOURCES})
+add_custom_target(cuda_jit_sources DEPENDS gen/cuda_jit_sources.h)
+add_dependencies(mlx cuda_jit_sources)
+target_include_directories(mlx PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/gen")
+
+# ------------------------ Compilation configs ------------------------
+
+target_compile_definitions(mlx PRIVATE MLX_USE_CUDA)
+
+# Enable defining device lambda functions.
+target_compile_options(mlx
+                       PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>")
+
+# Enable calling host constexpr functions from device. This is needed because
+# the constexpr version of isnan is host only.
+target_compile_options(
+  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")
+
+# Suppress nvcc warnings on C++ headers.
+target_compile_options(
+  mlx
+  PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=27,997,1394,20011,20208">
+)
+
+# Ignore some valid nvcc warnings, we might want to fix them in future.
+target_compile_options(
+  mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=177,550">)
+
+# Use stronger binaries compression. This feature was introduced in CUDA 12.8
+# and requires drivers released after CUDA 12.4.
+if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
+  target_compile_options(
+    mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--compress-mode=size>")
+endif()
+
+# Use native CUDA arch by default.
+if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
+  execute_process(
+    COMMAND __nvcc_device_query
+    OUTPUT_VARIABLE MLX_CUDA_ARCHITECTURES
+    OUTPUT_STRIP_TRAILING_WHITESPACE)
+  set(UPGRADABLE_ARCHITECTURES "90;100;121")
+  if(MLX_CUDA_ARCHITECTURES STREQUAL "")
+    message(
+      FATAL_ERROR
+        "Can not get native CUDA arch, must set MLX_CUDA_ARCHITECTURES")
+  elseif(MLX_CUDA_ARCHITECTURES IN_LIST UPGRADABLE_ARCHITECTURES)
+    # Use arch-specific compute capability whenever possible.
+    set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a")
+  endif()
+endif()
+message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
+set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
+                                     "${MLX_CUDA_ARCHITECTURES}")
+
+# Search CUDA libs from installed python packages.
+if(WIN32)
+  # Resolve paths of unfound DLL at runtime.
+  if(BUILD_SHARED_LIBS)
+    target_link_libraries(mlx PRIVATE "delayimp.lib")
+    target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp)
+  else()
+    # For static library the delayload must be compiled into final executables.
+    target_link_libraries(mlx PUBLIC "delayimp.lib")
+    target_sources(
+      mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp>)
+  endif()
+  # Get all the CUDA DLLs we could link with.
+  file(
+    GLOB CUDA_DLL_NAMES
+    RELATIVE "${CUDAToolkit_BIN_DIR}/x64"
+    "${CUDAToolkit_BIN_DIR}/x64/*.dll")
+  # Delay load CUDA and cuDNN libs.
+  foreach(CUDA_DLL ${CUDA_DLL_NAMES} ${CUDNN_DLL_NAMES})
+    target_link_options(mlx PUBLIC "/DELAYLOAD:${CUDA_DLL}")
+  endforeach()
+  # Pass the locations where CUDA DLLs are placed.
+  if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
+    target_compile_definitions(
+      mlx PUBLIC MLX_CUDA_BIN_DIR="${CUDAToolkit_BIN_DIR}/x64"
+                 MLX_CUDNN_BIN_DIR="${CUDNN_BIN_DIR}")
+  endif()
+else()
+  # For POSIX we rely on RPATH to search for CUDA libs.
+  if(MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
+    set_property(
+      TARGET mlx
+      APPEND
+      PROPERTY INSTALL_RPATH
+               # The paths here should match the install_requires in setup.py.
+               "$ORIGIN/../../nvidia/cublas/lib"
+               "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
+               "$ORIGIN/../../nvidia/cudnn/lib"
+               "$ORIGIN/../../nvidia/nccl/lib")
+  endif()
+endif()
+
+# ------------------------ Dependencies ------------------------
+
+# Use fixed version of CCCL.
+FetchContent_Declare(
+  cccl
+  URL "https://github.com/NVIDIA/cccl/releases/download/v3.1.3/cccl-v3.1.3.zip")
+FetchContent_MakeAvailable(cccl)
+target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")
+
+# Install CCCL headers for JIT.
+install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
+install(DIRECTORY ${cccl_SOURCE_DIR}/include/nv
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
+
+# The binary of C++ tests will not be installed so it can not find the CCCL
+# headers, and we have to hard-code the path.
+if(MLX_BUILD_TESTS)
+  target_compile_definitions(mlx
+                             PRIVATE MLX_CCCL_DIR="${cccl_SOURCE_DIR}/include")
+endif()
+
+# Use fixed version of NVTX.
+FetchContent_Declare(
+  nvtx3
+  GIT_REPOSITORY https://github.com/NVIDIA/NVTX.git
+  GIT_TAG v3.1.1
+  GIT_SHALLOW TRUE
+  SOURCE_SUBDIR c EXCLUDE_FROM_ALL)
+FetchContent_MakeAvailable(nvtx3)
+target_link_libraries(mlx PUBLIC $<BUILD_INTERFACE:nvtx3-cpp>)
+
+# Make cuda runtime APIs available in non-cuda files.
+target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+
+# Use cublasLt.
+target_link_libraries(mlx PRIVATE CUDA::cublasLt)
+
+# Use NVRTC and driver APIs.
+target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver)
+
+# Use the frontend APIs of cuDNN.
+FetchContent_Declare(
+  cudnn
+  GIT_REPOSITORY https://github.com/NVIDIA/cudnn-frontend.git
+  GIT_TAG v1.16.0
+  GIT_SHALLOW TRUE
+  EXCLUDE_FROM_ALL)
+set(CUDNN_FRONTEND_SKIP_JSON_LIB ON)
+set(CUDNN_FRONTEND_BUILD_SAMPLES OFF)
+set(CUDNN_FRONTEND_BUILD_TESTS OFF)
+set(CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS OFF)
+FetchContent_MakeAvailable(cudnn)
+target_link_libraries(mlx PRIVATE cudnn_frontend)
+# Link with the actual cuDNN libraries.
+target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
+
+# Use header-only CUTLASS.
+FetchContent_Declare(
+  cutlass
+  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+  GIT_TAG v4.3.5
+  GIT_SHALLOW TRUE
+  SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
+FetchContent_MakeAvailable(cutlass)
+target_include_directories(
+  mlx SYSTEM PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)
diff --git a/macos/third_party/mlx/mlx/backend/cuda/allocator.cpp b/macos/third_party/mlx/mlx/backend/cuda/allocator.cpp
new file mode 100644
index 0000000..0337a27
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/allocator.cpp
@@ -0,0 +1,454 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/memory.h"
+#include "mlx/scheduler.h"
+#include "mlx/utils.h"
+
+#include <cuda_runtime.h>
+#include <fmt/format.h>
+
+#include <cassert>
+#include <fstream>
+#include <string>
+
+namespace mlx::core {
+
+namespace cu {
+
+constexpr int page_size = 16384;
+
+// Any allocations smaller than this will try to use the small pool
+constexpr int small_block_size = 8;
+
+// The small pool size in bytes. This should be a multiple of the host page
+// size and small_block_size.
+constexpr int small_pool_size = 4 * page_size;
+
+// Check if running on Windows or Windows Subsystem for Linux
+bool is_windows() {
+#if defined(_WIN32)
+  return true;
+#elif defined(__linux__)
+  // WSL kernels contain "microsoft" or "WSL" in /proc/version
+  static bool is_wsl = []() {
+    std::ifstream version("/proc/version");
+    if (version.is_open()) {
+      std::string line;
+      std::getline(version, line);
+      return line.find("microsoft") != std::string::npos ||
+          line.find("Microsoft") != std::string::npos ||
+          line.find("WSL") != std::string::npos;
+    }
+    return false;
+  }();
+  return is_wsl;
+#else
+  return false;
+#endif
+}
+
+bool supports_managed_memory() {
+  static bool managed_memory = []() {
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      auto& d = cu::device(i);
+      if (!d.managed_memory()) {
+        return false;
+      }
+      // Empirically on Windows (and WSL) if there is no concurrentManagedAccess
+      // the managed memory also does not work.
+      if (is_windows() && !d.concurrent_managed_access()) {
+        return false;
+      }
+    }
+    return true;
+  }();
+  return managed_memory;
+}
+
+inline void* unified_malloc(size_t size) {
+  void* data = nullptr;
+  if (supports_managed_memory()) {
+    CHECK_CUDA_ERROR(cudaMallocManaged(&data, size));
+  } else {
+    CHECK_CUDA_ERROR(cudaMallocHost(&data, size));
+  }
+  return data;
+}
+
+inline void unified_free(void* data) {
+  if (supports_managed_memory()) {
+    CHECK_CUDA_ERROR(cudaFree(data));
+  } else {
+    CHECK_CUDA_ERROR(cudaFreeHost(data));
+  }
+}
+
+#if CUDART_VERSION >= 13000
+inline cudaMemLocation cuda_mem_loc(int i) {
+  cudaMemLocation loc;
+  loc.type = cudaMemLocationTypeDevice;
+  loc.id = i;
+  return loc;
+}
+#else
+inline int cuda_mem_loc(int i) {
+  return i;
+}
+#endif // CUDART_VERSION >= 13000
+
+SmallSizePool::SmallSizePool() {
+  auto num_blocks = small_pool_size / small_block_size;
+  buffer_ = new Block[num_blocks];
+  next_free_ = buffer_;
+
+  data_ = unified_malloc(small_pool_size);
+  if (supports_managed_memory()) {
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      if (device(i).concurrent_managed_access()) {
+        auto loc = cuda_mem_loc(i);
+        CHECK_CUDA_ERROR(cudaMemAdvise(
+            data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
+      }
+    }
+  }
+
+  auto curr = next_free_;
+  for (size_t i = 1; i < num_blocks; ++i) {
+    curr->next = buffer_ + i;
+    curr = curr->next;
+  }
+  curr->next = nullptr;
+}
+
+SmallSizePool::~SmallSizePool() {
+  unified_free(data_);
+  delete[] buffer_;
+}
+
+CudaBuffer* SmallSizePool::malloc() {
+  if (next_free_ == nullptr) {
+    return nullptr;
+  }
+  Block* b = next_free_;
+  uint64_t i = next_free_ - buffer_;
+  next_free_ = next_free_->next;
+  b->buf.data = static_cast<char*>(data_) + i * small_block_size;
+  b->buf.size = small_block_size;
+  b->buf.device = -1;
+  return &b->buf;
+}
+
+void SmallSizePool::free(CudaBuffer* buf) {
+  auto b = reinterpret_cast<Block*>(buf);
+  b->next = next_free_;
+  next_free_ = b;
+}
+
+bool SmallSizePool::in_pool(CudaBuffer* buf) {
+  constexpr int num_blocks = (small_pool_size / small_block_size);
+  auto b = reinterpret_cast<Block*>(buf);
+  int64_t block_num = b - buffer_;
+  return block_num >= 0 && block_num < num_blocks;
+}
+
+CudaAllocator::CudaAllocator()
+    : buffer_cache_(
+          page_size,
+          [](CudaBuffer* buf) { return buf->size; },
+          [this](CudaBuffer* buf) { free_cuda_buffer(buf); }) {
+  size_t free;
+  CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total_memory_));
+  memory_limit_ = total_memory_ * 0.95;
+  free_limit_ = total_memory_ - memory_limit_;
+  max_pool_size_ = memory_limit_;
+
+  int curr;
+  CHECK_CUDA_ERROR(cudaGetDevice(&curr));
+
+  int device_count = gpu::device_count();
+  free_streams_.resize(device_count);
+  mem_pools_.resize(device_count);
+  for (int i = 0; i < device_count; ++i) {
+    auto& d = device(i);
+    if (d.memory_pools()) {
+      free_streams_[i] = CudaStream(d);
+      CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pools_[i], i));
+    }
+  }
+  CHECK_CUDA_ERROR(cudaSetDevice(curr));
+}
+
+Buffer
+CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
+  if (size == 0) {
+    return Buffer{new CudaBuffer{nullptr, 0, -1}};
+  }
+
+  if (size <= small_block_size) {
+    size = 8;
+  } else if (size < page_size) {
+    size = next_power_of_2(size);
+  } else {
+    size = page_size * ((size + page_size - 1) / page_size);
+  }
+
+  if (size <= small_block_size || stream == nullptr) {
+    device = -1;
+  }
+
+  // Find available buffer from cache.
+  std::unique_lock lock(mutex_);
+  CudaBuffer* buf = buffer_cache_.reuse_from_cache(size);
+  if (!buf) {
+    // If we have a lot of memory pressure try to reclaim memory from the cache.
+    int64_t mem_to_free =
+        get_active_memory() + get_cache_memory() + size - memory_limit_;
+    if (mem_to_free > 0) {
+      buffer_cache_.release_cached_buffers(mem_to_free);
+    }
+
+    // Try the scalar pool first
+    if (size <= small_block_size) {
+      buf = scalar_pool_.malloc();
+    }
+    lock.unlock();
+    if (!buf) {
+      void* data = nullptr;
+      if (device == -1) {
+        data = unified_malloc(size);
+      } else {
+        if (mem_pools_[device]) { // supports memory pools
+          CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
+        } else {
+          CHECK_CUDA_ERROR(cudaMalloc(&data, size));
+        }
+      }
+      if (!data) {
+        std::ostringstream msg;
+        msg << "[malloc] Unable to allocate " << size << " bytes.";
+        throw std::runtime_error(msg.str());
+      }
+      buf = new CudaBuffer{data, size, device};
+    }
+    lock.lock();
+
+    // If any cuda memory pool has too much reserved memory, clear some
+    // memory from the cache. This prevents graph / kernel execution failing
+    // from OOM
+    if (get_cache_memory() > 0) {
+      for (auto p : mem_pools_) {
+        if (p) {
+          size_t used = 0;
+          CHECK_CUDA_ERROR(cudaMemPoolGetAttribute(
+              p, cudaMemPoolAttrReservedMemCurrent, &used));
+          if (used > (total_memory_ - free_limit_)) {
+            buffer_cache_.release_cached_buffers(free_limit_);
+            break;
+          }
+        }
+      }
+    }
+  }
+  active_memory_ += buf->size;
+  peak_memory_ = std::max(active_memory_, peak_memory_);
+
+  // Maintain the cache below the requested limit.
+  if (get_cache_memory() > max_pool_size_) {
+    buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_);
+  }
+  lock.unlock();
+  // Copy to unified memory here if the buffer is not on the right device.
+  if (buf->device >= 0 && buf->device != device) {
+    move_to_unified_memory(*buf, stream);
+  }
+  return Buffer{buf};
+}
+
+Buffer CudaAllocator::malloc(size_t size) {
+  return malloc_async(size, -1, nullptr);
+}
+
+void CudaAllocator::free(Buffer buffer) {
+  auto* buf = static_cast<CudaBuffer*>(buffer.ptr());
+  if (!buf) {
+    return;
+  }
+  if (buf->size == 0) {
+    delete buf;
+    return;
+  }
+
+  std::unique_lock lock(mutex_);
+  active_memory_ -= buf->size;
+  if (get_cache_memory() < max_pool_size_) {
+    buffer_cache_.recycle_to_cache(buf);
+  } else {
+    free_cuda_buffer(buf);
+  }
+}
+
+size_t CudaAllocator::size(Buffer buffer) const {
+  auto* buf = static_cast<CudaBuffer*>(buffer.ptr());
+  if (!buf) {
+    return 0;
+  }
+  return buf->size;
+}
+
+void CudaAllocator::move_to_unified_memory(
+    CudaBuffer& buf,
+    cudaStream_t stream) {
+  if (buf.device == -1) {
+    return;
+  }
+  void* data = unified_malloc(buf.size);
+  cudaMemcpyKind kind =
+      supports_managed_memory() ? cudaMemcpyDefault : cudaMemcpyDeviceToHost;
+  if (stream && mem_pools_[buf.device]) {
+    CHECK_CUDA_ERROR(cudaMemcpyAsync(data, buf.data, buf.size, kind, stream));
+    free_async(buf, stream);
+  } else {
+    CHECK_CUDA_ERROR(cudaMemcpy(data, buf.data, buf.size, kind));
+    free_async(buf);
+  }
+  buf.data = data;
+  buf.device = -1;
+}
+
+// This must be called with mutex_ aquired
+void CudaAllocator::free_cuda_buffer(CudaBuffer* buf) {
+  if (scalar_pool_.in_pool(buf)) {
+    scalar_pool_.free(buf);
+  } else {
+    free_async(*buf);
+    delete buf;
+  }
+}
+
+void CudaAllocator::free_async(CudaBuffer& buf, cudaStream_t stream) {
+  if (buf.device == -1) {
+    unified_free(buf.data);
+  } else {
+    // Free asynchronously when memory pools is supported.
+    if (mem_pools_[buf.device]) {
+      if (!stream) {
+        stream = free_streams_[buf.device];
+      }
+      CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream));
+    } else {
+      CHECK_CUDA_ERROR(cudaFree(buf.data));
+    }
+  }
+}
+
+size_t CudaAllocator::get_active_memory() const {
+  return active_memory_;
+}
+
+size_t CudaAllocator::get_peak_memory() const {
+  return peak_memory_;
+}
+
+void CudaAllocator::reset_peak_memory() {
+  std::lock_guard lock(mutex_);
+  peak_memory_ = 0;
+}
+
+size_t CudaAllocator::get_memory_limit() {
+  return memory_limit_;
+}
+
+size_t CudaAllocator::set_memory_limit(size_t limit) {
+  std::lock_guard lock(mutex_);
+  std::swap(limit, memory_limit_);
+  return limit;
+}
+
+size_t CudaAllocator::get_cache_memory() const {
+  return buffer_cache_.cache_size();
+}
+
+size_t CudaAllocator::set_cache_limit(size_t limit) {
+  std::lock_guard lk(mutex_);
+  std::swap(limit, max_pool_size_);
+  return limit;
+}
+
+void CudaAllocator::clear_cache() {
+  std::lock_guard lk(mutex_);
+  buffer_cache_.clear();
+}
+
+CudaAllocator& allocator() {
+  static auto* allocator_ = []() {
+    // Ensure scheduler is created before allocator.
+    scheduler::scheduler();
+    // By creating the |allocator_| on heap, the destructor of CudaAllocator
+    // will not be called on exit and buffers in the cache will be leaked. This
+    // can save some time at program exit.
+    return new CudaAllocator();
+  }();
+  return *allocator_;
+}
+
+Buffer malloc_async(size_t size, CommandEncoder& encoder) {
+  return allocator().malloc_async(
+      size, encoder.device().cuda_device(), encoder.stream());
+}
+
+} // namespace cu
+
+namespace allocator {
+
+Allocator& allocator() {
+  return cu::allocator();
+}
+
+void* Buffer::raw_ptr() {
+  if (!ptr_) {
+    return nullptr;
+  }
+  auto& cbuf = *static_cast<cu::CudaBuffer*>(ptr_);
+  cu::allocator().move_to_unified_memory(cbuf);
+  return cbuf.data;
+}
+
+} // namespace allocator
+
+size_t get_active_memory() {
+  return cu::allocator().get_active_memory();
+}
+size_t get_peak_memory() {
+  return cu::allocator().get_peak_memory();
+}
+void reset_peak_memory() {
+  return cu::allocator().reset_peak_memory();
+}
+size_t set_memory_limit(size_t limit) {
+  return cu::allocator().set_memory_limit(limit);
+}
+size_t get_memory_limit() {
+  return cu::allocator().get_memory_limit();
+}
+size_t get_cache_memory() {
+  return cu::allocator().get_cache_memory();
+}
+size_t set_cache_limit(size_t limit) {
+  return cu::allocator().set_cache_limit(limit);
+}
+void clear_cache() {
+  cu::allocator().clear_cache();
+}
+
+// Not supported in CUDA.
+size_t set_wired_limit(size_t) {
+  return 0;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/allocator.h b/macos/third_party/mlx/mlx/backend/cuda/allocator.h
new file mode 100644
index 0000000..af76ad9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/allocator.h
@@ -0,0 +1,94 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/buffer_cache.h"
+#include "mlx/backend/cuda/cuda_utils.h"
+
+#include <cuda_runtime.h>
+#include <mutex>
+#include <set>
+#include <utility>
+
+namespace mlx::core::cu {
+
+class CommandEncoder;
+
+using allocator::Buffer;
+
+// Stores cuda-managed unified memory.
+struct CudaBuffer {
+  void* data;
+  size_t size;
+  int device; // -1 for managed
+};
+
+class SmallSizePool {
+ private:
+  union Block {
+    Block* next;
+    CudaBuffer buf;
+  };
+
+  Block* buffer_{nullptr};
+  void* data_{nullptr};
+  Block* next_free_{nullptr};
+
+ public:
+  SmallSizePool();
+  ~SmallSizePool();
+
+  SmallSizePool(const SmallSizePool&) = delete;
+  SmallSizePool& operator=(const SmallSizePool&) = delete;
+
+  CudaBuffer* malloc();
+  void free(CudaBuffer* buf);
+  bool in_pool(CudaBuffer* buf);
+};
+
+class CudaAllocator : public allocator::Allocator {
+ public:
+  Buffer malloc(size_t size) override;
+  Buffer malloc_async(size_t size, int device, cudaStream_t stream);
+  void free(Buffer buffer) override;
+  size_t size(Buffer buffer) const override;
+
+  // Replace the memory of |buf| with unified memory (managed memory or pinned
+  // host memory), and copy the data over. Pass |stream| to copy asynchronously.
+  void move_to_unified_memory(CudaBuffer& buf, cudaStream_t stream = nullptr);
+
+  size_t get_active_memory() const;
+  size_t get_peak_memory() const;
+  void reset_peak_memory();
+  size_t get_memory_limit();
+  size_t set_memory_limit(size_t limit);
+  size_t get_cache_memory() const;
+  size_t set_cache_limit(size_t limit);
+  void clear_cache();
+
+ private:
+  void free_cuda_buffer(CudaBuffer* buf);
+  void free_async(CudaBuffer& buf, cudaStream_t stream = nullptr);
+
+  CudaAllocator();
+  friend CudaAllocator& allocator();
+
+  std::mutex mutex_;
+  size_t memory_limit_;
+  size_t free_limit_;
+  size_t total_memory_;
+  size_t max_pool_size_;
+  BufferCache<CudaBuffer> buffer_cache_;
+  size_t active_memory_{0};
+  size_t peak_memory_{0};
+  std::vector<CudaStream> free_streams_;
+  std::vector<cudaMemPool_t> mem_pools_;
+  SmallSizePool scalar_pool_;
+};
+
+CudaAllocator& allocator();
+
+Buffer malloc_async(size_t size, CommandEncoder& encoder);
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/arange.cu b/macos/third_party/mlx/mlx/backend/cuda/arange.cu
new file mode 100644
index 0000000..f4f4326
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/arange.cu
@@ -0,0 +1,68 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/fp16_math.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, typename IdxT, int N_WRITES>
+__global__ void arange(T* out, IdxT size, T start, T step) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_WRITES > size) {
+    for (IdxT i = index * N_WRITES; i < size; ++i) {
+      out[i] = start + i * step;
+    }
+  } else {
+    AlignedVector<T, N_WRITES> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_WRITES; ++i) {
+      out_vec[i] = start + (index * N_WRITES + i) * step;
+    }
+
+    store_vector<N_WRITES>(out, index, out_vec);
+  }
+}
+
+} // namespace cu
+
+void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Arange::eval_gpu");
+  if (out.size() == 0) {
+    return;
+  }
+  auto& encoder = cu::get_command_encoder(stream());
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  encoder.set_output_array(out);
+
+  dispatch_int_float_types(out.dtype(), "Arange", [&](auto type_tag) {
+    using CTYPE = MLX_GET_TYPE(type_tag);
+    using OutType = cuda_type_t<CTYPE>;
+    constexpr int N_WRITES = 16 / sizeof(OutType);
+    dispatch_bool(out.data_size() > INT32_MAX, [&](auto large) {
+      using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+      auto [num_blocks, block_dims] = get_launch_args(out, large(), N_WRITES);
+      encoder.add_kernel_node(
+          cu::arange<OutType, IdxT, N_WRITES>,
+          num_blocks,
+          block_dims,
+          0,
+          gpu_ptr<OutType>(out),
+          out.data_size(),
+          static_cast<CTYPE>(start_),
+          static_cast<CTYPE>(start_ + step_) - static_cast<CTYPE>(start_));
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/arg_reduce.cu b/macos/third_party/mlx/mlx/backend/cuda/arg_reduce.cu
new file mode 100644
index 0000000..bc891f5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/arg_reduce.cu
@@ -0,0 +1,189 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/fp16_math.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_reduce.cuh>
+
+#include <cassert>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T>
+struct IndexValPair {
+  uint32_t index;
+  T val;
+};
+
+template <typename T>
+struct ArgMin {
+  constexpr __device__ T init() {
+    return Limits<T>::max();
+  }
+
+  __device__ IndexValPair<T> operator()(
+      const IndexValPair<T>& best,
+      const IndexValPair<T>& current) {
+    if (best.val > current.val ||
+        (best.val == current.val && best.index > current.index)) {
+      return current;
+    } else {
+      return best;
+    }
+  }
+
+  template <int N>
+  __device__ IndexValPair<T> reduce_many(
+      IndexValPair<T> best,
+      const AlignedVector<T, N>& vals,
+      uint32_t offset) {
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      if (vals[i] < best.val) {
+        best.val = vals[i];
+        best.index = offset + i;
+      }
+    }
+    return best;
+  }
+};
+
+template <typename T>
+struct ArgMax {
+  constexpr __device__ T init() {
+    return Limits<T>::min();
+  }
+
+  __device__ IndexValPair<T> operator()(
+      const IndexValPair<T>& best,
+      const IndexValPair<T>& current) {
+    if (best.val < current.val ||
+        (best.val == current.val && best.index > current.index)) {
+      return current;
+    } else {
+      return best;
+    }
+  }
+
+  template <int N>
+  __device__ IndexValPair<T> reduce_many(
+      IndexValPair<T> best,
+      const AlignedVector<T, N>& vals,
+      uint32_t offset) {
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      if (vals[i] > best.val) {
+        best.val = vals[i];
+        best.index = offset + i;
+      }
+    }
+    return best;
+  }
+};
+
+template <typename T, typename Op, int BLOCK_DIM, int N_READS = 4>
+__global__ void arg_reduce_general(
+    const T* in,
+    uint32_t* out,
+    size_t size,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides in_strides,
+    const __grid_constant__ Strides out_strides,
+    int32_t ndim,
+    int64_t axis_stride,
+    int32_t axis_size) {
+  auto block = cg::this_thread_block();
+
+  int64_t index = cg::this_grid().block_rank();
+  if (index >= size) {
+    return;
+  }
+
+  int64_t in_idx = elem_to_loc(index, shape.data(), in_strides.data(), ndim);
+  int64_t out_idx = elem_to_loc(index, shape.data(), out_strides.data(), ndim);
+  in += in_idx;
+
+  Op op;
+  T init = op.init();
+  IndexValPair<T> best{0, init};
+
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto tid = r * BLOCK_DIM + block.thread_index().x;
+    auto vals = load_vector<N_READS>(in, tid, axis_size, axis_stride, init);
+    best = op.reduce_many(best, vals, tid * N_READS);
+  }
+
+  typedef cub::BlockReduce<IndexValPair<T>, BLOCK_DIM> BlockReduceT;
+  __shared__ typename BlockReduceT::TempStorage temp;
+
+  best = BlockReduceT(temp).Reduce(best, op);
+
+  if (block.thread_rank() == 0) {
+    out[out_idx] = best.index;
+  }
+}
+
+} // namespace cu
+
+void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("ArgReduce::eval_gpu");
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  // Prepare the shapes, strides and axis arguments.
+  Shape shape = remove_index(in.shape(), axis_);
+  Strides in_strides = remove_index(in.strides(), axis_);
+  Strides out_strides = out.ndim() == in.ndim()
+      ? remove_index(out.strides(), axis_)
+      : out.strides();
+  int64_t axis_stride = in.strides()[axis_];
+  int32_t axis_size = in.shape()[axis_];
+  int32_t ndim = shape.size();
+
+  // ArgReduce.
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_real_types(in.dtype(), "ArgReduce", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr uint32_t N_READS = 4;
+    dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
+      dim3 num_blocks = get_2d_grid_dims(out.shape(), out.strides());
+      auto kernel =
+          cu::arg_reduce_general<T, cu::ArgMax<T>, block_dim(), N_READS>;
+      if (reduce_type_ == ArgReduce::ArgMin) {
+        kernel = cu::arg_reduce_general<T, cu::ArgMin<T>, block_dim(), N_READS>;
+      }
+      encoder.add_kernel_node(
+          kernel,
+          num_blocks,
+          block_dim(),
+          0,
+          gpu_ptr<T>(in),
+          gpu_ptr<uint32_t>(out),
+          out.size(),
+          const_param(shape),
+          const_param(in_strides),
+          const_param(out_strides),
+          ndim,
+          axis_stride,
+          axis_size);
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/bin2h.cmake b/macos/third_party/mlx/mlx/backend/cuda/bin2h.cmake
new file mode 100644
index 0000000..b791d3d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/bin2h.cmake
@@ -0,0 +1,150 @@
+# Based on: https://github.com/sivachandran/cmake-bin2h
+#
+# Copyright 2020 Sivachandran Paramasivam
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+include(CMakeParseArguments)
+
+# Function to wrap a given string into multiple lines at the given column
+# position.
+#
+# Parameters:
+#
+# * VARIABLE - The name of the CMake variable holding the string.
+# * AT_COLUMN - The column position at which string will be wrapped.
+function(WRAP_STRING)
+  set(oneValueArgs VARIABLE AT_COLUMN)
+  cmake_parse_arguments(WRAP_STRING "${options}" "${oneValueArgs}" "" ${ARGN})
+
+  string(LENGTH ${${WRAP_STRING_VARIABLE}} stringLength)
+  math(EXPR offset "0")
+
+  while(stringLength GREATER 0)
+    if(stringLength GREATER ${WRAP_STRING_AT_COLUMN})
+      math(EXPR length "${WRAP_STRING_AT_COLUMN}")
+    else()
+      math(EXPR length "${stringLength}")
+    endif()
+
+    string(SUBSTRING ${${WRAP_STRING_VARIABLE}} ${offset} ${length} line)
+    set(lines "${lines}\n ${line}")
+
+    math(EXPR stringLength "${stringLength} - ${length}")
+    math(EXPR offset "${offset} + ${length}")
+  endwhile()
+
+  set(${WRAP_STRING_VARIABLE}
+      "${lines}"
+      PARENT_SCOPE)
+endfunction()
+
+# Function to embed contents of a file as byte array in C/C++ header file(.h).
+# The header file will contain a byte array and integer variable holding the
+# size of the array.
+#
+# Parameters:
+#
+# * SOURCE_FILES - The paths of source files whose contents will be embedded in
+#   the header file.
+# * VARIABLE_NAME - The name of the variable for the byte array. The string
+#   "_SIZE" will be append to this name and will be used a variable name for
+#   size variable.
+# * HEADER_FILE - The path of header file.
+# * APPEND - If specified appends to the header file instead of overwriting it
+# * HEADER_NAMESPACE - The namespace, where the array should be located in.
+# * NULL_TERMINATE - If specified a null byte(zero) will be append to the byte
+#   array.
+#
+# Usage:
+#
+# bin2h(SOURCE_FILE "Logo.png" HEADER_FILE "Logo.h" VARIABLE_NAME "LOGO_PNG")
+function(BIN2H)
+  set(options APPEND NULL_TERMINATE)
+  set(oneValueArgs VARIABLE_NAME HEADER_FILE HEADER_NAMESPACE)
+  set(multiValueArgs SOURCE_FILES)
+  cmake_parse_arguments(BIN2H "${options}" "${oneValueArgs}"
+                        "${multiValueArgs}" ${ARGN})
+
+  set(arrayDefinition "")
+  foreach(SOURCE_FILE IN LISTS BIN2H_SOURCE_FILES)
+    # get filename without extension
+    get_filename_component(FILE_NAME_WE ${SOURCE_FILE} NAME_WE)
+    # convert the filename to a valid C identifier
+    string(MAKE_C_IDENTIFIER "${FILE_NAME_WE}" VALID_FILE_NAME)
+
+    # reads source file contents as hex string
+    file(READ ${SOURCE_FILE} hexString HEX)
+
+    # append null
+    if(BIN2H_NULL_TERMINATE)
+      string(APPEND hexString "00")
+    endif()
+
+    # wraps the hex string into multiple lines
+    wrap_string(VARIABLE hexString AT_COLUMN 24)
+
+    # strip the © in source code
+    string(REGEX REPLACE "c2a9" "2020" arrayValues ${hexString})
+
+    string(REGEX REPLACE "([0-9a-f][0-9a-f])" " 0x\\1," arrayValues
+                         ${arrayValues})
+
+    # make a full variable name for the array
+    set(FULL_VARIABLE_NAME "${BIN2H_VARIABLE_NAME}_${VALID_FILE_NAME}")
+
+    # declares byte array and the length variables
+    string(APPEND arrayDefinition
+           "constexpr char ${FULL_VARIABLE_NAME}[] = {${arrayValues}\n};\n\n")
+  endforeach()
+
+  # add namespace wrapper if defined
+  if(DEFINED BIN2H_HEADER_NAMESPACE)
+    set(namespaceStart "namespace ${BIN2H_HEADER_NAMESPACE} {")
+    set(namespaceEnd "} // namespace ${BIN2H_HEADER_NAMESPACE}")
+    set(declarations "${namespaceStart}\n\n${arrayDefinition}${namespaceEnd}\n")
+  endif()
+
+  set(arrayIncludes "#pragma once")
+  string(PREPEND declarations "${arrayIncludes}\n\n")
+
+  if(BIN2H_APPEND)
+    file(APPEND ${BIN2H_HEADER_FILE} "${declarations}")
+  else()
+    file(WRITE ${BIN2H_HEADER_FILE} "${declarations}")
+  endif()
+endfunction()
+
+# ----------------------------- CLI args -----------------------------
+
+string(REPLACE ":" ";" MLX_JIT_SOURCES_LIST ${MLX_JIT_SOURCES})
+foreach(source ${MLX_JIT_SOURCES_LIST})
+  list(APPEND MLX_JIT_SOURCES_ABS "${MLX_SOURCE_ROOT}/${source}")
+endforeach()
+
+bin2h(
+  SOURCE_FILES
+  ${MLX_JIT_SOURCES_ABS}
+  NULL_TERMINATE
+  VARIABLE_NAME
+  "jit_source"
+  HEADER_NAMESPACE
+  "mlx::core"
+  HEADER_FILE
+  "${CMAKE_CURRENT_BINARY_DIR}/gen/cuda_jit_sources.h")
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/cuda/binary/CMakeLists.txt
new file mode 100644
index 0000000..bda289d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/CMakeLists.txt
@@ -0,0 +1,21 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/add.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arctan2.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/bitwise_binary.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/divide.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/equal.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/greater.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/greater_equal.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/less.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/less_equal.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/logical_and.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/logical_or.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/log_add_exp.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/minimum.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/maximum.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/multiply.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/power.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/remainder.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/not_equal.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/subtract.cu)
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/add.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/add.cu
new file mode 100644
index 0000000..87dfd7e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/add.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Add)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/arctan2.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/arctan2.cu
new file mode 100644
index 0000000..2fd7e39
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/arctan2.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(ArcTan2)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/binary.cuh b/macos/third_party/mlx/mlx/backend/cuda/binary/binary.cuh
new file mode 100644
index 0000000..d43027c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/binary.cuh
@@ -0,0 +1,383 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/binary.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/binary_ops.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (int i = index * N_READS; i < size; ++i) {
+      out[i] = Op{}(a[0], b[0]);
+    }
+  } else {
+    AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = Op{}(a[0], b[0]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      out[i] = Op{}(a[0], b[i]);
+    }
+  } else {
+    auto b_vec = load_vector<N_READS>(b, index);
+
+    AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = Op{}(a[0], b_vec[i]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      out[i] = Op{}(a[i], b[0]);
+    }
+  } else {
+    auto a_vec = load_vector<N_READS>(a, index);
+
+    AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = Op{}(a_vec[i], b[0]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void binary_vv(const In* a, const In* b, Out* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      out[i] = Op{}(a[i], b[i]);
+    }
+  } else {
+    auto a_vec = load_vector<N_READS>(a, index);
+    auto b_vec = load_vector<N_READS>(b, index);
+
+    AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = Op{}(a_vec[i], b_vec[i]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+template <
+    typename Op,
+    typename In,
+    typename Out,
+    typename IdxT,
+    int NDIM,
+    int N_READS>
+__global__ void binary_g_nd(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size_rest,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> a_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> b_strides) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[NDIM - 1];
+  auto a_stride_x = a_strides[NDIM - 1];
+  auto b_stride_x = b_strides[NDIM - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [a_idx, b_idx] = elem_to_loc_nd<NDIM>(
+      index_rest * shape_x, shape.data(), a_strides.data(), b_strides.data());
+  auto a_vec =
+      load_vector<N_READS>(a + a_idx, index_x, shape_x, a_stride_x, In(0));
+  auto b_vec =
+      load_vector<N_READS>(b + b_idx, index_x, shape_x, b_stride_x, In(0));
+
+  AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = Op{}(a_vec[i], b_vec[i]);
+  }
+  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void binary_g(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size_rest,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides a_strides,
+    const __grid_constant__ Strides b_strides,
+    int ndim) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[ndim - 1];
+  auto a_stride_x = a_strides[ndim - 1];
+  auto b_stride_x = b_strides[ndim - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [a_idx, b_idx] = elem_to_loc(
+      index_rest * shape_x,
+      shape.data(),
+      a_strides.data(),
+      b_strides.data(),
+      ndim);
+  auto a_vec =
+      load_vector<N_READS>(a + a_idx, index_x, shape_x, a_stride_x, In(0));
+  auto b_vec =
+      load_vector<N_READS>(b + b_idx, index_x, shape_x, b_stride_x, In(0));
+
+  AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = Op{}(a_vec[i], b_vec[i]);
+  }
+  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
+}
+
+template <typename Op, typename In, typename Out>
+constexpr bool supports_binary_op() {
+  if (std::is_same_v<Op, Add> || std::is_same_v<Op, Divide> ||
+      std::is_same_v<Op, Maximum> || std::is_same_v<Op, Minimum> ||
+      std::is_same_v<Op, Multiply> || std::is_same_v<Op, Subtract> ||
+      std::is_same_v<Op, Power> || std::is_same_v<Op, Remainder>) {
+    return std::is_same_v<In, Out>;
+  }
+  if (std::is_same_v<Op, Equal> || std::is_same_v<Op, Greater> ||
+      std::is_same_v<Op, GreaterEqual> || std::is_same_v<Op, Less> ||
+      std::is_same_v<Op, LessEqual> || std::is_same_v<Op, NotEqual>) {
+    return std::is_same_v<Out, bool>;
+  }
+  if (std::is_same_v<Op, LogicalAnd> || std::is_same_v<Op, LogicalOr>) {
+    return std::is_same_v<Out, bool> && std::is_same_v<In, bool>;
+  }
+  if (std::is_same_v<Op, NaNEqual>) {
+    return std::is_same_v<Out, bool> && is_inexact_v<In>;
+  }
+  if (std::is_same_v<Op, LogAddExp>) {
+    return std::is_same_v<In, Out> && is_inexact_v<In>;
+  }
+  if (std::is_same_v<Op, ArcTan2>) {
+    return std::is_same_v<In, Out> && is_floating_v<In>;
+  }
+  if (std::is_same_v<Op, BitwiseAnd> || std::is_same_v<Op, BitwiseOr> ||
+      std::is_same_v<Op, BitwiseXor>) {
+    return std::is_same_v<In, Out> && std::is_integral_v<In>;
+  }
+  if (std::is_same_v<Op, LeftShift> || std::is_same_v<Op, RightShift>) {
+    return std::is_same_v<In, Out> && std::is_integral_v<In> &&
+        !std::is_same_v<In, bool>;
+  }
+  return false;
+}
+
+} // namespace cu
+
+template <typename Op>
+void binary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  assert(inputs.size() > 1);
+  const auto& a = inputs[0];
+  const auto& b = inputs[1];
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cu::get_command_encoder(s);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  dispatch_all_types(a.dtype(), [&](auto in_type_tag) {
+    dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
+      using CTYPE_IN = MLX_GET_TYPE(in_type_tag);
+      using CTYPE_OUT = MLX_GET_TYPE(out_type_tag);
+      if constexpr (cu::supports_binary_op<Op, CTYPE_IN, CTYPE_OUT>()) {
+        using InType = cuda_type_t<CTYPE_IN>;
+        using OutType = cuda_type_t<CTYPE_OUT>;
+        auto bopt = get_binary_op_type(a, b);
+        if (bopt == BinaryOpType::General) {
+          dispatch_bool(
+              a.data_size() > INT32_MAX || b.data_size() > INT32_MAX ||
+                  out.data_size() > INT32_MAX,
+              [&](auto large) {
+                using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+                Shape shape;
+                std::vector<Strides> strides;
+                std::tie(shape, strides) = collapse_contiguous_dims(a, b, out);
+                auto& a_strides = strides[0];
+                auto& b_strides = strides[1];
+                int ndim = shape.size();
+                int work_per_thread = 1;
+                auto dim0 = ndim > 0 ? shape.back() : 1;
+                auto rest = out.size() / dim0;
+                if (dim0 >= 4) {
+                  work_per_thread = 4;
+                }
+                dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+                auto block_dims = get_block_dims(dim0, rest, 1);
+                uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x);
+                uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y);
+                if (ndim <= 3) {
+                  dispatch_1_2_3(ndim, [&](auto dims_constant) {
+                    auto kernel = cu::binary_g_nd<
+                        Op,
+                        InType,
+                        OutType,
+                        IdxT,
+                        dims_constant(),
+                        1>;
+                    if (work_per_thread == 4) {
+                      kernel = cu::binary_g_nd<
+                          Op,
+                          InType,
+                          OutType,
+                          IdxT,
+                          dims_constant(),
+                          4>;
+                    }
+                    encoder.add_kernel_node(
+                        kernel,
+                        {num_blocks_x, num_blocks_y},
+                        block_dims,
+                        0,
+                        gpu_ptr<InType>(a),
+                        gpu_ptr<InType>(b),
+                        gpu_ptr<OutType>(out),
+                        rest,
+                        const_param<dims_constant()>(shape),
+                        const_param<dims_constant()>(a_strides),
+                        const_param<dims_constant()>(b_strides));
+                  });
+                } else {
+                  auto kernel = cu::binary_g<Op, InType, OutType, IdxT, 1>;
+                  if (work_per_thread == 4) {
+                    kernel = cu::binary_g<Op, InType, OutType, IdxT, 4>;
+                  }
+                  encoder.add_kernel_node(
+                      kernel,
+                      {num_blocks_x, num_blocks_y},
+                      block_dims,
+                      0,
+                      gpu_ptr<InType>(a),
+                      gpu_ptr<InType>(b),
+                      gpu_ptr<OutType>(out),
+                      rest,
+                      const_param(shape),
+                      const_param(a_strides),
+                      const_param(b_strides),
+                      ndim);
+                }
+              });
+        } else {
+          dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
+            using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
+            constexpr int N_READS = 16 / sizeof(InType);
+            auto kernel = cu::binary_ss<Op, InType, OutType, IdxT, N_READS>;
+            if (bopt == BinaryOpType::ScalarVector) {
+              kernel = cu::binary_sv<Op, InType, OutType, IdxT, N_READS>;
+            } else if (bopt == BinaryOpType::VectorScalar) {
+              kernel = cu::binary_vs<Op, InType, OutType, IdxT, N_READS>;
+            } else if (bopt == BinaryOpType::VectorVector) {
+              kernel = cu::binary_vv<Op, InType, OutType, IdxT, N_READS>;
+            }
+            auto [num_blocks, block_dims] = get_launch_args(
+                out.data_size(), out.shape(), out.strides(), large(), N_READS);
+            encoder.add_kernel_node(
+                kernel,
+                num_blocks,
+                block_dims,
+                0,
+                gpu_ptr<InType>(a),
+                gpu_ptr<InType>(b),
+                gpu_ptr<OutType>(out),
+                out.data_size());
+          });
+        }
+      } else {
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do binary op {} on inputs of {} with result of {}.",
+                op,
+                dtype_to_string(a.dtype()),
+                dtype_to_string(out.dtype())));
+      }
+    });
+  });
+}
+
+template <typename Op>
+void binary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto bopt = get_binary_op_type(a, b);
+  auto& encoder = cu::get_command_encoder(s);
+
+  set_binary_op_output_data(
+      a, b, out, bopt, [&](auto n) { return cu::malloc_async(n, encoder); });
+  binary_op_gpu_inplace<Op>(inputs, out, op, s);
+}
+
+#define BINARY_GPU(func)                                              \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    nvtx3::scoped_range r(#func "::eval_gpu");                        \
+    auto& s = out.primitive().stream();                               \
+    binary_op_gpu<cu::func>(inputs, out, name(), s);                  \
+  }
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/bitwise_binary.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/bitwise_binary.cu
new file mode 100644
index 0000000..8025a3b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/bitwise_binary.cu
@@ -0,0 +1,27 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+void BitwiseBinary::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("BitwiseBinary::eval_gpu");
+  auto& s = out.primitive().stream();
+  switch (op_) {
+    case BitwiseBinary::And:
+      binary_op_gpu<cu::BitwiseAnd>(inputs, out, name(), s);
+      break;
+    case BitwiseBinary::Or:
+      binary_op_gpu<cu::BitwiseOr>(inputs, out, name(), s);
+      break;
+    case BitwiseBinary::Xor:
+      binary_op_gpu<cu::BitwiseXor>(inputs, out, name(), s);
+      break;
+    case BitwiseBinary::LeftShift:
+      binary_op_gpu<cu::LeftShift>(inputs, out, name(), s);
+      break;
+    case BitwiseBinary::RightShift:
+      binary_op_gpu<cu::RightShift>(inputs, out, name(), s);
+      break;
+  }
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/divide.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/divide.cu
new file mode 100644
index 0000000..fcf3dc7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/divide.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Divide)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/equal.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/equal.cu
new file mode 100644
index 0000000..559b3e8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/equal.cu
@@ -0,0 +1,15 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+void Equal::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Equal::eval_gpu");
+  auto& s = out.primitive().stream();
+  if (equal_nan_) {
+    binary_op_gpu<cu::NaNEqual>(inputs, out, name(), s);
+  } else {
+    binary_op_gpu<cu::Equal>(inputs, out, name(), s);
+  }
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/greater.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/greater.cu
new file mode 100644
index 0000000..c982020
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/greater.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Greater)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/greater_equal.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/greater_equal.cu
new file mode 100644
index 0000000..4666fb4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/greater_equal.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(GreaterEqual)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/less.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/less.cu
new file mode 100644
index 0000000..a2053fa
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/less.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Less)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/less_equal.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/less_equal.cu
new file mode 100644
index 0000000..7f9bc51
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/less_equal.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(LessEqual)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/log_add_exp.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/log_add_exp.cu
new file mode 100644
index 0000000..17614f8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/log_add_exp.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(LogAddExp)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/logical_and.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/logical_and.cu
new file mode 100644
index 0000000..6bbeb1a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/logical_and.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(LogicalAnd)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/logical_or.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/logical_or.cu
new file mode 100644
index 0000000..63afdb9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/logical_or.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(LogicalOr)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/maximum.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/maximum.cu
new file mode 100644
index 0000000..4f6cb6e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/maximum.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Maximum)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/minimum.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/minimum.cu
new file mode 100644
index 0000000..ec4c1ab
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/minimum.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Minimum)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/multiply.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/multiply.cu
new file mode 100644
index 0000000..bfc15fc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/multiply.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Multiply)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/not_equal.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/not_equal.cu
new file mode 100644
index 0000000..49f05c9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/not_equal.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(NotEqual)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/power.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/power.cu
new file mode 100644
index 0000000..cacdc75
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/power.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Power)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/remainder.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/remainder.cu
new file mode 100644
index 0000000..a55006b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/remainder.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Remainder)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary/subtract.cu b/macos/third_party/mlx/mlx/backend/cuda/binary/subtract.cu
new file mode 100644
index 0000000..37f3874
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary/subtract.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/binary/binary.cuh"
+
+namespace mlx::core {
+BINARY_GPU(Subtract)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/binary_two.cu b/macos/third_party/mlx/mlx/backend/cuda/binary_two.cu
new file mode 100644
index 0000000..4acd304
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/binary_two.cu
@@ -0,0 +1,412 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/binary.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/binary_ops.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void
+binary_two_ss(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      auto out = Op{}(a[0], b[0]);
+      out_a[i] = out[0];
+      out_b[i] = out[1];
+    }
+  } else {
+    AlignedVector<Out, N_READS> out_a_vec;
+    AlignedVector<Out, N_READS> out_b_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      auto out = Op{}(a[0], b[0]);
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
+    }
+
+    store_vector<N_READS>(out_a, index, out_a_vec);
+    store_vector<N_READS>(out_b, index, out_b_vec);
+  }
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void
+binary_two_sv(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      auto out = Op{}(a[0], b[i]);
+      out_a[i] = out[0];
+      out_b[i] = out[1];
+    }
+  } else {
+    auto b_vec = load_vector<N_READS>(b, index);
+
+    AlignedVector<Out, N_READS> out_a_vec;
+    AlignedVector<Out, N_READS> out_b_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      auto out = Op{}(a[0], b_vec[i]);
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
+    }
+
+    store_vector<N_READS>(out_a, index, out_a_vec);
+    store_vector<N_READS>(out_b, index, out_b_vec);
+  }
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void
+binary_two_vs(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      auto out = Op{}(a[i], b[0]);
+      out_a[i] = out[0];
+      out_b[i] = out[1];
+    }
+  } else {
+    auto a_vec = load_vector<N_READS>(a, index);
+
+    AlignedVector<Out, N_READS> out_a_vec;
+    AlignedVector<Out, N_READS> out_b_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      auto out = Op{}(a_vec[i], b[0]);
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
+    }
+
+    store_vector<N_READS>(out_a, index, out_a_vec);
+    store_vector<N_READS>(out_b, index, out_b_vec);
+  }
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void
+binary_two_vv(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      auto out = Op{}(a[i], b[i]);
+      out_a[i] = out[0];
+      out_b[i] = out[1];
+    }
+  } else {
+    auto a_vec = load_vector<N_READS>(a, index);
+    auto b_vec = load_vector<N_READS>(b, index);
+
+    AlignedVector<Out, N_READS> out_a_vec;
+    AlignedVector<Out, N_READS> out_b_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      auto out = Op{}(a_vec[i], b_vec[i]);
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
+    }
+
+    store_vector<N_READS>(out_a, index, out_a_vec);
+    store_vector<N_READS>(out_b, index, out_b_vec);
+  }
+}
+
+template <
+    typename Op,
+    typename In,
+    typename Out,
+    typename IdxT,
+    int NDIM,
+    int N_READS>
+__global__ void binary_two_g_nd(
+    const In* a,
+    const In* b,
+    Out* out_a,
+    Out* out_b,
+    IdxT size_rest,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> a_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> b_strides) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[NDIM - 1];
+  auto a_stride_x = a_strides[NDIM - 1];
+  auto b_stride_x = b_strides[NDIM - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [a_idx, b_idx] = elem_to_loc_nd<NDIM>(
+      index_rest * shape_x, shape.data(), a_strides.data(), b_strides.data());
+  auto a_vec =
+      load_vector<N_READS>(a + a_idx, index_x, shape_x, a_stride_x, In(0));
+  auto b_vec =
+      load_vector<N_READS>(b + b_idx, index_x, shape_x, b_stride_x, In(0));
+
+  AlignedVector<Out, N_READS> out_vec_a;
+  AlignedVector<Out, N_READS> out_vec_b;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    auto out = Op{}(a_vec[i], b_vec[i]);
+    out_vec_a[i] = out[0];
+    out_vec_b[i] = out[1];
+  }
+  store_vector(out_a + shape_x * index_rest, index_x, out_vec_a, shape_x);
+  store_vector(out_b + shape_x * index_rest, index_x, out_vec_b, shape_x);
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void binary_two_g(
+    const In* a,
+    const In* b,
+    Out* out_a,
+    Out* out_b,
+    IdxT size_rest,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides a_strides,
+    const __grid_constant__ Strides b_strides,
+    int ndim) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[ndim - 1];
+  auto a_stride_x = a_strides[ndim - 1];
+  auto b_stride_x = b_strides[ndim - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [a_idx, b_idx] = elem_to_loc(
+      index_rest * shape_x,
+      shape.data(),
+      a_strides.data(),
+      b_strides.data(),
+      ndim);
+  auto a_vec =
+      load_vector<N_READS>(a + a_idx, index_x, shape_x, a_stride_x, In(0));
+  auto b_vec =
+      load_vector<N_READS>(b + b_idx, index_x, shape_x, b_stride_x, In(0));
+
+  AlignedVector<Out, N_READS> out_vec_a;
+  AlignedVector<Out, N_READS> out_vec_b;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    auto out = Op{}(a_vec[i], b_vec[i]);
+    out_vec_a[i] = out[0];
+    out_vec_b[i] = out[1];
+  }
+  store_vector(out_a + shape_x * index_rest, index_x, out_vec_a, shape_x);
+  store_vector(out_b + shape_x * index_rest, index_x, out_vec_b, shape_x);
+}
+
+template <typename Op, typename In, typename Out>
+constexpr bool supports_binary_two_op() {
+  if (std::is_same_v<Op, DivMod>) {
+    return std::is_same_v<In, Out> &&
+        (std::is_integral_v<Out> || is_floating_v<Out>);
+  }
+  return false;
+}
+
+} // namespace cu
+
+template <typename Op>
+void binary_two_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const char* op,
+    const Stream& s) {
+  assert(inputs.size() > 1);
+  const auto& a = inputs[0];
+  const auto& b = inputs[1];
+  auto& out_a = outputs[0];
+  auto& out_b = outputs[1];
+  auto bopt = get_binary_op_type(a, b);
+  auto& encoder = cu::get_command_encoder(s);
+  set_binary_op_output_data(
+      a, b, out_a, bopt, [&](auto n) { return cu::malloc_async(n, encoder); });
+  set_binary_op_output_data(
+      a, b, out_b, bopt, [&](auto n) { return cu::malloc_async(n, encoder); });
+
+  if (out_a.size() == 0) {
+    return;
+  }
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out_a);
+  encoder.set_output_array(out_b);
+  dispatch_all_types(a.dtype(), [&](auto in_type_tag) {
+    dispatch_all_types(out_a.dtype(), [&](auto out_type_tag) {
+      using CTYPE_IN = MLX_GET_TYPE(in_type_tag);
+      using CTYPE_OUT = MLX_GET_TYPE(out_type_tag);
+      if constexpr (cu::supports_binary_two_op<Op, CTYPE_IN, CTYPE_OUT>()) {
+        using InType = cuda_type_t<CTYPE_IN>;
+        using OutType = cuda_type_t<CTYPE_OUT>;
+
+        auto bopt = get_binary_op_type(a, b);
+        if (bopt == BinaryOpType::General) {
+          dispatch_bool(
+              a.data_size() > INT32_MAX || b.data_size() > INT32_MAX ||
+                  out_a.data_size() > INT32_MAX,
+              [&](auto large) {
+                using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+                Shape shape;
+                std::vector<Strides> strides;
+                std::tie(shape, strides) =
+                    collapse_contiguous_dims(a, b, out_a);
+                auto& a_strides = strides[0];
+                auto& b_strides = strides[1];
+                int ndim = shape.size();
+                int work_per_thread = 1;
+                auto dim0 = ndim > 0 ? shape.back() : 1;
+                auto rest = out_a.size() / dim0;
+                if (dim0 >= 4) {
+                  work_per_thread = 4;
+                }
+                dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+                auto block_dims = get_block_dims(dim0, rest, 1);
+                uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x);
+                uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y);
+
+                if (ndim <= 3) {
+                  dispatch_1_2_3(ndim, [&](auto dims_constant) {
+                    auto kernel = cu::binary_two_g_nd<
+                        Op,
+                        InType,
+                        OutType,
+                        IdxT,
+                        dims_constant(),
+                        1>;
+                    if (work_per_thread == 4) {
+                      kernel = cu::binary_two_g_nd<
+                          Op,
+                          InType,
+                          OutType,
+                          IdxT,
+                          dims_constant(),
+                          4>;
+                    }
+                    encoder.add_kernel_node(
+                        kernel,
+                        {num_blocks_x, num_blocks_y},
+                        block_dims,
+                        0,
+                        gpu_ptr<InType>(a),
+                        gpu_ptr<InType>(b),
+                        gpu_ptr<OutType>(out_a),
+                        gpu_ptr<OutType>(out_b),
+                        rest,
+                        const_param<dims_constant()>(shape),
+                        const_param<dims_constant()>(a_strides),
+                        const_param<dims_constant()>(b_strides));
+                  });
+                } else {
+                  auto kernel = cu::binary_two_g<Op, InType, OutType, IdxT, 1>;
+                  if (work_per_thread == 4) {
+                    kernel = cu::binary_two_g<Op, InType, OutType, IdxT, 4>;
+                  }
+                  encoder.add_kernel_node(
+                      kernel,
+                      {num_blocks_x, num_blocks_y},
+                      block_dims,
+                      0,
+                      gpu_ptr<InType>(a),
+                      gpu_ptr<InType>(b),
+                      gpu_ptr<OutType>(out_a),
+                      gpu_ptr<OutType>(out_b),
+                      rest,
+                      const_param(shape),
+                      const_param(a_strides),
+                      const_param(b_strides),
+                      ndim);
+                }
+              });
+        } else {
+          dispatch_bool(out_a.data_size() > UINT32_MAX, [&](auto large) {
+            using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
+            constexpr int N_READS = 16 / sizeof(InType);
+            auto kernel = cu::binary_two_ss<Op, InType, OutType, IdxT, N_READS>;
+            if (bopt == BinaryOpType::ScalarVector) {
+              kernel = cu::binary_two_sv<Op, InType, OutType, IdxT, N_READS>;
+            } else if (bopt == BinaryOpType::VectorScalar) {
+              kernel = cu::binary_two_vs<Op, InType, OutType, IdxT, N_READS>;
+            } else if (bopt == BinaryOpType::VectorVector) {
+              kernel = cu::binary_two_vv<Op, InType, OutType, IdxT, N_READS>;
+            }
+            auto [num_blocks, block_dims] = get_launch_args(
+                out_a.data_size(),
+                out_a.shape(),
+                out_a.strides(),
+                large(),
+                N_READS);
+            encoder.add_kernel_node(
+                kernel,
+                num_blocks,
+                block_dims,
+                0,
+                gpu_ptr<InType>(a),
+                gpu_ptr<InType>(b),
+                gpu_ptr<OutType>(out_a),
+                gpu_ptr<OutType>(out_b),
+                out_a.data_size());
+          });
+        }
+      } else {
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do binary op {} on inputs of {} with result of {}.",
+                op,
+                dtype_to_string(a.dtype()),
+                dtype_to_string(out_a.dtype())));
+      }
+    });
+  });
+}
+
+template <typename Op>
+void binary_two_op_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const char* op,
+    const Stream& s) {
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, outputs[0], bopt);
+  set_binary_op_output_data(a, b, outputs[1], bopt);
+  binary_two_op_gpu_inplace<Op>(inputs, outputs, op, s);
+}
+
+void DivMod::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("DivMod::eval_gpu");
+  auto& s = outputs[0].primitive().stream();
+  binary_two_op_gpu<cu::DivMod>(inputs, outputs, name(), s);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/compiled.cpp b/macos/third_party/mlx/mlx/backend/cuda/compiled.cpp
new file mode 100644
index 0000000..29cf2e4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/compiled.cpp
@@ -0,0 +1,357 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/graph_utils.h"
+#include "mlx/primitives.h"
+
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+struct FusedKernelBuilder {
+  std::string os;
+  const std::string& kernel_name;
+  const std::vector<array>& inputs;
+  const std::vector<array>& outputs;
+  const std::vector<array>& tape;
+  const std::function<bool(size_t)>& is_constant;
+
+  void build(const char* name, bool contiguous) {
+    NodeNamer namer;
+
+    // Function parameters.
+    std::vector<std::string> params;
+    for (size_t i = 0; i < inputs.size(); ++i) {
+      if (is_constant(i)) {
+        continue;
+      }
+      const auto& x = inputs[i];
+      const std::string& xname = namer.get_name(x);
+      params.push_back(
+          fmt::format("const {}* {}", dtype_to_cuda_type(x.dtype()), xname));
+      if (!is_scalar(x) && !contiguous) {
+        params.push_back(
+            fmt::format(
+                "const __grid_constant__ cuda::std::array<int64_t, NDIM> {}_strides",
+                xname));
+      }
+    }
+    for (const auto& x : outputs) {
+      params.push_back(
+          fmt::format(
+              "{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x)));
+    }
+    if (!contiguous) {
+      params.push_back(
+          "const __grid_constant__ cuda::std::array<int32_t, NDIM> shape");
+    }
+    params.push_back("IdxT size");
+
+    // Build function signature.
+    if (contiguous) {
+      os += "template <typename IdxT = uint32_t, int work_per_thread = 1>\n";
+    } else {
+      os +=
+          "template <int NDIM, typename IdxT = uint32_t, int work_per_thread = 1>\n";
+    }
+    os += fmt::format("__global__ void {}(\n", kernel_name + name);
+    for (size_t i = 0; i < params.size(); ++i) {
+      os += "    ";
+      os += params[i];
+      if (i != params.size() - 1) {
+        os += ",\n";
+      }
+    }
+    os += ") {\n";
+
+    // Index. For non contiguous kernels we create a separate index
+    // variable per variable otherwise everyone uses `index`.
+    os +=
+        "  IdxT index = cg::this_grid().thread_rank() * work_per_thread;\n"
+        "  if (index >= size) {\n"
+        "    return;\n"
+        "  }\n";
+    if (!contiguous) {
+      for (size_t i = 0; i < inputs.size(); ++i) {
+        const auto& x = inputs[i];
+        const std::string& xname = namer.get_name(x);
+        if (is_scalar(x) || is_constant(i)) {
+          continue;
+        }
+        os += "  IdxT " + xname + "_idx = 0;\n";
+      }
+      os += "  {\n";
+      os += "    IdxT loc = index;\n";
+      os +=
+          "    #pragma unroll\n"
+          "    for (int i = NDIM - 1; i >= 0; i--) {\n";
+      for (size_t i = 0; i < inputs.size(); ++i) {
+        const auto& x = inputs[i];
+        const std::string& xname = namer.get_name(x);
+        if (is_scalar(x) || is_constant(i)) {
+          continue;
+        }
+        os += "      " + xname + "_idx += (loc \% shape[i]) * IdxT(" + xname +
+            "_strides[i]);\n";
+      }
+      os +=
+          "      loc /= shape[i];\n"
+          "    }\n"
+          "  }\n";
+    }
+
+    // Vectorized read loop
+    if (contiguous) {
+      for (size_t i = 0; i < inputs.size(); ++i) {
+        const auto& x = inputs[i];
+        if (is_scalar(x) || is_constant(i)) {
+          continue;
+        }
+        const std::string& xname = namer.get_name(x);
+        std::string type = dtype_to_cuda_type(x.dtype());
+        os += fmt::format(
+            "  auto vec_{0} = load_vector<work_per_thread, {1}>({0} + index, 0, size - index, 0);\n",
+            xname,
+            type);
+      }
+    }
+
+    // Create some space for the outputs
+    for (const auto& x : outputs) {
+      const std::string& xname = namer.get_name(x);
+      std::string type = dtype_to_cuda_type(x.dtype());
+      os += fmt::format(
+          "  AlignedVector<{}, work_per_thread> vec_{};\n", type, xname);
+    }
+
+    // Work loop
+    if (!contiguous) {
+      os +=
+          "\n"
+          "  for (int i = 0; i < work_per_thread && index < size; i++) {\n";
+    } else {
+      os +=
+          "\n"
+          "  #pragma unroll\n"
+          "  for (int i = 0; i < work_per_thread; i++) {\n";
+    }
+
+    // Read inputs.
+    for (size_t i = 0; i < inputs.size(); ++i) {
+      const auto& x = inputs[i];
+      const std::string& xname = namer.get_name(x);
+      std::string type = dtype_to_cuda_type(x.dtype());
+      std::string value;
+      if (is_constant(i)) {
+        std::ostringstream ss;
+        print_constant(ss, x);
+        value = fmt::format("static_cast<{}>({})", type, ss.str());
+      } else if (is_scalar(x)) {
+        value = fmt::format("{}[0]", xname);
+      } else if (contiguous) {
+        value = fmt::format("vec_{}[i]", xname);
+      } else {
+        value = fmt::format("{}[{}_idx]", xname, xname);
+      }
+      os += fmt::format("    {} tmp_{} = {};\n", type, xname, value);
+    }
+
+    // Write tape.
+    for (const auto& x : tape) {
+      const std::string& xname = namer.get_name(x);
+      std::string type = dtype_to_cuda_type(x.dtype());
+      std::string value;
+      if (is_static_cast(x.primitive())) {
+        value = fmt::format(
+            "static_cast<{}>(tmp_{})", type, namer.get_name(x.inputs()[0]));
+      } else {
+        value = x.primitive().name();
+        value += "{}(";
+        for (size_t i = 0; i < x.inputs().size() - 1; ++i) {
+          value += fmt::format("tmp_{}, ", namer.get_name(x.inputs()[i]));
+        }
+        value += fmt::format("tmp_{})", namer.get_name(x.inputs().back()));
+      }
+      os += fmt::format("    {} tmp_{} = {};\n", type, xname, value);
+    }
+
+    // Write output.
+    for (const auto& x : outputs) {
+      os += fmt::format("    vec_{0}[i] = tmp_{0};\n", namer.get_name(x));
+    }
+
+    // End of work loop
+    if (!contiguous) {
+      os += "\n";
+      for (size_t i = 0; i < inputs.size(); ++i) {
+        const auto& x = inputs[i];
+        const std::string& xname = namer.get_name(x);
+        if (is_scalar(x) || is_constant(i)) {
+          continue;
+        }
+        os += fmt::format("    {0}_idx += {0}_strides[NDIM - 1];\n", xname);
+      }
+    }
+    os += "  }\n";
+
+    // Store the output to global memory
+    for (const auto& x : outputs) {
+      os += fmt::format(
+          "  store_vector({0} + index, 0, vec_{0}, size - index);\n",
+          namer.get_name(x));
+    }
+
+    os += "}\n";
+  }
+};
+
+} // namespace cu
+
+constexpr const char* g_jit_includes = R"(
+#include "mlx/backend/cuda/device/binary_ops.cuh"
+#include "mlx/backend/cuda/device/ternary_ops.cuh"
+#include "mlx/backend/cuda/device/unary_ops.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+#define inf cuda::std::numeric_limits<float>::infinity()
+)";
+
+void Compiled::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("Compiled::eval_gpu");
+  auto& s = stream();
+
+  // Determine the work per thread for the vectorized reads/writes. We take it
+  // as 16 over the max itemsize for the outputs. Another heuristic could be
+  // over the max itemsize of all arrays.
+  int max_size = 1;
+  for (const auto& x : outputs) {
+    max_size = (max_size > x.itemsize()) ? max_size : x.itemsize();
+  }
+  int work_per_thread = 16 / max_size;
+
+  cu::JitModule& mod = cu::get_jit_module(s.device, lib_name(), [&]() {
+    // Build source code.
+    cu::FusedKernelBuilder builder{
+        g_jit_includes, lib_name(), inputs_, outputs_, tape_, is_constant_};
+    builder.os +=
+        "namespace mlx::core::cu {\n\n"
+        "namespace cg = cooperative_groups;\n\n";
+    builder.build("_contiguous", true);
+    builder.os += "\n";
+    builder.build("_strided", false);
+    builder.os += "\n} // namespace mlx::core::cu\n";
+    // Build kernel names.
+    std::vector<std::string> kernel_names;
+    kernel_names.push_back(
+        fmt::format(
+            "mlx::core::cu::{}_contiguous<uint32_t, {}>",
+            lib_name(),
+            work_per_thread));
+    kernel_names.push_back(
+        fmt::format(
+            "mlx::core::cu::{}_contiguous<int64_t, {}>",
+            lib_name(),
+            work_per_thread));
+    for (int wpt : {1, work_per_thread}) {
+      for (int i = 1; i <= MAX_NDIM; ++i) {
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::{}_strided<{}, uint32_t, {}>",
+                lib_name(),
+                i,
+                wpt));
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::{}_strided<{}, int64_t, {}>",
+                lib_name(),
+                i,
+                wpt));
+      }
+    }
+
+    return std::make_tuple(
+        false, std::move(builder.os), std::move(kernel_names));
+  });
+
+  // Collapse contiguous dims to route to a faster kernel if possible. Also
+  // handle all broadcasting.
+  auto [contiguous, shape, strides_vec] =
+      compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
+
+  // Whether to use large index.
+  bool large = compiled_use_large_index(inputs, outputs, contiguous);
+
+  cu::KernelArgs args;
+  // Put inputs.
+  int strides_index = 1;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    if (is_constant_(i)) {
+      continue;
+    }
+    const auto& x = inputs[i];
+    args.append(x);
+    if (!contiguous && !is_scalar(x)) {
+      args.append_ptr(strides_vec[strides_index++].data());
+    }
+  }
+
+  auto& encoder = cu::get_command_encoder(s);
+
+  // Put outputs.
+  compiled_allocate_outputs(
+      inputs, outputs, is_constant_, contiguous, [&](auto n) {
+        return cu::malloc_async(n, encoder);
+      });
+  for (auto& x : outputs) {
+    args.append(x);
+  }
+
+  // Put shape and size.
+  if (!contiguous) {
+    args.append_ptr(shape.data());
+  }
+  if (large) {
+    args.append<int64_t>(outputs[0].data_size());
+  } else {
+    args.append<uint32_t>(outputs[0].data_size());
+  }
+
+  // Choose work per thread
+  if (!contiguous && shape.back() % work_per_thread != 0) {
+    work_per_thread = 1;
+  }
+
+  // Launch kernel.
+  const char* index_type = large ? "int64_t" : "uint32_t";
+  std::string kernel_name = fmt::format("mlx::core::cu::{}", lib_name());
+  if (contiguous) {
+    kernel_name +=
+        fmt::format("_contiguous<{}, {}>", index_type, work_per_thread);
+  } else {
+    kernel_name += fmt::format(
+        "_strided<{}, {}, {}>", shape.size(), index_type, work_per_thread);
+  }
+  for (const auto& in : inputs) {
+    encoder.set_input_array(in);
+  }
+  for (const auto& out : outputs) {
+    encoder.set_output_array(out);
+  }
+
+  auto [kernel, max_block_dims] = mod.get_kernel_and_dims(kernel_name);
+  auto [num_blocks, block_dims] =
+      get_launch_args(outputs[0], large, work_per_thread, max_block_dims);
+  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/conv.cpp b/macos/third_party/mlx/mlx/backend/cuda/conv.cpp
new file mode 100644
index 0000000..3b86bf2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/conv.cpp
@@ -0,0 +1,403 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/conv/conv.h"
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/lru_cache.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+#include <cassert>
+
+namespace mlx::core {
+
+namespace {
+
+enum ConvBackendType {
+  CONV_FALLBACK,
+  CONV_FORWARD,
+  CONV_BACKWARD_INPUT,
+  CONV_BACKWARD_WEIGHT,
+};
+
+struct ConvCacheKey {
+  int device_id;
+  fe::DataType_t cudnn_dtype;
+  std::array<int, MAX_NDIM> input_shape;
+  std::array<int, MAX_NDIM> weight_shape;
+  std::array<int, MAX_NDIM> stride;
+  std::array<int, MAX_NDIM> padding_lo;
+  std::array<int, MAX_NDIM> padding_hi;
+  std::array<int, MAX_NDIM> dilation;
+  int groups;
+  bool flip;
+  uint8_t input_alignment;
+  uint8_t weight_alignment;
+  uint8_t output_alignment;
+};
+
+auto& conv_cache() {
+  static LRUBytesKeyCache<
+      ConvCacheKey,
+      std::pair<ConvBackendType, std::optional<DnnGraph>>>
+      cache("MLX_CUDA_CONV_CACHE_SIZE", /* default_capacity */ 128);
+  return cache;
+}
+
+auto get_conv_settings(
+    ConvBackendType backend_type,
+    array& x,
+    array& w,
+    array& y,
+    const std::vector<int>& kernel_strides,
+    const std::vector<int>& padding_lo_,
+    const std::vector<int>& padding_hi_,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation) {
+  auto padding_lo = convert_vector<int64_t>(padding_lo_);
+  auto padding_hi = convert_vector<int64_t>(padding_hi_);
+
+  if (backend_type == CONV_BACKWARD_INPUT) {
+    for (int i = 0; i < padding_lo.size(); ++i) {
+      int wt_size = 1 + kernel_dilation[i] * (w.shape(1 + i) - 1);
+      padding_lo[i] = wt_size - padding_lo[i] - 1;
+      int in_size = 1 + kernel_strides[i] * (y.shape(1 + i) - 1);
+      int out_size = 1 + input_dilation[i] * (x.shape(1 + i) - 1);
+      padding_hi[i] = out_size - in_size + padding_hi[i];
+    }
+    return std::make_tuple(
+        convert_vector<int64_t>(input_dilation),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_dilation));
+
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    padding_hi = padding_lo;
+    return std::make_tuple(
+        convert_vector<int64_t>(kernel_dilation),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_strides));
+
+  } else {
+    return std::make_tuple(
+        convert_vector<int64_t>(kernel_strides),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_dilation));
+  }
+}
+
+std::optional<DnnGraph> build_conv_graph(
+    cu::CommandEncoder& encoder,
+    ConvBackendType backend_type,
+    Dtype dtype,
+    array& x,
+    array& w,
+    array& y,
+    const std::vector<int64_t>& stride,
+    const std::vector<int64_t>& padding_lo,
+    const std::vector<int64_t>& padding_hi,
+    const std::vector<int64_t>& dilation) {
+  auto compute_dtype =
+      (dtype == float16 || dtype == bfloat16) ? float32 : dtype;
+  DnnGraph graph(encoder.device().get_cudnn_handle(), dtype, compute_dtype);
+  auto x_ = graph.tensor_nchw("X", 'x', x);
+  auto w_ = graph.tensor_nchw("W", 'w', w);
+
+  auto set_options = [&](auto& options) {
+    options.set_compute_data_type(dtype_to_cudnn_type(compute_dtype))
+        .set_convolution_mode(fe::ConvolutionMode_t::CROSS_CORRELATION)
+        .set_stride(stride)
+        .set_pre_padding(padding_lo)
+        .set_post_padding(padding_hi)
+        .set_dilation(dilation);
+  };
+
+  std::shared_ptr<fe::graph::Tensor_attributes> y_;
+  if (backend_type == CONV_FORWARD) {
+    auto options = fe::graph::Conv_fprop_attributes();
+    set_options(options);
+    y_ = graph.conv_fprop(x_, w_, options);
+  } else if (backend_type == CONV_BACKWARD_INPUT) {
+    auto options = fe::graph::Conv_dgrad_attributes();
+    set_options(options);
+    y_ = graph.conv_dgrad(x_, w_, options);
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    auto options = fe::graph::Conv_wgrad_attributes();
+    set_options(options);
+    y_ = graph.conv_wgrad(w_, x_, options);
+  }
+  graph.tensor_nchw(y_, 'y', y)->set_output(true);
+
+  if (graph.prepare().is_bad()) {
+    return std::nullopt;
+  }
+  graph.deselect_numeric_notes({fe::NumericalNote_t::DOWN_CONVERT_INPUTS});
+  if (dtype == float32 && !env::enable_tf32()) {
+    graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE});
+  }
+  CHECK_CUDNN_FE_ERROR(graph.build());
+  return graph;
+}
+
+// Transpose from (C_out, H, W, C_in / groups) to (C_in, H, W, C_out / groups).
+array group_transpose(
+    const array& x,
+    int groups,
+    int group_dim,
+    int axis1,
+    int axis2,
+    Stream s) {
+  if (groups == 1) {
+    return swapaxes_in_eval(x, axis1, axis2);
+  }
+  int ndim = x.ndim();
+  if (group_dim < 0) {
+    group_dim += ndim;
+  }
+  if (axis1 < 0) {
+    axis1 += ndim;
+  }
+  if (axis2 < 0) {
+    axis2 += ndim;
+  }
+  if (group_dim <= axis1) {
+    axis1 += 1;
+  }
+  if (group_dim <= axis2) {
+    axis2 += 1;
+  }
+  auto shape = x.shape();
+  shape.insert(shape.begin() + group_dim, groups);
+  shape[group_dim + 1] = shape[group_dim + 1] / groups;
+  array x_trans = reshape_in_eval(x, std::move(shape), s);
+  x_trans = swapaxes_in_eval(x_trans, axis1, axis2);
+  x_trans = flatten_in_eval(x_trans, group_dim, group_dim + 1, s);
+  return x_trans;
+}
+
+// Do necessary transposes and copies to prepare the inputs and outputs for
+// building the cuDNN conv op. It is safe to be called multiple times in one
+// eval_gpu, with cost of possible redundant copies.
+std::tuple<array, array, array> prepare_args(
+    cu::CommandEncoder& encoder,
+    ConvBackendType backend_type,
+    array in,
+    array wt,
+    array out,
+    int groups,
+    Stream s) {
+  // Transpose the args depending on the backend type.
+  // TODO: Handle groups.
+  if (backend_type == CONV_BACKWARD_INPUT) {
+    wt = group_transpose(wt, groups, 0, 0, -1, s);
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    in = group_transpose(in, groups, -1, 0, -1, s);
+    wt = swapaxes_in_eval(wt, 0, -1);
+    // Create a contiguous array that shares the data with |out|, but with dim
+    // C_in and C_out swapped.
+    Shape shape(out.shape());
+    std::swap(shape.front(), shape.back());
+    Strides strides(shape.size(), 1);
+    for (int i = shape.size() - 2; i >= 0; --i) {
+      strides[i] = shape[i + 1] * strides[i + 1];
+    }
+    array intermediate(std::move(shape), out.dtype(), nullptr, {});
+    intermediate.copy_shared_buffer(
+        out, std::move(strides), {true, true, false}, out.data_size());
+    out = intermediate;
+  }
+
+  // cuDNN requires contiguous input.
+  if (!in.flags().row_contiguous) {
+    in = contiguous_copy_gpu(in, s);
+    encoder.add_temporary(in);
+  }
+  if (!wt.flags().row_contiguous) {
+    wt = contiguous_copy_gpu(wt, s);
+    encoder.add_temporary(wt);
+  }
+
+  return {std::move(in), std::move(wt), std::move(out)};
+}
+
+// Register inputs and outputs before actually running conv op. Can only be
+// called once per eval_gpu.
+void register_args(
+    cu::CommandEncoder& encoder,
+    ConvBackendType backend_type,
+    array& in,
+    array& wt,
+    array& intermediate_out,
+    array& final_out) {
+  encoder.set_input_array(in);
+  encoder.set_input_array(wt);
+  encoder.set_output_array(final_out);
+
+  if (backend_type == CONV_BACKWARD_WEIGHT) {
+    // Turn |out| into a strided array, which will have C_in and C_out swapped
+    // in vjp and the final |grad_weight| will then be contiguous.
+    Strides strides = intermediate_out.strides();
+    std::swap(strides.front(), strides.back());
+    final_out.copy_shared_buffer(
+        intermediate_out,
+        std::move(strides),
+        {false, false, false},
+        intermediate_out.data_size());
+  }
+}
+
+} // namespace
+
+void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
+  nvtx3::scoped_range r("Convolution::eval_gpu");
+  if (out_.size() == 0) {
+    return;
+  }
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 2);
+  array in = inputs[0];
+  array wt = inputs[1];
+  array out = out_;
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  Dtype dtype = out.dtype();
+
+  // Search cache.
+  BytesKey<ConvCacheKey> cache_key;
+  cache_key.pod = {
+      encoder.device().cuda_device(),
+      dtype_to_cudnn_type(dtype),
+      vector_key(in.shape()),
+      vector_key(wt.shape()),
+      vector_key(kernel_strides_),
+      vector_key(padding_lo_),
+      vector_key(padding_hi_),
+      vector_key(kernel_dilation_),
+      groups_,
+      flip_,
+      get_alignment(in),
+      get_alignment(wt),
+      get_alignment(out)};
+  if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) {
+    auto& [backend_type, graph] = it->second;
+    if (graph) {
+      // Run cached graph.
+      std::tie(in, wt, out) =
+          prepare_args(encoder, backend_type, in, wt, out, groups_, s);
+      register_args(encoder, backend_type, in, wt, out, out_);
+      CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+          encoder,
+          {
+              {'x', gpu_ptr<void>(in)},
+              {'w', gpu_ptr<void>(wt)},
+              {'y', gpu_ptr<void>(out)},
+          }));
+    } else {
+      // Run fallback kernel.
+      gemm_conv(
+          encoder,
+          in,
+          wt,
+          out,
+          kernel_strides_,
+          padding_lo_,
+          kernel_dilation_,
+          input_dilation_,
+          groups_,
+          flip_,
+          s);
+    }
+    return;
+  }
+
+  // There is no reliable way to deduce the proper cuDNN backend for the
+  // convolution, so we make a best guess and then try.
+  SmallVector<ConvBackendType, 2> try_backends;
+  if (flip_) {
+    // When weight is flipped, we assume it is backward input convolution.
+    try_backends.push_back(CONV_BACKWARD_INPUT);
+  } else {
+    // Otherwise it could be backward weight convolution or forward convolution,
+    // mathematically there is no difference so we have to use heuristics.
+    // Empirically backward convolutions have large kernel dimensions, and
+    // usually have |in| and |wt| transposed.
+    if (!in.flags().row_contiguous && !wt.flags().row_contiguous &&
+        wt.shape(2) > out.shape(2)) {
+      try_backends = {CONV_BACKWARD_WEIGHT, CONV_FORWARD};
+    } else {
+      try_backends = {CONV_FORWARD, CONV_BACKWARD_WEIGHT};
+    }
+  }
+
+  // Try to build op graph.
+  ConvBackendType backend_type;
+  std::optional<DnnGraph> graph;
+  for (auto try_backend : try_backends) {
+    auto [x, w, y] =
+        prepare_args(encoder, try_backend, in, wt, out, groups_, s);
+    auto [stride, padding_lo, padding_hi, dilation] = get_conv_settings(
+        try_backend,
+        x,
+        w,
+        y,
+        kernel_strides_,
+        padding_lo_,
+        padding_hi_,
+        kernel_dilation_,
+        input_dilation_);
+    graph = build_conv_graph(
+        encoder,
+        try_backend,
+        dtype,
+        x,
+        w,
+        y,
+        stride,
+        padding_lo,
+        padding_hi,
+        dilation);
+    if (graph) {
+      backend_type = try_backend;
+      in = std::move(x);
+      wt = std::move(w);
+      out = std::move(y);
+      break;
+    }
+  }
+
+  if (graph) {
+    register_args(encoder, backend_type, in, wt, out, out_);
+    CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+        encoder,
+        {
+            {'x', gpu_ptr<void>(in)},
+            {'w', gpu_ptr<void>(wt)},
+            {'y', gpu_ptr<void>(out)},
+        }));
+    conv_cache().emplace(
+        cache_key, std::make_pair(backend_type, std::move(*graph)));
+    return;
+  }
+
+  // Use fallback kernel for settings not supported by cuDNN.
+  gemm_conv(
+      encoder,
+      in,
+      wt,
+      out,
+      kernel_strides_,
+      padding_lo_,
+      kernel_dilation_,
+      input_dilation_,
+      groups_,
+      flip_,
+      s);
+  conv_cache().emplace(cache_key, std::make_pair(CONV_FALLBACK, std::nullopt));
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/conv/conv.h b/macos/third_party/mlx/mlx/backend/cuda/conv/conv.h
new file mode 100644
index 0000000..62dc934
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/conv/conv.h
@@ -0,0 +1,126 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/gpu/copy.h"
+
+namespace mlx::core {
+
+template <int NDIM>
+struct ConvParams {
+  int N; // Batch size
+  int C; // In channels
+  int O; // Out channels
+  int strides[NDIM];
+  int padding[NDIM];
+  int kernel_dilation[NDIM];
+  int input_dilation[NDIM];
+  int groups;
+  bool flip;
+  int in_spatial_dims[NDIM];
+  int wt_spatial_dims[NDIM];
+  int out_spatial_dims[NDIM];
+  int64_t in_strides[NDIM + 2];
+
+  ConvParams(
+      const array& in,
+      const array& wt,
+      const array& out,
+      const std::vector<int>& strides,
+      const std::vector<int>& padding,
+      const std::vector<int>& kernel_dilation,
+      const std::vector<int>& input_dilation,
+      int groups,
+      bool flip)
+      : N(in.shape(0)),
+        C(in.shape(-1)),
+        O(wt.shape(0)),
+        groups(groups),
+        flip(flip) {
+    std::copy_n(strides.begin(), NDIM, this->strides);
+    std::copy_n(padding.begin(), NDIM, this->padding);
+    std::copy_n(kernel_dilation.begin(), NDIM, this->kernel_dilation);
+    std::copy_n(input_dilation.begin(), NDIM, this->input_dilation);
+    std::copy_n(in.shape().begin() + 1, NDIM, this->in_spatial_dims);
+    std::copy_n(wt.shape().begin() + 1, NDIM, this->wt_spatial_dims);
+    std::copy_n(out.shape().begin() + 1, NDIM, this->out_spatial_dims);
+    std::copy_n(in.strides().begin(), NDIM + 2, this->in_strides);
+  }
+};
+
+void gemm_grouped_conv(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    const std::vector<int>& strides,
+    const std::vector<int>& padding,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation,
+    int groups,
+    bool flip,
+    Stream s);
+
+void gemm_conv(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    const std::vector<int>& strides,
+    const std::vector<int>& padding,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation,
+    bool flip,
+    Stream s);
+
+inline void gemm_conv(
+    cu::CommandEncoder& encoder,
+    array in,
+    array wt,
+    array& out,
+    const std::vector<int>& strides,
+    const std::vector<int>& padding,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation,
+    int groups,
+    bool flip,
+    Stream s) {
+  if (!in.flags().row_contiguous) {
+    in = contiguous_copy_gpu(in, s);
+    encoder.add_temporary(in);
+  }
+  if (!wt.flags().row_contiguous) {
+    wt = contiguous_copy_gpu(wt, s);
+    encoder.add_temporary(wt);
+  }
+
+  if (groups == 1) {
+    gemm_conv(
+        encoder,
+        in,
+        wt,
+        out,
+        strides,
+        padding,
+        kernel_dilation,
+        input_dilation,
+        flip,
+        s);
+  } else {
+    gemm_grouped_conv(
+        encoder,
+        in,
+        wt,
+        out,
+        strides,
+        padding,
+        kernel_dilation,
+        input_dilation,
+        groups,
+        flip,
+        s);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/conv/gemm_conv.cu b/macos/third_party/mlx/mlx/backend/cuda/conv/gemm_conv.cu
new file mode 100644
index 0000000..a869542
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/conv/gemm_conv.cu
@@ -0,0 +1,217 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/conv/conv.h"
+#include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, int NDIM>
+__global__ void naive_unfold_nd(
+    const T* in,
+    T* out,
+    int filter_size,
+    int out_pixels,
+    const __grid_constant__ ConvParams<NDIM> params) {
+  auto block = cg::this_thread_block();
+  auto tid = block.group_index();
+  auto lid = block.thread_index();
+
+  int index_batch = tid.z / out_pixels; // [0, N)
+  int index_out_spatial = tid.z % out_pixels; // [0, H_out * W_out)
+  int index_wt_spatial =
+      tid.x * block.dim_threads().x + lid.x; // [0, H_wt * W_wt)
+
+  if (index_wt_spatial >= filter_size / params.C) {
+    return;
+  }
+
+  in += tid.y; // [0, C)
+  out += tid.z * filter_size + index_wt_spatial * params.C + tid.y;
+
+  bool valid = index_batch < params.N;
+
+  // Get the coordinates in input.
+  int index_in[NDIM] = {};
+#pragma unroll
+  for (int i = NDIM - 1; i >= 0; --i) {
+    int index_out = index_out_spatial % params.out_spatial_dims[i];
+    int index_wt = index_wt_spatial % params.wt_spatial_dims[i];
+
+    if (params.flip) {
+      index_wt = params.wt_spatial_dims[i] - index_wt - 1;
+    }
+
+    int index = index_out * params.strides[i] - params.padding[i] +
+        index_wt * params.kernel_dilation[i];
+    int index_max =
+        1 + params.input_dilation[i] * (params.in_spatial_dims[i] - 1);
+
+    valid &= (index >= 0) && (index < index_max) &&
+        (index % params.input_dilation[i] == 0);
+
+    index_in[i] = index / params.input_dilation[i];
+
+    index_out_spatial /= params.out_spatial_dims[i];
+    index_wt_spatial /= params.wt_spatial_dims[i];
+  }
+
+  if (valid) {
+    int in_offset = index_batch * params.in_strides[0];
+#pragma unroll
+    for (int i = 0; i < NDIM; ++i) {
+      in_offset += index_in[i] * params.in_strides[i + 1];
+    }
+    *out = in[in_offset];
+  } else {
+    *out = T{0};
+  }
+}
+
+} // namespace cu
+
+template <int NDIM>
+array unfold_inputs_nd(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    int mat_M,
+    int mat_K,
+    int mat_N,
+    ConvParams<NDIM>& params) {
+  array unfolded({mat_M, mat_K}, in.dtype(), nullptr, {});
+  unfolded.set_data(cu::malloc_async(unfolded.nbytes(), encoder));
+  encoder.add_temporary(unfolded);
+
+  int filter_size = params.C;
+#pragma unroll
+  for (int i = 0; i < NDIM; ++i) {
+    filter_size *= params.wt_spatial_dims[i];
+  }
+
+  int out_pixels = 1;
+#pragma unroll
+  for (int i = 0; i < NDIM; ++i) {
+    out_pixels *= params.out_spatial_dims[i];
+  }
+
+  int wt_spatial_size = mat_K / params.C;
+  dim3 block_dims;
+  block_dims.x = std::min(std::max(wt_spatial_size, 32), 1024);
+  dim3 num_blocks;
+  num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x);
+  num_blocks.y = params.C;
+  num_blocks.z = mat_M;
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(unfolded);
+  dispatch_float_types(in.dtype(), "unfold", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    encoder.add_kernel_node(
+        cu::naive_unfold_nd<DataType, NDIM>,
+        num_blocks,
+        block_dims,
+        0,
+        gpu_ptr<DataType>(in),
+        gpu_ptr<DataType>(unfolded),
+        filter_size,
+        out_pixels,
+        params);
+  });
+
+  return unfolded;
+}
+
+template <int NDIM>
+void gemm_conv_nd(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    ConvParams<NDIM>& params,
+    Stream s) {
+  // Get gemm shapes.
+  int mat_M = out.size() / params.O; // N * H_out * W_out
+  int mat_K = wt.size() / params.O; // C * H_wt * W_wt
+  int mat_N = params.O; // O
+
+  // Unfold input to (N * H_out * W_out, C * H_wt * W_wt) for gemm.
+  array in_unfolded =
+      unfold_inputs_nd<NDIM>(encoder, in, mat_M, mat_K, mat_N, params);
+
+  // Reshape weight to (C * H_wt * W_wt, O) for gemm.
+  array wt_reshaped({mat_K, mat_N}, wt.dtype(), nullptr, {});
+  wt_reshaped.copy_shared_buffer(
+      wt,
+      {1, mat_K},
+      {false, false, /* col_contiguous */ true},
+      wt.data_size());
+
+  // Single batch.
+  Shape batch_shape{1};
+  Strides a_batch_strides{0};
+  Strides b_batch_strides{0};
+
+  // Run matmul.
+  CublasGemm gemm(
+      encoder.device(),
+      in.dtype(),
+      false, // a_transposed
+      mat_M, // a_rows
+      mat_K, // a_cols
+      mat_K, // lda
+      true, // b_transposed
+      mat_K, // b_rows
+      mat_N, // b_cols
+      mat_K, // ldb
+      batch_shape.back(),
+      a_batch_strides.back(),
+      b_batch_strides.back());
+  gemm.run(
+      encoder,
+      out,
+      in_unfolded,
+      wt_reshaped,
+      batch_shape,
+      a_batch_strides,
+      b_batch_strides);
+}
+
+void gemm_conv(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    const std::vector<int>& strides,
+    const std::vector<int>& padding,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation,
+    bool flip,
+    Stream s) {
+  int conv_ndim = in.ndim() - 2;
+  if (conv_ndim < 1 || conv_ndim > 3) {
+    throw std::runtime_error(
+        fmt::format("[conv] Unsupported gemm_conv for {}D conv.", conv_ndim));
+  }
+  dispatch_1_2_3(conv_ndim, [&](auto ndim_constant) {
+    ConvParams<ndim_constant()> params(
+        in,
+        wt,
+        out,
+        strides,
+        padding,
+        kernel_dilation,
+        input_dilation,
+        1, // groups
+        flip);
+    gemm_conv_nd<ndim_constant()>(encoder, in, wt, out, params, s);
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu b/macos/third_party/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu
new file mode 100644
index 0000000..44e067c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu
@@ -0,0 +1,231 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/conv/conv.h"
+#include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, int NDIM>
+__global__ void naive_grouped_unfold_transpose_nd(
+    const T* in,
+    T* out,
+    int filter_size,
+    int out_pixels,
+    const __grid_constant__ ConvParams<NDIM> params) {
+  auto block = cg::this_thread_block();
+  auto tid = block.group_index();
+  auto lid = block.thread_index();
+
+  int index_batch = tid.z / out_pixels; // [0, N)
+  int index_out_spatial = tid.z % out_pixels; // [0, H_out * W_out)
+  int index_wt_spatial =
+      tid.x * block.dim_threads().x + lid.x; // [0, H_wt * W_wt)
+
+  if (index_wt_spatial >= filter_size / params.C) {
+    return;
+  }
+
+  in += tid.y; // [0, C)
+  out += tid.z * filter_size + tid.y * (filter_size / params.C);
+
+  bool valid = index_batch < params.N;
+
+  // Get the coordinates in input.
+  int index_in[NDIM] = {};
+  int wt_stride = 1;
+#pragma unroll
+  for (int i = NDIM - 1; i >= 0; --i) {
+    int index_out = index_out_spatial % params.out_spatial_dims[i];
+    int index_wt = index_wt_spatial % params.wt_spatial_dims[i];
+    out += index_wt * wt_stride;
+
+    if (params.flip) {
+      index_wt = params.wt_spatial_dims[i] - index_wt - 1;
+    }
+
+    int index = index_out * params.strides[i] - params.padding[i] +
+        index_wt * params.kernel_dilation[i];
+    int index_max =
+        1 + params.input_dilation[i] * (params.in_spatial_dims[i] - 1);
+
+    valid &= (index >= 0) && (index < index_max) &&
+        (index % params.input_dilation[i] == 0);
+
+    index_in[i] = index / params.input_dilation[i];
+
+    index_out_spatial /= params.out_spatial_dims[i];
+    index_wt_spatial /= params.wt_spatial_dims[i];
+    wt_stride *= params.wt_spatial_dims[i];
+  }
+
+  if (valid) {
+    int in_offset = index_batch * params.in_strides[0];
+#pragma unroll
+    for (int i = 0; i < NDIM; ++i) {
+      in_offset += index_in[i] * params.in_strides[i + 1];
+    }
+    *out = in[in_offset];
+  } else {
+    *out = T{0};
+  }
+}
+
+} // namespace cu
+
+template <int NDIM>
+array grouped_unfold_transpose_inputs_nd(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    int mat_M,
+    int mat_K,
+    int mat_N,
+    ConvParams<NDIM>& params) {
+  array unfolded({mat_M, mat_K * params.groups}, in.dtype(), nullptr, {});
+  unfolded.set_data(cu::malloc_async(unfolded.nbytes(), encoder));
+  encoder.add_temporary(unfolded);
+
+  int filter_size = params.C;
+#pragma unroll
+  for (int i = 0; i < NDIM; ++i) {
+    filter_size *= params.wt_spatial_dims[i];
+  }
+
+  int out_pixels = 1;
+#pragma unroll
+  for (int i = 0; i < NDIM; ++i) {
+    out_pixels *= params.out_spatial_dims[i];
+  }
+
+  int wt_spatial_size = (mat_K * params.groups) / params.C;
+  dim3 block_dims;
+  block_dims.x = std::min(std::max(wt_spatial_size, 32), 1024);
+  dim3 num_blocks;
+  num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x);
+  num_blocks.y = params.C;
+  num_blocks.z = mat_M;
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(unfolded);
+  dispatch_float_types(in.dtype(), "unfold", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    encoder.add_kernel_node(
+        cu::naive_grouped_unfold_transpose_nd<DataType, NDIM>,
+        num_blocks,
+        block_dims,
+        0,
+        gpu_ptr<DataType>(in),
+        gpu_ptr<DataType>(unfolded),
+        filter_size,
+        out_pixels,
+        params);
+  });
+
+  return unfolded;
+}
+
+template <int NDIM>
+void gemm_grouped_conv_nd(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    ConvParams<NDIM>& params,
+    Stream s) {
+  // Get gemm shapes.
+  int C_per_group = params.C / params.groups;
+  int O_per_group = params.O / params.groups;
+  int mat_M = out.size() / params.O; // N * H_out * W_out
+  int mat_K = wt.size() / params.O; // C_per_group * H_wt * W_wt
+  int mat_N = O_per_group; // O_per_group
+
+  // Unfold input to (N * H_out * W_out, C * H_wt * W_wt) for gemm.
+  array in_unfolded = grouped_unfold_transpose_inputs_nd<NDIM>(
+      encoder, in, mat_M, mat_K, mat_N, params);
+
+  // Reshape weight to (O, C_per_group, H_wt * W_wt) for gemm.
+  int wt_spatial_size = (wt.size() / wt.shape(0)) / wt.shape(-1);
+  array wt_view(
+      {params.O, C_per_group, wt_spatial_size}, wt.dtype(), nullptr, {});
+  wt_view.copy_shared_buffer(
+      wt, {wt.strides(0), 1, C_per_group}, wt.flags(), wt.size());
+  array wt_reshaped = contiguous_copy_gpu(wt_view, s);
+
+  // Batch with size of groups.
+  Shape batch_shape{params.groups};
+  Strides a_batch_strides{mat_K};
+  Strides b_batch_strides{mat_N * mat_K};
+
+  // Run matmul.
+  CublasGemm gemm(
+      encoder.device(),
+      in.dtype(),
+      false, // a_transposed
+      mat_M, // a_rows
+      mat_K, // a_cols
+      mat_K * params.groups, // lda
+      true, // b_transposed
+      mat_K, // b_rows
+      mat_N, // b_cols
+      mat_K, // ldb
+      batch_shape.back(),
+      a_batch_strides.back(),
+      b_batch_strides.back());
+  gemm.set_out(
+      out.dtype(),
+      false, // out_transposed
+      mat_M, // out_rows
+      mat_N, // out_cols
+      mat_N * params.groups, // out_ld
+      params.groups, // batch_count
+      mat_N); // batch_stride
+  gemm.run(
+      encoder,
+      out,
+      in_unfolded,
+      wt_reshaped,
+      batch_shape,
+      a_batch_strides,
+      b_batch_strides);
+}
+
+void gemm_grouped_conv(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    const std::vector<int>& strides,
+    const std::vector<int>& padding,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation,
+    int groups,
+    bool flip,
+    Stream s) {
+  int conv_ndim = in.ndim() - 2;
+  if (conv_ndim < 1 || conv_ndim > 3) {
+    throw std::runtime_error(
+        fmt::format("[conv] Unsupported gemm_conv for {}D conv.", conv_ndim));
+  }
+  dispatch_1_2_3(conv_ndim, [&](auto ndim_constant) {
+    ConvParams<ndim_constant()> params(
+        in,
+        wt,
+        out,
+        strides,
+        padding,
+        kernel_dilation,
+        input_dilation,
+        groups,
+        flip);
+    gemm_grouped_conv_nd<ndim_constant()>(encoder, in, wt, out, params, s);
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/copy.cu b/macos/third_party/mlx/mlx/backend/cuda/copy.cu
new file mode 100644
index 0000000..d3d67e6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/copy.cu
@@ -0,0 +1,132 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/copy/copy.cuh"
+
+namespace mlx::core {
+
+void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
+  auto& encoder = cu::get_command_encoder(s);
+  bool donated = set_copy_output_data(
+      in, out, ctype, [&](auto n) { return cu::malloc_async(n, encoder); });
+  if (donated && in.dtype() == out.dtype()) {
+    // If the output has the same type as the input then there is nothing to
+    // copy, just use the buffer.
+    return;
+  }
+  if (ctype == CopyType::GeneralGeneral) {
+    ctype = CopyType::General;
+  }
+  copy_gpu_inplace(in, out, ctype, s);
+}
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out,
+    int64_t offset_in,
+    int64_t offset_out,
+    CopyType ctype,
+    const Stream& s,
+    std::optional<array> dynamic_offset_in,
+    std::optional<array> dynamic_offset_out) {
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cu::get_command_encoder(s);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  if (ctype == CopyType::Scalar || ctype == CopyType::Vector) {
+    copy_contiguous(encoder, ctype, in, out, offset_in, offset_out);
+    return;
+  }
+
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    auto [shape_collapsed, strides_vec] = collapse_contiguous_dims(
+        shape, std::vector{strides_in, strides_out}, INT32_MAX);
+    if (ctype == CopyType::General) {
+      copy_general_input(
+          encoder,
+          ctype,
+          in,
+          out,
+          offset_in,
+          offset_out,
+          shape_collapsed,
+          strides_vec[0]);
+    } else {
+      if (dynamic_offset_in || dynamic_offset_out) {
+        if (!dynamic_offset_in) {
+          dynamic_offset_in = array(0, int64);
+          encoder.add_temporary(*dynamic_offset_in);
+        }
+        if (!dynamic_offset_out) {
+          dynamic_offset_out = array(0, int64);
+          encoder.add_temporary(*dynamic_offset_out);
+        }
+        encoder.set_input_array(*dynamic_offset_in);
+        encoder.set_input_array(*dynamic_offset_out);
+        copy_general_dynamic(
+            encoder,
+            ctype,
+            in,
+            out,
+            offset_in,
+            offset_out,
+            shape_collapsed,
+            strides_vec[0],
+            strides_vec[1],
+            *dynamic_offset_in,
+            *dynamic_offset_out);
+      } else {
+        copy_general(
+            encoder,
+            ctype,
+            in,
+            out,
+            offset_in,
+            offset_out,
+            shape_collapsed,
+            strides_vec[0],
+            strides_vec[1]);
+      }
+    }
+    return;
+  }
+}
+
+void fill_gpu(const array& in, array& out, const Stream& s) {
+  if (out.size() == 0) {
+    return;
+  }
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  copy_contiguous(encoder, CopyType::Scalar, in, out, 0, 0);
+}
+
+void reshape_gpu(const array& in, array& out, Stream s) {
+  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
+  if (copy_necessary) {
+    auto& encoder = cu::get_command_encoder(s);
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    copy_gpu_inplace(
+        in,
+        out,
+        in.shape(),
+        in.strides(),
+        make_contiguous_strides(in.shape()),
+        0,
+        0,
+        CopyType::General,
+        s);
+  } else {
+    shared_buffer_reshape(in, out_strides, out);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/copy/copy.cuh b/macos/third_party/mlx/mlx/backend/cuda/copy/copy.cuh
new file mode 100644
index 0000000..e80fdec
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/copy/copy.cuh
@@ -0,0 +1,55 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/cast_op.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core {
+
+void copy_contiguous(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out);
+
+void copy_general(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out);
+
+void copy_general_dynamic(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out,
+    const array& dynamic_offset_in,
+    const array& dynamic_offset_out);
+
+void copy_general_input(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/copy/copy_contiguous.cu b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_contiguous.cu
new file mode 100644
index 0000000..6f51734
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_contiguous.cu
@@ -0,0 +1,88 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/copy/copy.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename In, typename Out, typename IdxT, int N_READS>
+__global__ void copy_s(const In* in, Out* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      out[i] = cast_to<Out>(in[0]);
+    }
+  } else {
+    AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = cast_to<Out>(in[0]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+template <typename In, typename Out, typename IdxT, int N_READS>
+__global__ void copy_v(const In* in, Out* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      out[i] = cast_to<Out>(in[i]);
+    }
+  } else {
+    auto in_vec = load_vector<N_READS>(in, index);
+
+    AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = cast_to<Out>(in_vec[i]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+} // namespace cu
+
+void copy_contiguous(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t in_offset,
+    int64_t out_offset) {
+  dispatch_all_types(in.dtype(), [&](auto in_type_tag) {
+    dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
+      dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
+        using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
+        using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
+        using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
+        constexpr int N_READS = 16 / sizeof(InType);
+        auto kernel = cu::copy_s<InType, OutType, IdxT, N_READS>;
+        if (ctype == CopyType::Vector) {
+          kernel = cu::copy_v<InType, OutType, IdxT, N_READS>;
+        }
+        auto [num_blocks, block_dims] = get_launch_args(
+            out.data_size(), out.shape(), out.strides(), large(), N_READS);
+        encoder.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<InType>(in) + in_offset,
+            gpu_ptr<OutType>(out) + out_offset,
+            out.data_size());
+      });
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general.cu b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general.cu
new file mode 100644
index 0000000..fec110a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general.cu
@@ -0,0 +1,171 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/copy/copy.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename In, typename Out, typename IdxT, int NDIM, int N_READS>
+__global__ void copy_gg_nd(
+    const In* in,
+    Out* out,
+    IdxT size_rest,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> strides_in,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> strides_out) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[NDIM - 1];
+  auto in_stride_x = strides_in[NDIM - 1];
+  auto out_stride_x = strides_out[NDIM - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [idx_in, idx_out] = elem_to_loc_nd<NDIM>(
+      index_rest * shape_x,
+      shape.data(),
+      strides_in.data(),
+      strides_out.data());
+
+  auto in_vec =
+      load_vector<N_READS>(in + idx_in, index_x, shape_x, in_stride_x, In(0));
+  AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = CastOp<In, Out>{}(in_vec[i]);
+  }
+  store_vector(out + idx_out, index_x, out_vec, shape_x, out_stride_x);
+}
+
+template <typename In, typename Out, typename IdxT, int N_READS>
+__global__ void copy_gg(
+    const In* in,
+    Out* out,
+    IdxT size_rest,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides strides_in,
+    const __grid_constant__ Strides strides_out,
+    int ndim) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[ndim - 1];
+  auto in_stride_x = strides_in[ndim - 1];
+  auto out_stride_x = strides_out[ndim - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [idx_in, idx_out] = elem_to_loc(
+      index_rest * shape_x,
+      shape.data(),
+      strides_in.data(),
+      strides_out.data(),
+      ndim);
+
+  auto in_vec =
+      load_vector<N_READS>(in + idx_in, index_x, shape_x, in_stride_x, In(0));
+  AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = CastOp<In, Out>{}(in_vec[i]);
+  }
+  store_vector(out + idx_out, index_x, out_vec, shape_x, out_stride_x);
+}
+
+} // namespace cu
+
+void copy_general(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out) {
+  dispatch_all_types(in.dtype(), [&](auto in_type_tag) {
+    dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
+      dispatch_bool(
+          in.data_size() > INT32_MAX || out.data_size() > INT32_MAX,
+          [&](auto large) {
+            using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
+            using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
+            using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
+            OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
+            int ndim = shape.size();
+            size_t data_size = 1;
+            for (auto& s : shape)
+              data_size *= s;
+
+            int work_per_thread = 1;
+            auto dim0 = ndim > 0 ? shape.back() : 1;
+            auto rest = data_size / dim0;
+            if (dim0 >= 4) {
+              work_per_thread = 4;
+            }
+
+            dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+            auto block_dims = get_block_dims(dim0, rest, 1);
+            uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x);
+            uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y);
+
+            if (ndim <= 3) {
+              dispatch_1_2_3(ndim, [&](auto ndim_constant) {
+                auto kernel =
+                    cu::copy_gg_nd<InType, OutType, IdxT, ndim_constant(), 1>;
+                if (work_per_thread == 4) {
+                  kernel =
+                      cu::copy_gg_nd<InType, OutType, IdxT, ndim_constant(), 4>;
+                }
+                encoder.add_kernel_node(
+                    kernel,
+                    {num_blocks_x, num_blocks_y},
+                    block_dims,
+                    0,
+                    in_ptr,
+                    out_ptr,
+                    rest,
+                    const_param<ndim_constant()>(shape),
+                    const_param<ndim_constant()>(strides_in),
+                    const_param<ndim_constant()>(strides_out));
+              });
+            } else { // ndim >= 4
+              auto kernel = cu::copy_gg<InType, OutType, IdxT, 1>;
+              if (work_per_thread == 4) {
+                kernel = cu::copy_gg<InType, OutType, IdxT, 4>;
+              }
+              encoder.add_kernel_node(
+                  kernel,
+                  {num_blocks_x, num_blocks_y},
+                  block_dims,
+                  0,
+                  in_ptr,
+                  out_ptr,
+                  rest,
+                  const_param(shape),
+                  const_param(strides_in),
+                  const_param(strides_out),
+                  ndim);
+            }
+          });
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general_dynamic.cu b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general_dynamic.cu
new file mode 100644
index 0000000..96cf6fb
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general_dynamic.cu
@@ -0,0 +1,118 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/copy/copy.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename In, typename Out, typename IdxT, int NDIM>
+__global__ void copy_gg_dynamic_nd(
+    const In* in,
+    Out* out,
+    IdxT size,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> strides_in,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> strides_out,
+    const int64_t* offset_in,
+    const int64_t* offset_out) {
+  IdxT index = cg::this_grid().thread_rank();
+  if (index < size) {
+    auto [idx_in, idx_out] = elem_to_loc_nd<NDIM>(
+        index, shape.data(), strides_in.data(), strides_out.data());
+    out[idx_out + *offset_out] = CastOp<In, Out>{}(in[idx_in + *offset_in]);
+  }
+}
+
+template <typename In, typename Out, typename IdxT>
+__global__ void copy_gg_dynamic(
+    const In* in,
+    Out* out,
+    IdxT size,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides strides_in,
+    const __grid_constant__ Strides strides_out,
+    int ndim,
+    const int64_t* offset_in,
+    const int64_t* offset_out) {
+  IdxT index = cg::this_grid().thread_rank();
+  if (index < size) {
+    auto [idx_in, idx_out] = elem_to_loc(
+        index, shape.data(), strides_in.data(), strides_out.data(), ndim);
+    out[idx_out + *offset_out] = CastOp<In, Out>{}(in[idx_in + *offset_in]);
+  }
+}
+
+} // namespace cu
+
+void copy_general_dynamic(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out,
+    const array& dynamic_offset_in,
+    const array& dynamic_offset_out) {
+  dispatch_all_types(in.dtype(), [&](auto in_type_tag) {
+    dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
+      dispatch_bool(
+          in.data_size() > INT32_MAX || out.data_size() > INT32_MAX,
+          [&](auto large) {
+            using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
+            using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
+            using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
+            OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
+            int ndim = shape.size();
+            if (ndim <= 3) {
+              dispatch_1_2_3(ndim, [&](auto dims_constant) {
+                auto [num_blocks, block_dims] = get_launch_args(out, large());
+                encoder.add_kernel_node(
+                    cu::copy_gg_dynamic_nd<
+                        InType,
+                        OutType,
+                        IdxT,
+                        dims_constant()>,
+                    num_blocks,
+                    block_dims,
+                    0,
+                    in_ptr,
+                    out_ptr,
+                    out.size(),
+                    const_param<dims_constant()>(shape),
+                    const_param<dims_constant()>(strides_in),
+                    const_param<dims_constant()>(strides_out),
+                    gpu_ptr<int64_t>(dynamic_offset_in),
+                    gpu_ptr<int64_t>(dynamic_offset_out));
+              });
+            } else { // ndim >= 4
+              auto [num_blocks, block_dims] = get_launch_args(out, large());
+              encoder.add_kernel_node(
+                  cu::copy_gg_dynamic<InType, OutType, IdxT>,
+                  num_blocks,
+                  block_dims,
+                  0,
+                  in_ptr,
+                  out_ptr,
+                  out.size(),
+                  const_param(shape),
+                  const_param(strides_in),
+                  const_param(strides_out),
+                  ndim,
+                  gpu_ptr<int64_t>(dynamic_offset_in),
+                  gpu_ptr<int64_t>(dynamic_offset_out));
+            }
+          });
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general_input.cu b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general_input.cu
new file mode 100644
index 0000000..b1b7939
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/copy/copy_general_input.cu
@@ -0,0 +1,229 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/copy/copy.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+static constexpr int TILE_SIZE = 16;
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename In, typename Out, typename IdxT, int NDIM, int N_READS>
+__global__ void copy_g_nd(
+    const In* in,
+    Out* out,
+    IdxT size_rest,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> strides) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[NDIM - 1];
+  auto stride_x = strides[NDIM - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto idx =
+      elem_to_loc_nd<NDIM>(index_rest * shape_x, shape.data(), strides.data());
+  auto in_vec =
+      load_vector<N_READS>(in + idx, index_x, shape_x, stride_x, In(0));
+  AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = CastOp<In, Out>{}(in_vec[i]);
+  }
+  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
+}
+
+template <typename In, typename Out, typename IdxT, int N_READS>
+__global__ void copy_g(
+    const In* in,
+    Out* out,
+    IdxT size_rest,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides strides,
+    int ndim) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[ndim - 1];
+  auto stride_x = strides[ndim - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto idx =
+      elem_to_loc(index_rest * shape_x, shape.data(), strides.data(), ndim);
+  auto in_vec =
+      load_vector<N_READS>(in + idx, index_x, shape_x, stride_x, In(0));
+  AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = CastOp<In, Out>{}(in_vec[i]);
+  }
+  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
+}
+
+template <typename In, typename Out, int N_READS>
+__global__ void
+copy_col_row(const In* in, Out* out, int64_t rows, int64_t cols) {
+  __shared__ Out
+      tile[N_READS * TILE_SIZE][N_READS * TILE_SIZE + 4 / sizeof(Out)];
+
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+
+  auto tile_row = grid.block_index().x * TILE_SIZE * N_READS;
+  auto tile_col = grid.block_index().y * TILE_SIZE * N_READS;
+
+  auto tidx = block.thread_index().x;
+  auto tidy = N_READS * block.thread_index().y;
+
+  auto in_ptr = in + (tile_col + tidy) * rows + tile_row;
+
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    if ((tile_col + tidy + i) < cols) {
+      auto in_vec = load_vector<N_READS>(in_ptr, tidx, rows - tile_row, In(0));
+#pragma unroll
+      for (int j = 0; j < N_READS; ++j) {
+        tile[N_READS * tidx + j][tidy + i] = CastOp<In, Out>{}(in_vec[j]);
+      }
+      in_ptr += rows;
+    }
+  }
+
+  block.sync();
+
+  auto out_ptr = out + (tile_row + tidy) * cols + tile_col;
+
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    if ((tile_row + tidy + i) < rows) {
+      AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+      for (int j = 0; j < N_READS; ++j) {
+        out_vec[j] = tile[tidy + i][N_READS * tidx + j];
+      }
+      store_vector(out_ptr, tidx, out_vec, cols - tile_col);
+      out_ptr += cols;
+    }
+  }
+}
+
+} // namespace cu
+
+void copy_general_input(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in) {
+  dispatch_all_types(in.dtype(), [&](auto in_type_tag) {
+    dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
+      using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
+      using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
+      const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
+      OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
+      int ndim = shape.size();
+
+      // Column contiguous to row contiguous specialization
+      if (ndim == 2 && strides_in[0] == 1 && strides_in[1] == shape[0]) {
+        constexpr int work_per_thread =
+            std::min(static_cast<int>(16 / sizeof(OutType)), 8);
+        dim3 block_dims = {TILE_SIZE, TILE_SIZE};
+        uint32_t num_blocks_x =
+            cuda::ceil_div(shape[0], TILE_SIZE * work_per_thread);
+        uint32_t num_blocks_y =
+            cuda::ceil_div(shape[1], TILE_SIZE * work_per_thread);
+        auto kernel = cu::copy_col_row<InType, OutType, work_per_thread>;
+        encoder.add_kernel_node(
+            kernel,
+            {num_blocks_x, num_blocks_y},
+            block_dims,
+            0,
+            in_ptr,
+            out_ptr,
+            int64_t(shape[0]),
+            int64_t(shape[1]));
+        return;
+      }
+
+      dispatch_bool(
+          in.data_size() > INT32_MAX || out.data_size() > INT32_MAX,
+          [&](auto large) {
+            using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+
+            int work_per_thread = 8;
+            auto dim0 = ndim > 0 ? shape.back() : 1;
+            auto rest = out.size() / dim0;
+            if (dim0 >= 4 && dim0 < 8) {
+              work_per_thread = 4;
+            } else if (dim0 < 4) {
+              work_per_thread = 1;
+            }
+            dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+            auto block_dims = get_block_dims(dim0, rest, 1);
+            uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x);
+            uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y);
+
+            if (ndim <= 3) {
+              dispatch_1_2_3(ndim, [&](auto dims_constant) {
+                auto kernel =
+                    cu::copy_g_nd<InType, OutType, IdxT, dims_constant(), 1>;
+                if (work_per_thread == 8) {
+                  kernel =
+                      cu::copy_g_nd<InType, OutType, IdxT, dims_constant(), 8>;
+                } else if (work_per_thread == 4) {
+                  kernel =
+                      cu::copy_g_nd<InType, OutType, IdxT, dims_constant(), 4>;
+                }
+                encoder.add_kernel_node(
+                    kernel,
+                    {num_blocks_x, num_blocks_y},
+                    block_dims,
+                    0,
+                    in_ptr,
+                    out_ptr,
+                    rest,
+                    const_param<dims_constant()>(shape),
+                    const_param<dims_constant()>(strides_in));
+              });
+            } else { // ndim >= 4
+              auto kernel = cu::copy_g<InType, OutType, IdxT, 1>;
+              if (work_per_thread == 8) {
+                kernel = cu::copy_g<InType, OutType, IdxT, 8>;
+              } else if (work_per_thread == 4) {
+                kernel = cu::copy_g<InType, OutType, IdxT, 4>;
+              }
+              encoder.add_kernel_node(
+                  kernel,
+                  {num_blocks_x, num_blocks_y},
+                  block_dims,
+                  0,
+                  in_ptr,
+                  out_ptr,
+                  rest,
+                  const_param(shape),
+                  const_param(strides_in),
+                  ndim);
+            }
+          });
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/cublas_utils.cpp b/macos/third_party/mlx/mlx/backend/cuda/cublas_utils.cpp
new file mode 100644
index 0000000..1176bd4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/cublas_utils.cpp
@@ -0,0 +1,222 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+namespace cublas_utils {
+
+namespace {
+
+struct CublasPreference {
+  CublasPreference(cu::Device& device) {
+    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 MiB
+    // for Hopper+:
+    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
+    uint64_t MiB = 1024 * 1024;
+    uint64_t workspace_size =
+        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
+
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&pref_));
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
+        pref_,
+        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+        &workspace_size,
+        sizeof(uint64_t)));
+  }
+
+  ~CublasPreference() {
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref_));
+  }
+
+  cublasLtMatmulPreference_t pref_{nullptr};
+};
+
+} // namespace
+
+cublasLtMatmulPreference_t get_preference(cu::Device& device) {
+  static CublasPreference pref(device);
+  return pref.pref_;
+}
+
+cublasLtMatrixLayout_t create_matrix_layout(
+    cudaDataType_t type,
+    uint64_t rows,
+    uint64_t cols,
+    bool transposed,
+    int64_t ld,
+    int32_t batch_count,
+    int64_t batch_stride) {
+  cublasLtMatrixLayout_t desc;
+  if (transposed) {
+    std::swap(rows, cols);
+  }
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutCreate(&desc, type, rows, cols, ld));
+  if (batch_count > 1) {
+    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+        desc,
+        CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
+        &batch_count,
+        sizeof(int32_t)));
+    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+        desc,
+        CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
+        &batch_stride,
+        sizeof(int64_t)));
+  }
+  return desc;
+}
+
+} // namespace cublas_utils
+
+CublasMatmulBase::~CublasMatmulBase() {
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(a_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(b_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(c_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(out_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescDestroy(matmul_desc_));
+}
+
+void CublasMatmulBase::init_base(
+    cu::Device& device,
+    cudaDataType_t scale_type,
+    cublasComputeType_t compute_type,
+    cudaDataType_t data_type,
+    cudaDataType_t output_type,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride) {
+  M_ = a_rows;
+  N_ = b_cols;
+  scale_type_ = scale_type;
+  handle_ = device.get_cublaslt_handle();
+  pref_ = cublas_utils::get_preference(device);
+  heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;
+
+  CHECK_CUBLAS_ERROR(
+      cublasLtMatmulDescCreate(&matmul_desc_, compute_type, scale_type));
+
+  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(int32_t)));
+
+  // In cublasLt matrices use column-major layout, while it is possible to use
+  // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
+  // epilogue does not work with the option. So instead we swap A and B to make
+  // cublasLt return the row-major result, which works because:
+  // - the data of a matrix in row-major layout is identical to its transpose in
+  //   column-major layout
+  // - C^T = (A @ B)^T = B^T @ A^T
+  cublasOperation_t a_op = b_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_TRANSA,
+      &a_op,
+      sizeof(cublasOperation_t)));
+  cublasOperation_t b_op = a_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_TRANSB,
+      &b_op,
+      sizeof(cublasOperation_t)));
+
+  a_desc_ = cublas_utils::create_matrix_layout(
+      data_type,
+      b_cols,
+      b_rows,
+      b_transposed,
+      ldb,
+      batch_count,
+      b_batch_stride);
+  b_desc_ = cublas_utils::create_matrix_layout(
+      data_type,
+      a_cols,
+      a_rows,
+      a_transposed,
+      lda,
+      batch_count,
+      a_batch_stride);
+  out_desc_ = cublas_utils::create_matrix_layout(
+      output_type, b_cols, a_rows, false, b_cols, batch_count, b_cols * a_rows);
+}
+
+void CublasMatmulBase::execute_matmul(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* c,
+    const void* alpha_ptr,
+    const void* beta_ptr) {
+  if (heuristic_.state != CUBLAS_STATUS_SUCCESS) {
+    int ret = 0;
+    CHECK_CUBLAS_ERROR(cublasLtMatmulAlgoGetHeuristic(
+        handle_,
+        matmul_desc_,
+        a_desc_,
+        b_desc_,
+        c ? c_desc_ : out_desc_,
+        out_desc_,
+        pref_,
+        1,
+        &heuristic_,
+        &ret));
+    if (ret == 0) {
+      throw std::runtime_error("Can not find algorithm for matmul.");
+    }
+  }
+
+  void* workspace_ptr = allocate_workspace(encoder, heuristic_.workspaceSize);
+
+  // Execute matmul
+  auto capture = encoder.capture_context();
+  CHECK_CUBLAS_ERROR(cublasLtMatmul(
+      handle_,
+      matmul_desc_,
+      alpha_ptr,
+      b, // a and b are swapped for row-major layout
+      a_desc_,
+      a,
+      b_desc_,
+      beta_ptr,
+      c ? c : out,
+      c ? c_desc_ : out_desc_,
+      out,
+      out_desc_,
+      &heuristic_.algo,
+      workspace_ptr,
+      heuristic_.workspaceSize,
+      encoder.stream()));
+}
+
+void CublasMatmulBase::set_bias(
+    cu::CommandEncoder& encoder,
+    const array& bias) {
+  encoder.set_input_array(bias);
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_EPILOGUE,
+      &epilogue,
+      sizeof(epilogue)));
+  auto* bias_ptr = gpu_ptr<void>(bias);
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_BIAS_POINTER,
+      &bias_ptr,
+      sizeof(bias_ptr)));
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/cublas_utils.h b/macos/third_party/mlx/mlx/backend/cuda/cublas_utils.h
new file mode 100644
index 0000000..56702b1
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/cublas_utils.h
@@ -0,0 +1,95 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+#include <cublasLt.h>
+#include "mlx/array.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core {
+namespace cublas_utils {
+
+// Get the shared cublas preference for a device
+cublasLtMatmulPreference_t get_preference(cu::Device& device);
+
+cublasLtMatrixLayout_t create_matrix_layout(
+    cudaDataType_t type,
+    uint64_t rows,
+    uint64_t cols,
+    bool transposed,
+    int64_t ld,
+    int32_t batch_count,
+    int64_t batch_stride);
+
+inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) {
+  switch (dtype) {
+    case float16:
+      return CUDA_R_16F;
+    case bfloat16:
+      return CUDA_R_16BF;
+    case float32:
+      return CUDA_R_32F;
+    case float64:
+      return CUDA_R_64F;
+    case complex64:
+      return CUDA_C_32F;
+    default:
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype)));
+  }
+}
+
+} // namespace cublas_utils
+
+class CublasMatmulBase {
+ public:
+  virtual ~CublasMatmulBase();
+
+  void set_bias(cu::CommandEncoder& encoder, const array& bias);
+
+ protected:
+  CublasMatmulBase() = default;
+
+  // Common member variables shared by all matmul types
+  uint64_t M_;
+  uint64_t N_;
+  cudaDataType_t scale_type_;
+  cublasLtMatmulPreference_t pref_{nullptr};
+  cublasLtHandle_t handle_{nullptr};
+  cublasLtMatmulDesc_t matmul_desc_{nullptr};
+  cublasLtMatrixLayout_t a_desc_{nullptr};
+  cublasLtMatrixLayout_t b_desc_{nullptr};
+  cublasLtMatrixLayout_t c_desc_{nullptr};
+  cublasLtMatrixLayout_t out_desc_{nullptr};
+  cublasLtMatmulHeuristicResult_t heuristic_;
+
+  void init_base(
+      cu::Device& device,
+      cudaDataType_t scale_type,
+      cublasComputeType_t compute_type,
+      cudaDataType_t data_type,
+      cudaDataType_t output_type,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride);
+
+  void execute_matmul(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* c,
+      const void* alpha_ptr,
+      const void* beta_ptr);
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/cuda.h b/macos/third_party/mlx/mlx/backend/cuda/cuda.h
new file mode 100644
index 0000000..410bae4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/cuda.h
@@ -0,0 +1,21 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+
+namespace mlx::core::cu {
+
+/* Check if the CUDA backend is available. */
+MLX_API bool is_available();
+
+/* Get information about a CUDA device. */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(int device_index = 0);
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/cuda_utils.h b/macos/third_party/mlx/mlx/backend/cuda/cuda_utils.h
new file mode 100644
index 0000000..4c60fec
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/cuda_utils.h
@@ -0,0 +1,90 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cublasLt.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cudnn.h>
+
+namespace mlx::core {
+
+// Throw exception if the cuda API does not succeed.
+void check_cublas_error(const char* name, cublasStatus_t err);
+void check_cuda_error(const char* name, cudaError_t err);
+void check_cuda_error(const char* name, CUresult err);
+void check_cudnn_error(const char* name, cudnnStatus_t err);
+
+// The macro version that prints the command that failed.
+#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
+#define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
+#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
+
+// Base class for RAII managed CUDA resources.
+template <typename Handle, cudaError_t (*Destroy)(Handle)>
+class CudaHandle {
+ public:
+  CudaHandle(Handle handle = nullptr) : handle_(handle) {}
+
+  CudaHandle(CudaHandle&& other) : handle_(other.handle_) {
+    assert(this != &other);
+    other.handle_ = nullptr;
+  }
+
+  ~CudaHandle() {
+    // Skip if there was an error to avoid throwing in the destructors
+    if (cudaPeekAtLastError() != cudaSuccess) {
+      return;
+    }
+    reset();
+  }
+
+  CudaHandle(const CudaHandle&) = delete;
+  CudaHandle& operator=(const CudaHandle&) = delete;
+
+  CudaHandle& operator=(CudaHandle&& other) {
+    assert(this != &other);
+    reset();
+    std::swap(handle_, other.handle_);
+    return *this;
+  }
+
+  void reset() {
+    if (handle_ != nullptr) {
+      CHECK_CUDA_ERROR(Destroy(handle_));
+      handle_ = nullptr;
+    }
+  }
+
+  operator Handle() const {
+    return handle_;
+  }
+
+ protected:
+  Handle handle_;
+};
+
+namespace cu {
+class Device;
+}; // namespace cu
+
+// Wrappers of CUDA resources.
+class CudaGraph : public CudaHandle<cudaGraph_t, cudaGraphDestroy> {
+ public:
+  using CudaHandle::CudaHandle;
+  explicit CudaGraph(cu::Device& device);
+  void end_capture(cudaStream_t stream);
+};
+
+class CudaGraphExec : public CudaHandle<cudaGraphExec_t, cudaGraphExecDestroy> {
+ public:
+  void instantiate(cudaGraph_t graph);
+};
+
+class CudaStream : public CudaHandle<cudaStream_t, cudaStreamDestroy> {
+ public:
+  using CudaHandle::CudaHandle;
+  explicit CudaStream(cu::Device& device);
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/cudnn_utils.cpp b/macos/third_party/mlx/mlx/backend/cuda/cudnn_utils.cpp
new file mode 100644
index 0000000..af2647f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/cudnn_utils.cpp
@@ -0,0 +1,133 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+namespace {
+
+#define RETURN_IF_ERROR(cmd)          \
+  if (auto ret = cmd; ret.is_bad()) { \
+    return ret;                       \
+  }
+
+// In MLX a singleton dim (shape[dim] == 1) can have any stride, but in cuDNN
+// whether a tensor is contiguous is determined with:
+// shape[dim] == shape[dim + 1] * strides[dim + 1]
+// So a contiguous array with singleton dims in MLX may be mistakenly treated
+// as strided in cuDNN, and we work around it by normalizing the strides.
+std::vector<int64_t> normalized_strides(const array& x) {
+  std::vector<int64_t> strides(x.strides().begin(), x.strides().end());
+  if (std::all_of(
+          strides.begin(), strides.end(), [](int64_t s) { return s == 0; })) {
+    strides.back() = 1;
+    return strides;
+  }
+  if (!x.flags().row_contiguous || x.ndim() < 2) {
+    return strides;
+  }
+  for (int i = x.ndim() - 2; i >= 0; --i) {
+    if (x.shape(i) == 1) {
+      strides[i] = x.shape(i + 1) * strides[i + 1];
+    }
+  }
+  return strides;
+}
+
+// Return the shape and strides after transposing from NHWC to NCHW.
+inline auto nhwc_to_nchw(const array& x) {
+  auto shape = convert_vector<int64_t>(x.shape());
+  auto strides = normalized_strides(x);
+  assert(shape.size() >= 3);
+  shape.insert(shape.begin() + 1, shape.back());
+  shape.erase(shape.end() - 1);
+  strides.insert(strides.begin() + 1, strides.back());
+  strides.erase(strides.end() - 1);
+  return std::make_tuple(std::move(shape), std::move(strides));
+}
+
+} // namespace
+
+fe::error_t DnnGraph::prepare() {
+  RETURN_IF_ERROR(validate());
+  try {
+    RETURN_IF_ERROR(build_operation_graph(handle_));
+  } catch (cudnn_frontend::cudnnException& error) {
+    // cuDNN bug: they did not catch all exceptions in the API.
+    return {fe::error_code_t::CUDNN_BACKEND_API_FAILED, error.what()};
+  }
+  RETURN_IF_ERROR(create_execution_plans({fe::HeurMode_t::A}));
+  return {};
+}
+
+fe::error_t DnnGraph::build() {
+  RETURN_IF_ERROR(check_support(handle_));
+  RETURN_IF_ERROR(build_plans(handle_));
+  return {};
+}
+
+fe::error_t DnnGraph::encode_graph(
+    cu::CommandEncoder& encoder,
+    std::unordered_map<int64_t, void*> variant_pack) {
+  cudnnSetStream(handle_, encoder.stream());
+  CudaGraph cuda_graph(encoder.device());
+  RETURN_IF_ERROR(populate_cuda_graph(
+      handle_, variant_pack, prepare_workspace(encoder), cuda_graph));
+  encoder.add_graph_node(cuda_graph);
+  return {};
+}
+
+fe::error_t DnnGraph::encode_capturing(
+    cu::CommandEncoder& encoder,
+    std::unordered_map<int64_t, void*> variant_pack) {
+  auto* workspace_ptr = prepare_workspace(encoder);
+  auto capture = encoder.capture_context();
+  cudnnSetStream(handle_, encoder.stream());
+  auto ret = execute(handle_, variant_pack, workspace_ptr);
+  if (ret.is_bad()) {
+    capture.discard = true;
+  }
+  return ret;
+}
+
+void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) {
+  int64_t workspace_size = 0;
+  CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size));
+  return allocate_workspace(encoder, workspace_size);
+}
+
+void DnnGraph::set_tensor_attrs(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x,
+    const std::vector<int64_t>& shape,
+    const std::vector<int64_t>& strides) {
+  tensor->set_uid(uid)
+      .set_alignment(get_alignment(x))
+      .set_data_type(dtype_to_cudnn_type(x.dtype()))
+      .set_dim(shape)
+      .set_stride(strides);
+}
+
+void DnnGraph::set_tensor_attrs(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x) {
+  set_tensor_attrs(
+      tensor,
+      uid,
+      x,
+      convert_vector<int64_t>(x.shape()),
+      normalized_strides(x));
+}
+
+void DnnGraph::set_tensor_attrs_nchw(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x) {
+  auto [shape, strides] = nhwc_to_nchw(x);
+  set_tensor_attrs(tensor, uid, x, shape, strides);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/cudnn_utils.h b/macos/third_party/mlx/mlx/backend/cuda/cudnn_utils.h
new file mode 100644
index 0000000..ef93021
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/cudnn_utils.h
@@ -0,0 +1,173 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/config.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/dtype_utils.h"
+
+#include <cudnn_frontend.h>
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+namespace cu {
+class CommandEncoder;
+}
+
+namespace fe = cudnn_frontend;
+
+#define CHECK_CUDNN_FE_ERROR(cmd)                                    \
+  do {                                                               \
+    auto error = cmd;                                                \
+    if (!error.is_good()) {                                          \
+      throw std::runtime_error(                                      \
+          fmt::format("{} failed: {}.", #cmd, error.get_message())); \
+    }                                                                \
+  } while (0)
+
+// Return pointer alignment of |x|'s data.
+inline uint8_t get_alignment(const array& x) {
+  uint8_t alignment = 1;
+  uintptr_t address = reinterpret_cast<uintptr_t>(gpu_ptr<void>(x));
+  for (; alignment < 32; alignment *= 2) {
+    if (address % (alignment * 2)) {
+      return alignment;
+    }
+  }
+  return alignment;
+}
+
+// Convert the type of elements in |vec| to |T|.
+template <typename T, typename Vec>
+inline std::vector<T> convert_vector(const Vec& vec) {
+  return std::vector<T>(vec.begin(), vec.end());
+}
+
+// Map dtype to cudnn data type.
+inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
+  switch (dtype) {
+    case int8:
+      return fe::DataType_t::INT8;
+    case int32:
+      return fe::DataType_t::INT32;
+    case uint8:
+      return fe::DataType_t::UINT8;
+    case float16:
+      return fe::DataType_t::HALF;
+    case bfloat16:
+      return fe::DataType_t::BFLOAT16;
+    case float32:
+      return fe::DataType_t::FLOAT;
+    case float64:
+      return fe::DataType_t::DOUBLE;
+    default:
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype)));
+  }
+}
+
+// Return an array that can be used as map key for |vec| with size <= MAX_NDIM.
+//
+// There are 2 differences from the const_param util from kernel_utils.cuh:
+// 1. The rest of array is filled with 0.
+// 2. This util can be used in .cpp files.
+template <int NDIM = MAX_NDIM, typename Vec>
+inline std::array<typename Vec::value_type, NDIM> vector_key(const Vec& vec) {
+  if (vec.size() > NDIM) {
+    throw std::runtime_error(
+        fmt::format("ndim can not be larger than {}.", NDIM));
+  }
+  std::array<typename Vec::value_type, NDIM> result = {};
+  std::copy_n(vec.begin(), vec.size(), result.begin());
+  return result;
+}
+
+// Extends cuDNN graph with helpers.
+class DnnGraph : public fe::graph::Graph {
+ public:
+  DnnGraph(cudnnHandle_t handle, Dtype io_dtype, Dtype compute_dtype = float32)
+      : handle_(handle) {
+    set_io_data_type(dtype_to_cudnn_type(io_dtype));
+    set_intermediate_data_type(dtype_to_cudnn_type(compute_dtype));
+    set_compute_data_type(dtype_to_cudnn_type(compute_dtype));
+  }
+
+  // Create a cuDNN tensor description from MLX array |x|.
+  auto& tensor(
+      std::shared_ptr<fe::graph::Tensor_attributes>& attrs,
+      int64_t uid,
+      const array& x) {
+    set_tensor_attrs(attrs, uid, x);
+    return attrs;
+  }
+  auto tensor(const char* name, int64_t uid, const array& x) {
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    tensor(attrs, uid, x);
+    return attrs;
+  }
+
+  // Create a cuDNN tensor description from MLX array |x|, and transpose it from
+  // NHWC layout to NCHW.
+  auto& tensor_nchw(
+      std::shared_ptr<fe::graph::Tensor_attributes>& attrs,
+      int64_t uid,
+      const array& x) {
+    set_tensor_attrs_nchw(attrs, uid, x);
+    return attrs;
+  }
+  auto tensor_nchw(const char* name, int64_t uid, const array& x) {
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    tensor_nchw(attrs, uid, x);
+    return attrs;
+  }
+
+  // Create a cuDNN tensor for scalar.
+  auto scalar(const char* name, int64_t uid, Dtype dtype) {
+    return Graph::tensor(
+        fe::graph::Tensor_attributes()
+            .set_name(name)
+            .set_uid(uid)
+            .set_dim({1, 1, 1, 1})
+            .set_stride({1, 1, 1, 1})
+            .set_is_pass_by_value(true)
+            .set_data_type(dtype_to_cudnn_type(dtype)));
+  }
+
+  // Call this before setting notes.
+  fe::error_t prepare();
+  // Call this after setting notes.
+  fe::error_t build();
+
+  // Add cuDNN graph to CUDA graph, using native CUDA graph API.
+  fe::error_t encode_graph(
+      cu::CommandEncoder& encoder,
+      std::unordered_map<int64_t, void*> variant_pack);
+  // Add cuDNN graph to CUDA graph, using stream capture.
+  fe::error_t encode_capturing(
+      cu::CommandEncoder& encoder,
+      std::unordered_map<int64_t, void*> variant_pack);
+
+ private:
+  void* prepare_workspace(cu::CommandEncoder& encoder);
+
+  void set_tensor_attrs(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x,
+      const std::vector<int64_t>& shape,
+      const std::vector<int64_t>& strides);
+  void set_tensor_attrs(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x);
+  void set_tensor_attrs_nchw(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x);
+
+  cudnnHandle_t handle_;
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/custom_kernel.cpp b/macos/third_party/mlx/mlx/backend/cuda/custom_kernel.cpp
new file mode 100644
index 0000000..7d2d0d5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/custom_kernel.cpp
@@ -0,0 +1,379 @@
+// Copyright © 2025 Apple Inc.
+
+#include <iostream>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/fast.h"
+#include "mlx/fast_primitives.h"
+
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core::fast {
+
+namespace {
+
+constexpr const char* default_header = R"(
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+#define inf cuda::std::numeric_limits<float>::infinity()
+
+)";
+
+std::string template_arguments_hash(
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args) {
+  if (template_args.empty()) {
+    return "";
+  }
+
+  std::string hash;
+  hash.reserve(512);
+
+  for (const auto& [name, arg] : template_args) {
+    if (std::holds_alternative<int>(arg)) {
+      hash += fmt::format("_{}", std::get<int>(arg));
+    } else if (std::holds_alternative<bool>(arg)) {
+      hash += (std::get<bool>(arg)) ? "_t" : "_f";
+    } else if (std::holds_alternative<Dtype>(arg)) {
+      hash += "_";
+      hash += get_type_string(std::get<Dtype>(arg));
+    }
+  }
+
+  return hash;
+}
+
+std::string build_kernel(
+    const std::string& func_name,
+    const std::string& header,
+    const std::string& source,
+    const std::vector<std::string>& input_names,
+    const std::vector<array>& inputs,
+    const std::vector<std::string>& output_names,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args,
+    const std::vector<std::tuple<bool, bool, bool>>& shape_infos) {
+  std::string kernel_source;
+  kernel_source.reserve(header.size() + source.size() + 8192);
+  kernel_source += default_header;
+  kernel_source += header;
+  kernel_source +=
+      "namespace mlx::core::cu {\n\n"
+      "namespace cg = cooperative_groups;\n\n";
+
+  kernel_source += "__global__ void ";
+  kernel_source += func_name;
+  kernel_source += "(\n";
+
+  // Add inputs
+  for (int i = 0; i < inputs.size(); ++i) {
+    const auto& name = input_names[i];
+    const auto& arr = inputs[i];
+    kernel_source += "    const ";
+    kernel_source += dtype_to_cuda_type(arr.dtype());
+    kernel_source += "* ";
+    kernel_source += name;
+    kernel_source += ",\n";
+    // Add input shape, strides and ndim if present in the source
+    if (arr.ndim() > 0) {
+      if (std::get<0>(shape_infos[i])) {
+        kernel_source += "    const __grid_constant__ Shape ";
+        kernel_source += name;
+        kernel_source += "_shape,\n";
+      }
+      if (std::get<1>(shape_infos[i])) {
+        kernel_source += "    const __grid_constant__ Strides ";
+        kernel_source += name;
+        kernel_source += "_strides,\n";
+      }
+      if (std::get<2>(shape_infos[i])) {
+        kernel_source += "    const __grid_constant__ int ";
+        kernel_source += name;
+        kernel_source += "_ndim,\n";
+      }
+    }
+  }
+
+  // Add outputs
+  for (int i = 0; i < output_names.size(); ++i) {
+    const auto& name = output_names[i];
+    const auto& dtype = output_dtypes[i];
+    kernel_source += "    ";
+    kernel_source += dtype_to_cuda_type(dtype);
+    kernel_source += "* ";
+    kernel_source += name;
+    if (i < output_names.size() - 1) {
+      kernel_source += ",\n";
+    } else {
+      kernel_source += ") {\n";
+    }
+  }
+
+  // Set compile time constants
+  if (!template_args.empty()) {
+    for (const auto& [name, arg] : template_args) {
+      if (std::holds_alternative<int>(arg)) {
+        kernel_source +=
+            fmt::format("  constexpr int {} = {};\n", name, std::get<int>(arg));
+      } else if (std::holds_alternative<bool>(arg)) {
+        kernel_source += fmt::format(
+            "  constexpr bool {} = {};\n", name, std::get<bool>(arg));
+      } else {
+        kernel_source += fmt::format(
+            "  using {} = {};\n",
+            name,
+            dtype_to_cuda_type(std::get<Dtype>(arg)));
+      }
+    }
+    kernel_source += "\n";
+  }
+
+  kernel_source += source;
+  kernel_source += "\n}\n\n} // namespace mlx::core::cu\n";
+
+  return kernel_source;
+}
+
+} // namespace
+
+CustomKernelFunction cuda_kernel(
+    const std::string& name,
+    const std::vector<std::string>& input_names,
+    const std::vector<std::string>& output_names,
+    const std::string& source,
+    const std::string& header,
+    bool ensure_row_contiguous,
+    int shared_memory) {
+  if (output_names.empty()) {
+    throw std::invalid_argument(
+        "[custom_kernel] Must specify at least one output.");
+  }
+
+  std::vector<std::tuple<bool, bool, bool>> shape_infos;
+  for (auto& n : input_names) {
+    std::tuple<bool, bool, bool> shape_info;
+    std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos;
+    std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos;
+    std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos;
+    shape_infos.push_back(shape_info);
+  }
+
+  return [=, shape_infos = std::move(shape_infos)](
+             const std::vector<array>& inputs,
+             const std::vector<Shape>& output_shapes,
+             const std::vector<Dtype>& output_dtypes,
+             std::tuple<int, int, int> grid,
+             std::tuple<int, int, int> threadgroup,
+             const std::vector<std::pair<std::string, TemplateArg>>&
+                 template_args = {},
+             std::optional<float> init_value = std::nullopt,
+             bool verbose = false,
+             StreamOrDevice s_ = {}) {
+    if (inputs.size() != input_names.size()) {
+      std::ostringstream msg;
+      msg << "[custom_kernel] Expected `inputs` to have size "
+          << input_names.size() << " but got size " << inputs.size() << "."
+          << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_shapes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[custom_kernel] Expected `output_shapes` to have size "
+          << output_names.size() << " but got size " << output_shapes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_dtypes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[custom_kernel] Expected `output_dtypes` to have size "
+          << output_names.size() << " but got size " << output_dtypes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+
+    auto s = to_stream(s_);
+    if (s.device != Device::gpu) {
+      throw std::invalid_argument("[custom_kernel] Only supports the GPU.");
+    }
+
+    std::string kernel_name =
+        "custom_kernel_" + name + template_arguments_hash(template_args);
+    std::string kernel_source = build_kernel(
+        kernel_name,
+        header,
+        source,
+        input_names,
+        inputs,
+        output_names,
+        output_dtypes,
+        template_args,
+        shape_infos);
+
+    if (verbose) {
+      std::cout << "Generated source code for `" << kernel_name
+                << "`:" << std::endl
+                << "```" << std::endl
+                << kernel_source << std::endl
+                << "```" << std::endl;
+    }
+
+    return array::make_arrays(
+        std::move(output_shapes),
+        std::move(output_dtypes),
+        std::make_shared<CustomKernel>(
+            s,
+            std::move(kernel_name),
+            std::move(kernel_source),
+            grid,
+            threadgroup,
+            shape_infos,
+            ensure_row_contiguous,
+            init_value,
+            std::vector<ScalarArg>{},
+            false,
+            shared_memory),
+        std::move(inputs));
+  };
+}
+
+std::vector<array> precompiled_cuda_kernel(
+    const std::string& name,
+    const std::string& compiled_source,
+    const std::vector<array>& inputs,
+    const std::vector<Shape>& output_shapes,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<ScalarArg>& scalars,
+    std::tuple<int, int, int> grid,
+    std::tuple<int, int, int> threadgroup,
+    int shared_memory,
+    std::optional<float> init_value,
+    bool ensure_row_contiguous,
+    StreamOrDevice s) {
+  std::vector<std::tuple<bool, bool, bool>> shape_infos(
+      inputs.size(), {false, false, false});
+  return array::make_arrays(
+      output_shapes,
+      output_dtypes,
+      std::make_shared<CustomKernel>(
+          to_stream(s),
+          name,
+          compiled_source,
+          grid,
+          threadgroup,
+          shape_infos,
+          ensure_row_contiguous,
+          init_value,
+          scalars,
+          true,
+          shared_memory),
+      inputs);
+}
+
+void CustomKernel::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("CustomKernel::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  std::vector<array> copies;
+
+  // Allocate and initialize the output arrays
+  for (auto& out : outputs) {
+    if (init_value_) {
+      copies.emplace_back(init_value_.value(), out.dtype());
+      fill_gpu(copies.back(), out, s);
+    } else {
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    }
+  }
+
+  // Create the input arrays and copy if needed
+  auto check_input = [&copies, &s, this](const array& x) -> const array {
+    bool no_copy = x.flags().row_contiguous;
+    if (!ensure_row_contiguous_ || no_copy) {
+      return x;
+    } else {
+      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
+      copy_gpu(x, copies.back(), CopyType::General, s);
+      return copies.back();
+    }
+  };
+  std::vector<array> checked_inputs;
+  for (const array& in : inputs) {
+    checked_inputs.push_back(check_input(in));
+  }
+
+  // Compile the custom kernel
+  std::string kernel_name =
+      (is_precompiled_) ? name_ : "mlx::core::cu::" + name_;
+  cu::JitModule& mod = cu::get_jit_module(
+      s.device,
+      name_,
+      [&]() {
+        return std::make_tuple(
+            is_precompiled_, source_, std::vector{kernel_name});
+      },
+      false);
+
+  // Make the arguments
+  cu::KernelArgs args;
+  for (int i = 0; i < checked_inputs.size(); i++) {
+    const array& in = checked_inputs[i];
+    auto& shape_info = shape_infos_[i];
+    args.append(in);
+    if (std::get<0>(shape_info)) {
+      args.append_ndim(in.shape());
+    }
+    if (std::get<1>(shape_info)) {
+      args.append_ndim(in.strides());
+    }
+    if (std::get<2>(shape_info)) {
+      args.append<int32_t>(in.ndim());
+    }
+  }
+  for (auto& out : outputs) {
+    args.append(out);
+  }
+  for (auto& s : scalar_arguments_) {
+    if (std::holds_alternative<bool>(s)) {
+      args.append(std::get<bool>(s));
+    } else if (std::holds_alternative<int>(s)) {
+      args.append(std::get<int>(s));
+    } else if (std::holds_alternative<float>(s)) {
+      args.append(std::get<float>(s));
+    }
+  }
+
+  // Make the grid
+  const auto [tx, ty, tz] = threadgroup_;
+  const auto [gx, gy, gz] = grid_;
+  dim3 block(std::min(tx, gx), std::min(ty, gy), std::min(tz, gz));
+  dim3 grid((gx + tx - 1) / tx, (gy + ty - 1) / ty, (gz + tz - 1) / tz);
+
+  // Call the kernel
+  for (const auto& in : checked_inputs) {
+    encoder.set_input_array(in);
+  }
+  for (const auto& out : outputs) {
+    encoder.set_output_array(out);
+  }
+  for (const auto& t : copies) {
+    encoder.add_temporary(t);
+  }
+  auto kernel =
+      mod.get_kernel(kernel_name, [smem = shared_memory_](CUfunction kernel) {
+        if (smem > 0 && smem > 48000) {
+          cuFuncSetAttribute(
+              kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem);
+        }
+      });
+  encoder.add_kernel_node(kernel, grid, block, shared_memory_, args.args());
+}
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/backend/cuda/cutlass_utils.cuh b/macos/third_party/mlx/mlx/backend/cuda/cutlass_utils.cuh
new file mode 100644
index 0000000..c9bbf0d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/cutlass_utils.cuh
@@ -0,0 +1,46 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/dtype.h"
+
+#include <cutlass/bfloat16.h>
+#include <cutlass/half.h>
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+// Throw exception if the cutlass API does not succeed.
+inline void check_cutlass_error(const char* name, cutlass::Status status) {
+  if (status != cutlass::Status::kSuccess) {
+    throw std::runtime_error(
+        fmt::format(
+            "{} failed with code: {}.",
+            name,
+            cutlass::cutlassGetStatusString(status)));
+  }
+}
+
+// The macro version that prints the command that failed.
+#define CHECK_CUTLASS_ERROR(cmd) check_cutlass_error(#cmd, (cmd))
+
+// Maps CPU types to CUTLASS types.
+template <typename T>
+struct CTypeToCutlassType {
+  using type = T;
+};
+
+template <>
+struct CTypeToCutlassType<float16_t> {
+  using type = cutlass::half_t;
+};
+
+template <>
+struct CTypeToCutlassType<bfloat16_t> {
+  using type = cutlass::bfloat16_t;
+};
+
+template <typename T>
+using cutlass_type_t = typename CTypeToCutlassType<T>::type;
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/delayload.cpp b/macos/third_party/mlx/mlx/backend/cuda/delayload.cpp
new file mode 100644
index 0000000..0ac1cc5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/delayload.cpp
@@ -0,0 +1,80 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+
+// clang-format off
+#include <windows.h> // must be included first
+#include <delayimp.h>
+// clang-format on
+
+namespace mlx::core {
+
+namespace fs = std::filesystem;
+
+inline fs::path relative_to_current_binary(const char* relative) {
+  return fs::absolute(current_binary_dir() / relative);
+}
+
+inline fs::path cublas_bin_dir() {
+#if defined(MLX_CUDA_BIN_DIR)
+  return MLX_CUDA_BIN_DIR;
+#else
+  return relative_to_current_binary("../nvidia/cublas/bin");
+#endif
+}
+
+fs::path load_nvrtc() {
+#if defined(MLX_CUDA_BIN_DIR)
+  fs::path nvrtc_bin_dir = MLX_CUDA_BIN_DIR;
+#else
+  fs::path nvrtc_bin_dir =
+      relative_to_current_binary("../nvidia/cuda_nvrtc/bin");
+#endif
+  // Internally nvrtc loads some libs dynamically, add to search dirs.
+  ::AddDllDirectory(nvrtc_bin_dir.c_str());
+  return nvrtc_bin_dir;
+}
+
+fs::path load_cudnn() {
+#if defined(MLX_CUDNN_BIN_DIR)
+  fs::path cudnn_bin_dir = MLX_CUDNN_BIN_DIR;
+#else
+  fs::path cudnn_bin_dir = relative_to_current_binary("../nvidia/cudnn/bin");
+#endif
+  // Must load cudnn_graph64_9.dll before locating symbols, otherwise We would
+  // get errors like "Invalid handle. Cannot load symbol cudnnCreate".
+  for (const auto& dll : fs::directory_iterator(cudnn_bin_dir)) {
+    if (dll.path().filename().string().starts_with("cudnn_graph") &&
+        dll.path().extension() == ".dll") {
+      ::LoadLibraryW(dll.path().c_str());
+      break;
+    }
+  }
+  // Internally cuDNN loads some libs dynamically, add to search dirs.
+  load_nvrtc();
+  ::AddDllDirectory(cudnn_bin_dir.c_str());
+  ::AddDllDirectory(cublas_bin_dir().c_str());
+  return cudnn_bin_dir;
+}
+
+// Called by system when failed to locate a lazy-loaded DLL.
+FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) {
+  HMODULE mod = NULL;
+  if (dliNotify == dliNotePreLoadLibrary) {
+    std::string dll = pdli->szDll;
+    if (dll.starts_with("cudnn")) {
+      static auto cudnn_bin_dir = load_cudnn();
+      mod = ::LoadLibraryW((cudnn_bin_dir / dll).c_str());
+    } else if (dll.starts_with("cublas")) {
+      mod = ::LoadLibraryW((cublas_bin_dir() / dll).c_str());
+    } else if (dll.starts_with("nvrtc")) {
+      static auto nvrtc_bin_dir = load_nvrtc();
+      mod = ::LoadLibraryW((nvrtc_bin_dir / dll).c_str());
+    }
+  }
+  return reinterpret_cast<FARPROC>(mod);
+}
+
+} // namespace mlx::core
+
+extern "C" const PfnDliHook __pfnDliNotifyHook2 = mlx::core::delayload_helper;
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device.cpp b/macos/third_party/mlx/mlx/backend/cuda/device.cpp
new file mode 100644
index 0000000..1d7ad7d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device.cpp
@@ -0,0 +1,522 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/cuda/worker.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/utils.h"
+
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+#include <future>
+#include <unordered_set>
+
+namespace mlx::core::cu {
+
+namespace {
+
+bool use_cuda_graphs() {
+  static bool use_graphs = env::get_var("MLX_USE_CUDA_GRAPHS", true);
+  return use_graphs;
+}
+
+const char* save_cuda_graphs_dot_file() {
+  static const char* filename = []() -> const char* {
+    const char* env = std::getenv("MLX_SAVE_CUDA_GRAPHS_DOT_FILE");
+    if (env && std::strlen(env) == 0) {
+      return nullptr;
+    }
+    return env;
+  }();
+  return filename;
+}
+
+} // namespace
+
+Device::Device(int device) : device_(device) {
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &compute_capability_major_, cudaDevAttrComputeCapabilityMajor, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &compute_capability_minor_, cudaDevAttrComputeCapabilityMinor, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &concurrent_managed_access_,
+      cudaDevAttrConcurrentManagedAccess,
+      device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &host_native_atomic_, cudaDevAttrHostNativeAtomicSupported, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &managed_memory_, cudaDevAttrManagedMemory, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &memory_pools_, cudaDevAttrMemoryPoolsSupported, device_));
+}
+
+Device::~Device() {
+  if (cudnn_handle_) {
+    CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_handle_));
+  }
+  if (cublaslt_handle_) {
+    CHECK_CUBLAS_ERROR(cublasLtDestroy(cublaslt_handle_));
+  }
+}
+
+void Device::make_current() {
+  // We need to set/get current CUDA device very frequently, cache it to reduce
+  // actual calls of CUDA APIs.
+  static thread_local int current = 0;
+  if (current != device_) {
+    CHECK_CUDA_ERROR(cudaSetDevice(device_));
+    current = device_;
+  }
+}
+
+CommandEncoder& Device::get_command_encoder(Stream s) {
+  auto it = encoders_.find(s.index);
+  if (it == encoders_.end()) {
+    it = encoders_.try_emplace(s.index, *this).first;
+  }
+  return it->second;
+}
+
+cublasLtHandle_t Device::get_cublaslt_handle() {
+  if (!cublaslt_handle_) {
+    make_current();
+    CHECK_CUBLAS_ERROR(cublasLtCreate(&cublaslt_handle_));
+  }
+  return cublaslt_handle_;
+}
+
+cudnnHandle_t Device::get_cudnn_handle() {
+  if (!cudnn_handle_) {
+    make_current();
+    CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_handle_));
+  }
+  return cudnn_handle_;
+}
+
+CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) {
+  enc.device().make_current();
+  if (!use_cuda_graphs()) {
+    return;
+  }
+  CHECK_CUDA_ERROR(
+      cudaStreamBeginCapture(enc.stream(), cudaStreamCaptureModeThreadLocal));
+}
+
+CommandEncoder::CaptureContext::~CaptureContext() {
+  if (!use_cuda_graphs()) {
+    enc.node_count_++;
+    return;
+  }
+
+  graph.end_capture(enc.stream());
+  if (discard) {
+    return;
+  }
+  enc.add_graph_node(graph);
+}
+
+CommandEncoder::ConcurrentContext::ConcurrentContext(CommandEncoder& enc)
+    : enc(enc) {
+  enc.in_concurrent_ = true;
+}
+
+CommandEncoder::ConcurrentContext::~ConcurrentContext() {
+  enc.in_concurrent_ = false;
+  if (!use_cuda_graphs()) {
+    return;
+  }
+
+  // Use an empty graph node for synchronization
+  CommandEncoder::GraphNode empty{NULL, "E", std::to_string(enc.node_count_++)};
+  CHECK_CUDA_ERROR(cudaGraphAddEmptyNode(&empty.node, enc.graph_, NULL, 0));
+
+  // Insert the concurrent -> empty node dependencies
+  for (auto& from : enc.concurrent_nodes_) {
+    enc.from_nodes_.push_back(from.node);
+    enc.to_nodes_.push_back(empty.node);
+    enc.graph_deps_key_ += from.id;
+    enc.graph_deps_key_ += "-";
+    enc.graph_deps_key_ += empty.id;
+    enc.graph_deps_key_ += "-";
+  }
+
+  // Insert the input -> concurrent node dependencies without updating output
+  // nodes
+  auto outputs = std::move(enc.active_outputs_);
+  enc.insert_graph_dependencies(std::move(enc.concurrent_nodes_));
+
+  // Update output node to be the empty node
+  for (auto o : outputs) {
+    enc.node_map_.emplace(o, empty).first->second = empty;
+  }
+}
+
+void CommandEncoder::insert_graph_dependencies(GraphNode node) {
+  node.id = std::to_string(node_count_++);
+  if (in_concurrent_) {
+    concurrent_nodes_.push_back(std::move(node));
+  } else {
+    std::vector<GraphNode> nodes;
+    nodes.push_back(std::move(node));
+    insert_graph_dependencies(std::move(nodes));
+  }
+}
+
+void CommandEncoder::insert_graph_dependencies(std::vector<GraphNode> nodes) {
+  for (auto& node : nodes) {
+    graph_nodes_key_ += node.node_type;
+    graph_nodes_key_ += "-";
+  }
+  std::vector<GraphNode> deps;
+  {
+    // Dependencies must be added in the same order to produce a consistent
+    // topology
+    std::unordered_set<cudaGraphNode_t> set_deps;
+    for (auto d : active_deps_) {
+      if (auto it = node_map_.find(d); it != node_map_.end()) {
+        auto [_, inserted] = set_deps.insert(it->second.node);
+        if (inserted) {
+          deps.push_back(it->second);
+        }
+      }
+    }
+  }
+  active_deps_.clear();
+
+  for (auto o : active_outputs_) {
+    for (auto& node : nodes) {
+      node_map_.emplace(o, node).first->second = node;
+    }
+  }
+  active_outputs_.clear();
+
+  for (auto& from : deps) {
+    for (auto& to : nodes) {
+      from_nodes_.push_back(from.node);
+      to_nodes_.push_back(to.node);
+      graph_deps_key_ += from.id;
+      graph_deps_key_ += "-";
+      graph_deps_key_ += to.id;
+      graph_deps_key_ += "-";
+    }
+  }
+}
+
+// Can be tuned with MLX_MAX_OPS_PER_BUFFER, MLX_MAX_MB_PER_BUFFER
+std::pair<int, int> get_graph_limits(Device& d) {
+  auto cc =
+      d.compute_capability_major() * 100 + d.compute_capability_minor() * 10;
+  int ops = 20;
+  int mb = 100;
+  switch (cc) {
+    case 800: // A100
+      ops = 20;
+      mb = 400;
+      break;
+    case 900: // H100
+    case 1000: // B200
+    case 1200: // Consumer Blackwell
+      ops = 100;
+      mb = 1000;
+      break;
+    case 1210: // DGX Spark
+      ops = 20;
+      mb = 25;
+      break;
+  }
+  return {env::max_ops_per_buffer(ops), env::max_mb_per_buffer(mb)};
+}
+
+CommandEncoder::CommandEncoder(Device& d)
+    : device_(d),
+      stream_(d),
+      graph_(d),
+      worker_(d),
+      graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) {
+  std::tie(max_ops_per_graph_, max_mb_per_graph_) = get_graph_limits(d);
+}
+
+void CommandEncoder::add_completed_handler(std::function<void()> task) {
+  worker_.add_task(std::move(task));
+}
+
+void CommandEncoder::set_input_array(const array& arr) {
+  if (!use_cuda_graphs()) {
+    return;
+  }
+  bytes_in_graph_ += arr.data_size();
+  auto id = reinterpret_cast<std::uintptr_t>(arr.buffer().ptr());
+  active_deps_.push_back(id);
+}
+
+void CommandEncoder::set_output_array(const array& arr) {
+  if (!use_cuda_graphs()) {
+    return;
+  }
+
+  auto id = reinterpret_cast<std::uintptr_t>(arr.buffer().ptr());
+  active_deps_.push_back(id);
+  active_outputs_.push_back(id);
+}
+
+void CommandEncoder::add_kernel_node(
+    void* func,
+    dim3 grid_dim,
+    dim3 block_dim,
+    uint32_t smem_bytes,
+    void** params) {
+  if (!use_cuda_graphs()) {
+    node_count_++;
+    CHECK_CUDA_ERROR(cudaLaunchKernel(
+        func, grid_dim, block_dim, params, smem_bytes, stream()));
+    return;
+  }
+  cudaKernelNodeParams kernel_params = {0};
+  kernel_params.func = func;
+  kernel_params.gridDim = grid_dim;
+  kernel_params.blockDim = block_dim;
+  kernel_params.kernelParams = params;
+  kernel_params.sharedMemBytes = smem_bytes;
+  add_kernel_node(kernel_params);
+}
+
+void CommandEncoder::add_kernel_node(
+    CUfunction func,
+    dim3 grid_dim,
+    dim3 block_dim,
+    uint32_t smem_bytes,
+    void** params) {
+  if (!use_cuda_graphs()) {
+    node_count_++;
+    CHECK_CUDA_ERROR(cuLaunchKernel(
+        func,
+        grid_dim.x,
+        grid_dim.y,
+        grid_dim.z,
+        block_dim.x,
+        block_dim.y,
+        block_dim.z,
+        smem_bytes,
+        stream(),
+        params,
+        nullptr));
+    return;
+  }
+
+  CUDA_KERNEL_NODE_PARAMS kernel_params = {0};
+  kernel_params.func = func;
+  kernel_params.gridDimX = grid_dim.x;
+  kernel_params.gridDimY = grid_dim.y;
+  kernel_params.gridDimZ = grid_dim.z;
+  kernel_params.blockDimX = block_dim.x;
+  kernel_params.blockDimY = block_dim.y;
+  kernel_params.blockDimZ = block_dim.z;
+  kernel_params.kernelParams = params;
+  kernel_params.sharedMemBytes = smem_bytes;
+  add_kernel_node(kernel_params);
+}
+
+void CommandEncoder::add_kernel_node(const cudaKernelNodeParams& params) {
+  cudaGraphNode_t node;
+  CHECK_CUDA_ERROR(cudaGraphAddKernelNode(&node, graph_, NULL, 0, &params));
+  insert_graph_dependencies(GraphNode{node, "K"});
+}
+
+void CommandEncoder::add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params) {
+  CUgraphNode node;
+  CHECK_CUDA_ERROR(cuGraphAddKernelNode(&node, graph_, NULL, 0, &params));
+  insert_graph_dependencies(GraphNode{node, "K"});
+}
+
+std::pair<std::string, bool> subgraph_to_key(cudaGraph_t graph) {
+  // Constructs a key representing the nodes of a sub-graph.
+  // Also checks if the sub-graph is updatable as CUDA graphs do not get
+  // updated correctly if a kernel node getting updated has a different cluster
+  // shape than the node it's being updated with.
+  std::string key = "(";
+  size_t num_nodes = 0;
+  CHECK_CUDA_ERROR(cudaGraphGetNodes(graph, nullptr, &num_nodes));
+  if (num_nodes == 0) {
+    return {key + ")", true};
+  }
+  bool is_updatable = true;
+  std::vector<cudaGraphNode_t> nodes(num_nodes);
+  CHECK_CUDA_ERROR(cudaGraphGetNodes(graph, nodes.data(), &num_nodes));
+  for (const auto& node : nodes) {
+    if (!is_updatable) {
+      break;
+    }
+    cudaGraphNodeType type;
+    CHECK_CUDA_ERROR(cudaGraphNodeGetType(node, &type));
+    switch (type) {
+      case cudaGraphNodeTypeGraph: {
+        // Try to be updatable for a structure like graph -> graph -> kernel
+        cudaGraph_t child;
+        CHECK_CUDA_ERROR(cudaGraphChildGraphNodeGetGraph(node, &child));
+        auto [subkey, sub_is_updatable] = subgraph_to_key(child);
+        is_updatable &= sub_is_updatable;
+        key += subkey;
+        break;
+      }
+      case cudaGraphNodeTypeHost:
+        key += "H";
+        break;
+      case cudaGraphNodeTypeMemset:
+        key += "M";
+        break;
+      case cudaGraphNodeTypeKernel: {
+        cudaLaunchAttributeValue cluster_dim;
+        CHECK_CUDA_ERROR(cudaGraphKernelNodeGetAttribute(
+            node, cudaLaunchAttributeClusterDimension, &cluster_dim));
+        // Only allow dim.x to be greater than 1
+        if (cluster_dim.clusterDim.y > 1 || cluster_dim.clusterDim.z > 1) {
+          is_updatable = false;
+        } else {
+          key += "K";
+          key += std::to_string(cluster_dim.clusterDim.x);
+        }
+        break;
+      }
+      case cudaGraphNodeTypeWaitEvent:
+        key += "W";
+        break;
+      case cudaGraphNodeTypeEventRecord:
+        key += "R";
+        break;
+      default:
+        is_updatable = false;
+    }
+  }
+  key += ")";
+  return {key, is_updatable};
+}
+
+void CommandEncoder::add_graph_node(cudaGraph_t child) {
+  if (!use_cuda_graphs()) {
+    node_count_++;
+    CudaGraphExec graph_exec;
+    graph_exec.instantiate(child);
+    device_.make_current();
+    CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream()));
+    return;
+  }
+  cudaGraphNode_t node;
+  auto [sub_graph_key, is_updatable] = subgraph_to_key(child);
+  is_graph_updatable_ &= is_updatable;
+  CHECK_CUDA_ERROR(cudaGraphAddChildGraphNode(&node, graph_, NULL, 0, child));
+  insert_graph_dependencies(GraphNode{node, sub_graph_key});
+}
+
+bool CommandEncoder::needs_commit() {
+  return (node_count_ > max_ops_per_graph_) ||
+      ((bytes_in_graph_ >> 20) > max_mb_per_graph_);
+}
+
+void CommandEncoder::commit() {
+  nvtx3::scoped_range r("CommandEncoder::commit");
+  if (!temporaries_.empty()) {
+    add_completed_handler([temporaries = std::move(temporaries_)]() {});
+  }
+  if (use_cuda_graphs() && node_count_ > 0) {
+    if (!from_nodes_.empty()) {
+#if CUDART_VERSION >= 13000
+      CHECK_CUDA_ERROR(cudaGraphAddDependencies(
+          graph_,
+          from_nodes_.data(),
+          to_nodes_.data(),
+          nullptr, // edgeData
+          from_nodes_.size()));
+#else
+      CHECK_CUDA_ERROR(cudaGraphAddDependencies(
+          graph_, from_nodes_.data(), to_nodes_.data(), from_nodes_.size()));
+#endif
+    }
+
+    device_.make_current();
+
+    if (!is_graph_updatable_) {
+      CudaGraphExec graph_exec;
+      graph_exec.instantiate(graph_);
+      CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream_));
+    } else {
+      auto graph_key = graph_nodes_key_ + ":" + graph_deps_key_;
+      auto& graph_exec = graph_cache_[graph_key];
+
+      if (graph_exec != nullptr) {
+        cudaGraphExecUpdateResult update_result;
+#if CUDART_VERSION >= 12000
+        cudaGraphExecUpdateResultInfo info;
+        cudaGraphExecUpdate(graph_exec, graph_, &info);
+        update_result = info.result;
+#else
+        cudaGraphNode_t error_node;
+        cudaGraphExecUpdate(graph_exec, graph_, &error_node, &update_result);
+#endif // CUDART_VERSION >= 12000
+        if (update_result != cudaGraphExecUpdateSuccess) {
+          cudaGetLastError(); // reset error
+          graph_exec.reset();
+        }
+      }
+      if (graph_exec == nullptr) {
+        graph_exec.instantiate(graph_);
+      }
+
+      CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream_));
+    }
+
+    // Save cuda graph to dot file
+    if (const char* filename = save_cuda_graphs_dot_file(); filename) {
+      static int count = 0;
+      auto path = fmt::format("{}_{}.dot", filename, ++count);
+      CHECK_CUDA_ERROR(cudaGraphDebugDotPrint(graph_, path.c_str(), 0));
+    }
+
+    // Reset state
+    from_nodes_.clear();
+    to_nodes_.clear();
+    graph_deps_key_.clear();
+    graph_nodes_key_.clear();
+    node_map_.clear();
+    graph_ = CudaGraph(device_);
+    is_graph_updatable_ = true;
+  }
+
+  // Put completion handlers in a batch.
+  worker_.commit(stream_);
+  node_count_ = 0;
+  bytes_in_graph_ = 0;
+}
+
+void CommandEncoder::synchronize() {
+  CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_));
+  auto p = std::make_shared<std::promise<void>>();
+  std::future<void> f = p->get_future();
+  add_completed_handler([p = std::move(p)]() { p->set_value(); });
+  commit();
+  f.wait();
+}
+
+Device& device(int cuda_device) {
+  static auto devices = []() {
+    std::vector<Device> devices;
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      devices.emplace_back(i);
+    }
+    // Initialize the jit module cache here ensures it is not unloaded before
+    // any evaluation is done.
+    get_jit_module_cache();
+    return devices;
+  }();
+  return devices.at(cuda_device);
+}
+
+Device& device(mlx::core::Device d) {
+  return device(d.index);
+}
+
+CommandEncoder& get_command_encoder(Stream s) {
+  return device(s.device).get_command_encoder(s);
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device.h b/macos/third_party/mlx/mlx/backend/cuda/device.h
new file mode 100644
index 0000000..527827a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device.h
@@ -0,0 +1,195 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/lru_cache.h"
+#include "mlx/backend/cuda/worker.h"
+#include "mlx/stream.h"
+
+#include <cublasLt.h>
+#include <cuda.h>
+#include <cudnn.h>
+
+#include <unordered_map>
+
+namespace mlx::core::cu {
+
+class CommandEncoder {
+ public:
+  struct CaptureContext {
+    CaptureContext(CommandEncoder& enc);
+    ~CaptureContext();
+    CudaGraph graph;
+    CommandEncoder& enc;
+    bool discard{false};
+  };
+  struct ConcurrentContext {
+    ConcurrentContext(CommandEncoder& enc);
+    ~ConcurrentContext();
+    CommandEncoder& enc;
+  };
+
+  explicit CommandEncoder(Device& d);
+
+  CommandEncoder(const CommandEncoder&) = delete;
+  CommandEncoder& operator=(const CommandEncoder&) = delete;
+
+  CaptureContext capture_context() {
+    return CaptureContext{*this};
+  }
+  ConcurrentContext concurrent_context() {
+    return ConcurrentContext{*this};
+  }
+
+  void set_input_array(const array& arr);
+  void set_output_array(const array& arr);
+
+  template <typename F, typename... Params>
+  void add_kernel_node(
+      F* func,
+      dim3 grid_dim,
+      dim3 block_dim,
+      uint32_t smem_bytes,
+      Params&&... params) {
+    constexpr size_t num = sizeof...(Params);
+    void* ptrs[num];
+    size_t i = 0;
+    ([&](auto&& p) { ptrs[i++] = static_cast<void*>(&p); }(
+         std::forward<Params>(params)),
+     ...);
+    add_kernel_node((void*)func, grid_dim, block_dim, smem_bytes, ptrs);
+  }
+
+  void add_kernel_node(
+      CUfunction func,
+      dim3 grid_dim,
+      dim3 block_dim,
+      uint32_t smem_bytes,
+      void** params);
+
+  void add_kernel_node(
+      void* func,
+      dim3 grid_dim,
+      dim3 block_dim,
+      uint32_t smem_bytes,
+      void** params);
+
+  void add_graph_node(cudaGraph_t child);
+
+  void add_temporary(const array& arr) {
+    temporaries_.push_back(arr.data_shared_ptr());
+  }
+
+  void add_completed_handler(std::function<void()> task);
+  bool needs_commit();
+  void commit();
+
+  Device& device() {
+    return device_;
+  }
+
+  CudaStream& stream() {
+    return stream_;
+  }
+
+  // Wait until kernels and completion handlers are finished
+  void synchronize();
+
+ private:
+  void add_kernel_node(const cudaKernelNodeParams& params);
+  void add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params);
+
+  struct GraphNode {
+    cudaGraphNode_t node;
+    // K = kernel
+    // E = empty
+    // () = subgraph (with metadata)
+    // Symbols ':', '-' are reserved as separators
+    std::string node_type;
+    std::string id;
+  };
+
+  void insert_graph_dependencies(GraphNode node);
+  void insert_graph_dependencies(std::vector<GraphNode> nodes);
+
+  Device& device_;
+  CudaStream stream_;
+  CudaGraph graph_;
+  Worker worker_;
+  int node_count_{0};
+  bool in_concurrent_{false};
+  std::vector<cudaGraphNode_t> from_nodes_;
+  std::vector<cudaGraphNode_t> to_nodes_;
+  std::string graph_nodes_key_;
+  std::string graph_deps_key_;
+  std::vector<GraphNode> concurrent_nodes_;
+  std::vector<std::shared_ptr<array::Data>> temporaries_;
+  LRUCache<std::string, CudaGraphExec> graph_cache_;
+  std::vector<std::uintptr_t> active_deps_;
+  std::vector<std::uintptr_t> active_outputs_;
+  std::unordered_map<std::uintptr_t, GraphNode> node_map_;
+  size_t bytes_in_graph_{0};
+  bool is_graph_updatable_{true};
+  int max_ops_per_graph_;
+  int max_mb_per_graph_;
+};
+
+class Device {
+ public:
+  explicit Device(int device);
+  ~Device();
+
+  Device(Device&&) = default;
+  Device(const Device&) = delete;
+  Device& operator=(const Device&) = delete;
+
+  // Make this device the current cuda device, this method is thread-safe.
+  void make_current();
+
+  CommandEncoder& get_command_encoder(Stream s);
+  cublasLtHandle_t get_cublaslt_handle();
+  cudnnHandle_t get_cudnn_handle();
+
+  int cuda_device() const {
+    return device_;
+  }
+  int compute_capability_major() const {
+    return compute_capability_major_;
+  }
+  int compute_capability_minor() const {
+    return compute_capability_minor_;
+  }
+  bool concurrent_managed_access() const {
+    return concurrent_managed_access_ == 1;
+  }
+  bool host_native_atomic() const {
+    return host_native_atomic_ == 1;
+  }
+  bool managed_memory() const {
+    return managed_memory_ == 1;
+  }
+  bool memory_pools() const {
+    return memory_pools_ == 1;
+  }
+
+ private:
+  int device_;
+  int compute_capability_major_;
+  int compute_capability_minor_;
+  int concurrent_managed_access_;
+  int host_native_atomic_;
+  int managed_memory_;
+  int memory_pools_;
+  std::string device_name_;
+  cublasLtHandle_t cublaslt_handle_{nullptr};
+  cudnnHandle_t cudnn_handle_{nullptr};
+  std::unordered_map<int, CommandEncoder> encoders_;
+};
+
+Device& device(int cuda_device);
+Device& device(mlx::core::Device d);
+CommandEncoder& get_command_encoder(Stream s);
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/atomic_ops.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/atomic_ops.cuh
new file mode 100644
index 0000000..cfb185c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/atomic_ops.cuh
@@ -0,0 +1,63 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/complex.cuh"
+#include "mlx/backend/cuda/device/fp16_math.cuh"
+
+#include <cuda/atomic>
+
+namespace mlx::core::cu {
+
+template <typename T>
+inline __device__ void atomic_add(T* out, T val) {
+  cuda::atomic_ref<T, cuda::thread_scope_device> ref(*out);
+  ref += val;
+}
+
+template <typename T>
+inline __device__ void atomic_prod(T* out, T val) {
+  cuda::atomic_ref<T, cuda::thread_scope_device> ref(*out);
+  T old = ref.load();
+  while (!ref.compare_exchange_strong(old, old * val)) {
+  }
+}
+
+template <typename T>
+inline __device__ void atomic_max(T* out, T val) {
+  cuda::atomic_ref<T, cuda::thread_scope_device> ref(*out);
+  ref.fetch_max(val);
+}
+
+template <typename T>
+inline __device__ void atomic_min(T* out, T val) {
+  cuda::atomic_ref<T, cuda::thread_scope_device> ref(*out);
+  ref.fetch_min(val);
+}
+
+// Somehow cuda::atomic_ref does not provide atomic add for following types.
+template <typename T>
+inline __device__ void atomic_add_general(T* out, T val) {
+  cuda::atomic_ref<T, cuda::thread_scope_device> ref(*out);
+  T old = ref.load();
+  while (!ref.compare_exchange_strong(old, old + val)) {
+  }
+}
+
+inline __device__ void atomic_add(__half* out, __half val) {
+  atomicAdd(out, val);
+}
+
+inline __device__ void atomic_add(complex64_t* out, complex64_t val) {
+  atomic_add_general(out, val);
+}
+
+inline __device__ void atomic_add(__nv_bfloat16* out, __nv_bfloat16 val) {
+#if __CUDA_ARCH__ < 800
+  atomic_add_general(out, val);
+#else
+  atomicAdd(out, val);
+#endif
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/binary_ops.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/binary_ops.cuh
new file mode 100644
index 0000000..b0b7962
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/binary_ops.cuh
@@ -0,0 +1,300 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/unary_ops.cuh"
+
+#include <cuda/std/array>
+
+namespace mlx::core::cu {
+
+struct Add {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x + y;
+  }
+};
+
+struct FloorDivide {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    if constexpr (cuda::std::is_integral_v<T>) {
+      return x / y;
+    } else {
+      return cuda::std::trunc(x / y);
+    }
+  }
+};
+
+struct Divide {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x / y;
+  }
+};
+
+struct Remainder {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    if constexpr (cuda::std::is_integral_v<T>) {
+      if constexpr (cuda::std::is_signed_v<T>) {
+        auto r = x % y;
+        if (r != 0 && (r < 0 != y < 0)) {
+          r += y;
+        }
+        return r;
+      } else {
+        return x % y;
+      }
+    } else if constexpr (is_complex_v<T>) {
+      return x % y;
+    } else {
+      T r = cuda::std::fmod(x, y);
+      if (r != 0 && (r < 0 != y < 0)) {
+        r = r + y;
+      }
+      return r;
+    }
+  }
+};
+
+struct Equal {
+  template <typename T>
+  __device__ bool operator()(T x, T y) {
+    return x == y;
+  }
+};
+
+struct NaNEqual {
+  template <typename T>
+  __device__ bool operator()(T x, T y) {
+    using cuda::std::isnan;
+    if constexpr (is_complex_v<T>) {
+      return x == y ||
+          (isnan(x.real()) && isnan(y.real()) && isnan(x.imag()) &&
+           isnan(y.imag())) ||
+          (x.real() == y.real() && isnan(x.imag()) && isnan(y.imag())) ||
+          (isnan(x.real()) && isnan(y.real()) && x.imag() == y.imag());
+    } else {
+      return x == y || (isnan(x) && isnan(y));
+    }
+  }
+};
+
+struct Greater {
+  template <typename T>
+  __device__ bool operator()(T x, T y) {
+    return x > y;
+  }
+};
+
+struct GreaterEqual {
+  template <typename T>
+  __device__ bool operator()(T x, T y) {
+    return x >= y;
+  }
+};
+
+struct Less {
+  template <typename T>
+  __device__ bool operator()(T x, T y) {
+    return x < y;
+  }
+};
+
+struct LessEqual {
+  template <typename T>
+  __device__ bool operator()(T x, T y) {
+    return x <= y;
+  }
+};
+
+struct LogAddExp {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    if constexpr (is_complex_v<T>) {
+      if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag()) ||
+          cuda::std::isnan(y.real()) || cuda::std::isnan(y.imag())) {
+        return {
+            cuda::std::numeric_limits<float>::quiet_NaN(),
+            cuda::std::numeric_limits<float>::quiet_NaN()};
+      }
+      auto max = x.real() > y.real() ? x : y;
+      auto min = x.real() < y.real() ? x : y;
+      auto min_real = min.real();
+      auto max_real = max.real();
+      if (!cuda::std::isfinite(min_real) && (min_real == max_real)) {
+        if (min_real < 0) {
+          return min;
+        } else {
+          return Log{}(Exp{}(min) + Exp{}(max));
+        }
+      } else {
+        return Log1p{}(Exp{}(min - max)) + max;
+      }
+    } else {
+      if (cuda::std::isnan(x) || cuda::std::isnan(y)) {
+        return cuda::std::numeric_limits<T>::quiet_NaN();
+      }
+      T maxval = max(x, y);
+      T minval = min(x, y);
+      return (minval == -cuda::std::numeric_limits<T>::infinity() ||
+              maxval == cuda::std::numeric_limits<T>::infinity())
+          ? maxval
+          : T(maxval + cuda::std::log1p(cuda::std::exp(minval - maxval)));
+    }
+  };
+};
+
+struct Maximum {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    if constexpr (cuda::std::is_integral_v<T>) {
+      return max(x, y);
+    } else if constexpr (is_complex_v<T>) {
+      if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag())) {
+        return x;
+      }
+      return x > y ? x : y;
+    } else {
+      if (cuda::std::isnan(x)) {
+        return x;
+      }
+      return x > y ? x : y;
+    }
+  }
+};
+
+struct Minimum {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    if constexpr (cuda::std::is_integral_v<T>) {
+      return min(x, y);
+    } else if constexpr (is_complex_v<T>) {
+      if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag())) {
+        return x;
+      }
+      return x < y ? x : y;
+    } else {
+      if (cuda::std::isnan(x)) {
+        return x;
+      }
+      return x < y ? x : y;
+    }
+  }
+};
+
+struct Multiply {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x * y;
+  }
+};
+
+struct NotEqual {
+  template <typename T>
+  __device__ bool operator()(T x, T y) {
+    if constexpr (is_complex_v<T>) {
+      return x.real() != y.real() || x.imag() != y.imag();
+    } else {
+      return x != y;
+    }
+  }
+};
+
+struct Power {
+  template <typename T>
+  __device__ T operator()(T base, T exp) {
+    if constexpr (cuda::std::is_integral_v<T>) {
+      T res = 1;
+      // Raising an integer to a negative power is undefined
+      if constexpr (cuda::std::is_signed_v<T>) {
+        if (exp < 0) {
+          return 0;
+        }
+      }
+      while (exp) {
+        if (exp & 1) {
+          res *= base;
+        }
+        exp >>= 1;
+        base *= base;
+      }
+      return res;
+    } else if constexpr (is_complex_v<T>) {
+      return cuda::std::pow(base, exp);
+    } else {
+      return cuda::std::pow(base, exp);
+    }
+  }
+};
+
+struct Subtract {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x - y;
+  }
+};
+
+struct LogicalAnd {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x && y;
+  };
+};
+
+struct LogicalOr {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x || y;
+  };
+};
+
+struct BitwiseAnd {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x & y;
+  };
+};
+
+struct BitwiseOr {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x | y;
+  };
+};
+
+struct BitwiseXor {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x ^ y;
+  };
+};
+
+struct LeftShift {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x << y;
+  };
+};
+
+struct RightShift {
+  template <typename T>
+  __device__ T operator()(T x, T y) {
+    return x >> y;
+  };
+};
+
+struct ArcTan2 {
+  template <typename T>
+  __device__ T operator()(T y, T x) {
+    return cuda::std::atan2(y, x);
+  }
+};
+
+struct DivMod {
+  template <typename T>
+  __device__ cuda::std::array<T, 2> operator()(T x, T y) {
+    return {FloorDivide{}(x, y), Remainder{}(x, y)};
+  };
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/cast_op.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/cast_op.cuh
new file mode 100644
index 0000000..b85a40c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/cast_op.cuh
@@ -0,0 +1,118 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/complex.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+namespace mlx::core::cu {
+
+// An op that does static_cast, with custom conversions for some types.
+template <typename SrcT, typename DstT, typename = void>
+struct CastOp {
+  static constexpr bool is_castable = cuda::std::is_convertible_v<SrcT, DstT>;
+
+  __device__ DstT operator()(SrcT x) {
+    return static_cast<DstT>(x);
+  }
+};
+
+// Castings between complex and boolean.
+template <typename T>
+struct CastOp<complex_t<T>, bool> {
+  static constexpr bool is_castable = true;
+
+  __device__ bool operator()(complex_t<T> x) {
+    return x.real() != 0 && x.imag() != 0;
+  }
+};
+
+template <typename T>
+struct CastOp<bool, complex_t<T>> {
+  static constexpr bool is_castable = true;
+
+  __device__ complex_t<T> operator()(bool x) {
+    return x ? complex_t<T>{1, 1} : complex_t<T>{0, 0};
+  }
+};
+
+// Converting a complex number to real number discards the imaginary part.
+template <typename T, typename DstT>
+struct CastOp<complex_t<T>, DstT, cuda::std::enable_if_t<!is_complex_v<DstT>>> {
+  static constexpr bool is_castable = cuda::std::is_convertible_v<T, DstT>;
+
+  __device__ DstT operator()(complex_t<T> x) {
+    static_assert(!is_complex_v<DstT>);
+    return static_cast<DstT>(x.real());
+  }
+};
+
+// Allow converting a real number to complex number.
+template <typename SrcT, typename T>
+struct CastOp<SrcT, complex_t<T>, cuda::std::enable_if_t<!is_complex_v<SrcT>>> {
+  static constexpr bool is_castable = cuda::std::is_convertible_v<SrcT, T>;
+
+  __device__ complex_t<T> operator()(SrcT x) {
+    static_assert(!is_complex_v<SrcT>);
+    return complex_t<T>{static_cast<T>(x), 0};
+  }
+};
+
+// Do nothing when no casting is needed.
+template <typename SrcT, typename DstT>
+struct CastOp<
+    SrcT,
+    DstT,
+    cuda::std::enable_if_t<cuda::std::is_same_v<SrcT, DstT>>> {
+  static constexpr bool is_castable = true;
+
+  __device__ SrcT operator()(SrcT x) {
+    return x;
+  }
+};
+
+// In CUDA 11 the half types do not define conversions between some types,
+// provide fallbacks here.
+#if CUDART_VERSION < 12000
+template <typename SrcT, typename DstT>
+struct CastOp<
+    SrcT,
+    DstT,
+    cuda::std::enable_if_t<
+        !cuda::std::is_convertible_v<SrcT, DstT> && !is_complex_v<SrcT> &&
+        (cuda::std::is_same_v<DstT, __half> ||
+         cuda::std::is_same_v<DstT, __nv_bfloat16>)>> {
+  static constexpr bool is_castable = true;
+
+  __device__ DstT operator()(SrcT x) {
+    return DstT(static_cast<float>(x));
+  }
+};
+
+template <typename SrcT, typename DstT>
+struct CastOp<
+    SrcT,
+    DstT,
+    cuda::std::enable_if_t<
+        !cuda::std::is_convertible_v<SrcT, DstT> && !is_complex_v<SrcT> &&
+        !cuda::std::is_same_v<DstT, __half> &&
+        !cuda::std::is_same_v<DstT, __nv_bfloat16> &&
+        (cuda::std::is_same_v<SrcT, __half> ||
+         cuda::std::is_same_v<SrcT, __nv_bfloat16>)>> {
+  static constexpr bool is_castable = true;
+
+  __device__ DstT operator()(SrcT x) {
+    return DstT(static_cast<float>(x));
+  }
+};
+#endif // CUDART_VERSION < 12000
+
+// Helper to deduce the SrcT.
+template <typename DstT, typename SrcT>
+inline __host__ __device__ auto cast_to(SrcT x) {
+  return CastOp<SrcT, DstT>{}(x);
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/complex.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/complex.cuh
new file mode 100644
index 0000000..03a7bff
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/complex.cuh
@@ -0,0 +1,60 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+// Make multiplication and division faster.
+#define LIBCUDACXX_ENABLE_SIMPLIFIED_COMPLEX_OPERATIONS
+
+#include <cuda/std/complex>
+#include <cuda/std/type_traits>
+
+namespace mlx::core::cu {
+
+// TODO: Consider using a faster implementation as cuda::std::complex has to
+// conform to C++ standard.
+template <typename T>
+using complex_t = cuda::std::complex<T>;
+
+using complex64_t = complex_t<float>;
+using complex128_t = complex_t<double>;
+
+template <typename T>
+struct is_complex : cuda::std::false_type {};
+
+template <typename T>
+struct is_complex<cuda::std::complex<T>> : cuda::std::true_type {};
+
+template <typename T>
+inline constexpr bool is_complex_v = is_complex<T>::value;
+
+// cuda::std::complex is missing some operators.
+template <typename T>
+inline __host__ __device__ complex_t<T> operator%(
+    complex_t<T> a,
+    complex_t<T> b) {
+  T r = a.real() - floor(a.real() / b.real()) * b.real();
+  T i = a.imag() - floor(a.imag() / b.imag()) * b.imag();
+  return complex_t<T>{r, i};
+}
+
+template <typename T>
+inline __host__ __device__ bool operator>(complex_t<T> a, complex_t<T> b) {
+  return (a.real() > b.real()) || (a.real() == b.real() && a.imag() > b.imag());
+}
+
+template <typename T>
+inline __host__ __device__ bool operator<(complex_t<T> a, complex_t<T> b) {
+  return operator>(b, a);
+}
+
+template <typename T>
+inline __host__ __device__ bool operator<=(complex_t<T> a, complex_t<T> b) {
+  return !(a > b);
+}
+
+template <typename T>
+inline __host__ __device__ bool operator>=(complex_t<T> a, complex_t<T> b) {
+  return !(a < b);
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/config.h b/macos/third_party/mlx/mlx/backend/cuda/device/config.h
new file mode 100644
index 0000000..5a34029
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/config.h
@@ -0,0 +1,12 @@
+// Copyright © 2025 Apple Inc.
+
+// This file is used by both CUDA kernel code and host-only C++ code.
+
+#pragma once
+
+// The maximum dimensions of shape/strides passed as kernel parameters.
+#define MAX_NDIM 10
+
+// All existing NVIDIA hardware has a fixed 32 warp size. Though a built-in
+// warpSize variable exists, using it would prevent compile-time optimizations.
+#define WARP_SIZE 32
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/fp16_math.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/fp16_math.cuh
new file mode 100644
index 0000000..2799859
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/fp16_math.cuh
@@ -0,0 +1,96 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda/std/type_traits>
+
+namespace mlx::core::cu {
+
+///////////////////////////////////////////////////////////////////////////////
+// Binary ops for half types.
+///////////////////////////////////////////////////////////////////////////////
+
+#define MLX_DEFINE_BINARY_OP(NAME, HALF_OP)                        \
+  template <typename T>                                            \
+  __forceinline__ __device__ auto NAME(T x, T y) {                 \
+    if constexpr (cuda::std::is_same_v<T, __half>) {               \
+      return HALF_OP(x, y);                                        \
+    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) { \
+      return HALF_OP(x, y);                                        \
+    } else {                                                       \
+      return ::NAME(x, y);                                         \
+    }                                                              \
+  }
+
+MLX_DEFINE_BINARY_OP(max, __hmax)
+MLX_DEFINE_BINARY_OP(min, __hmin)
+
+#undef MLX_DEFINE_BINARY_OP
+
+///////////////////////////////////////////////////////////////////////////////
+// Additional C++ operator overrides between half types and native types.
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, typename U>
+constexpr bool is_integral_except =
+    cuda::std::is_integral_v<T> && !cuda::std::is_same_v<T, U>;
+
+template <typename T, typename U>
+constexpr bool is_arithmetic_except =
+    cuda::std::is_arithmetic_v<T> && !cuda::std::is_same_v<T, U>;
+
+#define MLX_DEFINE_HALF_OP(HALF, HALF2FLOAT, FLOAT2HALF, OP)          \
+  template <                                                          \
+      typename T,                                                     \
+      typename = cuda::std::enable_if_t<is_integral_except<T, HALF>>> \
+  __forceinline__ __device__ HALF operator OP(HALF x, T y) {          \
+    return FLOAT2HALF(HALF2FLOAT(x) OP static_cast<float>(y));        \
+  }                                                                   \
+  template <                                                          \
+      typename T,                                                     \
+      typename = cuda::std::enable_if_t<is_integral_except<T, HALF>>> \
+  __forceinline__ __device__ HALF operator OP(T x, HALF y) {          \
+    return FLOAT2HALF(static_cast<float>(x) OP HALF2FLOAT(y));        \
+  }
+
+#define MLX_DEFINE_HALF_CMP(HALF, HALF2FLOAT, OP)                       \
+  template <                                                            \
+      typename T,                                                       \
+      typename = cuda::std::enable_if_t<is_arithmetic_except<T, HALF>>> \
+  __forceinline__ __device__ bool operator OP(HALF x, T y) {            \
+    return HALF2FLOAT(x) OP static_cast<float>(y);                      \
+  }                                                                     \
+  template <                                                            \
+      typename T,                                                       \
+      typename = cuda::std::enable_if_t<is_arithmetic_except<T, HALF>>> \
+  __forceinline__ __device__ bool operator OP(T x, HALF y) {            \
+    return static_cast<float>(y) OP HALF2FLOAT(x);                      \
+  }
+
+MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, +)
+MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, -)
+MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, *)
+MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, /)
+MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, +)
+MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, -)
+MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, *)
+MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, /)
+MLX_DEFINE_HALF_CMP(__half, __half2float, <)
+MLX_DEFINE_HALF_CMP(__half, __half2float, >)
+MLX_DEFINE_HALF_CMP(__half, __half2float, <=)
+MLX_DEFINE_HALF_CMP(__half, __half2float, >=)
+MLX_DEFINE_HALF_CMP(__half, __half2float, ==)
+MLX_DEFINE_HALF_CMP(__half, __half2float, !=)
+MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, <)
+MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, >)
+MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, <=)
+MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, >=)
+MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, ==)
+MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, !=)
+
+#undef MLX_DEFINE_HALF_OP
+#undef MLX_DEFINE_HALF_CMP
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/gather.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/gather.cuh
new file mode 100644
index 0000000..7dbd84a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/gather.cuh
@@ -0,0 +1,53 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/indexing.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, typename IdxT, int NIDX, int IDX_NDIM, typename LocT>
+__global__ void gather(
+    const T* src,
+    T* out,
+    LocT size,
+    const __grid_constant__ Shape src_shape,
+    const __grid_constant__ Strides src_strides,
+    int32_t src_ndim,
+    const __grid_constant__ Shape slice_sizes,
+    uint32_t slice_size,
+    const __grid_constant__ cuda::std::array<int32_t, NIDX> axes,
+    const __grid_constant__ cuda::std::array<IdxT*, NIDX> indices,
+    const __grid_constant__ cuda::std::array<int32_t, NIDX * IDX_NDIM>
+        indices_shape,
+    const __grid_constant__ cuda::std::array<int64_t, NIDX * IDX_NDIM>
+        indices_strides) {
+  LocT out_idx = cg::this_grid().thread_rank();
+  if (out_idx >= size) {
+    return;
+  }
+
+  LocT src_elem = out_idx % slice_size;
+  LocT idx_elem = out_idx / slice_size;
+
+  LocT src_loc =
+      elem_to_loc(src_elem, slice_sizes.data(), src_strides.data(), src_ndim);
+
+#pragma unroll
+  for (int i = 0; i < NIDX; ++i) {
+    LocT idx_loc = elem_to_loc_nd<IDX_NDIM>(
+        idx_elem,
+        indices_shape.data() + i * IDX_NDIM,
+        indices_strides.data() + i * IDX_NDIM);
+    int32_t axis = axes[i];
+    LocT idx_val = absolute_index(indices[i][idx_loc], src_shape[axis]);
+    src_loc += idx_val * src_strides[axis];
+  }
+
+  out[out_idx] = src[src_loc];
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/gather_axis.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/gather_axis.cuh
new file mode 100644
index 0000000..f863b2d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/gather_axis.cuh
@@ -0,0 +1,65 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/indexing.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+template <
+    typename T,
+    typename IdxT,
+    int NDIM,
+    bool SrcC,
+    bool IdxC,
+    typename LocT>
+__global__ void gather_axis(
+    const T* src,
+    const IdxT* indices,
+    T* out,
+    LocT idx_size_pre,
+    LocT idx_size_axis,
+    LocT idx_size_post,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> src_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> idx_strides,
+    int32_t axis,
+    int32_t axis_size,
+    int64_t src_stride_axis,
+    int64_t idx_stride_axis) {
+  LocT index = cg::this_grid().thread_rank();
+  if (index >= idx_size_pre * idx_size_axis * idx_size_post) {
+    return;
+  }
+
+  auto [x, y, z] = index_to_dims(index, idx_size_axis, idx_size_pre);
+
+  LocT elem_idx = z * idx_size_post;
+
+  LocT idx_loc = y * idx_stride_axis;
+  if constexpr (IdxC) {
+    idx_loc += elem_idx * idx_size_axis + x;
+  } else {
+    idx_loc +=
+        elem_to_loc_nd<NDIM>(elem_idx + x, shape.data(), idx_strides.data());
+  }
+
+  auto idx_val = absolute_index(indices[idx_loc], axis_size);
+
+  LocT src_loc = idx_val * src_stride_axis;
+  if constexpr (SrcC) {
+    src_loc += elem_idx * axis_size + x;
+  } else {
+    src_loc +=
+        elem_to_loc_nd<NDIM>(elem_idx + x, shape.data(), src_strides.data());
+  }
+
+  LocT out_idx = y * idx_size_post + elem_idx * idx_size_axis + x;
+
+  out[out_idx] = src[src_loc];
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/indexing.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/indexing.cuh
new file mode 100644
index 0000000..31cba1a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/indexing.cuh
@@ -0,0 +1,30 @@
+// Copyright © 2025 Apple Inc.
+
+#include <cuda/std/tuple>
+#include <cuda/std/type_traits>
+
+namespace mlx::core::cu {
+
+// Convert an absolute index to positions in a 3d grid, assuming the index is
+// calculated with:
+// index = x * dim1 * dim2 + y * dim2 + z
+template <typename T>
+inline __host__ __device__ cuda::std::tuple<T, T, T>
+index_to_dims(T index, T dim1, T dim2) {
+  T x = index / (dim1 * dim2);
+  T y = (index % (dim1 * dim2)) / dim2;
+  T z = index % dim2;
+  return cuda::std::make_tuple(x, y, z);
+}
+
+// Get absolute index from possible negative index.
+template <typename IdxT>
+inline __host__ __device__ auto absolute_index(IdxT idx, int32_t size) {
+  if constexpr (cuda::std::is_unsigned_v<IdxT>) {
+    return idx;
+  } else {
+    return static_cast<int32_t>(idx < 0 ? idx + size : idx);
+  }
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/scatter.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/scatter.cuh
new file mode 100644
index 0000000..b2f6403
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/scatter.cuh
@@ -0,0 +1,68 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/indexing.cuh"
+#include "mlx/backend/cuda/device/scatter_ops.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+template <
+    typename T,
+    typename IdxT,
+    typename Op,
+    int NIDX,
+    int IDX_NDIM,
+    typename LocT>
+__global__ void scatter(
+    const T* upd,
+    T* out,
+    LocT size,
+    const __grid_constant__ Shape upd_shape,
+    const __grid_constant__ Strides upd_strides,
+    int32_t upd_ndim,
+    LocT upd_post_idx_size,
+    const __grid_constant__ Shape out_shape,
+    const __grid_constant__ Strides out_strides,
+    int32_t out_ndim,
+    const __grid_constant__ cuda::std::array<int32_t, NIDX> axes,
+    const __grid_constant__ cuda::std::array<IdxT*, NIDX> indices,
+    const __grid_constant__ cuda::std::array<int32_t, NIDX * IDX_NDIM>
+        indices_shape,
+    const __grid_constant__ cuda::std::array<int64_t, NIDX * IDX_NDIM>
+        indices_strides) {
+  LocT upd_idx = cg::this_grid().thread_rank();
+  if (upd_idx >= size) {
+    return;
+  }
+
+  LocT out_elem = upd_idx % upd_post_idx_size;
+  LocT idx_elem = upd_idx / upd_post_idx_size;
+
+  LocT out_idx = elem_to_loc(
+      out_elem, upd_shape.data() + IDX_NDIM, out_strides.data(), out_ndim);
+
+#pragma unroll
+  for (int i = 0; i < NIDX; ++i) {
+    LocT idx_loc = elem_to_loc_nd<IDX_NDIM>(
+        idx_elem,
+        indices_shape.data() + i * IDX_NDIM,
+        indices_strides.data() + i * IDX_NDIM);
+    int32_t axis = axes[i];
+    LocT idx_val = absolute_index(indices[i][idx_loc], out_shape[axis]);
+    out_idx += idx_val * out_strides[axis];
+  }
+
+  LocT upd_loc = elem_to_loc(
+      out_elem + idx_elem * upd_post_idx_size,
+      upd_shape.data(),
+      upd_strides.data(),
+      upd_ndim);
+
+  Op{}(out + out_idx, upd[upd_loc]);
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/scatter_axis.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/scatter_axis.cuh
new file mode 100644
index 0000000..1f30f2e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/scatter_axis.cuh
@@ -0,0 +1,67 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/indexing.cuh"
+#include "mlx/backend/cuda/device/scatter_ops.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+template <
+    typename T,
+    typename IdxT,
+    typename Op,
+    int NDIM,
+    bool UpdC,
+    bool IdxC,
+    typename LocT>
+__global__ void scatter_axis(
+    const T* upd,
+    const IdxT* indices,
+    T* out,
+    LocT idx_size_pre,
+    LocT idx_size_axis,
+    LocT idx_size_post,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> upd_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> idx_strides,
+    int32_t axis,
+    int32_t axis_size,
+    int64_t upd_stride_axis,
+    int64_t idx_stride_axis) {
+  LocT index = cg::this_grid().thread_rank();
+  if (index >= idx_size_pre * idx_size_axis * idx_size_post) {
+    return;
+  }
+
+  auto [x, y, z] = index_to_dims(index, idx_size_axis, idx_size_pre);
+
+  LocT elem_idx = z * idx_size_post;
+
+  LocT idx_loc = y * idx_stride_axis;
+  if constexpr (IdxC) {
+    idx_loc += elem_idx * idx_size_axis + x;
+  } else {
+    idx_loc +=
+        elem_to_loc_nd<NDIM>(elem_idx + x, shape.data(), idx_strides.data());
+  }
+
+  auto idx_val = absolute_index(indices[idx_loc], axis_size);
+
+  LocT upd_loc = y * upd_stride_axis;
+  if constexpr (UpdC) {
+    upd_loc += elem_idx * idx_size_axis + x;
+  } else {
+    upd_loc +=
+        elem_to_loc_nd<NDIM>(elem_idx + x, shape.data(), upd_strides.data());
+  }
+
+  LocT out_idx = idx_val * idx_size_post + elem_idx * axis_size + x;
+
+  Op{}(out + out_idx, upd[upd_loc]);
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/scatter_ops.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/scatter_ops.cuh
new file mode 100644
index 0000000..d88f896
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/scatter_ops.cuh
@@ -0,0 +1,44 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/atomic_ops.cuh"
+
+namespace mlx::core::cu {
+
+struct ScatterAssign {
+  template <typename T>
+  __device__ void operator()(T* out, T val) const {
+    *out = val;
+  }
+};
+
+struct ScatterSum {
+  template <typename T>
+  __device__ void operator()(T* out, T val) const {
+    atomic_add(out, val);
+  }
+};
+
+struct ScatterProd {
+  template <typename T>
+  __device__ void operator()(T* out, T val) const {
+    atomic_prod(out, val);
+  }
+};
+
+struct ScatterMax {
+  template <typename T>
+  __device__ void operator()(T* out, T val) const {
+    atomic_max(out, val);
+  }
+};
+
+struct ScatterMin {
+  template <typename T>
+  __device__ void operator()(T* out, T val) const {
+    atomic_min(out, val);
+  }
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/ternary_ops.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/ternary_ops.cuh
new file mode 100644
index 0000000..4418454
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/ternary_ops.cuh
@@ -0,0 +1,13 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+namespace mlx::core::cu {
+
+struct Select {
+  template <typename T>
+  __device__ T operator()(bool condition, T x, T y) {
+    return condition ? x : y;
+  }
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/unary_ops.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/unary_ops.cuh
new file mode 100644
index 0000000..8234cf8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/unary_ops.cuh
@@ -0,0 +1,350 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/fp16_math.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cuda_fp8.h>
+#include <math_constants.h>
+#include <cuda/std/cmath>
+
+namespace mlx::core::cu {
+
+struct Abs {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (cuda::std::is_unsigned_v<T>) {
+      return x;
+    } else {
+      return cuda::std::abs(x);
+    }
+  }
+};
+
+struct ArcCos {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::acos(x);
+  }
+};
+
+struct ArcCosh {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::acosh(x);
+  }
+};
+
+struct ArcSin {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::asin(x);
+  }
+};
+
+struct ArcSinh {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::asinh(x);
+  }
+};
+
+struct ArcTan {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::atan(x);
+  }
+};
+
+struct ArcTanh {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::atanh(x);
+  }
+};
+
+struct BitwiseInvert {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return ~x;
+  }
+};
+
+struct Ceil {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (cuda::std::is_integral_v<T>) {
+      return x;
+    } else if constexpr (is_complex_v<T>) {
+      return T{cuda::std::ceil(x.real()), cuda::std::ceil(x.imag())};
+    } else {
+      return cuda::std::ceil(x);
+    }
+  }
+};
+
+struct Conjugate {
+  template <typename T>
+  __device__ complex_t<T> operator()(complex_t<T> x) {
+    return cuda::std::conj(x);
+  }
+};
+
+struct Cos {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::cos(x);
+  }
+};
+
+struct Cosh {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::cosh(x);
+  }
+};
+
+struct Erf {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (cuda::std::is_same_v<T, __half>) {
+      return erf(__half2float(x));
+    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
+      return erf(__bfloat162float(x));
+    } else {
+      return erf(x);
+    }
+  }
+};
+
+struct ErfInv {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (cuda::std::is_same_v<T, __half>) {
+      return erfinv(__half2float(x));
+    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
+      return erfinv(__bfloat162float(x));
+    } else {
+      return erfinv(x);
+    }
+  }
+};
+
+struct Exp {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::exp(x);
+  }
+};
+
+struct Expm1 {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::expm1(x);
+  }
+};
+
+struct Floor {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (cuda::std::is_integral_v<T>) {
+      return x;
+    } else if constexpr (is_complex_v<T>) {
+      return T{cuda::std::floor(x.real()), cuda::std::floor(x.imag())};
+    } else {
+      return cuda::std::floor(x);
+    }
+  }
+};
+
+struct Imag {
+  template <typename T>
+  __device__ auto operator()(complex_t<T> x) {
+    return x.imag();
+  }
+};
+
+struct Log {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::log(x);
+  }
+};
+
+struct Log2 {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (is_complex_v<T>) {
+      auto y = Log{}(x);
+      return {y.real() / CUDART_LN2_F, y.imag() / CUDART_LN2_F};
+    } else {
+      return cuda::std::log2(x);
+    }
+  }
+};
+
+struct Log10 {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::log10(x);
+  }
+};
+
+struct Log1p {
+  template <typename T>
+  __device__ T operator()(T z) {
+    if constexpr (is_complex_v<T>) {
+      float x = z.real();
+      float y = z.imag();
+      float zabs = Abs{}(z).real();
+      float theta = atan2f(y, x + 1);
+      if (zabs < 0.5f) {
+        float r = x * (2 + x) + y * y;
+        if (r == 0) { // handle underflow
+          return {x, theta};
+        }
+        return {0.5f * log1pf(r), theta};
+      } else {
+        float z0 = hypotf(x + 1, y);
+        return {logf(z0), theta};
+      }
+    } else {
+      return cuda::std::log1p(z);
+    }
+  }
+};
+
+struct LogicalNot {
+  __device__ bool operator()(bool x) {
+    return !x;
+  }
+};
+
+struct Negative {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (is_complex_v<T>) {
+      return T{0, 0} - x;
+    } else {
+      return -x;
+    }
+  }
+};
+
+struct Real {
+  template <typename T>
+  __device__ auto operator()(complex_t<T> x) {
+    return x.real();
+  }
+};
+
+struct Round {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (is_complex_v<T>) {
+      return {cuda::std::rint(x.real()), cuda::std::rint(x.imag())};
+    } else {
+      return cuda::std::rint(x);
+    }
+  }
+};
+
+struct Sigmoid {
+  template <typename T>
+  __device__ T operator()(T x) {
+    T y = 1 / (1 + cuda::std::exp(cuda::std::abs(x)));
+    return (x < 0) ? y : 1 - y;
+  }
+};
+
+struct Sign {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (cuda::std::is_unsigned_v<T>) {
+      return x != 0;
+    } else if constexpr (is_complex_v<T>) {
+      if (x.real() == 0 && x.imag() == 0) {
+        return x;
+      } else {
+        return x / Abs()(x);
+      }
+    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
+      return static_cast<float>((x > T(0.f)) - (x < T(0.f)));
+    } else {
+      return (x > T(0)) - (x < T(0));
+    }
+  }
+};
+
+struct Sin {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::sin(x);
+  }
+};
+
+struct Sinh {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::sinh(x);
+  }
+};
+
+struct Square {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return x * x;
+  }
+};
+
+struct Sqrt {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::sqrt(x);
+  }
+};
+
+struct Rsqrt {
+  template <typename T>
+  __device__ T operator()(T x) {
+    if constexpr (is_complex_v<T>) {
+      return 1.0f / Sqrt{}(x);
+    } else if constexpr (cuda::std::is_same_v<T, __half>) {
+      return rsqrt(__half2float(x));
+    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
+      return rsqrt(__bfloat162float(x));
+    } else {
+      return rsqrt(x);
+    }
+  }
+};
+
+struct Tan {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::tan(x);
+  }
+};
+
+struct Tanh {
+  template <typename T>
+  __device__ T operator()(T x) {
+    return cuda::std::tanh(x);
+  }
+};
+
+struct ToFP8 {
+  template <typename T>
+  __device__ uint8_t operator()(T x) {
+    return __nv_fp8_e4m3(x).__x;
+  }
+};
+
+struct FromFP8 {
+  __device__ float operator()(uint8_t x) {
+    return float(*(__nv_fp8_e4m3*)(&x));
+  }
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device/utils.cuh b/macos/third_party/mlx/mlx/backend/cuda/device/utils.cuh
new file mode 100644
index 0000000..2e017ba
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device/utils.cuh
@@ -0,0 +1,464 @@
+// Copyright © 2025 Apple Inc.
+
+// This file must not include any host-only code, utilities that work under both
+// host and device can be put here.
+//
+// See more about the requirements at:
+// https://docs.nvidia.com/cuda/nvrtc/#language
+
+#pragma once
+
+#include "mlx/backend/cuda/device/complex.cuh"
+#include "mlx/backend/cuda/device/config.h"
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda/std/array>
+#include <cuda/std/limits>
+#include <cuda/std/tuple>
+
+namespace mlx::core::cu {
+
+///////////////////////////////////////////////////////////////////////////////
+// CUDA kernel utils
+///////////////////////////////////////////////////////////////////////////////
+
+// To pass shape/strides to kernels via constant memory, their size must be
+// known at compile time.
+using Shape = cuda::std::array<int32_t, MAX_NDIM>;
+using Strides = cuda::std::array<int64_t, MAX_NDIM>;
+
+// Vectorized load/store.
+template <typename T, int N>
+struct alignas(sizeof(T) * N) AlignedVector {
+  T val[N];
+
+  __device__ T& operator[](int i) {
+    return val[i];
+  }
+
+  __device__ T operator[](int i) const {
+    return val[i];
+  }
+};
+
+template <int N, typename T>
+inline __host__ __device__ bool is_aligned(T* x) {
+  return (reinterpret_cast<uintptr_t>(x) % (N * sizeof(T))) == 0;
+}
+
+template <int N, typename T>
+inline __device__ AlignedVector<T, N> unsafe_load_vector(
+    const T* ptr,
+    uint32_t offset) {
+  auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
+  return from[offset];
+}
+
+template <int N, typename T>
+inline __device__ AlignedVector<T, N> load_vector(
+    const T* ptr,
+    uint32_t offset) {
+  if (is_aligned<N>(ptr)) {
+    auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
+    return from[offset];
+  } else {
+    AlignedVector<T, N> v;
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      v[i] = ptr[offset * N + i];
+    }
+    return v;
+  }
+}
+
+template <int N, typename T, typename SizeT>
+inline __device__ AlignedVector<T, N>
+load_vector(const T* ptr, uint32_t offset, SizeT size, T fallback) {
+  if (is_aligned<N>(ptr) && (offset + 1) * N <= size) {
+    auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
+    return from[offset];
+  } else {
+    AlignedVector<T, N> v;
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      v[i] = (N * offset + i) < size ? ptr[offset * N + i] : fallback;
+    }
+    return v;
+  }
+}
+
+template <int N, typename T, typename SizeT>
+inline __device__ AlignedVector<T, N> load_vector(
+    const T* ptr,
+    uint32_t offset,
+    SizeT size,
+    int64_t stride,
+    T fallback) {
+  if (is_aligned<N>(ptr) && stride == 1 && (offset + 1) * N <= size) {
+    auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
+    return from[offset];
+  } else {
+    AlignedVector<T, N> v;
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      v[i] =
+          (N * offset + i) < size ? ptr[stride * (offset * N + i)] : fallback;
+    }
+    return v;
+  }
+}
+
+template <int N, typename T>
+inline __device__ void
+unsafe_store_vector(T* ptr, uint32_t offset, const AlignedVector<T, N>& vec) {
+  auto* to = reinterpret_cast<AlignedVector<T, N>*>(ptr);
+  to[offset] = vec;
+}
+
+template <int N, typename T>
+inline __device__ void
+store_vector(T* ptr, uint32_t offset, const AlignedVector<T, N>& vec) {
+  if (is_aligned<N>(ptr)) {
+    auto* to = reinterpret_cast<AlignedVector<T, N>*>(ptr);
+    to[offset] = vec;
+  } else {
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      ptr[offset * N + i] = vec[i];
+    }
+  }
+}
+
+template <int N, typename T, typename SizeT>
+inline __device__ void store_vector(
+    T* ptr,
+    uint32_t offset,
+    const AlignedVector<T, N>& vec,
+    SizeT size) {
+  if (is_aligned<N>(ptr) && (offset + 1) * N <= size) {
+    auto* to = reinterpret_cast<AlignedVector<T, N>*>(ptr);
+    to[offset] = vec;
+  } else {
+    for (int i = 0; (offset * N + i) < size && i < N; ++i) {
+      ptr[offset * N + i] = vec[i];
+    }
+  }
+}
+
+template <int N, typename T, typename SizeT>
+inline __device__ void store_vector(
+    T* ptr,
+    uint32_t offset,
+    const AlignedVector<T, N>& vec,
+    SizeT size,
+    int64_t stride) {
+  if (is_aligned<N>(ptr) && (offset + 1) * N <= size && stride == 1) {
+    auto* to = reinterpret_cast<AlignedVector<T, N>*>(ptr);
+    to[offset] = vec;
+  } else {
+    for (int i = 0; (offset * N + i) < size && i < N; ++i) {
+      ptr[stride * (offset * N + i)] = vec[i];
+    }
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Type limits utils
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, typename = void>
+struct Limits {
+  static constexpr __host__ __device__ T max() {
+    return cuda::std::numeric_limits<T>::max();
+  }
+  static constexpr __host__ __device__ T min() {
+    return cuda::std::numeric_limits<T>::min();
+  }
+  static constexpr __host__ __device__ T finite_max() {
+    return cuda::std::numeric_limits<T>::max();
+  }
+  static constexpr __host__ __device__ T finite_min() {
+    return cuda::std::numeric_limits<T>::min();
+  }
+};
+
+template <typename T>
+struct Limits<
+    T,
+    cuda::std::enable_if_t<
+        cuda::std::is_same_v<T, float> || cuda::std::is_same_v<T, double>>> {
+  static constexpr __host__ __device__ T max() {
+    return cuda::std::numeric_limits<T>::infinity();
+  }
+  static constexpr __host__ __device__ T min() {
+    return -cuda::std::numeric_limits<T>::infinity();
+  }
+  static constexpr __host__ __device__ T finite_max() {
+    return cuda::std::numeric_limits<T>::max();
+  }
+  static constexpr __host__ __device__ T finite_min() {
+    return cuda::std::numeric_limits<T>::lowest();
+  }
+};
+
+// CUDA 11 does not have host side arithmetic operators for half types.
+template <typename T>
+struct Limits<
+    T,
+    cuda::std::enable_if_t<
+        cuda::std::is_same_v<T, __half> ||
+        cuda::std::is_same_v<T, __nv_bfloat16>>> {
+  static constexpr __host__ __device__ T max() {
+    return cuda::std::numeric_limits<T>::infinity();
+  }
+  static constexpr __host__ __device__ T min() {
+#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
+    return -cuda::std::numeric_limits<float>::infinity();
+#else
+    return -cuda::std::numeric_limits<T>::infinity();
+#endif
+  }
+  static constexpr __host__ __device__ T finite_max() {
+    return cuda::std::numeric_limits<T>::max();
+  }
+  static constexpr __host__ __device__ T finite_min() {
+#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
+    return cuda::std::numeric_limits<float>::lowest();
+#else
+    return cuda::std::numeric_limits<T>::lowest();
+#endif
+  }
+};
+
+template <>
+struct Limits<bool> {
+  static constexpr __host__ __device__ bool max() {
+    return true;
+  }
+  static constexpr __host__ __device__ bool min() {
+    return false;
+  }
+};
+
+template <typename T>
+struct Limits<complex_t<T>> {
+  static constexpr __host__ __device__ complex_t<T> max() {
+    return {Limits<T>::max(), Limits<T>::max()};
+  }
+  static constexpr __host__ __device__ complex_t<T> min() {
+    return {Limits<T>::min(), Limits<T>::min()};
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Indexing utils
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename IdxT = int64_t>
+inline __host__ __device__ IdxT
+elem_to_loc(IdxT elem, const int* shape, const int64_t* strides, int ndim) {
+  IdxT loc = 0;
+  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
+    loc += (elem % shape[i]) * IdxT(strides[i]);
+    elem /= shape[i];
+  }
+  return loc;
+}
+
+// Optimize when the ndim is known at compile time.
+template <int NDIM, typename IdxT = int64_t>
+inline __host__ __device__ IdxT
+elem_to_loc_nd(IdxT elem, const int* shape, const int64_t* strides) {
+  IdxT loc = 0;
+#pragma unroll
+  for (int i = NDIM - 1; i >= 0; --i) {
+    loc += (elem % shape[i]) * IdxT(strides[i]);
+    elem /= shape[i];
+  }
+  return loc;
+}
+
+template <int NDIM, typename IdxT = int64_t>
+inline __host__ __device__ cuda::std::tuple<IdxT, IdxT> elem_to_loc_nd(
+    IdxT elem,
+    const int* shape,
+    const int64_t* a_strides,
+    const int64_t* b_strides) {
+  IdxT a_loc = 0;
+  IdxT b_loc = 0;
+#pragma unroll
+  for (int i = NDIM - 1; i >= 0; --i) {
+    int dim_idx = elem % shape[i];
+    a_loc += dim_idx * IdxT(a_strides[i]);
+    b_loc += dim_idx * IdxT(b_strides[i]);
+    elem /= shape[i];
+  }
+  return cuda::std::make_tuple(a_loc, b_loc);
+}
+
+template <int NDIM, typename IdxT = int64_t>
+inline __host__ __device__ cuda::std::tuple<IdxT, IdxT, IdxT> elem_to_loc_nd(
+    IdxT elem,
+    const int* shape,
+    const int64_t* a_strides,
+    const int64_t* b_strides,
+    const int64_t* c_strides) {
+  IdxT a_loc = 0;
+  IdxT b_loc = 0;
+  IdxT c_loc = 0;
+#pragma unroll
+  for (int i = NDIM - 1; i >= 0; --i) {
+    int dim_idx = elem % shape[i];
+    a_loc += dim_idx * IdxT(a_strides[i]);
+    b_loc += dim_idx * IdxT(b_strides[i]);
+    c_loc += dim_idx * IdxT(c_strides[i]);
+    elem /= shape[i];
+  }
+  return cuda::std::make_tuple(a_loc, b_loc, c_loc);
+}
+
+template <typename IdxT = int64_t>
+inline __host__ __device__ cuda::std::tuple<IdxT, IdxT> elem_to_loc(
+    IdxT elem,
+    const int* shape,
+    const int64_t* a_strides,
+    const int64_t* b_strides,
+    int ndim) {
+  IdxT a_loc = 0;
+  IdxT b_loc = 0;
+  for (int i = ndim - 1; i >= 0; --i) {
+    int dim_idx = elem % shape[i];
+    a_loc += dim_idx * IdxT(a_strides[i]);
+    b_loc += dim_idx * IdxT(b_strides[i]);
+    elem /= shape[i];
+  }
+  return cuda::std::make_tuple(a_loc, b_loc);
+}
+
+template <typename IdxT = int64_t>
+inline __host__ __device__ cuda::std::tuple<IdxT, IdxT, IdxT> elem_to_loc(
+    IdxT elem,
+    const int* shape,
+    const int64_t* a_strides,
+    const int64_t* b_strides,
+    const int64_t* c_strides,
+    int ndim) {
+  IdxT a_loc = 0;
+  IdxT b_loc = 0;
+  IdxT c_loc = 0;
+  for (int i = ndim - 1; i >= 0; --i) {
+    int dim_idx = elem % shape[i];
+    a_loc += dim_idx * IdxT(a_strides[i]);
+    b_loc += dim_idx * IdxT(b_strides[i]);
+    c_loc += dim_idx * IdxT(c_strides[i]);
+    elem /= shape[i];
+  }
+  return cuda::std::make_tuple(a_loc, b_loc, c_loc);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Elem to loc in a loop utils
+///////////////////////////////////////////////////////////////////////////////
+
+template <int DIM, bool General = true, typename OffsetT = size_t>
+struct LoopedElemToLoc {
+  int dim;
+  LoopedElemToLoc<DIM - 1, General, OffsetT> inner_looper;
+  OffsetT offset{0};
+  int index{0};
+
+  __device__ LoopedElemToLoc(int dim) : dim(dim), inner_looper(dim - 1) {}
+
+  __device__ void next(const int* shape, const int64_t* strides) {
+    if (dim == 0) {
+      return;
+    }
+    index++;
+    offset += OffsetT(strides[dim - 1]);
+    if (index >= shape[dim - 1]) {
+      index = 0;
+      inner_looper.next(shape, strides);
+      offset = inner_looper.offset;
+    }
+  }
+
+  __device__ void next(int n, const int* shape, const int64_t* strides) {
+    if (dim == 0) {
+      return;
+    }
+    index += n;
+    offset += n * OffsetT(strides[dim - 1]);
+
+    if (index >= shape[dim - 1]) {
+      int extra = index - shape[dim - 1];
+      if (extra >= shape[dim - 1]) {
+        inner_looper.next(1 + extra / shape[dim - 1], shape, strides);
+        extra = extra % shape[dim - 1];
+      } else {
+        inner_looper.next(shape, strides);
+      }
+      index = 0;
+      offset = inner_looper.offset;
+      if (extra > 0) {
+        next(extra, shape, strides);
+      }
+    }
+  }
+
+  __device__ OffsetT location() {
+    return offset;
+  }
+};
+
+template <typename OffsetT>
+struct LoopedElemToLoc<1, true, OffsetT> {
+  int dim;
+  OffsetT offset{0};
+  int index{0};
+
+  __device__ LoopedElemToLoc(int dim) : dim(dim) {}
+
+  __device__ void next(const int* shape, const int64_t* strides) {
+    index++;
+    if (dim > 1) {
+      offset = elem_to_loc<OffsetT>(index, shape, strides, dim);
+    } else {
+      offset += OffsetT(strides[0]);
+    }
+  }
+
+  __device__ void next(int n, const int* shape, const int64_t* strides) {
+    index += n;
+    if (dim > 1) {
+      offset = elem_to_loc<OffsetT>(index, shape, strides, dim);
+    } else {
+      offset = index * OffsetT(strides[0]);
+    }
+  }
+
+  __device__ OffsetT location() {
+    return offset;
+  }
+};
+
+template <typename OffsetT>
+struct LoopedElemToLoc<1, false, OffsetT> {
+  OffsetT offset{0};
+
+  __device__ LoopedElemToLoc(int) {}
+
+  __device__ void next(const int*, const int64_t* strides) {
+    offset += OffsetT(strides[0]);
+  }
+
+  __device__ void next(int n, const int*, const int64_t* strides) {
+    offset += n * OffsetT(strides[0]);
+  }
+
+  __device__ OffsetT location() {
+    return offset;
+  }
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/device_info.cpp b/macos/third_party/mlx/mlx/backend/cuda/device_info.cpp
new file mode 100644
index 0000000..a4a15e2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/device_info.cpp
@@ -0,0 +1,232 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/cuda/cuda.h"
+
+#include <cuda_runtime.h>
+#include <dlfcn.h>
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+namespace mlx::core {
+
+namespace {
+
+// NVML dynamic loading for accurate memory reporting
+// (cudaMemGetInfo only sees current process)
+
+typedef int nvmlReturn_t;
+typedef struct nvmlDevice_st* nvmlDevice_t;
+struct nvmlMemory_t {
+  unsigned long long total;
+  unsigned long long free;
+  unsigned long long used;
+};
+
+struct NVMLState {
+  void* handle = nullptr;
+  nvmlReturn_t (*nvmlInit_v2)() = nullptr;
+  nvmlReturn_t (*nvmlDeviceGetHandleByUUID)(const char*, nvmlDevice_t*) =
+      nullptr;
+  nvmlReturn_t (*nvmlDeviceGetMemoryInfo)(nvmlDevice_t, nvmlMemory_t*) =
+      nullptr;
+};
+
+bool nvml_init(NVMLState& nvml) {
+#ifdef _WIN32
+  nvml.handle = dlopen("nvml.dll", RTLD_LAZY);
+  if (!nvml.handle) {
+    nvml.handle = dlopen(
+        "C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvml.dll", RTLD_LAZY);
+  }
+#else
+  nvml.handle = dlopen("libnvidia-ml.so.1", RTLD_LAZY);
+#endif
+  if (!nvml.handle)
+    return false;
+
+  nvml.nvmlInit_v2 =
+      (decltype(nvml.nvmlInit_v2))dlsym(nvml.handle, "nvmlInit_v2");
+  nvml.nvmlDeviceGetHandleByUUID =
+      (decltype(nvml.nvmlDeviceGetHandleByUUID))dlsym(
+          nvml.handle, "nvmlDeviceGetHandleByUUID");
+  nvml.nvmlDeviceGetMemoryInfo = (decltype(nvml.nvmlDeviceGetMemoryInfo))dlsym(
+      nvml.handle, "nvmlDeviceGetMemoryInfo");
+
+  if (!nvml.nvmlInit_v2 || !nvml.nvmlDeviceGetHandleByUUID ||
+      !nvml.nvmlDeviceGetMemoryInfo) {
+    return false;
+  }
+  return nvml.nvmlInit_v2() == 0;
+}
+
+bool nvml_get_memory(
+    NVMLState& nvml,
+    const char* uuid,
+    size_t* free,
+    size_t* total) {
+  if (!nvml.handle)
+    return false;
+  nvmlDevice_t device;
+  if (nvml.nvmlDeviceGetHandleByUUID(uuid, &device) != 0)
+    return false;
+  nvmlMemory_t mem;
+  if (nvml.nvmlDeviceGetMemoryInfo(device, &mem) != 0)
+    return false;
+  *free = mem.free;
+  *total = mem.total;
+  return true;
+}
+
+std::string format_uuid(const cudaUUID_t& uuid) {
+  char buf[64];
+  snprintf(
+      buf,
+      sizeof(buf),
+      "GPU-%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x",
+      (unsigned char)uuid.bytes[0],
+      (unsigned char)uuid.bytes[1],
+      (unsigned char)uuid.bytes[2],
+      (unsigned char)uuid.bytes[3],
+      (unsigned char)uuid.bytes[4],
+      (unsigned char)uuid.bytes[5],
+      (unsigned char)uuid.bytes[6],
+      (unsigned char)uuid.bytes[7],
+      (unsigned char)uuid.bytes[8],
+      (unsigned char)uuid.bytes[9],
+      (unsigned char)uuid.bytes[10],
+      (unsigned char)uuid.bytes[11],
+      (unsigned char)uuid.bytes[12],
+      (unsigned char)uuid.bytes[13],
+      (unsigned char)uuid.bytes[14],
+      (unsigned char)uuid.bytes[15]);
+  return buf;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info_impl(int device_index) {
+  // Static cache of device properties including UUID (needed for NVML lookup)
+  static auto all_devices = []() {
+    // Get device count
+    int count = 0;
+    cudaGetDeviceCount(&count);
+
+    // Collect info for all devices
+    struct DeviceInfo {
+      std::unordered_map<std::string, std::variant<std::string, size_t>> info;
+      std::string uuid;
+    };
+
+    std::vector<DeviceInfo> devices;
+
+    for (int i = 0; i < count; ++i) {
+      cudaDeviceProp prop;
+      cudaGetDeviceProperties(&prop, i);
+
+      DeviceInfo dev;
+      dev.info["device_name"] = std::string(prop.name);
+      dev.uuid = format_uuid(prop.uuid);
+      dev.info["uuid"] = dev.uuid;
+
+      // Architecture string (e.g., "sm_89")
+      char arch[16];
+      snprintf(arch, sizeof(arch), "sm_%d%d", prop.major, prop.minor);
+      dev.info["architecture"] = std::string(arch);
+
+      // PCI bus ID (domain:bus:device.function)
+      char pci_id[32];
+      snprintf(
+          pci_id,
+          sizeof(pci_id),
+          "%04x:%02x:%02x.0",
+          prop.pciDomainID,
+          prop.pciBusID,
+          prop.pciDeviceID);
+      dev.info["pci_bus_id"] = std::string(pci_id);
+
+      // Compute capability as size_t (to match Metal's variant type)
+      dev.info["compute_capability_major"] = static_cast<size_t>(prop.major);
+      dev.info["compute_capability_minor"] = static_cast<size_t>(prop.minor);
+
+      devices.push_back(std::move(dev));
+    }
+    return devices;
+  }();
+
+  // Initialize NVML once for fresh memory reads
+  static NVMLState nvml;
+  static bool nvml_initialized = nvml_init(nvml);
+
+  if (device_index < 0 ||
+      device_index >= static_cast<int>(all_devices.size())) {
+    static auto empty =
+        std::unordered_map<std::string, std::variant<std::string, size_t>>();
+    return empty;
+  }
+
+  // Return a copy with fresh memory info
+  // Using thread_local to avoid locks while keeping free_memory fresh
+  thread_local auto device_info_copy =
+      std::unordered_map<std::string, std::variant<std::string, size_t>>();
+
+  device_info_copy = all_devices[device_index].info;
+
+  // Get fresh memory info - try NVML first (system-wide), fallback to
+  // cudaMemGetInfo (process-level)
+  size_t free_mem, total_mem;
+
+  if (nvml_initialized &&
+      nvml_get_memory(
+          nvml,
+          all_devices[device_index].uuid.c_str(),
+          &free_mem,
+          &total_mem)) {
+    // NVML succeeded - use system-wide memory
+  } else {
+    // Fallback to cudaMemGetInfo (process-scoped)
+    int prev_device;
+    cudaGetDevice(&prev_device);
+    cudaSetDevice(device_index);
+    cudaMemGetInfo(&free_mem, &total_mem);
+    cudaSetDevice(prev_device);
+  }
+
+  device_info_copy["free_memory"] = free_mem;
+  device_info_copy["total_memory"] = total_mem;
+
+  return device_info_copy;
+}
+
+} // anonymous namespace
+
+namespace gpu {
+
+bool is_available() {
+  return true;
+}
+
+int device_count() {
+  int count = 0;
+  cudaGetDeviceCount(&count);
+  return count;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index) {
+  return device_info_impl(device_index);
+}
+
+} // namespace gpu
+
+namespace cu {
+
+bool is_available() {
+  return true;
+}
+
+} // namespace cu
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/distributed.cu b/macos/third_party/mlx/mlx/backend/cuda/distributed.cu
new file mode 100644
index 0000000..ac79875
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/distributed.cu
@@ -0,0 +1,121 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/distributed/primitives.h"
+#include "mlx/primitives.h"
+
+#include <cassert>
+
+namespace mlx::core::distributed {
+void AllReduce::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  auto set_input_output = [&](const array& in,
+                              array& out) -> std::pair<array, array> {
+    if (!in.flags().row_contiguous) {
+      copy_gpu(in, out, CopyType::General, s);
+      return {out, out};
+    } else if (in.is_donatable()) {
+      out.copy_shared_buffer(in);
+      return {in, out};
+    } else {
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+      return {in, out};
+    }
+  };
+
+  auto [input, output] = set_input_output(inputs[0], outputs[0]);
+
+  encoder.set_input_array(input);
+  encoder.set_output_array(output);
+
+  auto capture = encoder.capture_context();
+
+  switch (reduce_type_) {
+    case Sum:
+      distributed::detail::all_sum(group(), input, output, s);
+      break;
+    case Max:
+      distributed::detail::all_max(group(), input, output, s);
+      break;
+    case Min:
+      distributed::detail::all_min(group(), input, output, s);
+      break;
+    default:
+      throw std::runtime_error(
+          "Only all reduce sum, max, and min are supported.");
+  }
+}
+
+void AllGather::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  auto ensure_contiguous = [&s, &encoder](const array& x) {
+    if (x.flags().row_contiguous) {
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      encoder.add_temporary(x_copy);
+      return x_copy;
+    }
+  };
+
+  auto input = ensure_contiguous(inputs[0]);
+  outputs[0].set_data(cu::malloc_async(outputs[0].nbytes(), encoder));
+
+  encoder.set_input_array(input);
+  encoder.set_output_array(outputs[0]);
+
+  auto capture = encoder.capture_context();
+  distributed::detail::all_gather(group(), input, outputs[0], s);
+}
+
+void ReduceScatter::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(inputs.size() == 1);
+  assert(outputs.size() == 1);
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  auto ensure_contiguous = [&s, &encoder](const array& x) {
+    if (x.flags().row_contiguous) {
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      encoder.add_temporary(x_copy);
+      return x_copy;
+    }
+  };
+
+  auto input = ensure_contiguous(inputs[0]);
+  outputs[0].set_data(cu::malloc_async(outputs[0].nbytes(), encoder));
+
+  encoder.set_input_array(input);
+  encoder.set_output_array(outputs[0]);
+
+  auto capture = encoder.capture_context();
+
+  switch (reduce_type_) {
+    case Sum:
+      distributed::detail::sum_scatter(group(), input, outputs[0], s);
+      break;
+    default:
+      throw std::runtime_error("Only sum scatter is supported. ");
+  }
+}
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/backend/cuda/eval.cpp b/macos/third_party/mlx/mlx/backend/cuda/eval.cpp
new file mode 100644
index 0000000..1d04c8f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/eval.cpp
@@ -0,0 +1,66 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/gpu/eval.h"
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core::gpu {
+
+void new_stream(Stream s) {
+  // Force initalization of CUDA, so CUDA runtime get destroyed at last.
+  cudaFree(nullptr);
+  // Make sure CUDA event pool get destroyed after device and stream.
+  cu::CudaEvent::init_pool();
+  // Ensure the static stream objects get created.
+  cu::get_command_encoder(s);
+}
+
+void eval(array& arr) {
+  nvtx3::scoped_range r("gpu::eval");
+  auto outputs = arr.outputs();
+  {
+    // If the array is a tracer hold a reference
+    // to its inputs so they don't get donated
+    std::vector<array> inputs;
+    if (arr.is_tracer()) {
+      inputs = arr.inputs();
+    }
+    arr.primitive().eval_gpu(arr.inputs(), outputs);
+  }
+
+  auto& stream = arr.primitive().stream();
+  auto& encoder = cu::get_command_encoder(stream);
+  // Keep used buffers alive until kernel finishes running.
+  for (auto& in : arr.inputs()) {
+    // Except for the donated one.
+    if (in.data_shared_ptr() != arr.data_shared_ptr()) {
+      encoder.add_temporary(in);
+    }
+  }
+  for (auto& s : arr.siblings()) {
+    encoder.add_temporary(s);
+  }
+
+  if (encoder.needs_commit()) {
+    scheduler::notify_new_task(stream);
+    encoder.add_completed_handler(
+        [stream]() { scheduler::notify_task_completion(stream); });
+    encoder.commit();
+  }
+}
+
+void finalize(Stream s) {
+  nvtx3::scoped_range r("gpu::finalize");
+  cu::get_command_encoder(s).commit();
+}
+
+void synchronize(Stream s) {
+  nvtx3::scoped_range r("gpu::synchronize");
+  cu::get_command_encoder(s).synchronize();
+}
+
+} // namespace mlx::core::gpu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/event.cu b/macos/third_party/mlx/mlx/backend/cuda/event.cu
new file mode 100644
index 0000000..7334c17
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/event.cu
@@ -0,0 +1,413 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/event.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/event.h"
+#include "mlx/scheduler.h"
+
+#include <map>
+#include <vector>
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+///////////////////////////////////////////////////////////////////////////////
+// CudaEvent implementations
+///////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// Manage cached cudaEvent_t objects.
+class CudaEventPool {
+ public:
+  CudaEventHandle create(Device& d, int flags) {
+    if (!on_creation_thread()) {
+      return CudaEventHandle(d, flags);
+    }
+    auto& cache = cache_for(d, flags);
+    if (cache.empty()) {
+      return CudaEventHandle(d, flags);
+    } else {
+      CudaEventHandle ret = std::move(cache.back());
+      cache.pop_back();
+      return ret;
+    }
+  }
+
+  void release(CudaEventHandle event) {
+    if (!on_creation_thread()) {
+      // Event will be destroyed directly instead of getting moved to cache.
+      return;
+    }
+    cache_for(event.device, event.flags).push_back(std::move(event));
+  }
+
+ private:
+  std::vector<CudaEventHandle>& cache_for(Device& d, int flags) {
+    return cache_[d.cuda_device()][flags];
+  }
+
+  bool on_creation_thread() {
+    return std::this_thread::get_id() == thread_id_;
+  }
+
+  // The CudaEvent may be created and destroyed on different threads (for
+  // example when waiting on GPU work in CPU stream), we don't want to make
+  // the cache thread-safe as it adds overhead, so we just skip cache when
+  // using events in worker threads.
+  std::thread::id thread_id_{std::this_thread::get_id()};
+
+  // {device: {flags: [events]}}
+  std::map<int, std::map<int, std::vector<CudaEventHandle>>> cache_;
+};
+
+CudaEventPool& cuda_event_pool() {
+  static CudaEventPool pool;
+  return pool;
+}
+
+} // namespace
+
+CudaEventHandle::CudaEventHandle(Device& d, int flags)
+    : device(d), flags(flags) {
+  device.make_current();
+  CHECK_CUDA_ERROR(cudaEventCreateWithFlags(&handle_, flags));
+  assert(handle_ != nullptr);
+}
+
+CudaEvent::CudaEvent(Device& d, int flags)
+    : event_(cuda_event_pool().create(d, flags)) {}
+
+CudaEvent::~CudaEvent() {
+  cuda_event_pool().release(std::move(event_));
+}
+
+void CudaEvent::wait() {
+  nvtx3::scoped_range r("cu::CudaEvent::wait");
+  event_.device.make_current();
+  cudaEventSynchronize(event_);
+}
+
+void CudaEvent::wait(cudaStream_t stream) {
+  event_.device.make_current();
+  cudaStreamWaitEvent(stream, event_);
+}
+
+void CudaEvent::record(cudaStream_t stream) {
+  event_.device.make_current();
+  cudaEventRecord(event_, stream);
+}
+
+bool CudaEvent::completed() const {
+  // Note: cudaEventQuery can be safely called from any device.
+  return cudaEventQuery(event_) == cudaSuccess;
+}
+
+// static
+void CudaEvent::init_pool() {
+  cuda_event_pool();
+}
+
+// Wraps CudaEvent with a few features:
+// 1. The class can be copied.
+// 2. Make wait/record work with CPU streams.
+// 3. Add checks for waiting on un-recorded event.
+class CopyableCudaEvent {
+ public:
+  explicit CopyableCudaEvent(Device& d)
+      : event_(
+            std::make_shared<CudaEvent>(
+                d,
+                cudaEventDisableTiming | cudaEventBlockingSync)) {}
+
+  void wait() {
+    event_->wait();
+  }
+
+  void wait(Stream s) {
+    if (s.device == mlx::core::Device::cpu) {
+      scheduler::enqueue(s, [*this]() mutable {
+        check_recorded();
+        event_->wait();
+      });
+    } else {
+      check_recorded();
+      auto& encoder = cu::get_command_encoder(s);
+      encoder.commit();
+      event_->wait(encoder.stream());
+    }
+  }
+
+  void record(Stream s) {
+    if (s.device == mlx::core::Device::cpu) {
+      throw std::runtime_error("CudaEvent can not wait on CPU stream.");
+    } else {
+      auto& encoder = cu::get_command_encoder(s);
+      encoder.commit();
+      event_->record(encoder.stream());
+      recorded_ = true;
+    }
+  }
+
+  bool is_signaled() const {
+    return recorded_ && event_->completed();
+  }
+
+ private:
+  void check_recorded() const {
+    if (!recorded_) {
+      throw std::runtime_error(
+          "Should not wait on a CudaEvent before recording.");
+    }
+  }
+
+  std::shared_ptr<CudaEvent> event_;
+  bool recorded_{false};
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// AtomicEvent implementations
+///////////////////////////////////////////////////////////////////////////////
+
+__host__ __device__ void event_wait(uint32_t* ptr, uint32_t value) {
+  cuda::atomic_ref<uint32_t> ac(*ptr);
+  uint32_t current;
+  while ((current = ac.load()) < value) {
+    ac.wait(current);
+  }
+}
+
+__host__ __device__ void event_signal(uint32_t* ptr, uint32_t value) {
+  cuda::atomic_ref<uint32_t> ac(*ptr);
+  ac.store(value);
+  ac.notify_all();
+}
+
+__global__ void event_wait_kernel(uint32_t* ptr, uint32_t value) {
+  event_wait(ptr, value);
+}
+
+__global__ void event_signal_kernel(uint32_t* ptr, uint32_t value) {
+  __threadfence_system();
+  event_signal(ptr, value);
+  __threadfence_system();
+}
+
+auto check_gpu_coherency() {
+  static auto coherency = []() {
+    int device_count = gpu::device_count();
+    bool concurrent_managed_access = true;
+    bool host_native_atomic = true;
+    for (int i = 0; i < device_count; ++i) {
+      auto& d = cu::device(i);
+      concurrent_managed_access &= d.concurrent_managed_access();
+      host_native_atomic &= d.host_native_atomic();
+    }
+    return std::make_tuple(concurrent_managed_access, host_native_atomic);
+  }();
+  return coherency;
+}
+
+AtomicEvent::AtomicEvent() {
+  void* buf;
+  cudaError_t (*cuda_free)(void*);
+  // There are 3 kinds of systems we are implementing for:
+  // 1. concurrentManagedAccess == true
+  //    => use cuda::atom_ref on managed memory
+  // 2. hostNativeAtomicSupported == true
+  //    => use cuda::atom_ref on pinned host memory
+  // 2. no hardware cpu/gpu coherency
+  //    => use cuda::atom_ref on device memory
+  auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency();
+  if (concurrent_managed_access) {
+    CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = true;
+  } else if (host_native_atomic) {
+    CHECK_CUDA_ERROR(cudaMallocHost(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFreeHost;
+    coherent_ = true;
+  } else {
+    CHECK_CUDA_ERROR(cudaMalloc(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = false;
+  }
+  buf_ = std::shared_ptr<void>(
+      buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); });
+  if (coherent_) {
+    *ptr() = 0;
+  } else {
+    CHECK_CUDA_ERROR(cudaMemset(buf, 0, sizeof(uint32_t)));
+  }
+}
+
+void AtomicEvent::wait(uint32_t value) {
+  nvtx3::scoped_range r("cu::AtomicEvent::wait");
+  if (coherent_) {
+    event_wait(ptr(), value);
+  } else {
+    while (!is_signaled(value)) {
+      std::this_thread::yield();
+    }
+  }
+}
+
+void AtomicEvent::wait(cudaStream_t stream, uint32_t value) {
+  event_wait_kernel<<<1, 1, 0, stream>>>(ptr(), value);
+}
+
+void AtomicEvent::wait(Stream s, uint32_t value) {
+  nvtx3::scoped_range r("cu::AtomicEvent::wait(s)");
+  if (s.device == mlx::core::Device::cpu) {
+    scheduler::enqueue(s, [*this, value]() mutable { wait(value); });
+  } else {
+    auto& encoder = get_command_encoder(s);
+    encoder.commit();
+    wait(encoder.stream(), value);
+    encoder.add_completed_handler([buf = buf_]() {});
+  }
+}
+
+void AtomicEvent::signal(uint32_t value) {
+  nvtx3::scoped_range r("cu::AtomicEvent::signal");
+  if (coherent_) {
+    event_signal(ptr(), value);
+  } else {
+    signal(signal_stream(), value);
+  }
+}
+
+void AtomicEvent::signal(cudaStream_t stream, uint32_t value) {
+  event_signal_kernel<<<1, 1, 0, stream>>>(ptr(), value);
+}
+
+void AtomicEvent::signal(Stream s, uint32_t value) {
+  nvtx3::scoped_range r("cu::AtomicEvent::signal(s)");
+  if (s.device == mlx::core::Device::cpu) {
+    // Signal through a GPU stream so the atomic is updated in GPU - updating
+    // the atomic in CPU sometimes does not get GPU notified.
+    scheduler::enqueue(
+        s, [*this, value]() mutable { signal(signal_stream(), value); });
+  } else {
+    auto& encoder = get_command_encoder(s);
+    encoder.commit();
+    signal(encoder.stream(), value);
+    encoder.add_completed_handler([buf = buf_]() {});
+  }
+}
+
+bool AtomicEvent::is_signaled(uint32_t val) const {
+  return value() >= val;
+}
+
+uint32_t AtomicEvent::value() const {
+  nvtx3::scoped_range r("cu::AtomicEvent::value");
+  if (coherent_) {
+    cuda::atomic_ref<uint32_t> ac(*ptr());
+    return ac.load();
+  } else {
+    uint32_t val;
+    CHECK_CUDA_ERROR(
+        cudaMemcpy(&val, ptr(), sizeof(uint32_t), cudaMemcpyDeviceToHost));
+    return val;
+  }
+}
+
+const CudaStream& AtomicEvent::signal_stream() {
+  static CudaStream stream(device(0));
+  return stream;
+}
+
+} // namespace cu
+
+///////////////////////////////////////////////////////////////////////////////
+// Event implementations
+///////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+struct EventImpl {
+  // CudaEvent is preferred when possible because it is fast, however we have
+  // to fallback to AtomicEvent in following cases:
+  // 1. the event is used to wait/signal a cpu stream;
+  // 2. signal value other than 1 has been specified.
+  std::unique_ptr<cu::CopyableCudaEvent> cuda;
+  std::unique_ptr<cu::AtomicEvent> atomic;
+
+  bool is_created() const {
+    return cuda || atomic;
+  }
+
+  void ensure_created(Stream s, uint64_t signal_value) {
+    if (is_created()) {
+      return;
+    }
+    if (s.device == mlx::core::Device::cpu || signal_value > 1) {
+      nvtx3::mark("Using slow AtomicEvent");
+      atomic = std::make_unique<cu::AtomicEvent>();
+    } else {
+      cuda = std::make_unique<cu::CopyableCudaEvent>(cu::device(s.device));
+    }
+  }
+};
+
+} // namespace
+
+Event::Event(Stream s) : stream_(s) {
+  event_ = std::shared_ptr<void>(
+      new EventImpl(), [](void* ptr) { delete static_cast<EventImpl*>(ptr); });
+}
+
+void Event::wait() {
+  auto* event = static_cast<EventImpl*>(event_.get());
+  assert(event->is_created());
+  if (event->cuda) {
+    assert(value() == 1);
+    event->cuda->wait();
+  } else {
+    event->atomic->wait(value());
+  }
+  CHECK_CUDA_ERROR(cudaPeekAtLastError());
+}
+
+void Event::wait(Stream s) {
+  auto* event = static_cast<EventImpl*>(event_.get());
+  assert(event->is_created());
+  if (event->cuda) {
+    assert(value() == 1);
+    event->cuda->wait(s);
+  } else {
+    event->atomic->wait(s, value());
+  }
+}
+
+void Event::signal(Stream s) {
+  auto* event = static_cast<EventImpl*>(event_.get());
+  event->ensure_created(s, value());
+  if (event->cuda) {
+    assert(value() == 1);
+    event->cuda->record(s);
+  } else {
+    event->atomic->signal(s, value());
+  }
+}
+
+bool Event::is_signaled() const {
+  auto* event = static_cast<EventImpl*>(event_.get());
+  if (!event->is_created()) {
+    return false;
+  }
+  if (event->cuda) {
+    assert(value() == 1);
+    return event->cuda->is_signaled();
+  } else {
+    return event->atomic->is_signaled(value());
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/event.h b/macos/third_party/mlx/mlx/backend/cuda/event.h
new file mode 100644
index 0000000..68aa4d8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/event.h
@@ -0,0 +1,79 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/stream.h"
+
+#include <memory>
+
+#include <cuda_runtime.h>
+#include <cuda/atomic>
+
+namespace mlx::core::cu {
+
+class Device;
+
+// RAII-managed move-only wrapper of cudaEvent_t.
+struct CudaEventHandle : public CudaHandle<cudaEvent_t, cudaEventDestroy> {
+  CudaEventHandle(Device& d, int flags);
+  Device& device;
+  int flags;
+};
+
+// Wrapper of native cuda event. It can synchronize between GPU streams, or wait
+// on GPU stream in CPU stream, but can not wait on CPU stream.
+class CudaEvent {
+ public:
+  CudaEvent(Device& d, int flags);
+  ~CudaEvent();
+
+  CudaEvent(CudaEvent&&) = default;
+  CudaEvent& operator=(CudaEvent&&) = default;
+
+  CudaEvent(const CudaEvent&) = delete;
+  CudaEvent& operator=(const CudaEvent&) = delete;
+
+  void wait();
+  void wait(cudaStream_t stream);
+  void record(cudaStream_t stream);
+
+  // Return whether the recorded kernels have completed. Note that this method
+  // returns true if record() has not been called.
+  bool completed() const;
+
+  // Internal: make sure event pool is initialized.
+  static void init_pool();
+
+ private:
+  CudaEventHandle event_;
+};
+
+// Event that can synchronize between CPU and GPU. It is much slower than
+// CudaEvent so the latter should always be preferred when possible.
+class AtomicEvent {
+ public:
+  AtomicEvent();
+
+  void wait(uint32_t value);
+  void wait(cudaStream_t stream, uint32_t value);
+  void wait(Stream s, uint32_t value);
+  void signal(uint32_t value);
+  void signal(cudaStream_t stream, uint32_t value);
+  void signal(Stream s, uint32_t value);
+  bool is_signaled(uint32_t value) const;
+  uint32_t value() const;
+
+ private:
+  const CudaStream& signal_stream();
+
+  uint32_t* ptr() const {
+    return static_cast<uint32_t*>(buf_.get());
+  }
+
+  bool coherent_;
+  std::shared_ptr<void> buf_;
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/fence.cpp b/macos/third_party/mlx/mlx/backend/cuda/fence.cpp
new file mode 100644
index 0000000..dff5f74
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/fence.cpp
@@ -0,0 +1,41 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/fence.h"
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/event.h"
+
+namespace mlx::core {
+
+struct FenceImpl {
+  uint32_t count;
+  cu::AtomicEvent event;
+};
+
+Fence::Fence(Stream s) {
+  fence_ = std::shared_ptr<void>(
+      new FenceImpl{0}, [](void* ptr) { delete static_cast<FenceImpl*>(ptr); });
+}
+
+void Fence::wait(Stream s, const array&) {
+  auto* fence = static_cast<FenceImpl*>(fence_.get());
+  fence->event.wait(fence->count);
+}
+
+void Fence::update(Stream s, const array& a, bool cross_device) {
+  auto* fence = static_cast<FenceImpl*>(fence_.get());
+  if (cross_device) {
+    // Move to managed memory if there is a device switch
+    auto& cbuf =
+        *static_cast<cu::CudaBuffer*>(const_cast<array&>(a).buffer().ptr());
+    if (cbuf.device != -1) {
+      auto& encoder = cu::get_command_encoder(s);
+      encoder.commit();
+      cu::allocator().move_to_unified_memory(cbuf, encoder.stream());
+    }
+  }
+  fence->count++;
+  fence->event.signal(s, fence->count);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm.cpp b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm.cpp
new file mode 100644
index 0000000..7834045
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm.cpp
@@ -0,0 +1,233 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/utils.h"
+
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+namespace {
+
+cublasComputeType_t dtype_to_compute_type(Dtype dtype) {
+  switch (dtype) {
+    case float16:
+      return CUBLAS_COMPUTE_32F;
+    case bfloat16:
+      return CUBLAS_COMPUTE_32F;
+    case float32:
+      return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32
+                                           : CUBLAS_COMPUTE_32F;
+    case float64:
+      return CUBLAS_COMPUTE_64F;
+    case complex64:
+      return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32
+                                           : CUBLAS_COMPUTE_32F;
+    default:
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype)));
+  }
+}
+
+} // namespace
+
+CublasGemm::CublasGemm(
+    cu::Device& device,
+    Dtype dtype,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride) {
+  scale_type_ = cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm");
+  if (dtype == bfloat16 || dtype == float16) {
+    scale_type_ = CUDA_R_32F;
+  }
+  cudaDataType_t cublas_dtype =
+      cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm");
+
+  init_base(
+      device,
+      scale_type_,
+      dtype_to_compute_type(dtype),
+      cublas_dtype,
+      cublas_dtype,
+      a_transposed,
+      a_rows,
+      a_cols,
+      lda,
+      b_transposed,
+      b_rows,
+      b_cols,
+      ldb,
+      batch_count,
+      a_batch_stride,
+      b_batch_stride);
+}
+
+CublasGemm::CublasGemm(
+    cu::Device& device,
+    Dtype dtype,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int64_t ldc,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride,
+    int64_t c_batch_stride)
+    : CublasGemm(
+          device,
+          dtype,
+          a_transposed,
+          a_rows,
+          a_cols,
+          lda,
+          b_transposed,
+          b_rows,
+          b_cols,
+          ldb,
+          batch_count,
+          a_batch_stride,
+          b_batch_stride) {
+  auto type = cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm");
+  c_desc_ = cublas_utils::create_matrix_layout(
+      type, b_cols, a_rows, false, ldc, batch_count, c_batch_stride);
+}
+
+void CublasGemm::set_out(
+    Dtype dtype,
+    bool transposed,
+    uint64_t rows,
+    uint64_t cols,
+    int64_t ld,
+    int32_t batch_count,
+    int64_t batch_stride) {
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(out_desc_));
+  out_desc_ = cublas_utils::create_matrix_layout(
+      cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm"),
+      cols,
+      rows,
+      transposed,
+      ld,
+      batch_count,
+      batch_stride);
+}
+
+void CublasGemm::run(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const Shape& batch_shape,
+    const Strides& a_batch_strides,
+    const Strides& b_batch_strides,
+    float alpha) {
+  int batch_count = out.size() / (M_ * N_);
+  if (batch_count / batch_shape.back() > 1) {
+    run_batched(
+        encoder,
+        out,
+        a,
+        b,
+        batch_shape,
+        a_batch_strides,
+        b_batch_strides,
+        alpha);
+    return;
+  }
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+
+  execute(
+      encoder,
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
+      nullptr,
+      alpha);
+}
+
+void CublasGemm::run(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& c,
+    const Shape& batch_shape,
+    const Strides& a_batch_strides,
+    const Strides& b_batch_strides,
+    const Strides& c_batch_strides,
+    float alpha,
+    float beta) {
+  int batch_count = out.size() / (M_ * N_);
+  if (batch_count / batch_shape.back() > 1) {
+    run_batched(
+        encoder,
+        out,
+        a,
+        b,
+        c,
+        batch_shape,
+        a_batch_strides,
+        b_batch_strides,
+        c_batch_strides,
+        alpha,
+        beta);
+    return;
+  }
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(c);
+  encoder.set_output_array(out);
+
+  execute(
+      encoder,
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
+      gpu_ptr<void>(c),
+      alpha,
+      beta);
+}
+
+void CublasGemm::execute(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* c,
+    float alpha /* = 1 */,
+    float beta /* = 0 */) {
+  const void* alpha_ptr = &alpha;
+  const void* beta_ptr = &beta;
+  complex64_t alpha_c, beta_c;
+  if (scale_type_ == CUDA_C_32F) {
+    alpha_c = complex64_t{alpha, 0.0f};
+    beta_c = complex64_t{beta, 0.0f};
+    alpha_ptr = &alpha_c;
+    beta_ptr = &beta_c;
+  }
+
+  execute_matmul(encoder, out, a, b, c, alpha_ptr, beta_ptr);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm.h b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm.h
new file mode 100644
index 0000000..1fad45e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm.h
@@ -0,0 +1,114 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/device.h"
+
+#include <cublasLt.h>
+
+namespace mlx::core {
+
+class CublasGemm : public CublasMatmulBase {
+ public:
+  CublasGemm(
+      cu::Device& device,
+      Dtype dtype,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride);
+
+  CublasGemm(
+      cu::Device& device,
+      Dtype dtype,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int64_t ldc,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride,
+      int64_t c_batch_stride);
+
+  // The output's descriptor is inferred from inputs by default, use this method
+  // for unusual output.
+  void set_out(
+      Dtype dtype,
+      bool transposed,
+      uint64_t rows,
+      uint64_t cols,
+      int64_t ld,
+      int32_t batch_count,
+      int64_t batch_stride);
+
+  void run(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const Shape& batch_shape,
+      const Strides& a_batch_strides,
+      const Strides& b_batch_strides,
+      float alpha = 1.0f);
+
+  void run(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const array& c,
+      const Shape& batch_shape,
+      const Strides& a_batch_strides,
+      const Strides& b_batch_strides,
+      const Strides& c_batch_strides,
+      float alpha,
+      float beta);
+
+ private:
+  void run_batched(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const Shape& batch_shape,
+      const Strides& a_batch_strides,
+      const Strides& b_batch_strides,
+      float alpha);
+
+  void run_batched(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const array& c,
+      const Shape& batch_shape,
+      const Strides& a_batch_strides,
+      const Strides& b_batch_strides,
+      const Strides& c_batch_strides,
+      float alpha,
+      float beta);
+
+  void execute(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* c,
+      float alpha = 1,
+      float beta = 0);
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp
new file mode 100644
index 0000000..67b8c27
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp
@@ -0,0 +1,77 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/cublas_gemm.h"
+
+namespace mlx::core {
+
+void CublasGemm::run_batched(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const Shape& batch_shape,
+    const Strides& a_batch_strides,
+    const Strides& b_batch_strides,
+    float alpha) {
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  auto nbatch = out.size() / (M_ * N_ * batch_shape.back());
+  ContiguousIterator a_it(batch_shape, a_batch_strides, batch_shape.size() - 1);
+  ContiguousIterator b_it(batch_shape, b_batch_strides, batch_shape.size() - 1);
+  auto concurrent = encoder.concurrent_context();
+  for (size_t i = 0; i < nbatch; ++i) {
+    execute(
+        encoder,
+        gpu_ptr<int8_t>(out) +
+            out.itemsize() * i * batch_shape.back() * M_ * N_,
+        gpu_ptr<int8_t>(a) + a.itemsize() * a_it.loc,
+        gpu_ptr<int8_t>(b) + b.itemsize() * b_it.loc,
+        nullptr,
+        alpha);
+    a_it.step();
+    b_it.step();
+  }
+}
+
+void CublasGemm::run_batched(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& c,
+    const Shape& batch_shape,
+    const Strides& a_batch_strides,
+    const Strides& b_batch_strides,
+    const Strides& c_batch_strides,
+    float alpha,
+    float beta) {
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(c);
+  encoder.set_output_array(out);
+
+  auto nbatch = out.size() / (M_ * N_ * batch_shape.back());
+  ContiguousIterator a_it(batch_shape, a_batch_strides, batch_shape.size() - 1);
+  ContiguousIterator b_it(batch_shape, b_batch_strides, batch_shape.size() - 1);
+  ContiguousIterator c_it(batch_shape, c_batch_strides, batch_shape.size() - 1);
+  auto concurrent = encoder.concurrent_context();
+  for (size_t i = 0; i < nbatch; ++i) {
+    execute(
+        encoder,
+        gpu_ptr<int8_t>(out) +
+            out.itemsize() * i * batch_shape.back() * M_ * N_,
+        gpu_ptr<int8_t>(a) + a.itemsize() * a_it.loc,
+        gpu_ptr<int8_t>(b) + b.itemsize() * b_it.loc,
+        gpu_ptr<int8_t>(c) + c.itemsize() * c_it.loc,
+        alpha,
+        beta);
+    a_it.step();
+    b_it.step();
+    c_it.step();
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu
new file mode 100644
index 0000000..e4faba6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu
@@ -0,0 +1,329 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <int NDIM>
+__global__ void set_mm_device_pointers_nd(
+    int8_t** pointers,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* out_start,
+    int item_size,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> batch_shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> a_batch_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> b_batch_strides,
+    int64_t batch_stride,
+    int batch_count) {
+  auto index = cg::this_grid().thread_rank();
+  if (index >= batch_count) {
+    return;
+  }
+  auto [a_offset, b_offset] = elem_to_loc_nd<NDIM>(
+      index,
+      batch_shape.data(),
+      a_batch_strides.data(),
+      b_batch_strides.data());
+  pointers[index] = a_start + item_size * a_offset;
+  pointers[index + batch_count] = b_start + item_size * b_offset;
+  pointers[index + 2 * batch_count] =
+      out_start + item_size * index * batch_stride;
+}
+
+__global__ void set_mm_device_pointers_g(
+    int8_t** pointers,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* out_start,
+    int item_size,
+    const __grid_constant__ Shape batch_shape,
+    const __grid_constant__ Strides a_batch_strides,
+    const __grid_constant__ Strides b_batch_strides,
+    int64_t batch_stride,
+    int batch_ndim,
+    int batch_count) {
+  auto index = cg::this_grid().thread_rank();
+  if (index >= batch_count) {
+    return;
+  }
+  auto [a_offset, b_offset] = elem_to_loc(
+      index,
+      batch_shape.data(),
+      a_batch_strides.data(),
+      b_batch_strides.data(),
+      batch_ndim);
+  pointers[index] = a_start + item_size * a_offset;
+  pointers[index + batch_count] = b_start + item_size * b_offset;
+  pointers[index + 2 * batch_count] =
+      out_start + item_size * index * batch_stride;
+}
+
+template <int NDIM>
+__global__ void set_addmm_device_pointers_nd(
+    int8_t** pointers,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* c_start,
+    int8_t* out_start,
+    int item_size,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> batch_shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> a_batch_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> b_batch_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> c_batch_strides,
+    int64_t batch_stride,
+    int batch_count) {
+  auto index = cg::this_grid().thread_rank();
+  if (index >= batch_count) {
+    return;
+  }
+  auto [a_offset, b_offset, c_offset] = elem_to_loc_nd<NDIM>(
+      index,
+      batch_shape.data(),
+      a_batch_strides.data(),
+      b_batch_strides.data(),
+      c_batch_strides.data());
+  pointers[index] = a_start + item_size * a_offset;
+  pointers[index + batch_count] = b_start + item_size * b_offset;
+  pointers[index + 2 * batch_count] = c_start + item_size * c_offset;
+  pointers[index + 3 * batch_count] =
+      out_start + item_size * index * batch_stride;
+}
+
+__global__ void set_addmm_device_pointers_g(
+    int8_t** pointers,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* c_start,
+    int8_t* out_start,
+    int item_size,
+    const __grid_constant__ Shape batch_shape,
+    const __grid_constant__ Strides a_batch_strides,
+    const __grid_constant__ Strides b_batch_strides,
+    const __grid_constant__ Strides c_batch_strides,
+    int64_t batch_stride,
+    int batch_ndim,
+    int batch_count) {
+  auto index = cg::this_grid().thread_rank();
+  if (index >= batch_count) {
+    return;
+  }
+  auto [a_offset, b_offset, c_offset] = elem_to_loc(
+      index,
+      batch_shape.data(),
+      a_batch_strides.data(),
+      b_batch_strides.data(),
+      c_batch_strides.data(),
+      batch_ndim);
+  pointers[index] = a_start + item_size * a_offset;
+  pointers[index + batch_count] = b_start + item_size * b_offset;
+  pointers[index + 2 * batch_count] = c_start + item_size * c_offset;
+  pointers[index + 3 * batch_count] =
+      out_start + item_size * index * batch_stride;
+}
+
+} // namespace cu
+
+namespace {
+
+void set_pointer_mode(cublasLtMatrixLayout_t desc, int batch_count) {
+  auto batch_mode = CUBLASLT_BATCH_MODE_POINTER_ARRAY;
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+      desc,
+      CUBLASLT_MATRIX_LAYOUT_BATCH_MODE,
+      &batch_mode,
+      sizeof(batch_mode)));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+      desc, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(int32_t)));
+}
+
+} // namespace
+
+void CublasGemm::run_batched(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const Shape& batch_shape,
+    const Strides& a_batch_strides,
+    const Strides& b_batch_strides,
+    float alpha) {
+  int batch_count = out.size() / (M_ * N_);
+  set_pointer_mode(a_desc_, batch_count);
+  set_pointer_mode(b_desc_, batch_count);
+  set_pointer_mode(out_desc_, batch_count);
+
+  // Launch kernel to set device offsets
+  auto pointers = array(
+      cu::malloc_async(batch_count * sizeof(void*) * 3, encoder),
+      {batch_count * 3},
+      uint64);
+
+  encoder.add_temporary(pointers);
+  encoder.set_output_array(pointers);
+
+  int block_dims = std::min(batch_count, 256);
+  int num_blocks = cuda::ceil_div(batch_count, block_dims);
+  int64_t batch_stride = M_ * N_;
+  int item_size = out.itemsize();
+
+  int ndim = batch_shape.size();
+  if (ndim <= 3) {
+    dispatch_1_2_3(ndim, [&](auto ndim_constant) {
+      encoder.add_kernel_node(
+          cu::set_mm_device_pointers_nd<ndim_constant()>,
+          num_blocks,
+          block_dims,
+          0,
+          gpu_ptr<int8_t*>(pointers),
+          gpu_ptr<int8_t>(a),
+          gpu_ptr<int8_t>(b),
+          gpu_ptr<int8_t>(out),
+          item_size,
+          const_param<ndim_constant()>(batch_shape),
+          const_param<ndim_constant()>(a_batch_strides),
+          const_param<ndim_constant()>(b_batch_strides),
+          batch_stride,
+          batch_count);
+    });
+  } else {
+    encoder.add_kernel_node(
+        cu::set_mm_device_pointers_g,
+        num_blocks,
+        block_dims,
+        0,
+        gpu_ptr<int8_t*>(pointers),
+        gpu_ptr<int8_t>(a),
+        gpu_ptr<int8_t>(b),
+        gpu_ptr<int8_t>(out),
+        item_size,
+        const_param(batch_shape),
+        const_param(a_batch_strides),
+        const_param(b_batch_strides),
+        batch_stride,
+        ndim,
+        batch_count);
+  }
+
+  // Run matmul
+  encoder.set_input_array(pointers);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+
+  auto a_pointers = gpu_ptr<int8_t*>(pointers);
+  auto b_pointers = a_pointers + batch_count;
+  auto out_pointers = b_pointers + batch_count;
+  execute(
+      encoder,
+      reinterpret_cast<void*>(out_pointers),
+      reinterpret_cast<void*>(a_pointers),
+      reinterpret_cast<void*>(b_pointers),
+      nullptr,
+      alpha);
+}
+
+void CublasGemm::run_batched(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& c,
+    const Shape& batch_shape,
+    const Strides& a_batch_strides,
+    const Strides& b_batch_strides,
+    const Strides& c_batch_strides,
+    float alpha,
+    float beta) {
+  int batch_count = out.size() / (M_ * N_);
+  set_pointer_mode(a_desc_, batch_count);
+  set_pointer_mode(b_desc_, batch_count);
+  set_pointer_mode(c_desc_, batch_count);
+  set_pointer_mode(out_desc_, batch_count);
+
+  // Launch kernel to set device offsets
+  auto pointers = array(
+      cu::malloc_async(batch_count * sizeof(uint64_t) * 4, encoder),
+      {batch_count * 4},
+      uint64);
+
+  encoder.add_temporary(pointers);
+  encoder.set_output_array(pointers);
+
+  int block_dims = std::min(batch_count, 256);
+  int num_blocks = cuda::ceil_div(batch_count, block_dims);
+  int64_t batch_stride = M_ * N_;
+  int item_size = out.itemsize();
+
+  int ndim = batch_shape.size();
+  if (ndim <= 3) {
+    dispatch_1_2_3(ndim, [&](auto ndim_constant) {
+      encoder.add_kernel_node(
+          cu::set_addmm_device_pointers_nd<ndim_constant()>,
+          num_blocks,
+          block_dims,
+          0,
+          gpu_ptr<int8_t*>(pointers),
+          gpu_ptr<int8_t>(a),
+          gpu_ptr<int8_t>(b),
+          gpu_ptr<int8_t>(c),
+          gpu_ptr<int8_t>(out),
+          item_size,
+          const_param<ndim_constant()>(batch_shape),
+          const_param<ndim_constant()>(a_batch_strides),
+          const_param<ndim_constant()>(b_batch_strides),
+          const_param<ndim_constant()>(c_batch_strides),
+          batch_stride,
+          batch_count);
+    });
+  } else {
+    encoder.add_kernel_node(
+        cu::set_addmm_device_pointers_g,
+        num_blocks,
+        block_dims,
+        0,
+        gpu_ptr<int8_t*>(pointers),
+        gpu_ptr<int8_t>(a),
+        gpu_ptr<int8_t>(b),
+        gpu_ptr<int8_t>(c),
+        gpu_ptr<int8_t>(out),
+        item_size,
+        const_param(batch_shape),
+        const_param(a_batch_strides),
+        const_param(b_batch_strides),
+        const_param(c_batch_strides),
+        batch_stride,
+        ndim,
+        batch_count);
+  }
+
+  // Run matmul
+  encoder.set_input_array(pointers);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(c);
+  encoder.set_output_array(out);
+
+  auto a_pointers = gpu_ptr<int8_t*>(pointers);
+  auto b_pointers = a_pointers + batch_count;
+  auto c_pointers = b_pointers + batch_count;
+  auto out_pointers = c_pointers + batch_count;
+  execute(
+      encoder,
+      reinterpret_cast<void*>(out_pointers),
+      reinterpret_cast<void*>(a_pointers),
+      reinterpret_cast<void*>(b_pointers),
+      reinterpret_cast<void*>(c_pointers),
+      alpha,
+      beta);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/gemv.cu b/macos/third_party/mlx/mlx/backend/cuda/gemms/gemv.cu
new file mode 100644
index 0000000..daa2752
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/gemv.cu
@@ -0,0 +1,327 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/gemms/gemv.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+static constexpr int rows_per_block = 8;
+
+// Accumulator type selection per input element type T.
+template <typename T>
+struct GemvAccType {
+  using type = T;
+};
+
+template <>
+struct GemvAccType<__half> {
+  using type = float;
+};
+
+template <>
+struct GemvAccType<__nv_bfloat16> {
+  using type = float;
+};
+
+template <>
+struct GemvAccType<float> {
+  using type = float;
+};
+
+template <>
+struct GemvAccType<double> {
+  using type = double;
+};
+
+template <>
+struct GemvAccType<cu::complex64_t> {
+  using type = cu::complex64_t;
+};
+
+template <typename T, int rows_per_block, int n_per_thread>
+__device__ void
+gemv_impl(const T* mat, const T* vec, T* out, int rows, int cols) {
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  auto g_idx = block.group_index();
+  auto t_idx = block.thread_index();
+  int row = g_idx.x * rows_per_block + t_idx.y;
+
+  if (row < rows) {
+    using Acc = typename GemvAccType<T>::type;
+    Acc sum = Acc(0);
+    for (int col = n_per_thread * warp.thread_rank(); col < cols;
+         col += (WARP_SIZE * n_per_thread)) {
+      auto local_mat =
+          unsafe_load_vector<n_per_thread>(mat + row * cols + col, 0);
+      auto local_vec = unsafe_load_vector<n_per_thread>(vec + col, 0);
+#pragma unroll
+      for (int j = 0; j < n_per_thread; ++j) {
+        sum += static_cast<Acc>(local_mat[j]) * static_cast<Acc>(local_vec[j]);
+      }
+    }
+
+    sum = cg::reduce(warp, sum, cg::plus<Acc>{});
+    if (warp.thread_rank() == 0) {
+      out[row] = static_cast<T>(sum);
+    }
+  }
+}
+
+template <typename T, int rows_per_block, int n_per_thread>
+__global__ void
+gemv_single(const T* mat, const T* vec, T* out, int rows, int cols) {
+  gemv_impl<T, rows_per_block, n_per_thread>(mat, vec, out, rows, cols);
+}
+
+template <typename T, int rows_per_block, int n_per_thread>
+__global__ void gemv_batched(
+    const T* mat,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols,
+    const __grid_constant__ Shape batch_shape,
+    const __grid_constant__ Strides mat_batch_strides,
+    const __grid_constant__ Strides vec_batch_strides,
+    int batch_ndim) {
+  auto block = cg::this_thread_block();
+  auto batch_idx = block.group_index().y;
+  auto [vec_offset, mat_offset] = elem_to_loc(
+      batch_idx,
+      batch_shape.data(),
+      vec_batch_strides.data(),
+      mat_batch_strides.data(),
+      batch_ndim);
+  gemv_impl<T, rows_per_block, n_per_thread>(
+      mat + mat_offset, vec + vec_offset, out + batch_idx * rows, rows, cols);
+}
+
+template <typename T, int rows_per_block, int n_per_thread>
+__global__ void gemv_gather(
+    const T* mat,
+    const T* vec,
+    T* out,
+    uint32_t* mat_indices,
+    uint32_t* vec_indices,
+    int rows,
+    int cols,
+    const __grid_constant__ Shape mat_batch_shape,
+    const __grid_constant__ Strides mat_batch_strides,
+    int mat_batch_ndim,
+    const __grid_constant__ Shape vec_batch_shape,
+    const __grid_constant__ Strides vec_batch_strides,
+    int vec_batch_ndim,
+    const __grid_constant__ Shape index_shape,
+    const __grid_constant__ Strides mat_index_strides,
+    const __grid_constant__ Strides vec_index_strides,
+    int index_batch_ndim) {
+  auto block = cg::this_thread_block();
+  auto indices_idx = block.group_index().y;
+  uint32_t index_mat, index_vec;
+  if (index_batch_ndim > 1) {
+    auto [mat_idx_offset, vec_idx_offset] = elem_to_loc(
+        indices_idx,
+        index_shape.data(),
+        mat_index_strides.data(),
+        vec_index_strides.data(),
+        index_batch_ndim);
+    index_mat = mat_indices[mat_idx_offset];
+    index_vec = vec_indices[vec_idx_offset];
+  } else {
+    index_mat = mat_indices[indices_idx * mat_index_strides[0]];
+    index_vec = vec_indices[indices_idx * vec_index_strides[0]];
+  }
+
+  int64_t mat_offset;
+  if (mat_batch_ndim > 1) {
+    mat_offset = elem_to_loc(
+        index_mat,
+        mat_batch_shape.data(),
+        mat_batch_strides.data(),
+        mat_batch_ndim);
+  } else {
+    mat_offset = index_mat * mat_batch_strides[0];
+  }
+
+  int64_t vec_offset;
+  if (vec_batch_ndim > 1) {
+    vec_offset = elem_to_loc(
+        index_vec,
+        vec_batch_shape.data(),
+        vec_batch_strides.data(),
+        vec_batch_ndim);
+  } else {
+    vec_offset = index_vec * vec_batch_strides[0];
+  }
+
+  gemv_impl<T, rows_per_block, n_per_thread>(
+      mat + mat_offset, vec + vec_offset, out + indices_idx * rows, rows, cols);
+}
+
+bool can_use_gemv(int M, int N, int K, bool a_transposed, bool b_transposed) {
+  return K % 32 == 0 && ((M == 1 && b_transposed) || (N == 1 && !a_transposed));
+}
+
+template <typename F>
+void dispatch_n_per_thread(int n_per_thread, F&& f) {
+  switch (n_per_thread) {
+    case 1:
+      f(std::integral_constant<int, 1>{});
+      break;
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 4:
+      f(std::integral_constant<int, 4>{});
+      break;
+  }
+}
+
+void gemv(
+    const array& a,
+    const array& b,
+    array& out,
+    int M,
+    int N,
+    int K,
+    uint32_t batch_count,
+    const mlx::core::Shape& batch_shape,
+    const mlx::core::Strides& a_batch_strides,
+    const mlx::core::Strides& b_batch_strides,
+    CommandEncoder& encoder) {
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  dispatch_inexact_types(out.dtype(), "gemv", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    dim3 block_dims{WARP_SIZE, rows_per_block};
+    const DataType* mat;
+    const DataType* vec;
+    int rows;
+    int cols = K;
+    auto mat_strides = const_param(a_batch_strides);
+    auto vec_strides = const_param(b_batch_strides);
+
+    if (M == 1) {
+      mat = gpu_ptr<DataType>(b);
+      vec = gpu_ptr<DataType>(a);
+      rows = N;
+      std::swap(mat_strides, vec_strides);
+    } else {
+      mat = gpu_ptr<DataType>(a);
+      vec = gpu_ptr<DataType>(b);
+      rows = M;
+    }
+    uint32_t num_blocks_x = (rows + rows_per_block - 1) / rows_per_block;
+    int n_per_t;
+    if (K % 128 == 0 && is_aligned<4>(mat) && is_aligned<4>(vec)) {
+      n_per_t = 4;
+    } else if (K % 64 == 0 && is_aligned<2>(mat) && is_aligned<2>(vec)) {
+      n_per_t = 2;
+    } else {
+      n_per_t = 1;
+    }
+    dispatch_n_per_thread(n_per_t, [&](auto n_per_thread) {
+      if (batch_count == 1) {
+        auto kernel = gemv_single<DataType, rows_per_block, n_per_thread()>;
+        encoder.add_kernel_node(
+            kernel,
+            num_blocks_x,
+            block_dims,
+            0,
+            mat,
+            vec,
+            gpu_ptr<DataType>(out),
+            rows,
+            cols);
+      } else {
+        auto kernel = gemv_batched<DataType, rows_per_block, n_per_thread()>;
+        encoder.add_kernel_node(
+            kernel,
+            dim3{num_blocks_x, batch_count},
+            block_dims,
+            0,
+            mat,
+            vec,
+            gpu_ptr<DataType>(out),
+            rows,
+            cols,
+            const_param(batch_shape),
+            mat_strides,
+            vec_strides,
+            batch_shape.size());
+      }
+    });
+  });
+}
+
+void gather_mv(
+    const array& mat_,
+    const array& vec_,
+    const array& mat_indices,
+    const array& vec_indices,
+    array& out,
+    int N,
+    int K,
+    CommandEncoder& encoder) {
+  encoder.set_input_array(mat_);
+  encoder.set_input_array(vec_);
+  encoder.set_input_array(mat_indices);
+  encoder.set_input_array(vec_indices);
+  encoder.set_output_array(out);
+  dispatch_inexact_types(out.dtype(), "gather_mv", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    dim3 block_dims{WARP_SIZE, rows_per_block};
+    int rows = N;
+    int cols = K;
+    uint32_t batch_size = static_cast<uint32_t>(out.size() / N);
+    const DataType* mat = gpu_ptr<DataType>(mat_);
+    const DataType* vec = gpu_ptr<DataType>(vec_);
+
+    uint32_t num_blocks_x = (rows + rows_per_block - 1) / rows_per_block;
+    int n_per_t;
+    if (K % 128 == 0 && is_aligned<4>(mat) && is_aligned<4>(vec)) {
+      n_per_t = 4;
+    } else if (K % 64 == 0 && is_aligned<2>(mat) && is_aligned<2>(vec)) {
+      n_per_t = 2;
+    } else {
+      n_per_t = 1;
+    }
+
+    dispatch_n_per_thread(n_per_t, [&](auto n_per_thread) {
+      auto kernel = gemv_gather<DataType, rows_per_block, n_per_thread()>;
+      encoder.add_kernel_node(
+          kernel,
+          dim3{num_blocks_x, batch_size},
+          block_dims,
+          0,
+          mat,
+          vec,
+          gpu_ptr<DataType>(out),
+          gpu_ptr<uint32_t>(mat_indices),
+          gpu_ptr<uint32_t>(vec_indices),
+          rows,
+          cols,
+          const_param(mat_.shape()),
+          const_param(mat_.strides()),
+          mat_.ndim() - 2,
+          const_param(vec_.shape()),
+          const_param(vec_.strides()),
+          vec_.ndim() - 2,
+          const_param(mat_indices.shape()),
+          const_param(mat_indices.strides()),
+          const_param(vec_indices.strides()),
+          mat_indices.ndim());
+    });
+  });
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/gemv.h b/macos/third_party/mlx/mlx/backend/cuda/gemms/gemv.h
new file mode 100644
index 0000000..b0dd49d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/gemv.h
@@ -0,0 +1,34 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core::cu {
+
+bool can_use_gemv(int M, int N, int K, bool a_transposed, bool b_transposed);
+
+void gemv(
+    const array& a,
+    const array& b,
+    array& out,
+    int M,
+    int N,
+    int K,
+    uint32_t batch_count,
+    const mlx::core::Shape& batch_shape,
+    const mlx::core::Strides& a_batch_strides,
+    const mlx::core::Strides& b_batch_strides,
+    CommandEncoder& encoder);
+
+void gather_mv(
+    const array& mat,
+    const array& vec,
+    const array& mat_indices,
+    const array& vec_indices,
+    array& out,
+    int N,
+    int K,
+    CommandEncoder& encoder);
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/grouped_gemm.h b/macos/third_party/mlx/mlx/backend/cuda/gemms/grouped_gemm.h
new file mode 100644
index 0000000..308a1ba
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/grouped_gemm.h
@@ -0,0 +1,25 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+namespace mlx::core {
+
+namespace cu {
+class CommandEncoder;
+}
+
+class array;
+
+void cutlass_grouped_gemm_unaligned(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int group_count,
+    const array& a,
+    const array& b,
+    const array& indices,
+    array& out,
+    cu::CommandEncoder& encoder);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu b/macos/third_party/mlx/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu
new file mode 100644
index 0000000..0c1bb9f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu
@@ -0,0 +1,358 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/cutlass_utils.cuh"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cutlass/gemm/device/default_gemm_configuration.h>
+#include <cutlass/gemm/device/gemm_grouped.h>
+#include <cutlass/gemm/kernel/default_gemm_grouped.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+using ProblemSize = cutlass::gemm::GemmCoord;
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <int N_READS>
+__global__ void prepare_grouped_mm_data(
+    const uint32_t* indices,
+    size_t size,
+    int group_count,
+    int K,
+    int N,
+    int lda,
+    int ldb,
+    int item_size,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* out_start,
+    int a_batch_stride,
+    int b_batch_stride,
+    int out_batch_stride,
+    ProblemSize* problem_sizes,
+    int64_t* a_lds,
+    int64_t* b_lds,
+    int64_t* out_lds,
+    void** a_ptrs,
+    void** b_ptrs,
+    void** out_ptrs) {
+  auto block = cg::this_thread_block();
+
+  // cumsum(histogram(indices)) - offset for each group.
+  extern __shared__ uint32_t cum_histo[];
+
+  int group = block.thread_rank();
+  if (group < group_count) {
+    cum_histo[group] = 0;
+  }
+
+  block.sync();
+
+  // Since |indices| is sorted, the position where element changes would be its
+  // cumulative histogram.
+  size_t elems_per_block = block.num_threads() * N_READS;
+  for (int r = 0; r < cuda::ceil_div(size, elems_per_block); ++r) {
+    // TODO: Use vectorized read.
+    for (int i = 0; i < N_READS; ++i) {
+      size_t pos = r * elems_per_block + group * N_READS + i;
+      if (pos >= size) {
+        break;
+      }
+      auto elem = indices[pos];
+      auto next = pos < size - 1 ? indices[pos + 1] : group_count;
+      while (elem < next) {
+        cum_histo[elem] = pos + 1;
+        elem++;
+      }
+    }
+  }
+
+  block.sync();
+
+  if (group < group_count) {
+    // Fill shapes.
+    int delta =
+        group == 0 ? cum_histo[0] : cum_histo[group] - cum_histo[group - 1];
+    problem_sizes[group] = {delta, N, K};
+    a_lds[group] = lda;
+    b_lds[group] = ldb;
+    out_lds[group] = N;
+    // Fill pointers.
+    auto offset = group == 0 ? 0 : cum_histo[group - 1];
+    a_ptrs[group] = a_start + offset * item_size * a_batch_stride;
+    b_ptrs[group] = b_start + group * item_size * b_batch_stride;
+    out_ptrs[group] = out_start + offset * item_size * out_batch_stride;
+  }
+}
+
+} // namespace cu
+
+namespace {
+
+// Shared GEMM configuration for every type and arch.
+template <typename T, typename ArchTag, int kAlignmentC>
+struct CommonGemmConfiguration {
+  using Element = T;
+  using Arch = ArchTag;
+  using Accumulator = std::conditional_t<(sizeof(T) < 4), float, T>;
+  using EpilogueOutputOp = cutlass::epilogue::thread::
+      LinearCombination<T, kAlignmentC, Accumulator, Accumulator>;
+};
+
+// Slow GEMM configuration as fallback.
+template <
+    typename T,
+    typename Arch,
+    int kAlignmentC = 1,
+    bool kEnableTF32 = false,
+    typename Enable = void>
+struct GemmConfiguration : public CommonGemmConfiguration<T, Arch, 1> {
+  using OpClass = cutlass::arch::OpClassSimt;
+  using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 8>;
+  using WarpShape = cutlass::gemm::GemmShape<32, 64, 8>;
+  using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>;
+  static const int kAlignmentAB = 1;
+  static const int kStages = 2;
+};
+
+// Specialized GEMM configuration for sm80 and later.
+template <typename T, typename Arch, int kAlignmentC>
+struct GemmConfiguration<
+    T,
+    Arch,
+    kAlignmentC,
+    true,
+    std::enable_if_t<Arch::kMinComputeCapability >= 80 && sizeof(T) <= 4>>
+    : public CommonGemmConfiguration<T, cutlass::arch::Sm80, kAlignmentC> {
+  using OpClass = cutlass::arch::OpClassTensorOp;
+  using ThreadblockShape = cutlass::gemm::GemmShape<256, 128, 32>;
+  using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32 / sizeof(T)>;
+  static const int kAlignmentAB = 1;
+  static const int kStages = 2;
+};
+
+// Specialized GEMM configuration for tf32 on sm80.
+template <int kAlignmentC>
+struct GemmConfiguration<float, cutlass::arch::Sm80, kAlignmentC, true>
+    : public CommonGemmConfiguration<float, cutlass::arch::Sm80, kAlignmentC> {
+  using OpClass = cutlass::arch::OpClassTensorOp;
+  using ThreadblockShape = cutlass::gemm::GemmShape<256, 128, 32>;
+  using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static const int kAlignmentAB = 1;
+  static const int kStages = 3; // use SM80_CP_ASYNC
+};
+
+// Get direct access to kernel.
+template <typename GemmKernel>
+class GemmGroupedEncoder
+    : public cutlass::gemm::device::GemmGrouped<GemmKernel> {
+ public:
+  void encode(cu::CommandEncoder& encoder) {
+    encoder.add_kernel_node(
+        cutlass::Kernel<GemmKernel>,
+        {static_cast<uint32_t>(this->params_.threadblock_count), 1, 1},
+        {GemmKernel::kThreadCount, 1, 1},
+        sizeof(typename GemmKernel::SharedStorage),
+        this->params_);
+  }
+};
+
+// Invoke the grouped GEMM of CUTLASS 2.x API, which supports small alignments.
+template <typename GemmConfiguration>
+void grouped_gemm_v2(
+    bool a_transposed,
+    bool b_transposed,
+    int group_count,
+    ProblemSize* problem_sizes,
+    int64_t* a_lds,
+    int64_t* b_lds,
+    int64_t* out_lds,
+    void* a_ptrs,
+    void* b_ptrs,
+    void* out_ptrs,
+    cu::CommandEncoder& encoder) {
+  dispatch_bool(a_transposed, [&](auto a_transposed_tag) {
+    dispatch_bool(b_transposed, [&](auto b_transposed_tag) {
+      using LayoutA = std::conditional_t<
+          a_transposed_tag.value,
+          cutlass::layout::ColumnMajor,
+          cutlass::layout::RowMajor>;
+      using LayoutB = std::conditional_t<
+          b_transposed_tag.value,
+          cutlass::layout::ColumnMajor,
+          cutlass::layout::RowMajor>;
+      using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped<
+          typename GemmConfiguration::Element,
+          LayoutA,
+          cutlass::ComplexTransform::kNone,
+          GemmConfiguration::kAlignmentAB,
+          typename GemmConfiguration::Element,
+          LayoutB,
+          cutlass::ComplexTransform::kNone,
+          GemmConfiguration::kAlignmentAB,
+          typename GemmConfiguration::Element,
+          cutlass::layout::RowMajor,
+          typename GemmConfiguration::Accumulator,
+          typename GemmConfiguration::OpClass,
+          typename GemmConfiguration::Arch,
+          typename GemmConfiguration::ThreadblockShape,
+          typename GemmConfiguration::WarpShape,
+          typename GemmConfiguration::InstructionShape,
+          typename GemmConfiguration::EpilogueOutputOp,
+          cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
+          GemmConfiguration::kStages>::GemmKernel;
+      using GemmGrouped = GemmGroupedEncoder<GemmKernel>;
+
+      static int threadblock_count = GemmGrouped::sufficient();
+      typename GemmGrouped::Arguments args(
+          problem_sizes,
+          group_count,
+          threadblock_count,
+          {/* alpha */ 1, /* beta */ 0},
+          reinterpret_cast<typename GemmGrouped::ElementA**>(a_ptrs),
+          reinterpret_cast<typename GemmGrouped::ElementB**>(b_ptrs),
+          reinterpret_cast<typename GemmGrouped::ElementC**>(out_ptrs),
+          reinterpret_cast<typename GemmGrouped::ElementC**>(out_ptrs),
+          a_lds,
+          b_lds,
+          out_lds,
+          out_lds);
+
+      GemmGrouped gemm;
+      CHECK_CUTLASS_ERROR(gemm.initialize(
+          args,
+          allocate_workspace(encoder, gemm.get_workspace_size(args)),
+          encoder.stream()));
+      gemm.encode(encoder);
+    });
+  });
+}
+
+template <typename F>
+void dispatch_cutlass_arch(cu::Device& device, F&& f) {
+  if (device.compute_capability_major() < 8) {
+    f(type_identity<cutlass::arch::Sm75>{});
+  } else if (device.compute_capability_major() == 8) {
+    f(type_identity<cutlass::arch::Sm80>{});
+  } else {
+    f(type_identity<cutlass::arch::Sm90>{});
+  }
+}
+
+auto* get_grouped_mm_funcion(Dtype dtype, int N, cu::Device& device) {
+  auto* fun = grouped_gemm_v2<GemmConfiguration<float, cutlass::arch::Sm75>>;
+  dispatch_float_types(dtype, "grouped_gemm_v2", [&](auto type_tag) {
+    using DataType = cutlass_type_t<MLX_GET_TYPE(type_tag)>;
+    dispatch_cutlass_arch(device, [&](auto arch_tag) {
+      using Arch = MLX_GET_TYPE(arch_tag);
+      dispatch_bool(N % 8 == 0, [&](auto is_out_aligned) {
+        constexpr int kAlignmentC = is_out_aligned ? 8 : 1;
+        dispatch_bool(env::enable_tf32(), [&](auto kEnableTF32) {
+          fun = grouped_gemm_v2<
+              GemmConfiguration<DataType, Arch, kAlignmentC, kEnableTF32>>;
+        });
+      });
+    });
+  });
+  return fun;
+}
+
+} // namespace
+
+void cutlass_grouped_gemm_unaligned(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int group_count,
+    const array& a,
+    const array& b,
+    const array& indices,
+    array& out,
+    cu::CommandEncoder& encoder) {
+  int K = a.shape(-1);
+  int N = b.shape(-1);
+
+  // Prepare device pointers for matmul.
+  int problem_sizes_nbytes =
+      group_count * cuda::ceil_div(sizeof(ProblemSize), 8) * 8;
+  int nbytes = problem_sizes_nbytes +
+      group_count * (3 * sizeof(void*) + 3 * sizeof(int64_t));
+  nbytes = cuda::ceil_div(nbytes, 256) * 256;
+  array gemm_args(cu::malloc_async(nbytes, encoder), {nbytes}, int8);
+  encoder.add_temporary(gemm_args);
+
+  ProblemSize* problem_sizes = gpu_ptr<ProblemSize>(gemm_args);
+  int64_t* a_lds = gpu_ptr<int64_t>(gemm_args) + problem_sizes_nbytes / 8;
+  int64_t* b_lds = a_lds + group_count;
+  int64_t* out_lds = b_lds + group_count;
+  void** a_ptrs = reinterpret_cast<void**>(out_lds + group_count);
+  void** b_ptrs = a_ptrs + group_count;
+  void** out_ptrs = b_ptrs + group_count;
+
+  // Fill the pointers by computing offsets from indices.
+  constexpr int N_READS = 4;
+  int n_threads = cuda::ceil_div(indices.size(), N_READS);
+  n_threads = group_count < n_threads ? n_threads : group_count;
+  dim3 block_dims(std::min(n_threads, 1024));
+  dim3 num_blocks(1);
+
+  encoder.set_input_array(indices);
+  encoder.set_output_array(gemm_args);
+  encoder.add_kernel_node(
+      cu::prepare_grouped_mm_data<N_READS>,
+      num_blocks,
+      block_dims,
+      group_count * sizeof(uint32_t), // sizeof(cum_histo)
+      gpu_ptr<uint32_t>(indices),
+      indices.size(),
+      group_count,
+      K,
+      N,
+      lda,
+      ldb,
+      out.itemsize(),
+      gpu_ptr<int8_t>(a),
+      gpu_ptr<int8_t>(b),
+      gpu_ptr<int8_t>(out),
+      a.shape(-2) * a.shape(-1), // a_batch_stride
+      b.shape(-2) * b.shape(-1), // b_batch_stride
+      out.shape(-2) * out.shape(-1), // out_batch_stride
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs);
+
+  // Invoke grouped GEMM.
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(gemm_args);
+  encoder.set_output_array(out);
+  auto* fun = get_grouped_mm_funcion(a.dtype(), N, encoder.device());
+  fun(a_transposed,
+      b_transposed,
+      group_count,
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      encoder);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/indexing.cpp b/macos/third_party/mlx/mlx/backend/cuda/indexing.cpp
new file mode 100644
index 0000000..f855cd8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/indexing.cpp
@@ -0,0 +1,434 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include "cuda_jit_sources.h"
+
+#include <cuda.h>
+#include <fmt/format.h>
+#include <nvrtc.h>
+#include <nvtx3/nvtx3.hpp>
+
+#include <cassert>
+#include <numeric>
+
+namespace mlx::core {
+
+namespace {
+
+constexpr const char* g_scatter_ops[] = {"Max", "Min", "Sum", "Prod", "Assign"};
+
+void append_indices_arg(
+    cu::KernelArgs& args,
+    const std::vector<array>& inputs,
+    int nidx,
+    int idx_ndim) {
+  SmallVector<const void*> indices(nidx);
+  for (int i = 0; i < nidx; ++i) {
+    indices[i] = gpu_ptr<void>(inputs[i + 1]);
+  }
+  args.append(std::move(indices));
+  SmallVector<int32_t> indices_shape(nidx * idx_ndim);
+  for (int i = 0; i < nidx; ++i) {
+    std::copy_n(
+        inputs[i + 1].shape().begin(),
+        idx_ndim,
+        indices_shape.data() + i * idx_ndim);
+  }
+  args.append(std::move(indices_shape));
+  SmallVector<int64_t> indices_strides(nidx * idx_ndim);
+  for (int i = 0; i < nidx; ++i) {
+    std::copy_n(
+        inputs[i + 1].strides().begin(),
+        idx_ndim,
+        indices_strides.data() + i * idx_ndim);
+  }
+  args.append(std::move(indices_strides));
+}
+
+} // namespace
+
+void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Gather::eval_gpu");
+  assert(inputs.size() > 0);
+  const auto& src = inputs[0];
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  if (out.size() == 0) {
+    return;
+  }
+
+  int nidx = inputs.size() - 1;
+  Dtype idx_dtype = nidx > 0 ? inputs[1].dtype() : int32;
+  int32_t idx_ndim = nidx > 0 ? inputs[1].ndim() : 0;
+
+  bool large = (nidx > 0 && inputs[1].size() > INT32_MAX) ||
+      (src.size() > INT32_MAX) || (out.size() > INT32_MAX);
+
+  uint32_t slice_size = std::accumulate(
+      slice_sizes_.begin(), slice_sizes_.end(), 1, std::multiplies<uint32_t>());
+
+  std::string module_name = fmt::format(
+      "gather_{}_{}_{}",
+      dtype_to_string(out.dtype()),
+      dtype_to_string(idx_dtype),
+      nidx);
+
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names;
+    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
+      for (int large = 0; large <= 1; ++large) {
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::gather<{}, {}, {}, {}, {}>",
+                dtype_to_cuda_type(out.dtype()),
+                dtype_to_cuda_type(idx_dtype),
+                nidx,
+                ndim,
+                large ? "int64_t" : "int32_t"));
+      }
+    }
+    return std::make_tuple(false, jit_source_gather, std::move(kernel_names));
+  });
+
+  cu::KernelArgs args;
+  args.append(src);
+  args.append(out);
+  if (large) {
+    args.append<int64_t>(out.size());
+  } else {
+    args.append<int32_t>(out.size());
+  }
+  args.append_ndim(src.shape());
+  args.append_ndim(src.strides());
+  args.append<int32_t>(src.ndim());
+  args.append_ndim(slice_sizes_);
+  args.append(slice_size);
+  args.append(axes_);
+  append_indices_arg(args, inputs, nidx, idx_ndim);
+
+  std::string kernel_name = fmt::format(
+      "mlx::core::cu::gather<{}, {}, {}, {}, {}>",
+      dtype_to_cuda_type(out.dtype()),
+      dtype_to_cuda_type(idx_dtype),
+      nidx,
+      idx_ndim,
+      large ? "int64_t" : "int32_t");
+
+  for (const auto& in : inputs) {
+    encoder.set_input_array(in);
+  }
+  encoder.set_output_array(out);
+
+  auto kernel = mod.get_kernel(kernel_name);
+  auto [num_blocks, block_dims] = get_launch_args(out, large);
+  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+}
+
+void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Gather::eval_gpu");
+  assert(inputs.size() > 1);
+  auto& upd = inputs.back();
+
+  // Copy src into out.
+  CopyType copy_type;
+  if (inputs[0].data_size() == 1) {
+    copy_type = CopyType::Scalar;
+  } else if (inputs[0].flags().row_contiguous) {
+    copy_type = CopyType::Vector;
+  } else {
+    copy_type = CopyType::General;
+  }
+  copy_gpu(inputs[0], out, copy_type);
+
+  // Empty update.
+  if (upd.size() == 0) {
+    return;
+  }
+
+  int nidx = axes_.size();
+  Dtype idx_dtype = nidx > 0 ? inputs[1].dtype() : int32;
+  int32_t idx_ndim = nidx > 0 ? inputs[1].ndim() : 0;
+
+  bool large = (nidx > 0 && inputs[1].size() > INT32_MAX) ||
+      (upd.size() > INT32_MAX) || (out.size() > INT32_MAX);
+
+  int32_t upd_post_idx_size = std::accumulate(
+      upd.shape().begin() + idx_ndim,
+      upd.shape().end(),
+      1,
+      std::multiplies<int32_t>());
+
+  const char* op = g_scatter_ops[reduce_type_];
+  std::string module_name = fmt::format(
+      "scatter_{}_{}_{}_{}",
+      dtype_to_string(out.dtype()),
+      dtype_to_string(idx_dtype),
+      op,
+      nidx);
+
+  auto& s = stream();
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names;
+    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
+      for (int large = 0; large <= 1; ++large) {
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>",
+                dtype_to_cuda_type(out.dtype()),
+                dtype_to_cuda_type(idx_dtype),
+                op,
+                nidx,
+                ndim,
+                large ? "int64_t" : "int32_t"));
+      }
+    }
+    return std::make_tuple(false, jit_source_scatter, std::move(kernel_names));
+  });
+
+  cu::KernelArgs args;
+  args.append(upd);
+  args.append(out);
+  if (large) {
+    args.append<int64_t>(upd.size());
+  } else {
+    args.append<int32_t>(upd.size());
+  }
+  args.append_ndim(upd.shape());
+  args.append_ndim(upd.strides());
+  args.append<int32_t>(upd.ndim());
+  if (large) {
+    args.append<int64_t>(upd_post_idx_size);
+  } else {
+    args.append<int32_t>(upd_post_idx_size);
+  }
+  args.append_ndim(out.shape());
+  args.append_ndim(out.strides());
+  args.append<int32_t>(out.ndim());
+  args.append(axes_);
+  append_indices_arg(args, inputs, nidx, idx_ndim);
+
+  std::string kernel_name = fmt::format(
+      "mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>",
+      dtype_to_cuda_type(out.dtype()),
+      dtype_to_cuda_type(idx_dtype),
+      op,
+      nidx,
+      idx_ndim,
+      large ? "int64_t" : "int32_t");
+
+  auto& encoder = cu::get_command_encoder(s);
+  for (const auto& in : inputs) {
+    encoder.set_input_array(in);
+  }
+  encoder.set_output_array(out);
+  auto kernel = mod.get_kernel(kernel_name);
+  auto [num_blocks, block_dims] = get_launch_args(upd, large);
+  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+}
+
+void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("GatherAxis::eval_gpu");
+  assert(inputs.size() > 1);
+  const auto& src = inputs[0];
+  const auto& idx = inputs[1];
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  if (out.size() == 0) {
+    return;
+  }
+
+  bool large = idx.size() > INT32_MAX || src.size() > INT32_MAX;
+
+  std::string module_name = fmt::format(
+      "gather_axis_{}_{}",
+      dtype_to_string(out.dtype()),
+      dtype_to_string(idx.dtype()));
+
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names;
+    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
+      for (int contiguous = 0; contiguous < 4; ++contiguous) {
+        for (int large = 0; large <= 1; ++large) {
+          kernel_names.push_back(
+              fmt::format(
+                  "mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>",
+                  dtype_to_cuda_type(out.dtype()),
+                  dtype_to_cuda_type(idx.dtype()),
+                  ndim,
+                  contiguous & 1 ? true : false,
+                  contiguous & 2 ? true : false,
+                  large ? "int64_t" : "int32_t"));
+        }
+      }
+    }
+    return std::make_tuple(
+        false, jit_source_gather_axis, std::move(kernel_names));
+  });
+
+  size_t idx_size_pre = 1;
+  size_t idx_size_post = 1;
+  for (int i = 0; i < axis_; ++i) {
+    idx_size_pre *= idx.shape(i);
+  }
+  for (int i = axis_ + 1; i < idx.ndim(); ++i) {
+    idx_size_post *= idx.shape(i);
+  }
+  size_t idx_size_axis = idx.shape(axis_);
+
+  cu::KernelArgs args;
+  args.append(src);
+  args.append(idx);
+  args.append(out);
+  if (large) {
+    args.append<int64_t>(idx_size_pre);
+    args.append<int64_t>(idx_size_axis);
+    args.append<int64_t>(idx_size_post);
+  } else {
+    args.append<int32_t>(idx_size_pre);
+    args.append<int32_t>(idx_size_axis);
+    args.append<int32_t>(idx_size_post);
+  }
+  args.append(remove_index(idx.shape(), axis_));
+  args.append(remove_index(src.strides(), axis_));
+  args.append(remove_index(idx.strides(), axis_));
+  args.append<int32_t>(axis_);
+  args.append(src.shape(axis_));
+  args.append(src.strides(axis_));
+  args.append(idx.strides(axis_));
+
+  std::string kernel_name = fmt::format(
+      "mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>",
+      dtype_to_cuda_type(out.dtype()),
+      dtype_to_cuda_type(idx.dtype()),
+      src.ndim() - 1,
+      src.flags().row_contiguous,
+      idx.flags().row_contiguous,
+      large ? "int64_t" : "int32_t");
+
+  for (const auto& in : inputs) {
+    encoder.set_input_array(in);
+  }
+  encoder.set_output_array(out);
+  auto kernel = mod.get_kernel(kernel_name);
+  auto [num_blocks, block_dims] = get_launch_args(idx, large);
+  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+}
+
+void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("ScatterAxis::eval_gpu");
+  assert(inputs.size() > 2);
+  const auto& src = inputs[0];
+  const auto& idx = inputs[1];
+  const auto& upd = inputs[2];
+
+  // Copy src into out.
+  CopyType copy_type;
+  if (src.data_size() == 1) {
+    copy_type = CopyType::Scalar;
+  } else if (src.flags().row_contiguous) {
+    copy_type = CopyType::Vector;
+  } else {
+    copy_type = CopyType::General;
+  }
+  copy_gpu(src, out, copy_type);
+
+  // Empty update.
+  if (upd.size() == 0) {
+    return;
+  }
+
+  bool large = idx.size() > INT32_MAX || src.size() > INT32_MAX;
+
+  const char* op = reduce_type_ == ScatterAxis::Sum ? "Sum" : "Assign";
+  std::string module_name = fmt::format(
+      "scatter_axis_{}_{}_{}",
+      dtype_to_string(out.dtype()),
+      dtype_to_string(idx.dtype()),
+      op);
+
+  auto& s = stream();
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names;
+    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
+      for (int contiguous = 0; contiguous < 4; ++contiguous) {
+        for (int large = 0; large <= 1; ++large) {
+          kernel_names.push_back(
+              fmt::format(
+                  "mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>",
+                  dtype_to_cuda_type(out.dtype()),
+                  dtype_to_cuda_type(idx.dtype()),
+                  op,
+                  ndim,
+                  contiguous & 1 ? true : false,
+                  contiguous & 2 ? true : false,
+                  large ? "int64_t" : "int32_t"));
+        }
+      }
+    }
+    return std::make_tuple(
+        false, jit_source_scatter_axis, std::move(kernel_names));
+  });
+
+  size_t idx_size_pre = 1;
+  size_t idx_size_post = 1;
+  for (int i = 0; i < axis_; ++i) {
+    idx_size_pre *= idx.shape(i);
+  }
+  for (int i = axis_ + 1; i < idx.ndim(); ++i) {
+    idx_size_post *= idx.shape(i);
+  }
+  size_t idx_size_axis = idx.shape(axis_);
+
+  cu::KernelArgs args;
+  args.append(upd);
+  args.append(idx);
+  args.append(out);
+  if (large) {
+    args.append<int64_t>(idx_size_pre);
+    args.append<int64_t>(idx_size_axis);
+    args.append<int64_t>(idx_size_post);
+  } else {
+    args.append<int32_t>(idx_size_pre);
+    args.append<int32_t>(idx_size_axis);
+    args.append<int32_t>(idx_size_post);
+  }
+  args.append(remove_index(idx.shape(), axis_));
+  args.append(remove_index(upd.strides(), axis_));
+  args.append(remove_index(idx.strides(), axis_));
+  args.append<int32_t>(axis_);
+  args.append(out.shape(axis_));
+  args.append(upd.strides(axis_));
+  args.append(idx.strides(axis_));
+
+  std::string kernel_name = fmt::format(
+      "mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>",
+      dtype_to_cuda_type(out.dtype()),
+      dtype_to_cuda_type(idx.dtype()),
+      op,
+      idx.ndim() - 1,
+      upd.flags().row_contiguous,
+      idx.flags().row_contiguous,
+      large ? "int64_t" : "int32_t");
+
+  auto& encoder = cu::get_command_encoder(s);
+  for (const auto& in : inputs) {
+    encoder.set_input_array(in);
+  }
+  encoder.set_output_array(out);
+  auto kernel = mod.get_kernel(kernel_name);
+  auto [num_blocks, block_dims] = get_launch_args(idx, large);
+  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/jit_module.cpp b/macos/third_party/mlx/mlx/backend/cuda/jit_module.cpp
new file mode 100644
index 0000000..b0ebb40
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/jit_module.cpp
@@ -0,0 +1,443 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/version.h"
+
+#include "cuda_jit_sources.h"
+
+#include <cstdlib>
+#include <filesystem>
+#include <fstream>
+
+#include <fmt/format.h>
+#include <nvrtc.h>
+
+namespace mlx::core::cu {
+
+namespace {
+
+#define CHECK_NVRTC_ERROR(cmd) check_nvrtc_error(#cmd, (cmd))
+
+void check_nvrtc_error(const char* name, nvrtcResult err) {
+  if (err != NVRTC_SUCCESS) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}", name, nvrtcGetErrorString(err)));
+  }
+}
+
+// Return the --include-path args used for invoking NVRTC.
+const std::vector<std::string>& include_path_args() {
+  static std::vector<std::string> cached_args = []() {
+    std::vector<std::string> args;
+    // Add path to bundled CCCL headers.
+    auto root_dir = current_binary_dir();
+#if !defined(_WIN32)
+    root_dir = root_dir.parent_path();
+#endif
+    auto path = root_dir / "include" / "cccl";
+#if defined(MLX_CCCL_DIR)
+    if (!std::filesystem::exists(path)) {
+      path = MLX_CCCL_DIR;
+    }
+#endif
+    if (std::filesystem::exists(path)) {
+      args.push_back(fmt::format("--include-path={}", path.string()));
+    }
+    // Add path to CUDA runtime headers, try local-installed python package
+    // first and then system-installed headers.
+    path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include";
+    if (std::filesystem::exists(path)) {
+      args.push_back(fmt::format("--include-path={}", path.string()));
+    } else {
+      const char* home = std::getenv("CUDA_HOME");
+      if (!home) {
+        home = std::getenv("CUDA_PATH");
+      }
+#if defined(__linux__)
+      if (!home) {
+        home = "/usr/local/cuda";
+      }
+#endif
+      if (home && std::filesystem::exists(home)) {
+        args.push_back(fmt::format("--include-path={}/include", home));
+      } else {
+        throw std::runtime_error(
+            "Can not find locations of CUDA headers, please set environment "
+            "variable CUDA_HOME or CUDA_PATH.");
+      }
+    }
+    return args;
+  }();
+  return cached_args;
+}
+
+// Get the cache directory for storing compiled results.
+const std::filesystem::path& ptx_cache_dir() {
+  static std::filesystem::path cache = []() -> std::filesystem::path {
+    std::filesystem::path cache;
+    if (auto c = std::getenv("MLX_PTX_CACHE_DIR"); c) {
+      cache = c;
+    } else {
+      cache =
+          std::filesystem::temp_directory_path() / "mlx" / version() / "ptx";
+    }
+
+#if defined(_WIN32)
+    // Add "\\?\" prefix to support long file path.
+    const wchar_t* long_path_prefix = L"\\\\?\\";
+    if (cache.is_relative()) {
+      cache = std::filesystem::absolute(cache);
+    }
+    if (!cache.native().starts_with(long_path_prefix)) {
+      cache = long_path_prefix + cache.native();
+    }
+#endif
+
+    if (!std::filesystem::exists(cache)) {
+      std::error_code error;
+      if (!std::filesystem::create_directories(cache, error)) {
+        return std::filesystem::path();
+      }
+    }
+    return cache;
+  }();
+  return cache;
+}
+
+std::filesystem::path get_ptx_path(
+    const std::filesystem::path& cache_dir,
+    const std::string& module_name) {
+  constexpr int max_file_name_length = 245;
+  if (module_name.size() <= max_file_name_length) {
+    return cache_dir / (module_name + ".ptx");
+  }
+
+  auto ptx_path = cache_dir;
+  int offset = 0;
+  while (module_name.size() - offset > max_file_name_length) {
+    ptx_path /= module_name.substr(offset, max_file_name_length);
+    offset += max_file_name_length;
+  }
+  ptx_path /= module_name.substr(offset) + ".ptx";
+
+  return ptx_path;
+}
+
+// Try to read the cached |ptx| and |ptx_kernels| from |cache_dir|.
+bool read_cached_ptx(
+    const std::filesystem::path& cache_dir,
+    const std::string& module_name,
+    std::string& ptx,
+    std::vector<std::pair<std::string, std::string>>& ptx_kernels) {
+  if (cache_dir.empty()) {
+    return false;
+  }
+
+  auto ptx_path = get_ptx_path(cache_dir, module_name);
+  std::error_code error;
+  auto ptx_size = std::filesystem::file_size(ptx_path, error);
+  if (error) {
+    return false;
+  }
+  std::ifstream ptx_file(ptx_path, std::ios::binary);
+  if (!ptx_file.good()) {
+    return false;
+  }
+  ptx.resize(ptx_size);
+  ptx_file.read(ptx.data(), ptx_size);
+
+  std::ifstream txt_file(ptx_path.replace_extension(".txt"), std::ios::binary);
+  std::string line;
+  while (std::getline(txt_file, line)) {
+    auto tab = line.find('\t');
+    if (tab != std::string::npos) {
+      ptx_kernels.emplace_back(line.substr(0, tab), line.substr(tab + 1));
+    }
+  }
+  return true;
+}
+
+// Write the |ptx| and |ptx_kernels| to |cache_dir| with |name|.
+void write_cached_ptx(
+    const std::filesystem::path& cache_dir,
+    const std::string& module_name,
+    const std::string& ptx,
+    const std::vector<std::pair<std::string, std::string>>& ptx_kernels,
+    const std::string& source_code) {
+  if (cache_dir.empty()) {
+    return;
+  }
+
+  auto ptx_path = get_ptx_path(cache_dir, module_name);
+
+  // Ensure that the directory exists
+  auto parent = ptx_path.parent_path();
+  if (parent != cache_dir) {
+    std::filesystem::create_directories(parent);
+  }
+
+  // Write the compiled code and mangled names
+  std::ofstream ptx_file(ptx_path, std::ios::binary);
+  if (!ptx.empty()) {
+    ptx_file.write(&ptx.front(), ptx.size());
+  }
+  std::ofstream txt_file(ptx_path.replace_extension(".txt"), std::ios::binary);
+  for (const auto& [name, mangled] : ptx_kernels) {
+    txt_file << name << "\t" << mangled << std::endl;
+  }
+
+  // Write the generated code
+  std::ofstream source_file(ptx_path.replace_extension(".cu"));
+  source_file << source_code;
+}
+
+// Return if |device|'s version is not newer than |major|.|minor| version.
+inline bool version_lower_equal(Device& device, int major, int minor) {
+  if (device.compute_capability_major() < major) {
+    return true;
+  } else if (device.compute_capability_major() == major) {
+    return device.compute_capability_minor() <= minor;
+  } else {
+    return false;
+  }
+}
+
+// Return whether NVRTC supports compiling to |device|'s SASS code.
+bool compiler_supports_device_sass(Device& device) {
+  int nvrtc_major, nvrtc_minor;
+  CHECK_NVRTC_ERROR(nvrtcVersion(&nvrtc_major, &nvrtc_minor));
+  if (nvrtc_major < 9) {
+    return false;
+  } else if (nvrtc_major == 9) {
+    return version_lower_equal(device, 7, 2);
+  } else if (nvrtc_major == 10) {
+    return version_lower_equal(device, 7, 5);
+  } else if (nvrtc_major == 11 && nvrtc_minor == 0) {
+    return version_lower_equal(device, 8, 0);
+  } else if (nvrtc_major == 11 && nvrtc_minor < 8) {
+    return version_lower_equal(device, 8, 6);
+  } else {
+    return true;
+  }
+}
+
+#define INCLUDE_PREFIX "mlx/backend/cuda/device/"
+
+constexpr const char* g_include_names[] = {
+    INCLUDE_PREFIX "atomic_ops.cuh",
+    INCLUDE_PREFIX "binary_ops.cuh",
+    INCLUDE_PREFIX "cast_op.cuh",
+    INCLUDE_PREFIX "config.h",
+    INCLUDE_PREFIX "complex.cuh",
+    INCLUDE_PREFIX "fp16_math.cuh",
+    INCLUDE_PREFIX "indexing.cuh",
+    INCLUDE_PREFIX "scatter_ops.cuh",
+    INCLUDE_PREFIX "unary_ops.cuh",
+    INCLUDE_PREFIX "ternary_ops.cuh",
+    INCLUDE_PREFIX "utils.cuh",
+};
+
+#undef INCLUDE_PREFIX
+
+constexpr const char* g_headers[] = {
+    jit_source_atomic_ops,
+    jit_source_binary_ops,
+    jit_source_cast_op,
+    jit_source_config,
+    jit_source_complex,
+    jit_source_fp16_math,
+    jit_source_indexing,
+    jit_source_scatter_ops,
+    jit_source_unary_ops,
+    jit_source_ternary_ops,
+    jit_source_utils,
+};
+
+void compile(
+    Device& device,
+    const std::string& module_name,
+    const std::string& source,
+    const std::vector<std::string>& kernel_names,
+    std::string& ptx,
+    std::vector<std::pair<std::string, std::string>>& ptx_kernels) {
+  // Create the program
+  nvrtcProgram prog;
+  CHECK_NVRTC_ERROR(nvrtcCreateProgram(
+      &prog,
+      source.c_str(),
+      (module_name + ".cu").c_str(),
+      std::size(g_headers),
+      g_headers,
+      g_include_names));
+  std::unique_ptr<nvrtcProgram, void (*)(nvrtcProgram*)> prog_freer(
+      &prog,
+      [](nvrtcProgram* p) { CHECK_NVRTC_ERROR(nvrtcDestroyProgram(p)); });
+  for (const auto& name : kernel_names) {
+    CHECK_NVRTC_ERROR(nvrtcAddNameExpression(prog, name.c_str()));
+  }
+
+  // Compile program.
+  std::vector<const char*> args;
+  bool use_sass = compiler_supports_device_sass(device);
+  auto cc = device.compute_capability_major();
+  std::string arch_tag = (cc >= 9) ? "a" : "";
+  std::string compute = fmt::format(
+      "--gpu-architecture={}_{}{}{}",
+      use_sass ? "sm" : "compute",
+      cc,
+      device.compute_capability_minor(),
+      arch_tag);
+  args.push_back(compute.c_str());
+  for (const auto& include : include_path_args()) {
+    args.push_back(include.c_str());
+  }
+  nvrtcResult compile_result =
+      nvrtcCompileProgram(prog, args.size(), args.data());
+  if (compile_result != NVRTC_SUCCESS) {
+    size_t log_size;
+    CHECK_NVRTC_ERROR(nvrtcGetProgramLogSize(prog, &log_size));
+    std::vector<char> log(log_size + 1, 0);
+    CHECK_NVRTC_ERROR(nvrtcGetProgramLog(prog, log.data()));
+    throw std::runtime_error(
+        fmt::format("Failed to compile kernel: {}.", log.data()));
+  }
+
+  // Get mangled names of kernel names.
+  for (const auto& name : kernel_names) {
+    const char* mangled;
+    CHECK_NVRTC_ERROR(nvrtcGetLoweredName(prog, name.c_str(), &mangled));
+    ptx_kernels.emplace_back(name, mangled);
+  }
+
+  // Get ptx data.
+  size_t ptx_size;
+  if (use_sass) {
+    CHECK_NVRTC_ERROR(nvrtcGetCUBINSize(prog, &ptx_size));
+  } else {
+    CHECK_NVRTC_ERROR(nvrtcGetPTXSize(prog, &ptx_size));
+  }
+  ptx.resize(ptx_size);
+  if (use_sass) {
+    CHECK_NVRTC_ERROR(nvrtcGetCUBIN(prog, ptx.data()));
+  } else {
+    CHECK_NVRTC_ERROR(nvrtcGetPTX(prog, ptx.data()));
+  }
+}
+
+void load_module(
+    const std::string& module_name,
+    const std::string& ptx,
+    const std::vector<std::pair<std::string, std::string>>& ptx_kernels,
+    CUmodule& module_,
+    std::unordered_map<std::string, std::tuple<CUfunction, bool, uint32_t>>&
+        kernels) {
+  // Load module.
+  char jit_log[4089] = {};
+  CUjit_option options[] = {
+      CU_JIT_ERROR_LOG_BUFFER, CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES};
+  void* values[] = {jit_log, reinterpret_cast<void*>(std::size(jit_log) - 1)};
+  CUresult jit_result = cuModuleLoadDataEx(
+      &module_, ptx.data(), std::size(options), options, values);
+  if (jit_result != CUDA_SUCCESS) {
+    throw std::runtime_error(
+        fmt::format(
+            "Failed to load compiled {} kernel: {}.", module_name, jit_log));
+  }
+
+  // Load kernels.
+  for (const auto& [name, mangled] : ptx_kernels) {
+    CUfunction kernel;
+    CHECK_CUDA_ERROR(cuModuleGetFunction(&kernel, module_, mangled.c_str()));
+    kernels[name] = std::make_tuple(kernel, false, 0);
+  }
+}
+
+} // namespace
+
+JitModule::JitModule(
+    Device& device,
+    const std::string& module_name,
+    const KernelBuilder& builder,
+    bool use_disk_cache) {
+  // Will hold the actual device executable source code and kernel names
+  std::string ptx;
+  std::vector<std::pair<std::string, std::string>> ptx_kernels;
+
+  // Try to load them from the file cache
+  if (!read_cached_ptx(ptx_cache_dir(), module_name, ptx, ptx_kernels)) {
+    auto [precompiled, source_code, kernel_names] = builder();
+
+    // Get the PTX or cubin
+    if (precompiled) {
+      ptx = std::move(source_code);
+      for (auto& name : kernel_names) {
+        ptx_kernels.emplace_back(name, name);
+      }
+    } else {
+      compile(device, module_name, source_code, kernel_names, ptx, ptx_kernels);
+    }
+
+    // If requested save them in the file cache for the next launch
+    if (use_disk_cache) {
+      write_cached_ptx(
+          ptx_cache_dir(), module_name, ptx, ptx_kernels, source_code);
+    }
+  }
+
+  // Load the module
+  load_module(module_name, ptx, ptx_kernels, module_, kernels_);
+}
+
+JitModule::~JitModule() {
+  CHECK_CUDA_ERROR(cuModuleUnload(module_));
+}
+
+std::pair<CUfunction, uint32_t> JitModule::get_kernel_and_dims(
+    const std::string& kernel_name,
+    std::function<void(CUfunction)> configure_kernel) {
+  auto it = kernels_.find(kernel_name);
+  if (it == kernels_.end()) {
+    throw std::runtime_error(
+        fmt::format("There is no kernel named {}.", kernel_name));
+  }
+
+  // If it is the first time we run this kernel then configure it. Do it only
+  // once!
+  auto kernel = std::get<0>(it->second);
+  if (!std::get<1>(it->second)) {
+    if (configure_kernel) {
+      configure_kernel(kernel);
+    }
+    std::get<1>(it->second) = true;
+    std::get<2>(it->second) = max_occupancy_block_dim(kernel);
+  }
+
+  return {kernel, std::get<2>(it->second)};
+}
+
+CUfunction JitModule::get_kernel(
+    const std::string& kernel_name,
+    std::function<void(CUfunction)> configure_kernel) {
+  return get_kernel_and_dims(kernel_name, std::move(configure_kernel)).first;
+}
+
+std::unordered_map<std::string, JitModule>& get_jit_module_cache() {
+  static std::unordered_map<std::string, JitModule> map;
+  return map;
+}
+
+JitModule& get_jit_module(
+    const mlx::core::Device& device,
+    const std::string& name,
+    const KernelBuilder& builder,
+    bool cache) {
+  auto& map = get_jit_module_cache();
+  auto it = map.find(name);
+  if (it == map.end()) {
+    it = map.try_emplace(name, cu::device(device), name, builder, cache).first;
+  }
+  return it->second;
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/jit_module.h b/macos/third_party/mlx/mlx/backend/cuda/jit_module.h
new file mode 100644
index 0000000..3849d8d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/jit_module.h
@@ -0,0 +1,120 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/config.h"
+
+#include <deque>
+#include <unordered_map>
+#include <utility>
+#include <variant>
+
+#include <cuda.h>
+#include <fmt/format.h>
+
+namespace mlx::core::cu {
+
+class Device;
+
+using KernelBuilderResult = std::tuple<
+    /* precompiled */ bool,
+    /* source code */ std::string,
+    /* kernel names */ std::vector<std::string>>;
+using KernelBuilder = std::function<KernelBuilderResult()>;
+
+struct KernelArgs {
+  void** args() {
+    return args_.data();
+  }
+
+  void append(const array& a) {
+    append(reinterpret_cast<CUdeviceptr>(gpu_ptr<void>(a)));
+  }
+
+  template <typename T>
+  void append(T val) {
+    storage_.emplace_back(val);
+    append_ptr(&storage_.back());
+  }
+
+  template <typename T>
+  void append(SmallVector<T> vec) {
+    storage_.emplace_back(std::move(vec));
+    append_ptr(std::get<SmallVector<T>>(storage_.back()).data());
+  }
+
+  template <typename T>
+  void append(const std::vector<T>& vec) {
+    append(SmallVector<T>(vec.begin(), vec.end()));
+  }
+
+  // Make sure the arg is copied to an array with size of NDIM.
+  template <size_t NDIM = MAX_NDIM, typename T>
+  void append_ndim(SmallVector<T> vec) {
+    if (vec.size() > NDIM) {
+      throw std::runtime_error(
+          fmt::format("ndim can not be larger than {}.", NDIM));
+    }
+    vec.resize(NDIM);
+    append(std::move(vec));
+  }
+
+  void append_ptr(const void* v) {
+    args_.push_back(const_cast<void*>(v));
+  }
+
+ private:
+  std::vector<void*> args_;
+
+  // The cuGraphAddKernelNode API requires passing pointers to arguments so
+  // store temporary values until the node is created.
+  using Arg = std::variant<
+      std::monostate,
+      CUdeviceptr,
+      bool,
+      int32_t,
+      uint32_t,
+      int64_t,
+      float,
+      SmallVector<const void*>,
+      SmallVector<int32_t>,
+      SmallVector<int64_t>>;
+  std::deque<Arg> storage_;
+};
+
+class JitModule {
+ public:
+  JitModule(
+      Device& device,
+      const std::string& module_name,
+      const KernelBuilder& builder,
+      bool cache);
+  ~JitModule();
+
+  JitModule(const JitModule&) = delete;
+  JitModule& operator=(const JitModule&) = delete;
+  CUfunction get_kernel(
+      const std::string& kernel_name,
+      std::function<void(CUfunction)> configure_kernel = nullptr);
+  std::pair<CUfunction, uint32_t> get_kernel_and_dims(
+      const std::string& kernel_name,
+      std::function<void(CUfunction)> configure_kernel = nullptr);
+
+ private:
+  CUmodule module_{nullptr};
+  std::unordered_map<std::string, std::tuple<CUfunction, bool, uint32_t>>
+      kernels_;
+};
+
+std::unordered_map<std::string, JitModule>& get_jit_module_cache();
+
+JitModule& get_jit_module(
+    const mlx::core::Device& device,
+    const std::string& name,
+    const KernelBuilder& builder,
+    bool use_disk_cache = true);
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/kernel_utils.cu b/macos/third_party/mlx/mlx/backend/cuda/kernel_utils.cu
new file mode 100644
index 0000000..b8cef13
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/kernel_utils.cu
@@ -0,0 +1,52 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+
+namespace mlx::core {
+
+dim3 get_block_dims(int dim0, int dim1, int dim2, int pow2) {
+  Dims dims = get_block_dims_common(dim0, dim1, dim2, pow2);
+  return dim3(std::get<0>(dims), std::get<1>(dims), std::get<2>(dims));
+}
+
+dim3 get_2d_grid_dims(const Shape& shape, const Strides& strides) {
+  Dims dims = get_2d_grid_dims_common(shape, strides);
+  return dim3(std::get<0>(dims), std::get<1>(dims), std::get<2>(dims));
+}
+
+dim3 get_2d_grid_dims(
+    const Shape& shape,
+    const Strides& strides,
+    size_t divisor) {
+  Dims dims = get_2d_grid_dims_common(shape, strides, divisor);
+  return dim3(std::get<0>(dims), std::get<1>(dims), std::get<2>(dims));
+}
+
+std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2) {
+  auto [grid, block] = get_grid_and_block_common(dim0, dim1, dim2);
+  auto [gx, gy, gz] = grid;
+  auto [bx, by, bz] = block;
+  return std::make_pair(dim3(gx, gy, gz), dim3(bx, by, bz));
+}
+
+std::tuple<dim3, uint32_t> get_launch_args(
+    size_t size,
+    const Shape& shape,
+    const Strides& strides,
+    bool large,
+    int work_per_thread /* = 1 */,
+    uint32_t max_block_dim /* = 1024 */) {
+  size_t nthreads = cuda::ceil_div(size, work_per_thread);
+  uint32_t block_dim = max_block_dim < nthreads ? max_block_dim : nthreads;
+  dim3 num_blocks;
+  if (large) {
+    num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
+    num_blocks.x = cuda::ceil_div(num_blocks.x, block_dim);
+  } else {
+    num_blocks.x = cuda::ceil_div(nthreads, block_dim);
+  }
+  return std::make_tuple(num_blocks, block_dim);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/kernel_utils.cuh b/macos/third_party/mlx/mlx/backend/cuda/kernel_utils.cuh
new file mode 100644
index 0000000..1507a0a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/kernel_utils.cuh
@@ -0,0 +1,148 @@
+// Copyright © 2025 Apple Inc.
+
+// This file includes host-only utilities for writing CUDA kernels, the
+// difference from backend/cuda/device/utils.cuh is that the latter file only
+// include device-only code.
+
+#pragma once
+
+#include <type_traits>
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <fmt/format.h>
+#include <cuda/cmath>
+
+namespace mlx::core {
+
+template <typename F>
+void dispatch_1_2_3(int n, F&& f) {
+  switch (n) {
+    case 1:
+      f(std::integral_constant<int, 1>{});
+      break;
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 3:
+      f(std::integral_constant<int, 3>{});
+      break;
+  }
+}
+
+template <typename F>
+void dispatch_bool(bool v, F&& f) {
+  if (v) {
+    f(std::true_type{});
+  } else {
+    f(std::false_type{});
+  }
+}
+
+template <typename F>
+void dispatch_block_dim(int threads, F&& f) {
+  if (threads <= WARP_SIZE) {
+    f(std::integral_constant<int, WARP_SIZE>{});
+  } else if (threads <= WARP_SIZE * 2) {
+    f(std::integral_constant<int, WARP_SIZE * 2>{});
+  } else if (threads <= WARP_SIZE * 4) {
+    f(std::integral_constant<int, WARP_SIZE * 4>{});
+  } else if (threads <= WARP_SIZE * 8) {
+    f(std::integral_constant<int, WARP_SIZE * 8>{});
+  } else if (threads <= WARP_SIZE * 16) {
+    f(std::integral_constant<int, WARP_SIZE * 16>{});
+  } else {
+    f(std::integral_constant<int, WARP_SIZE * 32>{});
+  }
+}
+
+// Maps CPU types to CUDA types.
+template <typename T>
+struct CTypeToCudaType {
+  using type = T;
+};
+
+template <>
+struct CTypeToCudaType<float16_t> {
+  using type = __half;
+};
+
+template <>
+struct CTypeToCudaType<bfloat16_t> {
+  using type = __nv_bfloat16;
+};
+
+template <>
+struct CTypeToCudaType<complex64_t> {
+  using type = cu::complex64_t;
+};
+
+template <typename T>
+using cuda_type_t = typename CTypeToCudaType<T>::type;
+
+// Type traits for detecting floating numbers.
+template <typename T>
+inline constexpr bool is_floating_v =
+    cuda::std::is_same_v<T, float> || cuda::std::is_same_v<T, double> ||
+    cuda::std::is_same_v<T, float16_t> || cuda::std::is_same_v<T, bfloat16_t>;
+
+// Type traits for detecting complex numbers.
+template <typename T>
+inline constexpr bool is_complex_v = cuda::std::is_same_v<T, complex64_t> ||
+    cuda::std::is_same_v<T, complex128_t>;
+
+// Type traits for detecting complex or real floating point numbers.
+template <typename T>
+inline constexpr bool is_inexact_v = is_floating_v<T> || is_complex_v<T>;
+
+// Utility to copy data from vector to array in host.
+template <int NDIM = MAX_NDIM, typename T = int32_t>
+inline cuda::std::array<T, NDIM> const_param(const SmallVector<T>& vec) {
+  if (vec.size() > NDIM) {
+    throw std::runtime_error(
+        fmt::format("ndim can not be larger than {}.", NDIM));
+  }
+  cuda::std::array<T, NDIM> result;
+  std::copy_n(vec.begin(), vec.size(), result.begin());
+  return result;
+}
+
+// Compute the grid and block dimensions, check backend/common/utils.h for docs.
+dim3 get_block_dims(int dim0, int dim1, int dim2, int pow2 = 10);
+dim3 get_2d_grid_dims(const Shape& shape, const Strides& strides);
+dim3 get_2d_grid_dims(
+    const Shape& shape,
+    const Strides& strides,
+    size_t divisor);
+std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2);
+
+// Get the num_blocks and block_dims assuming each thread handles
+// |work_per_thread| elements of |arr|.
+std::tuple<dim3, uint32_t> get_launch_args(
+    size_t size,
+    const Shape& shape,
+    const Strides& strides,
+    bool large,
+    int work_per_thread = 1,
+    uint32_t max_block_dim = 1024);
+
+inline std::tuple<dim3, uint32_t> get_launch_args(
+    const array& arr,
+    bool large,
+    int work_per_thread = 1,
+    uint32_t max_block_dim = 1024) {
+  return get_launch_args(
+      arr.size(),
+      arr.shape(),
+      arr.strides(),
+      large,
+      work_per_thread,
+      max_block_dim);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/layer_norm.cu b/macos/third_party/mlx/mlx/backend/cuda/layer_norm.cu
new file mode 100644
index 0000000..da554a1
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/layer_norm.cu
@@ -0,0 +1,417 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/reduce/reduce.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/fast_primitives.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+inline __device__ float3 plus_f3(const float3& a, const float3& b) {
+  return {a.x + b.x, a.y + b.y, a.z + b.z};
+}
+
+// Similar to cub::BlockReduce, but result is broadcasted to every thread.
+template <typename T, int BLOCK_DIM>
+struct BlockBroadcastReduce {
+  static_assert(WARP_SIZE <= BLOCK_DIM && BLOCK_DIM <= WARP_SIZE * WARP_SIZE);
+  static_assert(BLOCK_DIM % WARP_SIZE == 0);
+  using TempStorage = T[BLOCK_DIM / WARP_SIZE];
+
+  cg::thread_block& block;
+  TempStorage& temp;
+
+  template <typename Op>
+  __device__ T Reduce(const T& input, const Op& op, const T& init_value) {
+    auto warp = cg::tiled_partition<WARP_SIZE>(block);
+    T x = cg::reduce(warp, input, op);
+    if (warp.thread_rank() == 0) {
+      temp[warp.meta_group_rank()] = x;
+    }
+    block.sync();
+    x = warp.thread_rank() < warp.meta_group_size() ? temp[warp.thread_rank()]
+                                                    : init_value;
+    return cg::reduce(warp, x, op);
+  }
+
+  __device__ T Sum(const T& input) {
+    return Reduce(input, cg::plus<T>{}, T{});
+  }
+};
+
+template <typename T, int BLOCK_DIM, int N_READS = 4>
+__global__ void layer_norm(
+    const T* x,
+    const T* w,
+    const T* b,
+    T* out,
+    float eps,
+    int32_t axis_size,
+    int64_t w_stride,
+    int64_t b_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceT = BlockBroadcastReduce<float, BLOCK_DIM>;
+  __shared__ typename BlockReduceT::TempStorage temp;
+
+  x += grid.block_rank() * axis_size;
+  out += grid.block_rank() * axis_size;
+
+  // Sum.
+  float sum = 0;
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      sum += static_cast<float>(xn[i]);
+    }
+  }
+  sum = BlockReduceT{block, temp}.Sum(sum);
+
+  // Mean.
+  float mean = sum / axis_size;
+
+  // Normalizer.
+  float normalizer = 0;
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    if ((index + 1) * N_READS <= axis_size) {
+      auto xn = load_vector<N_READS>(x, index);
+#pragma unroll
+      for (int i = 0; i < N_READS; ++i) {
+        float t = static_cast<float>(xn[i]) - mean;
+        normalizer += t * t;
+      }
+    } else {
+      for (int i = index * N_READS; i < axis_size; ++i) {
+        float t = static_cast<float>(x[i]) - mean;
+        normalizer += t * t;
+      }
+    }
+  }
+  normalizer = BlockReduceT{block, temp}.Sum(normalizer);
+  normalizer = rsqrt(normalizer / axis_size + eps);
+
+  // Outputs.
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+    auto bn = load_vector<N_READS>(b, index, axis_size, b_stride, T(0));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      float norm = (static_cast<float>(xn[i]) - mean) * normalizer;
+      xn[i] = wn[i] * static_cast<T>(norm) + bn[i];
+    }
+    store_vector<N_READS>(out, index, xn, axis_size);
+  }
+}
+
+template <typename T, bool HAS_W, int BLOCK_DIM, int N_READS = 4>
+__global__ void layer_norm_vjp(
+    const T* x,
+    const T* w,
+    const T* g,
+    T* gx,
+    T* gw,
+    float eps,
+    int32_t axis_size,
+    int64_t w_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceF = BlockBroadcastReduce<float, BLOCK_DIM>;
+  using BlockReduceF3 = BlockBroadcastReduce<float3, BLOCK_DIM>;
+  __shared__ union {
+    typename BlockReduceF::TempStorage f;
+    typename BlockReduceF3::TempStorage f3;
+  } temp;
+
+  x += grid.block_rank() * axis_size;
+  g += grid.block_rank() * axis_size;
+  gx += grid.block_rank() * axis_size;
+  gw += grid.block_rank() * axis_size;
+
+  // Sum.
+  float sum = 0;
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      sum += static_cast<float>(xn[i]);
+    }
+  }
+  sum = BlockReduceF{block, temp.f}.Sum(sum);
+
+  // Mean.
+  float mean = sum / axis_size;
+
+  // Normalizer.
+  float3 factors = {};
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+
+    if ((index + 1) * N_READS <= axis_size) {
+      auto xn = load_vector<N_READS>(x, index);
+#pragma unroll
+      for (int i = 0; i < N_READS; ++i) {
+        float t = static_cast<float>(xn[i]) - mean;
+        float wi = wn[i];
+        float gi = gn[i];
+        float wg = wi * gi;
+        factors = plus_f3(factors, {wg, wg * t, t * t});
+      }
+    } else {
+      for (int i = index * N_READS; i < axis_size; ++i) {
+        float t = static_cast<float>(x[i]) - mean;
+        float wi = wn[i];
+        float gi = gn[i];
+        float wg = wi * gi;
+        factors = plus_f3(factors, {wg, wg * t, t * t});
+      }
+    }
+  }
+  factors = BlockReduceF3{block, temp.f3}.Reduce(factors, plus_f3, {});
+  float meanwg = factors.x / axis_size;
+  float meanwgxc = factors.y / axis_size;
+  float normalizer2 = 1 / (factors.z / axis_size + eps);
+  float normalizer = sqrt(normalizer2);
+
+  // Outputs.
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+
+    for (int i = 0; i < N_READS; i++) {
+      float xi = (static_cast<float>(xn[i]) - mean) * normalizer;
+      float wi = wn[i];
+      float gi = gn[i];
+      xn[i] = normalizer * (wi * gi - meanwg) - xi * meanwgxc * normalizer2;
+      if constexpr (HAS_W) {
+        wn[i] = gi * xi;
+      }
+    }
+    store_vector<N_READS>(gx, index, xn, axis_size);
+    if constexpr (HAS_W) {
+      store_vector<N_READS>(gw, index, wn, axis_size);
+    }
+  }
+}
+
+} // namespace cu
+
+namespace fast {
+
+bool LayerNorm::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+// TODO: There are duplicate code with backend/metal/normalization.cpp
+void LayerNorm::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("LayerNorm::eval_gpu");
+  auto& s = stream();
+  auto& out = outputs[0];
+  auto& encoder = cu::get_command_encoder(s);
+
+  // Make sure that the last dimension is contiguous.
+  auto set_output = [&s, &out, &encoder](const array& x) {
+    bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
+    if (no_copy && x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back());
+    }
+    if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            cu::malloc_async(x.data_size() * x.itemsize(), encoder),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  const array x = set_output(inputs[0]);
+  const array& w = inputs[1];
+  const array& b = inputs[2];
+
+  int32_t axis_size = x.shape().back();
+  int32_t n_rows = x.data_size() / axis_size;
+  int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+  int64_t b_stride = (b.ndim() == 1) ? b.strides()[0] : 0;
+
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(b);
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "layernorm", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
+    dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
+      auto kernel = cu::layer_norm<DataType, block_dim(), N_READS>;
+      encoder.add_kernel_node(
+          kernel,
+          n_rows,
+          block_dim(),
+          0,
+          gpu_ptr<DataType>(x),
+          gpu_ptr<DataType>(w),
+          gpu_ptr<DataType>(b),
+          gpu_ptr<DataType>(out),
+          eps_,
+          axis_size,
+          w_stride,
+          b_stride);
+    });
+  });
+}
+
+void LayerNormVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("LayerNormVJP::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  // Ensure row contiguity. We could relax this step by checking that the array
+  // is contiguous (no broadcasts or holes) and that the input strides are the
+  // same as the cotangent strides but for now this is simpler.
+  auto check_input = [&s](const array& x, bool& copied) {
+    if (x.flags().row_contiguous) {
+      copied = false;
+      return x;
+    }
+    copied = true;
+    return contiguous_copy_gpu(x, s);
+  };
+  bool donate_x = inputs[0].is_donatable();
+  bool donate_g = inputs[3].is_donatable();
+  bool copied;
+  auto x = check_input(inputs[0], copied);
+  donate_x |= copied;
+  const array& w = inputs[1];
+  const array& b = inputs[2];
+  bool g_copied;
+  auto g = check_input(inputs[3], g_copied);
+  donate_g |= g_copied;
+  array& gx = outputs[0];
+  array& gw = outputs[1];
+  array& gb = outputs[2];
+
+  // Check whether we had a weight.
+  bool has_w = w.ndim() != 0;
+
+  // Allocate space for the outputs.
+  bool g_in_gx = false;
+  if (donate_x) {
+    gx.copy_shared_buffer(x);
+  } else if (donate_g) {
+    gx.copy_shared_buffer(g);
+    g_in_gx = true;
+  } else {
+    gx.set_data(cu::malloc_async(gx.nbytes(), encoder));
+  }
+  if (g_copied && !g_in_gx) {
+    encoder.add_temporary(g);
+  }
+
+  int32_t axis_size = x.shape().back();
+  int32_t n_rows = x.data_size() / axis_size;
+  int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+
+  // Allocate a temporary to store the gradients for w and allocate the output
+  // gradient accumulators.
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
+  bool g_in_gw = false;
+  if (has_w) {
+    if (!g_in_gx && donate_g) {
+      g_in_gw = true;
+      gw_temp.copy_shared_buffer(g);
+    } else {
+      gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder));
+      encoder.add_temporary(gw_temp);
+    }
+  }
+
+  // The gradient for b in case we had a b.
+  bool has_gb = (gb.ndim() == 1 && gb.size() == axis_size);
+  if (has_gb) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    col_reduce(encoder, g, gb, Reduce::ReduceType::Sum, {0}, plan);
+  }
+
+  // Insert dependency if `g` was donated
+  if ((g_in_gx || g_in_gw) && has_gb) {
+    encoder.set_input_array(gb);
+  }
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(g);
+  encoder.set_output_array(gx);
+  encoder.set_output_array(gw_temp);
+  dispatch_float_types(gx.dtype(), "layernorm_vjp", [&](auto type_tag) {
+    dispatch_bool(has_w, [&](auto has_w_constant) {
+      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      constexpr int N_READS = 16 / sizeof(DataType);
+      dispatch_block_dim(
+          cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
+            auto kernel = cu::layer_norm_vjp<
+                DataType,
+                has_w_constant.value,
+                block_dim(),
+                N_READS>;
+            encoder.add_kernel_node(
+                kernel,
+                n_rows,
+                block_dim(),
+                0,
+                gpu_ptr<DataType>(x),
+                gpu_ptr<DataType>(w),
+                gpu_ptr<DataType>(g),
+                gpu_ptr<DataType>(gx),
+                gpu_ptr<DataType>(gw_temp),
+                eps_,
+                axis_size,
+                w_stride);
+          });
+    });
+  });
+
+  if (has_w) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    col_reduce(encoder, gw_temp, gw, Reduce::ReduceType::Sum, {0}, plan);
+  }
+}
+
+} // namespace fast
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/load.cpp b/macos/third_party/mlx/mlx/backend/cuda/load.cpp
new file mode 100644
index 0000000..2acc52f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/load.cpp
@@ -0,0 +1,60 @@
+// Copyright © 2023 Apple Inc.
+
+#include <algorithm>
+#include <utility>
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/primitives.h"
+
+namespace {
+
+template <const uint8_t scalar_size>
+void swap_endianness(uint8_t* data_bytes, size_t N) {
+  struct Elem {
+    uint8_t bytes[scalar_size];
+  };
+
+  Elem* data = reinterpret_cast<Elem*>(data_bytes);
+
+  for (size_t i = 0; i < N; i++) {
+    for (size_t j = 0; j < (scalar_size / 2); j++) {
+      std::swap(data[i].bytes[j], data[i].bytes[scalar_size - j - 1]);
+    }
+  }
+}
+
+} // namespace
+
+namespace mlx::core {
+
+void Load::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& encoder = cu::get_command_encoder(stream());
+  auto size = out.size();
+  auto nbytes = size * out.itemsize();
+  out.set_data(cu::malloc_async(nbytes, encoder));
+  auto out_ptr = malloc(nbytes);
+  reader_->read(static_cast<char*>(out_ptr), nbytes, offset_);
+  if (swap_endianness_) {
+    switch (out.itemsize()) {
+      case 2:
+        swap_endianness<2>(reinterpret_cast<uint8_t*>(out_ptr), size);
+        break;
+      case 4:
+        swap_endianness<4>(reinterpret_cast<uint8_t*>(out_ptr), size);
+        break;
+      case 8:
+        swap_endianness<8>(reinterpret_cast<uint8_t*>(out_ptr), size);
+        break;
+    }
+  }
+  CHECK_CUDA_ERROR(cudaMemcpyAsync(
+      gpu_ptr<void>(out),
+      out_ptr,
+      nbytes,
+      cudaMemcpyDefault,
+      encoder.stream()));
+  CHECK_CUDA_ERROR(cudaLaunchHostFunc(encoder.stream(), free, out_ptr));
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/logsumexp.cu b/macos/third_party/mlx/mlx/backend/cuda/logsumexp.cu
new file mode 100644
index 0000000..a55a7d0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/logsumexp.cu
@@ -0,0 +1,161 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/cast_op.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <nvtx3/nvtx3.hpp>
+#include <cub/block/block_load.cuh>
+
+#include <cassert>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T>
+inline __device__ T softmax_exp(T x) {
+  // Softmax doesn't need high precision exponential cause x is gonna be in
+  // (-oo, 0] anyway and subsequently it will be divided by sum(exp(x_i)).
+  return __expf(x);
+}
+
+template <typename T, typename AccT, int BLOCK_DIM, int N_READS = 4>
+__global__ void logsumexp(const T* in, T* out, int axis_size) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  in += grid.block_rank() * axis_size;
+
+  cg::greater<AccT> max_op;
+  cg::plus<AccT> plus_op;
+
+  // Thread reduce.
+  AccT prevmax;
+  AccT maxval = Limits<AccT>::finite_min();
+  AccT normalizer = 0;
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto vals = load_vector<N_READS>(in, index, axis_size, Limits<T>::min());
+    prevmax = maxval;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      maxval = max_op(maxval, static_cast<AccT>(vals[i]));
+    }
+    // Online normalizer calculation for softmax:
+    // https://github.com/NVIDIA/online-softmax
+    normalizer = normalizer * softmax_exp(prevmax - maxval);
+    for (int i = 0; i < N_READS; i++) {
+      normalizer =
+          normalizer + softmax_exp(static_cast<AccT>(vals[i]) - maxval);
+    }
+  }
+
+  // First warp reduce.
+  prevmax = maxval;
+  maxval = cg::reduce(warp, maxval, max_op);
+  normalizer = normalizer * softmax_exp(prevmax - maxval);
+  normalizer = cg::reduce(warp, normalizer, plus_op);
+
+  __shared__ AccT local_max[WARP_SIZE];
+  __shared__ AccT local_normalizer[WARP_SIZE];
+
+  // Write to shared memory and do second warp reduce.
+  prevmax = maxval;
+  if (warp.thread_rank() == 0) {
+    local_max[warp.meta_group_rank()] = maxval;
+  }
+  block.sync();
+  maxval = warp.thread_rank() < warp.meta_group_size()
+      ? local_max[warp.thread_rank()]
+      : Limits<AccT>::finite_min();
+  maxval = cg::reduce(warp, maxval, max_op);
+  normalizer = normalizer * softmax_exp(prevmax - maxval);
+  if (warp.thread_rank() == 0) {
+    local_normalizer[warp.meta_group_rank()] = normalizer;
+  }
+  block.sync();
+  normalizer = warp.thread_rank() < warp.meta_group_size()
+      ? local_normalizer[warp.thread_rank()]
+      : AccT{};
+  normalizer = cg::reduce(warp, normalizer, plus_op);
+
+  // Write output.
+  if (block.thread_rank() == 0) {
+    out[grid.block_rank()] = isinf(maxval) ? maxval : log(normalizer) + maxval;
+  }
+}
+
+} // namespace cu
+
+void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("LogSumExp::eval_gpu");
+  assert(inputs.size() == 1);
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  // Make sure that the last dimension is contiguous.
+  auto ensure_contiguous = [&s, &encoder](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      encoder.add_temporary(x_copy);
+      return x_copy;
+    }
+  };
+
+  auto in = ensure_contiguous(inputs[0]);
+  if (in.flags().row_contiguous) {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  } else {
+    auto n = in.shape(-1);
+    auto flags = in.flags();
+    auto strides = in.strides();
+    for (auto& s : strides) {
+      s /= n;
+    }
+    bool col_contig = strides[0] == 1;
+    for (int i = 1; col_contig && i < strides.size(); ++i) {
+      col_contig &=
+          (out.shape(i) == 1 || strides[i - 1] == out.shape(i) * strides[i]);
+    }
+    flags.col_contiguous = col_contig;
+    out.set_data(
+        cu::malloc_async(in.nbytes() / n, encoder),
+        in.data_size() / n,
+        std::move(strides),
+        flags);
+  }
+
+  int axis_size = in.shape().back();
+  int n_rows = in.data_size() / axis_size;
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "logsumexp", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
+    dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
+      auto kernel = cu::logsumexp<DataType, float, block_dim(), N_READS>;
+      encoder.add_kernel_node(
+          kernel,
+          n_rows,
+          block_dim(),
+          0,
+          gpu_ptr<DataType>(in),
+          gpu_ptr<DataType>(out),
+          axis_size);
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/lru_cache.h b/macos/third_party/mlx/mlx/backend/cuda/lru_cache.h
new file mode 100644
index 0000000..f23f3ff
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/lru_cache.h
@@ -0,0 +1,190 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/utils.h"
+
+#include <cstring>
+#include <list>
+#include <unordered_map>
+#include <utility>
+
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+template <
+    typename K,
+    typename V,
+    template <typename...> typename M = std::unordered_map>
+class LRUCache {
+ public:
+  using value_type = std::pair<K, V>;
+  using list_type = std::list<value_type>;
+  using iterator = typename list_type::iterator;
+  using const_iterator = typename list_type::const_iterator;
+  using map_type = M<K, iterator>;
+
+  explicit LRUCache(size_t capacity) : capacity_(capacity) {
+    if (capacity == 0) {
+      throw std::runtime_error("LRUCache requires capacity > 0.");
+    }
+  }
+
+  // Initialize with capacity read from |env_name|.
+  LRUCache(const char* env_name, int default_capacity)
+      : LRUCache(env::get_var(env_name, default_capacity)) {
+    if (env::get_var("MLX_ENABLE_CACHE_THRASHING_CHECK", 1)) {
+      env_name_ = env_name;
+    }
+  }
+
+  size_t size() const {
+    return map_.size();
+  }
+  size_t capacity() const {
+    return capacity_;
+  }
+  bool empty() const {
+    return vlist_.empty();
+  }
+
+  void resize(size_t new_capacity) {
+    capacity_ = new_capacity;
+    trim();
+  }
+
+  iterator begin() {
+    return vlist_.begin();
+  }
+  const_iterator begin() const {
+    return vlist_.begin();
+  }
+  iterator end() {
+    return vlist_.end();
+  }
+  const_iterator end() const {
+    return vlist_.end();
+  }
+
+  void clear() {
+    map_.clear();
+    vlist_.clear();
+  }
+
+  iterator find(const K& key) {
+    auto it = map_.find(key);
+    if (it == map_.end())
+      return end();
+    vlist_.splice(vlist_.begin(), vlist_, it->second);
+    return it->second;
+  }
+
+  template <typename U>
+  std::pair<iterator, bool> emplace(const K& key, U&& value) {
+    auto it = map_.find(key);
+    if (it != map_.end()) {
+      vlist_.splice(vlist_.begin(), vlist_, it->second);
+      return {it->second, false};
+    }
+
+    if (env_name_ && ++cache_misses_ > 2 * capacity_) {
+      throw std::runtime_error(
+          fmt::format(
+              "Cache thrashing is happening, please set the environment variable "
+              "{} to a larger value than {} to fix degraded performance.",
+              env_name_,
+              capacity_));
+    }
+
+    vlist_.emplace_front(key, std::forward<U>(value));
+    map_[key] = vlist_.begin();
+
+    trim();
+
+    return {vlist_.begin(), true};
+  }
+
+  iterator erase(iterator pos) {
+    map_.erase(pos->first);
+    return vlist_.erase(pos);
+  }
+
+  V& operator[](const K& key) {
+    auto it = find(key);
+    if (it == end()) {
+      it = emplace(key, V{}).first;
+    }
+    return it->second;
+  }
+
+ private:
+  void trim() {
+    while (map_.size() > capacity_) {
+      auto last = std::prev(vlist_.end());
+      map_.erase(last->first);
+      vlist_.pop_back();
+    }
+  }
+
+  const char* env_name_{nullptr};
+  size_t cache_misses_{0};
+
+  list_type vlist_;
+  map_type map_;
+  size_t capacity_;
+};
+
+// Turn a POD struct into a container key by doing bytes compare.
+//
+// Usage:
+//   BytesKey<MyKey> key;
+//   key.pod = { ... };
+template <typename T>
+struct BytesKey {
+  T pod;
+  static_assert(std::is_standard_layout_v<T>, "T is not POD");
+
+  BytesKey() {
+    // Make sure the paddings between members are filled with 0.
+    memset(&pod, 0, sizeof(T));
+  }
+
+  BytesKey(const BytesKey& other) {
+    memcpy(&pod, &other.pod, sizeof(T));
+  }
+
+  BytesKey(BytesKey&& other) {
+    memcpy(&pod, &other.pod, sizeof(T));
+  }
+
+  bool operator==(const BytesKey& other) const {
+    auto* ptr1 = reinterpret_cast<const uint8_t*>(&pod);
+    auto* ptr2 = reinterpret_cast<const uint8_t*>(&other.pod);
+    return memcmp(ptr1, ptr2, sizeof(T)) == 0;
+  }
+};
+
+// Compute hash according to the bytes value of T.
+template <typename T>
+struct BytesHash {
+  static_assert(std::is_standard_layout_v<T>, "T is not POD");
+
+  size_t operator()(const T& pod) const {
+    auto* ptr = reinterpret_cast<const uint8_t*>(&pod);
+    uint32_t value = 0x811C9DC5;
+    for (int i = 0; i < sizeof(T); ++i) {
+      value ^= ptr[i];
+      value *= 0x01000193;
+    }
+    return value;
+  }
+};
+
+template <typename K, typename V>
+using BytesKeyHashMap = std::unordered_map<K, V, BytesHash<K>>;
+
+template <typename K, typename V>
+using LRUBytesKeyCache = LRUCache<BytesKey<K>, V, BytesKeyHashMap>;
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/matmul.cpp b/macos/third_party/mlx/mlx/backend/cuda/matmul.cpp
new file mode 100644
index 0000000..725f1b4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/matmul.cpp
@@ -0,0 +1,373 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/matmul.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/gemms/gemv.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+#include <numeric>
+
+namespace mlx::core {
+
+namespace {
+
+std::tuple<bool, int64_t, array>
+check_transpose(cu::CommandEncoder& enc, const Stream& s, const array& arr) {
+  auto stx = arr.strides()[arr.ndim() - 2];
+  auto sty = arr.strides()[arr.ndim() - 1];
+  if (sty == 1 && stx == arr.shape(-1)) {
+    return std::make_tuple(false, stx, arr);
+  } else if (stx == 1 && sty == arr.shape(-2)) {
+    return std::make_tuple(true, sty, arr);
+  } else {
+    array arr_copy = contiguous_copy_gpu(arr, s);
+    enc.add_temporary(arr_copy);
+    return std::make_tuple(false, arr.shape(-1), arr_copy);
+  }
+}
+
+std::tuple<bool, int64_t, array>
+ensure_batch_contiguous(const array& x, cu::CommandEncoder& encoder, Stream s) {
+  if (x.flags().row_contiguous) {
+    return std::make_tuple(false, x.strides(-2), x);
+  }
+
+  bool rc = true;
+  for (int i = 0; i < x.ndim() - 3; i++) {
+    rc &= (x.strides(i + 1) * x.shape(i)) == x.strides(i);
+  }
+  if (rc) {
+    return check_transpose(encoder, s, x);
+  }
+
+  array x_copy = contiguous_copy_gpu(x, s);
+  encoder.add_temporary(x_copy);
+  return std::make_tuple(false, x_copy.strides(-2), x_copy);
+}
+
+array ensure_row_contiguous(
+    const array& x,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    encoder.add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+void gemm_and_bias(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const std::optional<array>& bias = std::nullopt,
+    float alpha = 1.0f) {
+  // Check and collapse batch dimensions
+  auto [batch_shape, a_batch_strides, b_batch_strides] = collapse_batches(a, b);
+
+  auto batch_count = out.size() / (M * N);
+
+  // Collapse batches into M if needed
+  if (batch_count > 1 && !a_transposed && batch_shape.size() == 1 &&
+      a.strides()[a.ndim() - 2] == K && a_batch_strides.back() == M * K &&
+      b_batch_strides.back() == 0) {
+    M *= batch_shape.back();
+    batch_count = 1;
+
+    a_batch_strides = {0};
+    b_batch_strides = {0};
+    batch_shape = {1};
+  }
+
+  // Use gemmv when possible
+  if (!bias && cu::can_use_gemv(M, N, K, a_transposed, b_transposed)) {
+    cu::gemv(
+        a,
+        b,
+        out,
+        M,
+        N,
+        K,
+        batch_count,
+        batch_shape,
+        a_batch_strides,
+        b_batch_strides,
+        encoder);
+    return;
+  }
+
+  // Invoke cublasLt
+  CublasGemm gemm(
+      encoder.device(),
+      a.dtype(),
+      a_transposed,
+      M,
+      K,
+      lda,
+      b_transposed,
+      K,
+      N,
+      ldb,
+      batch_shape.back(),
+      a_batch_strides.back(),
+      b_batch_strides.back());
+  if (bias) {
+    if (a.dtype() == complex64) {
+      throw std::runtime_error(
+          "[gemm_and_bias] complex64 bias epilogue isn’t supported in cublasLtMatmul.");
+    }
+    gemm.set_bias(encoder, *bias);
+  }
+  gemm.run(
+      encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides, alpha);
+}
+
+void gather_mm_rhs(
+    const array& a_,
+    const array& b_,
+    const array& indices_,
+    array& out,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  if (a_.size() / a_.shape(-2) / a_.shape(-1) != indices_.size()) {
+    throw std::runtime_error("[gather_mm] Broadcasting lhs is not supported.");
+  }
+
+  int group_count = b_.size() / b_.shape(-1) / b_.shape(-2);
+  if (group_count > 1024) {
+    throw std::runtime_error(
+        "[gather_mm] Group count can not be larger than 1024.");
+  }
+
+  auto [a_transposed, lda, a] = ensure_batch_contiguous(a_, encoder, s);
+  auto [b_transposed, ldb, b] = ensure_batch_contiguous(b_, encoder, s);
+  auto indices = ensure_row_contiguous(indices_, encoder, s);
+
+  cutlass_grouped_gemm_unaligned(
+      a_transposed,
+      lda,
+      b_transposed,
+      ldb,
+      group_count,
+      a,
+      b,
+      indices,
+      out,
+      encoder);
+}
+
+} // namespace
+
+void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Matmul::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 2);
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+  // Return 0s if either input is empty.
+  if (a_pre.size() == 0 || b_pre.size() == 0) {
+    array zero(0, a_pre.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int K = a_pre.shape(-1);
+
+  // Keep a vector with copies to be cleared in the completed buffer to release
+  // the arrays
+  auto [a_transposed, lda, a] = check_transpose(encoder, s, a_pre);
+  auto [b_transposed, ldb, b] = check_transpose(encoder, s, b_pre);
+
+  gemm_and_bias(
+      encoder, M, N, K, a_transposed, lda, b_transposed, ldb, out, a, b);
+}
+
+void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("AddMM::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 3);
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+  auto c = inputs[2];
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Init checks and prep
+
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int K = a_pre.shape(-1);
+
+  // Keep a vector with copies to be cleared in the completed buffer to release
+  // the arrays
+  auto [a_transposed, lda, a] = check_transpose(encoder, s, a_pre);
+  auto [b_transposed, ldb, b] = check_transpose(encoder, s, b_pre);
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Dispatch to GEMM with epilogue or AddMM
+
+  if (beta_ == 1 && a.dtype() != complex64 && c.strides(-1) == 1 &&
+      c.data_size() == out.shape(-1)) {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    gemm_and_bias(
+        encoder,
+        M,
+        N,
+        K,
+        a_transposed,
+        lda,
+        b_transposed,
+        ldb,
+        out,
+        a,
+        b,
+        c,
+        alpha_);
+    return;
+  }
+
+  int64_t ldc;
+  {
+    auto stx = c.strides()[c.ndim() - 2];
+    auto sty = c.strides()[c.ndim() - 1];
+    if (sty == 1 && stx == c.shape(-1)) {
+      ldc = stx;
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    } else if (sty == 1 && stx == 0) {
+      ldc = 0;
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    } else {
+      // Copy C into out and set C to out
+      ldc = c.shape(-1);
+      copy_gpu(c, out, CopyType::General, s);
+      c = out;
+    }
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Check and collapse batch dimensions
+
+  auto [batch_shape, a_batch_strides, b_batch_strides, c_batch_strides] =
+      collapse_batches(a, b, c);
+
+  auto batch_count = out.size() / (M * N);
+
+  // Collapse batches into M if needed
+  if (batch_count > 1 && !a_transposed && batch_shape.size() == 1 &&
+      a.strides()[a.ndim() - 2] == K && a_batch_strides.back() == M * K &&
+      c_batch_strides.back() == M * c.strides()[c.ndim() - 2] &&
+      b_batch_strides.back() == 0) {
+    M *= batch_shape.back();
+    batch_count = 1;
+
+    a_batch_strides = {0};
+    b_batch_strides = {0};
+    c_batch_strides = {0};
+    batch_shape = {1};
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Invoke cublasLt with AddMM settings
+
+  CublasGemm gemm(
+      cu::device(s.device),
+      a.dtype(),
+      a_transposed,
+      M,
+      K,
+      lda,
+      b_transposed,
+      K,
+      N,
+      ldb,
+      ldc,
+      batch_shape.back(),
+      a_batch_strides.back(),
+      b_batch_strides.back(),
+      c_batch_strides.back());
+  gemm.run(
+      encoder,
+      out,
+      a,
+      b,
+      c,
+      batch_shape,
+      a_batch_strides,
+      b_batch_strides,
+      c_batch_strides,
+      alpha_,
+      beta_);
+}
+
+void GatherMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("GatherMM::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 4);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& lhs_indices = inputs[2];
+  auto& rhs_indices = inputs[3];
+
+  // Return 0s if either input is empty.
+  if (a.size() == 0 || b.size() == 0) {
+    array zero(0, a.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  // Extract shapes from inputs.
+  int M = a.shape(-2);
+  int N = b.shape(-1);
+  int K = a.shape(-1);
+
+  // We are walking a in order and b is also in order so we can batch up the
+  // matmuls and reuse reading a and b.
+  if (M == 1 && right_sorted_ == true) {
+    gather_mm_rhs(a, b, rhs_indices, out, encoder, s);
+    return;
+  }
+
+  auto [transposed_a, lda, a_] = check_transpose(encoder, s, a);
+  auto [transposed_b, ldb, b_] = check_transpose(encoder, s, b);
+  auto use_gemv = cu::can_use_gemv(M, N, K, transposed_a, transposed_b);
+  if (M == 1 && use_gemv) {
+    gather_mv(b_, a_, rhs_indices, lhs_indices, out, N, K, encoder);
+    return;
+  }
+
+  if (N == 1 && use_gemv) {
+    gather_mv(a_, b_, lhs_indices, rhs_indices, out, M, K, encoder);
+    return;
+  }
+
+  throw std::runtime_error("NYI");
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/no_cuda.cpp b/macos/third_party/mlx/mlx/backend/cuda/no_cuda.cpp
new file mode 100644
index 0000000..175a505
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/no_cuda.cpp
@@ -0,0 +1,47 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/fast.h"
+
+namespace mlx::core {
+
+namespace cu {
+
+bool is_available() {
+  return false;
+}
+
+} // namespace cu
+
+namespace fast {
+
+CustomKernelFunction cuda_kernel(
+    const std::string&,
+    const std::vector<std::string>&,
+    const std::vector<std::string>&,
+    const std::string&,
+    const std::string&,
+    bool,
+    int) {
+  throw std::runtime_error("[cuda_kernel] No CUDA back-end.");
+}
+
+std::vector<array> precompiled_cuda_kernel(
+    const std::string&,
+    const std::string&,
+    const std::vector<array>&,
+    const std::vector<Shape>&,
+    const std::vector<Dtype>&,
+    const std::vector<ScalarArg>&,
+    std::tuple<int, int, int>,
+    std::tuple<int, int, int>,
+    int shared_memory,
+    std::optional<float> init_value,
+    bool ensure_row_contiguous,
+    StreamOrDevice) {
+  throw std::runtime_error("[cuda_kernel] No CUDA back-end.");
+}
+
+} // namespace fast
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/primitives.cpp b/macos/third_party/mlx/mlx/backend/cuda/primitives.cpp
new file mode 100644
index 0000000..0caac8d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/primitives.cpp
@@ -0,0 +1,46 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/distributed/primitives.h"
+#include <cuda_runtime.h>
+#include "mlx/fast_primitives.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+#define NO_GPU_MULTI(func)                                             \
+  void func::eval_gpu(                                                 \
+      const std::vector<array>& inputs, std::vector<array>& outputs) { \
+    throw std::runtime_error(#func " has no CUDA implementation.");    \
+  }
+
+#define NO_GPU_USE_FALLBACK(func)     \
+  bool func::use_fallback(Stream s) { \
+    return true;                      \
+  }                                   \
+  NO_GPU_MULTI(func)
+
+#define NO_GPU(func)                                                  \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    throw std::runtime_error(#func " has no CUDA implementation.");   \
+  }
+
+NO_GPU(BlockMaskedMM)
+NO_GPU(FFT)
+NO_GPU(GatherQMM)
+NO_GPU(Hadamard)
+NO_GPU_MULTI(LUF)
+NO_GPU_MULTI(QRF)
+NO_GPU(SegmentedMM)
+NO_GPU_MULTI(SVD)
+NO_GPU(Inverse)
+NO_GPU(Cholesky)
+NO_GPU_MULTI(Eig)
+NO_GPU_MULTI(Eigh)
+NO_GPU(MaskedScatter)
+
+namespace distributed {
+NO_GPU_MULTI(Send)
+NO_GPU_MULTI(Recv)
+} // namespace distributed
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/affine_quantize.cu b/macos/third_party/mlx/mlx/backend/cuda/quantized/affine_quantize.cu
new file mode 100644
index 0000000..105547b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/affine_quantize.cu
@@ -0,0 +1,329 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core {
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, int group_size, int bits>
+__global__ void
+affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) {
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+  constexpr float eps = 1e-7;
+  constexpr int simd_size = WARP_SIZE;
+  constexpr float n_bins = (1 << bits) - 1;
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+  constexpr int values_per_reduce = group_size / simd_size;
+  constexpr int writes_per_reduce = pack_factor / values_per_reduce;
+  constexpr int writes_per_pack =
+      writes_per_reduce > 1 ? 1 : values_per_reduce / pack_factor;
+  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
+
+  size_t offset = tidx + grid_dim_x * size_t(tidy);
+  size_t in_index = offset * values_per_reduce;
+  if (in_index >= size) {
+    return;
+  }
+  size_t out_index = power_of_2_bits
+      ? offset * writes_per_pack
+      : offset * bytes_per_pack / writes_per_reduce;
+
+  float w_thread[values_per_reduce];
+  float w_min = Limits<float>::max();
+  float w_max = 0;
+
+#pragma clang loop unroll(full)
+  for (int i = 0; i < values_per_reduce; i++) {
+    float val = w[in_index + i];
+    w_thread[i] = val;
+    w_min = min(w_min, val);
+    w_max = max(w_max, val);
+  }
+
+  cg::greater<float> max_op;
+  cg::less<float> min_op;
+  auto warp = cg::tiled_partition<WARP_SIZE>(cg::this_thread_block());
+
+  w_min = cg::reduce(warp, w_min, min_op);
+  w_max = cg::reduce(warp, w_max, max_op);
+
+  float scale = max((w_max - w_min) / n_bins, eps);
+  bool side = abs(w_min) > abs(w_max);
+  scale = side ? scale : -scale;
+  float edge = side ? w_min : w_max;
+  float q0 = round(edge / scale);
+  bool at_zero = q0 == 0.0f;
+  scale = at_zero ? scale : edge / q0;
+  float bias = at_zero ? 0 : edge;
+
+  // Write out the scales and biases
+  size_t gindex = in_index / group_size;
+  if (in_index % group_size == 0) {
+    scales[gindex] = static_cast<T>(scale);
+    biases[gindex] = static_cast<T>(bias);
+  }
+
+  using OutType = std::conditional_t<bits == 5, uint64_t, uint32_t>;
+  OutType output = 0;
+
+#pragma clang loop unroll(full)
+  for (int i = 0; i < values_per_reduce; i++) {
+    uint8_t val = min(round((w_thread[i] - bias) / scale), n_bins);
+    if (bits == 8) {
+      output = val;
+    } else {
+      output |= val << (bits * (i % pack_factor));
+    }
+
+    if (pack_factor < values_per_reduce && i % pack_factor == pack_factor - 1) {
+      out[out_index + i / pack_factor] = output;
+      output = 0;
+    } else {
+#pragma clang loop unroll(full)
+      for (int j = 1; j < writes_per_reduce; j++) {
+        uint8_t sval = warp.shfl_down(val, j);
+        output |= static_cast<OutType>(sval)
+            << (bits * (j * values_per_reduce + i));
+      }
+    }
+  }
+  if constexpr (bits == 3 || bits == 6) {
+    if (in_index % pack_factor == 0 && out_index % bytes_per_pack == 0) {
+      out[out_index] = output & 0xff;
+      out[out_index + 1] = (output & 0xff00) >> 8;
+      out[out_index + 2] = (output & 0xff0000) >> 16;
+    }
+  } else if constexpr (bits == 5) {
+    if (in_index % pack_factor == 0 && out_index % bytes_per_pack == 0) {
+      out[out_index] = output & 0xff;
+      out[out_index + 1] = (output & 0xff00) >> 8;
+      out[out_index + 2] = (output & 0xff0000) >> 16;
+      out[out_index + 3] = (output & 0xff000000) >> 24;
+      out[out_index + 4] = (output & 0xff00000000) >> 32;
+    }
+  } else {
+    if constexpr (writes_per_reduce > 0) {
+      if (out_index % writes_per_reduce == 0) {
+        out[out_index / writes_per_reduce] = output;
+      }
+    }
+  }
+}
+
+template <typename T, int group_size, int bits>
+__global__ void affine_dequantize(
+    const uint8_t* w,
+    const T* scales,
+    const T* biases,
+    T* out,
+    size_t size) {
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+
+  size_t offset = tidx + grid_dim_x * size_t(tidy);
+  size_t oindex = offset * pack_factor;
+
+  if (oindex >= size) {
+    return;
+  }
+
+  size_t gindex = oindex / group_size;
+  T scale = scales[gindex];
+  T bias = biases[gindex];
+  out += oindex;
+
+  if constexpr (bits == 3) {
+    w += offset * bytes_per_pack;
+    out[0] = static_cast<T>(w[0] & 0x7) * scale + bias;
+    out[1] = static_cast<T>((w[0] & 0x38) >> 3) * scale + bias;
+    out[2] = (static_cast<T>((w[0] & 0xc0) >> 6) +
+              static_cast<T>((w[1] & 0x1) << 2)) *
+            scale +
+        bias;
+    out[3] = static_cast<T>((w[1] & 0xe) >> 1) * scale + bias;
+    out[4] = static_cast<T>((w[1] & 0x70) >> 4) * scale + bias;
+    out[5] = (static_cast<T>((w[1] & 0x80) >> 7) +
+              static_cast<T>((w[2] & 0x3) << 1)) *
+            scale +
+        bias;
+    out[6] = static_cast<T>((w[2] & 0x1c) >> 2) * scale + bias;
+    out[7] = static_cast<T>((w[2] & 0xe0) >> 5) * scale + bias;
+  } else if constexpr (bits == 5) {
+    w += offset * bytes_per_pack;
+    out[0] = static_cast<T>(w[0] & 0x1f) * scale + bias;
+    out[1] = (static_cast<T>((w[0] & 0xe0) >> 5) +
+              static_cast<T>((w[1] & 0x3) << 3)) *
+            scale +
+        bias;
+    out[2] = static_cast<T>((w[1] & 0x7c) >> 2) * scale + bias;
+    out[3] = (static_cast<T>((w[1] & 0x80) >> 7) +
+              static_cast<T>((w[2] & 0xf) << 1)) *
+            scale +
+        bias;
+    out[4] = (static_cast<T>((w[2] & 0xf0) >> 4) +
+              static_cast<T>((w[3] & 0x1) << 4)) *
+            scale +
+        bias;
+    out[5] = static_cast<T>((w[3] & 0x3e) >> 1) * scale + bias;
+    out[6] = (static_cast<T>((w[3] & 0xc0) >> 6) +
+              static_cast<T>((w[4] & 0x7) << 2)) *
+            scale +
+        bias;
+    out[7] = static_cast<T>((w[4] & 0xf8) >> 3) * scale + bias;
+  } else if constexpr (bits == 6) {
+    w += offset * bytes_per_pack;
+    out[0] = static_cast<T>(w[0] & 0x3f) * scale + bias;
+    out[1] = (static_cast<T>((w[0] >> 6) & 0x03) +
+              static_cast<T>((w[1] & 0x0f) << 2)) *
+            scale +
+        bias;
+    out[2] = (static_cast<T>((w[1] >> 4) & 0x0f) +
+              static_cast<T>((w[2] & 0x03) << 4)) *
+            scale +
+        bias;
+    out[3] = static_cast<T>((w[2] >> 2) & 0x3f) * scale + bias;
+  } else {
+    uint32_t val = w[offset];
+#pragma clang loop unroll(full)
+    for (int i = 0; i < pack_factor; i++) {
+      uint8_t d;
+      if (bits == 2) {
+        d = (val >> (bits * i)) & 0x03;
+      } else if (bits == 4) {
+        d = (val >> (bits * i)) & 0x0f;
+      } else if (bits == 8) {
+        d = val;
+      }
+      out[i] = scale * static_cast<T>(d) + bias;
+    }
+  }
+}
+
+} // namespace cu
+
+void affine_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    array& biases,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  // Calculate the number of elements per thread
+  int per_thread = group_size_ / WARP_SIZE;
+  size_t size = w.size() / per_thread;
+
+  // Calculate the thread grid that we need to launch
+  bool large = size > UINT_MAX;
+  auto grid_shape = w.shape();
+  grid_shape.back() /= per_thread;
+
+  enc.set_input_array(w);
+  enc.set_output_array(wq);
+  enc.set_output_array(scales);
+  enc.set_output_array(biases);
+  dispatch_float_types(w.dtype(), "affine_quantize", [&](auto type_tag) {
+    dispatch_groups(group_size_, [&](auto group_size) {
+      dispatch_bits(bits_, [&](auto bits) {
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        auto kernel = cu::affine_quantize<T, group_size.value, bits.value>;
+        auto [num_blocks, block_dims] =
+            get_launch_args(size, grid_shape, w.strides(), large);
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<T>(w),
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<T>(scales),
+            gpu_ptr<T>(biases),
+            w.size());
+      });
+    });
+  });
+}
+
+void affine_dequantize(
+    const array& wq,
+    const array& scales,
+    const array& biases,
+    array& w,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  // Calculate how many numbers we pack together. For 2, 4, 8 bits we pack in
+  // one uint8, for 3, 6 in 3 uint8 and for 5 in 5 uint8.
+  constexpr int uint8_per_uint32 = 4;
+  int packs_per_int;
+  switch (bits_) {
+    case 3:
+    case 5:
+      packs_per_int = 8;
+      break;
+    case 6:
+      packs_per_int = 4;
+      break;
+    default:
+      packs_per_int = 8 / bits_;
+  }
+
+  size_t size = w.size() / packs_per_int;
+  bool large = size > UINT_MAX;
+  auto grid_shape = w.shape();
+  grid_shape.back() *= uint8_per_uint32;
+
+  enc.set_input_array(wq);
+  enc.set_input_array(scales);
+  enc.set_input_array(biases);
+  enc.set_output_array(w);
+  dispatch_float_types(w.dtype(), "affine_dequantize", [&](auto type_tag) {
+    dispatch_groups(group_size_, [&](auto group_size) {
+      dispatch_bits(bits_, [&](auto bits) {
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        auto kernel = cu::affine_dequantize<T, group_size.value, bits.value>;
+        auto [num_blocks, block_dims] =
+            get_launch_args(size, grid_shape, w.strides(), large);
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<T>(scales),
+            gpu_ptr<T>(biases),
+            gpu_ptr<T>(w),
+            w.size());
+      });
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/convert_fp8.cu b/macos/third_party/mlx/mlx/backend/cuda/quantized/convert_fp8.cu
new file mode 100644
index 0000000..c0be2e3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/convert_fp8.cu
@@ -0,0 +1,19 @@
+// Copyright © 2025 Apple Inc.
+#include "mlx/backend/cuda/unary/unary.cuh"
+#include "mlx/fast_primitives.h"
+
+namespace mlx::core {
+void fast::ConvertFP8::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("ConvertFP8::eval_gpu");
+  auto& in = inputs[0];
+  auto& out = outputs[0];
+  auto& s = out.primitive().stream();
+  if (to_fp8_) {
+    unary_op_gpu<cu::ToFP8>(inputs, out, name(), s);
+  } else {
+    unary_op_gpu<cu::FromFP8>(inputs, out, name(), s);
+  }
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/cublas_qqmm.cpp b/macos/third_party/mlx/mlx/backend/cuda/quantized/cublas_qqmm.cpp
new file mode 100644
index 0000000..959ae93
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/cublas_qqmm.cpp
@@ -0,0 +1,206 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+
+#include <fmt/format.h>
+#include "mlx/backend/cuda/cublas_utils.h"
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+// Currently cublas supports only mxfp8 and nvfp4
+// quantization modes for block scaled quantization
+cudaDataType_t qmode_to_cublas_scale_dtype(std::string mode) {
+  if (mode == "mxfp8") {
+    return CUDA_R_8F_UE8M0;
+  } else if (mode == "nvfp4") {
+    return CUDA_R_8F_UE4M3;
+  } else {
+    throw std::runtime_error(
+        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+  }
+}
+
+cudaDataType_t qmode_to_cublas_dtype(std::string mode) {
+  if (mode == "mxfp8") {
+    return CUDA_R_8F_E4M3;
+  } else if (mode == "nvfp4") {
+    return CUDA_R_4F_E2M1;
+  } else {
+    throw std::runtime_error(
+        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+  }
+}
+
+cublasLtMatmulMatrixScale_t qmode_to_cublas_scale_mode(std::string mode) {
+  if (mode == "mxfp8") {
+    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;
+  } else if (mode == "nvfp4") {
+    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
+  } else {
+    throw std::runtime_error(
+        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+  }
+}
+
+} // namespace
+
+CublasQQMM::CublasQQMM(
+    cu::Device& device,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride,
+    Dtype out_dtype,
+    std::string qmode) {
+  // The compute type must be CUBLAS_COMPUTE_32F.
+  // The scale type must be CUDA_R_32F.
+  cudaDataType_t scale_type = CUDA_R_32F;
+  cublasComputeType_t gemm_compute_type = CUBLAS_COMPUTE_32F;
+  cudaDataType_t output_type =
+      cublas_utils::dtype_to_cublas_type(out_dtype, "CublasQQMM");
+  cudaDataType_t data_type = qmode_to_cublas_dtype(qmode);
+  quantization_mode_ = std::string(qmode);
+
+  init_base(
+      device,
+      scale_type,
+      gemm_compute_type,
+      data_type,
+      output_type,
+      a_transposed,
+      a_rows,
+      a_cols,
+      lda,
+      b_transposed,
+      b_rows,
+      b_cols,
+      ldb,
+      batch_count,
+      a_batch_stride,
+      b_batch_stride);
+
+  a_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
+  b_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
+
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_B_SCALE_MODE,
+      &a_scale_mode_,
+      sizeof(a_scale_mode_)));
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_A_SCALE_MODE,
+      &b_scale_mode_,
+      sizeof(b_scale_mode_)));
+}
+
+CublasQQMM::CublasQQMM(
+    cu::Device& device,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int64_t ldc,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride,
+    int64_t c_batch_stride,
+    Dtype out_dtype,
+    std::string qmode)
+    : CublasQQMM(
+          device,
+          a_transposed,
+          a_rows,
+          a_cols,
+          lda,
+          b_transposed,
+          b_rows,
+          b_cols,
+          ldb,
+          batch_count,
+          a_batch_stride,
+          b_batch_stride,
+          out_dtype,
+          qmode) {
+  auto type = cublas_utils::dtype_to_cublas_type(
+      out_dtype, "CublasQQMM"); // must match the output type
+  c_desc_ = cublas_utils::create_matrix_layout(
+      type,
+      b_transposed ? b_rows : b_cols,
+      a_transposed ? a_cols : a_rows,
+      false,
+      ldc,
+      batch_count,
+      c_batch_stride);
+}
+
+void CublasQQMM::run(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    float alpha) {
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(a_scale);
+  encoder.set_input_array(b_scale);
+  encoder.set_output_array(out);
+
+  execute(
+      encoder,
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
+      gpu_ptr<void>(a_scale),
+      gpu_ptr<void>(b_scale),
+      nullptr,
+      alpha);
+}
+
+void CublasQQMM::execute(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* a_scale,
+    const void* b_scale,
+    const void* c,
+    float alpha /* = 1 */,
+    float beta /* = 0 */) {
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
+      &b_scale,
+      sizeof(b_scale)));
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
+      &a_scale,
+      sizeof(a_scale)));
+
+  const void* alpha_ptr = &alpha;
+  const void* beta_ptr = &beta;
+
+  execute_matmul(encoder, out, a, b, c, alpha_ptr, beta_ptr);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/cublas_qqmm.h b/macos/third_party/mlx/mlx/backend/cuda/quantized/cublas_qqmm.h
new file mode 100644
index 0000000..0a710f6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/cublas_qqmm.h
@@ -0,0 +1,88 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/device.h"
+
+#include <cublasLt.h>
+
+namespace mlx::core {
+
+class CublasQQMM : public CublasMatmulBase {
+ public:
+  CublasQQMM(
+      cu::Device& device,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride,
+      Dtype out_dtype,
+      std::string quantization_mode);
+
+  CublasQQMM(
+      cu::Device& device,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int64_t ldc,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride,
+      int64_t c_batch_stride,
+      Dtype out_dtype,
+      std::string quantization_mode);
+
+  void run(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const array& a_scale,
+      const array& b_scale,
+      float alpha = 1.0f);
+
+ private:
+  void run_batched(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const array& a_scale,
+      const array& b_scale,
+      const Shape& batch_shape,
+      const Strides& a_batch_strides,
+      const Strides& b_batch_strides,
+      float alpha);
+
+  void execute(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* a_scale,
+      const void* b_scale,
+      const void* c,
+      float alpha = 1,
+      float beta = 0);
+
+  std::string quantization_mode_;
+  cublasLtMatmulMatrixScale_t a_scale_mode_;
+  cublasLtMatmulMatrixScale_t b_scale_mode_;
+  cublasLtMatmulMatrixScale_t c_scale_mode_;
+  cublasLtMatmulMatrixScale_t out_scale_mode_;
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/cuda_fp4.h b/macos/third_party/mlx/mlx/backend/cuda/quantized/cuda_fp4.h
new file mode 100644
index 0000000..4500be6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/cuda_fp4.h
@@ -0,0 +1,100 @@
+#pragma once
+
+struct __nv_fp8_e8m0 {
+  __device__ __nv_fp8_e8m0(float x) {
+    if (!std::isfinite(x)) {
+      __x = 0xFF;
+      return;
+    }
+    if (x < 0.0f) {
+      __x = 0x00;
+      return;
+    }
+    float le = std::log2f(x);
+    int n = static_cast<int>(std::nearbyintf(le));
+
+    n = n < -127 ? -127 : n;
+    n = n > 127 ? 127 : n;
+    __x = static_cast<uint8_t>(n + 127);
+  }
+
+  __device__ operator float() {
+    if (__x == 0xFF) {
+      return std::numeric_limits<float>::quiet_NaN();
+    }
+    return std::ldexp(1.0f, static_cast<int>(__x) - 127);
+  }
+
+  uint8_t __x{0};
+};
+
+struct __nv_fp4_e2m1 {
+  __device__ __nv_fp4_e2m1(float x) {
+    if (std::isnan(x)) {
+      __x = 0x7;
+      return;
+    }
+
+    const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0;
+    x = std::abs(x);
+
+    if (x > 5.0f) {
+      __x = 0x7;
+    } else if (x >= 3.5f) {
+      __x = 0x6;
+    } else if (x > 2.5f) {
+      __x = 0x5;
+    } else if (x >= 1.75f) {
+      __x = 0x4;
+    } else if (x > 1.25f) {
+      __x = 0x3;
+    } else if (x >= 0.75f) {
+      __x = 0x2;
+    } else if (x > 0.25f) {
+      __x = 0x1;
+    } else {
+      __x = 0x0;
+    }
+    __x |= sign_bit;
+  }
+
+  __device__ operator float() {
+    static const float LUT[16] = {
+        0.0f,
+        0.5f,
+        1.0f,
+        1.5f,
+        2.0f,
+        3.0f,
+        4.0f,
+        6.0f,
+        -0.0f,
+        -0.5f,
+        -1.0f,
+        -1.5f,
+        -2.0f,
+        -3.0f,
+        -4.0f,
+        -6.0f};
+
+    return LUT[__x];
+  }
+  uint8_t __x{0};
+};
+
+struct __nv_fp4x4_e2m1 {
+  __device__ operator float4() {
+    float4 out;
+    auto bits = __high & 0xf;
+    out.x = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__high >> 4) & 0xf;
+    out.y = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low) & 0xf;
+    out.z = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low >> 4) & 0xf;
+    out.w = float(*(__nv_fp4_e2m1*)(&bits));
+    return out;
+  }
+  uint8_t __high{0};
+  uint8_t __low{0};
+};
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/fp_quantize.cu b/macos/third_party/mlx/mlx/backend/cuda/quantized/fp_quantize.cu
new file mode 100644
index 0000000..3b4e96e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/fp_quantize.cu
@@ -0,0 +1,496 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/mxfp8_quantize.cuh"
+#include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh"
+#include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/backend/cuda/vector_types.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core {
+namespace cu {
+
+template <int bits>
+struct Dequantize {
+  __device__ float operator()(uint8_t x) {
+    if constexpr (bits == 8) {
+      return float(*(__nv_fp8_e4m3*)(&x));
+    } else {
+      return float(*(__nv_fp4_e2m1*)(&x));
+    }
+  }
+};
+
+namespace cg = cooperative_groups;
+
+template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
+__global__ void fp_quantize_dequantize(T* w, T* out, size_t size) {
+  using Tx2 = Vector2_t<T>;
+  using Tx4 = Vector4_t<T>;
+  uint32_t rbits = 0; // reserved bits for future use
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+
+  size_t thread_idx = tidx + grid_dim_x * size_t(tidy);
+  size_t base_idx = thread_idx * group_size;
+
+  if (base_idx >= size) {
+    return;
+  }
+
+  auto w_tile = load_vector<group_size, T>(w, thread_idx);
+  float scale = 0.0f;
+
+  Tx2 amax_2x = Tx2{0.0f, 0.0f};
+
+#pragma unroll
+  for (int i = 0; i < group_size; i += 2) {
+    auto pair = Tx2{w_tile[i], w_tile[i + 1]};
+    abs_max_x2<Tx2>(amax_2x, amax_2x, pair);
+  }
+
+  scale = static_cast<float>(
+      max(fabsf(static_cast<float>(amax_2x.x)),
+          fabsf(static_cast<float>(amax_2x.y))));
+
+  scale /= bits == 4 ? 6.0f : 448.0f;
+  // Convert to mx scale or nv scale
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto s = ScaleType(scale);
+  scale = float(s);
+  AlignedVector<T, group_size> w_hat;
+
+#pragma unroll
+  for (int i = 0; i < group_size / 4; i++) {
+    Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
+    float4 dq;
+    if constexpr (bits == 8) {
+      uint32_t quantized_val =
+          scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      dq = dequant_fp8(quantized_val);
+    } else {
+      uint16_t quantized_val =
+          scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      dq = dequant_fp4(quantized_val);
+    }
+    w_hat[i * 4] = static_cast<T>(dq.x * scale);
+    w_hat[i * 4 + 1] = static_cast<T>(dq.y * scale);
+    w_hat[i * 4 + 2] = static_cast<T>(dq.z * scale);
+    w_hat[i * 4 + 3] = static_cast<T>(dq.w * scale);
+  }
+  store_vector<group_size>(out, thread_idx, w_hat);
+}
+
+template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
+__global__ void
+fp_quantize_rowwise(T* w, uint8_t* out, uint8_t* scales, size_t size) {
+  using Tx2 = Vector2_t<T>;
+  using Tx4 = Vector4_t<T>;
+  uint32_t rbits = 0; // reserved bits for future use
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+
+  size_t thread_idx = tidx + grid_dim_x * size_t(tidy);
+  size_t base_idx = thread_idx * group_size;
+
+  if (base_idx >= size) {
+    return;
+  }
+
+  auto w_tile = load_vector<group_size, T>(w, thread_idx);
+  float scale = 0.0f;
+
+  Tx2 amax_2x = Tx2{0.0f, 0.0f};
+
+#pragma unroll
+  for (int i = 0; i < group_size; i += 2) {
+    auto pair = Tx2{w_tile[i], w_tile[i + 1]};
+    abs_max_x2<Tx2>(amax_2x, amax_2x, pair);
+  }
+
+  scale = static_cast<float>(
+      max(fabsf(static_cast<float>(amax_2x.x)),
+          fabsf(static_cast<float>(amax_2x.y))));
+
+  scale /= bits == 4 ? 6.0f : 448.0f;
+  // Convert to mx scale or nv scale
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto s = ScaleType(scale);
+  uint8_t q_scale = s.__x;
+  scale = float(s);
+
+  scales[thread_idx] = q_scale;
+  constexpr int elem_per_byte = bits == 8 ? 1 : 2;
+  AlignedVector<uint8_t, group_size / elem_per_byte> quantized;
+
+#pragma unroll
+  for (int i = 0; i < group_size / 4; i++) {
+    Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
+    if constexpr (bits == 8) {
+      uint32_t quantized_val =
+          scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      *reinterpret_cast<uint32_t*>(&quantized[i * 4]) = quantized_val;
+    } else {
+      uint16_t quantized_val =
+          scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      *reinterpret_cast<uint16_t*>(&quantized[i * 2]) = quantized_val;
+    }
+  }
+  store_vector<group_size / elem_per_byte>(out, thread_idx, quantized);
+}
+
+template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
+__global__ void fp_quantize_columnwise(
+    T* w,
+    uint8_t* out,
+    uint8_t* scales,
+    size_t size,
+    int M,
+    int K) {
+  // Input: [M, K] with strides [1, M] (M-major)
+  // Quantized output: [M, K/elem_per_byte] row-major (K-major)
+  // Scales: [M, K/group_size] row-major (K-major)
+  // Quantize along K (last dimension, groups of group_size elements)
+
+  using Tx2 = Vector2_t<T>;
+  using Tx4 = Vector4_t<T>;
+  uint32_t rbits = 0;
+
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  constexpr int BLOCK_X = 32;
+  constexpr int BLOCK_Y = 32;
+  constexpr int elem_per_byte = (bits == 8) ? 1 : 2;
+  constexpr int bytes_per_group = group_size / elem_per_byte;
+
+  constexpr int rows_per_block = BLOCK_X;
+  constexpr int cols_per_block = BLOCK_Y * group_size;
+  constexpr int local_cols = cols_per_block / elem_per_byte;
+  constexpr int bytes_per_block = rows_per_block * local_cols;
+
+  constexpr int SMEM_PAD = 4;
+  constexpr int padded_local_cols = local_cols + SMEM_PAD;
+
+  auto tidx = idx_in_block.x;
+  auto tidy = idx_in_block.y;
+
+  int num_col_blocks = (K + cols_per_block - 1) / cols_per_block;
+  auto bidx = block_idx.x % num_col_blocks;
+  auto bidy = block_idx.x / num_col_blocks;
+
+  T thread_data[group_size];
+
+  __shared__ uint8_t quantized_smem[rows_per_block * padded_local_cols];
+  __shared__ uint8_t scales_smem[BLOCK_X][BLOCK_Y + SMEM_PAD];
+
+  int row_base = bidy * rows_per_block + tidx;
+  int col_base = bidx * cols_per_block + tidy * group_size;
+
+  bool valid = (row_base < M) && (col_base + group_size <= K);
+  if (valid) {
+#pragma unroll
+    for (int i = 0; i < group_size; i++) {
+      auto index = row_base + (col_base + i) * M;
+      thread_data[i] = w[index];
+    }
+
+    // Compute scale
+    Tx2 amax_2x = Tx2{0.0f, 0.0f};
+#pragma unroll
+    for (int r = 0; r < group_size; r += 2) {
+      auto pair = Tx2{thread_data[r], thread_data[r + 1]};
+      abs_max_x2<Tx2>(amax_2x, amax_2x, pair);
+    }
+    float scale =
+        max(fabsf(static_cast<float>(amax_2x.x)),
+            fabsf(static_cast<float>(amax_2x.y)));
+    scale /= (bits == 4) ? 6.0f : 448.0f;
+    using ScaleType =
+        std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+    auto s = ScaleType(scale);
+    scale = float(s);
+    scales_smem[tidx][tidy] = s.__x;
+
+    int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group;
+
+#pragma unroll
+    for (int j = 0; j < group_size / 4; j++) {
+      Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&thread_data[j * 4]);
+      if constexpr (bits == 8) {
+        uint32_t quantized_val =
+            scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+        *reinterpret_cast<uint32_t*>(&quantized_smem[shared_idx + j * 4]) =
+            quantized_val;
+      } else {
+        uint16_t quantized_val =
+            scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+        *reinterpret_cast<uint16_t*>(&quantized_smem[shared_idx + j * 2]) =
+            quantized_val;
+      }
+    }
+  }
+  __syncthreads();
+
+  int output_cols = K / elem_per_byte;
+  int num_groups_per_row = K / group_size;
+  int linear_tid = tidx + tidy * BLOCK_X;
+  // Write back quantized values
+#pragma unroll
+  for (int i = linear_tid; i < bytes_per_block; i += BLOCK_X * BLOCK_Y) {
+    int local_row = i / local_cols;
+    int local_col = i % local_cols;
+
+    int global_row = bidy * rows_per_block + local_row;
+    int global_col = bidx * local_cols + local_col;
+
+    if (global_row < M && global_col < output_cols) {
+      int physical_idx = local_row * padded_local_cols + local_col;
+      out[global_row * output_cols + global_col] = quantized_smem[physical_idx];
+    }
+  }
+  // Write back scales
+  constexpr int num_scales = BLOCK_X * BLOCK_Y;
+#pragma unroll
+  for (int i = linear_tid; i < num_scales; i += BLOCK_X * BLOCK_Y) {
+    int local_row = i / BLOCK_Y;
+    int local_col = i % BLOCK_Y;
+
+    int global_row = bidy * BLOCK_X + local_row;
+    int global_col = bidx * BLOCK_Y + local_col;
+
+    if (global_row < M && global_col < num_groups_per_row) {
+      scales[global_row * num_groups_per_row + global_col] =
+          scales_smem[local_row][local_col];
+    }
+  }
+}
+
+template <typename T, int group_size, int bits, bool use_mx_scale>
+__global__ void
+fp_dequantize(const uint8_t* w, const uint8_t* scales, T* out, size_t size) {
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+
+  constexpr int pack_factor = bits == 8 ? 1 : 2;
+  size_t offset = tidx + grid_dim_x * size_t(tidy);
+  size_t oindex = offset * pack_factor;
+
+  if (oindex >= size) {
+    return;
+  }
+
+  size_t gindex = oindex / group_size;
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto scale = float(((ScaleType*)(scales))[gindex]);
+
+  out += oindex;
+
+  uint32_t val = w[offset];
+#pragma clang loop unroll(full)
+  for (int i = 0; i < pack_factor; i++) {
+    uint8_t d;
+    if (bits == 4) {
+      d = (val >> (bits * i)) & 0x0f;
+    } else if (bits == 8) {
+      d = val;
+    }
+    out[i] = static_cast<T>(scale * Dequantize<bits>{}(d));
+  }
+}
+
+inline std::tuple<dim3, dim3>
+get_columnwise_quantize_launch_args(size_t size, int group_size, int M, int K) {
+  constexpr int BLOCK_X = 32;
+  constexpr int BLOCK_Y = 32;
+  int rows_per_block = BLOCK_X;
+  int cols_per_block = BLOCK_Y * group_size;
+
+  dim3 grid;
+  grid.x =
+      cuda::ceil_div(K, cols_per_block) * cuda::ceil_div(M, rows_per_block);
+  grid.y = 1;
+  grid.z = 1;
+
+  dim3 block(BLOCK_X, BLOCK_Y);
+
+  return std::make_tuple(grid, block);
+}
+
+} // namespace cu
+
+void fp_quantize_dequantize(
+    const array& w,
+    array& what,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  enc.set_input_array(w);
+  enc.set_output_array(what);
+  dispatch_float_types(w.dtype(), "fp_quantize_dequantize", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    if constexpr (!std::is_same_v<T, double>) {
+      auto kernel = cu::fp_quantize_dequantize<T, 32, 4, true, false>;
+      if (bits == 8) {
+        kernel = cu::fp_quantize_dequantize<T, 32, 8, true, false>;
+      } else if (group_size == 16) {
+        kernel = cu::fp_quantize_dequantize<T, 16, 4, false, false>;
+      }
+      bool large = w.size() > UINT_MAX;
+      auto [num_blocks, block_dims] =
+          get_launch_args(w.size(), w.shape(), w.strides(), large, group_size);
+
+      enc.add_kernel_node(
+          kernel,
+          num_blocks,
+          block_dims,
+          0,
+          gpu_ptr<T>(w),
+          gpu_ptr<T>(what),
+          w.size());
+    }
+  });
+}
+
+void fp_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  enc.set_input_array(w);
+  enc.set_output_array(wq);
+  enc.set_output_array(scales);
+  if (w.strides().back() != 1) {
+    dispatch_float_types(w.dtype(), "fp_quantize_columnwise", [&](auto type_tag) {
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      if constexpr (!std::is_same_v<T, double>) {
+        auto M = w.shape(-2);
+        auto K = w.shape(-1);
+        auto kernel = cu::fp_quantize_columnwise<T, 32, 4, true, false>;
+        if (bits == 8) {
+          kernel = cu::fp_quantize_columnwise<T, 32, 8, true, false>;
+        } else if (group_size == 16) {
+          kernel = cu::fp_quantize_columnwise<T, 16, 4, false, false>;
+        }
+        auto [num_blocks, block_dims] =
+            cu::get_columnwise_quantize_launch_args(w.size(), group_size, M, K);
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<T>(w),
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<uint8_t>(scales),
+            w.size(),
+            M,
+            K);
+      } else {
+        throw std::runtime_error(
+            "[Quantize::eval_gpu] Can not quantize input with type float64.");
+      }
+    });
+  } else {
+    dispatch_float_types(w.dtype(), "fp_quantize_rowwise", [&](auto type_tag) {
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      if constexpr (!std::is_same_v<T, double>) {
+        auto kernel = cu::fp_quantize_rowwise<T, 32, 4, true, false>;
+        if (bits == 8) {
+          kernel = cu::fp_quantize_rowwise<T, 32, 8, true, false>;
+        } else if (group_size == 16) {
+          kernel = cu::fp_quantize_rowwise<T, 16, 4, false, false>;
+        }
+        bool large = w.size() > UINT_MAX;
+        auto [num_blocks, block_dims] = get_launch_args(
+            w.size(), w.shape(), w.strides(), large, group_size);
+
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<T>(w),
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<uint8_t>(scales),
+            w.size());
+      } else {
+        throw std::runtime_error(
+            "[Quantize::eval_gpu] Can not quantize input with type float64.");
+      }
+    });
+  }
+}
+
+void fp_dequantize(
+    const array& wq,
+    const array& scales,
+    array& w,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  constexpr int uint8_per_uint32 = 4;
+  int packs_per_int = 8 / bits;
+
+  size_t size = w.size() / packs_per_int;
+  bool large = size > UINT_MAX;
+  auto grid_shape = w.shape();
+  grid_shape.back() *= uint8_per_uint32;
+
+  enc.set_input_array(wq);
+  enc.set_input_array(scales);
+  enc.set_output_array(w);
+  dispatch_float_types(w.dtype(), "fp_dequantize", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    if constexpr (!std::is_same_v<T, double>) {
+      auto kernel = cu::fp_dequantize<T, 32, 4, true>;
+      if (bits == 8) {
+        kernel = cu::fp_dequantize<T, 32, 8, true>;
+      } else if (group_size == 16) {
+        kernel = cu::fp_dequantize<T, 16, 4, false>;
+      }
+      auto [num_blocks, block_dims] =
+          get_launch_args(size, grid_shape, w.strides(), large);
+      enc.add_kernel_node(
+          kernel,
+          num_blocks,
+          block_dims,
+          0,
+          gpu_ptr<uint8_t>(wq),
+          gpu_ptr<uint8_t>(scales),
+          gpu_ptr<T>(w),
+          w.size());
+    } else {
+      throw std::runtime_error(
+          "[Quantize::eval_gpu] Can not dequantize to output with type float64.");
+    }
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/mxfp8_quantize.cuh b/macos/third_party/mlx/mlx/backend/cuda/quantized/mxfp8_quantize.cuh
new file mode 100644
index 0000000..1edf364
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/mxfp8_quantize.cuh
@@ -0,0 +1,32 @@
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include "mlx/backend/cuda/vector_types.cuh"
+
+namespace mlx::core::cu {
+
+// TODO implement fast path
+template <typename T>
+__device__ __forceinline__ uint32_t
+scale_cvt_Tx4_to_fp8x4_fallback(const Vector4_t<T> input, const float scale) {
+  uint32_t out_fp8x4 = 0;
+  float4 scaled;
+  scaled.x = static_cast<float>(input.x) * scale;
+  scaled.y = static_cast<float>(input.y) * scale;
+  scaled.z = static_cast<float>(input.z) * scale;
+  scaled.w = static_cast<float>(input.w) * scale;
+  out_fp8x4 = __nv_fp8x4_e4m3(scaled).__x;
+  return out_fp8x4;
+}
+
+// Place holder for future fast path implementation
+template <typename T, bool USE_SR>
+__device__ __forceinline__ uint32_t scale_cvt_Tx4_to_fp8x4(
+    const Vector4_t<T> input,
+    const float scale,
+    uint32_t rbits) {
+  return scale_cvt_Tx4_to_fp8x4_fallback(input, scale);
+}
+} // namespace mlx::core::cu
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp b/macos/third_party/mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp
new file mode 100644
index 0000000..5552e24
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp
@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder&,
+    int,
+    int,
+    int,
+    bool,
+    int64_t,
+    bool,
+    int64_t,
+    array&,
+    const array&,
+    const array&,
+    const array&,
+    const array&,
+    Dtype,
+    QuantizationMode,
+    float) {
+  throw std::runtime_error(
+      "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher.");
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/nvfp4_quantize.cuh b/macos/third_party/mlx/mlx/backend/cuda/quantized/nvfp4_quantize.cuh
new file mode 100644
index 0000000..da0df42
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/nvfp4_quantize.cuh
@@ -0,0 +1,334 @@
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp4.h>
+#include <cuda_runtime.h>
+#include "mlx/backend/cuda/vector_types.cuh"
+
+namespace mlx::core::cu {
+
+using bf16x4 = Vector4_t<__nv_bfloat16>;
+using fp16x4 = Vector4_t<__half>;
+using f32x4 = Vector4_t<float>;
+
+template <typename T>
+__device__ __forceinline__ uint16_t
+scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t<T> input, const float scale) {
+  // Fallback implementation for architectures that do not support cvt
+  // instructions or for cuda versions with no fp4 support (< 12.8) -> scalar
+  uint16_t out_fp4x4 = 0;
+  fp32x4 scaled;
+  scaled.x = static_cast<float>(input.x) * scale;
+  scaled.y = static_cast<float>(input.y) * scale;
+  scaled.z = static_cast<float>(input.z) * scale;
+  scaled.w = static_cast<float>(input.w) * scale;
+  uint8_t q0 = __nv_fp4_e2m1(scaled.x).__x;
+  uint8_t q1 = __nv_fp4_e2m1(scaled.y).__x;
+  uint8_t q2 = __nv_fp4_e2m1(scaled.z).__x;
+  uint8_t q3 = __nv_fp4_e2m1(scaled.w).__x;
+  out_fp4x4 = (static_cast<uint16_t>(q3) << 12) |
+      (static_cast<uint16_t>(q2) << 8) | (static_cast<uint16_t>(q1) << 4) |
+      static_cast<uint16_t>(q0);
+  return out_fp4x4;
+}
+
+#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
+    defined(__CUDA_ARCH_SPECIFIC__)
+
+__device__ __forceinline__ uint16_t
+scale_cvt_bf16x4_to_fp4x4_rn(const bf16x4 input_bf16x4, const float2 scale) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_bf16; \n\t" // first bf16
+      ".reg.b16 x1_bf16; \n\t" // second bf16
+      ".reg.b16 x2_bf16; \n\t" // third bf16
+      ".reg.b16 x3_bf16; \n\t" // fourth bf16
+      ".reg.b32 x0; \n\t" // to hold scaled first
+      ".reg.b32 x1; \n\t" // to hold scaled second
+      ".reg.b32 x2; \n\t" // to hold scaled third
+      ".reg.b32 x3; \n\t" // to hold scaled fourth
+      ".reg.b64 x01; \n\t" // to hold vector mul
+      ".reg.b64 x23; \n\t"
+      ".reg.b8 q0; \n\t" // output byte fp4x2 (first pair)
+      ".reg.b8 q1; \n\t" // output byte fp4x2 (second pair)
+      "mov.b64 {x0_bf16, x1_bf16, x2_bf16, x3_bf16} , %1; \n\t" // unpack bf16
+      "cvt.f32.bf16 x0, x0_bf16; \n\t" // convert to f32
+      "cvt.f32.bf16 x1, x1_bf16; \n\t"
+      "cvt.f32.bf16 x2, x2_bf16; \n\t"
+      "cvt.f32.bf16 x3, x3_bf16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t" // scale first pair
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t" // scale second pair
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q0, x1, x0; \n\t" // convert to fp4x2 first
+                                                     // pair
+      "cvt.rn.satfinite.e2m1x2.f32 q1, x3, x2; \n\t" // convert to fp4x2 second
+                                                     // pair
+      "mov.b16 %0, {q0, q1}; \n\t" // pack to output
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_bf16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(
+            scale))); // here cast is needed becuase an asm operand must have
+                      // scalar type
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_bf16x4_to_fp4x4_rs(
+    const bf16x4 input_bf16x4,
+    const float2 scale,
+    uint32_t rbits) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_bf16; \n\t"
+      ".reg.b16 x1_bf16; \n\t"
+      ".reg.b16 x2_bf16; \n\t"
+      ".reg.b16 x3_bf16; \n\t"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b16 q0; \n\t"
+      "mov.b64 {x0_bf16, x1_bf16, x2_bf16, x3_bf16} , %1; \n\t"
+      "cvt.f32.bf16 x0, x0_bf16; \n\t"
+      "cvt.f32.bf16 x1, x1_bf16; \n\t"
+      "cvt.f32.bf16 x2, x2_bf16; \n\t"
+      "cvt.f32.bf16 x3, x3_bf16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t"
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rs.satfinite.e2m1x4.f32 q0, {x3, x2, x1, x0}, %3; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_bf16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(scale)),
+        "r"(rbits));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_fp32x4_to_fp4x4_rn(
+    const float2 input_fp32x2_0,
+    const float2 input_fp32x2_1,
+    const float2 scale) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b8 q0; \n\t"
+      ".reg.b8 q1; \n\t"
+      "mov.b64 x01, {%1, %2}; \n\t"
+      "mul.f32x2 x01, x01, %5; \n\t"
+      "mov.b64 x23, {%3, %4}; \n\t"
+      "mul.f32x2 x23, x23, %5; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q0, x1, x0; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q1, x3, x2; \n\t"
+      "mov.b16 %0, {q0, q1}; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "f"(input_fp32x2_0.x),
+        "f"(input_fp32x2_0.y),
+        "f"(input_fp32x2_1.x),
+        "f"(input_fp32x2_1.y),
+        "l"(reinterpret_cast<const uint64_t&>(scale)));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_fp32x4_to_fp4x4_rs(
+    const float2 input_fp32x2_0,
+    const float2 input_fp32x2_1,
+    const float2 scale,
+    uint32_t rbits) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b16 q0; \n\t"
+      "mov.b64 x01, {%1, %2}; \n\t"
+      "mul.f32x2 x01, x01, %5; \n\t"
+      "mov.b64 x23, {%3, %4}; \n\t"
+      "mul.f32x2 x23, x23, %5; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rs.satfinite.e2m1x4.f32 q0, {x3, x2, x1, x0}, %6; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "f"(input_fp32x2_0.x),
+        "f"(input_fp32x2_0.y),
+        "f"(input_fp32x2_1.x),
+        "f"(input_fp32x2_1.y),
+        "l"(reinterpret_cast<const uint64_t&>(scale)),
+        "r"(rbits));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t
+scale_cvt_fp16x4_to_fp4x4_rn(const fp16x4 input_fp16x4, const float2 scale) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_fp16; \n\t"
+      ".reg.b16 x1_fp16; \n\t"
+      ".reg.b16 x2_fp16; \n\t"
+      ".reg.b16 x3_fp16; \n\t"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b8 q0; \n\t"
+      ".reg.b8 q1; \n\t"
+      "mov.b64 {x0_fp16, x1_fp16, x2_fp16, x3_fp16} , %1; \n\t"
+      "cvt.f32.f16 x0, x0_fp16; \n\t"
+      "cvt.f32.f16 x1, x1_fp16; \n\t"
+      "cvt.f32.f16 x2, x2_fp16; \n\t"
+      "cvt.f32.f16 x3, x3_fp16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t"
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q0, x1, x0; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q1, x3, x2; \n\t"
+      "mov.b16 %0, {q0, q1}; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_fp16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(scale)));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_fp16x4_to_fp4x4_rs(
+    const fp16x4 input_fp16x4,
+    const float2 scale,
+    uint32_t rbits) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_fp16; \n\t"
+      ".reg.b16 x1_fp16; \n\t"
+      ".reg.b16 x2_fp16; \n\t"
+      ".reg.b16 x3_fp16; \n\t"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b16 q0; \n\t"
+      "mov.b64 {x0_fp16, x1_fp16, x2_fp16, x3_fp16} , %1; \n\t"
+      "cvt.f32.f16 x0, x0_fp16; \n\t"
+      "cvt.f32.f16 x1, x1_fp16; \n\t"
+      "cvt.f32.f16 x2, x2_fp16; \n\t"
+      "cvt.f32.f16 x3, x3_fp16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t"
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rs.satfinite.e2m1x4.f32 q0, {x3, x2, x1, x0}, %3; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_fp16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(scale)),
+        "r"(rbits));
+  return out_fp4x4;
+}
+
+template <bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_bf16x4_to_fp4x4(
+    const bf16x4 input,
+    const float scale,
+    uint32_t rbits) {
+  float2 scale_fp32x2 = make_float2(scale, scale);
+  if constexpr (USE_SR) {
+    return scale_cvt_bf16x4_to_fp4x4_rs(input, scale_fp32x2, rbits);
+  } else {
+    return scale_cvt_bf16x4_to_fp4x4_rn(input, scale_fp32x2);
+  }
+}
+
+template <bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_fp16x4_to_fp4x4(
+    const fp16x4 input,
+    const float scale,
+    uint32_t rbits) {
+  float2 scale_fp32x2 = make_float2(scale, scale);
+  if constexpr (USE_SR) {
+    return scale_cvt_fp16x4_to_fp4x4_rs(input, scale_fp32x2, rbits);
+  } else {
+    return scale_cvt_fp16x4_to_fp4x4_rn(input, scale_fp32x2);
+  }
+}
+
+template <bool USE_SR>
+__device__ __forceinline__ uint16_t
+scale_cvt_f32x4_to_fp4x4(const f32x4 input, const float scale, uint32_t rbits) {
+  float2 scale_fp32x2 = make_float2(scale, scale);
+  float2 input_fp32x2_0 = make_float2(input.x, input.y);
+  float2 input_fp32x2_1 = make_float2(input.z, input.w);
+
+  if constexpr (USE_SR) {
+    return scale_cvt_fp32x4_to_fp4x4_rs(
+        input_fp32x2_0, input_fp32x2_1, scale_fp32x2, rbits);
+  } else {
+    return scale_cvt_fp32x4_to_fp4x4_rn(
+        input_fp32x2_0, input_fp32x2_1, scale_fp32x2);
+  }
+}
+
+template <typename T, bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4_fast(
+    const Vector4_t<T> input,
+    const float scale,
+    uint32_t rbits) {
+  if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+    return scale_cvt_bf16x4_to_fp4x4<USE_SR>(input, scale, rbits);
+  } else if constexpr (std::is_same<T, __half>::value) {
+    return scale_cvt_fp16x4_to_fp4x4<USE_SR>(input, scale, rbits);
+  } else {
+    return scale_cvt_f32x4_to_fp4x4<USE_SR>(input, scale, rbits);
+  }
+}
+#endif // (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) &&
+       // (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)
+
+template <typename T, bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4(
+    const Vector4_t<T> input,
+    const float scale,
+    uint32_t rbits) {
+#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
+    (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)
+  return scale_cvt_Tx4_to_fp4x4_fast<T, USE_SR>(input, scale, rbits);
+#else
+  static_assert(
+      !USE_SR,
+      "Stochastic rounding (USE_SR=true) requires CUDA >= 12.8 and compute capability >= 1000.");
+  return scale_cvt_Tx4_to_fp4x4_fallback(input, scale);
+#endif
+}
+} // namespace mlx::core::cu
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/qmv.cu b/macos/third_party/mlx/mlx/backend/cuda/quantized/qmv.cu
new file mode 100644
index 0000000..0cc0046
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/qmv.cu
@@ -0,0 +1,304 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/utils.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+static constexpr int rows_per_block = 8;
+
+template <typename T>
+__device__ void adjust_matrix_offsets(
+    const T*& x,
+    const uint32_t*& w,
+    const uint8_t*& scales,
+    T*& y,
+    int output_stride,
+    const int& x_batch_ndims,
+    const Shape x_shape,
+    const Strides x_strides,
+    const int& w_batch_ndims,
+    const Shape w_shape,
+    const Strides w_strides,
+    const Strides s_strides) {
+  uint32_t idx = cg::this_grid().block_index().z;
+  if (x_batch_ndims == 1) {
+    x += idx * x_strides[0];
+  } else {
+    x += elem_to_loc(idx, x_shape.data(), x_strides.data(), x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += idx * w_strides[0];
+    scales += idx * s_strides[0];
+  } else {
+    auto [w_idx, s_idx] = elem_to_loc(
+        idx, w_shape.data(), w_strides.data(), s_strides.data(), w_batch_ndims);
+    w += w_idx;
+    scales += s_idx;
+  }
+  y += idx * output_stride;
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__device__ void fp_qmv_impl(
+    const uint32_t* mat,
+    const uint8_t* scales_,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols) {
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  constexpr int vals_per_item = bits == 8 ? 4 : 8;
+  constexpr int nv_per_thread = vals_per_item * n_per_thread;
+  auto g_idx = block.group_index();
+  auto t_idx = block.thread_index();
+  int row = g_idx.y * rows_per_block + t_idx.y;
+
+  vec += g_idx.x * cols;
+  out += g_idx.x * rows;
+
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto scales = (ScaleType*)(scales_);
+  auto packed_cols = cols / vals_per_item;
+
+  if (row < rows) {
+    constexpr int scales_per_step = std::max(nv_per_thread / group_size, 1);
+    constexpr int scale_step = (WARP_SIZE * nv_per_thread) / group_size;
+    constexpr int n_per_step = n_per_thread / scales_per_step;
+    // Offset scales to correct row
+    scales += row * (cols / group_size) +
+        (warp.thread_rank() * nv_per_thread) / group_size;
+    float sum = 0.0f;
+    for (int col = n_per_thread * warp.thread_rank(); col < packed_cols;
+         col += (WARP_SIZE * n_per_thread)) {
+      auto local_vec =
+          unsafe_load_vector<nv_per_thread>(vec + vals_per_item * col, 0);
+      auto local_mat =
+          unsafe_load_vector<n_per_thread>(mat + row * packed_cols + col, 0);
+#pragma unroll
+      for (int i = 0; i < scales_per_step; ++i) {
+        float2 local_sum = {0.0f, 0.0f};
+#pragma unroll
+        for (int j = 0; j < n_per_step; ++j) {
+          int k = n_per_step * i + j;
+          if constexpr (bits == 8) {
+            auto v = dequant_fp8(local_mat[k]);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+            local_sum.x +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+            local_sum.y +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
+          } else {
+            auto v = dequant_fp4(local_mat[k]);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+            local_sum.y +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+            local_sum.x +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
+
+            v = dequant_fp4(local_mat[k] >> 16);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k + 4]);
+            local_sum.y +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 5]);
+            local_sum.x +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 6]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 7]);
+          }
+        }
+        sum += (local_sum.x + local_sum.y) * float(scales[i]);
+      }
+      scales += scale_step;
+    }
+
+    sum = cg::reduce(warp, sum, cg::plus<float>{});
+    if (warp.thread_rank() == 0) {
+      out[row] = static_cast<T>(sum);
+    }
+  }
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__global__ void fp_qmv_single(
+    const uint32_t* mat,
+    const uint8_t* scales,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols) {
+  fp_qmv_impl<T, rows_per_block, n_per_thread, bits, group_size, use_mx_scale>(
+      mat, scales, vec, out, rows, cols);
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__global__ void fp_qmv_batched(
+    const uint32_t* mat,
+    const uint8_t* scales,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols,
+    int vec_batch_ndims,
+    const __grid_constant__ Shape vec_shape,
+    const __grid_constant__ Strides vec_strides,
+    int mat_batch_ndims,
+    const __grid_constant__ Shape mat_shape,
+    const __grid_constant__ Strides mat_strides,
+    const __grid_constant__ Strides scales_strides) {
+  adjust_matrix_offsets<T>(
+      vec,
+      mat,
+      scales,
+      out,
+      rows * vec_shape[vec_batch_ndims],
+      vec_batch_ndims,
+      vec_shape,
+      vec_strides,
+      mat_batch_ndims,
+      mat_shape,
+      mat_strides,
+      scales_strides);
+  fp_qmv_impl<T, rows_per_block, n_per_thread, bits, group_size, use_mx_scale>(
+      mat, scales, vec, out, rows, cols);
+}
+
+template <typename F>
+void dispatch_1_2_4(int n, F&& f) {
+  switch (n) {
+    case 1:
+      f(std::integral_constant<int, 1>{});
+      break;
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 4:
+      f(std::integral_constant<int, 4>{});
+      break;
+  }
+}
+
+void fp_qmv(
+    const array& mat,
+    const array& scales,
+    const array& vec,
+    array& out,
+    int bits,
+    int group_size,
+    int M,
+    int N,
+    int K,
+    CommandEncoder& encoder) {
+  encoder.set_input_array(mat);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(vec);
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "qmv", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    if constexpr (!std::is_same_v<T, double>) {
+      dim3 block_dims{WARP_SIZE, rows_per_block};
+      uint32_t B = out.size() / (M * N);
+      uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block;
+      const uint32_t* mat_ptr = gpu_ptr<uint32_t>(mat);
+      const T* vec_ptr = gpu_ptr<T>(vec);
+      int n = 1;
+      if (K % 32 == 0 && cu::is_aligned<4>(mat_ptr) &&
+          ((bits == 4 && cu::is_aligned<8>(vec_ptr)) ||
+           cu::is_aligned<4>(vec_ptr))) {
+        n = 4;
+      } else if (
+          cu::is_aligned<2>(mat_ptr) &&
+          ((bits == 4 && cu::is_aligned<4>(vec_ptr)) ||
+           cu::is_aligned<2>(vec_ptr))) {
+        n = 2;
+      }
+      dispatch_1_2_4(n, [&](auto n) {
+        dispatch_bool(B > 1, [&](auto batched) {
+          if (!batched.value) {
+            auto kernel =
+                fp_qmv_single<T, rows_per_block, n.value, 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_single<T, rows_per_block, n.value, 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_single<T, rows_per_block, n.value, 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint32_t>(M), blocks_y},
+                block_dims,
+                0,
+                mat_ptr,
+                gpu_ptr<uint8_t>(scales),
+                vec_ptr,
+                gpu_ptr<T>(out),
+                N,
+                K);
+          } else {
+            auto kernel =
+                fp_qmv_batched<T, rows_per_block, n.value, 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_batched<T, rows_per_block, n.value, 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_batched<T, rows_per_block, n.value, 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint32_t>(M), blocks_y, B},
+                block_dims,
+                0,
+                mat_ptr,
+                gpu_ptr<uint8_t>(scales),
+                vec_ptr,
+                gpu_ptr<T>(out),
+                N,
+                K,
+                vec.ndim() - 2,
+                const_param(vec.shape()),
+                const_param(vec.strides()),
+                mat.ndim() - 2,
+                const_param(mat.shape()),
+                const_param(mat.strides()),
+                const_param(scales.strides()));
+          }
+        });
+      });
+    }
+  });
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/qmv.h b/macos/third_party/mlx/mlx/backend/cuda/quantized/qmv.h
new file mode 100644
index 0000000..d4f0419
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/qmv.h
@@ -0,0 +1,21 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core::cu {
+
+void fp_qmv(
+    const array& w,
+    const array& scales,
+    const array& vec,
+    array& out,
+    int bits,
+    int group_size,
+    int M,
+    int N,
+    int K,
+    CommandEncoder& encoder);
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm.cpp b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm.cpp
new file mode 100644
index 0000000..3b5be77
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm.cpp
@@ -0,0 +1,158 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+#include "mlx/backend/cuda/quantized/qqmm_utils.h"
+#include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace {
+
+array pad_and_swizzle_scales(
+    const array& scale,
+    cu::CommandEncoder& encoder,
+    const Stream& s) {
+  // Compute padded dimensions for full tiles (128 rows × 4 cols)
+  auto [pad_outer, pad_inner] =
+      get_padded_scale_dims(scale.shape(-2), scale.shape(-1));
+  // cuBLAS requirements for scale factor layout:
+  // 1. Dimensions must be padded to full tiles (128 rows × 4 cols)
+  // 2. Out-of-bounds values must be filled with zeros
+  // 3. Starting addresses must be 16-byte aligned
+  //
+  // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
+  // Note: cu::malloc_async already provides 256-byte alignment
+  array scale_tiled(
+      cu::malloc_async(pad_outer * pad_inner, encoder),
+      Shape{pad_outer, pad_inner},
+      scale.dtype());
+  swizzle_scales(scale, scale_tiled, encoder, s);
+
+  encoder.add_temporary(scale_tiled);
+  return scale_tiled;
+}
+
+} // namespace
+
+void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(
+      (inputs.size() == 3 && inputs[1].dtype() == uint32) ||
+      (inputs.size() == 2));
+  nvtx3::scoped_range r("QQMatmul::eval_gpu");
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  auto& device = encoder.device();
+
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+    bool donate_x = inputs[0].is_donatable();
+    array x = ensure_row_contiguous(inputs[0], encoder, s);
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype());
+    if (!donate_x) {
+      encoder.add_temporary(xhat);
+    }
+    fp_quantize_dequantize(x, xhat, group_size_, bits_, encoder, s);
+
+    // Make sure the last two dims of w and s are contiguous
+    array w = ensure_row_contiguous_matrix(inputs[1], encoder, s);
+    array scales = ensure_row_contiguous_matrix(inputs[2], encoder, s);
+
+    bool non_batched = w.ndim() == 2;
+    int K = x.shape(-1);
+    int M = non_batched ? x.size() / K : x.shape(-2);
+    int N = out.shape(-1);
+
+    fp_qmv(w, scales, xhat, out, bits_, group_size_, M, N, K, encoder);
+    return;
+  }
+
+  auto cc = device.compute_capability_major() * 100 +
+      device.compute_capability_minor() * 10;
+  if (cc < 1000) {
+    throw std::runtime_error(
+        "[QQMatmul::eval_gpu] QQMM is only supported on GPUs with compute capability 10.0 or higher.");
+  }
+  auto quantize = [&](const array& input,
+                      cu::CommandEncoder& encoder,
+                      const Stream& s) -> std::pair<array, array> {
+    auto x = ensure_contiguous(input, encoder, s);
+    auto xq_shape = x.shape();
+    xq_shape.back() = x.shape(-1) * bits_ / 32;
+
+    auto sshape = x.shape();
+    const int64_t scales_inner = x.shape(-1) / group_size_;
+    auto [pad_outer, pad_inner] =
+        get_padded_scale_dims(x.shape(-2), scales_inner);
+    sshape[x.ndim() - 2] = pad_outer;
+    sshape[x.ndim() - 1] = pad_inner;
+    sshape.back() = scales_inner;
+
+    // Allocate outputs
+    const int64_t xq_bytes = x.size() * bits_ / 8;
+    const int64_t batch = x.size() / (x.shape(-2) * x.shape(-1));
+    const int64_t scales_bytes = batch * (pad_outer * pad_inner);
+
+    array x_q(cu::malloc_async(xq_bytes, encoder), std::move(xq_shape), uint32);
+    array scales_x(
+        cu::malloc_async(scales_bytes, encoder), std::move(sshape), uint8);
+
+    fp_quantize(x, x_q, scales_x, group_size_, bits_, encoder, s);
+
+    encoder.add_temporary(x_q);
+    encoder.add_temporary(scales_x);
+    return {x_q, scales_x};
+  };
+  auto [x_q, scale_x_pre] = quantize(inputs[0], encoder, s);
+  auto [w_q, scale_w_pre] = !w_quantized ? quantize(inputs[1], encoder, s)
+                                         : std::make_pair(inputs[1], inputs[2]);
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  auto out_dtype = out.dtype();
+
+  int M = x_q.shape(-2);
+  int N = w_q.shape(-2); // always transposed
+  int K_packed = x_q.shape(-1);
+  int K = K_packed * (32 / bits_);
+
+  // Repack scales from linear to tiled layout for tensor cores
+  array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s);
+  array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s);
+
+  bool x_transposed = false;
+  bool w_transposed = true; // always transposed
+  int64_t lda = K;
+  int64_t ldb = K;
+
+  qqmm_impl(
+      encoder,
+      M,
+      N,
+      K,
+      x_transposed,
+      lda,
+      w_transposed,
+      ldb,
+      out,
+      x_q,
+      w_q,
+      scale_x,
+      scale_w,
+      out_dtype,
+      mode_);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp
new file mode 100644
index 0000000..dd9407d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp
@@ -0,0 +1,50 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+
+namespace mlx::core {
+
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha) {
+  // Invoke CublasQQMM
+  std::string qmode = quantization_mode_to_string(mode);
+
+  // Currently only supports non-batched QQMM operations
+  // that covers all use cases for training, we will just collapse (batch,
+  // seq_len) into (tokens)
+  CublasQQMM qqmm(
+      encoder.device(),
+      a_transposed,
+      M,
+      K,
+      lda,
+      b_transposed,
+      K,
+      N,
+      ldb,
+      1, // batch_count
+      0, // a_batch_stride
+      0, // b_batch_stride
+      out_dtype,
+      qmode);
+
+  qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_impl.h b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_impl.h
new file mode 100644
index 0000000..7288e2f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_impl.h
@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha = 1.0f);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu
new file mode 100644
index 0000000..c876470
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu
@@ -0,0 +1,227 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qqmm_utils.h"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cg = cooperative_groups;
+
+constexpr int TILE_ROWS = 128;
+constexpr int TILE_COLS = 4;
+constexpr int TILES_PER_LANE = 1;
+constexpr int LANES_PER_BLOCK = 32;
+
+// To pass scales to tensor cores, they need to be repacked into a tiled layout
+// https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
+// Tiled layout for scale factors is very well described in CUTLASS
+// documentation:
+// https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/blackwell_functionality.md#scale-factor-layouts
+// Conceptually, it should be like this:
+// q_w = mx.zeros(shape=(M, N)) <-- zeros just for an example
+// s.shape = (M, N // 16) -- packed in row contigous order, group_size = 16
+// cbg_cnt = N // 16 // 4
+// rb_cnt = M // 128
+// tmp = x.reshape(rb_cnt, 4, 32, cbg_cnt, 4)
+// repacked_scales = tmp.transpose(0, 3, 2, 1, 4)
+// example: indecis of intial tile 128 x 4 of scales (packed in row major tensor
+// (M, K // 16), where M = 128, K = 64): array([[0, 1, 2, 3],
+//       [4, 5, 6, 7],
+//       [8, 9, 10, 11],
+//       ...,
+//       [500, 501, 502, 503],
+//       [504, 505, 506, 507],
+//       [508, 509, 510, 511]]
+// packed scales within tile 128 x 4:
+// array([[[[[0, 1, 2, 3], <-- s_0,0..s_0,3 scales
+//          [128, 129, 130, 131], <-- s_32,0..s_32,3 scales
+//          [256, 257, 258, 259], <-- s_64,0..s_64,3 scales
+//          [384, 385, 386, 387]], <-- s_96,0..s_96,3 scales
+//         [[4, 5, 6, 7], <-- s_1,0..s_1,3 scales
+//          [132, 133, 134, 135], ...
+//          [260, 261, 262, 263],
+//          [388, 389, 390, 391]],
+//         [[124, 125, 126, 127],
+//          [252, 253, 254, 255],
+//          [380, 381, 382, 383],
+//          [508, 509, 510, 511]]]]],
+
+inline std::tuple<dim3, dim3> get_swizzle_launch_args(
+    size_t M_swizzled,
+    size_t K_swizzled) {
+  constexpr int tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+  constexpr int warps_per_block = TILE_ROWS / 4; // 128 / 4 = 32
+
+  const int num_tiles_k = K_swizzled / TILE_COLS;
+  const int num_tiles_m = M_swizzled / TILE_ROWS;
+
+  dim3 grid;
+  grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block);
+  grid.y = num_tiles_m;
+  grid.z = 1;
+  // Block is always (32, 32) = 1024 threads
+  dim3 block(LANES_PER_BLOCK, warps_per_block, 1);
+
+  return std::make_tuple(grid, block);
+}
+
+namespace cu {
+
+__global__ void swizzle_scales(
+    const uint8_t* scales_linear,
+    uint8_t* scales_swizzled,
+    const size_t M,
+    const size_t K,
+    const size_t M_swizzled,
+    const size_t K_swizzled) {
+  constexpr int tile_size = TILE_ROWS * TILE_COLS;
+  constexpr int num_tile_rows_per_thread = 4;
+  constexpr int max_tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+
+  constexpr int tile_stride = tile_size / 16; // 32 int4s per tile
+
+  // Each thread loads 16 scales from 4 rows (stride 32) and packs them into
+  // int4. For example: thread (0, 0) loads scales at rows 0,32,64,96 of tile 0,
+  // thread (1, 0) loads rows 0,32,64,96 of of tile 1, etc.
+  // The store is strided within a warp (stride 32 int4s), so we first
+  // write to shared memory, then do a coalesced store from shared to global
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  auto tidx = idx_in_block.x;
+  auto tidy = idx_in_block.y;
+  auto linear_tid = tidy * block_size.x + tidx;
+
+  const int bid_x = block_idx.x;
+  const int bid_y = block_idx.y;
+
+  const int K_int = K_swizzled / 4;
+
+  const size_t output_offset = static_cast<size_t>(bid_y) * TILE_ROWS * K_int +
+      static_cast<size_t>(bid_x) * max_tiles_per_block * tile_size / 4;
+  int* output_block = reinterpret_cast<int*>(scales_swizzled) + output_offset;
+
+  const int grid_dim_x = cg::this_grid().dim_blocks().x;
+  const int grid_dim_y = cg::this_grid().dim_blocks().y;
+
+  int remaining = K_int - bid_x * max_tiles_per_block;
+  int tiles_in_block = min(remaining, max_tiles_per_block);
+  bool valid_tile = tidx * TILES_PER_LANE < tiles_in_block;
+
+  __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride];
+
+  // Initialize to zero for padding
+  int thread_tile_rows[num_tile_rows_per_thread] = {0};
+
+  if (valid_tile) {
+    const size_t col_base =
+        static_cast<size_t>(bid_x) * max_tiles_per_block * TILE_COLS +
+        tidx * TILE_COLS;
+
+    const bool aligned_k = (K % 4 == 0);
+
+    if (aligned_k) {
+      // fast path: K is aligned, use vectorized loads with stride K/4
+      const int K_stride = K / 4;
+      const size_t block_offset =
+          static_cast<size_t>(bid_y) * TILE_ROWS * K_stride +
+          static_cast<size_t>(bid_x) * max_tiles_per_block;
+      const int* input_block =
+          reinterpret_cast<const int*>(scales_linear) + block_offset;
+// load
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        const int thread_offset =
+            (i * block_size.x + tidy) * K_stride + tidx * TILES_PER_LANE;
+        if (row < M && col_base + TILE_COLS <= K) {
+          thread_tile_rows[i] = __ldg(input_block + thread_offset);
+        } else if (row < M) {
+// partial tile at K boundary: load byte-by-byte
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row * K + col_base + c];
+            }
+          }
+        }
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        if (row < M) {
+          const size_t row_start = row * K;
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row_start + col_base + c];
+            }
+          }
+        }
+      }
+    }
+    // store to shared with XOR swizzle to avoid bank conflicts
+    int base_idx = tidx * tile_stride + tidy;
+    int xor_bits = (tidy >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    strided_scales_thread[swizzled_idx] =
+        *reinterpret_cast<int4*>(thread_tile_rows);
+  }
+
+  cg::thread_block block = cg::this_thread_block();
+  cg::sync(block);
+
+  const int total_int4s = tiles_in_block * tile_stride;
+#pragma unroll
+  for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) {
+    int tile_idx = i / tile_stride;
+    int row_idx = i % tile_stride;
+    int base_idx = tile_idx * tile_stride + row_idx;
+    int xor_bits = (row_idx >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    reinterpret_cast<int4*>(output_block)[i] =
+        strided_scales_thread[swizzled_idx];
+  }
+}
+} // namespace cu
+
+void swizzle_scales(
+    const array& scales,
+    array& scales_tiled,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  enc.set_input_array(scales);
+  enc.set_output_array(scales_tiled);
+  // Note: scales_tiled is padded to full tiles so if num_rows or num_cols
+  // are not multiples of tile sizes
+  size_t input_rows = scales.shape(-2);
+  size_t input_cols = scales.shape(-1);
+
+  size_t output_rows = scales_tiled.shape(-2);
+  size_t output_cols = scales_tiled.shape(-1);
+
+  auto [num_blocks, block_dims] =
+      get_swizzle_launch_args(output_rows, output_cols);
+  enc.add_kernel_node(
+      cu::swizzle_scales,
+      num_blocks,
+      block_dims,
+      0,
+      gpu_ptr<uint8_t>(scales),
+      gpu_ptr<uint8_t>(scales_tiled),
+      input_rows,
+      input_cols,
+      output_rows,
+      output_cols);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_utils.h b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_utils.h
new file mode 100644
index 0000000..0a9a78f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/qqmm_utils.h
@@ -0,0 +1,30 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+// Compute padded dimensions for tiled layout
+// Tiles are 128 rows × 4 columns, must allocate full tiles
+inline std::pair<int, int> get_padded_scale_dims(int num_rows, int num_cols) {
+  constexpr int rows_per_tile = 128;
+  constexpr int cols_per_tile = 4;
+
+  int padded_rows =
+      ((num_rows + rows_per_tile - 1) / rows_per_tile) * rows_per_tile;
+  int padded_cols =
+      ((num_cols + cols_per_tile - 1) / cols_per_tile) * cols_per_tile;
+
+  return {padded_rows, padded_cols};
+}
+
+void swizzle_scales(
+    const array& scales,
+    array& scales_tiled,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized.cpp b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized.cpp
new file mode 100644
index 0000000..30662d5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized.cpp
@@ -0,0 +1,85 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("QuantizedMatmul::eval_gpu");
+  auto& s = stream();
+  auto& d = cu::device(s.device);
+  auto& enc = d.get_command_encoder(s);
+
+  out.set_data(cu::malloc_async(out.nbytes(), enc));
+
+  // Make sure the last two dims of x and w, s, b are contiguous. This should
+  // be relaxed for x.
+  array x = ensure_row_contiguous_matrix(inputs[0], enc, s);
+  array w = ensure_row_contiguous_matrix(inputs[1], enc, s);
+  array scales = ensure_row_contiguous_matrix(inputs[2], enc, s);
+  std::optional<array> biases = std::nullopt;
+  if (inputs.size() == 4) {
+    biases = ensure_row_contiguous_matrix(inputs[3], enc, s);
+  }
+
+  bool non_batched = w.ndim() == 2 && x.flags().row_contiguous;
+  int K = x.shape(-1);
+  int M = non_batched ? x.size() / K : x.shape(-2);
+  int N = out.shape(-1);
+
+  if (M > 8 || !transpose_ || mode_ == QuantizationMode::Affine) {
+    throw std::runtime_error("QMM NYI");
+  }
+
+  if (transpose_) {
+    fp_qmv(w, scales, x, out, bits_, group_size_, M, N, K, enc);
+    return;
+  }
+}
+
+void fast::Quantize::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("Quantize::eval_gpu");
+  auto& s = stream();
+  auto& d = cu::device(s.device);
+  auto& enc = d.get_command_encoder(s);
+
+  if (dequantize_) {
+    auto wq = ensure_row_contiguous(inputs[0], enc, s);
+    auto scales = ensure_row_contiguous(inputs[1], enc, s);
+    auto& w = outputs[0];
+
+    w.set_data(cu::malloc_async(w.nbytes(), enc));
+
+    if (mode_ == QuantizationMode::Affine) {
+      auto biases = ensure_row_contiguous(inputs[2], enc, s);
+      affine_dequantize(wq, scales, biases, w, group_size_, bits_, enc, s);
+    } else {
+      fp_dequantize(wq, scales, w, group_size_, bits_, enc, s);
+    }
+  } else {
+    auto w = ensure_contiguous(inputs[0], enc, s);
+    auto& wq = outputs[0];
+    auto& scales = outputs[1];
+
+    wq.set_data(cu::malloc_async(wq.nbytes(), enc));
+    scales.set_data(cu::malloc_async(scales.nbytes(), enc));
+    if (mode_ == QuantizationMode::Affine) {
+      auto& biases = outputs[2];
+      biases.set_data(cu::malloc_async(biases.nbytes(), enc));
+      affine_quantize(w, wq, scales, biases, group_size_, bits_, enc, s);
+    } else {
+      fp_quantize(w, wq, scales, group_size_, bits_, enc, s);
+    }
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized.h b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized.h
new file mode 100644
index 0000000..744a12f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized.h
@@ -0,0 +1,53 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+void affine_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    array& biases,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void affine_dequantize(
+    const array& wq,
+    const array& scales,
+    const array& biases,
+    array& w,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void fp_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void fp_dequantize(
+    const array& wq,
+    const array& scales,
+    array& w,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void fp_quantize_dequantize(
+    const array& w,
+    array& what,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized_utils.cuh b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized_utils.cuh
new file mode 100644
index 0000000..93d8329
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized_utils.cuh
@@ -0,0 +1,88 @@
+// Copyright © 2025 Apple Inc.
+
+#include <cuda_fp4.h>
+#include <cuda_fp8.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+inline __device__ float4 dequant_fp8(uint32_t bits) {
+  auto out = *(__nv_fp8x4_e4m3*)(&bits);
+  return out.operator float4();
+}
+
+inline __device__ float4 dequant_fp4(uint16_t bits) {
+  auto out = *(__nv_fp4x4_e2m1*)(&bits);
+  return out.operator float4();
+}
+
+template <int bits, int wsize = 8>
+inline constexpr __device__ short get_pack_factor() {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+template <int bits, int wsize = 8>
+inline constexpr __device__ short get_bytes_per_pack() {
+  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
+template <typename T>
+__device__ __forceinline__ void abs_max_x2(T& out, const T& x1, const T& x2) {
+  if constexpr (
+      (std::is_same<T, __nv_bfloat162>::value) ||
+      (std::is_same<T, __half2>::value)) {
+    T a = x1;
+    T b = x2;
+    out = __hmax2(__habs2(a), __habs2(b));
+  } else if constexpr (std::is_same<T, float2>::value) {
+    float2 a = x1;
+    float2 b = x2;
+    out.x = fmaxf(fabsf(a.x), fabsf(b.x));
+    out.y = fmaxf(fabsf(a.y), fabsf(b.y));
+  }
+}
+
+} // namespace cu
+
+template <typename F>
+void dispatch_groups(int group_size, F&& f) {
+  switch (group_size) {
+    case 32:
+      f(std::integral_constant<int, 32>{});
+      break;
+    case 64:
+      f(std::integral_constant<int, 64>{});
+      break;
+    case 128:
+      f(std::integral_constant<int, 128>{});
+      break;
+  }
+}
+
+template <typename F>
+void dispatch_bits(int bits, F&& f) {
+  switch (bits) {
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 3:
+      f(std::integral_constant<int, 3>{});
+      break;
+    case 4:
+      f(std::integral_constant<int, 4>{});
+      break;
+    case 5:
+      f(std::integral_constant<int, 5>{});
+      break;
+    case 6:
+      f(std::integral_constant<int, 6>{});
+      break;
+    case 8:
+      f(std::integral_constant<int, 8>{});
+      break;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized_utils.h b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized_utils.h
new file mode 100644
index 0000000..66be768
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/quantized/quantized_utils.h
@@ -0,0 +1,50 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/gpu/copy.h"
+
+namespace mlx::core {
+inline array ensure_row_contiguous(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    enc.add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+inline array ensure_row_contiguous_matrix(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  if (x.ndim() < 2) {
+    if (x.strides()[0] == 1) {
+      return x;
+    }
+  } else {
+    auto stride_0 = x.strides()[x.ndim() - 2];
+    auto stride_1 = x.strides()[x.ndim() - 1];
+    if (stride_0 == x.shape(-1) && stride_1 == 1) {
+      return x;
+    }
+  }
+  array x_copy = contiguous_copy_gpu(x, s);
+  enc.add_temporary(x_copy);
+  return x_copy;
+}
+
+inline array
+ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) {
+  if (x.flags().row_contiguous || x.flags().col_contiguous) {
+    return x;
+  }
+  array x_copy = contiguous_copy_gpu(x, s);
+  enc.add_temporary(x_copy);
+  return x_copy;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/random.cu b/macos/third_party/mlx/mlx/backend/cuda/random.cu
new file mode 100644
index 0000000..e9a310d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/random.cu
@@ -0,0 +1,202 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+#include <cassert>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+__constant__ constexpr uint32_t rotations[2][4] = {
+    {13, 15, 26, 6},
+    {17, 29, 16, 24}};
+
+union rbits {
+  uint2 val;
+  uint8_t bytes[2][4];
+};
+
+__device__ rbits threefry2x32_hash(uint2 key, uint2 count) {
+  uint32_t ks[] = {key.x, key.y, key.x ^ key.y ^ 0x1BD11BDA};
+
+  rbits v;
+  v.val.x = count.x + ks[0];
+  v.val.y = count.y + ks[1];
+
+  for (int i = 0; i < 5; ++i) {
+    for (auto r : rotations[i % 2]) {
+      v.val.x += v.val.y;
+      v.val.y = (v.val.y << r) | (v.val.y >> (32 - r));
+      v.val.y ^= v.val.x;
+    }
+    v.val.x += ks[(i + 1) % 3];
+    v.val.y += ks[(i + 2) % 3] + i + 1;
+  }
+
+  return v;
+}
+
+__global__ void rbitsc(
+    const uint32_t* keys,
+    uint8_t* out,
+    dim3 grid_dims,
+    bool odd,
+    uint32_t bytes_per_key) {
+  auto grid = cg::this_grid();
+  uint32_t thread_index = grid.thread_rank();
+  uint32_t index_x = thread_index % grid_dims.x;
+  uint32_t index_y = thread_index / grid_dims.x;
+  if (index_x >= grid_dims.x || index_y >= grid_dims.y) {
+    return;
+  }
+
+  auto kidx = 2 * index_x;
+  auto key = uint2{keys[kidx], keys[kidx + 1]};
+  auto half_size = grid_dims.y - odd;
+  out += index_x * bytes_per_key;
+  bool drop_last = odd && (index_y == half_size);
+  auto bits = threefry2x32_hash(
+      key, uint2{index_y, drop_last ? 0 : index_y + grid_dims.y});
+  size_t idx = size_t(index_y) << 2;
+  for (int i = 0; i < 4; ++i) {
+    out[idx + i] = bits.bytes[0][i];
+  }
+  if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index_y) + grid_dims.y) << 2;
+    if ((index_y + 1) == half_size && (bytes_per_key % 4) > 0) {
+      int edge_bytes = (bytes_per_key % 4);
+      for (int i = 0; i < edge_bytes; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    } else {
+      for (int i = 0; i < 4; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    }
+  }
+}
+
+__global__ void rbits(
+    const uint32_t* keys,
+    uint8_t* out,
+    dim3 grid_dims,
+    bool odd,
+    uint32_t bytes_per_key,
+    int32_t ndim,
+    const __grid_constant__ Shape key_shape,
+    const __grid_constant__ Strides key_strides) {
+  auto grid = cg::this_grid();
+  uint32_t thread_index = grid.thread_rank();
+  uint32_t index_x = thread_index % grid_dims.x;
+  uint32_t index_y = thread_index / grid_dims.x;
+  if (index_x >= grid_dims.x || index_y >= grid_dims.y) {
+    return;
+  }
+
+  auto kidx = 2 * index_x;
+  auto k1_elem = elem_to_loc(kidx, key_shape.data(), key_strides.data(), ndim);
+  auto k2_elem =
+      elem_to_loc(kidx + 1, key_shape.data(), key_strides.data(), ndim);
+  auto key = uint2{keys[k1_elem], keys[k2_elem]};
+  auto half_size = grid_dims.y - odd;
+  out += size_t(index_x) * bytes_per_key;
+  bool drop_last = odd && (index_y == half_size);
+  auto bits = threefry2x32_hash(
+      key, uint2{index_y, drop_last ? 0 : index_y + grid_dims.y});
+  size_t idx = size_t(index_y) << 2;
+  for (int i = 0; i < 4; ++i) {
+    out[idx + i] = bits.bytes[0][i];
+  }
+  if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index_y) + grid_dims.y) << 2;
+    if ((index_y + 1) == half_size && (bytes_per_key % 4) > 0) {
+      int edge_bytes = (bytes_per_key % 4);
+      for (int i = 0; i < edge_bytes; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    } else {
+      for (int i = 0; i < 4; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    }
+  }
+}
+
+} // namespace cu
+
+void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("RandomBits::eval_gpu");
+  assert(inputs.size() == 1);
+
+  // keys has shape (N1, ..., NK, 2)
+  // out has shape (N1, ..., NK, M1, M2, ...)
+  auto& keys = inputs[0];
+  size_t num_keys = keys.size() / 2;
+
+  size_t elems_per_key = out.size() / num_keys;
+  size_t bytes_per_key = out.itemsize() * elems_per_key;
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  if (out.size() == 0) {
+    return;
+  }
+
+  size_t out_per_key = (bytes_per_key + 4 - 1) / 4;
+  size_t half_size = out_per_key / 2;
+
+  bool odd = out_per_key % 2;
+  if ((half_size + odd) >= UINT32_MAX || num_keys >= UINT32_MAX) {
+    throw std::runtime_error("[RandomBits::eval_gpu] Large size unsupported");
+  }
+
+  encoder.set_input_array(keys);
+  encoder.set_output_array(out);
+  int64_t total = num_keys * (half_size + odd);
+  uint32_t threads_y = 1;
+  while ((total / threads_y) >= UINT_MAX) {
+    threads_y *= 2;
+  }
+  uint32_t threads_x = cuda::ceil_div(total, threads_y);
+
+  dim3 grid_dims{
+      static_cast<uint32_t>(num_keys), static_cast<uint32_t>(half_size + odd)};
+  auto [grid, block] = get_grid_and_block(threads_x, threads_y, 1);
+  auto& stream = encoder.stream();
+  if (keys.flags().row_contiguous) {
+    encoder.add_kernel_node(
+        cu::rbitsc,
+        grid,
+        block,
+        0,
+        gpu_ptr<uint32_t>(keys),
+        gpu_ptr<uint8_t>(out),
+        grid_dims,
+        odd,
+        bytes_per_key);
+  } else {
+    encoder.add_kernel_node(
+        cu::rbits,
+        grid,
+        block,
+        0,
+        gpu_ptr<uint32_t>(keys),
+        gpu_ptr<uint8_t>(out),
+        grid_dims,
+        odd,
+        bytes_per_key,
+        keys.ndim(),
+        const_param(keys.shape()),
+        const_param(keys.strides()));
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce.cu b/macos/third_party/mlx/mlx/backend/cuda/reduce.cu
new file mode 100644
index 0000000..269efc0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce.cu
@@ -0,0 +1,73 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/reduce/reduce.cuh"
+#include "mlx/backend/gpu/copy.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+#include <cassert>
+
+namespace mlx::core {
+
+void Reduce::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Reduce::eval_gpu");
+  assert(inputs.size() == 1);
+  array in = inputs[0];
+
+  // Make sure no identity reductions trickle down here.
+  assert(!axes_.empty());
+  assert(out.size() != in.size());
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  if (in.size() == 0) {
+    init_reduce(encoder, in, out, reduce_type_);
+    return;
+  }
+
+  // Reduce.
+  ReductionPlan plan = get_reduction_plan(in, axes_);
+
+  // If it is a general reduce then copy the input to a contiguous array and
+  // recompute the plan.
+  //
+  // TODO: Instead of copying we can use elem-to-loc to deal with broadcasting
+  //       like we do in Metal. When it comes to broadcasted reduction axes
+  //       some can be ignored eg for min/max.
+  bool broadcasted = false;
+  for (int i = 0, j = 0; i < in.ndim() && !broadcasted; i++) {
+    if (j < axes_.size() && axes_[j] == i) {
+      j++;
+    } else {
+      broadcasted = in.strides(i) == 0;
+    }
+  }
+  if (plan.type == GeneralReduce || broadcasted || !in.flags().contiguous) {
+    array in_copy = contiguous_copy_gpu(in, s);
+    encoder.add_temporary(in_copy);
+    in = in_copy;
+    plan = get_reduction_plan(in, axes_);
+  }
+
+  if (plan.type == ContiguousAllReduce) {
+    all_reduce(encoder, in, out, reduce_type_);
+    return;
+  }
+
+  if (plan.type == ContiguousReduce || plan.type == GeneralContiguousReduce) {
+    row_reduce(encoder, in, out, reduce_type_, axes_, plan);
+    return;
+  }
+
+  if (plan.type == ContiguousStridedReduce ||
+      plan.type == GeneralStridedReduce) {
+    col_reduce(encoder, in, out, reduce_type_, axes_, plan);
+    return;
+  }
+
+  throw std::runtime_error("No plan reached in reduce.");
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce/all_reduce.cu b/macos/third_party/mlx/mlx/backend/cuda/reduce/all_reduce.cu
new file mode 100644
index 0000000..962e80d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce/all_reduce.cu
@@ -0,0 +1,159 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/reduce/reduce.cuh"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <cub/block/block_load.cuh>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, typename U, typename ReduceOp, int N = 4>
+__global__ void all_reduce(T* in, U* out, size_t block_step, size_t size) {
+  // TODO: Process multiple "rows" in each thread
+  constexpr int M = 1;
+
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  const U init = cu::ReduceInit<ReduceOp, T>::value();
+  ReduceOp op;
+
+  T vals[N];
+  U accs[M];
+  accs[0] = init;
+
+  size_t start = grid.block_rank() * block_step;
+  size_t end = start + block_step;
+  size_t check = min(end, size);
+
+  size_t i = start;
+  for (; i + block.size() * N <= check; i += block.size() * N) {
+    cub::LoadDirectBlockedVectorized<T, N>(block.thread_rank(), in + i, vals);
+    for (int j = 0; j < N; j++) {
+      accs[0] = op(accs[0], cast_to<U>(vals[j]));
+    }
+  }
+
+  if (i < check) {
+    cub::LoadDirectBlocked(
+        block.thread_rank(), in + i, vals, check - i, cast_to<T>(init));
+    for (int i = 0; i < N; i++) {
+      accs[0] = op(accs[0], cast_to<U>(vals[i]));
+    }
+  }
+
+  __shared__ U shared_accumulators[32];
+  block_reduce(block, warp, accs, shared_accumulators, op, init);
+
+  if (block.thread_rank() == 0) {
+    out[grid.block_rank()] = accs[0];
+  }
+}
+
+} // namespace cu
+
+void all_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type) {
+  constexpr int N_READS = 8;
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  auto get_args = [](int size, int N) {
+    int threads = std::min(512, (size + N - 1) / N);
+    threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+    int reductions_per_step = threads * N;
+    size_t steps_needed =
+        (size + reductions_per_step - 1) / reductions_per_step;
+
+    int blocks;
+    if (steps_needed < 32) {
+      blocks = 1;
+    } else if (steps_needed < 128) {
+      blocks = 32;
+    } else if (steps_needed < 512) {
+      blocks = 128;
+    } else if (steps_needed < 1024) {
+      blocks = 512;
+    } else {
+      blocks = 1024;
+    }
+
+    size_t steps_per_block = (steps_needed + blocks - 1) / blocks;
+    size_t block_step = steps_per_block * reductions_per_step;
+
+    return std::make_tuple(blocks, threads, block_step);
+  };
+
+  int blocks, threads;
+  size_t block_step;
+  size_t insize = in.size();
+  Dtype dt = in.dtype();
+
+  // Cub doesn't like const pointers for load (sigh).
+  void* indata = const_cast<void*>(gpu_ptr<void>(in));
+
+  // Large array so allocate an intermediate and accumulate there
+  std::tie(blocks, threads, block_step) = get_args(insize, N_READS);
+  encoder.set_input_array(in);
+  if (blocks > 1) {
+    array intermediate({blocks}, out.dtype(), nullptr, {});
+    intermediate.set_data(cu::malloc_async(intermediate.nbytes(), encoder));
+    encoder.add_temporary(intermediate);
+    encoder.set_output_array(intermediate);
+    dispatch_all_types(dt, [&](auto type_tag) {
+      dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+        using OP = MLX_GET_TYPE(reduce_type_tag);
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        using U = typename cu::ReduceResult<OP, T>::type;
+        auto kernel = cu::all_reduce<T, U, OP, N_READS>;
+        encoder.add_kernel_node(
+            kernel,
+            blocks,
+            threads,
+            0,
+            static_cast<T*>(indata),
+            gpu_ptr<U>(intermediate),
+            block_step,
+            insize);
+      });
+    });
+
+    // Set the input for the next step and recalculate the blocks
+    indata = gpu_ptr<void>(intermediate);
+    dt = intermediate.dtype();
+    insize = intermediate.size();
+    std::tie(blocks, threads, block_step) = get_args(insize, N_READS);
+    encoder.set_input_array(intermediate);
+  }
+
+  encoder.set_output_array(out);
+  dispatch_all_types(dt, [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      using OP = MLX_GET_TYPE(reduce_type_tag);
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      using U = typename cu::ReduceResult<OP, T>::type;
+      auto kernel = cu::all_reduce<T, U, OP, N_READS>;
+      encoder.add_kernel_node(
+          kernel,
+          blocks,
+          threads,
+          0,
+          static_cast<T*>(indata),
+          gpu_ptr<U>(out),
+          block_step,
+          insize);
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce/col_reduce.cu b/macos/third_party/mlx/mlx/backend/cuda/reduce/col_reduce.cu
new file mode 100644
index 0000000..e33551d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce/col_reduce.cu
@@ -0,0 +1,510 @@
+// Copyright © 2025 Apple Inc.
+
+#include <numeric>
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/reduce/reduce.cuh"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <cub/block/block_load.cuh>
+#include <cub/cub.cuh>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+struct ColReduceArgs {
+  // The size of the contiguous column reduction.
+  size_t reduction_size;
+  int64_t reduction_stride;
+
+  // Input shape and strides excluding the reduction axes.
+  Shape shape;
+  Strides strides;
+  int ndim;
+
+  // Input shape and strides of the reduction axes (including last dimension).
+  Shape reduce_shape;
+  Strides reduce_strides;
+  int reduce_ndim;
+
+  // The number of column we are reducing. Namely prod(reduce_shape).
+  size_t non_col_reductions;
+
+  ColReduceArgs(
+      const array& in,
+      const ReductionPlan& plan,
+      const std::vector<int>& axes) {
+    using ShapeVector = decltype(plan.shape);
+    using StridesVector = decltype(plan.strides);
+
+    ShapeVector shape_vec;
+    StridesVector strides_vec;
+
+    assert(!plan.shape.empty());
+    reduction_size = plan.shape.back();
+    reduction_stride = plan.strides.back();
+
+    int64_t stride_back = 1;
+    std::tie(shape_vec, strides_vec) = shapes_without_reduction_axes(in, axes);
+    while (!shape_vec.empty() && stride_back < reduction_stride) {
+      stride_back *= shape_vec.back();
+      shape_vec.pop_back();
+      strides_vec.pop_back();
+    }
+    std::vector<int> indices(shape_vec.size());
+    std::iota(indices.begin(), indices.end(), 0);
+    std::sort(indices.begin(), indices.end(), [&](int left, int right) {
+      return strides_vec[left] > strides_vec[right];
+    });
+    ShapeVector sorted_shape;
+    StridesVector sorted_strides;
+    for (auto idx : indices) {
+      sorted_shape.push_back(shape_vec[idx]);
+      sorted_strides.push_back(strides_vec[idx]);
+    }
+    std::tie(shape_vec, strides_vec) =
+        collapse_contiguous_dims(sorted_shape, sorted_strides);
+    shape = const_param(shape_vec);
+    strides = const_param(strides_vec);
+    ndim = shape_vec.size();
+
+    reduce_shape = const_param(plan.shape);
+    reduce_strides = const_param(plan.strides);
+    reduce_ndim = plan.shape.size();
+
+    non_col_reductions = 1;
+    for (int i = 0; i < reduce_ndim - 1; i++) {
+      non_col_reductions *= reduce_shape[i];
+    }
+  }
+};
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int NDIM,
+    int BM,
+    int BN,
+    int N_READS = 4,
+    int BLOCKS = 1>
+__global__ void col_reduce_looped(
+    T* in,
+    U* out,
+    const __grid_constant__ ColReduceArgs args,
+    int64_t out_size) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  constexpr int threads_per_row = BN / N_READS;
+
+  // Compute the indices for the tile
+  size_t tile_idx = grid.block_rank();
+  size_t tile_x = tile_idx % ((args.reduction_stride + BN - 1) / BN);
+  size_t tile_y = tile_idx / ((args.reduction_stride + BN - 1) / BN);
+  size_t tile_out = tile_y / out_size;
+  tile_y = tile_y % out_size;
+
+  // Compute the indices for the thread within the tile
+  short thread_x = block.thread_rank() % threads_per_row;
+  short thread_y = block.thread_rank() / threads_per_row;
+
+  // Move the input pointer
+  in += elem_to_loc(tile_y, args.shape.data(), args.strides.data(), args.ndim) +
+      tile_x * BN;
+
+  // Initialize the running totals
+  Op op;
+  U totals[N_READS];
+  for (int i = 0; i < N_READS; i++) {
+    totals[i] = ReduceInit<Op, T>::value();
+  }
+
+  size_t total = args.non_col_reductions * args.reduction_size;
+  size_t per_block, start, end;
+  if constexpr (BLOCKS > 1) {
+    per_block = (total + BLOCKS - 1) / BLOCKS;
+    start = tile_out * per_block + thread_y;
+    end = min((tile_out + 1) * per_block, total);
+  } else {
+    per_block = total;
+    start = thread_y;
+    end = total;
+  }
+
+  LoopedElemToLoc<NDIM, (NDIM > 2)> loop(args.reduce_ndim);
+  loop.next(start, args.reduce_shape.data(), args.reduce_strides.data());
+  if (tile_x * BN + BN <= args.reduction_stride) {
+    if (args.reduction_stride % N_READS == 0) {
+      for (size_t r = start; r < end; r += BM) {
+        T vals[N_READS];
+        cub::LoadDirectBlockedVectorized(thread_x, in + loop.location(), vals);
+        for (int i = 0; i < N_READS; i++) {
+          totals[i] = op(totals[i], cast_to<U>(vals[i]));
+        }
+        loop.next(BM, args.reduce_shape.data(), args.reduce_strides.data());
+      }
+    } else {
+      for (size_t r = start; r < end; r += BM) {
+        T vals[N_READS];
+        cub::LoadDirectBlocked(thread_x, in + loop.location(), vals);
+        for (int i = 0; i < N_READS; i++) {
+          totals[i] = op(totals[i], cast_to<U>(vals[i]));
+        }
+        loop.next(BM, args.reduce_shape.data(), args.reduce_strides.data());
+      }
+    }
+  } else {
+    for (size_t r = start; r < end; r += BM) {
+      T vals[N_READS];
+      cub::LoadDirectBlocked(
+          thread_x,
+          in + loop.location(),
+          vals,
+          args.reduction_stride - tile_x * BN,
+          cast_to<T>(ReduceInit<Op, T>::value()));
+      for (int i = 0; i < N_READS; i++) {
+        totals[i] = op(totals[i], cast_to<U>(vals[i]));
+      }
+      loop.next(BM, args.reduce_shape.data(), args.reduce_strides.data());
+    }
+  }
+
+  // Do warp reduce for each output.
+  constexpr int n_outputs = BN / threads_per_row;
+  static_assert(BM == 32 && n_outputs == N_READS);
+  __shared__ U shared_vals[BM * BN];
+  short s_idx = thread_y * BN + thread_x * N_READS;
+  for (int i = 0; i < N_READS; i++) {
+    shared_vals[s_idx + i] = totals[i];
+  }
+  block.sync();
+  s_idx = warp.thread_rank() * BN + warp.meta_group_rank() * n_outputs;
+  for (int i = 0; i < n_outputs; i++) {
+    totals[i] = cg::reduce(warp, shared_vals[s_idx + i], op);
+  }
+
+  // Write result.
+  if (warp.thread_rank() == 0) {
+    if (BLOCKS > 1) {
+      out += tile_out * out_size * args.reduction_stride;
+    }
+    cub::StoreDirectBlocked(
+        warp.meta_group_rank(),
+        out + tile_y * args.reduction_stride + tile_x * BN,
+        totals,
+        args.reduction_stride - tile_x * BN);
+  }
+}
+
+template <typename T, typename U, typename Op, int N_READS = 4>
+__global__ void col_reduce_small(
+    const T* in,
+    U* out,
+    const __grid_constant__ ColReduceArgs args,
+    size_t total) {
+  Op op;
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  const auto idx = grid.thread_rank() * N_READS;
+  const auto before_axis = idx / args.reduction_stride;
+  const auto after_axis = idx % args.reduction_stride;
+  const auto offset =
+      before_axis * args.reduction_stride * args.reduction_size + after_axis;
+
+  if (idx >= total) {
+    return;
+  }
+
+  in += offset;
+  out += idx;
+
+  AlignedVector<U, N_READS> accumulator;
+  for (int i = 0; i < N_READS; i++) {
+    accumulator[i] = ReduceInit<Op, T>::value();
+  }
+
+  for (int i = 0; i < args.reduction_size; i++) {
+    auto values = load_vector<N_READS>(in, 0);
+
+    for (int j = 0; j < N_READS; j++) {
+      accumulator[j] = op(accumulator[j], cast_to<U>(values[j]));
+    }
+
+    in += args.reduction_stride;
+  }
+
+  store_vector(out, 0, accumulator);
+}
+
+} // namespace cu
+
+inline auto output_grid_for_col_reduce(
+    const array& out,
+    const cu::ColReduceArgs& args,
+    int bn,
+    int outer = 1) {
+  int gx, gy = 1;
+  size_t n_inner_blocks = cuda::ceil_div(args.reduction_stride, bn);
+  size_t n_outer_blocks = out.size() / args.reduction_stride;
+  size_t n_blocks = n_outer_blocks * n_inner_blocks * outer;
+  while (n_blocks / gy > INT32_MAX) {
+    gy *= 2;
+  }
+  gx = cuda::ceil_div(n_blocks, gy);
+
+  return dim3(gx, gy, 1);
+}
+
+void col_reduce_looped(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan,
+    const cu::ColReduceArgs& args) {
+  // Allocate data for the output using in's layout to access them as
+  // contiguously as possible.
+  allocate_same_layout(out, in, axes, encoder);
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      dispatch_reduce_ndim(args.reduce_ndim, [&](auto reduce_ndim) {
+        using OP = MLX_GET_TYPE(reduce_type_tag);
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        using U = typename cu::ReduceResult<OP, T>::type;
+        // Cub doesn't like const pointers for vectorized loads. (sigh)
+        T* indata = const_cast<T*>(gpu_ptr<T>(in));
+
+        constexpr int N_READS = 4;
+        constexpr int BM = 32;
+        constexpr int BN = 32;
+        dim3 grid = output_grid_for_col_reduce(out, args, BN);
+        int blocks = BM * BN / N_READS;
+        auto kernel =
+            cu::col_reduce_looped<T, U, OP, reduce_ndim(), BM, BN, N_READS>;
+        encoder.add_kernel_node(
+            kernel,
+            grid,
+            blocks,
+            0,
+            indata,
+            gpu_ptr<U>(out),
+            static_cast<cu::ColReduceArgs>(args),
+            out.size() / args.reduction_stride);
+      });
+    });
+  });
+}
+
+void col_reduce_small(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan,
+    const cu::ColReduceArgs& args) {
+  // Allocate data for the output using in's layout to access them as
+  // contiguously as possible.
+  allocate_same_layout(out, in, axes, encoder);
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      using OP = MLX_GET_TYPE(reduce_type_tag);
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      using U = typename cu::ReduceResult<OP, T>::type;
+
+      constexpr int N_READS = 16 / sizeof(T);
+      auto tmp_grid = get_2d_grid_dims(out.shape(), out.strides());
+      auto [grid, block] = get_grid_and_block(tmp_grid.x, tmp_grid.y, 1);
+      auto kernel = cu::col_reduce_small<T, U, OP, N_READS>;
+      encoder.add_kernel_node(
+          kernel,
+          grid,
+          block,
+          0,
+          gpu_ptr<T>(in),
+          gpu_ptr<U>(out),
+          static_cast<cu::ColReduceArgs>(args),
+          out.size());
+    });
+  });
+}
+
+void col_reduce_two_pass(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan,
+    const cu::ColReduceArgs& args) {
+  // Allocate data for the output using in's layout to access them as
+  // contiguously as possible.
+  allocate_same_layout(out, in, axes, encoder);
+
+  // Allocate an intermediate array to hold the 1st pass result
+  constexpr int outer = 32;
+
+  Shape intermediate_shape;
+  intermediate_shape.push_back(outer);
+  intermediate_shape.insert(
+      intermediate_shape.end(), out.shape().begin(), out.shape().end());
+
+  Strides intermediate_strides;
+  intermediate_strides.push_back(out.size());
+  intermediate_strides.insert(
+      intermediate_strides.end(), out.strides().begin(), out.strides().end());
+
+  array intermediate(intermediate_shape, out.dtype(), nullptr, {});
+  auto [data_size, rc, cc] =
+      check_contiguity(intermediate_shape, intermediate_strides);
+  auto fl = out.flags();
+  fl.row_contiguous = rc;
+  fl.col_contiguous = cc;
+  fl.contiguous = true;
+  intermediate.set_data(
+      cu::malloc_async(intermediate.nbytes(), encoder),
+      data_size,
+      intermediate_strides,
+      fl,
+      allocator::free);
+
+  encoder.add_temporary(intermediate);
+  encoder.set_input_array(in);
+  encoder.set_output_array(intermediate);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      dispatch_reduce_ndim(args.reduce_ndim, [&](auto reduce_ndim) {
+        using OP = MLX_GET_TYPE(reduce_type_tag);
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        using U = typename cu::ReduceResult<OP, T>::type;
+        // Cub doesn't like const pointers for vectorized loads. (sigh)
+        T* indata = const_cast<T*>(gpu_ptr<T>(in));
+
+        constexpr int N_READS = 4;
+        constexpr int BM = 32;
+        constexpr int BN = 32;
+        dim3 grid = output_grid_for_col_reduce(out, args, BN, outer);
+        int blocks = BM * BN / N_READS;
+        auto kernel = cu::
+            col_reduce_looped<T, U, OP, reduce_ndim(), BM, BN, N_READS, outer>;
+        encoder.add_kernel_node(
+            kernel,
+            grid,
+            blocks,
+            0,
+            indata,
+            gpu_ptr<U>(intermediate),
+            static_cast<cu::ColReduceArgs>(args),
+            out.size() / args.reduction_stride);
+      });
+    });
+  });
+
+  // Prepare the reduction arguments for the 2nd pass
+  cu::ColReduceArgs second_args = args;
+  second_args.reduction_size = outer;
+  second_args.reduction_stride = out.size();
+  second_args.ndim = 0;
+  second_args.reduce_shape[0] = outer;
+  second_args.reduce_strides[0] = out.size();
+  second_args.reduce_ndim = 1;
+  second_args.non_col_reductions = 1;
+
+  encoder.set_input_array(intermediate);
+  encoder.set_output_array(out);
+  dispatch_all_types(intermediate.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      dispatch_reduce_ndim(second_args.reduce_ndim, [&](auto reduce_ndim) {
+        using OP = MLX_GET_TYPE(reduce_type_tag);
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        using U = typename cu::ReduceResult<OP, T>::type;
+
+        constexpr int N_READS = 4;
+        constexpr int BM = 32;
+        constexpr int BN = 32;
+        dim3 grid = output_grid_for_col_reduce(out, second_args, BN);
+        int blocks = BM * BN / N_READS;
+        auto kernel =
+            cu::col_reduce_looped<T, U, OP, reduce_ndim(), BM, BN, N_READS>;
+        encoder.add_kernel_node(
+            kernel,
+            grid,
+            blocks,
+            0,
+            gpu_ptr<T>(intermediate),
+            gpu_ptr<U>(out),
+            second_args,
+            second_args.reduction_stride);
+      });
+    });
+  });
+}
+
+void col_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan) {
+  // Current col reduce options
+  //
+  // - col_reduce_looped
+  //
+  //   It is a general strided reduce. Each threadblock computes the output for
+  //   a subrow of the fast moving axis. For instance 32 elements.
+  //
+  // - col_reduce_small
+  //
+  //  It is a column reduce for small columns. Each thread loops over the whole
+  //  column without communicating with any other thread.
+  //
+  // - col_reduce_two_pass
+  //
+  //  It is a reduce for long columns. To increase parallelism, we split the
+  //  reduction in two passes. First we do a column reduce where many
+  //  threadblocks operate on different parts of the reduced axis. Then we
+  //  perform a final column reduce.
+  //
+  // Notes: As in row reduce we opt to read as much in order as possible and
+  //        leave transpositions as they are (contrary to our Metal backend).
+  //
+  //        Moreover we need different kernels for short rows and tuning
+
+  // Make the args struct to help route to the best kernel
+  cu::ColReduceArgs args(in, plan, axes);
+
+  // Small col reduce with a single or contiguous reduction axis
+  if (args.non_col_reductions == 1 && args.reduction_size <= 32 &&
+      args.reduction_stride % (16 / in.itemsize()) == 0) {
+    col_reduce_small(encoder, in, out, reduce_type, axes, plan, args);
+    return;
+  }
+
+  // Long column with smallish row
+  size_t total_sums = args.non_col_reductions * args.reduction_size;
+  size_t approx_threads = out.size();
+  if (total_sums / approx_threads > 32) {
+    col_reduce_two_pass(encoder, in, out, reduce_type, axes, plan, args);
+    return;
+  }
+
+  // Fallback col reduce
+  col_reduce_looped(encoder, in, out, reduce_type, axes, plan, args);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce/init_reduce.cu b/macos/third_party/mlx/mlx/backend/cuda/reduce/init_reduce.cu
new file mode 100644
index 0000000..df6f906
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce/init_reduce.cu
@@ -0,0 +1,50 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/reduce/reduce.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, typename U, typename Op>
+__global__ void init_reduce(U* out, size_t size) {
+  auto index = cg::this_grid().thread_rank();
+  if (index < size) {
+    out[index] = ReduceInit<Op, T>::value();
+  }
+}
+
+} // namespace cu
+
+void init_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type) {
+  // Allocate if needed
+  if (out.data_shared_ptr() == nullptr) {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  }
+
+  encoder.set_output_array(out);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      using OP = MLX_GET_TYPE(reduce_type_tag);
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      using U = typename cu::ReduceResult<OP, T>::type;
+      auto kernel = cu::init_reduce<T, U, OP>;
+      dim3 grid = get_2d_grid_dims(out.shape(), out.strides());
+      dim3 block(grid.x < 1024 ? grid.x : 1024, 1, 1);
+      grid.x = (grid.x + 1023) / 1024;
+      encoder.add_kernel_node(
+          kernel, grid, block, 0, gpu_ptr<U>(out), out.size());
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce.cuh b/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce.cuh
new file mode 100644
index 0000000..02e4955
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce.cuh
@@ -0,0 +1,71 @@
+// Copyright © 2025 Apple Inc.
+
+#include <type_traits>
+
+#include "mlx/backend/common/reduce.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/reduce/reduce_ops.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+template <typename F>
+void dispatch_reduce_ndim(int ndim, F&& f) {
+  if (ndim == 1) {
+    f(std::integral_constant<int, 1>{});
+  } else if (ndim == 2) {
+    f(std::integral_constant<int, 2>{});
+  } else {
+    f(std::integral_constant<int, 5>{});
+  }
+}
+
+template <typename F>
+void dispatch_reduce_ops(Reduce::ReduceType reduce_type, F&& f) {
+  if (reduce_type == Reduce::ReduceType::And) {
+    f(type_identity<cu::And>{});
+  } else if (reduce_type == Reduce::ReduceType::Or) {
+    f(type_identity<cu::Or>{});
+  } else if (reduce_type == Reduce::ReduceType::Sum) {
+    f(type_identity<cu::Sum>{});
+  } else if (reduce_type == Reduce::ReduceType::Prod) {
+    f(type_identity<cu::Prod>{});
+  } else if (reduce_type == Reduce::ReduceType::Max) {
+    f(type_identity<cu::Max>{});
+  } else if (reduce_type == Reduce::ReduceType::Min) {
+    f(type_identity<cu::Min>{});
+  } else {
+    throw std::invalid_argument("Unknown reduce type.");
+  }
+}
+
+void all_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type);
+
+void row_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan);
+
+void col_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan);
+
+void init_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce_ops.cuh b/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce_ops.cuh
new file mode 100644
index 0000000..6c6b182
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce_ops.cuh
@@ -0,0 +1,211 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/atomic_ops.cuh"
+#include "mlx/backend/cuda/device/cast_op.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+#include "mlx/backend/cuda/reduce/reduce_utils.cuh"
+
+namespace mlx::core::cu {
+
+// Reduce ops.
+struct And {
+  __device__ __forceinline__ bool operator()(bool a, bool b) {
+    return a && b;
+  }
+
+  __device__ void atomic_update(bool* x, bool y) {
+    atomic_reduce<bool, And>(x, y);
+  }
+};
+
+struct Or {
+  __device__ __forceinline__ bool operator()(bool a, bool b) {
+    return a || b;
+  }
+
+  __device__ void atomic_update(bool* x, bool y) {
+    atomic_reduce<bool, Or>(x, y);
+  }
+};
+
+struct Sum {
+  template <typename T>
+  __device__ __forceinline__ T operator()(T a, T b) {
+    return a + b;
+  }
+
+  template <typename T>
+  __device__ void atomic_update(T* x, T y) {
+    atomic_reduce<T, Sum>(x, y);
+  }
+
+  __device__ void atomic_update(__nv_bfloat16* x, __nv_bfloat16 y) {
+    atomic_add(x, y);
+  }
+
+  __device__ void atomic_update(int* x, int y) {
+    atomic_add(x, y);
+  }
+
+  __device__ void atomic_update(float* x, float y) {
+    atomic_add(x, y);
+  }
+};
+
+struct Prod {
+  template <typename T>
+  __device__ __forceinline__ T operator()(T a, T b) {
+    return a * b;
+  }
+
+  template <typename T>
+  __device__ void atomic_update(T* x, T y) {
+    atomic_reduce<T, Prod>(x, y);
+  }
+};
+
+struct Min {
+  template <typename T>
+  __device__ __forceinline__ T operator()(T a, T b) {
+    if constexpr (is_complex_v<T>) {
+      if (cuda::std::isnan(a.real()) || cuda::std::isnan(a.imag())) {
+        return a;
+      }
+      if (cuda::std::isnan(b.real()) || cuda::std::isnan(b.imag())) {
+        return b;
+      }
+    } else if constexpr (!cuda::std::is_integral_v<T>) {
+      if (cuda::std::isnan(a) || cuda::std::isnan(b)) {
+        return cuda::std::numeric_limits<float>::quiet_NaN();
+      }
+    }
+    return a < b ? a : b;
+  }
+
+  template <typename T>
+  __device__ void atomic_update(T* x, T y) {
+    atomic_reduce<T, Min>(x, y);
+  }
+};
+
+struct Max {
+  template <typename T>
+  __device__ __forceinline__ T operator()(T a, T b) {
+    if constexpr (is_complex_v<T>) {
+      if (cuda::std::isnan(a.real()) || cuda::std::isnan(a.imag())) {
+        return a;
+      }
+      if (cuda::std::isnan(b.real()) || cuda::std::isnan(b.imag())) {
+        return b;
+      }
+    } else if constexpr (!cuda::std::is_integral_v<T>) {
+      if (cuda::std::isnan(a) || cuda::std::isnan(b)) {
+        return cuda::std::numeric_limits<float>::quiet_NaN();
+      }
+    }
+    return a > b ? a : b;
+  }
+
+  template <typename T>
+  __device__ void atomic_update(T* x, T y) {
+    atomic_reduce<T, Max>(x, y);
+  }
+};
+
+// Traits to get the result type of reduce op.
+template <typename Op, typename T>
+struct ReduceResult;
+
+template <typename T>
+struct ReduceResult<And, T> {
+  using type = bool;
+};
+
+template <typename T>
+struct ReduceResult<Or, T> {
+  using type = bool;
+};
+
+template <typename T>
+struct ReduceResult<Sum, T> {
+  using type = cuda::std::conditional_t<
+      (cuda::std::is_integral_v<T> && sizeof(T) <= 4),
+      int32_t,
+      T>;
+};
+
+template <typename T>
+struct ReduceResult<Prod, T> {
+  using type = cuda::std::conditional_t<
+      (cuda::std::is_integral_v<T> && sizeof(T) <= 4),
+      int32_t,
+      T>;
+};
+
+template <typename T>
+struct ReduceResult<Min, T> {
+  using type = T;
+};
+
+template <typename T>
+struct ReduceResult<Max, T> {
+  using type = T;
+};
+
+// Traits to get the init value of reduce op.
+template <typename Op, typename T>
+struct ReduceInit;
+
+template <typename T>
+struct ReduceInit<And, T> {
+  static constexpr __host__ __device__ bool value() {
+    return true;
+  }
+};
+
+template <typename T>
+struct ReduceInit<Or, T> {
+  static constexpr __host__ __device__ bool value() {
+    return false;
+  }
+};
+
+template <typename T>
+struct ReduceInit<Sum, T> {
+  static constexpr __host__ __device__ auto value() {
+    if constexpr (is_complex_v<T>) {
+      return T{0, 0};
+    } else {
+      return cast_to<typename ReduceResult<Sum, T>::type>(0);
+    }
+  }
+};
+
+template <typename T>
+struct ReduceInit<Prod, T> {
+  static constexpr __host__ __device__ auto value() {
+    if constexpr (is_complex_v<T>) {
+      return T{1, 0};
+    } else {
+      return cast_to<typename ReduceResult<Prod, T>::type>(1);
+    }
+  }
+};
+
+template <typename T>
+struct ReduceInit<Min, T> {
+  static constexpr __host__ __device__ T value() {
+    return Limits<T>::max();
+  }
+};
+
+template <typename T>
+struct ReduceInit<Max, T> {
+  static constexpr __host__ __device__ T value() {
+    return Limits<T>::min();
+  }
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce_utils.cuh b/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce_utils.cuh
new file mode 100644
index 0000000..fdbe723
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce/reduce_utils.cuh
@@ -0,0 +1,145 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <numeric>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <size_t N>
+struct uint_by_size;
+template <>
+struct uint_by_size<2> {
+  using type = uint16_t;
+};
+template <>
+struct uint_by_size<4> {
+  using type = uint32_t;
+};
+template <>
+struct uint_by_size<8> {
+  using type = unsigned long long int;
+};
+
+template <typename T, typename Op>
+__device__ void atomic_reduce(T* x, T y) {
+  if constexpr (sizeof(T) == 1) {
+    using U = uint16_t;
+    U* x_int = (U*)((char*)x - ((size_t)x % 2));
+    int shift = ((char*)x - (char*)x_int) * 8;
+    int mask = 0xff << shift;
+    U old_val, new_val;
+    do {
+      old_val = *x_int;
+      T result = Op{}(static_cast<T>((old_val >> shift) & 0xff), y);
+      new_val = (old_val & ~mask) | (result << shift);
+    } while (atomicCAS(x_int, old_val, new_val) != old_val);
+  } else {
+    using U = typename uint_by_size<sizeof(T)>::type;
+    U* x_int = (U*)(x);
+    U old_val, new_val;
+    do {
+      old_val = *x_int;
+      T result = Op{}(*((T*)&old_val), y);
+      new_val = *((U*)&result);
+    } while (atomicCAS(x_int, old_val, new_val) != old_val);
+  }
+}
+
+template <typename T, int N, typename Block, typename Warp, typename Op>
+inline __device__ void
+block_reduce(Block block, Warp warp, T (&vals)[N], T* smem, Op op, T init) {
+  // First reduce in the current warp
+  for (int i = 0; i < N; i++) {
+    vals[i] = cg::reduce(warp, vals[i], op);
+  }
+
+  // Reduce across warps
+  if (warp.meta_group_size() > 1) {
+    if (warp.thread_rank() == 0) {
+      for (int i = 0; i < N; i++) {
+        smem[warp.meta_group_rank() * N + i] = vals[i];
+      }
+    }
+    block.sync();
+    if (warp.thread_rank() < warp.meta_group_size()) {
+      for (int i = 0; i < N; i++) {
+        vals[i] = smem[warp.thread_rank() * N + i];
+      }
+    } else {
+      for (int i = 0; i < N; i++) {
+        vals[i] = init;
+      }
+    }
+    for (int i = 0; i < N; i++) {
+      vals[i] = cg::reduce(warp, vals[i], op);
+    }
+  }
+}
+
+} // namespace cu
+
+inline void allocate_same_layout(
+    array& out,
+    const array& in,
+    const std::vector<int>& axes,
+    cu::CommandEncoder& encoder) {
+  if (in.flags().row_contiguous) {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    return;
+  }
+
+  if (out.ndim() < in.ndim()) {
+    throw std::runtime_error(
+        "Reduction without keepdims only supported for row-contiguous inputs");
+  }
+
+  // Calculate the transpositions applied to in in order to apply them to out.
+  std::vector<int> axis_order(in.ndim());
+  std::iota(axis_order.begin(), axis_order.end(), 0);
+  std::sort(axis_order.begin(), axis_order.end(), [&](int left, int right) {
+    return in.strides(left) > in.strides(right);
+  });
+
+  // Transpose the shape and calculate the strides
+  Shape out_shape(in.ndim());
+  Strides out_strides(in.ndim(), 1);
+  for (int i = 0; i < in.ndim(); i++) {
+    out_shape[i] = out.shape(axis_order[i]);
+  }
+  for (int i = in.ndim() - 2; i >= 0; i--) {
+    out_strides[i] = out_shape[i + 1] * out_strides[i + 1];
+  }
+
+  // Reverse the axis order to get the final strides
+  Strides final_strides(in.ndim());
+  for (int i = 0; i < in.ndim(); i++) {
+    final_strides[axis_order[i]] = out_strides[i];
+  }
+
+  // Calculate the resulting contiguity and do the memory allocation
+  auto [data_size, rc, cc] = check_contiguity(out.shape(), final_strides);
+  auto fl = in.flags();
+  fl.row_contiguous = rc;
+  fl.col_contiguous = cc;
+  fl.contiguous = true;
+  out.set_data(
+      cu::malloc_async(out.nbytes(), encoder),
+      data_size,
+      final_strides,
+      fl,
+      allocator::free);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/reduce/row_reduce.cu b/macos/third_party/mlx/mlx/backend/cuda/reduce/row_reduce.cu
new file mode 100644
index 0000000..ea99e11
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/reduce/row_reduce.cu
@@ -0,0 +1,361 @@
+// Copyright © 2025 Apple Inc.
+
+#include <numeric>
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/reduce/reduce.cuh"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+struct RowReduceArgs {
+  // The size of the row being reduced, i.e. the size of last dimension.
+  int row_size;
+
+  // Input shape and strides excluding the reduction axes.
+  Shape shape;
+  Strides strides;
+  int ndim;
+
+  // Input shape and strides of the reduction axes excluding last dimension.
+  Shape reduce_shape;
+  Strides reduce_strides;
+  int reduce_ndim;
+
+  // The number of rows we are reducing. Namely prod(reduce_shape).
+  size_t non_row_reductions;
+
+  RowReduceArgs(
+      const array& in,
+      const ReductionPlan& plan,
+      const std::vector<int>& axes) {
+    assert(!plan.shape.empty());
+    row_size = plan.shape.back();
+
+    auto [shape_vec, strides_vec] = shapes_without_reduction_axes(in, axes);
+    std::tie(shape_vec, strides_vec) =
+        collapse_contiguous_dims(shape_vec, strides_vec);
+    shape = const_param(shape_vec);
+    strides = const_param(strides_vec);
+    ndim = shape_vec.size();
+
+    reduce_shape = const_param(plan.shape);
+    reduce_strides = const_param(plan.strides);
+    reduce_ndim = plan.shape.size() - 1;
+
+    non_row_reductions = 1;
+    for (int i = 0; i < reduce_ndim; i++) {
+      non_row_reductions *= reduce_shape[i];
+    }
+  }
+
+  // Convert shape and strides as if in was contiguous
+  void sort_access_pattern(const array& in, const std::vector<int>& axes) {
+    auto shape_vec = in.shape();
+    auto strides_vec = in.strides();
+    std::tie(shape_vec, strides_vec) =
+        shapes_without_reduction_axes(shape_vec, strides_vec, axes);
+    std::vector<int> indices(shape_vec.size());
+    std::iota(indices.begin(), indices.end(), 0);
+    std::sort(indices.begin(), indices.end(), [&](int left, int right) {
+      return strides_vec[left] > strides_vec[right];
+    });
+    decltype(shape_vec) sorted_shape;
+    decltype(strides_vec) sorted_strides;
+    for (auto idx : indices) {
+      sorted_shape.push_back(shape_vec[idx]);
+      sorted_strides.push_back(strides_vec[idx]);
+    }
+    std::tie(shape_vec, strides_vec) =
+        collapse_contiguous_dims(sorted_shape, sorted_strides);
+    shape = const_param(shape_vec);
+    strides = const_param(strides_vec);
+    ndim = shape_vec.size();
+  }
+};
+
+template <typename T, typename U, typename ReduceOp, int N = 4, int M = 1>
+__global__ void
+row_reduce_simple(const T* in, U* out, size_t n_rows, int size) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  const U init = cu::ReduceInit<ReduceOp, T>::value();
+  ReduceOp op;
+
+  AlignedVector<T, N> vals[M];
+  AlignedVector<U, M> accs;
+  for (int i = 0; i < M; i++) {
+    accs[i] = init;
+  }
+
+  const size_t start_row =
+      min(n_rows - M, static_cast<size_t>(grid.block_rank() * M));
+  const size_t full_blocks = size / (block.size() * N);
+  const size_t final_offset = full_blocks * (block.size() * N);
+  in += start_row * size + block.thread_rank() * N;
+  out += start_row;
+
+  for (size_t r = 0; r < full_blocks; r++) {
+    for (int k = 0; k < M; k++) {
+      vals[k] = load_vector<N>(in + k * size, 0);
+    }
+    for (int k = 0; k < M; k++) {
+      for (int j = 0; j < N; j++) {
+        accs[k] = op(accs[k], cast_to<U>(vals[k][j]));
+      }
+    }
+
+    in += block.size() * N;
+  }
+
+  if (final_offset < size) {
+    for (int k = 0; k < M; k++) {
+      for (int i = 0; i < N; i++) {
+        vals[k][i] = ((final_offset + block.thread_rank() * N + i) < size)
+            ? in[k * size + i]
+            : cast_to<T>(init);
+      }
+    }
+    for (int k = 0; k < M; k++) {
+      for (int j = 0; j < N; j++) {
+        accs[k] = op(accs[k], cast_to<U>(vals[k][j]));
+      }
+    }
+  }
+
+  __shared__ U shared_accumulators[32 * M];
+  block_reduce(block, warp, accs.val, shared_accumulators, op, init);
+
+  if (block.thread_rank() == 0) {
+    if (grid.block_rank() * M + M <= n_rows) {
+      store_vector(out, 0, accs);
+    } else {
+      short offset = grid.block_rank() * M + M - n_rows;
+      for (int i = offset; i < M; i++) {
+        out[i] = accs[i];
+      }
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int NDIM, int N_READS = 4>
+__global__ void row_reduce_looped(
+    const T* in,
+    U* out,
+    const __grid_constant__ RowReduceArgs args) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  size_t out_idx = grid.block_rank();
+
+  Op op;
+
+  U total[1];
+  U init = ReduceInit<Op, T>::value();
+  total[0] = init;
+  LoopedElemToLoc<NDIM, (NDIM > 2)> loop(args.reduce_ndim);
+  const size_t full_blocks = args.row_size / (block.size() * N_READS);
+  const size_t final_offset = full_blocks * (block.size() * N_READS);
+
+  in += elem_to_loc(out_idx, args.shape.data(), args.strides.data(), args.ndim);
+  in += block.thread_rank() * N_READS;
+
+  // Unaligned reduce
+  if (final_offset < args.row_size) {
+    bool mask[N_READS];
+    for (int i = 0; i < N_READS; i++) {
+      mask[i] =
+          (final_offset + block.thread_rank() * N_READS + i) < args.row_size;
+    }
+
+    for (size_t n = 0; n < args.non_row_reductions; n++) {
+      const T* inlocal = in + loop.location();
+
+      for (size_t r = 0; r < full_blocks; r++) {
+        auto vals = load_vector<N_READS>(inlocal, 0);
+        for (int i = 0; i < N_READS; i++) {
+          total[0] = op(total[0], cast_to<U>(vals[i]));
+        }
+        inlocal += block.size() * N_READS;
+      }
+
+      {
+        T vals[N_READS];
+        for (int i = 0; i < N_READS; i++) {
+          vals[i] = mask[i] ? inlocal[i] : cast_to<T>(init);
+        }
+        for (int i = 0; i < N_READS; i++) {
+          total[0] = op(total[0], cast_to<U>(vals[i]));
+        }
+      }
+
+      loop.next(args.reduce_shape.data(), args.reduce_strides.data());
+    }
+  }
+
+  // Aligned case
+  else {
+    for (size_t n = 0; n < args.non_row_reductions; n++) {
+      const T* inlocal = in + loop.location();
+
+      for (size_t r = 0; r < full_blocks; r++) {
+        auto vals = load_vector<N_READS>(inlocal, 0);
+        for (int i = 0; i < N_READS; i++) {
+          total[0] = op(total[0], cast_to<U>(vals[i]));
+        }
+        inlocal += block.size() * N_READS;
+      }
+
+      loop.next(args.reduce_shape.data(), args.reduce_strides.data());
+    }
+  }
+
+  __shared__ U shared_accumulators[32];
+  block_reduce(block, warp, total, shared_accumulators, op, init);
+
+  if (block.thread_rank() == 0) {
+    out[out_idx] = total[0];
+  }
+}
+
+} // namespace cu
+
+void row_reduce_simple(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan) {
+  // Allocate data for the output using in's layout to avoid elem_to_loc in the
+  // kernel.
+  allocate_same_layout(out, in, axes, encoder);
+
+  // TODO: If out.size() < 1024 which will be a common case then write this in
+  //       2 passes. Something like 32 * out.size() and then do a warp reduce.
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      using OP = MLX_GET_TYPE(reduce_type_tag);
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      using U = typename cu::ReduceResult<OP, T>::type;
+
+      constexpr int N_READS = 16 / sizeof(T);
+
+      // Calculate the grid and block dims
+      size_t reductions = (plan.shape.back() + N_READS - 1) / N_READS;
+      dim3 grid = get_2d_grid_dims(out.shape(), out.strides());
+      int warps = (reductions + WARP_SIZE - 1) / WARP_SIZE;
+      warps /= 4;
+      warps = std::max(std::min(warps, 32), 1);
+      int threads = warps * WARP_SIZE;
+      dim3 block(threads, 1, 1);
+
+      // Pick the kernel
+      auto kernel = cu::row_reduce_simple<T, U, OP, N_READS>;
+      if (grid.x >= 1024) {
+        grid.x = (grid.x + 1) / 2;
+        kernel = cu::row_reduce_simple<T, U, OP, N_READS, 2>;
+      }
+
+      T* indata = const_cast<T*>(gpu_ptr<T>(in));
+      int size = plan.shape.back();
+      encoder.add_kernel_node(
+          kernel, grid, block, 0, indata, gpu_ptr<U>(out), out.size(), size);
+    });
+  });
+}
+
+void row_reduce_looped(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan,
+    cu::RowReduceArgs args) {
+  // Allocate data for the output using in's layout to access them as
+  // contiguously as possible.
+  allocate_same_layout(out, in, axes, encoder);
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      using OP = MLX_GET_TYPE(reduce_type_tag);
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      using U = typename cu::ReduceResult<OP, T>::type;
+
+      constexpr int N_READS = 16 / sizeof(T);
+
+      // Calculate the grid and block dims
+      args.sort_access_pattern(in, axes);
+      dim3 grid = get_2d_grid_dims(out.shape(), out.strides());
+      size_t reductions = (args.row_size + N_READS - 1) / N_READS;
+      int warps = (reductions + WARP_SIZE - 1) / WARP_SIZE;
+      warps /= 4;
+      warps = std::max(std::min(warps, 32), 1);
+      int threads = warps * WARP_SIZE;
+      dim3 block(threads, 1, 1);
+
+      // Pick the kernel
+      auto kernel = cu::row_reduce_looped<T, U, OP, 1, N_READS>;
+      dispatch_reduce_ndim(args.reduce_ndim, [&](auto reduce_ndim) {
+        kernel = cu::row_reduce_looped<T, U, OP, reduce_ndim.value, N_READS>;
+      });
+
+      encoder.add_kernel_node(
+          kernel, grid, block, 0, gpu_ptr<T>(in), gpu_ptr<U>(out), args);
+    });
+  });
+}
+
+void row_reduce(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan) {
+  // Current row reduction options
+  //
+  // - row_reduce_simple
+  //
+  //   That means that we are simply reducing across the fastest moving axis.
+  //   We are reducing 1 or 2 rows per threadblock depending on the size of
+  //   output.
+  //
+  // - row_reduce_looped
+  //
+  //   It is a general row reduction. We are computing 1 output per
+  //   threadblock. We read the fastest moving axis vectorized and loop over
+  //   the rest of the axes.
+  //
+  // Notes: We opt to read as much in order as possible and leave
+  //        transpositions as they are (contrary to our Metal backend).
+
+  // Simple row reduce means that we have 1 axis that we are reducing over and
+  // it has stride 1.
+  if (plan.shape.size() == 1) {
+    row_reduce_simple(encoder, in, out, reduce_type, axes, plan);
+    return;
+  }
+
+  // Make the args struct to help route to the best kernel
+  cu::RowReduceArgs args(in, plan, axes);
+
+  // Fallback row reduce
+  row_reduce_looped(encoder, in, out, reduce_type, axes, plan, std::move(args));
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/rms_norm.cu b/macos/third_party/mlx/mlx/backend/cuda/rms_norm.cu
new file mode 100644
index 0000000..a9fc79f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/rms_norm.cu
@@ -0,0 +1,536 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/reduce/reduce.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/fast_primitives.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+inline __device__ float2 plus_f2(const float2& a, const float2& b) {
+  return {a.x + b.x, a.y + b.y};
+}
+
+// Similar to cub::BlockReduce, but result is broadcasted to every thread.
+template <typename T, int BLOCK_DIM, int GROUP_DIM = WARP_SIZE>
+struct BlockBroadcastReduce {
+  using TempStorage = T[std::max(BLOCK_DIM / WARP_SIZE, 1)];
+
+  cg::thread_block& block;
+  TempStorage& temp;
+
+  template <typename Op>
+  __device__ T Reduce(const T& input, const Op& op, const T& init_value) {
+    auto warp = cg::tiled_partition<GROUP_DIM>(block);
+    T x = cg::reduce(warp, input, op);
+    if constexpr (BLOCK_DIM > GROUP_DIM) {
+      if (warp.thread_rank() == 0) {
+        temp[warp.meta_group_rank()] = x;
+      }
+      block.sync();
+      x = warp.thread_rank() < warp.meta_group_size() ? temp[warp.thread_rank()]
+                                                      : init_value;
+      return cg::reduce(warp, x, op);
+    } else {
+      return x;
+    }
+  }
+
+  __device__ T Sum(const T& input) {
+    return Reduce(input, cg::plus<T>{}, T{});
+  }
+};
+
+template <typename T, int BLOCK_DIM, int REDUCE_DIM, int N_READS = 4>
+__global__ void rms_norm_small(
+    const T* x,
+    const T* w,
+    T* out,
+    float eps,
+    uint32_t axis_size,
+    uint32_t n_rows,
+    int64_t w_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceT = BlockBroadcastReduce<float, BLOCK_DIM, REDUCE_DIM>;
+  __shared__ typename BlockReduceT::TempStorage temp;
+
+  auto row =
+      (grid.block_rank() * block.dim_threads().y) + block.thread_index().y;
+  if (row >= n_rows) {
+    return;
+  }
+  x += row * axis_size;
+  out += row * axis_size;
+
+  // Normalizer.
+  float normalizer = 0;
+  auto index = block.thread_index().x;
+  auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    float t = static_cast<float>(xn[i]);
+    normalizer += t * t;
+  }
+
+  normalizer = BlockReduceT{block, temp}.Sum(normalizer);
+  normalizer = rsqrt(normalizer / axis_size + eps);
+
+  // Outputs.
+  auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    float y = static_cast<float>(xn[i]) * normalizer;
+    xn[i] = wn[i] * static_cast<T>(y);
+  }
+  store_vector<N_READS>(out, index, xn, axis_size);
+}
+
+template <typename T, int BLOCK_DIM, int N_READS = 4>
+__global__ void rms_norm(
+    const T* x,
+    const T* w,
+    T* out,
+    float eps,
+    uint32_t axis_size,
+    int64_t w_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceT = BlockBroadcastReduce<float, BLOCK_DIM>;
+  __shared__ typename BlockReduceT::TempStorage temp;
+
+  x += grid.block_rank() * axis_size;
+  out += grid.block_rank() * axis_size;
+
+  // Normalizer.
+  float normalizer = 0;
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      float t = static_cast<float>(xn[i]);
+      normalizer += t * t;
+    }
+  }
+  normalizer = BlockReduceT{block, temp}.Sum(normalizer);
+  normalizer = rsqrt(normalizer / axis_size + eps);
+
+  // Outputs.
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      float y = static_cast<float>(xn[i]) * normalizer;
+      xn[i] = wn[i] * static_cast<T>(y);
+    }
+    store_vector<N_READS>(out, index, xn, axis_size);
+  }
+}
+
+template <
+    typename T,
+    bool HAS_W,
+    int BLOCK_DIM,
+    int REDUCE_DIM,
+    int N_READS = 4>
+__global__ void rms_norm_vjp_small(
+    const T* x,
+    const T* w,
+    const T* g,
+    T* gx,
+    T* gw,
+    float eps,
+    int32_t axis_size,
+    int32_t n_rows,
+    int64_t w_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceF2 = BlockBroadcastReduce<float2, BLOCK_DIM, REDUCE_DIM>;
+  __shared__ typename BlockReduceF2::TempStorage temp;
+
+  auto row =
+      (grid.block_rank() * block.dim_threads().y) + block.thread_index().y;
+  if (row >= n_rows) {
+    return;
+  }
+
+  x += row * axis_size;
+  g += row * axis_size;
+  gx += row * axis_size;
+  gw += row * axis_size;
+
+  // Normalizer.
+  float2 factors = {};
+  auto index = block.thread_index().x;
+  auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+  auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+  auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+  for (int i = 0; i < N_READS; i++) {
+    float t = static_cast<float>(xn[i]);
+    float wi = wn[i];
+    float gi = gn[i];
+    float wg = wi * gi;
+    factors = plus_f2(factors, {wg * t, t * t});
+  }
+
+  factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {});
+  float meangwx = factors.x / axis_size;
+  float normalizer = rsqrt(factors.y / axis_size + eps);
+  float normalizer3 = normalizer * normalizer * normalizer;
+
+  // Outputs.
+  for (int i = 0; i < N_READS; i++) {
+    float xi = xn[i];
+    float wi = wn[i];
+    float gi = gn[i];
+    xn[i] = static_cast<T>(normalizer * wi * gi - xi * meangwx * normalizer3);
+    if constexpr (HAS_W) {
+      wn[i] = static_cast<T>(gi * xi * normalizer);
+    }
+  }
+  store_vector<N_READS>(gx, index, xn, axis_size);
+  if constexpr (HAS_W) {
+    store_vector<N_READS>(gw, index, wn, axis_size);
+  }
+}
+
+template <typename T, bool HAS_W, int BLOCK_DIM, int N_READS = 4>
+__global__ void rms_norm_vjp(
+    const T* x,
+    const T* w,
+    const T* g,
+    T* gx,
+    T* gw,
+    float eps,
+    int32_t axis_size,
+    int64_t w_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceF2 = BlockBroadcastReduce<float2, BLOCK_DIM>;
+  __shared__ typename BlockReduceF2::TempStorage temp;
+
+  x += grid.block_rank() * axis_size;
+  g += grid.block_rank() * axis_size;
+  gx += grid.block_rank() * axis_size;
+  gw += grid.block_rank() * axis_size;
+
+  // Normalizer.
+  float2 factors = {};
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+    for (int i = 0; i < N_READS; i++) {
+      float t = static_cast<float>(xn[i]);
+      float wi = wn[i];
+      float gi = gn[i];
+      float wg = wi * gi;
+      factors = plus_f2(factors, {wg * t, t * t});
+    }
+  }
+  factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {});
+  float meangwx = factors.x / axis_size;
+  float normalizer = rsqrt(factors.y / axis_size + eps);
+  float normalizer3 = normalizer * normalizer * normalizer;
+
+  // Outputs.
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+    for (int i = 0; i < N_READS; i++) {
+      float xi = xn[i];
+      float wi = wn[i];
+      float gi = gn[i];
+      xn[i] = static_cast<T>(normalizer * wi * gi - xi * meangwx * normalizer3);
+      if constexpr (HAS_W) {
+        wn[i] = static_cast<T>(gi * xi * normalizer);
+      }
+    }
+    store_vector<N_READS>(gx, index, xn, axis_size);
+    if constexpr (HAS_W) {
+      store_vector<N_READS>(gw, index, wn, axis_size);
+    }
+  }
+}
+
+} // namespace cu
+
+namespace fast {
+
+bool RMSNorm::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+template <int n_per_thread, typename F>
+void dispatch_group_dim(int axis_size, F&& f) {
+  if (axis_size <= n_per_thread * 8) {
+    f(std::integral_constant<int, 8>{},
+      std::integral_constant<int, 1>(),
+      std::integral_constant<int, 16>());
+  } else if (axis_size <= n_per_thread * 16) {
+    f(std::integral_constant<int, 16>{},
+      std::integral_constant<int, 1>(),
+      std::integral_constant<int, 8>());
+  } else if (axis_size <= n_per_thread * 32) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 1>(),
+      std::integral_constant<int, 4>());
+  } else if (axis_size <= n_per_thread * 32 * 2) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 2>(),
+      std::integral_constant<int, 2>());
+  } else if (axis_size <= n_per_thread * 32 * 4) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 4>(),
+      std::integral_constant<int, 1>());
+  } else if (axis_size <= n_per_thread * 32 * 8) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 8>(),
+      std::integral_constant<int, 1>());
+  } else if (axis_size <= n_per_thread * 32 * 16) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 16>(),
+      std::integral_constant<int, 1>());
+  } else {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 32>(),
+      std::integral_constant<int, 1>());
+  }
+}
+
+// TODO: There are duplicate code with backend/metal/normalization.cpp
+void RMSNorm::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("RMSNorm::eval_gpu");
+  auto& s = stream();
+  auto& out = outputs[0];
+  auto& encoder = cu::get_command_encoder(s);
+
+  // Make sure that the last dimension is contiguous.
+  auto set_output = [&s, &out, &encoder](const array& x) {
+    bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
+    if (no_copy && x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back());
+    }
+    if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            cu::malloc_async(x.data_size() * x.itemsize(), encoder),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  const array x = set_output(inputs[0]);
+  const array& w = inputs[1];
+
+  int32_t axis_size = x.shape().back();
+  int32_t n_rows = x.data_size() / axis_size;
+  int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "rms_norm", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
+    if (axis_size <= N_READS * 1024) {
+      dispatch_group_dim<N_READS>(
+          axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) {
+            constexpr int block_dim = n_groups() * group_dim();
+            auto kernel =
+                cu::rms_norm_small<DataType, block_dim, group_dim(), N_READS>;
+            auto n_blocks =
+                (n_rows + groups_per_block() - 1) / groups_per_block();
+            encoder.add_kernel_node(
+                kernel,
+                n_blocks,
+                {block_dim, groups_per_block()},
+                0,
+                gpu_ptr<DataType>(x),
+                gpu_ptr<DataType>(w),
+                gpu_ptr<DataType>(out),
+                eps_,
+                axis_size,
+                n_rows,
+                w_stride);
+          });
+    } else {
+      auto kernel = cu::rms_norm<DataType, 1024, N_READS>;
+      encoder.add_kernel_node(
+          kernel,
+          n_rows,
+          1024,
+          0,
+          gpu_ptr<DataType>(x),
+          gpu_ptr<DataType>(w),
+          gpu_ptr<DataType>(out),
+          eps_,
+          axis_size,
+          w_stride);
+    }
+  });
+}
+
+void RMSNormVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("RMSNormVJP::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  // Ensure row contiguity. We could relax this step by checking that the array
+  // is contiguous (no broadcasts or holes) and that the input strides are the
+  // same as the cotangent strides but for now this is simpler.
+  auto check_input = [&s](const array& x, bool& copied) {
+    if (x.flags().row_contiguous) {
+      copied = false;
+      return x;
+    }
+    copied = true;
+    return contiguous_copy_gpu(x, s);
+  };
+  bool donate_x = inputs[0].is_donatable();
+  bool donate_g = inputs[2].is_donatable();
+  bool copied;
+  auto x = check_input(inputs[0], copied);
+  donate_x |= copied;
+  const array& w = inputs[1];
+  bool g_copied;
+  auto g = check_input(inputs[2], g_copied);
+  donate_g |= g_copied;
+  array& gx = outputs[0];
+  array& gw = outputs[1];
+
+  // Check whether we had a weight.
+  bool has_w = w.ndim() != 0;
+
+  // Allocate space for the outputs.
+  bool g_in_gx = false;
+  if (donate_x) {
+    gx.copy_shared_buffer(x);
+  } else if (donate_g) {
+    gx.copy_shared_buffer(g);
+    g_in_gx = true;
+  } else {
+    gx.set_data(cu::malloc_async(gx.nbytes(), encoder));
+  }
+  if (g_copied && !g_in_gx) {
+    encoder.add_temporary(g);
+  }
+
+  int32_t axis_size = x.shape().back();
+  int32_t n_rows = x.data_size() / axis_size;
+  int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+
+  // Allocate a temporary to store the gradients for w and allocate the output
+  // gradient accumulators.
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
+  if (has_w) {
+    if (!g_in_gx && donate_g) {
+      gw_temp.copy_shared_buffer(g);
+    } else {
+      gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder));
+      encoder.add_temporary(gw_temp);
+    }
+  }
+
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(g);
+  encoder.set_output_array(gx);
+  encoder.set_output_array(gw_temp);
+  dispatch_float_types(gx.dtype(), "rms_norm_vjp", [&](auto type_tag) {
+    dispatch_bool(has_w, [&](auto has_w_constant) {
+      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      constexpr int N_READS = 16 / sizeof(DataType);
+      if (axis_size <= N_READS * 1024) {
+        dispatch_group_dim<N_READS>(
+            axis_size,
+            [&](auto group_dim, auto n_groups, auto groups_per_block) {
+              constexpr int block_dim = group_dim() * n_groups();
+              auto kernel = cu::rms_norm_vjp_small<
+                  DataType,
+                  has_w_constant.value,
+                  block_dim,
+                  group_dim(),
+                  N_READS>;
+              auto n_blocks =
+                  (n_rows + groups_per_block() - 1) / groups_per_block();
+              encoder.add_kernel_node(
+                  kernel,
+                  n_blocks,
+                  {block_dim, groups_per_block()},
+                  0,
+                  gpu_ptr<DataType>(x),
+                  gpu_ptr<DataType>(w),
+                  gpu_ptr<DataType>(g),
+                  gpu_ptr<DataType>(gx),
+                  gpu_ptr<DataType>(gw_temp),
+                  eps_,
+                  axis_size,
+                  n_rows,
+                  w_stride);
+            });
+      } else {
+        auto kernel =
+            cu::rms_norm_vjp<DataType, has_w_constant.value, 1024, N_READS>;
+        encoder.add_kernel_node(
+            kernel,
+            n_rows,
+            1024,
+            0,
+            gpu_ptr<DataType>(x),
+            gpu_ptr<DataType>(w),
+            gpu_ptr<DataType>(g),
+            gpu_ptr<DataType>(gx),
+            gpu_ptr<DataType>(gw_temp),
+            eps_,
+            axis_size,
+            w_stride);
+      }
+    });
+  });
+
+  if (has_w) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    col_reduce(encoder, gw_temp, gw, Reduce::ReduceType::Sum, {0}, plan);
+  }
+}
+
+} // namespace fast
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/rope.cu b/macos/third_party/mlx/mlx/backend/cuda/rope.cu
new file mode 100644
index 0000000..b0106d4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/rope.cu
@@ -0,0 +1,429 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/fast_primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+template <typename T, bool traditional, bool forward>
+__device__ void rope_single_impl(
+    const T* in,
+    T* out,
+    int32_t offset,
+    float inv_freq,
+    float scale,
+    int64_t stride,
+    uint2 pos,
+    uint2 dims) {
+  float L = scale * static_cast<float>(offset);
+
+  // Compute costheta, sintheta
+  float theta = L * inv_freq;
+  float costheta = cos(theta);
+  float sintheta = sin(theta);
+
+  // Compute the input and output indices
+  uint32_t index_1, index_2;
+  if (traditional) {
+    index_1 = 2 * pos.x + pos.y * stride;
+    index_2 = index_1 + 1;
+  } else {
+    index_1 = pos.x + pos.y * stride;
+    index_2 = index_1 + dims.x;
+  }
+
+  // Read and write the output
+  float x1 = static_cast<float>(in[index_1]);
+  float x2 = static_cast<float>(in[index_2]);
+  float rx1;
+  float rx2;
+  if (forward) {
+    rx1 = x1 * costheta - x2 * sintheta;
+    rx2 = x1 * sintheta + x2 * costheta;
+  } else {
+    rx1 = x2 * sintheta + x1 * costheta;
+    rx2 = x2 * costheta - x1 * sintheta;
+  }
+  out[index_1] = static_cast<T>(rx1);
+  out[index_2] = static_cast<T>(rx2);
+}
+
+template <typename T, bool traditional, bool forward>
+__global__ void rope_single(
+    const T* in,
+    T* out,
+    const int32_t* offset,
+    float scale,
+    float base,
+    int64_t stride,
+    uint2 dims) {
+  uint2 pos = make_uint2(
+      blockIdx.x * blockDim.x + threadIdx.x,
+      blockIdx.y * blockDim.y + threadIdx.y);
+  if (pos.x >= dims.x || pos.y >= dims.y) {
+    return;
+  }
+
+  float d = static_cast<float>(pos.x) / static_cast<float>(dims.x);
+  float inv_freq = exp2(-d * base);
+  rope_single_impl<T, traditional, forward>(
+      in, out, *offset, inv_freq, scale, stride, pos, dims);
+}
+
+template <typename T, bool traditional, bool forward>
+__global__ void rope_single_freqs(
+    const T* in,
+    T* out,
+    const int32_t* offset,
+    const float* freqs,
+    float scale,
+    int64_t stride,
+    uint2 dims,
+    int64_t freq_stride) {
+  uint2 pos = make_uint2(
+      blockIdx.x * blockDim.x + threadIdx.x,
+      blockIdx.y * blockDim.y + threadIdx.y);
+  if (pos.x >= dims.x || pos.y >= dims.y) {
+    return;
+  }
+
+  float inv_freq = 1.0 / freqs[freq_stride * pos.x];
+  rope_single_impl<T, traditional, forward>(
+      in, out, *offset, inv_freq, scale, stride, pos, dims);
+}
+
+template <typename T, bool traditional, bool forward, int N = 4>
+__device__ void rope_impl(
+    const T* in,
+    T* out,
+    const int* offset,
+    float inv_freq,
+    float scale,
+    const cuda::std::array<int64_t, 3> strides,
+    const cuda::std::array<int64_t, 3> out_strides,
+    int64_t offset_stride,
+    int n_head,
+    uint3 pos,
+    uint3 dims) {
+  auto n_head_up = N * ((n_head + N - 1) / N);
+  auto head_idx = static_cast<int>((pos.z * N) % n_head_up);
+  auto batch_idx = (pos.z * N) / n_head_up;
+  auto batch_offset = offset[batch_idx * offset_stride];
+  float L = scale * static_cast<float>(pos.y + batch_offset);
+  auto mat_idx = batch_idx * n_head + head_idx;
+
+  // Compute costheta, sintheta
+  float theta = L * inv_freq;
+  float costheta = cos(theta);
+  float sintheta = sin(theta);
+
+  // Compute the input and output indices
+  size_t in_index_1, in_index_2;
+  size_t out_index_1, out_index_2;
+  if (traditional) {
+    out_index_1 = 2 * pos.x * out_strides[2] + pos.y * out_strides[1] +
+        mat_idx * out_strides[0];
+    out_index_2 = out_index_1 + 1;
+    in_index_1 =
+        2 * pos.x * strides[2] + pos.y * strides[1] + mat_idx * strides[0];
+    in_index_2 = in_index_1 + strides[2];
+  } else {
+    out_index_1 = pos.x * out_strides[2] + pos.y * out_strides[1] +
+        mat_idx * out_strides[0];
+    out_index_2 = out_index_1 + dims.x * out_strides[2];
+    in_index_1 = pos.x * strides[2] + pos.y * strides[1] + mat_idx * strides[0];
+    in_index_2 = in_index_1 + dims.x * strides[2];
+  }
+  for (int i = 0; i < N && head_idx + i < n_head; ++i) {
+    // Read and write the output
+    float x1 = static_cast<float>(in[in_index_1]);
+    float x2 = static_cast<float>(in[in_index_2]);
+    float rx1;
+    float rx2;
+    if (forward) {
+      rx1 = x1 * costheta - x2 * sintheta;
+      rx2 = x1 * sintheta + x2 * costheta;
+    } else {
+      rx1 = x2 * sintheta + x1 * costheta;
+      rx2 = x2 * costheta - x1 * sintheta;
+    }
+    out[out_index_1] = static_cast<T>(rx1);
+    out[out_index_2] = static_cast<T>(rx2);
+    in_index_1 += strides[0];
+    in_index_2 += strides[0];
+    out_index_1 += out_strides[0];
+    out_index_2 += out_strides[0];
+  }
+}
+
+template <typename T, bool traditional, bool forward>
+__global__ void rope(
+    const T* in,
+    T* out,
+    const int32_t* offset,
+    float scale,
+    float base,
+    const __grid_constant__ cuda::std::array<int64_t, 3> strides,
+    const __grid_constant__ cuda::std::array<int64_t, 3> out_strides,
+    int64_t offset_stride,
+    int n_head,
+    uint3 dims) {
+  uint3 pos = make_uint3(
+      blockIdx.x * blockDim.x + threadIdx.x,
+      blockIdx.y * blockDim.y + threadIdx.y,
+      blockIdx.z * blockDim.z + threadIdx.z);
+  if (pos.x >= dims.x || pos.y >= dims.y || pos.z >= dims.z) {
+    return;
+  }
+
+  float d = static_cast<float>(pos.x) / static_cast<float>(dims.x);
+  float inv_freq = exp2(-d * base);
+  rope_impl<T, traditional, forward>(
+      in,
+      out,
+      offset,
+      inv_freq,
+      scale,
+      strides,
+      out_strides,
+      offset_stride,
+      n_head,
+      pos,
+      dims);
+}
+
+template <typename T, bool traditional, bool forward>
+__global__ void rope_freqs(
+    const T* in,
+    T* out,
+    const int32_t* offset,
+    const float* freqs,
+    float scale,
+    float base,
+    const __grid_constant__ cuda::std::array<int64_t, 3> strides,
+    const __grid_constant__ cuda::std::array<int64_t, 3> out_strides,
+    int64_t offset_stride,
+    int n_head,
+    uint3 dims,
+    int64_t freq_stride) {
+  uint3 pos = make_uint3(
+      blockIdx.x * blockDim.x + threadIdx.x,
+      blockIdx.y * blockDim.y + threadIdx.y,
+      blockIdx.z * blockDim.z + threadIdx.z);
+  if (pos.x >= dims.x || pos.y >= dims.y || pos.z >= dims.z) {
+    return;
+  }
+
+  float inv_freq = 1.0 / freqs[freq_stride * pos.x];
+  rope_impl<T, traditional, forward>(
+      in,
+      out,
+      offset,
+      inv_freq,
+      scale,
+      strides,
+      out_strides,
+      offset_stride,
+      n_head,
+      pos,
+      dims);
+}
+
+} // namespace cu
+
+namespace fast {
+
+bool RoPE::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+void RoPE::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("RoPE::eval_gpu");
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  auto& in = inputs[0];
+  auto& offset = inputs[1];
+  auto& out = outputs[0];
+
+  cuda::std::array<int64_t, 3> strides;
+  cuda::std::array<int64_t, 3> out_strides;
+  bool donated = false;
+  int ndim = in.ndim();
+
+  int B = in.shape(0);
+  int T = in.shape(-2);
+  int D = in.shape(-1);
+  size_t mat_size = T * D;
+  int dispatch_ndim = ndim;
+  while (in.shape(-dispatch_ndim) == 1 && dispatch_ndim > 3) {
+    dispatch_ndim--;
+  }
+
+  int N = 1;
+  for (int i = 1; i < (ndim - 2); ++i) {
+    N *= in.shape(i);
+  }
+
+  // We apply rope to less that the whole vector so copy to output and then
+  // apply in-place.
+  if (dims_ < D) {
+    donated = true;
+    auto ctype =
+        (in.flags().row_contiguous) ? CopyType::Vector : CopyType::General;
+    copy_gpu(in, out, ctype, s);
+    strides[0] = mat_size;
+    strides[1] = out.strides()[ndim - 2];
+    strides[2] = out.strides()[ndim - 1];
+  }
+
+  // Either copy or apply in-place
+  else if (in.flags().row_contiguous) {
+    if (in.is_donatable()) {
+      donated = true;
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    }
+    strides[0] = mat_size;
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
+  } else if (dispatch_ndim == 3) {
+    // Handle non-contiguous 3D inputs
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    strides[0] = in.strides()[ndim - 3];
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
+  } else {
+    // Copy non-contiguous > 3D inputs into the output and treat
+    // input as donated
+    donated = true;
+    copy_gpu(in, out, CopyType::General, s);
+    strides[0] = mat_size;
+    strides[1] = out.strides()[ndim - 2];
+    strides[2] = out.strides()[ndim - 1];
+  }
+  out_strides[0] = mat_size;
+  out_strides[1] = out.strides()[ndim - 2];
+  out_strides[2] = out.strides()[ndim - 1];
+
+  // Some flags to help us dispatch below
+  bool single = in.flags().row_contiguous && B == 1 && T == 1;
+  bool with_freqs = inputs.size() == 3;
+
+  encoder.set_input_array(donated ? out : in);
+  encoder.set_input_array(offset);
+  if (with_freqs) {
+    encoder.set_input_array(inputs[2]);
+  }
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "rope", [&](auto type_tag) {
+    dispatch_bool(traditional_, [&](auto traditional) {
+      dispatch_bool(forward_, [&](auto forward) {
+        using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        if (single && !with_freqs) {
+          auto kernel =
+              cu::rope_single<DataType, traditional.value, forward.value>;
+          uint2 dims = make_uint2(dims_ / 2, N);
+          auto [grid, block] = get_grid_and_block(dims.x, dims.y, 1);
+          encoder.add_kernel_node(
+              kernel,
+              grid,
+              block,
+              0,
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
+              scale_,
+              std::log2(base_),
+              mat_size,
+              dims);
+        } else if (single) {
+          auto kernel =
+              cu::rope_single_freqs<DataType, traditional.value, forward.value>;
+          uint2 dims = make_uint2(dims_ / 2, N);
+          auto [grid, block] = get_grid_and_block(dims.x, dims.y, 1);
+          encoder.add_kernel_node(
+              kernel,
+              grid,
+              block,
+              0,
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
+              gpu_ptr<float>(inputs[2]),
+              scale_,
+              mat_size,
+              dims,
+              inputs[2].strides(0));
+        } else if (with_freqs) {
+          auto kernel =
+              cu::rope_freqs<DataType, traditional.value, forward.value>;
+          int n_per_thread = 4;
+          uint32_t dimz = B * ((N + n_per_thread - 1) / n_per_thread);
+          uint3 dims = make_uint3(dims_ / 2, T, dimz);
+          auto [grid, block] = get_grid_and_block(dims.x, dims.y, dims.z);
+          int64_t offset_stride = 0;
+          if (inputs[1].ndim() > 0) {
+            offset_stride = inputs[1].strides()[0];
+          }
+          encoder.add_kernel_node(
+              kernel,
+              grid,
+              block,
+              0,
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
+              gpu_ptr<float>(inputs[2]),
+              scale_,
+              std::log2(base_),
+              strides,
+              out_strides,
+              offset_stride,
+              N,
+              dims,
+              inputs[2].strides(0));
+        } else {
+          auto kernel = cu::rope<DataType, traditional.value, forward.value>;
+          int n_per_thread = 4;
+          uint32_t dimz = B * ((N + n_per_thread - 1) / n_per_thread);
+          uint3 dims = make_uint3(dims_ / 2, T, dimz);
+          auto [grid, block] = get_grid_and_block(dims.x, dims.y, dims.z);
+          int64_t offset_stride = 0;
+          if (inputs[1].ndim() > 0) {
+            offset_stride = inputs[1].strides()[0];
+          }
+          encoder.add_kernel_node(
+              kernel,
+              grid,
+              block,
+              0,
+              gpu_ptr<DataType>(donated ? out : in),
+              gpu_ptr<DataType>(out),
+              gpu_ptr<int32_t>(offset),
+              scale_,
+              std::log2(base_),
+              strides,
+              out_strides,
+              offset_stride,
+              N,
+              dims);
+        }
+      });
+    });
+  });
+}
+
+} // namespace fast
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp b/macos/third_party/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp
new file mode 100644
index 0000000..54700bd
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp
@@ -0,0 +1,506 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/lru_cache.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/fast_primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace {
+
+array prepare_sdpa_input(const array& x, Stream s) {
+  // SDPA kernel's requirements on inputs:
+  // 1. last dim's stride be 1;
+  // 2. pointer be aligned.
+  if (x.strides(-1) != 1 || get_alignment(x) < 16) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    auto& encoder = cu::get_command_encoder(s);
+    encoder.add_temporary(x_copy);
+    return x_copy;
+  }
+  return x;
+}
+
+void malloc_with_same_layout(
+    cu::CommandEncoder& encoder,
+    array& o,
+    const array& q) {
+  if (q.flags().row_contiguous) {
+    o.set_data(cu::malloc_async(o.nbytes(), encoder));
+    return;
+  }
+  // fill_order = argsort(q.strides())
+  Shape fill_order(q.ndim());
+  std::iota(fill_order.begin(), fill_order.end(), 0);
+  std::stable_sort(
+      fill_order.begin(), fill_order.end(), [&q](int idx1, int idx2) {
+        auto s1 = q.strides(idx1) > 0 ? q.strides(idx1) : 1;
+        auto s2 = q.strides(idx2) > 0 ? q.strides(idx2) : 1;
+        return s1 < s2;
+      });
+  // Generate o_strides with fill_order
+  Strides o_strides(q.ndim());
+  int64_t stride = 1;
+  for (int i : fill_order) {
+    o_strides[i] = stride;
+    stride *= o.shape(i);
+  }
+  // o is a transposed contiguous array
+  o.set_data(
+      cu::malloc_async(o.nbytes(), encoder),
+      o.size(),
+      o_strides,
+      {true, false, false});
+}
+
+constexpr int QKV_NDIM = 4;
+
+struct SDPACacheKey {
+  int device_id;
+  fe::DataType_t cudnn_dtype;
+  std::array<int, QKV_NDIM> q_shape;
+  std::array<int, QKV_NDIM> k_shape;
+  std::array<int, QKV_NDIM> v_shape;
+  std::array<int64_t, QKV_NDIM> q_strides;
+  std::array<int64_t, QKV_NDIM> k_strides;
+  std::array<int64_t, QKV_NDIM> v_strides;
+  bool do_causal;
+  std::array<int, QKV_NDIM> mask_shape;
+  std::array<int64_t, QKV_NDIM> mask_strides;
+  bool output_logsumexp;
+};
+
+inline BytesKey<SDPACacheKey> build_sdpa_cache_key(
+    cu::CommandEncoder& encoder,
+    const array& q,
+    const array& k,
+    const array& v,
+    bool do_causal,
+    const std::optional<array>& mask_arr,
+    bool output_logsumexp = true) {
+  BytesKey<SDPACacheKey> cache_key;
+  cache_key.pod = {
+      encoder.device().cuda_device(),
+      dtype_to_cudnn_type(q.dtype()),
+      vector_key<QKV_NDIM>(q.shape()),
+      vector_key<QKV_NDIM>(k.shape()),
+      vector_key<QKV_NDIM>(v.shape()),
+      vector_key<QKV_NDIM>(q.strides()),
+      vector_key<QKV_NDIM>(k.strides()),
+      vector_key<QKV_NDIM>(v.strides()),
+      do_causal,
+      {},
+      {},
+      output_logsumexp,
+  };
+  if (mask_arr) {
+    cache_key.pod.mask_shape = vector_key<QKV_NDIM>(mask_arr->shape());
+    cache_key.pod.mask_strides = vector_key<QKV_NDIM>(mask_arr->strides());
+  }
+  return cache_key;
+}
+
+auto& sdpa_cache() {
+  static LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
+      "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 64);
+  return cache;
+}
+
+auto& sdpa_backward_cache() {
+  static LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
+      "MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE", /* default_capacity */ 64);
+  return cache;
+}
+
+enum UIDS {
+  Q,
+  K,
+  V,
+  SCALE,
+  BIAS,
+  O,
+  STATS,
+  // Backward graph:
+  D_Q,
+  D_K,
+  D_V,
+  D_O,
+};
+
+DnnGraph build_sdpa_graph(
+    cudnnHandle_t handle,
+    const array& q,
+    const array& k,
+    const array& v,
+    bool do_causal,
+    const std::optional<array>& mask_arr,
+    bool output_logsumexp,
+    const array& o,
+    const std::optional<array>& stats) {
+  DnnGraph graph(handle, q.dtype());
+
+  auto q_ = graph.tensor("Q", Q, q);
+  auto k_ = graph.tensor("K", K, k);
+  auto v_ = graph.tensor("V", V, v);
+
+  auto options = fe::graph::SDPA_attributes()
+                     .set_name("sdpa_cudnn")
+                     .set_attn_scale(graph.scalar("Scale", SCALE, float32))
+                     .set_generate_stats(output_logsumexp);
+  if (do_causal) {
+    options.set_causal_mask_bottom_right(do_causal);
+  }
+  if (mask_arr) {
+    options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr));
+  }
+
+  auto [o_, stats_] = graph.sdpa(q_, k_, v_, options);
+  graph.tensor(o_, O, o)->set_output(true);
+  if (output_logsumexp) {
+    graph.tensor(stats_, STATS, *stats)->set_output(true);
+  }
+
+  CHECK_CUDNN_FE_ERROR(graph.prepare());
+  graph.select_behavior_notes(
+      {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
+  CHECK_CUDNN_FE_ERROR(graph.build());
+  return graph;
+}
+
+DnnGraph build_sdpa_backward_graph(
+    cudnnHandle_t handle,
+    const array& q,
+    const array& k,
+    const array& v,
+    bool do_causal,
+    const std::optional<array>& mask_arr,
+    const array& o,
+    const array& d_o,
+    const array& stats,
+    array& d_q,
+    array& d_k,
+    array& d_v) {
+  DnnGraph graph(handle, q.dtype());
+
+  auto q_ = graph.tensor("Q", Q, q);
+  auto k_ = graph.tensor("K", K, k);
+  auto v_ = graph.tensor("V", V, v);
+  auto o_ = graph.tensor("O", O, o);
+  auto d_o_ = graph.tensor("D_O", D_O, d_o);
+  auto stats_ = graph.tensor("STATS", STATS, stats);
+
+  auto options = fe::graph::SDPA_backward_attributes()
+                     .set_name("sdpa_backward_cudnn")
+                     .set_attn_scale(graph.scalar("Scale", SCALE, float32));
+  if (do_causal) {
+    options.set_causal_mask_bottom_right(do_causal);
+  }
+  if (mask_arr) {
+    options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr));
+  }
+
+  auto [d_q_, d_k_, d_v_] =
+      graph.sdpa_backward(q_, k_, v_, o_, d_o_, stats_, options);
+  graph.tensor(d_q_, D_Q, d_q)->set_output(true);
+  graph.tensor(d_k_, D_K, d_k)->set_output(true);
+  graph.tensor(d_v_, D_V, d_v)->set_output(true);
+
+  CHECK_CUDNN_FE_ERROR(graph.prepare());
+  graph.select_behavior_notes(
+      {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
+  CHECK_CUDNN_FE_ERROR(graph.build());
+  return graph;
+}
+
+} // namespace
+
+bool supports_sdpa_cudnn(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool do_causal,
+    Stream s) {
+  static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SPDA", 1);
+  if (!enabled) {
+    return false;
+  }
+
+  // cuDNN SDPA requires Ampere and later.
+  if (cu::device(s.device).compute_capability_major() < 8) {
+    return false;
+  }
+
+  // Only use cuDNN for prefilling (T_q > 1) and training (T_q == T_kv).
+  if ((q.shape(2) == 1) && (q.shape(2) != k.shape(2))) {
+    return false;
+  }
+
+  // cuDNN does not support bottom right mask when T_q > T_kv.
+  if (do_causal && (q.shape(2) > k.shape(2))) {
+    return false;
+  }
+
+  // D_qk and D_v must be a multiple of 8 with maximum value 128.
+  if ((q.shape(-1) % 8 != 0) || (q.shape(-1) > 128) || (v.shape(-1) % 8 != 0) ||
+      (v.shape(-1) > 128)) {
+    return false;
+  }
+
+  Dtype dtype = q.dtype();
+  return dtype == float16 || dtype == bfloat16;
+}
+
+void sdpa_cudnn(
+    const array& q,
+    const array& k,
+    const array& v,
+    float scale,
+    array& o,
+    std::optional<array>& stats,
+    bool do_causal,
+    const std::optional<array>& mask_arr,
+    bool output_logsumexp,
+    Stream s) {
+  auto& encoder = cu::get_command_encoder(s);
+  auto handle = encoder.device().get_cudnn_handle();
+
+  malloc_with_same_layout(encoder, o, q);
+
+  encoder.set_input_array(q);
+  encoder.set_input_array(k);
+  encoder.set_input_array(v);
+  encoder.set_output_array(o);
+  if (mask_arr) {
+    encoder.set_input_array(*mask_arr);
+  }
+  if (output_logsumexp) {
+    stats->set_data(cu::malloc_async(stats->nbytes(), encoder));
+    encoder.set_output_array(*stats);
+  }
+
+  // Search cache.
+  auto cache_key = build_sdpa_cache_key(
+      encoder, q, k, v, do_causal, mask_arr, output_logsumexp);
+  auto it = sdpa_cache().find(cache_key);
+  if (it == sdpa_cache().end()) {
+    auto graph = build_sdpa_graph(
+        handle, q, k, v, do_causal, mask_arr, output_logsumexp, o, stats);
+    it = sdpa_cache().emplace(cache_key, std::move(graph)).first;
+  }
+  auto& graph = it->second;
+
+  std::unordered_map<int64_t, void*> variant_pack{
+      {Q, gpu_ptr<void>(q)},
+      {K, gpu_ptr<void>(k)},
+      {V, gpu_ptr<void>(v)},
+      {SCALE, &scale},
+      {O, gpu_ptr<void>(o)}};
+  if (mask_arr) {
+    variant_pack[BIAS] = gpu_ptr<void>(*mask_arr);
+  }
+  if (output_logsumexp) {
+    variant_pack[STATS] = gpu_ptr<void>(*stats);
+  }
+
+  CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
+}
+
+void sdpa_backward_cudnn(
+    const array& q,
+    const array& k,
+    const array& v,
+    float scale,
+    const array& o,
+    const array& stats,
+    bool do_causal,
+    const std::optional<array>& mask_arr,
+    const array& d_o,
+    array& d_q,
+    array& d_k,
+    array& d_v,
+    Stream s) {
+  auto& encoder = cu::get_command_encoder(s);
+  auto handle = encoder.device().get_cudnn_handle();
+
+  malloc_with_same_layout(encoder, d_q, q);
+  malloc_with_same_layout(encoder, d_k, k);
+  malloc_with_same_layout(encoder, d_v, v);
+
+  encoder.set_input_array(q);
+  encoder.set_input_array(k);
+  encoder.set_input_array(v);
+  encoder.set_input_array(o);
+  encoder.set_input_array(stats);
+  encoder.set_input_array(d_o);
+  encoder.set_output_array(d_q);
+  encoder.set_output_array(d_k);
+  encoder.set_output_array(d_v);
+  if (mask_arr) {
+    encoder.set_input_array(*mask_arr);
+  }
+
+  // Search cache.
+  auto cache_key = build_sdpa_cache_key(encoder, q, k, v, do_causal, mask_arr);
+  auto it = sdpa_backward_cache().find(cache_key);
+  if (it == sdpa_backward_cache().end()) {
+    auto graph = build_sdpa_backward_graph(
+        handle, q, k, v, do_causal, mask_arr, o, d_o, stats, d_q, d_k, d_v);
+    it = sdpa_backward_cache().emplace(cache_key, std::move(graph)).first;
+  }
+  auto& graph = it->second;
+
+  std::unordered_map<int64_t, void*> variant_pack{
+      {Q, gpu_ptr<void>(q)},
+      {K, gpu_ptr<void>(k)},
+      {V, gpu_ptr<void>(v)},
+      {SCALE, &scale},
+      {O, gpu_ptr<void>(o)},
+      {STATS, gpu_ptr<void>(stats)},
+      {D_O, gpu_ptr<void>(d_o)},
+      {D_Q, gpu_ptr<void>(d_q)},
+      {D_K, gpu_ptr<void>(d_k)},
+      {D_V, gpu_ptr<void>(d_v)}};
+  if (mask_arr) {
+    variant_pack[BIAS] = gpu_ptr<void>(*mask_arr);
+  }
+
+  CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
+}
+
+// Defined in scaled_dot_product_attention.cu file.
+bool supports_sdpa_vector(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool output_logsumexp);
+void sdpa_vector(
+    const array& q,
+    const array& k,
+    const array& v,
+    float scale,
+    array& o,
+    bool do_causal,
+    const std::optional<array>& sinks,
+    Stream s);
+
+namespace fast {
+
+bool ScaledDotProductAttention::use_fallback(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool is_training,
+    bool output_logsumexp,
+    Stream s) {
+  if (s.device == Device::cpu) {
+    return true;
+  }
+
+  return !supports_sdpa_vector(
+             q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp) &&
+      !supports_sdpa_cudnn(q, k, v, do_causal, s);
+}
+
+bool ScaledDotProductAttention::supports_bool_mask() {
+  return false;
+}
+
+void ScaledDotProductAttention::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("ScaledDotProductAttention::eval_gpu");
+
+  auto& s = stream();
+
+  array q = prepare_sdpa_input(inputs[0], s);
+  array k = prepare_sdpa_input(inputs[1], s);
+  array v = prepare_sdpa_input(inputs[2], s);
+  array& out = outputs[0];
+  bool has_mask = inputs.size() - has_sinks_ > 3;
+  bool has_arr_mask = has_mask && !do_causal_;
+
+  std::optional<array> mask_arr;
+  if (has_arr_mask) {
+    mask_arr = prepare_sdpa_input(inputs[3], s);
+  }
+  std::optional<array> stats;
+  if (output_logsumexp_) {
+    stats = outputs[1];
+  }
+
+  if (supports_sdpa_vector(
+          q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) {
+    if (has_sinks_) {
+      sdpa_vector(q, k, v, scale_, out, do_causal_, inputs.back(), s);
+    } else {
+      sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s);
+    }
+  } else {
+    sdpa_cudnn(
+        q,
+        k,
+        v,
+        scale_,
+        out,
+        stats,
+        do_causal_,
+        mask_arr,
+        output_logsumexp_,
+        s);
+  }
+}
+
+bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) {
+  // The frontend adds a padding mask when sequence length is not a multiple of
+  // tile size.
+  if (q.shape(2) % 128 != 0) {
+    return true;
+  }
+  return s.device == Device::cpu;
+}
+
+void ScaledDotProductAttentionVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("ScaledDotProductAttentionVJP::eval_gpu");
+
+  auto& s = stream();
+
+  assert(inputs.size() >= 6);
+  int primals_size = inputs.size() - 3;
+  bool has_arr_mask = primals_size > 3 + has_sinks_;
+
+  array q = prepare_sdpa_input(inputs[0], s);
+  array k = prepare_sdpa_input(inputs[1], s);
+  array v = prepare_sdpa_input(inputs[2], s);
+  array o = prepare_sdpa_input(inputs[primals_size], s);
+  array stats = prepare_sdpa_input(inputs[primals_size + 1], s);
+  array d_o = prepare_sdpa_input(inputs[primals_size + 2], s);
+
+  std::optional<array> mask_arr;
+  if (has_arr_mask) {
+    mask_arr = prepare_sdpa_input(inputs[3], s);
+  }
+
+  assert(outputs.size() == 3);
+  auto& d_q = outputs[0];
+  auto& d_k = outputs[1];
+  auto& d_v = outputs[2];
+
+  sdpa_backward_cudnn(
+      q, k, v, scale_, o, stats, do_causal_, mask_arr, d_o, d_q, d_k, d_v, s);
+}
+
+} // namespace fast
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/scaled_dot_product_attention.cu b/macos/third_party/mlx/mlx/backend/cuda/scaled_dot_product_attention.cu
new file mode 100644
index 0000000..111cfee
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/scaled_dot_product_attention.cu
@@ -0,0 +1,798 @@
+// Copyright © 2025 Apple Inc.
+
+// Required for using M_LOG2E in MSVC.
+#define _USE_MATH_DEFINES
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/config.h"
+#include "mlx/backend/cuda/device/utils.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+#define PRAGMA_LOOP_UNROLL #pragma unroll
+
+struct AttnParams {
+  int B;
+  int H;
+  int D;
+
+  int qL;
+  int kL;
+
+  int gqa_factor;
+  float scale;
+
+  int64_t Q_strides[3];
+  int64_t K_strides[3];
+  int64_t V_strides[3];
+  int64_t O_strides[3];
+};
+
+template <typename T, bool do_causal, int D>
+__global__ void kernel_sdpav_1pass(
+    const T* Q,
+    const T* K,
+    const T* V,
+    T* O,
+    const T* sinks,
+    __grid_constant__ const AttnParams params) {
+  constexpr int BN = 32;
+  constexpr int BD = 32;
+
+  constexpr int v_per_thread = D / BD;
+
+  const int inner_k_stride = BN * int(params.K_strides[2]);
+  const int inner_v_stride = BN * int(params.V_strides[2]);
+
+  typedef float U;
+
+  U q[v_per_thread];
+  U k[v_per_thread];
+  U o[v_per_thread];
+
+  __shared__ U outputs[BN][BD + 1];
+  __shared__ U max_scores[BN];
+  __shared__ U sum_exp_scores[BN];
+
+  const U scale_log2 = params.scale * M_LOG2E;
+
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<32>(block);
+
+  const int lane_idx = warp.thread_rank();
+  const int warp_idx = warp.meta_group_rank();
+
+  // Adjust to thread block and thread
+  const int batch_idx = blockIdx.z;
+  const int head_idx = blockIdx.x;
+  const int kv_head_idx = head_idx / params.gqa_factor;
+
+  const int q_seq_idx = blockIdx.y;
+  const int kv_seq_idx = warp_idx;
+
+  Q += batch_idx * params.Q_strides[0] + // Batch
+      head_idx * params.Q_strides[1] + // Head
+      q_seq_idx * params.Q_strides[2]; // Sequence
+
+  K += batch_idx * params.K_strides[0] + // Batch
+      kv_head_idx * params.K_strides[1] + // Head
+      kv_seq_idx * params.K_strides[2]; // Sequence
+
+  V += batch_idx * params.V_strides[0] + // Batch
+      kv_head_idx * params.V_strides[1] + // Head
+      kv_seq_idx * params.V_strides[2]; // Sequence
+
+  O += batch_idx * params.O_strides[0] + // Batch
+      head_idx * params.O_strides[1] + // Head
+      q_seq_idx * params.O_strides[2]; // Sequence
+
+  // Read the query and 0 the output accumulator
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    q[i] = scale_log2 * static_cast<U>(Q[v_per_thread * lane_idx + i]);
+  }
+
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    o[i] = 0.f;
+  }
+
+  U max_score = Limits<U>::finite_min();
+  U sum_exp_score = 0.f;
+  if (sinks && warp_idx == 0) {
+    max_score = M_LOG2E * static_cast<U>(sinks[head_idx]);
+    sum_exp_score = 1.f;
+  }
+
+  // For each key
+  for (int i = kv_seq_idx; i < params.kL; i += BN) {
+    bool use_key = true;
+    if constexpr (do_causal) {
+      use_key = i <= (params.kL - params.qL + q_seq_idx);
+    }
+
+    if (use_key) {
+      // Read the key
+      PRAGMA_LOOP_UNROLL
+      for (int j = 0; j < v_per_thread; j++) {
+        k[j] = K[v_per_thread * lane_idx + j];
+      }
+
+      // Compute the i-th score
+      U score = 0.f;
+      PRAGMA_LOOP_UNROLL
+      for (int j = 0; j < v_per_thread; j++) {
+        score += q[j] * k[j];
+      }
+
+      // Warp sum
+      score = cg::reduce(warp, score, cg::plus<U>());
+
+      // Update the accumulators
+      U new_max = max(max_score, score);
+      U factor = exp2f(max_score - new_max);
+      U exp_score = exp2f(score - new_max);
+
+      max_score = new_max;
+      sum_exp_score = sum_exp_score * factor + exp_score;
+
+      // Update the output accumulator
+      PRAGMA_LOOP_UNROLL
+      for (int j = 0; j < v_per_thread; j++) {
+        o[j] = o[j] * factor +
+            exp_score * static_cast<U>(V[v_per_thread * lane_idx + j]);
+      }
+    }
+
+    // Move the pointers to the next kv
+    K += inner_k_stride;
+    V += inner_v_stride;
+  }
+
+  if (lane_idx == 0) {
+    max_scores[warp_idx] = max_score;
+    sum_exp_scores[warp_idx] = sum_exp_score;
+  }
+  block.sync();
+
+  max_score = max_scores[lane_idx];
+  U new_max = cg::reduce(warp, max_score, cg::greater<U>());
+  U factor = exp2f(max_score - new_max);
+  sum_exp_score =
+      cg::reduce(warp, sum_exp_scores[lane_idx] * factor, cg::plus<U>());
+  sum_exp_score = sum_exp_score == 0 ? 0 : __frcp_rn(sum_exp_score);
+
+  // Now we need to aggregate all the outputs
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    outputs[lane_idx][warp_idx] = o[i];
+    block.sync();
+    U ot = outputs[warp_idx][lane_idx] * factor;
+    o[i] = cg::reduce(warp, ot, cg::plus<U>()) * sum_exp_score;
+    block.sync();
+  }
+
+  // And write the output
+  if (lane_idx == 0) {
+    PRAGMA_LOOP_UNROLL
+    for (int i = 0; i < v_per_thread; i++) {
+      O[v_per_thread * warp_idx + i] = static_cast<T>(o[i]);
+    }
+  }
+}
+
+template <typename T, bool do_causal, int D>
+__global__ void kernel_sdpav_2pass_1(
+    const T* Q,
+    const T* K,
+    const T* V,
+    const T* sinks,
+    float* partials,
+    float* sums,
+    float* maxs,
+    __grid_constant__ const AttnParams params) {
+  constexpr int BN = 8;
+  constexpr int BD = 32;
+  constexpr int blocks = 32;
+
+  constexpr int v_per_thread = D / BD;
+
+  const int inner_k_stride = blocks * BN * int(params.K_strides[2]);
+  const int inner_v_stride = blocks * BN * int(params.V_strides[2]);
+
+  typedef float U;
+
+  U q[v_per_thread];
+  U k[v_per_thread];
+  U o[v_per_thread];
+
+  __shared__ U outputs[BN][BD + 1];
+  __shared__ U max_scores[BN];
+  __shared__ U sum_exp_scores[BN];
+
+  const U scale_log2 = params.scale * 1.44269504089f;
+
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<32>(block);
+
+  const int lane_idx = warp.thread_rank();
+  const int warp_idx = warp.meta_group_rank();
+
+  // Adjust to thread block and thread
+  const int batch_idx = blockIdx.z / blocks;
+  const int block_idx = blockIdx.z % blocks;
+  const int head_idx = blockIdx.x;
+  const int kv_head_idx = head_idx / params.gqa_factor;
+
+  const int q_seq_idx = blockIdx.y;
+  const int kv_seq_idx = block_idx * BN + warp_idx;
+
+  Q += batch_idx * params.Q_strides[0] + // Batch
+      head_idx * params.Q_strides[1] + // Head
+      q_seq_idx * params.Q_strides[2]; // Sequence
+
+  K += batch_idx * params.K_strides[0] + // Batch
+      kv_head_idx * params.K_strides[1] + // Head
+      kv_seq_idx * params.K_strides[2]; // Sequence
+
+  V += batch_idx * params.V_strides[0] + // Batch
+      kv_head_idx * params.V_strides[1] + // Head
+      kv_seq_idx * params.V_strides[2]; // Sequence
+
+  const int p_stride_s = blocks;
+  const int p_stride_h = params.qL * p_stride_s;
+  const int p_stride_b = params.H * p_stride_h;
+  const int p_offset = batch_idx * p_stride_b + // Batch
+      head_idx * p_stride_h + // Head
+      q_seq_idx * p_stride_s + // Sequence
+      block_idx; // Block
+
+  partials += p_offset * D;
+  sums += p_offset;
+  maxs += p_offset;
+
+  // Read the query and 0 the output accumulator
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    q[i] = scale_log2 * static_cast<U>(Q[v_per_thread * lane_idx + i]);
+  }
+
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    o[i] = 0.f;
+  }
+
+  U max_score = Limits<U>::finite_min();
+  U sum_exp_score = 0.f;
+  if (sinks && warp_idx == 0 && block_idx == 0) {
+    max_score = M_LOG2E * static_cast<U>(sinks[head_idx]);
+    sum_exp_score = 1.f;
+  }
+
+  // For each key
+  for (int i = kv_seq_idx; i < params.kL; i += blocks * BN) {
+    bool use_key = true;
+    if constexpr (do_causal) {
+      use_key = i <= (params.kL - params.qL + q_seq_idx);
+    }
+
+    if (use_key) {
+      // Read the key
+      PRAGMA_LOOP_UNROLL
+      for (int j = 0; j < v_per_thread; j++) {
+        k[j] = K[v_per_thread * lane_idx + j];
+      }
+
+      // Compute the i-th score
+      U score = 0.f;
+      PRAGMA_LOOP_UNROLL
+      for (int j = 0; j < v_per_thread; j++) {
+        score += q[j] * k[j];
+      }
+
+      // Warp sum
+      score = cg::reduce(warp, score, cg::plus<U>());
+
+      // Update the accumulators
+      U new_max = max(max_score, score);
+      U factor = exp2f(max_score - new_max);
+      U exp_score = exp2f(score - new_max);
+
+      max_score = new_max;
+      sum_exp_score = sum_exp_score * factor + exp_score;
+
+      // Update the output accumulator
+      PRAGMA_LOOP_UNROLL
+      for (int j = 0; j < v_per_thread; j++) {
+        o[j] = o[j] * factor +
+            exp_score * static_cast<U>(V[v_per_thread * lane_idx + j]);
+      }
+    }
+
+    // Move the pointers to the next kv
+    K += inner_k_stride;
+    V += inner_v_stride;
+  }
+
+  if (lane_idx == 0) {
+    max_scores[warp_idx] = max_score;
+    sum_exp_scores[warp_idx] = sum_exp_score;
+  }
+
+  block.sync();
+
+  max_score = (lane_idx < BN) ? max_scores[lane_idx] : -1e9;
+  U new_max = cg::reduce(warp, max_score, cg::greater<U>());
+  U factor = exp2f(max_score - new_max);
+  sum_exp_score = (lane_idx < BN) ? sum_exp_scores[lane_idx] : 0.f;
+  sum_exp_score = cg::reduce(warp, sum_exp_score * factor, cg::plus<U>());
+
+  // Write the sum and new max
+  if (warp_idx == 0) {
+    sums[0] = sum_exp_score;
+    maxs[0] = new_max;
+  }
+
+  // Now we need to aggregate all the outputs
+  auto ff = exp2f(max_scores[warp_idx] - new_max);
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    outputs[warp_idx][lane_idx] = o[i] * ff;
+    block.sync();
+
+    if (warp_idx == 0) {
+      U ot = outputs[0][lane_idx];
+      PRAGMA_LOOP_UNROLL
+      for (int j = 1; j < BN; j++) {
+        ot += outputs[j][lane_idx];
+        warp.sync();
+      }
+      o[i] = ot;
+    }
+    block.sync();
+  }
+
+  if (warp_idx == 0) {
+    PRAGMA_LOOP_UNROLL
+    for (int i = 0; i < v_per_thread; i++) {
+      partials[v_per_thread * lane_idx + i] = o[i];
+    }
+  }
+}
+
+template <typename T, bool do_causal, int D>
+__global__ void kernel_sdpav_2pass_2(
+    const float* partials,
+    const float* sums,
+    const float* maxs,
+    T* O,
+    __grid_constant__ const AttnParams params) {
+  constexpr int BN = 32;
+  constexpr int BD = 32;
+  constexpr int blocks = 32;
+
+  constexpr int v_per_thread = D / BD;
+
+  typedef float U;
+
+  U o[v_per_thread];
+  __shared__ U outputs[BN][BD + 1];
+
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<32>(block);
+
+  const int lane_idx = warp.thread_rank();
+  const int warp_idx = warp.meta_group_rank();
+
+  // Adjust to thread block and thread
+  const int batch_idx = blockIdx.z;
+  const int head_idx = blockIdx.x;
+  const int q_seq_idx = blockIdx.y;
+
+  const int p_stride_s = blocks;
+  const int p_stride_h = params.qL * p_stride_s;
+  const int p_stride_b = params.H * p_stride_h;
+  const int p_offset = batch_idx * p_stride_b + // Batch
+      head_idx * p_stride_h + // Head
+      q_seq_idx * p_stride_s; // Sequence
+
+  partials += p_offset * D + warp_idx * D;
+  sums += p_offset;
+  maxs += p_offset;
+
+  O += batch_idx * params.O_strides[0] + // Batch
+      head_idx * params.O_strides[1] + // Head
+      q_seq_idx * params.O_strides[2]; // Sequence
+
+  U max_score = maxs[lane_idx];
+  U new_max = cg::reduce(warp, max_score, cg::greater<U>());
+  U factor = exp2f(max_score - new_max);
+  U sum_exp_score = cg::reduce(warp, sums[lane_idx] * factor, cg::plus<U>());
+  sum_exp_score = sum_exp_score == 0 ? 0 : __frcp_rn(sum_exp_score);
+
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    o[i] = partials[v_per_thread * lane_idx + i];
+  }
+
+  // Now we need to aggregate all the outputs
+  PRAGMA_LOOP_UNROLL
+  for (int i = 0; i < v_per_thread; i++) {
+    outputs[lane_idx][warp_idx] = o[i];
+    block.sync();
+    U ot = outputs[warp_idx][lane_idx] * factor;
+    o[i] = cg::reduce(warp, ot, cg::plus<U>()) * sum_exp_score;
+    block.sync();
+  }
+
+  // And write the output
+  if (lane_idx == 0) {
+    PRAGMA_LOOP_UNROLL
+    for (int i = 0; i < v_per_thread; i++) {
+      O[v_per_thread * warp_idx + i] = static_cast<T>(o[i]);
+    }
+  }
+}
+
+} // namespace cu
+
+namespace {
+
+template <typename F>
+void dispatch_headdim(int n, F&& f) {
+  switch (n) {
+    case 64:
+      f(std::integral_constant<int, 64>{});
+      break;
+    case 96:
+      f(std::integral_constant<int, 96>{});
+      break;
+    case 128:
+      f(std::integral_constant<int, 128>{});
+      break;
+  }
+}
+
+void sdpa_vector_1pass_fallback(
+    const Stream& s,
+    cu::CommandEncoder& encoder,
+    const array& q,
+    const array& k,
+    const array& v,
+    const float scale,
+    array& o,
+    bool do_causal,
+    const std::optional<array>& sinks) {
+  encoder.set_input_array(q);
+  encoder.set_input_array(k);
+  encoder.set_input_array(v);
+  if (sinks) {
+    encoder.set_input_array(*sinks);
+  }
+  encoder.set_output_array(o);
+
+  cu::AttnParams params{
+      /* int B = */ q.shape(0),
+      /* int H = */ q.shape(1),
+      /* int D = */ q.shape(3),
+
+      /* int qL = */ q.shape(2),
+      /* int kL = */ k.shape(2),
+
+      /* int gqa_factor = */ q.shape(1) / k.shape(1),
+      /* float scale = */ scale,
+
+      /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)},
+      /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)},
+      /* int64_t V_strides[3] = */ {v.strides(0), v.strides(1), v.strides(2)},
+      /* int64_t O_strides[3] = */ {o.strides(0), o.strides(1), o.strides(2)}};
+
+  dim3 grid_dim(params.H, params.qL, params.B);
+  dim3 block_dim(1024, 1, 1);
+
+  dispatch_float_types(o.dtype(), "kernel_sdpav_1pass", [&](auto type_tag) {
+    dispatch_bool(do_causal, [&](auto do_causal) {
+      dispatch_headdim(params.D, [&](auto headdim) {
+        using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+
+        auto kernel =
+            cu::kernel_sdpav_1pass<DataType, do_causal.value, headdim.value>;
+        encoder.add_kernel_node(
+            kernel,
+            grid_dim,
+            block_dim,
+            0,
+            gpu_ptr<DataType>(q),
+            gpu_ptr<DataType>(k),
+            gpu_ptr<DataType>(v),
+            gpu_ptr<DataType>(o),
+            sinks ? gpu_ptr<DataType>(*sinks) : nullptr,
+            params);
+      });
+    });
+  });
+}
+
+void sdpa_vector_2pass_fallback(
+    const Stream& s,
+    cu::CommandEncoder& encoder,
+    const array& q,
+    const array& k,
+    const array& v,
+    const float scale,
+    array& o,
+    bool do_causal,
+    const std::optional<array>& sinks) {
+  cu::AttnParams params{
+      /* int B = */ q.shape(0),
+      /* int H = */ q.shape(1),
+      /* int D = */ q.shape(3),
+
+      /* int qL = */ q.shape(2),
+      /* int kL = */ k.shape(2),
+
+      /* int gqa_factor = */ q.shape(1) / k.shape(1),
+      /* float scale = */ scale,
+
+      /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)},
+      /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)},
+      /* int64_t V_strides[3] = */ {v.strides(0), v.strides(1), v.strides(2)},
+      /* int64_t O_strides[3] = */ {o.strides(0), o.strides(1), o.strides(2)}};
+
+  // Allocate the intermediates
+  int blocks = 32;
+
+  Shape intermediate_shape;
+  intermediate_shape.reserve(o.ndim() + 1);
+  intermediate_shape.insert(
+      intermediate_shape.end(), o.shape().begin(), o.shape().end() - 1);
+  intermediate_shape.push_back(blocks);
+  intermediate_shape.push_back(o.shape().back());
+
+  array intermediate(intermediate_shape, float32, nullptr, {});
+  intermediate_shape.pop_back();
+  array sums(intermediate_shape, float32, nullptr, {});
+  array maxs(std::move(intermediate_shape), float32, nullptr, {});
+
+  intermediate.set_data(cu::malloc_async(intermediate.nbytes(), encoder));
+  sums.set_data(cu::malloc_async(sums.nbytes(), encoder));
+  maxs.set_data(cu::malloc_async(maxs.nbytes(), encoder));
+
+  encoder.add_temporary(intermediate);
+  encoder.add_temporary(sums);
+  encoder.add_temporary(maxs);
+
+  dispatch_float_types(o.dtype(), "kernel_sdpav_2pass", [&](auto type_tag) {
+    dispatch_bool(do_causal, [&](auto do_causal) {
+      dispatch_headdim(params.D, [&](auto headdim) {
+        using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+
+        {
+          auto kernel = cu::
+              kernel_sdpav_2pass_1<DataType, do_causal.value, headdim.value>;
+
+          encoder.set_input_array(q);
+          encoder.set_input_array(k);
+          encoder.set_input_array(v);
+          if (sinks) {
+            encoder.set_input_array(*sinks);
+          }
+
+          encoder.set_output_array(intermediate);
+          encoder.set_output_array(sums);
+          encoder.set_output_array(maxs);
+
+          dim3 grid_dim(params.H, params.qL, params.B * 32);
+          dim3 block_dim(8 * 32, 1, 1);
+
+          encoder.add_kernel_node(
+              kernel,
+              grid_dim,
+              block_dim,
+              0,
+              gpu_ptr<DataType>(q),
+              gpu_ptr<DataType>(k),
+              gpu_ptr<DataType>(v),
+              sinks ? gpu_ptr<DataType>(*sinks) : nullptr,
+              gpu_ptr<float>(intermediate),
+              gpu_ptr<float>(sums),
+              gpu_ptr<float>(maxs),
+              params);
+        }
+
+        {
+          auto kernel = cu::
+              kernel_sdpav_2pass_2<DataType, do_causal.value, headdim.value>;
+
+          encoder.set_input_array(intermediate);
+          encoder.set_input_array(sums);
+          encoder.set_input_array(maxs);
+          encoder.set_output_array(o);
+
+          dim3 grid_dim(params.H, params.qL, params.B);
+          dim3 block_dim(1024, 1, 1);
+
+          encoder.add_kernel_node(
+              kernel,
+              grid_dim,
+              block_dim,
+              0,
+              gpu_ptr<float>(intermediate),
+              gpu_ptr<float>(sums),
+              gpu_ptr<float>(maxs),
+              gpu_ptr<DataType>(o),
+              params);
+        }
+      });
+    });
+  });
+}
+
+void sdpa_vector_fallback(
+    const Stream& s,
+    cu::CommandEncoder& encoder,
+    const array& q,
+    const array& k,
+    const array& v,
+    const float scale,
+    array& o,
+    bool do_causal,
+    const std::optional<array>& sinks) {
+  int kL = k.shape(2);
+
+  if (kL > 1024) {
+    return sdpa_vector_2pass_fallback(
+        s, encoder, q, k, v, scale, o, do_causal, sinks);
+  } else {
+    return sdpa_vector_1pass_fallback(
+        s, encoder, q, k, v, scale, o, do_causal, sinks);
+  }
+}
+
+} // namespace
+
+bool supports_sdpa_vector(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool output_logsumexp) {
+  if (output_logsumexp) {
+    return false;
+  }
+
+  const int value_head_dim = v.shape(-1);
+  const int query_head_dim = q.shape(-1);
+  const int query_sequence_length = q.shape(2);
+  const int key_sequence_length = k.shape(2);
+
+  const bool sdpa_supported_head_dim = query_head_dim == value_head_dim &&
+      (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128);
+
+  const bool supported_vector_config =
+      sdpa_supported_head_dim && query_sequence_length < 4;
+
+  return supported_vector_config && !has_arr_mask;
+}
+
+void sdpa_vector(
+    const array& q_pre,
+    const array& k_pre,
+    const array& v_pre,
+    float scale,
+    array& o,
+    bool do_causal,
+    const std::optional<array>& sinks_pre,
+    Stream s) {
+  auto& encoder = cu::get_command_encoder(s);
+  std::vector<array> copies;
+
+  // Define some copy functions to ensure the layout of the inputs is as
+  // expected.
+  copies.reserve(4);
+  auto copy_unless = [&copies, &s](
+                         auto predicate, const array& arr) -> const array& {
+    if (!predicate(arr)) {
+      array arr_copy = contiguous_copy_gpu(arr, s);
+      copies.push_back(std::move(arr_copy));
+      return copies.back();
+    } else {
+      return arr;
+    }
+  };
+
+  // Checks that the headdim dimension has stride 1.
+  auto is_matrix_contiguous = [](const array& arr) {
+    return arr.strides(-1) == 1;
+  };
+
+  std::optional<array> sinks = std::nullopt;
+  if (sinks_pre) {
+    sinks = copy_unless(is_matrix_contiguous, sinks_pre.value());
+  }
+
+  // We are in vector mode ie single query
+  if (q_pre.shape(2) < 4) {
+    auto q_copy_unless = [](const array& arr) {
+      if (arr.flags().row_contiguous) {
+        return true;
+      }
+      auto& strides = arr.strides();
+      auto& shape = arr.shape();
+      if (shape[0] == 1 || shape[1] == 1) {
+        // If either the batch or head dimension is a singleton, the other can
+        // be transposed with the sequence dimension
+        auto bidx = shape[0] == 1 ? 1 : 0;
+        return (strides[3] == 1) && (strides[2] == shape[3] * shape[bidx]) &&
+            (strides[bidx] == shape[3]);
+      }
+      return false;
+    };
+
+    auto kv_copy_unless = [](const array& arr) {
+      // keys and values should be copied if:
+      // - the last dimension is not contiguous
+      // - the batch and head dim are not contiguous
+      auto& strides = arr.strides();
+      auto& shape = arr.shape();
+      if (strides.back() != 1) {
+        return false;
+      }
+      if (shape[0] == 1 || shape[1] == 1) {
+        return true;
+      }
+      return (strides[0] == strides[1] * shape[1]);
+    };
+
+    const auto& q = copy_unless(q_copy_unless, q_pre);
+    const auto& k = copy_unless(kv_copy_unless, k_pre);
+    const auto& v = copy_unless(kv_copy_unless, v_pre);
+
+    // Donate the query if possible
+    if (q.is_donatable() && q.flags().row_contiguous && q.size() == o.size()) {
+      o.copy_shared_buffer(q);
+    } else {
+      int64_t str_oD = 1;
+      int64_t str_oH = o.shape(3);
+      int64_t str_oL = o.shape(1) * str_oH;
+      int64_t str_oB = o.shape(2) * str_oL;
+
+      array::Flags flags{
+          /* bool contiguous = */ 1,
+          /* bool row_contiguous = */ o.shape(2) == 1,
+          /* bool col_contiguous = */ o.size() == o.shape(3),
+      };
+
+      o.set_data(
+          cu::malloc_async(o.nbytes(), encoder),
+          o.size(),
+          {str_oB, str_oH, str_oL, str_oD},
+          flags);
+    }
+
+    for (const auto& cp : copies) {
+      encoder.add_temporary(cp);
+    }
+
+    sdpa_vector_fallback(s, encoder, q, k, v, scale, o, do_causal, sinks);
+  }
+
+  // Full attention mode should never reach here
+  else {
+    throw std::runtime_error("Doesn't support matrix yet.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/scan.cu b/macos/third_party/mlx/mlx/backend/cuda/scan.cu
new file mode 100644
index 0000000..b4d6cf3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/scan.cu
@@ -0,0 +1,468 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/binary_ops.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/reduce/reduce_ops.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/scan.h>
+#include <nvtx3/nvtx3.hpp>
+
+#include <cassert>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename Op, typename T>
+struct ScanResult {
+  using type = T;
+};
+
+template <>
+struct ScanResult<Sum, bool> {
+  using type = int32_t;
+};
+
+template <typename T>
+struct ReduceInit<LogAddExp, T> {
+  static constexpr __host__ __device__ T value() {
+    return Limits<T>::min();
+  }
+};
+
+template <bool reverse, typename T, typename U, int N_READS>
+inline __device__ void
+load_values(int index, const T* in, U (&values)[N_READS], int size, U init) {
+  int remaining = size - index * N_READS;
+  if constexpr (reverse) {
+    in += remaining - N_READS;
+    if (remaining < N_READS) {
+      for (int i = 0; i < N_READS; ++i) {
+        values[N_READS - i - 1] =
+            (N_READS - i - 1 < remaining) ? cast_to<U>(in[i]) : init;
+      }
+    } else {
+      for (int i = 0; i < N_READS; ++i) {
+        values[N_READS - i - 1] = cast_to<U>(in[i]);
+      }
+    }
+  } else {
+    in += index * N_READS;
+    if (remaining < N_READS) {
+      for (int i = 0; i < N_READS; ++i) {
+        values[i] = (i < remaining) ? cast_to<U>(in[i]) : init;
+      }
+    } else {
+      for (int i = 0; i < N_READS; ++i) {
+        values[i] = cast_to<U>(in[i]);
+      }
+    }
+  }
+}
+
+template <bool reverse, int offset, typename T, int N_READS>
+inline __device__ void
+store_values(int index, T* out, T (&values)[N_READS], int size) {
+  int start = index * N_READS + offset;
+  int remaining = size - start;
+  if constexpr (reverse) {
+    out += remaining - N_READS;
+    if (remaining < N_READS) {
+      for (int i = 0; i < N_READS; ++i) {
+        if (N_READS - i - 1 < remaining) {
+          out[i] = values[N_READS - i - 1];
+        }
+      }
+    } else {
+      for (int i = 0; i < N_READS; ++i) {
+        out[i] = values[N_READS - i - 1];
+      }
+    }
+  } else {
+    out += start;
+    if (remaining < N_READS) {
+      for (int i = 0; i < N_READS; ++i) {
+        if (i < remaining) {
+          out[i] = values[i];
+        }
+      }
+    } else {
+      for (int i = 0; i < N_READS; ++i) {
+        out[i] = values[i];
+      }
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS,
+    bool inclusive,
+    bool reverse>
+__global__ void contiguous_scan(const T* in, U* out, int32_t axis_size) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  in += grid.block_rank() * axis_size;
+  out += grid.block_rank() * axis_size;
+
+  __shared__ U warp_sums[WARP_SIZE];
+
+  Op op;
+  U init = ReduceInit<Op, T>::value();
+  U prefix = init;
+
+  // Scan per block.
+  for (int r = 0; r < cuda::ceil_div(axis_size, block.size() * N_READS); ++r) {
+    int32_t index = r * block.size() + block.thread_rank();
+    U values[N_READS];
+    load_values<reverse>(index, in, values, axis_size, init);
+
+    // Compute an inclusive scan per thread.
+    for (int i = 1; i < N_READS; ++i) {
+      values[i] = op(values[i], values[i - 1]);
+    }
+
+    // Compute exclusive scan of thread sums.
+    U prev_thread_sum = cg::exclusive_scan(warp, values[N_READS - 1], op);
+    if (warp.thread_rank() == 0) {
+      prev_thread_sum = init;
+    }
+
+    // Write wrap's sum to shared memory.
+    if (warp.thread_rank() == WARP_SIZE - 1) {
+      warp_sums[warp.meta_group_rank()] =
+          op(prev_thread_sum, values[N_READS - 1]);
+    }
+    block.sync();
+
+    // Compute exclusive scan of warp sums.
+    if (warp.meta_group_rank() == 0) {
+      U prev_warp_sum =
+          cg::exclusive_scan(warp, warp_sums[warp.thread_rank()], op);
+      if (warp.thread_rank() == 0) {
+        prev_warp_sum = init;
+      }
+      warp_sums[warp.thread_rank()] = prev_warp_sum;
+    }
+    block.sync();
+
+    // Compute the output.
+    for (int i = 0; i < N_READS; ++i) {
+      values[i] = op(values[i], prefix);
+      values[i] = op(values[i], warp_sums[warp.meta_group_rank()]);
+      values[i] = op(values[i], prev_thread_sum);
+    }
+
+    // Write the values.
+    if (inclusive) {
+      store_values<reverse, 0>(index, out, values, axis_size);
+    } else {
+      store_values<reverse, 1>(index, out, values, axis_size);
+      if (reverse) {
+        if (block.thread_rank() == 0 && index == 0) {
+          out[axis_size - 1] = init;
+        }
+      } else {
+        if (block.thread_rank() == 0 && index == 0) {
+          out[0] = init;
+        }
+      }
+    }
+    block.sync();
+
+    // Share the prefix.
+    if ((warp.meta_group_rank() == warp.meta_group_size() - 1) &&
+        (warp.thread_rank() == WARP_SIZE - 1)) {
+      warp_sums[0] = values[N_READS - 1];
+    }
+    block.sync();
+    prefix = warp_sums[0];
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS,
+    int BM,
+    int BN,
+    bool inclusive,
+    bool reverse>
+__global__ void strided_scan(
+    const T* in,
+    U* out,
+    int32_t axis_size,
+    int64_t stride,
+    int64_t stride_blocks) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  constexpr int BN_pad = WARP_SIZE + 16 / sizeof(U);
+  constexpr int n_warps = BN / N_READS;
+  constexpr int n_scans = BN / n_warps;
+
+  __shared__ U read_buffer[BM * BN_pad];
+
+  Op op;
+  U init = ReduceInit<Op, T>::value();
+  U values[n_scans];
+  U prefix[n_scans];
+  for (int i = 0; i < n_scans; ++i) {
+    prefix[i] = init;
+  }
+
+  // Compute offsets.
+  int64_t offset = (grid.block_rank() / stride_blocks) * axis_size * stride;
+  int64_t global_index_x = (grid.block_rank() % stride_blocks) * BN;
+  uint32_t read_offset_y = (block.thread_rank() * N_READS) / BN;
+  uint32_t read_offset_x = (block.thread_rank() * N_READS) % BN;
+  uint32_t scan_offset_y = warp.thread_rank();
+  uint32_t scan_offset_x = warp.meta_group_rank() * n_scans;
+
+  uint32_t stride_limit = stride - global_index_x;
+  in += offset + global_index_x + read_offset_x;
+  out += offset + global_index_x + read_offset_x;
+  U* read_into = read_buffer + read_offset_y * BN_pad + read_offset_x;
+  U* read_from = read_buffer + scan_offset_y * BN_pad + scan_offset_x;
+
+  for (uint32_t j = 0; j < axis_size; j += BM) {
+    // Calculate the indices for the current thread.
+    uint32_t index_y = j + read_offset_y;
+    uint32_t check_index_y = index_y;
+    if (reverse) {
+      index_y = axis_size - 1 - index_y;
+    }
+
+    // Read in SM.
+    if (check_index_y < axis_size && (read_offset_x + N_READS) < stride_limit) {
+      for (int i = 0; i < N_READS; ++i) {
+        read_into[i] = in[index_y * stride + i];
+      }
+    } else {
+      for (int i = 0; i < N_READS; ++i) {
+        if (check_index_y < axis_size && (read_offset_x + i) < stride_limit) {
+          read_into[i] = in[index_y * stride + i];
+        } else {
+          read_into[i] = init;
+        }
+      }
+    }
+    block.sync();
+
+    // Read strided into registers.
+    for (int i = 0; i < n_scans; ++i) {
+      values[i] = read_from[i];
+    }
+
+    // Perform the scan.
+    for (int i = 0; i < n_scans; ++i) {
+      values[i] = cg::inclusive_scan(warp, values[i], op);
+      values[i] = op(values[i], prefix[i]);
+      prefix[i] = warp.shfl(values[i], WARP_SIZE - 1);
+    }
+
+    // Write to SM.
+    for (int i = 0; i < n_scans; ++i) {
+      read_from[i] = values[i];
+    }
+    block.sync();
+
+    // Write to device memory.
+    if (!inclusive) {
+      if (check_index_y == 0) {
+        if ((read_offset_x + N_READS) < stride_limit) {
+          for (int i = 0; i < N_READS; ++i) {
+            out[index_y * stride + i] = init;
+          }
+        } else {
+          for (int i = 0; i < N_READS; ++i) {
+            if ((read_offset_x + i) < stride_limit) {
+              out[index_y * stride + i] = init;
+            }
+          }
+        }
+      }
+      if (reverse) {
+        index_y -= 1;
+        check_index_y += 1;
+      } else {
+        index_y += 1;
+        check_index_y += 1;
+      }
+    }
+    if (check_index_y < axis_size && (read_offset_x + N_READS) < stride_limit) {
+      for (int i = 0; i < N_READS; ++i) {
+        out[index_y * stride + i] = read_into[i];
+      }
+    } else {
+      for (int i = 0; i < N_READS; ++i) {
+        if (check_index_y < axis_size && (read_offset_x + i) < stride_limit) {
+          out[index_y * stride + i] = read_into[i];
+        }
+      }
+    }
+  }
+}
+
+} // namespace cu
+
+template <typename F>
+void dispatch_scan_ops(Scan::ReduceType scan_op, F&& f) {
+  if (scan_op == Scan::ReduceType::Max) {
+    f(type_identity<cu::Max>{});
+  } else if (scan_op == Scan::ReduceType::Min) {
+    f(type_identity<cu::Min>{});
+  } else if (scan_op == Scan::ReduceType::Sum) {
+    f(type_identity<cu::Sum>{});
+  } else if (scan_op == Scan::ReduceType::Prod) {
+    f(type_identity<cu::Prod>{});
+  } else if (scan_op == Scan::ReduceType::LogAddExp) {
+    f(type_identity<cu::LogAddExp>{});
+  } else {
+    throw std::invalid_argument("Unknown reduce type.");
+  }
+}
+
+template <typename Op>
+const char* op_to_string() {
+  if (cuda::std::is_same_v<Op, cu::Max>) {
+    return "Max";
+  } else if (cuda::std::is_same_v<Op, cu::Min>) {
+    return "Min";
+  } else if (cuda::std::is_same_v<Op, cu::Sum>) {
+    return "Sum";
+  } else if (cuda::std::is_same_v<Op, cu::Prod>) {
+    return "Prod";
+  } else if (cuda::std::is_same_v<Op, cu::LogAddExp>) {
+    return "LogAddExp";
+  } else {
+    throw std::invalid_argument("Unknown op.");
+  }
+}
+
+template <typename Op, typename T>
+constexpr bool supports_scan_op() {
+  if constexpr (cuda::std::is_same_v<Op, LogAddExp>) {
+    return is_inexact_v<T>;
+  } else {
+    return true;
+  }
+}
+
+void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Scan::eval_gpu");
+  assert(inputs.size() == 1);
+  auto in = inputs[0];
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  if (in.flags().contiguous && in.strides()[axis_] != 0) {
+    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(
+          cu::malloc_async(in.data_size() * out.itemsize(), encoder),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+    }
+  } else {
+    in = contiguous_copy_gpu(in, s);
+    out.copy_shared_buffer(in);
+  }
+
+  constexpr int N_READS = 4;
+  int32_t axis_size = in.shape(axis_);
+  bool contiguous = in.strides()[axis_] == 1;
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    dispatch_scan_ops(reduce_type_, [&](auto scan_op_tag) {
+      using Op = MLX_GET_TYPE(scan_op_tag);
+      if constexpr (supports_scan_op<Op, T>()) {
+        using U = typename cu::ScanResult<Op, T>::type;
+        dispatch_bool(inclusive_, [&](auto inclusive) {
+          dispatch_bool(reverse_, [&](auto reverse) {
+            if (contiguous) {
+              auto kernel = cu::contiguous_scan<
+                  T,
+                  U,
+                  Op,
+                  N_READS,
+                  inclusive.value,
+                  reverse.value>;
+              int block_dim = cuda::ceil_div(axis_size, N_READS);
+              block_dim = cuda::ceil_div(block_dim, WARP_SIZE) * WARP_SIZE;
+              block_dim = std::min(block_dim, WARP_SIZE * WARP_SIZE);
+              encoder.add_kernel_node(
+                  kernel,
+                  in.data_size() / axis_size,
+                  block_dim,
+                  0,
+                  gpu_ptr<T>(in),
+                  gpu_ptr<U>(out),
+                  axis_size);
+            } else {
+              constexpr int BM = WARP_SIZE;
+              constexpr int BN = WARP_SIZE;
+              auto kernel = cu::strided_scan<
+                  T,
+                  U,
+                  Op,
+                  N_READS,
+                  BM,
+                  BN,
+                  inclusive.value,
+                  reverse.value>;
+              int64_t stride = in.strides()[axis_];
+              int64_t stride_blocks = cuda::ceil_div(stride, BN);
+              dim3 num_blocks = get_2d_grid_dims(
+                  in.shape(), in.strides(), axis_size * stride);
+              if (num_blocks.x * stride_blocks <= UINT32_MAX) {
+                num_blocks.x *= stride_blocks;
+              } else {
+                num_blocks.y *= stride_blocks;
+              }
+              int block_dim = (BN / N_READS) * WARP_SIZE;
+              encoder.add_kernel_node(
+                  kernel,
+                  num_blocks,
+                  block_dim,
+                  0,
+                  gpu_ptr<T>(in),
+                  gpu_ptr<U>(out),
+                  axis_size,
+                  stride,
+                  stride_blocks);
+            }
+          });
+        });
+      } else {
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do scan op {} on inputs of {} with result of {}.",
+                op_to_string<Op>(),
+                dtype_to_string(in.dtype()),
+                dtype_to_string(out.dtype())));
+      }
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/slicing.cpp b/macos/third_party/mlx/mlx/backend/cuda/slicing.cpp
new file mode 100644
index 0000000..d45b16d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/slicing.cpp
@@ -0,0 +1,111 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/slicing.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/gpu/slicing.h"
+#include "mlx/dtype_utils.h"
+
+#include <numeric>
+
+namespace mlx::core {
+
+void concatenate_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    int axis,
+    const Stream& s) {
+  std::vector<int> sizes;
+  sizes.push_back(0);
+  for (auto& p : inputs) {
+    sizes.push_back(p.shape(axis));
+  }
+  std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
+
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  auto strides = out.strides();
+  auto flags = out.flags();
+  flags.row_contiguous = false;
+  flags.col_contiguous = false;
+  flags.contiguous = false;
+  auto concurrent = encoder.concurrent_context();
+  for (int i = 0; i < inputs.size(); i++) {
+    array out_slice(inputs[i].shape(), out.dtype(), nullptr, {});
+    size_t data_offset = strides[axis] * sizes[i];
+    out_slice.copy_shared_buffer(
+        out, strides, flags, out_slice.size(), data_offset);
+    copy_gpu_inplace(inputs[i], out_slice, CopyType::GeneralGeneral, s);
+  }
+}
+
+array compute_dynamic_offset(
+    const array& indices,
+    const Strides& strides,
+    const std::vector<int>& axes,
+    const Stream& s) {
+  Dtype dtype = indices.dtype();
+  int nidx = axes.size();
+
+  std::string module_name =
+      fmt::format("compute_dynamic_offset_{}_{}", dtype_to_string(dtype), nidx);
+  std::string kernel_name = fmt::format(
+      "mlx::core::cu::compute_dynamic_offset<{}, {}>",
+      dtype_to_cuda_type(dtype),
+      nidx);
+
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::string source = R"(
+        #include "mlx/backend/cuda/device/utils.cuh"
+
+        namespace mlx::core::cu {
+
+        template <typename T, int NIDX>
+        __global__ void compute_dynamic_offset(
+            const T* indices,
+            int64_t* offset,
+            const __grid_constant__ Strides strides,
+            const __grid_constant__ cuda::std::array<int, NIDX> axes) {
+          int64_t acc = 0;
+          #pragma unroll
+          for (int i = 0; i < NIDX; ++i) {
+            acc += indices[i] * strides[axes[i]];
+          }
+          *offset = acc;
+        }
+
+        } // namespace mlx::core::cu
+    )";
+    return std::make_tuple(false, std::move(source), std::vector{kernel_name});
+  });
+
+  auto& encoder = cu::get_command_encoder(s);
+  // Prepare output.
+  array offset({1}, int64, nullptr, {});
+  bool donate = indices.is_donatable() &&
+      (indices.data_size() * indices.itemsize()) >= offset.itemsize();
+  if (donate) {
+    offset.copy_shared_buffer(indices);
+  } else {
+    offset.set_data(cu::malloc_async(offset.itemsize(), encoder));
+  }
+
+  encoder.add_temporary(offset);
+  encoder.set_input_array(indices);
+  encoder.set_output_array(offset);
+
+  cu::KernelArgs args;
+  args.append(indices);
+  args.append(offset);
+  args.append_ndim(strides);
+  args.append(axes);
+
+  auto kernel = mod.get_kernel(kernel_name);
+  encoder.add_kernel_node(kernel, 1, 1, 0, args.args());
+
+  return offset;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/softmax.cu b/macos/third_party/mlx/mlx/backend/cuda/softmax.cu
new file mode 100644
index 0000000..4b2a35b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/softmax.cu
@@ -0,0 +1,162 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/cast_op.cuh"
+#include "mlx/backend/cuda/device/fp16_math.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include <nvtx3/nvtx3.hpp>
+
+#include <cassert>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T>
+inline __device__ T softmax_exp(T x) {
+  // Softmax doesn't need high precision exponential cause x is gonna be in
+  // (-oo, 0] anyway and subsequently it will be divided by sum(exp(x_i)).
+  return __expf(x);
+}
+
+template <typename T, typename AccT, int BLOCK_DIM, int N_READS = 4>
+__global__ void softmax(const T* in, T* out, int axis_size) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  in += grid.block_rank() * axis_size;
+  out += grid.block_rank() * axis_size;
+
+  cg::greater<AccT> max_op;
+  cg::plus<AccT> plus_op;
+
+  // Thread reduce.
+  AccT prevmax;
+  AccT maxval = Limits<AccT>::finite_min();
+  AccT normalizer = cast_to<AccT>(0);
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto vals = load_vector<N_READS>(in, index, axis_size, Limits<T>::min());
+    prevmax = maxval;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      maxval = max_op(maxval, static_cast<AccT>(vals[i]));
+    }
+
+    // Online normalizer calculation for softmax:
+    // https://github.com/NVIDIA/online-softmax
+    normalizer = normalizer * softmax_exp(prevmax - maxval);
+#pragma unroll
+    for (int i = 0; i < N_READS; i++) {
+      normalizer =
+          normalizer + softmax_exp(static_cast<AccT>(vals[i]) - maxval);
+    }
+  }
+
+  // First warp reduce.
+  prevmax = maxval;
+  maxval = cg::reduce(warp, maxval, max_op);
+  normalizer = normalizer * softmax_exp(prevmax - maxval);
+  normalizer = cg::reduce(warp, normalizer, plus_op);
+
+  __shared__ AccT local_max[WARP_SIZE];
+  __shared__ AccT local_normalizer[WARP_SIZE];
+
+  // Write to shared memory and do second warp reduce.
+  prevmax = maxval;
+  if (warp.thread_rank() == 0) {
+    local_max[warp.meta_group_rank()] = maxval;
+  }
+  block.sync();
+  maxval = warp.thread_rank() < warp.meta_group_size()
+      ? local_max[warp.thread_rank()]
+      : Limits<AccT>::min();
+  maxval = cg::reduce(warp, maxval, max_op);
+  normalizer = normalizer * softmax_exp(prevmax - maxval);
+  if (warp.thread_rank() == 0) {
+    local_normalizer[warp.meta_group_rank()] = normalizer;
+  }
+  block.sync();
+  normalizer = warp.thread_rank() < warp.meta_group_size()
+      ? local_normalizer[warp.thread_rank()]
+      : AccT{};
+  normalizer = cg::reduce(warp, normalizer, plus_op);
+  normalizer = 1 / normalizer;
+
+  // Write output.
+  for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) {
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto vals = load_vector<N_READS>(in, index, axis_size, T(0));
+    for (int i = 0; i < N_READS; i++) {
+      vals[i] = softmax_exp(static_cast<AccT>(vals[i]) - maxval) * normalizer;
+    }
+    store_vector<N_READS>(out, index, vals, axis_size);
+  }
+}
+
+} // namespace cu
+
+void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Softmax::eval_gpu");
+  assert(inputs.size() == 1);
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  // Make sure that the last dimension is contiguous.
+  auto set_output = [&s, &out, &encoder](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            cu::malloc_async(x.data_size() * x.itemsize(), encoder),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  array in = set_output(inputs[0]);
+  bool precise = in.dtype() != float32 && precise_;
+
+  int axis_size = in.shape().back();
+  int n_rows = in.data_size() / axis_size;
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "softmax", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
+    dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
+      auto kernel = cu::softmax<DataType, DataType, block_dim(), N_READS>;
+      if (precise) {
+        kernel = cu::softmax<DataType, float, block_dim(), N_READS>;
+      }
+      encoder.add_kernel_node(
+          kernel,
+          n_rows,
+          block_dim(),
+          0,
+          gpu_ptr<DataType>(in),
+          gpu_ptr<DataType>(out),
+          axis_size);
+    });
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/sort.cu b/macos/third_party/mlx/mlx/backend/cuda/sort.cu
new file mode 100644
index 0000000..c258c01
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/sort.cu
@@ -0,0 +1,1076 @@
+// Copyright © 2025 Apple Inc.
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/fp16_math.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+#include <cuda/std/limits>
+#include <cuda/std/type_traits>
+
+namespace mlx::core {
+
+constexpr int N_PER_THREAD = 8;
+
+namespace cu {
+
+template <typename T>
+__device__ __forceinline__ T nan_value();
+
+template <>
+__device__ __forceinline__ float nan_value<float>() {
+  return cuda::std::numeric_limits<float>::quiet_NaN();
+}
+
+template <>
+__device__ __forceinline__ double nan_value<double>() {
+  return cuda::std::numeric_limits<double>::quiet_NaN();
+}
+
+template <>
+__device__ __forceinline__ __half nan_value<__half>() {
+  return __float2half(cuda::std::numeric_limits<float>::quiet_NaN());
+}
+
+template <>
+__device__ __forceinline__ __nv_bfloat16 nan_value<__nv_bfloat16>() {
+  return __float2bfloat16(cuda::std::numeric_limits<float>::quiet_NaN());
+}
+
+template <typename T, typename = void>
+struct InitValue {
+  __device__ __forceinline__ static T value() {
+    return Limits<T>::max();
+  }
+};
+
+template <typename T>
+struct InitValue<T, cuda::std::enable_if_t<std::is_floating_point_v<T>>> {
+  __device__ __forceinline__ static T value() {
+    return nan_value<T>();
+  }
+};
+
+template <typename T>
+__device__ __forceinline__ void thread_swap(T& a, T& b) {
+  T w = a;
+  a = b;
+  b = w;
+}
+
+template <typename T>
+struct LessThan {
+  __device__ __forceinline__ static T init() {
+    return InitValue<T>::value();
+  }
+
+  __device__ __forceinline__ bool operator()(T a, T b) const {
+    if constexpr (std::is_floating_point_v<T>) {
+      bool an = cuda::std::isnan(a);
+      bool bn = cuda::std::isnan(b);
+      if (an | bn) {
+        return (!an) & bn;
+      }
+    }
+    return a < b;
+  }
+};
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    int N_PER_THREAD,
+    typename CompareOp>
+struct ThreadSort {
+  __device__ __forceinline__ static void sort(
+      ValT (&vals)[N_PER_THREAD],
+      IdxT (&idxs)[N_PER_THREAD]) {
+    CompareOp op;
+#pragma unroll
+    for (int i = 0; i < N_PER_THREAD; ++i) {
+#pragma unroll
+      for (int j = i & 1; j < N_PER_THREAD - 1; j += 2) {
+        if (op(vals[j + 1], vals[j])) {
+          thread_swap(vals[j + 1], vals[j]);
+          if constexpr (ARG_SORT) {
+            thread_swap(idxs[j + 1], idxs[j]);
+          }
+        }
+      }
+    }
+  }
+};
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD,
+    typename CompareOp>
+struct BlockMergeSort {
+  using thread_sort_t =
+      ThreadSort<ValT, IdxT, ARG_SORT, N_PER_THREAD, CompareOp>;
+
+  __device__ __forceinline__ static int merge_partition(
+      const ValT* As,
+      const ValT* Bs,
+      int A_sz,
+      int B_sz,
+      int sort_md) {
+    CompareOp op;
+
+    int A_st = max(0, sort_md - B_sz);
+    int A_ed = min(sort_md, A_sz);
+
+    while (A_st < A_ed) {
+      int md = A_st + (A_ed - A_st) / 2;
+      auto a = As[md];
+      auto b = Bs[sort_md - 1 - md];
+
+      if (op(b, a)) {
+        A_ed = md;
+      } else {
+        A_st = md + 1;
+      }
+    }
+
+    return A_ed;
+  }
+
+  __device__ __forceinline__ static void merge_step(
+      const ValT* As,
+      const ValT* Bs,
+      const IdxT* As_idx,
+      const IdxT* Bs_idx,
+      int A_sz,
+      int B_sz,
+      ValT (&vals)[N_PER_THREAD],
+      IdxT (&idxs)[N_PER_THREAD]) {
+    CompareOp op;
+    int a_idx = 0;
+    int b_idx = 0;
+
+#pragma unroll
+    for (int i = 0; i < N_PER_THREAD; ++i) {
+      auto a = (a_idx < A_sz) ? As[a_idx] : ValT(CompareOp::init());
+      auto b = (b_idx < B_sz) ? Bs[b_idx] : ValT(CompareOp::init());
+      bool pred = (b_idx < B_sz) && (a_idx >= A_sz || op(b, a));
+
+      vals[i] = pred ? b : a;
+      if constexpr (ARG_SORT) {
+        if (pred) {
+          idxs[i] = Bs_idx[b_idx];
+        } else {
+          idxs[i] = (a_idx < A_sz) ? As_idx[a_idx] : IdxT(0);
+        }
+      }
+
+      b_idx += int(pred);
+      a_idx += int(!pred);
+    }
+  }
+
+  __device__ __forceinline__ static void
+  sort(ValT* tgp_vals, IdxT* tgp_idxs, int size_sorted_axis) {
+    int idx = threadIdx.x * N_PER_THREAD;
+
+    ValT thread_vals[N_PER_THREAD];
+    IdxT thread_idxs[N_PER_THREAD];
+#pragma unroll
+    for (int i = 0; i < N_PER_THREAD; ++i) {
+      thread_vals[i] = tgp_vals[idx + i];
+      if constexpr (ARG_SORT) {
+        thread_idxs[i] = tgp_idxs[idx + i];
+      }
+    }
+
+    if (idx < size_sorted_axis) {
+      thread_sort_t::sort(thread_vals, thread_idxs);
+    }
+
+    for (int merge_threads = 2; merge_threads <= BLOCK_THREADS;
+         merge_threads *= 2) {
+      __syncthreads();
+#pragma unroll
+      for (int i = 0; i < N_PER_THREAD; ++i) {
+        tgp_vals[idx + i] = thread_vals[i];
+        if constexpr (ARG_SORT) {
+          tgp_idxs[idx + i] = thread_idxs[i];
+        }
+      }
+      __syncthreads();
+
+      int merge_group = threadIdx.x / merge_threads;
+      int merge_lane = threadIdx.x % merge_threads;
+
+      int sort_sz = N_PER_THREAD * merge_threads;
+      int sort_st = N_PER_THREAD * merge_threads * merge_group;
+
+      int A_st = sort_st;
+      int A_ed = sort_st + sort_sz / 2;
+      int B_st = sort_st + sort_sz / 2;
+      int B_ed = sort_st + sort_sz;
+
+      const ValT* As = tgp_vals + A_st;
+      const ValT* Bs = tgp_vals + B_st;
+      int A_sz = A_ed - A_st;
+      int B_sz = B_ed - B_st;
+
+      int sort_md = N_PER_THREAD * merge_lane;
+      int partition = merge_partition(As, Bs, A_sz, B_sz, sort_md);
+
+      As += partition;
+      Bs += sort_md - partition;
+
+      A_sz -= partition;
+      B_sz -= sort_md - partition;
+
+      const IdxT* As_idx = ARG_SORT ? tgp_idxs + A_st + partition : nullptr;
+      const IdxT* Bs_idx =
+          ARG_SORT ? tgp_idxs + B_st + sort_md - partition : nullptr;
+
+      merge_step(As, Bs, As_idx, Bs_idx, A_sz, B_sz, thread_vals, thread_idxs);
+    }
+
+    __syncthreads();
+#pragma unroll
+    for (int i = 0; i < N_PER_THREAD; ++i) {
+      tgp_vals[idx + i] = thread_vals[i];
+      if constexpr (ARG_SORT) {
+        tgp_idxs[idx + i] = thread_idxs[i];
+      }
+    }
+  }
+};
+
+template <
+    typename T,
+    typename U,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD,
+    typename CompareOp = LessThan<T>>
+struct KernelMergeSort {
+  using ValT = T;
+  using IdxT = uint32_t;
+  using block_merge_sort_t = BlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD,
+      CompareOp>;
+
+  static constexpr int N_PER_BLOCK = BLOCK_THREADS * N_PER_THREAD;
+
+  __device__ __forceinline__ static void block_sort(
+      const T* inp,
+      U* out,
+      int size_sorted_axis,
+      int64_t in_stride_sorted_axis,
+      int64_t out_stride_sorted_axis,
+      int64_t in_stride_segment_axis,
+      int64_t out_stride_segment_axis,
+      ValT* tgp_vals,
+      IdxT* tgp_idxs) {
+    inp += blockIdx.y * in_stride_segment_axis;
+    out += blockIdx.y * out_stride_segment_axis;
+
+    for (int i = threadIdx.x; i < N_PER_BLOCK; i += BLOCK_THREADS) {
+      tgp_vals[i] = i < size_sorted_axis ? inp[i * in_stride_sorted_axis]
+                                         : ValT(CompareOp::init());
+      if constexpr (ARG_SORT) {
+        tgp_idxs[i] = i;
+      }
+    }
+
+    __syncthreads();
+    block_merge_sort_t::sort(tgp_vals, tgp_idxs, size_sorted_axis);
+    __syncthreads();
+
+    for (int i = threadIdx.x; i < size_sorted_axis; i += BLOCK_THREADS) {
+      if constexpr (ARG_SORT) {
+        out[i * out_stride_sorted_axis] = tgp_idxs[i];
+      } else {
+        out[i * out_stride_sorted_axis] = tgp_vals[i];
+      }
+    }
+  }
+};
+
+template <
+    typename T,
+    typename U,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD>
+__global__ void block_sort_kernel(
+    const T* inp,
+    U* out,
+    int size_sorted_axis,
+    int64_t in_stride_sorted_axis,
+    int64_t out_stride_sorted_axis,
+    int64_t in_stride_segment_axis,
+    int64_t out_stride_segment_axis) {
+  using sort_kernel =
+      KernelMergeSort<T, U, ARG_SORT, BLOCK_THREADS, N_PER_THREAD>;
+  using ValT = typename sort_kernel::ValT;
+  using IdxT = typename sort_kernel::IdxT;
+
+  if constexpr (ARG_SORT) {
+    __shared__ ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    __shared__ IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        in_stride_segment_axis,
+        out_stride_segment_axis,
+        tgp_vals,
+        tgp_idxs);
+  } else {
+    __shared__ ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        in_stride_segment_axis,
+        out_stride_segment_axis,
+        tgp_vals,
+        nullptr);
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD>
+__global__ void block_sort_nc_kernel(
+    const T* inp,
+    U* out,
+    int size_sorted_axis,
+    int64_t in_stride_sorted_axis,
+    int64_t out_stride_sorted_axis,
+    const __grid_constant__ Shape nc_shape,
+    const __grid_constant__ Strides in_nc_strides,
+    const __grid_constant__ Strides out_nc_strides,
+    int nc_dim) {
+  using sort_kernel =
+      KernelMergeSort<T, U, ARG_SORT, BLOCK_THREADS, N_PER_THREAD>;
+  using ValT = typename sort_kernel::ValT;
+  using IdxT = typename sort_kernel::IdxT;
+
+  int64_t in_block_idx = elem_to_loc(
+      int64_t(blockIdx.y), nc_shape.data(), in_nc_strides.data(), nc_dim);
+  int64_t out_block_idx = elem_to_loc(
+      int64_t(blockIdx.y), nc_shape.data(), out_nc_strides.data(), nc_dim);
+
+  inp += in_block_idx;
+  out += out_block_idx;
+
+  if constexpr (ARG_SORT) {
+    __shared__ ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    __shared__ IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        0,
+        0,
+        tgp_vals,
+        tgp_idxs);
+  } else {
+    __shared__ ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        0,
+        0,
+        tgp_vals,
+        nullptr);
+  }
+}
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD,
+    typename CompareOp = LessThan<ValT>>
+struct KernelMultiBlockMergeSort {
+  using block_merge_sort_t = BlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD,
+      CompareOp>;
+
+  static constexpr int N_PER_BLOCK = BLOCK_THREADS * N_PER_THREAD;
+
+  __device__ __forceinline__ static void block_sort(
+      const ValT* inp,
+      ValT* out_vals,
+      IdxT* out_idxs,
+      int size_sorted_axis,
+      int64_t stride_sorted_axis,
+      ValT* tgp_vals,
+      IdxT* tgp_idxs) {
+    int base_idx = blockIdx.x * N_PER_BLOCK;
+
+    for (int i = threadIdx.x; i < N_PER_BLOCK; i += BLOCK_THREADS) {
+      int idx = base_idx + i;
+      tgp_vals[i] = idx < size_sorted_axis ? inp[idx * stride_sorted_axis]
+                                           : ValT(CompareOp::init());
+      tgp_idxs[i] = idx;
+    }
+
+    __syncthreads();
+    block_merge_sort_t::sort(tgp_vals, tgp_idxs, size_sorted_axis);
+    __syncthreads();
+
+    for (int i = threadIdx.x; i < N_PER_BLOCK; i += BLOCK_THREADS) {
+      int idx = base_idx + i;
+      if (idx < size_sorted_axis) {
+        out_vals[idx] = tgp_vals[i];
+        out_idxs[idx] = tgp_idxs[i];
+      }
+    }
+  }
+
+  __device__ __forceinline__ static int merge_partition(
+      const ValT* As,
+      const ValT* Bs,
+      int A_sz,
+      int B_sz,
+      int sort_md) {
+    CompareOp op;
+
+    int A_st = max(0, sort_md - B_sz);
+    int A_ed = min(sort_md, A_sz);
+
+    while (A_st < A_ed) {
+      int md = A_st + (A_ed - A_st) / 2;
+      auto a = As[md];
+      auto b = Bs[sort_md - 1 - md];
+
+      if (op(b, a)) {
+        A_ed = md;
+      } else {
+        A_st = md + 1;
+      }
+    }
+
+    return A_ed;
+  }
+};
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD>
+__global__ void mb_block_sort_kernel(
+    const ValT* inp,
+    ValT* out_vals,
+    IdxT* out_idxs,
+    int size_sorted_axis,
+    int64_t stride_sorted_axis,
+    const __grid_constant__ Shape nc_shape,
+    const __grid_constant__ Strides nc_strides,
+    int nc_dim) {
+  using sort_kernel = KernelMultiBlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD>;
+
+  int64_t block_idx = elem_to_loc(
+      int64_t(blockIdx.y), nc_shape.data(), nc_strides.data(), nc_dim);
+
+  inp += block_idx;
+  out_vals += blockIdx.y * size_sorted_axis;
+  out_idxs += blockIdx.y * size_sorted_axis;
+
+  __shared__ ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+  __shared__ IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+
+  sort_kernel::block_sort(
+      inp,
+      out_vals,
+      out_idxs,
+      size_sorted_axis,
+      stride_sorted_axis,
+      tgp_vals,
+      tgp_idxs);
+}
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD>
+__global__ void mb_block_partition_kernel(
+    IdxT* block_partitions,
+    const ValT* dev_vals,
+    const IdxT* dev_idxs,
+    int size_sorted_axis,
+    int merge_tiles,
+    int n_blocks) {
+  using sort_kernel = KernelMultiBlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD>;
+
+  (void)dev_idxs;
+
+  block_partitions += blockIdx.y * blockDim.x;
+  dev_vals += blockIdx.y * size_sorted_axis;
+  dev_idxs += blockIdx.y * size_sorted_axis;
+
+  for (int i = threadIdx.x; i <= n_blocks; i += blockDim.x) {
+    int merge_group = i / merge_tiles;
+    int merge_lane = i % merge_tiles;
+
+    int sort_sz = sort_kernel::N_PER_BLOCK * merge_tiles;
+    int sort_st = sort_kernel::N_PER_BLOCK * merge_tiles * merge_group;
+
+    int A_st = min(size_sorted_axis, sort_st);
+    int A_ed = min(size_sorted_axis, sort_st + sort_sz / 2);
+    int B_st = A_ed;
+    int B_ed = min(size_sorted_axis, B_st + sort_sz / 2);
+
+    int partition_at = min(B_ed - A_st, sort_kernel::N_PER_BLOCK * merge_lane);
+    int partition = sort_kernel::merge_partition(
+        dev_vals + A_st,
+        dev_vals + B_st,
+        A_ed - A_st,
+        B_ed - B_st,
+        partition_at);
+
+    block_partitions[i] = A_st + partition;
+  }
+}
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    int BLOCK_THREADS,
+    int N_PER_THREAD,
+    typename CompareOp = LessThan<ValT>>
+__global__ void mb_block_merge_kernel(
+    const IdxT* block_partitions,
+    const ValT* dev_vals_in,
+    const IdxT* dev_idxs_in,
+    ValT* dev_vals_out,
+    IdxT* dev_idxs_out,
+    int size_sorted_axis,
+    int merge_tiles,
+    int num_tiles) {
+  using sort_kernel = KernelMultiBlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD,
+      CompareOp>;
+
+  using block_sort_t = typename sort_kernel::block_merge_sort_t;
+
+  block_partitions += blockIdx.y * (num_tiles + 1);
+  dev_vals_in += blockIdx.y * size_sorted_axis;
+  dev_idxs_in += blockIdx.y * size_sorted_axis;
+  dev_vals_out += blockIdx.y * size_sorted_axis;
+  dev_idxs_out += blockIdx.y * size_sorted_axis;
+
+  int block_idx = blockIdx.x;
+  int merge_group = block_idx / merge_tiles;
+  int sort_st = sort_kernel::N_PER_BLOCK * merge_tiles * merge_group;
+  int sort_sz = sort_kernel::N_PER_BLOCK * merge_tiles;
+  int sort_md = sort_kernel::N_PER_BLOCK * block_idx - sort_st;
+
+  int A_st = block_partitions[block_idx + 0];
+  int A_ed = block_partitions[block_idx + 1];
+  int B_st = min(size_sorted_axis, 2 * sort_st + sort_sz / 2 + sort_md - A_st);
+  int B_ed = min(
+      size_sorted_axis,
+      2 * sort_st + sort_sz / 2 + sort_md + sort_kernel::N_PER_BLOCK - A_ed);
+
+  if ((block_idx % merge_tiles) == merge_tiles - 1) {
+    A_ed = min(size_sorted_axis, sort_st + sort_sz / 2);
+    B_ed = min(size_sorted_axis, sort_st + sort_sz);
+  }
+
+  int A_sz = A_ed - A_st;
+  int B_sz = B_ed - B_st;
+
+  ValT thread_vals[N_PER_THREAD];
+  IdxT thread_idxs[N_PER_THREAD];
+#pragma unroll
+  for (int i = 0; i < N_PER_THREAD; i++) {
+    int idx = BLOCK_THREADS * i + threadIdx.x;
+    if (idx < (A_sz + B_sz)) {
+      thread_vals[i] = (idx < A_sz) ? dev_vals_in[A_st + idx]
+                                    : dev_vals_in[B_st + idx - A_sz];
+      thread_idxs[i] = (idx < A_sz) ? dev_idxs_in[A_st + idx]
+                                    : dev_idxs_in[B_st + idx - A_sz];
+    } else {
+      thread_vals[i] = CompareOp::init();
+      thread_idxs[i] = 0;
+    }
+  }
+
+  __shared__ ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+  __shared__ IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+  __syncthreads();
+#pragma unroll
+  for (int i = 0; i < N_PER_THREAD; i++) {
+    int idx = BLOCK_THREADS * i + threadIdx.x;
+    tgp_vals[idx] = thread_vals[i];
+    tgp_idxs[idx] = thread_idxs[i];
+  }
+  __syncthreads();
+
+  int sort_md_local = min(A_sz + B_sz, N_PER_THREAD * int(threadIdx.x));
+
+  int A_st_local = block_sort_t::merge_partition(
+      tgp_vals, tgp_vals + A_sz, A_sz, B_sz, sort_md_local);
+  int A_ed_local = A_sz;
+
+  int B_st_local = sort_md_local - A_st_local;
+  int B_ed_local = B_sz;
+
+  int A_sz_local = A_ed_local - A_st_local;
+  int B_sz_local = B_ed_local - B_st_local;
+
+  block_sort_t::merge_step(
+      tgp_vals + A_st_local,
+      tgp_vals + A_ed_local + B_st_local,
+      tgp_idxs + A_st_local,
+      tgp_idxs + A_ed_local + B_st_local,
+      A_sz_local,
+      B_sz_local,
+      thread_vals,
+      thread_idxs);
+
+  __syncthreads();
+#pragma unroll
+  for (int i = 0; i < N_PER_THREAD; ++i) {
+    int idx = threadIdx.x * N_PER_THREAD;
+    tgp_vals[idx + i] = thread_vals[i];
+    tgp_idxs[idx + i] = thread_idxs[i];
+  }
+
+  __syncthreads();
+  int base_idx = blockIdx.x * sort_kernel::N_PER_BLOCK;
+  for (int i = threadIdx.x; i < sort_kernel::N_PER_BLOCK; i += BLOCK_THREADS) {
+    int idx = base_idx + i;
+    if (idx < size_sorted_axis) {
+      dev_vals_out[idx] = tgp_vals[i];
+      dev_idxs_out[idx] = tgp_idxs[i];
+    }
+  }
+}
+
+} // namespace cu
+
+namespace {
+
+void single_block_sort(
+    const Stream& s,
+    const array& in,
+    array& out,
+    int axis,
+    int bn,
+    bool argsort) {
+  int n_rows = in.size() / in.shape(axis);
+
+  auto in_nc_str = in.strides();
+  in_nc_str.erase(in_nc_str.begin() + axis);
+
+  auto out_nc_str = out.strides();
+  out_nc_str.erase(out_nc_str.begin() + axis);
+
+  auto nc_shape = in.shape();
+  nc_shape.erase(nc_shape.begin() + axis);
+
+  int nc_dim = nc_shape.size();
+
+  int size_sorted_axis = in.shape(axis);
+  int64_t in_stride_sorted_axis = in.strides()[axis];
+  int64_t out_stride_sorted_axis = out.strides()[axis];
+
+  bool contiguous = in.flags().contiguous;
+  auto check_strides = [](const array& x, int64_t sort_stride) {
+    int64_t min_stride =
+        *std::min_element(x.strides().begin(), x.strides().end());
+    int64_t max_stride =
+        *std::max_element(x.strides().begin(), x.strides().end());
+    return sort_stride == min_stride || sort_stride == max_stride;
+  };
+  contiguous &= check_strides(in, in_stride_sorted_axis);
+  contiguous &= check_strides(out, out_stride_sorted_axis);
+
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    using CTYPE = MLX_GET_TYPE(type_tag);
+    if constexpr (!std::is_same_v<CTYPE, complex64_t>) {
+      using ValT = cuda_type_t<CTYPE>;
+      dispatch_block_dim(bn, [&](auto block_dim) {
+        constexpr int BLOCK_THREADS = block_dim();
+        if constexpr (BLOCK_THREADS < 1024) {
+          dim3 grid(1, n_rows, 1);
+          dim3 block(BLOCK_THREADS, 1, 1);
+
+          dispatch_bool(argsort, [&](auto arg_tag) {
+            constexpr bool ARG_SORT = decltype(arg_tag)::value;
+            using OutT = std::conditional_t<ARG_SORT, uint32_t, ValT>;
+
+            if (contiguous) {
+              auto kernel = cu::block_sort_kernel<
+                  ValT,
+                  OutT,
+                  ARG_SORT,
+                  BLOCK_THREADS,
+                  N_PER_THREAD>;
+              int64_t in_stride_segment_axis = INT64_MAX;
+              int64_t out_stride_segment_axis = INT64_MAX;
+              for (int i = 0; i < nc_shape.size(); i++) {
+                if (nc_shape[i] == 1) {
+                  continue;
+                }
+                if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) {
+                  throw std::runtime_error(
+                      "[Sort::eval_gpu] Stride too large.");
+                }
+                in_stride_segment_axis =
+                    std::min(in_stride_segment_axis, in_nc_str[i]);
+                out_stride_segment_axis =
+                    std::min(out_stride_segment_axis, out_nc_str[i]);
+              }
+              encoder.add_kernel_node(
+                  kernel,
+                  grid,
+                  block,
+                  0,
+                  gpu_ptr<ValT>(in),
+                  gpu_ptr<OutT>(out),
+                  size_sorted_axis,
+                  in_stride_sorted_axis,
+                  out_stride_sorted_axis,
+                  in_stride_segment_axis,
+                  out_stride_segment_axis);
+            } else {
+              auto kernel = cu::block_sort_nc_kernel<
+                  ValT,
+                  OutT,
+                  ARG_SORT,
+                  BLOCK_THREADS,
+                  N_PER_THREAD>;
+              auto nc_shape_param = const_param(nc_shape);
+              auto in_nc_strides_param = const_param(in_nc_str);
+              auto out_nc_strides_param = const_param(out_nc_str);
+              encoder.add_kernel_node(
+                  kernel,
+                  grid,
+                  block,
+                  0,
+                  gpu_ptr<ValT>(in),
+                  gpu_ptr<OutT>(out),
+                  size_sorted_axis,
+                  in_stride_sorted_axis,
+                  out_stride_sorted_axis,
+                  nc_shape_param,
+                  in_nc_strides_param,
+                  out_nc_strides_param,
+                  nc_dim);
+            }
+          });
+        }
+      });
+    } else {
+      throw std::runtime_error(
+          "CUDA backend does not support sorting complex numbers");
+    }
+  });
+}
+
+void multi_block_sort(
+    const Stream& s,
+    const array& in,
+    array& out,
+    int axis,
+    int n_blocks,
+    bool argsort) {
+  int n_rows = in.size() / in.shape(axis);
+
+  auto nc_str = in.strides();
+  nc_str.erase(nc_str.begin() + axis);
+
+  auto nc_shape = in.shape();
+  nc_shape.erase(nc_shape.begin() + axis);
+
+  int nc_dim = nc_shape.size();
+
+  if (nc_dim == 0) {
+    nc_shape = {0};
+    nc_str = {1};
+  }
+
+  int size_sorted_axis = in.shape(axis);
+  int64_t stride_sorted_axis = in.strides()[axis];
+
+  array dev_vals_in({n_rows, size_sorted_axis}, in.dtype(), nullptr, {});
+  array dev_vals_out({n_rows, size_sorted_axis}, in.dtype(), nullptr, {});
+
+  array dev_idxs_in({n_rows, size_sorted_axis}, uint32, nullptr, {});
+  array dev_idxs_out({n_rows, size_sorted_axis}, uint32, nullptr, {});
+
+  array block_partitions({n_rows, n_blocks + 1}, uint32, nullptr, {});
+
+  auto& encoder = cu::get_command_encoder(s);
+
+  dev_vals_in.set_data(cu::malloc_async(dev_vals_in.nbytes(), encoder));
+  dev_vals_out.set_data(cu::malloc_async(dev_vals_out.nbytes(), encoder));
+  dev_idxs_in.set_data(cu::malloc_async(dev_idxs_in.nbytes(), encoder));
+  dev_idxs_out.set_data(cu::malloc_async(dev_idxs_out.nbytes(), encoder));
+  block_partitions.set_data(
+      cu::malloc_async(block_partitions.nbytes(), encoder));
+
+  encoder.add_temporary(block_partitions);
+
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    using CTYPE = MLX_GET_TYPE(type_tag);
+    if constexpr (!std::is_same_v<CTYPE, complex64_t>) {
+      using ValT = cuda_type_t<CTYPE>;
+      using IdxT = uint32_t;
+      constexpr int BLOCK_THREADS = sizeof(ValT) == 8 ? 256 : 512;
+      dim3 grid(n_blocks, n_rows, 1);
+      dim3 block(BLOCK_THREADS, 1, 1);
+
+      dispatch_bool(argsort, [&](auto arg_tag) {
+        constexpr bool ARG_SORT = decltype(arg_tag)::value;
+        auto nc_shape_param = const_param(nc_shape);
+        auto nc_strides_param = const_param(nc_str);
+
+        auto block_sort_kernel = cu::mb_block_sort_kernel<
+            ValT,
+            IdxT,
+            ARG_SORT,
+            BLOCK_THREADS,
+            N_PER_THREAD>;
+        encoder.set_input_array(in);
+        encoder.set_output_array(dev_vals_in);
+        encoder.set_output_array(dev_idxs_in);
+        encoder.add_kernel_node(
+            block_sort_kernel,
+            grid,
+            block,
+            0,
+            gpu_ptr<ValT>(in),
+            gpu_ptr<ValT>(dev_vals_in),
+            gpu_ptr<IdxT>(dev_idxs_in),
+            size_sorted_axis,
+            stride_sorted_axis,
+            nc_shape_param,
+            nc_strides_param,
+            nc_dim);
+
+        int n_thr_per_group = (n_blocks + 1) < 1024 ? (n_blocks + 1) : 1024;
+
+        for (int merge_tiles = 2; (merge_tiles / 2) < n_blocks;
+             merge_tiles *= 2) {
+          auto partition_kernel = cu::mb_block_partition_kernel<
+              ValT,
+              IdxT,
+              ARG_SORT,
+              BLOCK_THREADS,
+              N_PER_THREAD>;
+
+          encoder.set_input_array(dev_vals_in);
+          encoder.set_input_array(dev_idxs_in);
+          encoder.set_output_array(block_partitions);
+
+          encoder.add_kernel_node(
+              partition_kernel,
+              dim3(1, n_rows, 1),
+              dim3(n_thr_per_group, 1, 1),
+              0,
+              gpu_ptr<IdxT>(block_partitions),
+              gpu_ptr<ValT>(dev_vals_in),
+              gpu_ptr<IdxT>(dev_idxs_in),
+              size_sorted_axis,
+              merge_tiles,
+              n_blocks);
+
+          auto merge_kernel = cu::mb_block_merge_kernel<
+              ValT,
+              IdxT,
+              ARG_SORT,
+              BLOCK_THREADS,
+              N_PER_THREAD>;
+
+          encoder.set_input_array(dev_vals_in);
+          encoder.set_input_array(dev_idxs_in);
+          encoder.set_input_array(block_partitions);
+          encoder.set_output_array(dev_vals_out);
+          encoder.set_output_array(dev_idxs_out);
+
+          encoder.add_kernel_node(
+              merge_kernel,
+              dim3(n_blocks, n_rows, 1),
+              dim3(BLOCK_THREADS, 1, 1),
+              0,
+              gpu_ptr<IdxT>(block_partitions),
+              gpu_ptr<ValT>(dev_vals_in),
+              gpu_ptr<IdxT>(dev_idxs_in),
+              gpu_ptr<ValT>(dev_vals_out),
+              gpu_ptr<IdxT>(dev_idxs_out),
+              size_sorted_axis,
+              merge_tiles,
+              n_blocks);
+          std::swap(dev_vals_in, dev_vals_out);
+          std::swap(dev_idxs_in, dev_idxs_out);
+        }
+      });
+    } else {
+      throw std::runtime_error(
+          "CUDA backend does not support sorting complex numbers");
+    }
+  });
+
+  encoder.add_temporary(dev_vals_out);
+  encoder.add_temporary(dev_idxs_out);
+  encoder.add_temporary(argsort ? dev_vals_in : dev_idxs_in);
+  if (axis == in.ndim() - 1) {
+    // Copy buffer to out, no need for temporary
+    out.copy_shared_buffer(
+        argsort ? dev_idxs_in : dev_vals_in,
+        out.strides(),
+        out.flags(),
+        out.size());
+  } else {
+    encoder.add_temporary(argsort ? dev_idxs_in : dev_vals_in);
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    auto strides = out.strides();
+    for (int ax = axis + 1; ax < strides.size(); ax++) {
+      strides[ax] *= out.shape(axis);
+    }
+    strides[axis] = 1;
+    copy_gpu_inplace(
+        (argsort) ? dev_idxs_in : dev_vals_in,
+        out,
+        out.shape(),
+        strides,
+        out.strides(),
+        0,
+        0,
+        CopyType::General,
+        s);
+  }
+}
+
+void gpu_merge_sort(
+    const Stream& s,
+    const array& in,
+    array& out,
+    int axis_,
+    bool argsort) {
+  int axis = axis_ < 0 ? axis_ + in.ndim() : axis_;
+  int size_sorted_axis = in.shape(axis);
+
+  constexpr int tn = N_PER_THREAD;
+  int potential_bn = (size_sorted_axis + tn - 1) / tn;
+
+  int bn;
+  if (potential_bn > 256) {
+    bn = 512;
+  } else if (potential_bn > 128) {
+    bn = 256;
+  } else if (potential_bn > 64) {
+    bn = 128;
+  } else if (potential_bn > 32) {
+    bn = 64;
+  } else {
+    bn = 32;
+  }
+
+  if (bn == 512 && size_of(in.dtype()) > 4) {
+    bn = 256;
+  }
+
+  int n_per_block = bn * tn;
+  int n_blocks = (size_sorted_axis + n_per_block - 1) / n_per_block;
+
+  if (n_blocks > 1) {
+    return multi_block_sort(s, in, out, axis, n_blocks, argsort);
+  }
+  return single_block_sort(s, in, out, axis, bn, argsort);
+}
+
+void gpu_sort(
+    const Stream& s,
+    const array& in,
+    array& out,
+    int axis,
+    bool argsort) {
+  auto& encoder = cu::get_command_encoder(s);
+  gpu_merge_sort(s, in, out, axis, argsort);
+}
+
+} // namespace
+
+void ArgSort::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("ArgSort::eval_gpu");
+  assert(inputs.size() == 1);
+  gpu_sort(stream(), inputs[0], out, axis_, true);
+}
+
+void Sort::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Sort::eval_gpu");
+  assert(inputs.size() == 1);
+  gpu_sort(stream(), inputs[0], out, axis_, false);
+}
+
+void ArgPartition::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("ArgPartition::eval_gpu");
+  gpu_sort(stream(), inputs[0], out, axis_, true);
+}
+
+void Partition::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Partition::eval_gpu");
+  gpu_sort(stream(), inputs[0], out, axis_, false);
+}
+
+} // namespace mlx::core
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/cuda/steel/defines.cuh b/macos/third_party/mlx/mlx/backend/cuda/steel/defines.cuh
new file mode 100644
index 0000000..bf92042
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/steel/defines.cuh
@@ -0,0 +1,9 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#define MLX_UNROLL _Pragma("unroll")
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+#define MLX_CUDA_SM_80_ENABLED
+#endif
diff --git a/macos/third_party/mlx/mlx/backend/cuda/steel/gemm.cuh b/macos/third_party/mlx/mlx/backend/cuda/steel/gemm.cuh
new file mode 100644
index 0000000..99580d2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/steel/gemm.cuh
@@ -0,0 +1,101 @@
+
+#include "mlx/backend/cuda/steel/mma.cuh"
+#include "mlx/backend/cuda/steel/tiles.cuh"
+
+namespace mlx::core::cu {
+
+/**
+ * An example gemm written with the utils.
+ *
+ * Computes A @ B.T when A and B are all aligned with the block sizes.
+ */
+template <typename T, int BM, int BN, int BK>
+__global__ void ab_t_aligned(const T* a, const T* b, T* y, int N, int K) {
+  constexpr int WARPS_M = 2;
+  constexpr int WARPS_N = 2;
+  constexpr int NUM_WARPS = WARPS_M * WARPS_N;
+  constexpr int WARP_STEP_M = BM / WARPS_M;
+  constexpr int WARP_STEP_N = BN / WARPS_N;
+
+  // Precompute some offsets for each thread
+  const int warpid = threadIdx.x / 32;
+  const int laneid = threadIdx.x % 32;
+  const int wm = warpid / WARPS_N;
+  const int wn = warpid % WARPS_N;
+  const int offset_m = wm * WARP_STEP_M;
+  const int offset_n = wn * WARP_STEP_N;
+
+  // Allocate shared memory
+  extern __shared__ char shmem[];
+  SharedTile<T, BM, BK>(&as)[2] = *(SharedTile<T, BM, BK>(*)[2])(&shmem[0]);
+  SharedTile<T, BN, BK>(&bs)[2] =
+      *(SharedTile<T, BN, BK>(*)[2])(&shmem[sizeof(T) * 2 * BM * BK]);
+
+  // Allocate registers for the MMA
+  RegisterTile<float, BM / WARPS_M, BN / WARPS_N> C;
+  RegisterTile<T, BM / WARPS_M, 16> A;
+  RegisterTile<T, BN / WARPS_N, 16> B;
+
+  // Move the global pointers to the tile
+  a += blockIdx.y * BM * K;
+  b += blockIdx.x * BN * K;
+  y += blockIdx.y * BM * N + blockIdx.x * BN;
+
+  // Zero the accumulators
+  C.fill(0);
+
+  // Start the SM pipeline
+  load_async<NUM_WARPS>(as[0], as[0].base_addr(), a, K);
+  load_async<NUM_WARPS>(bs[0], bs[0].base_addr(), b, K);
+  cp_async_commit();
+
+  int tic = 0;
+  for (int k_block = BK; k_block < K; k_block += BK) {
+    load_async<NUM_WARPS>(as[tic ^ 1], as[tic ^ 1].base_addr(), a + k_block, K);
+    load_async<NUM_WARPS>(bs[tic ^ 1], bs[tic ^ 1].base_addr(), b + k_block, K);
+    cp_async_commit();
+    cp_async_wait<1>();
+    __syncthreads();
+
+    MLX_UNROLL
+    for (int k = 0; k < BK / 16; k++) {
+      A.load(
+          as[tic],
+          as[tic].base_addr(),
+          offset_m + laneid % 16,
+          k * 16 + laneid / 16 * 8);
+      B.load(
+          bs[tic],
+          bs[tic].base_addr(),
+          offset_n + laneid % 16,
+          k * 16 + laneid / 16 * 8);
+
+      mma_t(C, A, B);
+    }
+
+    tic ^= 1;
+  }
+
+  // Empty the pipeline
+  cp_async_wait_all();
+  __syncthreads();
+  MLX_UNROLL
+  for (int k = 0; k < BK / 16; k++) {
+    A.load(
+        as[tic],
+        as[tic].base_addr(),
+        offset_m + laneid % 16,
+        k * 16 + laneid / 16 * 8);
+    B.load(
+        bs[tic],
+        bs[tic].base_addr(),
+        offset_n + laneid % 16,
+        k * 16 + laneid / 16 * 8);
+
+    mma_t(C, A, B);
+  }
+
+  C.store_global(y, N, offset_m, offset_n);
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/steel/mma.cuh b/macos/third_party/mlx/mlx/backend/cuda/steel/mma.cuh
new file mode 100644
index 0000000..94e3149
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/steel/mma.cuh
@@ -0,0 +1,117 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/steel/defines.cuh"
+#include "mlx/backend/cuda/steel/tiles.cuh"
+
+namespace mlx::core::cu {
+
+/**
+ * Fallback mma.
+ *
+ * We should probably a) implement a fallback or complain about it to the
+ * compiler.
+ */
+template <typename U, typename T>
+__device__ inline void
+mma_t(Tile16x16<U>& C, Tile16x16<T>& A, Tile16x16<T>& B) {}
+
+/**
+ * Multiply the 16x16 bfloat16 tiles and accumulate the result in one 16x16
+ * float tile.
+ *
+ * We actually perform C += A @ B.T
+ */
+__device__ __forceinline__ void mma_t(
+    Tile16x16<float>& C,
+    Tile16x16<__nv_bfloat16>& A,
+    Tile16x16<__nv_bfloat16>& B) {
+#if defined(MLX_CUDA_SM_80_ENABLED)
+  asm volatile(
+      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+      "{%0, %1, %2, %3}, "
+      "{%4, %5, %6, %7}, "
+      "{%8, %9}, "
+      "{%10, %11, %12, %13};"
+
+      // D matrix
+      : "+f"(C.values[0].x),
+        "+f"(C.values[0].y),
+        "+f"(C.values[1].x),
+        "+f"(C.values[1].y)
+
+      // A matrix
+      : "r"(*(uint32_t*)(&A.values[0])),
+        "r"(*(uint32_t*)(&A.values[1])),
+        "r"(*(uint32_t*)(&A.values[2])),
+        "r"(*(uint32_t*)(&A.values[3])),
+
+        // B matrix
+        "r"(*(uint32_t*)(&B.values[0])),
+        "r"(*(uint32_t*)(&B.values[2])),
+
+        // C matrix
+        "f"(C.values[0].x),
+        "f"(C.values[0].y),
+        "f"(C.values[1].x),
+        "f"(C.values[1].y));
+  asm volatile(
+      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+      "{%0, %1, %2, %3}, "
+      "{%4, %5, %6, %7}, "
+      "{%8, %9}, "
+      "{%10, %11, %12, %13};"
+
+      // D matrix
+      : "+f"(C.values[2].x),
+        "+f"(C.values[2].y),
+        "+f"(C.values[3].x),
+        "+f"(C.values[3].y)
+
+      // A matrix
+      : "r"(*(uint32_t*)(&A.values[0])),
+        "r"(*(uint32_t*)(&A.values[1])),
+        "r"(*(uint32_t*)(&A.values[2])),
+        "r"(*(uint32_t*)(&A.values[3])),
+
+        // B matrix
+        "r"(*(uint32_t*)(&B.values[1])),
+        "r"(*(uint32_t*)(&B.values[3])),
+
+        // C matrix
+        "f"(C.values[2].x),
+        "f"(C.values[2].y),
+        "f"(C.values[3].x),
+        "f"(C.values[3].y));
+#endif
+}
+
+/**
+ * Multiply larger register tiles by delegating to mma_t.
+ */
+template <typename U, typename T, int M, int N, int K>
+__device__ __forceinline__ void mma_t(
+    RegisterTile<U, M, N>& C,
+    RegisterTile<T, M, K>& A,
+    RegisterTile<T, N, K>& B) {
+  constexpr int TILES_M = RegisterTile<T, M, K>::TILES_Y;
+  constexpr int TILES_K = RegisterTile<T, M, K>::TILES_X;
+  constexpr int TILES_N = RegisterTile<T, N, K>::TILES_Y;
+
+  MLX_UNROLL
+  for (int k = 0; k < TILES_K; k++) {
+    MLX_UNROLL
+    for (int m = 0; m < TILES_M; m++) {
+      MLX_UNROLL
+      for (int n = 0; n < TILES_N; n++) {
+        mma_t(
+            C.data[m * TILES_N + n],
+            A.data[m * TILES_K + k],
+            B.data[n * TILES_K + k]);
+      }
+    }
+  }
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/steel/tiles.cuh b/macos/third_party/mlx/mlx/backend/cuda/steel/tiles.cuh
new file mode 100644
index 0000000..4583802
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/steel/tiles.cuh
@@ -0,0 +1,450 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/steel/utils.cuh"
+#include "mlx/backend/cuda/vector_types.cuh"
+
+namespace mlx::core::cu {
+
+/**
+ * The basic building block for Ampere mmas. A 16x16 tile distributed across
+ * the warp.
+ *
+ * Each thread holds 8 values. They are distributed according to
+ * https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-fragment-mma-16816-float
+ *
+ * For use instructions see the individual methods eg load().
+ */
+template <typename T>
+struct Tile16x16 {
+  using T2 = Vector2_t<T>;
+
+  T2 values[4];
+
+  __device__ inline void fill(T v) {
+    T2 v2 = {v, v};
+    for (int i = 0; i < 4; i++) {
+      values[i] = v2;
+    }
+  }
+
+  /**
+   * Load a 16x16 tile from shared memory.
+   *
+   * The instruction is a bit weird in the sense that the address provided by
+   * each thread and the elements loaded are not the same.
+   *
+   * We load 4 8x8 tiles. The tile rows are stored contiguously in memory. As a
+   * result the warp provides 4*8 = 32 addresses one per row.
+   *
+   * Threads 0-7 provide the addresses for the first tile, 8-15 for the second
+   * and so on. For instance to load a non swizzled tile we would do
+   *
+   *    base_addr + (laneid % 16) * BK + (laneid / 2) * 8
+   *
+   * See
+   * https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-ldmatrix
+   */
+  __device__ __forceinline__ void load(uint32_t row_address) {
+    if constexpr (
+        std::is_same_v<T2, __nv_bfloat162> || std::is_same_v<T2, __half2>) {
+      asm volatile(
+          "ldmatrix.sync.aligned.m8n8.x4.shared::cta.b16 {%0, %1, %2, %3}, [%4];\n"
+          : "=r"(*(uint32_t*)&(values[0])),
+            "=r"(*(uint32_t*)&(values[1])),
+            "=r"(*(uint32_t*)&(values[2])),
+            "=r"(*(uint32_t*)&(values[3]))
+          : "r"(row_address));
+    }
+  }
+
+  /**
+   * Store the tile to the address pointed to by `x`.
+   *
+   * The provided pointer is a generic pointer but this is meant to be used to
+   * store to global memory. For storing to shared memory we should use
+   * `stmatrix`.
+   *
+   * This also showcases the format of the tile quite nicely. Each register is
+   * holding to adjacent values. The indices are
+   *
+   *    row + 0, col + 0
+   *    row + 8, col + 0
+   *    row + 0, col + 8
+   *    row + 8, col + 8
+   *
+   * Given that we are dealing with Vector2_t<U> the column offsets are 4
+   * instead of 8.
+   */
+  template <typename U>
+  __device__ inline void store_global(U* x, int N) {
+    using U2 = Vector2_t<U>;
+    U2* x2 = reinterpret_cast<U2*>(x);
+    const int laneid = threadIdx.x % 32;
+    const int row = laneid / 4;
+    const int col = laneid % 4;
+    if constexpr (std::is_same_v<U2, T2>) {
+      x2[(row + 0) * (N / 2) + col + 0] = values[0];
+      x2[(row + 0) * (N / 2) + col + 4] = values[2];
+      x2[(row + 8) * (N / 2) + col + 0] = values[1];
+      x2[(row + 8) * (N / 2) + col + 4] = values[3];
+    } else if constexpr (
+        std::is_same_v<T2, float2> && std::is_same_v<U, __nv_bfloat16>) {
+      x2[(row + 0) * (N / 2) + col + 0] =
+          __floats2bfloat162_rn(values[0].x, values[0].y);
+      x2[(row + 0) * (N / 2) + col + 4] =
+          __floats2bfloat162_rn(values[2].x, values[2].y);
+      x2[(row + 8) * (N / 2) + col + 0] =
+          __floats2bfloat162_rn(values[1].x, values[1].y);
+      x2[(row + 8) * (N / 2) + col + 4] =
+          __floats2bfloat162_rn(values[3].x, values[3].y);
+    }
+  }
+
+  template <typename U>
+  __device__ inline void store_global_safe(U* x, int N, int max_rows) {
+    const int laneid = threadIdx.x % 32;
+    const int row = laneid / 4;
+    const int col = laneid % 4;
+    if (row < max_rows) {
+      x[(row + 0) * N + 2 * col + 0] = static_cast<U>(values[0].x);
+      x[(row + 0) * N + 2 * col + 1] = static_cast<U>(values[0].y);
+      x[(row + 0) * N + 2 * col + 8] = static_cast<U>(values[2].x);
+      x[(row + 0) * N + 2 * col + 9] = static_cast<U>(values[2].y);
+    }
+    if (row + 8 < max_rows) {
+      x[(row + 8) * N + 2 * col + 0] = static_cast<U>(values[1].x);
+      x[(row + 8) * N + 2 * col + 1] = static_cast<U>(values[1].y);
+      x[(row + 8) * N + 2 * col + 8] = static_cast<U>(values[3].x);
+      x[(row + 8) * N + 2 * col + 9] = static_cast<U>(values[3].y);
+    }
+  }
+};
+
+/**
+ * A simple container of multiple Tile16x16.
+ *
+ * Provides utility functions for loading and manipulating collections of basic
+ * tiles.
+ */
+template <typename T, int ROWS_, int COLS_>
+struct RegisterTile {
+  static constexpr int ROWS = ROWS_;
+  static constexpr int COLS = COLS_;
+  static constexpr int TILES_X = COLS / 16;
+  static constexpr int TILES_Y = ROWS / 16;
+
+  Tile16x16<T> data[TILES_X * TILES_Y];
+
+  __device__ inline void fill(T v) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        data[i * TILES_X + j].fill(v);
+      }
+    }
+  }
+
+  template <typename Tile>
+  __device__ __forceinline__ void
+  load(Tile& tile, uint32_t base_address, int row, int col) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        data[i * TILES_X + j].load(
+            tile.loc(base_address, row + i * 16, col + j * 16));
+      }
+    }
+  }
+
+  template <typename Tile, typename F>
+  __device__ __forceinline__ void
+  load(Tile& tile, F f, uint32_t base_address, int row, int col) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        f(data[i * TILES_X + j],
+          tile,
+          base_address,
+          row + i * 16,
+          col + j * 16);
+      }
+    }
+  }
+
+  template <typename U>
+  __device__ inline void store_global(U* x, int N, int row, int col) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        data[i * TILES_X + j].store_global(
+            x + (row + i * 16) * N + col + j * 16, N);
+      }
+    }
+  }
+
+  template <typename U>
+  __device__ inline void
+  store_global_safe(U* x, int N, int row, int col, int max_rows) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        data[i * TILES_X + j].store_global_safe(
+            x + (row + i * 16) * N + col + j * 16, N, max_rows - row - i * 16);
+      }
+    }
+  }
+};
+
+/**
+ * A simple container of multiple Tile16x16.
+ *
+ * Provides utility functions for loading and manipulating collections of basic
+ * tiles.
+ */
+template <typename T, int ROWS_, int COLS_>
+struct RegisterTile {
+  static constexpr int ROWS = ROWS_;
+  static constexpr int COLS = COLS_;
+  static constexpr int TILES_X = COLS / 16;
+  static constexpr int TILES_Y = ROWS / 16;
+
+  Tile16x16<T> data[TILES_X * TILES_Y];
+
+  __device__ inline void fill(T v) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        data[i * TILES_X + j].fill(v);
+      }
+    }
+  }
+
+  template <typename Tile>
+  __device__ inline void
+  load(Tile& tile, uint32_t base_address, int row, int col) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        data[i * TILES_X + j].load(
+            tile.loc(base_address, row + i * 16, col + j * 16));
+      }
+    }
+  }
+
+  template <typename U>
+  __device__ inline void store_global(U* x, int N, int row, int col) {
+    MLX_UNROLL
+    for (int i = 0; i < TILES_Y; i++) {
+      MLX_UNROLL
+      for (int j = 0; j < TILES_X; j++) {
+        data[i * TILES_X + j].store_global(
+            x + (row + i * 16) * N + col + j * 16, N);
+      }
+    }
+  }
+};
+
+template <typename T, int ROWS_, int COLS_>
+struct SharedTile {
+  static constexpr int ROWS = ROWS_;
+  static constexpr int COLS = COLS_;
+  static constexpr int TILES_X = COLS / 16;
+  static constexpr int TILES_Y = ROWS / 16;
+  static constexpr int NUMEL = ROWS * COLS;
+
+  // Swizzle taken from ThunderKittens. Should be changed when we switch to
+  // cute Layouts.
+  //
+  // See inludes/types/shared/st.cuh
+  //
+  // I do feel that it is too math heavy and can be improved. Also the math is
+  // done every time although the addresses don't change from load to load. I
+  // guess we are expecting the compiler to figure that out.
+  static constexpr int swizzle_bytes =
+      (sizeof(T) == 2 ? (TILES_X % 4 == 0 ? 128 : (TILES_X % 2 == 0 ? 64 : 32))
+                      : (sizeof(T) == 4 ? (TILES_X % 2 == 0 ? 128 : 64) : 0));
+
+  T data[ROWS * COLS];
+
+  __device__ inline uint32_t base_addr() const {
+    return __cvta_generic_to_shared(&data[0]);
+  }
+
+  // Return a pointer to the element at (row, col) using the swizzle.
+  __device__ static inline T* ptr(T* ptr, int row, int col) {
+    if constexpr (swizzle_bytes > 0) {
+      static constexpr int swizzle_repeat = swizzle_bytes * 8;
+      static constexpr int subtile_cols = swizzle_bytes / sizeof(T);
+      const int outer_idx = col / subtile_cols;
+      const uint64_t addr =
+          (uint64_t)(&ptr
+                         [outer_idx * ROWS * subtile_cols + row * subtile_cols +
+                          col % subtile_cols]);
+      const int swizzle = ((addr % swizzle_repeat) >> 7) << 4;
+      return (T*)(addr ^ swizzle);
+    } else {
+      return ptr + row * COLS + col;
+    }
+  }
+
+  // Return the location of the element at (row, col) using the swizzle.
+  __device__ static inline uint32_t loc(uint32_t ptr, int row, int col) {
+    if constexpr (swizzle_bytes > 0) {
+      static constexpr int swizzle_repeat = swizzle_bytes * 8;
+      static constexpr int subtile_cols = swizzle_bytes / sizeof(T);
+      const int outer_idx = col / subtile_cols;
+      const uint32_t addr = ptr +
+          sizeof(T) *
+              (outer_idx * ROWS * subtile_cols + row * subtile_cols +
+               col % subtile_cols);
+      const int swizzle = ((addr % swizzle_repeat) >> 7) << 4;
+      return (addr ^ swizzle);
+    } else {
+      return ptr + sizeof(T) * (row * COLS + col);
+    }
+  }
+
+  // Convenience functions to edit elements going through the swizzle.
+  __device__ inline T& operator()(int row, int col) {
+    return *ptr(data, row, col);
+  }
+  __device__ inline void store(float4& v, int row, int col) {
+    *(reinterpret_cast<float4*>(ptr(data, row, col))) = v;
+  }
+  __device__ inline void store(float2& v, int row, int col) {
+    *(reinterpret_cast<float2*>(ptr(data, row, col))) = v;
+  }
+  __device__ inline void store(float& v, int row, int col) {
+    *(reinterpret_cast<float*>(ptr(data, row, col))) = v;
+  }
+  template <int N>
+  __device__ inline void store(T (&v)[N], int row, int col) {
+    if constexpr (sizeof(T) * N == 4) {
+      store(*(reinterpret_cast<float*>(&v[0])), row, col);
+    } else if constexpr (sizeof(T) * N == 8) {
+      store(*(reinterpret_cast<float2*>(&v[0])), row, col);
+    } else if constexpr (sizeof(T) * N == 16) {
+      store(*(reinterpret_cast<float4*>(&v[0])), row, col);
+    } else {
+      MLX_UNROLL
+      for (int i = 0; i < N; i++) {
+        *ptr(data, row, col + i) = v[i];
+      }
+    }
+  }
+};
+
+/**
+ * Load the tile from global memory by loading 16 bytes at a time and storing
+ * them immediately.
+ *
+ * Can also be used as a fallback for architectures before sm_80.
+ */
+template <int NUM_WARPS, typename T, typename Tile>
+__device__ inline void load(Tile& tile, const T* x, int N) {
+  constexpr int NUM_THREADS = NUM_WARPS * 32;
+  constexpr int ELEMENTS_PER_LOAD = sizeof(float4) / sizeof(T);
+  constexpr int NUM_LOADS = Tile::NUMEL / ELEMENTS_PER_LOAD;
+  constexpr int NUM_LOADS_PER_THREAD = NUM_LOADS / NUM_THREADS;
+  constexpr int NUM_LOADS_PER_ROW = Tile::COLS / ELEMENTS_PER_LOAD;
+  constexpr int STEP_ROWS = NUM_THREADS / NUM_LOADS_PER_ROW;
+
+  const int row = threadIdx.x / NUM_LOADS_PER_ROW;
+  const int col = threadIdx.x % NUM_LOADS_PER_ROW;
+
+  x += row * N + col * ELEMENTS_PER_LOAD;
+
+  MLX_UNROLL
+  for (int i = 0; i < NUM_LOADS_PER_THREAD; i++) {
+    float4 tmp;
+    tmp = *(reinterpret_cast<const float4*>(&x[i * STEP_ROWS * N]));
+    tile.store(tmp, row + i * STEP_ROWS, col * ELEMENTS_PER_LOAD);
+  }
+}
+
+/**
+ * The asynchronous equivalent of load.
+ *
+ * Loads the tile from global memory by submitting a bunch of async copy
+ * instructions. The copy won't start until commit is called and we don't have
+ * a guarantee it will finish until wait is called.
+ *
+ * It should be used as follows
+ *
+ *    load(...)
+ *    load(...)
+ *    cp_async_commit()
+ *    do_other_stuff()
+ *    cp_async_wait_all()
+ *    do_stuff_with_shmem()
+ */
+template <int NUM_WARPS, typename T, typename Tile>
+__device__ inline void
+load_async(Tile& tile, uint32_t base_address, const T* x, int N) {
+  constexpr int NUM_THREADS = NUM_WARPS * 32;
+  constexpr int ELEMENTS_PER_LOAD = sizeof(float4) / sizeof(T);
+  constexpr int NUM_LOADS = Tile::NUMEL / ELEMENTS_PER_LOAD;
+  constexpr int NUM_LOADS_PER_THREAD = NUM_LOADS / NUM_THREADS;
+  constexpr int NUM_LOADS_PER_ROW = Tile::COLS / ELEMENTS_PER_LOAD;
+  constexpr int STEP_ROWS = NUM_THREADS / NUM_LOADS_PER_ROW;
+
+  const int row = threadIdx.x / NUM_LOADS_PER_ROW;
+  const int col = threadIdx.x % NUM_LOADS_PER_ROW;
+
+  x += row * N + col * ELEMENTS_PER_LOAD;
+
+  MLX_UNROLL
+  for (int i = 0; i < NUM_LOADS_PER_THREAD; i++) {
+    cp_async<16>(
+        tile.loc(base_address, row + i * STEP_ROWS, col * ELEMENTS_PER_LOAD),
+        x + i * STEP_ROWS * N);
+  }
+}
+
+/**
+ * Same as load_async but checks if we can load the row.
+ *
+ * NOTE: It should be changed to use a predicated cp async instead.
+ */
+template <int NUM_WARPS, typename T, typename Tile>
+__device__ inline void load_async_safe(
+    Tile& tile,
+    uint32_t base_address,
+    const T* x,
+    int N,
+    int max_rows) {
+  constexpr int NUM_THREADS = NUM_WARPS * 32;
+  constexpr int ELEMENTS_PER_LOAD = sizeof(float4) / sizeof(T);
+  constexpr int NUM_LOADS = Tile::NUMEL / ELEMENTS_PER_LOAD;
+  constexpr int NUM_LOADS_PER_THREAD = NUM_LOADS / NUM_THREADS;
+  constexpr int NUM_LOADS_PER_ROW = Tile::COLS / ELEMENTS_PER_LOAD;
+  constexpr int STEP_ROWS = NUM_THREADS / NUM_LOADS_PER_ROW;
+
+  const int row = threadIdx.x / NUM_LOADS_PER_ROW;
+  const int col = threadIdx.x % NUM_LOADS_PER_ROW;
+
+  x += row * N + col * ELEMENTS_PER_LOAD;
+
+  MLX_UNROLL
+  for (int i = 0; i < NUM_LOADS_PER_THREAD; i++) {
+    if (row + i * STEP_ROWS < max_rows) {
+      cp_async<16>(
+          tile.loc(base_address, row + i * STEP_ROWS, col * ELEMENTS_PER_LOAD),
+          x + i * STEP_ROWS * N);
+    } else {
+      float4 tmp = {0, 0, 0, 0};
+      tile.store(tmp, row + i * STEP_ROWS, col * ELEMENTS_PER_LOAD);
+    }
+  }
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/steel/utils.cuh b/macos/third_party/mlx/mlx/backend/cuda/steel/utils.cuh
new file mode 100644
index 0000000..0957c09
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/steel/utils.cuh
@@ -0,0 +1,89 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/utils.cuh"
+#include "mlx/backend/cuda/steel/defines.cuh"
+
+namespace mlx::core::cu {
+
+/**
+ * Copy bytes from the global memory address pointed to by x to the smem
+ * address pointed to by row_address.
+ *
+ * A simple wrapper over the PTX.
+ */
+template <int N, typename T>
+__device__ inline void cp_async(uint32_t row_address, const T* x) {
+  static_assert(
+      N == 16 || N == 8 || N == 4,
+      "cp.async is only supported for N in {4, 8, 16}.");
+#if defined(MLX_CUDA_SM_80_ENABLED)
+  if constexpr (N == 16) {
+    asm volatile(
+        "cp.async.ca.shared::cta.global [%0], [%1], 16;\n" ::"r"(row_address),
+        "l"(reinterpret_cast<const int4*>(x)));
+  } else if constexpr (N == 8) {
+    asm volatile(
+        "cp.async.ca.shared::cta.global [%0], [%1], 8;\n" ::"r"(row_address),
+        "l"(reinterpret_cast<const int2*>(x)));
+  } else if constexpr (N == 4) {
+    asm volatile(
+        "cp.async.ca.shared::cta.global [%0], [%1], 4;\n" ::"r"(row_address),
+        "l"(reinterpret_cast<const int*>(x)));
+  }
+#endif
+}
+
+/**
+ * Submit all the previous async copies to be executed.
+ */
+__device__ inline void cp_async_commit() {
+#if defined(MLX_CUDA_SM_80_ENABLED)
+  asm volatile("cp.async.commit_group;\n" ::);
+#endif
+}
+
+/**
+ * Wait for all but N of the async copies to finish.
+ */
+template <int N>
+__device__ inline void cp_async_wait() {
+#if defined(MLX_CUDA_SM_80_ENABLED)
+  if constexpr (N == 0) {
+    asm volatile("cp.async.wait_all;\n" ::);
+  } else {
+    asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
+  }
+#endif
+}
+
+/**
+ * Wait for all the async copies to finish.
+ */
+__device__ inline void cp_async_wait_all() {
+  cp_async_wait<0>();
+}
+
+/**
+ * Extract ``bits`` bits from the 32 bit value.
+ *
+ * Single instruction shift and mask.
+ */
+template <int bits>
+__device__ inline uint32_t extract_bits(uint32_t value, int start_bit) {
+  static_assert(
+      bits == 2 || bits == 4 || bits == 8,
+      "extract_bits only supports 2, 4, 8 for now.");
+  uint32_t result;
+  if constexpr (bits == 2) {
+    asm("bfe.u32 %0, %1, %2, 2;" : "=r"(result) : "r"(value), "r"(start_bit));
+  } else if constexpr (bits == 4) {
+    asm("bfe.u32 %0, %1, %2, 4;" : "=r"(result) : "r"(value), "r"(start_bit));
+  } else if constexpr (bits == 8) {
+    asm("bfe.u32 %0, %1, %2, 8;" : "=r"(result) : "r"(value), "r"(start_bit));
+  }
+  return result;
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/ternary.cu b/macos/third_party/mlx/mlx/backend/cuda/ternary.cu
new file mode 100644
index 0000000..e13ef64
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/ternary.cu
@@ -0,0 +1,271 @@
+// Copyright © 2025 Apple Inc.
+#include "mlx/backend/common/ternary.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/ternary_ops.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename Op, typename T, typename IdxT, int N_READS>
+__global__ void
+ternary_v(const bool* a, const T* b, const T* c, T* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      out[i] = Op{}(a[i], b[i], c[i]);
+    }
+  } else {
+    auto a_vec = load_vector<N_READS>(a, index);
+    auto b_vec = load_vector<N_READS>(b, index);
+    auto c_vec = load_vector<N_READS>(c, index);
+
+    AlignedVector<T, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = Op{}(a_vec[i], b_vec[i], c_vec[i]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+template <typename Op, typename T, typename IdxT, int NDIM, int N_READS>
+__global__ void ternary_g_nd(
+    const bool* a,
+    const T* b,
+    const T* c,
+    T* out,
+    IdxT size_rest,
+    const __grid_constant__ cuda::std::array<int32_t, NDIM> shape,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> a_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> b_strides,
+    const __grid_constant__ cuda::std::array<int64_t, NDIM> c_strides) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[NDIM - 1];
+  auto a_stride_x = a_strides[NDIM - 1];
+  auto b_stride_x = b_strides[NDIM - 1];
+  auto c_stride_x = c_strides[NDIM - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [a_idx, b_idx, c_idx] = elem_to_loc_nd<NDIM>(
+      index_rest * shape_x,
+      shape.data(),
+      a_strides.data(),
+      b_strides.data(),
+      c_strides.data());
+  auto a_vec =
+      load_vector<N_READS>(a + a_idx, index_x, shape_x, a_stride_x, false);
+  auto b_vec =
+      load_vector<N_READS>(b + b_idx, index_x, shape_x, b_stride_x, T(0));
+  auto c_vec =
+      load_vector<N_READS>(c + c_idx, index_x, shape_x, c_stride_x, T(0));
+
+  AlignedVector<T, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = Op{}(a_vec[i], b_vec[i], c_vec[i]);
+  }
+  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
+}
+
+template <typename Op, typename T, typename IdxT, int N_READS>
+__global__ void ternary_g(
+    const bool* a,
+    const T* b,
+    const T* c,
+    T* out,
+    IdxT size_rest,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides a_strides,
+    const __grid_constant__ Strides b_strides,
+    const __grid_constant__ Strides c_strides,
+    int ndim) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[ndim - 1];
+  auto a_stride_x = a_strides[ndim - 1];
+  auto b_stride_x = b_strides[ndim - 1];
+  auto c_stride_x = c_strides[ndim - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto [a_idx, b_idx, c_idx] = elem_to_loc(
+      index_rest * shape_x,
+      shape.data(),
+      a_strides.data(),
+      b_strides.data(),
+      c_strides.data(),
+      ndim);
+  auto a_vec =
+      load_vector<N_READS>(a + a_idx, index_x, shape_x, a_stride_x, false);
+  auto b_vec =
+      load_vector<N_READS>(b + b_idx, index_x, shape_x, b_stride_x, T(0));
+  auto c_vec =
+      load_vector<N_READS>(c + c_idx, index_x, shape_x, c_stride_x, T(0));
+
+  AlignedVector<T, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = Op{}(a_vec[i], b_vec[i], c_vec[i]);
+  }
+  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
+}
+
+} // namespace cu
+
+template <typename Op>
+void ternary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const Stream& s) {
+  const auto& a = inputs[0];
+  const auto& b = inputs[1];
+  const auto& c = inputs[2];
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cu::get_command_encoder(s);
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(c);
+  encoder.set_output_array(out);
+  dispatch_all_types(out.dtype(), [&](auto type_tag) {
+    using DType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+
+    auto topt = get_ternary_op_type(a, b, c);
+    if (topt == TernaryOpType::VectorVectorVector ||
+        topt == TernaryOpType::ScalarScalarScalar) {
+      dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
+        using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
+        constexpr int N_READS = 16 / sizeof(DType);
+        auto [num_blocks, block_dims] = get_launch_args(
+            out.data_size(), out.shape(), out.strides(), large(), N_READS);
+        encoder.add_kernel_node(
+            cu::ternary_v<Op, DType, IdxT, N_READS>,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<bool>(a),
+            gpu_ptr<DType>(b),
+            gpu_ptr<DType>(c),
+            gpu_ptr<DType>(out),
+            out.data_size());
+      });
+    } else {
+      dispatch_bool(
+          a.data_size() > INT32_MAX || b.data_size() > INT32_MAX ||
+              c.data_size() > INT32_MAX || out.data_size() > INT32_MAX,
+          [&](auto large) {
+            using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+            Shape shape;
+            std::vector<Strides> strides;
+            std::tie(shape, strides) = collapse_contiguous_dims(a, b, c, out);
+            auto& a_strides = strides[0];
+            auto& b_strides = strides[1];
+            auto& c_strides = strides[2];
+            int ndim = shape.size();
+            int work_per_thread = 1;
+            auto dim0 = ndim > 0 ? shape.back() : 1;
+            auto rest = out.size() / dim0;
+            if (dim0 >= 4) {
+              work_per_thread = 4;
+            }
+            dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+            auto block_dims = get_block_dims(dim0, rest, 1);
+            uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x);
+            uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y);
+
+            if (ndim <= 3) {
+              dispatch_1_2_3(ndim, [&](auto dims_constant) {
+                auto kernel =
+                    cu::ternary_g_nd<Op, DType, IdxT, dims_constant(), 1>;
+                if (work_per_thread == 4) {
+                  kernel =
+                      cu::ternary_g_nd<Op, DType, IdxT, dims_constant(), 4>;
+                }
+                encoder.add_kernel_node(
+                    kernel,
+                    {num_blocks_x, num_blocks_y},
+                    block_dims,
+                    0,
+                    gpu_ptr<bool>(a),
+                    gpu_ptr<DType>(b),
+                    gpu_ptr<DType>(c),
+                    gpu_ptr<DType>(out),
+                    rest,
+                    const_param<dims_constant()>(shape),
+                    const_param<dims_constant()>(a_strides),
+                    const_param<dims_constant()>(b_strides),
+                    const_param<dims_constant()>(c_strides));
+              });
+            } else {
+              auto kernel = cu::ternary_g<Op, DType, IdxT, 1>;
+              if (work_per_thread == 4) {
+                kernel = cu::ternary_g<Op, DType, IdxT, 4>;
+              }
+              encoder.add_kernel_node(
+                  kernel,
+                  {num_blocks_x, num_blocks_y},
+                  block_dims,
+                  0,
+                  gpu_ptr<bool>(a),
+                  gpu_ptr<DType>(b),
+                  gpu_ptr<DType>(c),
+                  gpu_ptr<DType>(out),
+                  rest,
+                  const_param(shape),
+                  const_param(a_strides),
+                  const_param(b_strides),
+                  const_param(c_strides),
+                  ndim);
+            }
+          });
+    }
+  });
+}
+
+template <typename Op>
+void ternary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const Stream& s) {
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& c = inputs[2];
+  auto topt = get_ternary_op_type(a, b, c);
+  auto& encoder = cu::get_command_encoder(s);
+  set_ternary_op_output_data(
+      a, b, c, out, topt, [&](auto n) { return cu::malloc_async(n, encoder); });
+  ternary_op_gpu_inplace<Op>(inputs, out, s);
+}
+
+void Select::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Select::eval_gpu");
+  auto& s = out.primitive().stream();
+  ternary_op_gpu<cu::Select>(inputs, out, s);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/cuda/unary/CMakeLists.txt
new file mode 100644
index 0000000..532c564
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/CMakeLists.txt
@@ -0,0 +1,34 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/abs.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arccos.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arccosh.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arcsin.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arcsinh.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arctan.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/arctanh.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/bitwise_invert.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/ceil.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/conjugate.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/cos.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/cosh.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/erf.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/erf_inv.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/exp.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/expm1.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/floor.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/imag.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/log.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/log1p.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/logical_not.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/negative.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/real.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/round.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/sigmoid.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/sign.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/sin.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/sinh.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/sqrt.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/square.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tan.cu
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tanh.cu)
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/abs.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/abs.cu
new file mode 100644
index 0000000..90b197d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/abs.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Abs)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/arccos.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/arccos.cu
new file mode 100644
index 0000000..3884997
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/arccos.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(ArcCos)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/arccosh.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/arccosh.cu
new file mode 100644
index 0000000..0ef0738
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/arccosh.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(ArcCosh)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/arcsin.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/arcsin.cu
new file mode 100644
index 0000000..07956ee
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/arcsin.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(ArcSin)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/arcsinh.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/arcsinh.cu
new file mode 100644
index 0000000..a7ab63e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/arcsinh.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(ArcSinh)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/arctan.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/arctan.cu
new file mode 100644
index 0000000..78639af
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/arctan.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(ArcTan)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/arctanh.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/arctanh.cu
new file mode 100644
index 0000000..488268c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/arctanh.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(ArcTanh)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/bitwise_invert.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/bitwise_invert.cu
new file mode 100644
index 0000000..77b88f3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/bitwise_invert.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(BitwiseInvert)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/ceil.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/ceil.cu
new file mode 100644
index 0000000..5ee300f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/ceil.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Ceil)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/conjugate.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/conjugate.cu
new file mode 100644
index 0000000..1d1d60e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/conjugate.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Conjugate)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/cos.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/cos.cu
new file mode 100644
index 0000000..cfceb86
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/cos.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Cos)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/cosh.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/cosh.cu
new file mode 100644
index 0000000..d5fcc70
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/cosh.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Cosh)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/erf.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/erf.cu
new file mode 100644
index 0000000..c785932
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/erf.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Erf)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/erf_inv.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/erf_inv.cu
new file mode 100644
index 0000000..16bbaba
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/erf_inv.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(ErfInv)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/exp.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/exp.cu
new file mode 100644
index 0000000..5a56669
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/exp.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Exp)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/expm1.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/expm1.cu
new file mode 100644
index 0000000..15e6ce4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/expm1.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Expm1)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/floor.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/floor.cu
new file mode 100644
index 0000000..a8c7ab0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/floor.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Floor)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/imag.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/imag.cu
new file mode 100644
index 0000000..9e3c05c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/imag.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Imag)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/log.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/log.cu
new file mode 100644
index 0000000..1fd2aa6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/log.cu
@@ -0,0 +1,21 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+void Log::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Log::eval_gpu");
+  auto& s = out.primitive().stream();
+  switch (base_) {
+    case Base::e:
+      unary_op_gpu<cu::Log>(inputs, out, name(), s);
+      break;
+    case Base::two:
+      unary_op_gpu<cu::Log2>(inputs, out, name(), s);
+      break;
+    case Base::ten:
+      unary_op_gpu<cu::Log10>(inputs, out, name(), s);
+      break;
+  }
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/log1p.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/log1p.cu
new file mode 100644
index 0000000..5396c3d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/log1p.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Log1p)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/logical_not.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/logical_not.cu
new file mode 100644
index 0000000..7f39870
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/logical_not.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(LogicalNot)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/negative.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/negative.cu
new file mode 100644
index 0000000..9c7e576
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/negative.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Negative)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/real.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/real.cu
new file mode 100644
index 0000000..361ffd3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/real.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Real)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/round.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/round.cu
new file mode 100644
index 0000000..4e80fdb
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/round.cu
@@ -0,0 +1,18 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+void Round::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Round::eval_gpu");
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  auto& s = out.primitive().stream();
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_op_gpu<cu::Round>(inputs, out, name(), s);
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/sigmoid.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/sigmoid.cu
new file mode 100644
index 0000000..3d94372
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/sigmoid.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Sigmoid)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/sign.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/sign.cu
new file mode 100644
index 0000000..d586d82
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/sign.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Sign)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/sin.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/sin.cu
new file mode 100644
index 0000000..47a5adc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/sin.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Sin)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/sinh.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/sinh.cu
new file mode 100644
index 0000000..7a73b7f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/sinh.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Sinh)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/sqrt.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/sqrt.cu
new file mode 100644
index 0000000..21f5f08
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/sqrt.cu
@@ -0,0 +1,15 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+void Sqrt::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Sqrt::eval_gpu");
+  auto& s = out.primitive().stream();
+  if (recip_) {
+    unary_op_gpu<cu::Rsqrt>(inputs, out, "Rsqrt", s);
+  } else {
+    unary_op_gpu<cu::Sqrt>(inputs, out, "Sqrt", s);
+  }
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/square.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/square.cu
new file mode 100644
index 0000000..bbb5f51
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/square.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Square)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/tan.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/tan.cu
new file mode 100644
index 0000000..3039dcd
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/tan.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Tan)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/tanh.cu b/macos/third_party/mlx/mlx/backend/cuda/unary/tanh.cu
new file mode 100644
index 0000000..ae69a51
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/tanh.cu
@@ -0,0 +1,7 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/unary/unary.cuh"
+
+namespace mlx::core {
+UNARY_GPU(Tanh)
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/unary/unary.cuh b/macos/third_party/mlx/mlx/backend/cuda/unary/unary.cuh
new file mode 100644
index 0000000..cade3f8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/unary/unary.cuh
@@ -0,0 +1,224 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/unary.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/unary_ops.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void unary_v(const In* in, Out* out, IdxT size) {
+  IdxT index = cg::this_grid().thread_rank();
+
+  if ((index + 1) * N_READS > size) {
+    for (IdxT i = index * N_READS; i < size; ++i) {
+      out[i] = Op{}(in[i]);
+    }
+  } else {
+    auto in_vec = load_vector<N_READS>(in, index);
+
+    AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      out_vec[i] = Op{}(in_vec[i]);
+    }
+
+    store_vector<N_READS>(out, index, out_vec);
+  }
+}
+
+template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
+__global__ void unary_g(
+    const In* in,
+    Out* out,
+    IdxT size_rest,
+    const __grid_constant__ Shape shape,
+    const __grid_constant__ Strides strides,
+    int ndim) {
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+  IdxT index_rest =
+      grid.block_index().y * block.dim_threads().y + block.thread_index().y;
+  if (index_rest >= size_rest) {
+    return;
+  }
+
+  auto shape_x = shape[ndim - 1];
+  auto stride_x = strides[ndim - 1];
+  IdxT index_x =
+      grid.block_index().x * block.dim_threads().x + block.thread_index().x;
+  auto idx =
+      elem_to_loc(index_rest * shape_x, shape.data(), strides.data(), ndim);
+  auto in_vec =
+      load_vector<N_READS>(in + idx, index_x, shape_x, stride_x, In(0));
+  AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    out_vec[i] = Op{}(in_vec[i]);
+  }
+  store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
+}
+
+template <typename Op, typename In, typename Out>
+constexpr bool supports_unary_op() {
+  if (std::is_same_v<Op, Abs> || std::is_same_v<Op, Negative> ||
+      std::is_same_v<Op, Sign> || std::is_same_v<Op, Square>) {
+    return std::is_same_v<In, Out>;
+  }
+  if (std::is_same_v<Op, ArcCosh> || std::is_same_v<Op, ArcSinh> ||
+      std::is_same_v<Op, ArcTanh> || std::is_same_v<Op, Erf> ||
+      std::is_same_v<Op, ErfInv> || std::is_same_v<Op, Expm1> ||
+      std::is_same_v<Op, Sigmoid>) {
+    return std::is_same_v<In, Out> && is_floating_v<In>;
+  }
+  if (std::is_same_v<Op, BitwiseInvert>) {
+    return std::is_same_v<In, Out> && std::is_integral_v<In> &&
+        !std::is_same_v<In, bool>;
+  }
+  if (std::is_same_v<Op, Ceil> || std::is_same_v<Op, Floor>) {
+    return std::is_same_v<In, Out> && !mlx::core::is_complex_v<In>;
+  }
+  if (std::is_same_v<Op, Conjugate>) {
+    return std::is_same_v<In, Out> && mlx::core::is_complex_v<In>;
+  }
+  if (std::is_same_v<Op, ArcCos> || std::is_same_v<Op, ArcSin> ||
+      std::is_same_v<Op, ArcTan> || std::is_same_v<Op, Cos> ||
+      std::is_same_v<Op, Cosh> || std::is_same_v<Op, Exp> ||
+      std::is_same_v<Op, Log> || std::is_same_v<Op, Log2> ||
+      std::is_same_v<Op, Log10> || std::is_same_v<Op, Log1p> ||
+      std::is_same_v<Op, Round> || std::is_same_v<Op, Rsqrt> ||
+      std::is_same_v<Op, Sqrt> || std::is_same_v<Op, Sin> ||
+      std::is_same_v<Op, Sinh> || std::is_same_v<Op, Tan> ||
+      std::is_same_v<Op, Tanh>) {
+    return std::is_same_v<In, Out> && is_inexact_v<In>;
+  }
+  if (std::is_same_v<Op, Imag> || std::is_same_v<Op, Real>) {
+    return mlx::core::is_complex_v<In> && std::is_same_v<Out, float>;
+  }
+  if (std::is_same_v<Op, LogicalNot>) {
+    return std::is_same_v<In, Out> && std::is_same_v<In, bool>;
+  }
+  if (std::is_same_v<Op, ToFP8>) {
+    return std::is_same_v<Out, uint8_t> && is_floating_v<In>;
+  }
+  if (std::is_same_v<Op, FromFP8>) {
+    return std::is_same_v<In, uint8_t> && is_floating_v<Out>;
+  }
+  return false;
+}
+
+} // namespace cu
+
+template <typename Op>
+void unary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  auto& in = inputs[0];
+  if (in.size() == 0) {
+    return;
+  }
+  bool contig = in.flags().contiguous;
+  bool large;
+  if (!contig) {
+    large = in.data_size() > INT32_MAX || out.size() > INT32_MAX;
+  } else {
+    large = in.data_size() > UINT32_MAX;
+  }
+
+  auto& encoder = cu::get_command_encoder(s);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  dispatch_all_types(in.dtype(), [&](auto in_type_tag) {
+    dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
+      using CTYPE_IN = MLX_GET_TYPE(in_type_tag);
+      using CTYPE_OUT = MLX_GET_TYPE(out_type_tag);
+      if constexpr (cu::supports_unary_op<Op, CTYPE_IN, CTYPE_OUT>()) {
+        dispatch_bool(large, [&](auto large) {
+          using InType = cuda_type_t<CTYPE_IN>;
+          using OutType = cuda_type_t<CTYPE_OUT>;
+          if (contig) {
+            using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
+            constexpr int N_READS = 16 / sizeof(OutType);
+            auto [num_blocks, block_dims] = get_launch_args(
+                out.data_size(), out.shape(), out.strides(), large, N_READS);
+            encoder.add_kernel_node(
+                cu::unary_v<Op, InType, OutType, IdxT, N_READS>,
+                num_blocks,
+                block_dims,
+                0,
+                gpu_ptr<InType>(in),
+                gpu_ptr<OutType>(out),
+                out.data_size());
+          } else {
+            using IdxT = std::conditional_t<large(), int64_t, int32_t>;
+            auto [shape, strides] = collapse_contiguous_dims(in);
+            auto ndim = shape.size();
+            int work_per_thread = 1;
+            auto kernel = cu::unary_g<Op, InType, OutType, IdxT, 1>;
+            auto dim0 = ndim > 0 ? shape.back() : 1;
+            auto rest = out.size() / dim0;
+            if (dim0 >= 4) {
+              kernel = cu::unary_g<Op, InType, OutType, IdxT, 4>;
+              work_per_thread = 4;
+            }
+            dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+            auto block_dims = get_block_dims(dim0, rest, 1);
+            uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x);
+            uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y);
+            encoder.add_kernel_node(
+                kernel,
+                {num_blocks_x, num_blocks_y},
+                block_dims,
+                0,
+                gpu_ptr<InType>(in),
+                gpu_ptr<OutType>(out),
+                rest,
+                const_param(shape),
+                const_param(strides),
+                ndim);
+          }
+        });
+      } else {
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do unary op {} on input of {} with output of {}.",
+                op,
+                dtype_to_string(in.dtype()),
+                dtype_to_string(out.dtype())));
+      }
+    });
+  });
+}
+
+template <typename Op>
+void unary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  auto& encoder = cu::get_command_encoder(s);
+  set_unary_output_data(
+      inputs[0], out, [&](auto n) { return cu::malloc_async(n, encoder); });
+  unary_op_gpu_inplace<Op>(inputs, out, op, s);
+}
+
+#define UNARY_GPU(func)                                               \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    nvtx3::scoped_range r(#func "::eval_gpu");                        \
+    auto& s = out.primitive().stream();                               \
+    unary_op_gpu<cu::func>(inputs, out, name(), s);                   \
+  }
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/utils.cpp b/macos/third_party/mlx/mlx/backend/cuda/utils.cpp
new file mode 100644
index 0000000..08f0f3d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/utils.cpp
@@ -0,0 +1,116 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/dtype_utils.h"
+
+#include <fmt/format.h>
+#include <cuda/cmath>
+#include <vector>
+
+namespace mlx::core {
+
+void check_cublas_error(const char* name, cublasStatus_t err) {
+  if (err != CUBLAS_STATUS_SUCCESS) {
+    // TODO: Use cublasGetStatusString when it is widely available.
+    throw std::runtime_error(
+        fmt::format("{} failed with code: {}.", name, static_cast<int>(err)));
+  }
+}
+
+void check_cuda_error(const char* name, cudaError_t err) {
+  if (err != cudaSuccess) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}", name, cudaGetErrorString(err)));
+  }
+}
+
+void check_cuda_error(const char* name, CUresult err) {
+  if (err != CUDA_SUCCESS) {
+    const char* err_str = "Unknown error";
+    cuGetErrorString(err, &err_str);
+    throw std::runtime_error(fmt::format("{} failed: {}", name, err_str));
+  }
+}
+
+void check_cudnn_error(const char* name, cudnnStatus_t err) {
+  if (err != CUDNN_STATUS_SUCCESS) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}.", name, cudnnGetErrorString(err)));
+  }
+}
+
+const char* dtype_to_cuda_type(const Dtype& dtype) {
+  switch (dtype) {
+    case bool_:
+      return "bool";
+    case int8:
+      return "int8_t";
+    case int16:
+      return "int16_t";
+    case int32:
+      return "int32_t";
+    case int64:
+      return "int64_t";
+    case uint8:
+      return "uint8_t";
+    case uint16:
+      return "uint16_t";
+    case uint32:
+      return "uint32_t";
+    case uint64:
+      return "uint64_t";
+    case float16:
+      return "__half";
+    case bfloat16:
+      return "__nv_bfloat16";
+    case float32:
+      return "float";
+    case float64:
+      return "double";
+    case complex64:
+      return "mlx::core::cu::complex64_t";
+    default:
+      return "unknown";
+  }
+}
+
+CudaGraph::CudaGraph(cu::Device& device) {
+  device.make_current();
+  CHECK_CUDA_ERROR(cudaGraphCreate(&handle_, 0));
+}
+
+void CudaGraph::end_capture(cudaStream_t stream) {
+  CHECK_CUDA_ERROR(cudaStreamEndCapture(stream, &handle_));
+}
+
+void CudaGraphExec::instantiate(cudaGraph_t graph) {
+  assert(handle_ == nullptr);
+  CHECK_CUDA_ERROR(cudaGraphInstantiate(&handle_, graph, nullptr, nullptr, 0));
+}
+
+CudaStream::CudaStream(cu::Device& device) {
+  device.make_current();
+  CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&handle_, cudaStreamNonBlocking));
+}
+
+void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size) {
+  if (workspace_size == 0) {
+    return nullptr;
+  }
+
+  // Workspace allocation should not be captured.
+#ifndef NDEBUG
+  cudaStreamCaptureStatus status;
+  CHECK_CUDA_ERROR(cudaStreamIsCapturing(encoder.stream(), &status));
+  assert(status == cudaStreamCaptureStatusNone);
+#endif
+
+  // Ensure workspace is 256-byte aligned.
+  int nbytes = cuda::ceil_div(workspace_size, 256) * 256;
+  array workspace(cu::malloc_async(nbytes, encoder), {nbytes}, int8);
+  encoder.add_temporary(workspace);
+  return gpu_ptr<void>(workspace);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/utils.h b/macos/third_party/mlx/mlx/backend/cuda/utils.h
new file mode 100644
index 0000000..387e79a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/utils.h
@@ -0,0 +1,49 @@
+// Copyright © 2025 Apple Inc.
+
+// This file include utilities that are used by C++ code (i.e. .cpp files).
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/cuda_utils.h"
+
+namespace mlx::core {
+
+template <typename T>
+inline uint32_t max_occupancy_block_dim(T kernel) {
+  int _, block_dim;
+  if constexpr (std::is_same_v<T, CUfunction>) {
+    CHECK_CUDA_ERROR(
+        cuOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel, 0, 0, 0));
+  } else {
+    CHECK_CUDA_ERROR(
+        cudaOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel));
+  }
+  return block_dim;
+}
+
+template <typename T>
+inline T* gpu_ptr(array& arr) {
+  return reinterpret_cast<T*>(
+      static_cast<char*>(
+          static_cast<cu::CudaBuffer*>(arr.buffer().ptr())->data) +
+      arr.offset());
+}
+
+// For const array, keep constness in pointer unless it is untyped.
+template <typename T>
+inline std::conditional_t<std::is_same_v<T, void>, void*, const T*> gpu_ptr(
+    const array& arr) {
+  return gpu_ptr<T>(const_cast<array&>(arr));
+}
+
+struct Dtype;
+
+// Convert Dtype to CUDA C++ types.
+const char* dtype_to_cuda_type(const Dtype& dtype);
+
+// Allocate an empty array and add it as temporary.
+void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/cuda/vector_types.cuh b/macos/third_party/mlx/mlx/backend/cuda/vector_types.cuh
new file mode 100644
index 0000000..2a7c092
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/vector_types.cuh
@@ -0,0 +1,48 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+namespace mlx::core::cu {
+
+template <typename T>
+struct Vector2;
+
+template <>
+struct Vector2<double> {
+  using type = double2;
+};
+
+template <>
+struct Vector2<float> {
+  using type = float2;
+};
+
+template <>
+struct Vector2<__half> {
+  using type = __half2;
+};
+
+template <>
+struct Vector2<__nv_bfloat16> {
+  using type = __nv_bfloat162;
+};
+
+template <typename T>
+using Vector2_t = typename Vector2<T>::type;
+
+template <typename T>
+struct Vector4 {
+  T x, y, z, w;
+};
+
+template <typename T>
+using Vector4_t = Vector4<T>;
+
+using bf16x4 = Vector4_t<__nv_bfloat16>;
+using fp16x4 = Vector4_t<__half>;
+using fp32x4 = Vector4_t<float>;
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/worker.cpp b/macos/third_party/mlx/mlx/backend/cuda/worker.cpp
new file mode 100644
index 0000000..495204d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/worker.cpp
@@ -0,0 +1,79 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/worker.h"
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core::cu {
+
+Worker::Worker(Device& d)
+    : signal_stream_(d),
+      signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync),
+      worker_(&Worker::thread_fn, this) {}
+
+Worker::~Worker() {
+  {
+    std::lock_guard lock(mtx_);
+    stop_ = true;
+  }
+  cond_.notify_one();
+  worker_.join();
+}
+
+void Worker::add_task(std::function<void()> task) {
+  pending_tasks_.push_back(std::move(task));
+}
+
+void Worker::signal(void* data) {
+  auto w = static_cast<Worker*>(data);
+  {
+    std::lock_guard lock(w->mtx_);
+    w->signaled_batch_++;
+  }
+  w->cond_.notify_one();
+}
+
+void Worker::commit(cudaStream_t stream) {
+  // Move pending tasks into tasks
+  if (pending_tasks_.empty()) {
+    return;
+  }
+  {
+    std::lock_guard lock(mtx_);
+    // Move pending tasks into ready tasks
+    worker_tasks_[++committed_batch_] = std::move(pending_tasks_);
+  }
+  signal_event_.record(stream);
+  signal_event_.wait(signal_stream_);
+  CHECK_CUDA_ERROR(cudaLaunchHostFunc(signal_stream_, signal, this));
+}
+
+void Worker::thread_fn() {
+  while (!stop_) {
+    uint64_t current_batch = 0;
+    Tasks tasks;
+    {
+      std::unique_lock<std::mutex> lk(mtx_);
+      cond_.wait(lk, [this, &current_batch] {
+        return this->signaled_batch_ > current_batch || this->stop_;
+      });
+      current_batch = signaled_batch_;
+      auto end = worker_tasks_.upper_bound(current_batch);
+      for (auto it = worker_tasks_.begin(); it != end; ++it) {
+        if (tasks.empty()) {
+          tasks = std::move(it->second);
+        } else {
+          std::move(
+              it->second.begin(), it->second.end(), std::back_inserter(tasks));
+        }
+      }
+      worker_tasks_.erase(worker_tasks_.begin(), end);
+    }
+    // Make sure tasks are cleared before the next wait
+    for (int i = 0; i < tasks.size(); ++i) {
+      auto task = std::move(tasks[i]);
+      task();
+    }
+  }
+}
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/cuda/worker.h b/macos/third_party/mlx/mlx/backend/cuda/worker.h
new file mode 100644
index 0000000..8f05e7b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/cuda/worker.h
@@ -0,0 +1,55 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/event.h"
+
+#include <condition_variable>
+#include <functional>
+#include <map>
+#include <mutex>
+#include <thread>
+
+namespace mlx::core::cu {
+
+// Run tasks in worker thread, synchronized with cuda stream.
+class Worker {
+ public:
+  explicit Worker(Device& d);
+  ~Worker();
+
+  Worker(const Worker&) = delete;
+  Worker& operator=(const Worker&) = delete;
+
+  // Add a pending |task| that will run when consumed or commited.
+  void add_task(std::function<void()> task);
+
+  // Inform worker thread to run current batches after kernels in |stream|
+  // finish running.
+  void commit(cudaStream_t stream);
+
+ private:
+  static void signal(void*);
+
+  void thread_fn();
+  std::mutex mtx_;
+  std::condition_variable cond_;
+
+  uint64_t committed_batch_{0};
+  uint64_t signaled_batch_{0};
+
+  // Cuda stream and event for signaling kernel completion.
+  CudaStream signal_stream_;
+  CudaEvent signal_event_;
+
+  bool stop_{false};
+
+  // Tasks are put in |pending_tasks_| first, and then moved to
+  // |worker_tasks_| when end_batch() is called.
+  using Tasks = std::vector<std::function<void()>>;
+  Tasks pending_tasks_;
+  std::map<uint64_t, Tasks> worker_tasks_;
+  std::thread worker_;
+};
+
+} // namespace mlx::core::cu
diff --git a/macos/third_party/mlx/mlx/backend/gpu/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/gpu/CMakeLists.txt
new file mode 100644
index 0000000..0396ae0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/CMakeLists.txt
@@ -0,0 +1,5 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp)
diff --git a/macos/third_party/mlx/mlx/backend/gpu/copy.cpp b/macos/third_party/mlx/mlx/backend/gpu/copy.cpp
new file mode 100644
index 0000000..403d07d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/copy.cpp
@@ -0,0 +1,89 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/primitives.h"
+
+#include <cassert>
+
+namespace mlx::core {
+
+void copy_gpu(const array& in, array& out, CopyType ctype) {
+  copy_gpu(in, out, ctype, out.primitive().stream());
+}
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    CopyType ctype,
+    const Stream& s) {
+  assert(in.shape() == out.shape());
+  return copy_gpu_inplace(
+      in, out, in.shape(), in.strides(), out.strides(), 0, 0, ctype, s);
+}
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const Strides& i_strides,
+    int64_t i_offset,
+    CopyType ctype,
+    const Stream& s) {
+  assert(in.shape() == out.shape());
+  return copy_gpu_inplace(
+      in, out, in.shape(), i_strides, out.strides(), i_offset, 0, ctype, s);
+}
+
+array contiguous_copy_gpu(const array& arr, const Stream& s) {
+  array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
+  copy_gpu(arr, arr_copy, CopyType::General, s);
+  return arr_copy;
+}
+
+array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s) {
+  int ndim = x.ndim();
+  if (start_axis < 0) {
+    start_axis += ndim;
+  }
+  if (end_axis < 0) {
+    end_axis += ndim;
+  }
+  start_axis = std::max(0, start_axis);
+  end_axis = std::min(ndim - 1, end_axis);
+
+  return reshape_in_eval(x, Flatten::output_shape(x, start_axis, end_axis), s);
+}
+
+array reshape_in_eval(const array& x, Shape shape, Stream s) {
+  array out(std::move(shape), x.dtype(), nullptr, {});
+  reshape_gpu(x, out, s);
+  return out;
+}
+
+array swapaxes_in_eval(const array& x, int axis1, int axis2) {
+  int ndim = x.ndim();
+  if (axis1 < 0) {
+    axis1 += ndim;
+  }
+  if (axis2 < 0) {
+    axis2 += ndim;
+  }
+
+  auto shape = x.shape();
+  std::swap(shape[axis1], shape[axis2]);
+  auto strides = x.strides();
+  std::swap(strides[axis1], strides[axis2]);
+
+  auto [data_size, row_contiguous, col_contiguous] =
+      check_contiguity(shape, strides);
+  bool contiguous = data_size == x.data_size();
+
+  array out(std::move(shape), x.dtype(), nullptr, {});
+  out.copy_shared_buffer(
+      x,
+      std::move(strides),
+      {contiguous, row_contiguous, col_contiguous},
+      x.data_size());
+  return out;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/gpu/copy.h b/macos/third_party/mlx/mlx/backend/gpu/copy.h
new file mode 100644
index 0000000..6e6bc79
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/copy.h
@@ -0,0 +1,57 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/copy.h"
+#include "mlx/stream.h"
+
+#include <optional>
+
+namespace mlx::core {
+
+// Generic copy inplace
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const Shape& data_shape,
+    const Strides& i_strides,
+    const Strides& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype,
+    const Stream& s,
+    std::optional<array> dynamic_i_offset = std::nullopt,
+    std::optional<array> dynamic_o_offset = std::nullopt);
+
+void copy_gpu(const array& src, array& out, CopyType ctype, const Stream& s);
+void copy_gpu(const array& src, array& out, CopyType ctype);
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    CopyType ctype,
+    const Stream& s);
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const Strides& i_strides,
+    int64_t i_offset,
+    CopyType ctype,
+    const Stream& s);
+
+// Fill the output with the scalar val
+void fill_gpu(const array& val, array& out, const Stream& s);
+
+// Return a contiguous array with same shape that copies the data of |arr|.
+array contiguous_copy_gpu(const array& arr, const Stream& s);
+
+// Copy data from |in| and transpose to |out|'s shape.
+void reshape_gpu(const array& in, array& out, Stream s);
+
+// Like the normal ops but safe to call in eval_gpu.
+array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s);
+array reshape_in_eval(const array& x, Shape shape, Stream s);
+array swapaxes_in_eval(const array& x, int axis1, int axis2);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/gpu/device_info.h b/macos/third_party/mlx/mlx/backend/gpu/device_info.h
new file mode 100644
index 0000000..7adb7f0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/device_info.h
@@ -0,0 +1,36 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+
+namespace mlx::core::gpu {
+
+MLX_API bool is_available();
+
+/**
+ * Get the number of available GPU devices.
+ */
+MLX_API int device_count();
+
+/**
+ * Get information about a GPU device.
+ *
+ * Returns a map of device properties. Keys vary by backend:
+ *   - device_name (string): Device name
+ *   - architecture (string): Architecture identifier
+ *   - total_memory/memory_size (size_t): Total device memory
+ *   - free_memory (size_t): Available memory (CUDA only)
+ *   - uuid (string): Device UUID (CUDA only)
+ *   - pci_bus_id (string): PCI bus ID (CUDA only)
+ *   - compute_capability_major/minor (size_t): Compute capability (CUDA only)
+ */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(int device_index = 0);
+
+} // namespace mlx::core::gpu
diff --git a/macos/third_party/mlx/mlx/backend/gpu/eval.h b/macos/third_party/mlx/mlx/backend/gpu/eval.h
new file mode 100644
index 0000000..f646c2e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/eval.h
@@ -0,0 +1,18 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <future>
+#include <memory>
+
+#include "mlx/array.h"
+#include "mlx/stream.h"
+
+namespace mlx::core::gpu {
+
+void new_stream(Stream stream);
+void eval(array& arr);
+void finalize(Stream s);
+void synchronize(Stream s);
+
+} // namespace mlx::core::gpu
diff --git a/macos/third_party/mlx/mlx/backend/gpu/primitives.cpp b/macos/third_party/mlx/mlx/backend/gpu/primitives.cpp
new file mode 100644
index 0000000..0138928
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/primitives.cpp
@@ -0,0 +1,307 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/primitives.h"
+#include "mlx/backend/common/slicing.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/gpu/slicing.h"
+
+#if defined(MLX_USE_CUDA)
+#include <nvtx3/nvtx3.hpp>
+#endif
+
+#include <cassert>
+
+#if defined(MLX_USE_CUDA)
+#define MLX_PROFILER_RANGE(message) nvtx3::scoped_range r(message)
+#else
+#define MLX_PROFILER_RANGE(message)
+#endif
+
+namespace mlx::core {
+
+void AsStrided::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("AsStrided::eval_gpu");
+  eval(inputs, out);
+}
+
+void AsType::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("AsType::eval_gpu");
+  CopyType ctype =
+      inputs[0].flags().contiguous ? CopyType::Vector : CopyType::General;
+  copy_gpu(inputs[0], out, ctype);
+}
+
+void Broadcast::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Broadcast::eval_gpu");
+  eval(inputs, out);
+}
+
+void BroadcastAxes::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("BroadcastAxes::eval_gpu");
+  eval(inputs, out);
+}
+
+void Concatenate::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Concatenate::eval_gpu");
+  concatenate_gpu(inputs, out, axis_, stream());
+}
+
+void Contiguous::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Contiguous::eval_gpu");
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  constexpr size_t extra_bytes = 16384;
+  if (in.buffer_size() <= out.nbytes() + extra_bytes &&
+      (in.flags().row_contiguous ||
+       (allow_col_major_ && in.flags().col_contiguous))) {
+    out.copy_shared_buffer(in);
+  } else {
+    copy_gpu(in, out, CopyType::General);
+  }
+}
+
+void Copy::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Copy::eval_gpu");
+  eval(inputs, out);
+}
+
+void CustomTransforms::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  MLX_PROFILER_RANGE("CustomTransforms::eval_gpu");
+  eval(inputs, outputs);
+}
+
+void Depends::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  MLX_PROFILER_RANGE("Depends::eval_gpu");
+  eval(inputs, outputs);
+}
+
+void DynamicSlice::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("DynamicSlice::eval_gpu");
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& start = inputs[1];
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto s = stream();
+  auto in_offset = compute_dynamic_offset(start, in.strides(), axes_, s);
+  copy_gpu_inplace(
+      /* const array& src = */ in,
+      /* array& dst = */ out,
+      /* const Shape& data_shape = */ out.shape(),
+      /* const Strides& i_strides = */ in.strides(),
+      /* const Strides& o_strides = */ out.strides(),
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ 0,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      /* const Stream& s = */ s,
+      /* std::optional<array> dynamic_i_offset = */ std::move(in_offset),
+      /* std::optional<array> dynamic_o_offset = */ std::nullopt);
+}
+
+void DynamicSliceUpdate::eval_gpu(
+    const std::vector<array>& inputs,
+    array& out) {
+  MLX_PROFILER_RANGE("DynamicSliceUpdate::eval_gpu");
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+  auto& start_indices = inputs[2];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  // Copy or donate input to output
+  auto s = stream();
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_gpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, s);
+
+  auto out_offset =
+      compute_dynamic_offset(start_indices, out.strides(), axes_, s);
+  copy_gpu_inplace(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const Shape& data_shape = */ upd.shape(),
+      /* const Strides& i_strides = */ upd.strides(),
+      /* const Strides& o_strides = */ out.strides(),
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ 0,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      /* const Stream& s = */ s,
+      /* std::optional<array> dynamic_i_offset = */ std::nullopt,
+      /* std::optional<array> dynamic_o_offset = */ std::move(out_offset));
+}
+
+void ExpandDims::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("ExpandDims::eval_gpu");
+  eval(inputs, out);
+}
+
+void Full::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Full::eval_gpu");
+  auto in = inputs[0];
+  CopyType ctype;
+  if (in.data_size() == 1) {
+    ctype = CopyType::Scalar;
+  } else if (in.flags().contiguous) {
+    ctype = CopyType::Vector;
+  } else {
+    ctype = CopyType::General;
+  }
+  copy_gpu(in, out, ctype);
+}
+
+void Flatten::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Flatten::eval_gpu");
+  reshape_gpu(inputs[0], out, stream());
+}
+
+void NumberOfElements::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("NumberOfElements::eval_gpu");
+  eval(inputs, out);
+}
+
+void Pad::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Pad::eval_gpu");
+  // Inputs must be base input array and scalar val array
+  assert(inputs.size() == 2);
+  auto& in = inputs[0];
+  auto& val = inputs[1];
+
+  // Padding value must be a scalar
+  assert(val.size() == 1);
+
+  // Padding value, input and output must be of the same type
+  assert(val.dtype() == in.dtype() && in.dtype() == out.dtype());
+
+  pad_gpu(in, val, out, axes_, low_pad_size_, stream());
+}
+
+void Reshape::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Reshape::eval_gpu");
+  reshape_gpu(inputs[0], out, stream());
+}
+
+void Split::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  MLX_PROFILER_RANGE("Split::eval_gpu");
+  eval(inputs, outputs);
+}
+
+void Slice::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Slice::eval_gpu");
+  assert(inputs.size() == 1);
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  slice_gpu(in, out, start_indices_, strides_, stream());
+}
+
+void SliceUpdate::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_gpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+  auto [data_offset, out_strides] =
+      prepare_slice(out, start_indices_, strides_);
+
+  // Do copy
+  copy_gpu_inplace(
+      /* const array& src = */ upd,
+      /* array& dst = */ out,
+      /* const Shape& data_shape = */ upd.shape(),
+      /* const Strides& i_strides = */ upd.strides(),
+      /* const Strides& o_strides = */ out_strides,
+      /* int64_t i_offset = */ 0,
+      /* int64_t o_offset = */ data_offset,
+      /* CopyType ctype = */ CopyType::GeneralGeneral,
+      /* const Stream& s = */ stream());
+}
+
+void Squeeze::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Squeeze::eval_gpu");
+  eval(inputs, out);
+}
+
+void StopGradient::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("StopGradient::eval_gpu");
+  eval(inputs, out);
+}
+
+void Transpose::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Transpose::eval_gpu");
+  eval(inputs, out);
+}
+
+void Unflatten::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Unflatten::eval_gpu");
+  reshape_gpu(inputs[0], out, stream());
+}
+
+void View::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("View::eval_gpu");
+  auto& in = inputs[0];
+  auto ibytes = size_of(in.dtype());
+  auto obytes = size_of(out.dtype());
+  // Conditions for buffer copying (disjunction):
+  // - type size is the same
+  // - type size is smaller and the last axis is contiguous
+  // - the entire array is row contiguous
+  if (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) ||
+      in.flags().row_contiguous) {
+    auto strides = in.strides();
+    for (int i = 0; i < static_cast<int>(strides.size()) - 1; ++i) {
+      strides[i] *= ibytes;
+      strides[i] /= obytes;
+    }
+    out.copy_shared_buffer(
+        in, strides, in.flags(), in.data_size() * ibytes / obytes);
+  } else {
+    auto tmp = array(in.shape(), in.dtype(), nullptr, {});
+    tmp.set_data(allocator::malloc(tmp.nbytes()));
+    copy_gpu_inplace(in, tmp, CopyType::General, stream());
+
+    auto flags = out.flags();
+    flags.contiguous = true;
+    flags.row_contiguous = true;
+    auto max_dim = std::max_element(out.shape().begin(), out.shape().end());
+    flags.col_contiguous = out.size() <= 1 || out.size() == *max_dim;
+    out.copy_shared_buffer(tmp, out.strides(), flags, out.size());
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/gpu/slicing.cpp b/macos/third_party/mlx/mlx/backend/gpu/slicing.cpp
new file mode 100644
index 0000000..7f0fec2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/slicing.cpp
@@ -0,0 +1,44 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/slicing.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/gpu/slicing.h"
+
+namespace mlx::core {
+
+void slice_gpu(
+    const array& in,
+    array& out,
+    const Shape& start_indices,
+    const Shape& strides,
+    const Stream&) {
+  slice(in, out, start_indices, strides);
+}
+
+void pad_gpu(
+    const array& in,
+    const array& val,
+    array& out,
+    const std::vector<int>& axes,
+    const Shape& low_pad_size,
+    const Stream& s) {
+  // Fill output with val
+  fill_gpu(val, out, s);
+
+  // Find offset for start of input values
+  size_t data_offset = 0;
+  for (int i = 0; i < axes.size(); i++) {
+    auto ax = axes[i] < 0 ? out.ndim() + axes[i] : axes[i];
+    data_offset += out.strides()[ax] * low_pad_size[i];
+  }
+
+  // Extract slice from output where input will be pasted
+  array out_slice(in.shape(), out.dtype(), nullptr, {});
+  out_slice.copy_shared_buffer(
+      out, out.strides(), out.flags(), out_slice.size(), data_offset);
+
+  // Copy input values into the slice
+  copy_gpu_inplace(in, out_slice, CopyType::GeneralGeneral, s);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/gpu/slicing.h b/macos/third_party/mlx/mlx/backend/gpu/slicing.h
new file mode 100644
index 0000000..596f7af
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/gpu/slicing.h
@@ -0,0 +1,36 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+void slice_gpu(
+    const array& in,
+    array& out,
+    const Shape& start_indices,
+    const Shape& strides,
+    const Stream& s);
+
+void concatenate_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    int axis,
+    const Stream& s);
+
+void pad_gpu(
+    const array& in,
+    const array& val,
+    array& out,
+    const std::vector<int>& axes,
+    const Shape& low_pad_size,
+    const Stream& s);
+
+array compute_dynamic_offset(
+    const array& indices,
+    const Strides& strides,
+    const std::vector<int>& axes,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/metal/CMakeLists.txt
new file mode 100644
index 0000000..4074e7b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/CMakeLists.txt
@@ -0,0 +1,144 @@
+function(make_jit_source SRC_FILE)
+  # This function takes a metal header file, runs the C preprocessesor on it,
+  # and makes the processed contents available as a string in a C++ function
+  # mlx::core::metal::${SRC_NAME}()
+  #
+  # To use the function, declare it in jit/includes.h and include
+  # jit/includes.h.
+  #
+  # Additional arguments to this function are treated as dependencies in the
+  # Cmake build system.
+  get_filename_component(SRC_NAME ${SRC_FILE} NAME)
+  add_custom_command(
+    OUTPUT jit/${SRC_NAME}.cpp
+    COMMAND
+      bash ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
+      ${CMAKE_CURRENT_BINARY_DIR}/jit ${CMAKE_C_COMPILER} ${PROJECT_SOURCE_DIR}
+      ${SRC_FILE}
+    DEPENDS make_compiled_preamble.sh kernels/${SRC_FILE}.h ${ARGN})
+  add_custom_target(${SRC_NAME} DEPENDS jit/${SRC_NAME}.cpp)
+  add_dependencies(mlx ${SRC_NAME})
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/jit/${SRC_NAME}.cpp)
+endfunction(make_jit_source)
+
+make_jit_source(utils kernels/bf16.h kernels/bf16_math.h kernels/complex.h
+                kernels/defines.h kernels/logging.h)
+make_jit_source(unary_ops kernels/erf.h kernels/expm1f.h kernels/fp8.h)
+make_jit_source(binary_ops)
+make_jit_source(ternary_ops)
+make_jit_source(reduce_utils kernels/atomic.h kernels/reduction/ops.h)
+make_jit_source(indexing/scatter kernels/indexing/indexing.h)
+make_jit_source(indexing/masked_scatter)
+make_jit_source(indexing/gather kernels/indexing/indexing.h)
+make_jit_source(indexing/gather_front kernels/indexing/indexing.h)
+make_jit_source(indexing/gather_axis)
+make_jit_source(indexing/scatter_axis)
+make_jit_source(hadamard)
+
+if(MLX_METAL_JIT)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jit_kernels.cpp)
+  make_jit_source(arange)
+  make_jit_source(copy)
+  make_jit_source(unary)
+  make_jit_source(binary)
+  make_jit_source(binary_two)
+  make_jit_source(fft kernels/fft/radix.h kernels/fft/readwrite.h)
+  make_jit_source(logsumexp)
+  make_jit_source(ternary)
+  make_jit_source(softmax)
+  make_jit_source(scan)
+  make_jit_source(sort)
+  make_jit_source(
+    reduce kernels/reduction/reduce_all.h kernels/reduction/reduce_col.h
+    kernels/reduction/reduce_row.h kernels/reduction/reduce_init.h)
+  make_jit_source(
+    steel/gemm/gemm kernels/steel/utils.h kernels/steel/gemm/loader.h
+    kernels/steel/gemm/mma.h kernels/steel/gemm/params.h
+    kernels/steel/gemm/transforms.h)
+  make_jit_source(steel/gemm/kernels/steel_gemm_fused)
+  make_jit_source(steel/gemm/kernels/steel_gemm_masked kernels/steel/defines.h)
+  make_jit_source(steel/gemm/kernels/steel_gemm_gather)
+  make_jit_source(steel/gemm/kernels/steel_gemm_splitk)
+  make_jit_source(steel/gemm/kernels/steel_gemm_segmented)
+  make_jit_source(
+    steel/conv/conv
+    kernels/steel/utils.h
+    kernels/steel/defines.h
+    kernels/steel/gemm/mma.h
+    kernels/steel/gemm/transforms.h
+    kernels/steel/conv/params.h
+    kernels/steel/conv/loader.h
+    kernels/steel/conv/loaders/loader_channel_l.h
+    kernels/steel/conv/loaders/loader_channel_n.h)
+  make_jit_source(steel/conv/kernels/steel_conv)
+  make_jit_source(steel/conv/kernels/steel_conv_general kernels/steel/defines.h
+                  kernels/steel/conv/loaders/loader_general.h)
+
+  make_jit_source(quantized_utils)
+  make_jit_source(quantized kernels/quantized_utils.h)
+  make_jit_source(fp_quantized kernels/quantized_utils.h kernels/fp8.h
+                  kernels/fp4.h)
+  make_jit_source(gemv_masked)
+
+  make_jit_source(steel/attn/kernels/steel_attention)
+
+  make_jit_source(
+    steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h
+    kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h)
+  make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax)
+  make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax)
+  make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax)
+
+  make_jit_source(quantized_nax kernels/quantized_utils.h)
+  make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h
+                  kernels/fp4.h)
+
+  make_jit_source(steel/attn/kernels/steel_attention_nax)
+
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp)
+endif()
+
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/normalization.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/ternary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/unary.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/resident.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
+
+if(NOT MLX_METAL_PATH)
+  set(MLX_METAL_PATH ${CMAKE_CURRENT_BINARY_DIR}/kernels/)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels)
+
+target_compile_definitions(mlx
+                           PRIVATE METAL_PATH="${MLX_METAL_PATH}/mlx.metallib")
diff --git a/macos/third_party/mlx/mlx/backend/metal/allocator.cpp b/macos/third_party/mlx/mlx/backend/metal/allocator.cpp
new file mode 100644
index 0000000..06e3d9a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/allocator.cpp
@@ -0,0 +1,278 @@
+// Copyright © 2023-2024 Apple Inc.
+#include "mlx/backend/metal/allocator.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/backend/metal/resident.h"
+#include "mlx/memory.h"
+
+#include <mach/vm_page_size.h>
+#include <unistd.h>
+#include <cstdlib>
+
+namespace mlx::core {
+
+constexpr size_t resource_options =
+    MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeUntracked;
+
+namespace allocator {
+
+Allocator& allocator() {
+  return metal::allocator();
+}
+
+void* Buffer::raw_ptr() {
+  if (!ptr_) {
+    return nullptr;
+  }
+  return static_cast<MTL::Buffer*>(ptr_)->contents();
+}
+
+} // namespace allocator
+
+namespace metal {
+
+MetalAllocator::MetalAllocator()
+    : device_(device(mlx::core::Device::gpu).mtl_device()),
+      buffer_cache_(
+          vm_page_size,
+          [](MTL::Buffer* buf) { return buf->length(); },
+          [this](MTL::Buffer* buf) {
+            if (!buf->heap()) {
+              residency_set_.erase(buf);
+            }
+            buf->release();
+          }),
+      residency_set_(device_) {
+  auto pool = metal::new_scoped_memory_pool();
+  const auto& info = gpu::device_info(0);
+  auto memsize = std::get<size_t>(info.at("memory_size"));
+  auto max_rec_size =
+      std::get<size_t>(info.at("max_recommended_working_set_size"));
+  resource_limit_ = std::get<size_t>(info.at("resource_limit"));
+  block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize);
+  gc_limit_ = std::min(static_cast<size_t>(0.95 * max_rec_size), block_limit_);
+  max_pool_size_ = block_limit_;
+  device(mlx::core::Device::gpu)
+      .set_residency_set(residency_set_.mtl_residency_set());
+  bool is_vm = std::get<std::string>(info.at("device_name")) ==
+      "Apple Paravirtual device";
+  if (is_vm) {
+    return;
+  }
+  auto heap_desc = MTL::HeapDescriptor::alloc()->init();
+  heap_desc->setResourceOptions(resource_options);
+  heap_desc->setSize(heap_size_);
+  heap_ = device_->newHeap(heap_desc);
+  heap_desc->release();
+  residency_set_.insert(heap_);
+}
+
+MetalAllocator::~MetalAllocator() {
+  auto pool = metal::new_scoped_memory_pool();
+  if (heap_) {
+    heap_->release();
+  }
+  buffer_cache_.clear();
+}
+
+size_t MetalAllocator::set_cache_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
+  std::swap(limit, max_pool_size_);
+  return limit;
+};
+
+size_t MetalAllocator::set_memory_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
+  std::swap(limit, block_limit_);
+  gc_limit_ = std::min(
+      block_limit_,
+      static_cast<size_t>(0.95 * device_->recommendedMaxWorkingSetSize()));
+  return limit;
+};
+
+size_t MetalAllocator::get_memory_limit() {
+  return block_limit_;
+}
+
+size_t MetalAllocator::set_wired_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
+  std::swap(limit, wired_limit_);
+  residency_set_.resize(wired_limit_);
+  return limit;
+};
+
+Buffer MetalAllocator::malloc(size_t size) {
+  // Metal doesn't like empty buffers
+  if (size == 0) {
+    return Buffer{nullptr};
+  }
+
+  // More helpful message if maximum buffer length is exceeded
+  if (size > device_->maxBufferLength()) {
+    std::ostringstream msg;
+    msg << "[metal::malloc] Attempting to allocate " << size
+        << " bytes which is greater than"
+        << " the maximum allowed buffer size of " << device_->maxBufferLength()
+        << " bytes.";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Align up memory
+  if (size > vm_page_size) {
+    size = vm_page_size * ((size + vm_page_size - 1) / vm_page_size);
+  }
+
+  // Try the cache
+  std::unique_lock lk(mutex_);
+  MTL::Buffer* buf = buffer_cache_.reuse_from_cache(size);
+  if (!buf) {
+    size_t mem_required = get_active_memory() + get_cache_memory() + size;
+
+    auto pool = metal::new_scoped_memory_pool();
+
+    // If we have a lot of memory pressure try to reclaim memory from the cache
+    if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) {
+      num_resources_ -=
+          buffer_cache_.release_cached_buffers(mem_required - gc_limit_);
+    }
+
+    // Allocate new buffer if needed
+    if (num_resources_ >= resource_limit_) {
+      std::ostringstream msg;
+      msg << "[metal::malloc] Resource limit (" << resource_limit_
+          << ") exceeded.";
+      throw std::runtime_error(msg.str());
+    }
+    lk.unlock();
+    if (size < small_size_ && heap_) {
+      buf = heap_->newBuffer(size, resource_options);
+    }
+    if (!buf) {
+      buf = device_->newBuffer(size, resource_options);
+    }
+    if (!buf) {
+      std::ostringstream msg;
+      msg << "[malloc] Unable to allocate " << size << " bytes.";
+      throw std::runtime_error(msg.str());
+    }
+    lk.lock();
+    num_resources_++;
+    if (!buf->heap()) {
+      residency_set_.insert(buf);
+    }
+  }
+
+  active_memory_ += buf->length();
+  peak_memory_ = std::max(peak_memory_, active_memory_);
+
+  // Maintain the cache below the requested limit
+  if (get_cache_memory() > max_pool_size_) {
+    auto pool = metal::new_scoped_memory_pool();
+    num_resources_ -= buffer_cache_.release_cached_buffers(
+        get_cache_memory() - max_pool_size_);
+  }
+
+  return Buffer{static_cast<void*>(buf)};
+}
+
+void MetalAllocator::clear_cache() {
+  std::unique_lock lk(mutex_);
+  auto pool = metal::new_scoped_memory_pool();
+  num_resources_ -= buffer_cache_.clear();
+}
+
+void MetalAllocator::free(Buffer buffer) {
+  auto buf = static_cast<MTL::Buffer*>(buffer.ptr());
+  if (buf == nullptr) {
+    return;
+  }
+  std::unique_lock lk(mutex_);
+  active_memory_ -= buf->length();
+  if (get_cache_memory() < max_pool_size_) {
+    buffer_cache_.recycle_to_cache(buf);
+  } else {
+    num_resources_--;
+    if (!buf->heap()) {
+      residency_set_.erase(buf);
+    }
+    lk.unlock();
+    auto pool = metal::new_scoped_memory_pool();
+    buf->release();
+  }
+}
+
+size_t MetalAllocator::size(Buffer buffer) const {
+  return static_cast<MTL::Buffer*>(buffer.ptr())->length();
+}
+
+Buffer MetalAllocator::make_buffer(void* ptr, size_t size) {
+  auto buf = device_->newBuffer(ptr, size, resource_options, nullptr);
+  if (!buf) {
+    return Buffer{nullptr};
+  }
+  std::unique_lock lk(mutex_);
+  residency_set_.insert(buf);
+  active_memory_ += buf->length();
+  peak_memory_ = std::max(peak_memory_, active_memory_);
+  num_resources_++;
+  return Buffer{static_cast<void*>(buf)};
+}
+
+void MetalAllocator::release(Buffer buffer) {
+  auto buf = static_cast<MTL::Buffer*>(buffer.ptr());
+  if (buf == nullptr) {
+    return;
+  }
+  std::unique_lock lk(mutex_);
+  active_memory_ -= buf->length();
+  num_resources_--;
+  lk.unlock();
+  auto pool = metal::new_scoped_memory_pool();
+  buf->release();
+}
+
+MetalAllocator& allocator() {
+  // By creating the |allocator_| on heap, the destructor of MetalAllocator
+  // will not be called on exit and buffers in the cache will be leaked. This
+  // can save some time at program exit.
+  static MetalAllocator* allocator_ = new MetalAllocator;
+  return *allocator_;
+}
+
+} // namespace metal
+
+size_t set_cache_limit(size_t limit) {
+  return metal::allocator().set_cache_limit(limit);
+}
+size_t set_memory_limit(size_t limit) {
+  return metal::allocator().set_memory_limit(limit);
+}
+size_t get_memory_limit() {
+  return metal::allocator().get_memory_limit();
+}
+size_t set_wired_limit(size_t limit) {
+  if (limit > std::get<size_t>(
+                  gpu::device_info(0).at("max_recommended_working_set_size"))) {
+    throw std::invalid_argument(
+        "[metal::set_wired_limit] Setting a wired limit larger than "
+        "the maximum working set size is not allowed.");
+  }
+  return metal::allocator().set_wired_limit(limit);
+}
+size_t get_active_memory() {
+  return metal::allocator().get_active_memory();
+}
+size_t get_peak_memory() {
+  return metal::allocator().get_peak_memory();
+}
+void reset_peak_memory() {
+  metal::allocator().reset_peak_memory();
+}
+size_t get_cache_memory() {
+  return metal::allocator().get_cache_memory();
+}
+void clear_cache() {
+  return metal::allocator().clear_cache();
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/allocator.h b/macos/third_party/mlx/mlx/backend/metal/allocator.h
new file mode 100644
index 0000000..5e177b3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/allocator.h
@@ -0,0 +1,79 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <map>
+#include <mutex>
+#include <vector>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/buffer_cache.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/resident.h"
+
+namespace mlx::core::metal {
+
+using allocator::Buffer;
+
+class MetalAllocator : public allocator::Allocator {
+  /** Allocator for Metal GPUs. */
+ public:
+  virtual Buffer malloc(size_t size) override;
+  virtual void free(Buffer buffer) override;
+  virtual size_t size(Buffer buffer) const override;
+  virtual Buffer make_buffer(void* ptr, size_t size) override;
+  virtual void release(Buffer buffer) override;
+
+  size_t get_active_memory() {
+    return active_memory_;
+  };
+  size_t get_peak_memory() {
+    return peak_memory_;
+  };
+  void reset_peak_memory() {
+    std::unique_lock lk(mutex_);
+    peak_memory_ = 0;
+  };
+  size_t get_cache_memory() {
+    return buffer_cache_.cache_size();
+  };
+  size_t set_cache_limit(size_t limit);
+  size_t set_memory_limit(size_t limit);
+  size_t get_memory_limit();
+  size_t set_wired_limit(size_t limit);
+  void clear_cache();
+
+ private:
+  MTL::Device* device_;
+
+  // The size of allocations which go on the heap until it is full. This size
+  // is chosen because it is the actual minimum size of a buffer allocated from
+  // the heap, a heap can have at most heap.size() / 256 buffers.
+  static constexpr int small_size_ = 256;
+  static constexpr int heap_size_ = 1 << 20;
+  MTL::Heap* heap_;
+  MetalAllocator();
+  ~MetalAllocator();
+  friend MetalAllocator& allocator();
+
+  // Caching allocator
+  BufferCache<MTL::Buffer> buffer_cache_;
+
+  ResidencySet residency_set_;
+
+  // Allocation stats
+  size_t block_limit_;
+  size_t gc_limit_;
+  size_t active_memory_{0};
+  size_t peak_memory_{0};
+  size_t max_pool_size_;
+  size_t wired_limit_{0};
+  size_t num_resources_{0};
+  size_t resource_limit_{0};
+
+  std::mutex mutex_;
+};
+
+MetalAllocator& allocator();
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/binary.cpp b/macos/third_party/mlx/mlx/backend/metal/binary.cpp
new file mode 100644
index 0000000..8c0e8c3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/binary.cpp
@@ -0,0 +1,257 @@
+// Copyright © 2024 Apple Inc.
+#include "mlx/backend/common/binary.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+#define BINARY_GPU(func)                                              \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    binary_op_gpu(inputs, out, name());                               \
+  }
+
+#define BINARY_GPU_MULTI(func)                                         \
+  void func::eval_gpu(                                                 \
+      const std::vector<array>& inputs, std::vector<array>& outputs) { \
+    binary_op_gpu(inputs, outputs, name());                            \
+  }
+
+namespace mlx::core {
+
+std::string get_kernel_name(
+    BinaryOpType bopt,
+    const char* op,
+    const array& a,
+    bool large,
+    int ndim,
+    int work_per_thread) {
+  std::string kname;
+  switch (bopt) {
+    case BinaryOpType::ScalarScalar:
+      kname = "ss";
+      break;
+    case BinaryOpType::ScalarVector:
+      kname = "sv";
+      break;
+    case BinaryOpType::VectorScalar:
+      kname = "vs";
+      break;
+    case BinaryOpType::VectorVector:
+      kname = "vv";
+      break;
+    case BinaryOpType::General:
+      kname = "g";
+      if (ndim <= 3) {
+        kname += std::to_string(ndim);
+      } else {
+        concatenate(kname, "n", std::to_string(work_per_thread));
+      }
+      if (large) {
+        kname += "large";
+      }
+      break;
+  }
+  if (bopt != BinaryOpType::General && bopt != BinaryOpType::ScalarScalar) {
+    if (large) {
+      kname += "2";
+    } else if (work_per_thread > 1) {
+      kname += "n";
+    }
+  }
+  concatenate(kname, "_", op, type_to_name(a));
+  return kname;
+}
+
+void binary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const char* op,
+    const Stream& s) {
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto bopt = get_binary_op_type(a, b);
+
+  auto& out = outputs[0];
+  if (out.size() == 0) {
+    return;
+  }
+
+  // Try to collapse contiguous dims
+  auto maybe_collapse = [bopt, &a, &b, &out]() {
+    if (bopt == BinaryOpType::General) {
+      auto [shape, strides] = collapse_contiguous_dims(a, b, out);
+      return std::make_tuple(shape, strides[0], strides[1], strides[2]);
+    } else {
+      decltype(a.strides()) e{};
+      return std::make_tuple(decltype(a.shape()){}, e, e, e);
+    }
+  };
+  auto [shape, strides_a, strides_b, strides_out] = maybe_collapse();
+
+  bool large;
+  auto ndim = shape.size();
+  int work_per_thread;
+  if (bopt == BinaryOpType::General) {
+    large = a.data_size() > INT32_MAX || b.data_size() > INT32_MAX ||
+        out.size() > INT32_MAX;
+    work_per_thread = large ? 4 : 2;
+  } else {
+    large = out.data_size() > UINT32_MAX;
+    work_per_thread = get_work_per_thread(a.dtype(), out.data_size());
+  }
+  std::string kernel_name =
+      get_kernel_name(bopt, op, a, large, shape.size(), work_per_thread);
+  auto& d = metal::device(s.device);
+
+  auto kernel = outputs.size() == 2
+      ? get_binary_two_kernel(d, kernel_name, a.dtype(), out.dtype(), op)
+      : get_binary_kernel(d, kernel_name, a.dtype(), out.dtype(), op);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int arg_idx = 0;
+  compute_encoder.set_input_array(a, arg_idx++);
+  compute_encoder.set_input_array(b, arg_idx++);
+  compute_encoder.set_output_array(outputs[0], arg_idx++);
+  if (outputs.size() == 2) {
+    compute_encoder.set_output_array(outputs[1], arg_idx++);
+  }
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (bopt == BinaryOpType::General) {
+    // Launch up to 3D grid of threads
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = out.size() / (dim0 * dim1);
+
+    if (ndim > 3) {
+      compute_encoder.set_vector_bytes(shape, arg_idx++);
+      compute_encoder.set_vector_bytes(strides_a, arg_idx++);
+      compute_encoder.set_vector_bytes(strides_b, arg_idx++);
+      compute_encoder.set_bytes<int>(ndim, arg_idx++);
+      dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    } else {
+      // The shape is implicit in the grid for <= 3D
+      compute_encoder.set_vector_bytes(strides_a, arg_idx++);
+      compute_encoder.set_vector_bytes(strides_b, arg_idx++);
+    }
+
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::binary] Must use 1024 sized block");
+    }
+    auto group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    // Launch a 1D or 2D grid of threads
+    size_t nthreads = ceildiv(out.data_size(), work_per_thread);
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    MTL::Size grid_dims;
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(out.data_size(), arg_idx++);
+      grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+    } else {
+      compute_encoder.set_bytes<int>(out.data_size(), arg_idx++);
+      grid_dims = MTL::Size(nthreads, 1, 1);
+    }
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void binary_op_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const char* op,
+    const Stream& s) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, outputs[0], bopt);
+  set_binary_op_output_data(a, b, outputs[1], bopt);
+  binary_op_gpu_inplace(inputs, outputs, op, s);
+}
+
+void binary_op_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const char* op) {
+  auto& s = outputs[0].primitive().stream();
+  binary_op_gpu(inputs, outputs, op, s);
+}
+
+void binary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  std::vector<array> outputs = {out};
+  binary_op_gpu_inplace(inputs, outputs, op, s);
+}
+
+void binary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  assert(inputs.size() == 2);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto bopt = get_binary_op_type(a, b);
+  set_binary_op_output_data(a, b, out, bopt);
+  binary_op_gpu_inplace(inputs, out, op, s);
+}
+
+void binary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op) {
+  auto& s = out.primitive().stream();
+  binary_op_gpu(inputs, out, op, s);
+}
+
+BINARY_GPU(Add)
+BINARY_GPU(ArcTan2)
+BINARY_GPU(Divide)
+BINARY_GPU_MULTI(DivMod)
+BINARY_GPU(Remainder)
+BINARY_GPU(Equal)
+BINARY_GPU(Greater)
+BINARY_GPU(GreaterEqual)
+BINARY_GPU(Less)
+BINARY_GPU(LessEqual)
+BINARY_GPU(LogicalAnd)
+BINARY_GPU(LogicalOr)
+BINARY_GPU(LogAddExp)
+BINARY_GPU(Maximum)
+BINARY_GPU(Minimum)
+BINARY_GPU(Multiply)
+BINARY_GPU(NotEqual)
+BINARY_GPU(Power)
+BINARY_GPU(Subtract)
+
+void BitwiseBinary::eval_gpu(const std::vector<array>& inputs, array& out) {
+  switch (op_) {
+    case BitwiseBinary::And:
+      binary_op_gpu(inputs, out, name());
+      break;
+    case BitwiseBinary::Or:
+      binary_op_gpu(inputs, out, name());
+      break;
+    case BitwiseBinary::Xor:
+      binary_op_gpu(inputs, out, name());
+      break;
+    case BitwiseBinary::LeftShift:
+      binary_op_gpu(inputs, out, name());
+      break;
+    case BitwiseBinary::RightShift:
+      binary_op_gpu(inputs, out, name());
+      break;
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/binary.h b/macos/third_party/mlx/mlx/backend/metal/binary.h
new file mode 100644
index 0000000..0341a2f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/binary.h
@@ -0,0 +1,33 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+void binary_op_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const char* op,
+    const Stream& s);
+
+void binary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+void binary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const char* op,
+    const Stream& s);
+
+void binary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/compiled.cpp b/macos/third_party/mlx/mlx/backend/metal/compiled.cpp
new file mode 100644
index 0000000..eb51ab7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/compiled.cpp
@@ -0,0 +1,471 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <fmt/format.h>
+#include <sstream>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/jit/includes.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/graph_utils.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+inline void build_kernel(
+    std::string& os,
+    const std::string& kernel_name,
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    const std::vector<array>& tape,
+    const std::function<bool(size_t)>& is_constant,
+    bool contiguous,
+    int ndim,
+    bool dynamic_dims,
+    bool use_big_index = false,
+    int work_per_thread = 1) {
+  NodeNamer namer;
+  bool add_indices = false;
+  int cnt = 0;
+
+  // Start the kernel
+  os += fmt::format(
+      "[[host_name(\"{0}\")]]\n[[kernel]] void {0}(\n", kernel_name);
+
+  // Add the input arguments
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    // Skip constants from the input list
+    if (is_constant(i)) {
+      continue;
+    }
+
+    const auto& x = inputs[i];
+    auto& xname = namer.get_name(x);
+
+    // Scalars and contiguous need no strides
+    if (!is_scalar(x) && !contiguous) {
+      add_indices = true;
+    }
+    os += fmt::format(
+        "    device const {0}* {1} [[buffer({2})]],\n",
+        get_type_string(x.dtype()),
+        xname,
+        cnt++);
+  }
+
+  std::string idx_type = use_big_index ? "int64_t" : "uint";
+  if (add_indices) {
+    os += fmt::format(
+        "    constant const int64_t* in_strides [[buffer({0})]],\n", cnt++);
+  }
+
+  // Add the output arguments
+  for (auto& x : outputs) {
+    os += fmt::format(
+        "    device {0}* {1} [[buffer({2})]],\n",
+        get_type_string(x.dtype()),
+        namer.get_name(x),
+        cnt++);
+  }
+  // Add output strides and shape to extract the indices.
+  if (!contiguous) {
+    os += fmt::format(
+        "    constant const int* output_shape [[buffer({0})]],\n", cnt++);
+  } else {
+    os += fmt::format(
+        "    constant const {0}& size [[buffer({1})]],\n", idx_type, cnt++);
+  }
+  if (dynamic_dims) {
+    os += fmt::format("    constant const int& ndim [[buffer({0})]],\n", cnt++);
+  }
+
+  // The thread index in the whole grid
+  os += "    uint3 pos [[thread_position_in_grid]],\n";
+  os += "    uint3 grid [[threads_per_grid]]) {\n";
+
+  os += fmt::format("  constexpr int N_ = {0};\n", work_per_thread);
+  if (contiguous && use_big_index) {
+    // This is only used for contiguous kernels which don't have
+    // a third grid dimension
+    os += "  int64_t index = N_ * (pos.x + grid.x * int64_t(pos.y));\n";
+  } else if (contiguous) {
+    os += "  uint index = N_ * pos.x;\n";
+  } else if (work_per_thread > 1) {
+    os += fmt::format(
+        "  int xshape = output_shape[{0}];\n",
+        dynamic_dims ? "ndim - 1" : std::to_string(ndim - 1));
+    os += fmt::format(
+        "  {0} index = N_ * pos.x + xshape * (pos.y + {0}(grid.y) * pos.z);\n",
+        idx_type);
+  } else {
+    os += fmt::format(
+        "  {0} index = pos.x + grid.x * (pos.y + {0}(grid.y) * pos.z);\n",
+        idx_type);
+  }
+  if (work_per_thread > 1 && contiguous) {
+    os += "  for (int i = 0; i < N_ && index < size; ++i) {\n";
+  }
+
+  // Read constant / contiguous inputs in tmps
+  std::vector<array> nc_inputs;
+  for (int i = 0; i < inputs.size(); ++i) {
+    auto& x = inputs[i];
+    auto& xname = namer.get_name(x);
+
+    if (is_constant(i)) {
+      auto type_str = get_type_string(x.dtype());
+      std::ostringstream ss;
+      print_constant(ss, x);
+      os += fmt::format(
+          "  auto tmp_{0} = static_cast<{1}>({2});\n",
+          xname,
+          get_type_string(x.dtype()),
+          ss.str());
+    } else if (is_scalar(x)) {
+      os += fmt::format(
+          "  {0} tmp_{1} = {1}[0];\n", get_type_string(x.dtype()), xname);
+    } else if (contiguous) {
+      os += fmt::format(
+          "  {0} tmp_{1} = {1}[index];\n", get_type_string(x.dtype()), xname);
+    } else {
+      nc_inputs.push_back(x);
+    }
+  }
+
+  // Initialize the indices for non-contiguous inputs
+  for (int i = 0; i < nc_inputs.size(); ++i) {
+    auto& xname = namer.get_name(nc_inputs[i]);
+    os += fmt::format("  {0} index_{1} = ", idx_type, xname);
+    if (ndim == 1) {
+      int offset = i * ndim;
+      os +=
+          fmt::format("elem_to_loc_1<uint>(pos.x, in_strides[{0}]);\n", offset);
+    } else if (ndim == 2) {
+      int offset = i * ndim;
+      os += fmt::format(
+          "elem_to_loc_2<{0}>({{pos.x, pos.y}}, in_strides + {1});\n",
+          idx_type,
+          offset);
+    } else if (ndim == 3) {
+      int offset = i * ndim;
+      os += fmt::format(
+          "elem_to_loc_3<{0}>(pos, in_strides + {1});\n", idx_type, offset);
+    } else if (!dynamic_dims) {
+      int offset = (i + 1) * ndim;
+      os += fmt::format(
+          "N_ * pos.x * {0}(in_strides[{1}]) + pos.y * {0}(in_strides[{2}]);\n",
+          idx_type,
+          offset - 1,
+          offset - 2);
+    } else {
+      os += fmt::format(
+          "N_ * pos.x * {0}(in_strides[ndim * {1} + ndim - 1]) + pos.y * {0}(in_strides[ndim * {1} + ndim - 2]);\n",
+          idx_type,
+          i);
+    }
+  }
+
+  if (!nc_inputs.empty() && (ndim > 3 || dynamic_dims)) {
+    os += "  uint zpos = pos.z;\n";
+    if (dynamic_dims) {
+      os += "  for (int d = ndim - 3; d >= 0; --d) {\n";
+    } else {
+      os += fmt::format("  for (int d = {0}; d >= 0; --d) {{\n", ndim - 3);
+    }
+    os += "    uint l = zpos % output_shape[d];\n";
+    for (int i = 0; i < nc_inputs.size(); ++i) {
+      auto& xname = namer.get_name(nc_inputs[i]);
+      os += fmt::format("    index_{0} += ", xname);
+      if (dynamic_dims) {
+        os +=
+            fmt::format("l * {0}(in_strides[{1} * ndim + d]);\n", idx_type, i);
+      } else {
+        os +=
+            fmt::format("l * {0}(in_strides[{1} + d]);\n", idx_type, i * ndim);
+      }
+    }
+    os += "    zpos /= output_shape[d];\n  }\n";
+  }
+
+  // Open per-thread loop
+  if (work_per_thread > 1 && !contiguous) {
+    os +=
+        "  for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {\n";
+  }
+
+  // Read non-contiguous inputs into tmps
+  for (int i = 0; i < nc_inputs.size(); ++i) {
+    auto& x = nc_inputs[i];
+    auto& xname = namer.get_name(x);
+    os += fmt::format(
+        "  {0} tmp_{1} = {1}[index_{1}];\n", get_type_string(x.dtype()), xname);
+  }
+
+  // Actually write the computation
+  for (auto& x : tape) {
+    os += fmt::format(
+        "  {0} tmp_{1} = ", get_type_string(x.dtype()), namer.get_name(x));
+    if (is_static_cast(x.primitive())) {
+      os += fmt::format(
+          "static_cast<{0}>(tmp_{1});\n",
+          get_type_string(x.dtype()),
+          namer.get_name(x.inputs()[0]));
+    } else {
+      os += x.primitive().name();
+      os += "()(";
+      for (int i = 0; i < x.inputs().size() - 1; i++) {
+        os += fmt::format("tmp_{0}, ", namer.get_name(x.inputs()[i]));
+      }
+      os += fmt::format("tmp_{0});\n", namer.get_name(x.inputs().back()));
+    }
+  }
+
+  // Write the outputs from tmps
+  for (auto& x : outputs) {
+    os += fmt::format("  {0}[index] = tmp_{0};\n", namer.get_name(x));
+  }
+  // Increment indices and close per thread loop
+  if (work_per_thread > 1) {
+    for (int i = 0; i < nc_inputs.size(); ++i) {
+      auto& x = nc_inputs[i];
+      auto& xname = namer.get_name(x);
+      if (!dynamic_dims) {
+        os += fmt::format(
+            "  index_{0} += in_strides[{1}];\n", xname, i * ndim + ndim - 1);
+      } else {
+        os += fmt::format(
+            "  index_{0} += in_strides[{1} * ndim + ndim - 1];\n", xname, i);
+      }
+    }
+    os += "  index++;\n  }\n";
+  }
+
+  // Finish the kernel
+  os += "}\n";
+
+  if (cnt > 31) {
+    std::ostringstream msg;
+    msg << "[compile] Too many inputs/outputs fused in the Metal Compiled "
+        << "primitive which exhausted the available argument buffers for "
+        << "the kernel. Please file an issue with the function that results "
+        << "in this error. The name of the kernel is '" << kernel_name << "'";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+void Compiled::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  // Get the kernel if someone else built it already
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto lib = d.get_library(kernel_lib_, [&]() {
+    int work_per_thread = get_work_per_thread(outputs_[0].dtype());
+    std::string kernel = metal::utils();
+    concatenate(
+        kernel, metal::unary_ops(), metal::binary_ops(), metal::ternary_ops());
+    build_kernel(
+        kernel,
+        kernel_lib_ + "_contiguous",
+        inputs_,
+        outputs_,
+        tape_,
+        is_constant_,
+        /* contiguous = */ true,
+        /* ndim = */ 0,
+        /* dynamic_dims = */ false,
+        /* use_big_index = */ false,
+        /* work_per_thread = */ 1);
+    if (work_per_thread > 1) {
+      build_kernel(
+          kernel,
+          kernel_lib_ + "_contiguous_n",
+          inputs_,
+          outputs_,
+          tape_,
+          is_constant_,
+          /* contiguous = */ true,
+          /* ndim = */ 0,
+          /* dynamic_dims = */ false,
+          /* use_big_index = */ false,
+          /* work_per_thread = */ work_per_thread);
+    }
+    build_kernel(
+        kernel,
+        kernel_lib_ + "_contiguous_large",
+        inputs_,
+        outputs_,
+        tape_,
+        is_constant_,
+        /* contiguous = */ true,
+        /* ndim = */ 0,
+        /* dynamic_dims = */ false,
+        /* use_big_index = */ true,
+        /* work_per_thread = */ work_per_thread);
+    for (int i = 1; i < 8; i++) {
+      build_kernel(
+          kernel,
+          kernel_lib_ + "_strided_" + std::to_string(i),
+          inputs_,
+          outputs_,
+          tape_,
+          is_constant_,
+          /* contiguous = */ false,
+          /* ndim = */ i,
+          /* dynamic_dims = */ false,
+          /* use_big_index = */ false,
+          /* work_per_thread = */ i > 3 ? 2 : 1);
+      if (i > 1) {
+        build_kernel(
+            kernel,
+            kernel_lib_ + "_strided_" + std::to_string(i) + "_large",
+            inputs_,
+            outputs_,
+            tape_,
+            is_constant_,
+            /* contiguous = */ false,
+            /* ndim = */ i,
+            /* dynamic_dims = */ false,
+            /* use_big_index = */ true,
+            /* work_per_thread = */ i > 3 ? 4 : 1);
+      }
+    }
+    build_kernel(
+        kernel,
+        kernel_lib_ + "_strided_dynamic",
+        inputs_,
+        outputs_,
+        tape_,
+        is_constant_,
+        /* contiguous = */ false,
+        /* ndim = */ 0,
+        /* dynamic_dims = */ true,
+        /* use_big_index = */ false,
+        /* work_per_thread = */ 2);
+    build_kernel(
+        kernel,
+        kernel_lib_ + "_strided_dynamic_large",
+        inputs_,
+        outputs_,
+        tape_,
+        is_constant_,
+        /* contiguous = */ false,
+        /* ndim = */ 0,
+        /* dynamic_dims = */ true,
+        /* use_big_index = */ true,
+        /* work_per_thread = */ 4);
+    return kernel;
+  });
+
+  // Collapse contiguous dims to route to a faster kernel if possible. Also
+  // handle all broadcasting.
+  auto [contiguous, shape, strides] =
+      compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_);
+
+  // Whether to use large index.
+  bool large = compiled_use_large_index(inputs, outputs, contiguous);
+
+  // Get the kernel from the lib
+  int ndim = shape.size();
+  bool dynamic = ndim >= 8;
+  auto kernel_name = kernel_lib_ + (contiguous ? "_contiguous" : "_strided_");
+  int work_per_thread = 1;
+  if (!contiguous) {
+    if (dynamic) {
+      kernel_name += "dynamic";
+    } else {
+      kernel_name += std::to_string(shape.size());
+    }
+    work_per_thread = ndim > 3 ? (large ? 4 : 2) : 1;
+  } else {
+    work_per_thread =
+        get_work_per_thread(outputs[0].dtype(), outputs[0].data_size());
+    if (work_per_thread > 1 && !large) {
+      kernel_name += "_n";
+    }
+  }
+  if (large) {
+    kernel_name += "_large";
+  }
+  auto kernel = d.get_kernel(kernel_name, lib);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Put the inputs in
+  int cnt = 0;
+  int stride_idx = 1; // idx 0 is the output strides
+  Strides in_strides;
+  for (int i = 0; i < inputs.size(); i++) {
+    if (is_constant_(i)) {
+      continue;
+    }
+    auto& x = inputs[i];
+    compute_encoder.set_input_array(x, cnt++);
+    if (!contiguous && !is_scalar(x)) {
+      in_strides.insert(
+          in_strides.end(),
+          strides[stride_idx].begin(),
+          strides[stride_idx].end());
+      stride_idx++;
+    }
+  }
+  if (!in_strides.empty()) {
+    compute_encoder.set_vector_bytes(in_strides, cnt++);
+  }
+
+  compiled_allocate_outputs(inputs, outputs, is_constant_, contiguous);
+
+  // Put the outputs in
+  for (auto& x : outputs) {
+    compute_encoder.set_output_array(x, cnt++);
+  }
+
+  // Put the output shape and strides in
+  if (!contiguous) {
+    compute_encoder.set_vector_bytes(shape, cnt++);
+  } else {
+    auto size = outputs[0].data_size();
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(size, cnt++);
+    } else {
+      compute_encoder.set_bytes<int>(size, cnt++);
+    }
+  }
+
+  // Put the number of dims in if it is dynamic
+  if (dynamic) {
+    compute_encoder.set_bytes(ndim, cnt++);
+  }
+
+  // Launch the kernel
+  if (contiguous) {
+    size_t nthreads = ceildiv(outputs[0].data_size(), work_per_thread);
+    MTL::Size group_dims(
+        std::min(nthreads, kernel->maxTotalThreadsPerThreadgroup()), 1, 1);
+    MTL::Size grid_dims = large
+        ? get_2d_grid_dims(
+              outputs[0].shape(), outputs[0].strides(), work_per_thread)
+        : MTL::Size(nthreads, 1, 1);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = outputs[0].size() / (dim0 * dim1);
+    dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    int pow2;
+    if (thread_group_size == 1024) {
+      pow2 = 10;
+    } else if (thread_group_size > 512) {
+      pow2 = 9;
+    } else {
+      throw std::runtime_error("[Metal::compiled] Must use > 512 sized block");
+    }
+    auto group_dims = get_block_dims(dim0, dim1, rest, pow2);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/conv.cpp b/macos/third_party/mlx/mlx/backend/metal/conv.cpp
new file mode 100644
index 0000000..b4a674f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/conv.cpp
@@ -0,0 +1,1118 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <algorithm>
+#include <cassert>
+#include <numeric>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+#include "mlx/backend/metal/matmul.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+using namespace mlx::steel;
+
+namespace mlx::core {
+
+namespace {
+
+template <int N>
+void explicit_gemm_conv_ND_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const MLXConvParams<N>& conv_params) {
+  // Get gemm shapes
+  int implicit_M = out.size() / conv_params.O;
+  int implicit_K = wt.size() / conv_params.O;
+  int implicit_N = conv_params.O;
+  // Prepare unfolding array
+  Shape unfolded_shape{implicit_M, implicit_K};
+  array in_unfolded(unfolded_shape, in.dtype(), nullptr, {});
+
+  in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes()));
+
+  // Prepare unfolding kernel
+  std::string kname;
+  kname.reserve(32);
+  concatenate(kname, "naive_unfold_nd_", type_to_name(in_unfolded), "_", N);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(in_unfolded, 1);
+
+  compute_encoder.set_bytes(conv_params, 2);
+
+  // Launch unfolding kernel
+  size_t tgp_x = std::min(conv_params.C, 64);
+  tgp_x = 32 * ((tgp_x + 32 - 1) / 32);
+  size_t tgp_y = 256 / tgp_x;
+
+  MTL::Size grid_dims = MTL::Size(
+      conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]);
+  MTL::Size group_dims = MTL::Size(
+      std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1);
+
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+  // Reshape weight
+  Shape wt_reshape{implicit_K, implicit_N};
+  Strides wt_restride{1, implicit_K};
+  array wt_reshaped(wt_reshape, wt.dtype(), nullptr, {});
+  auto wt_flags = wt.flags();
+  wt_flags.row_contiguous = false;
+  wt_flags.col_contiguous = true;
+  wt_reshaped.copy_shared_buffer(wt, wt_restride, wt_flags, wt.data_size());
+
+  // Perform gemm
+  std::vector<array> copies = {in_unfolded};
+  return steel_matmul(
+      s,
+      d,
+      /*a = */ in_unfolded,
+      /*b = */ wt_reshaped,
+      /*c = */ out,
+      /*M = */ implicit_M,
+      /*N = */ implicit_N,
+      /*K = */ implicit_K,
+      /*batch_size_out = */ 1,
+      /*a_cols = */ implicit_K,
+      /*b_cols = */ implicit_K,
+      /*a_transposed = */ false,
+      /*b_transposed = */ true,
+      /*copies = */ copies);
+}
+
+template <int N>
+void explicit_gemm_conv_group_ND_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const MLXConvParams<N>& conv_params) {
+  const int groups = conv_params.groups;
+  const int C_per_group = conv_params.C / conv_params.groups;
+  const int O_per_group = conv_params.O / conv_params.groups;
+  // Get gemm shapes
+  const int implicit_M = out.size() / conv_params.O;
+  const int implicit_K = wt.size() / conv_params.O;
+  const int implicit_N = O_per_group;
+
+  int kernel_size = 1;
+  for (int i = 0; i < N; ++i) {
+    kernel_size *= conv_params.wS[i];
+  }
+
+  // Prepare unfolding array
+  Shape unfolded_shape{implicit_M, implicit_K * groups};
+  array in_unfolded(unfolded_shape, in.dtype(), nullptr, {});
+  in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes()));
+
+  // Prepare unfolding kernel
+  std::string kname;
+  kname.reserve(32);
+  concatenate(
+      kname, "naive_unfold_transpose_nd_", type_to_name(in_unfolded), "_", N);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(in_unfolded, 1);
+
+  compute_encoder.set_bytes(conv_params, 2);
+
+  // Launch unfolding kernel
+  size_t tgp_x = std::min(conv_params.C, 64);
+  tgp_x = 32 * ((tgp_x + 32 - 1) / 32);
+  size_t tgp_y = 256 / tgp_x;
+
+  MTL::Size grid_dims = MTL::Size(
+      conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]);
+  MTL::Size group_dims = MTL::Size(
+      std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1);
+
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+  // Transpose kernel weights so that we can slice them by contiguous chunks
+  // of channel groups.
+  array wt_view(
+      {wt.shape(0), C_per_group, kernel_size}, wt.dtype(), nullptr, {});
+  wt_view.copy_shared_buffer(
+      wt, {wt.strides(0), 1, C_per_group}, wt.flags(), wt.size());
+
+  // Materialize
+  array wt_transpose = contiguous_copy_gpu(wt_view, s);
+
+  // Perform gemm
+  std::vector<array> copies = {in_unfolded, wt_transpose};
+  return steel_matmul_regular(
+      /* const Stream& s = */ s,
+      /* Device& d = */ d,
+      /* const array& a = */ in_unfolded,
+      /* const array& b = */ wt_transpose,
+      /* array& c = */ out,
+      /* int M = */ implicit_M,
+      /* int N = */ implicit_N,
+      /* int K = */ implicit_K,
+      /* int batch_size_out = */ groups,
+      /* int lda = */ implicit_K * groups,
+      /* int ldb = */ implicit_K,
+      /* int ldd = */ implicit_N * groups,
+      /* bool transpose_a = */ false,
+      /* bool transpose_b = */ true,
+      /* std::vector<array>& copies = */ copies,
+      /* Shape batch_shape = */ {1},
+      /* Strides batch_strides = */ {0},
+      /* int64_t A_batch_strides = */ int64_t(implicit_K),
+      /* int64_t B_batch_strides = */ int64_t(implicit_N) * implicit_K,
+      /* int64_t matrix_stride_out = */ int64_t(implicit_N));
+}
+
+void implicit_gemm_conv_2D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const MLXConvParams<2>& conv_params) {
+  const int groups = conv_params.groups;
+  const int C_per_group = conv_params.C / conv_params.groups;
+  const int O_per_group = conv_params.O / conv_params.groups;
+
+  // Deduce implicit gemm size
+  const int implicit_M = conv_params.N * conv_params.oS[0] * conv_params.oS[1];
+  const int implicit_N = O_per_group;
+  const int implicit_K = conv_params.wS[0] * conv_params.wS[1] * C_per_group;
+
+  // Determine block and warp tiles
+  int wm = 2, wn = 2;
+
+  int bm = implicit_M >= 8192 && C_per_group >= 64 ? 64 : 32;
+  int bn = (bm == 64 || implicit_N >= 64) ? 64 : 32;
+  int bk = 16;
+
+  if (implicit_N <= 16) {
+    bn = 8;
+    wm = 4;
+    wn = 1;
+  }
+
+  int tn = (implicit_N + bn - 1) / bn;
+  int tm = (implicit_M + bm - 1) / bm;
+  int swizzle_log = 0;
+
+  // Fix small channel specialization
+  int n_channel_specialization = 0;
+  int channel_k_iters = ((C_per_group + bk - 1) / bk);
+  int gemm_k_iters = conv_params.wS[0] * conv_params.wS[1] * channel_k_iters;
+
+  if (C_per_group <= 2) {
+    gemm_k_iters = (implicit_K + bk - 1) / bk;
+    n_channel_specialization = C_per_group;
+  } else if (C_per_group <= 4) {
+    gemm_k_iters = ((conv_params.wS[0] * conv_params.wS[1] * 4) + bk - 1) / bk;
+    n_channel_specialization = C_per_group;
+  }
+
+  bool small_filter = (!n_channel_specialization) &&
+      (conv_params.wS[0] <= 16 && conv_params.wS[1] <= 16);
+
+  // Fix host side helper params
+  int sign = (conv_params.flip ? -1 : 1);
+  int ijw = conv_params.in_strides[2] * conv_params.kdil[1];
+  int ijh = conv_params.in_strides[1] * conv_params.kdil[0];
+
+  int inp_jump_w = sign * ijw;
+  int inp_jump_h = sign * (ijh - (conv_params.wS[1] - 1) * ijw);
+  int inp_jump_c = bk - sign * (conv_params.wS[0] - 1) * ijh -
+      sign * (conv_params.wS[1] - 1) * ijw;
+
+  // Build implicit gemm params
+  ImplicitGemmConv2DParams gemm_params{
+      /* const int M = */ implicit_M,
+      /* const int N = */ implicit_N,
+      /* const int K = */ implicit_K,
+
+      /* const int gemm_k_iterations = */ gemm_k_iters,
+
+      /* const int inp_jump_w = */ inp_jump_w,
+      /* const int inp_jump_h = */ inp_jump_h,
+      /* const int inp_jump_c = */ inp_jump_c,
+
+      /* const int tiles_n = */ tn,
+      /* const int tiles_m = */ tm,
+      /* const int swizzle_log = */ swizzle_log};
+
+  // Determine kernel
+  std::string kname;
+  kname.reserve(64);
+  concatenate(
+      kname,
+      "implicit_gemm_conv_2d_",
+      type_to_name(out),
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn,
+      "_channel_",
+      n_channel_specialization ? std::to_string(n_channel_specialization) : "l",
+      "_filter_",
+      small_filter ? 's' : 'l');
+
+  // Encode and dispatch kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_conv_kernel(
+      d,
+      kname,
+      out,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn,
+      n_channel_specialization,
+      small_filter);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Deduce grid launch dimensions
+  int tile = 1 << swizzle_log;
+  size_t grid_dim_y = (tm + tile - 1) / tile;
+  size_t grid_dim_x = tn * tile;
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(grid_dim_x, grid_dim_y, groups);
+
+  // Encode arrays
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_input_array(wt, 1);
+  compute_encoder.set_output_array(out, 2);
+
+  // Encode params
+  compute_encoder.set_bytes(conv_params, 3);
+  compute_encoder.set_bytes(gemm_params, 4);
+
+  // Launch kernel
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void implicit_gemm_conv_2D_general_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const MLXConvParams<2>& conv_params) {
+  // Deduce implicit gemm size
+  int implicit_M = conv_params.N * conv_params.oS[0] * conv_params.oS[1];
+  int implicit_N = conv_params.O;
+  int implicit_K = conv_params.wS[0] * conv_params.wS[1] * conv_params.C;
+
+  // Determine block and warp tiles
+  int wm = 2, wn = 2;
+
+  // Make jump params
+  int f_wgt_jump_h =
+      std::lcm(conv_params.idil[0], conv_params.kdil[0]) / conv_params.kdil[0];
+  int f_wgt_jump_w =
+      std::lcm(conv_params.idil[1], conv_params.kdil[1]) / conv_params.kdil[1];
+
+  int f_out_jump_h =
+      std::lcm(conv_params.idil[0], conv_params.str[0]) / conv_params.str[0];
+  int f_out_jump_w =
+      std::lcm(conv_params.idil[1], conv_params.str[1]) / conv_params.str[1];
+
+  int adj_out_h = (conv_params.oS[0] + f_out_jump_h - 1) / f_out_jump_h;
+  int adj_out_w = (conv_params.oS[1] + f_out_jump_w - 1) / f_out_jump_w;
+  int adj_out_hw = adj_out_h * adj_out_w;
+  int adj_implicit_m = conv_params.N * adj_out_hw;
+
+  Conv2DGeneralJumpParams jump_params{
+      /* const int f_wgt_jump_h = */ f_wgt_jump_h,
+      /* const int f_wgt_jump_w = */ f_wgt_jump_w,
+
+      /* const int f_out_jump_h = */ f_out_jump_h,
+      /* const int f_out_jump_w = */ f_out_jump_w,
+
+      /* const int adj_out_h = */ adj_out_h,
+      /* const int adj_out_w = */ adj_out_w,
+      /* const int adj_out_hw = */ adj_out_hw,
+      /* const int adj_implicit_m = */ adj_implicit_m};
+
+  // Make base info
+  std::vector<Conv2DGeneralBaseInfo> base_h(f_out_jump_h);
+  std::vector<Conv2DGeneralBaseInfo> base_w(f_out_jump_w);
+
+  int jump_h = conv_params.flip ? -conv_params.kdil[0] : conv_params.kdil[0];
+  int jump_w = conv_params.flip ? -conv_params.kdil[1] : conv_params.kdil[1];
+
+  int init_h =
+      (conv_params.flip ? (conv_params.wS[0] - 1) * conv_params.kdil[0] : 0);
+  int init_w =
+      (conv_params.flip ? (conv_params.wS[1] - 1) * conv_params.kdil[1] : 0);
+
+  for (int i = 0; i < f_out_jump_h; ++i) {
+    int ih_loop = i * conv_params.str[0] - conv_params.pad[0] + init_h;
+
+    int wh_base = 0;
+    while (wh_base < conv_params.wS[0] && ih_loop % conv_params.idil[0] != 0) {
+      wh_base++;
+      ih_loop += jump_h;
+    }
+
+    int wh_size =
+        ((conv_params.wS[0] - wh_base) + f_wgt_jump_h - 1) / f_wgt_jump_h;
+    base_h[i] = {wh_base, wh_size};
+  }
+
+  for (int j = 0; j < f_out_jump_w; ++j) {
+    int iw_loop = j * conv_params.str[1] - conv_params.pad[1] + init_w;
+
+    int ww_base = 0;
+    while (ww_base < conv_params.wS[1] && iw_loop % conv_params.idil[1] != 0) {
+      ww_base++;
+      iw_loop += jump_w;
+    }
+
+    int ww_size =
+        ((conv_params.wS[1] - ww_base) + f_wgt_jump_w - 1) / f_wgt_jump_w;
+    base_w[j] = {ww_base, ww_size};
+  }
+
+  // Collect block sizes
+  int bm = adj_implicit_m >= 8192 && conv_params.C >= 64 ? 64 : 32;
+  int bn = (bm == 64 && implicit_N >= 64) ? 64 : 32;
+  int bk = 16;
+
+  int tn = (implicit_N + bn - 1) / bn;
+  int tm = (adj_implicit_m + bm - 1) / bm;
+  int swizzle_log = 0;
+
+  // Get channel iteration info
+  int channel_k_iters = ((conv_params.C + bk - 1) / bk);
+  int gemm_k_iters = channel_k_iters;
+  bool align_C = conv_params.C % bk == 0;
+
+  // Fix host side helper params
+  int sign = (conv_params.flip ? -1 : 1);
+  int ijw = conv_params.in_strides[2] * conv_params.kdil[1];
+  int ijh = conv_params.in_strides[1] * conv_params.kdil[0];
+
+  int inp_jump_w = sign * ijw;
+  int inp_jump_h = sign * (ijh - (conv_params.wS[1] - 1) * ijw);
+  int inp_jump_c = bk - sign * (conv_params.wS[0] - 1) * ijh -
+      sign * (conv_params.wS[1] - 1) * ijw;
+
+  // Build implicit gemm params
+  ImplicitGemmConv2DParams gemm_params{
+      /* const int M = */ implicit_M,
+      /* const int N = */ implicit_N,
+      /* const int K = */ implicit_K,
+
+      /* const int gemm_k_iterations = */ gemm_k_iters,
+
+      /* const int inp_jump_w = */ inp_jump_w,
+      /* const int inp_jump_h = */ inp_jump_h,
+      /* const int inp_jump_c = */ inp_jump_c,
+
+      /* const int tiles_n = */ tn,
+      /* const int tiles_m = */ tm,
+      /* const int swizzle_log = */ swizzle_log};
+
+  // Determine kernel
+  std::string kname;
+  kname.reserve(64);
+  concatenate(
+      kname,
+      "implicit_gemm_conv_2d_general_",
+      type_to_name(out),
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn);
+  std::string hash_name;
+  hash_name.reserve(64);
+  concatenate(hash_name, kname, "_alC_", align_C);
+  metal::MTLFCList func_consts = {
+      {&align_C, MTL::DataType::DataTypeBool, 200},
+  };
+
+  // Encode and dispatch kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_conv_general_kernel(
+      d, kname, hash_name, func_consts, out, bm, bn, bk, wm, wn);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Deduce grid launch dimensions
+  int tile = 1 << swizzle_log;
+  size_t grid_dim_y = (tm + tile - 1) / tile;
+  size_t grid_dim_x = tn * tile;
+  size_t grid_dim_z = f_out_jump_h * f_out_jump_w;
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(grid_dim_x, grid_dim_y, grid_dim_z);
+
+  // Encode arrays
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_input_array(wt, 1);
+  compute_encoder.set_output_array(out, 2);
+
+  // Encode params
+  compute_encoder.set_bytes(conv_params, 3);
+  compute_encoder.set_bytes(gemm_params, 4);
+  compute_encoder.set_bytes(jump_params, 5);
+
+  compute_encoder.set_vector_bytes(base_h, 6);
+  compute_encoder.set_vector_bytes(base_w, 7);
+
+  // Launch kernel
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void winograd_conv_2D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const MLXConvParams<2>& conv_params,
+    std::vector<array>& copies_w) {
+  Shape padded_shape = {
+      conv_params.N,
+      conv_params.iS[0] + 2 * conv_params.pad[0],
+      conv_params.iS[1] + 2 * conv_params.pad[1],
+      conv_params.C};
+
+  padded_shape[1] = 6 * ((padded_shape[1] - 2 + 5) / 6) + 2;
+  padded_shape[2] = 6 * ((padded_shape[2] - 2 + 5) / 6) + 2;
+
+  array in_padded(std::move(padded_shape), in.dtype(), nullptr, {});
+
+  // Fill with zeros
+  array zero_arr = array(0, in.dtype());
+  fill_gpu(zero_arr, in_padded, s);
+  copies_w.push_back(zero_arr);
+
+  // Pick input slice from padded
+  size_t data_offset = conv_params.pad[0] * in_padded.strides()[1] +
+      conv_params.pad[1] * in_padded.strides()[2];
+  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
+  in_padded_slice.copy_shared_buffer(
+      in_padded,
+      in_padded.strides(),
+      in_padded.flags(),
+      in_padded_slice.size(),
+      data_offset);
+
+  // Copy input values into the slice
+  copy_gpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, s);
+
+  copies_w.push_back(in_padded_slice);
+  copies_w.push_back(in_padded);
+
+  MLXConvParams<2> conv_params_updated{
+      /* const int  N = */ static_cast<int>(in_padded.shape(0)),
+      /* const int  C = */ static_cast<int>(in_padded.shape(3)),
+      /* const int  O = */ static_cast<int>(wt.shape(0)),
+      /* const int iS[NDIM] = */
+      {static_cast<int>(in_padded.shape(1)),
+       static_cast<int>(in_padded.shape(2))},
+      /* const int wS[NDIM] = */
+      {static_cast<int>(wt.shape(1)), static_cast<int>(wt.shape(2))},
+      /* const int oS[NDIM] = */
+      {static_cast<int>(out.shape(1)), static_cast<int>(out.shape(2))},
+      /* const int str[NDIM] = */ {1, 1},
+      /* const int pad[NDIM] = */ {0, 0},
+      /* const int kdil[NDIM] = */ {1, 1},
+      /* const int idil[NDIM] = */ {1, 1},
+      /* const size_t in_strides[NDIM + 2] = */
+      {in_padded.strides()[0],
+       in_padded.strides()[1],
+       in_padded.strides()[2],
+       in_padded.strides()[3]},
+      /* const size_t wt_strides[NDIM + 2] = */
+      {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]},
+      /* const size_t out_strides[NDIM + 2] = */
+      {out.strides()[0], out.strides()[1], out.strides()[2], out.strides()[3]},
+      /* const int groups = */ 1,
+      /* const bool flip = */ false,
+  };
+
+  int O_c = conv_params.O;
+  int C_c = conv_params.C;
+
+  int N_tiles_n = conv_params.N;
+  int N_tiles_h = (conv_params.oS[0] + 5) / 6;
+  int N_tiles_w = (conv_params.oS[1] + 5) / 6;
+  int N_tiles = N_tiles_n * N_tiles_h * N_tiles_w;
+
+  // Do filter transform
+  Shape filt_wg_shape = {8 * 8, conv_params.C, conv_params.O};
+  array filt_wg(std::move(filt_wg_shape), wt.dtype(), nullptr, {});
+  filt_wg.set_data(allocator::malloc(filt_wg.nbytes()));
+  copies_w.push_back(filt_wg);
+  {
+    int bc = 32;
+    int bo = 4;
+    std::string kname;
+    kname.reserve(32);
+    concatenate(
+        kname,
+        "winograd_conv_2d_weight_transform_",
+        type_to_name(out),
+        "_bc",
+        bc);
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kname);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    compute_encoder.set_input_array(wt, 0);
+    compute_encoder.set_output_array(filt_wg, 1);
+
+    compute_encoder.set_bytes(C_c, 2);
+    compute_encoder.set_bytes(O_c, 3);
+
+    MTL::Size group_dims = MTL::Size(32, bo, 1);
+    MTL::Size grid_dims = MTL::Size(O_c / bo, 1, 1);
+
+    compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+  }
+
+  // Do input transform
+  Shape inp_wg_shape = {8 * 8, N_tiles, conv_params.C};
+  array inp_wg(std::move(inp_wg_shape), in.dtype(), nullptr, {});
+  inp_wg.set_data(allocator::malloc(inp_wg.nbytes()));
+  copies_w.push_back(inp_wg);
+  {
+    int bc = 32;
+    int wm = 2;
+    int wn = 2;
+    std::string kname;
+    kname.reserve(32);
+    concatenate(
+        kname,
+        "winograd_conv_2d_input_transform_",
+        type_to_name(out),
+        "_bc",
+        bc);
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kname);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    compute_encoder.set_input_array(in_padded, 0);
+    compute_encoder.set_output_array(inp_wg, 1);
+
+    compute_encoder.set_bytes(conv_params_updated, 2);
+
+    MTL::Size group_dims = MTL::Size(32, wn, wm);
+    MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n);
+
+    compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+  }
+
+  // Do batched gemm
+  Shape out_wg_shape = {8 * 8, N_tiles, conv_params.O};
+  array out_wg(std::move(out_wg_shape), in.dtype(), nullptr, {});
+  out_wg.set_data(allocator::malloc(out_wg.nbytes()));
+  copies_w.push_back(out_wg);
+  {
+    std::vector<array> empty_copies;
+    steel_matmul(
+        s,
+        d,
+        /*a = */ inp_wg,
+        /*b = */ filt_wg,
+        /*c = */ out_wg,
+        /*M = */ N_tiles,
+        /*N = */ conv_params.O,
+        /*K = */ conv_params.C,
+        /*batch_size_out = */ 8 * 8,
+        /*a_cols = */ conv_params.C,
+        /*b_cols = */ conv_params.O,
+        /*a_transposed = */ false,
+        /*b_transposed = */ false,
+        /*copies = */ empty_copies);
+  }
+
+  // Do output transform
+  {
+    int bc = 32;
+    int wm = 2;
+    int wn = 2;
+    std::string kname;
+    kname.reserve(32);
+    concatenate(
+        kname,
+        "winograd_conv_2d_output_transform_",
+        type_to_name(out),
+        "_bo",
+        bc);
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kname);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    compute_encoder.set_input_array(out_wg, 0);
+    compute_encoder.set_output_array(out, 1);
+
+    compute_encoder.set_bytes(conv_params_updated, 2);
+
+    MTL::Size group_dims = MTL::Size(32, wn, wm);
+    MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n);
+
+    compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+  }
+}
+
+void depthwise_conv_2D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const MLXConvParams<2>& conv_params) {
+  std::string base_name;
+  base_name.reserve(32);
+  concatenate(base_name, "depthwise_conv_2d_", type_to_name(out));
+
+  const int N = conv_params.N;
+  const int ker_h = conv_params.wS[0];
+  const int ker_w = conv_params.wS[1];
+  const int str_h = conv_params.str[0];
+  const int str_w = conv_params.str[1];
+  const int tc = 8;
+  const int tw = 8;
+  const int th = 4;
+  const bool do_flip = conv_params.flip;
+
+  metal::MTLFCList func_consts = {
+      {&ker_h, MTL::DataType::DataTypeInt, 00},
+      {&ker_w, MTL::DataType::DataTypeInt, 01},
+      {&str_h, MTL::DataType::DataTypeInt, 10},
+      {&str_w, MTL::DataType::DataTypeInt, 11},
+      {&th, MTL::DataType::DataTypeInt, 100},
+      {&tw, MTL::DataType::DataTypeInt, 101},
+      {&do_flip, MTL::DataType::DataTypeBool, 200},
+  };
+
+  // clang-format off
+  std::string hash_name;
+  hash_name.reserve(64);
+  concatenate(
+      hash_name,
+      base_name,
+  "_ker_h_", ker_h,
+  "_ker_w_", ker_w,
+  "_str_h_", str_h,
+  "_str_w_", str_w,
+  "_tgp_h_", th,
+  "_tgp_w_", tw,
+  "_do_flip_", do_flip ? 't' : 'n'); // clang-format on
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(base_name, hash_name, func_consts);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_input_array(wt, 1);
+  compute_encoder.set_output_array(out, 2);
+
+  compute_encoder.set_bytes(conv_params, 3);
+
+  MTL::Size group_dims = MTL::Size(tc, tw, th);
+  MTL::Size grid_dims = MTL::Size(
+      conv_params.C / tc, conv_params.oS[1] / tw, (conv_params.oS[0] / th) * N);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void dispatch_conv_2D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const MLXConvParams<2>& conv_params,
+    std::vector<array>& copies) {
+  bool is_stride_one = conv_params.str[0] == 1 && conv_params.str[1] == 1;
+  bool is_kdil_one = conv_params.kdil[0] == 1 && conv_params.kdil[1] == 1;
+  bool is_idil_one = conv_params.idil[0] == 1 && conv_params.idil[1] == 1;
+
+  if (is_idil_one && conv_params.groups > 1) {
+    const int C_per_group = conv_params.C / conv_params.groups;
+    const int O_per_group = conv_params.O / conv_params.groups;
+
+    if (C_per_group == 1 && O_per_group == 1 && is_kdil_one &&
+        conv_params.wS[0] <= 7 && conv_params.wS[1] <= 7 &&
+        conv_params.str[0] <= 2 && conv_params.str[1] <= 2 &&
+        conv_params.oS[0] % 8 == 0 && conv_params.oS[1] % 8 == 0 &&
+        conv_params.wt_strides[1] == conv_params.wS[1] &&
+        conv_params.C % 16 == 0 && conv_params.C == conv_params.O) {
+      return depthwise_conv_2D_gpu(s, d, in, wt, out, conv_params);
+    }
+
+    if ((C_per_group <= 4 || C_per_group % 16 == 0) &&
+        (O_per_group <= 16 || O_per_group % 16 == 0)) {
+      return implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params);
+    } else {
+      return explicit_gemm_conv_group_ND_gpu(s, d, in, wt, out, conv_params);
+    }
+  }
+
+  // Direct to winograd conv
+  bool inp_large =
+      (conv_params.N * conv_params.iS[0] * conv_params.iS[1]) >= 4096;
+  bool channels_large = (conv_params.C + conv_params.O) >= 256;
+  bool out_large =
+      (conv_params.N * conv_params.oS[0] * conv_params.oS[1]) >= 256;
+  if (!conv_params.flip && is_stride_one && is_kdil_one && is_idil_one &&
+      conv_params.wS[0] == 3 && conv_params.wS[1] == 3 &&
+      conv_params.C % 32 == 0 && conv_params.O % 32 == 0 && inp_large &&
+      channels_large) {
+    return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
+  }
+
+  // Direct to implicit gemm conv
+  if (is_idil_one && (conv_params.C <= 4 || conv_params.C % 16 == 0) &&
+      (conv_params.O <= 16 || conv_params.O % 16 == 0)) {
+    return implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params);
+  }
+
+  else if ((conv_params.C % 16 == 0 && conv_params.O % 16 == 0) || out_large) {
+    return implicit_gemm_conv_2D_general_gpu(s, d, in, wt, out, conv_params);
+  }
+
+  // Direct to explicit gemm conv
+  else {
+    return explicit_gemm_conv_ND_gpu(s, d, in, wt, out, conv_params);
+  }
+}
+
+void depthwise_conv_1D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    array wt,
+    array out) {
+  bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX;
+  std::string base_name;
+  base_name.reserve(32);
+  concatenate(
+      base_name,
+      "depthwise_conv_1d_",
+      large ? "_large" : "",
+      type_to_name(out));
+
+  if (!wt.flags().row_contiguous) {
+    wt = contiguous_copy_gpu(wt, s);
+    d.add_temporary(wt, s.index);
+  }
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(base_name);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  auto B = in.shape(0);
+  auto Tout = out.shape(1);
+  auto D = in.shape(2);
+  auto K = wt.shape(1);
+
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_input_array(wt, 1);
+  compute_encoder.set_output_array(out, 2);
+  if (large) {
+    int64_t strides[3] = {in.strides(0), in.strides(1), in.strides(2)};
+    compute_encoder.set_bytes(strides, 3, 3);
+
+  } else {
+    int strides[3] = {
+        static_cast<int>(in.strides(0)),
+        static_cast<int>(in.strides(1)),
+        static_cast<int>(in.strides(2))};
+    compute_encoder.set_bytes(strides, 3, 3);
+  }
+
+  compute_encoder.set_bytes(K, 4);
+  auto group_dims = get_block_dims(D, Tout, B);
+  MTL::Size grid_dims = MTL::Size(D, Tout, B);
+
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void conv_1D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    int groups,
+    bool flip,
+    std::vector<array>& copies) {
+  bool is_idil_one = in_dilation[0] == 1;
+  int C = in.shape(2);
+  int O = wt.shape(0);
+  // Fast path for fully separable 1D convolution
+  if (is_idil_one && (groups == C) && groups == O && wt_strides[0] == 1 &&
+      wt_dilation[0] == 1 && padding[0] == 0 && !flip) {
+    depthwise_conv_1D_gpu(s, d, in, wt, out);
+    return;
+  }
+
+  const int C_per_group = C / groups;
+  const int O_per_group = O / groups;
+
+  // Direct to implicit gemm conv
+  if (is_idil_one && (C_per_group <= 4 || C_per_group % 16 == 0) &&
+      (O_per_group <= 16 || O_per_group % 16 == 0)) {
+    MLXConvParams<2> conv_params{
+        /* const int  N = */ static_cast<int>(in.shape(0)),
+        /* const int  C = */ C,
+        /* const int  O = */ O,
+        /* const int iS[NDIM] = */ {static_cast<int>(in.shape(1)), 1},
+        /* const int wS[NDIM] = */ {static_cast<int>(wt.shape(1)), 1},
+        /* const int oS[NDIM] = */ {static_cast<int>(out.shape(1)), 1},
+        /* const int str[NDIM] = */ {wt_strides[0], 1},
+        /* const int pad[NDIM] = */ {padding[0], 0},
+        /* const int kdil[NDIM] = */ {wt_dilation[0], 1},
+        /* const int idil[NDIM] = */ {in_dilation[0], 1},
+        /* const size_t in_strides[NDIM + 2] = */
+        {in.strides()[0], in.strides()[1], 0, in.strides()[2]},
+        /* const size_t wt_strides[NDIM + 2] = */
+        {wt.strides()[0], wt.strides()[1], 0, wt.strides()[2]},
+        /* const size_t out_strides[NDIM + 2] = */
+        {out.strides()[0], out.strides()[1], 0, out.strides()[2]},
+        /* const int groups = */ groups,
+        /* const bool flip = */ flip};
+
+    dispatch_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
+    return;
+  }
+
+  // Make conv params
+  MLXConvParams<1> conv_params{
+      /* const int  N = */ static_cast<int>(in.shape(0)),
+      /* const int  C = */ static_cast<int>(in.shape(2)),
+      /* const int  O = */ static_cast<int>(wt.shape(0)),
+      /* const int iS[NDIM] = */ {static_cast<int>(in.shape(1))},
+      /* const int wS[NDIM] = */ {static_cast<int>(wt.shape(1))},
+      /* const int oS[NDIM] = */ {static_cast<int>(out.shape(1))},
+      /* const int str[NDIM] = */ {wt_strides[0]},
+      /* const int pad[NDIM] = */ {padding[0]},
+      /* const int kdil[NDIM] = */ {wt_dilation[0]},
+      /* const int idil[NDIM] = */ {in_dilation[0]},
+      /* const size_t in_strides[NDIM + 2] = */
+      {in.strides()[0], in.strides()[1], in.strides()[2]},
+      /* const size_t wt_strides[NDIM + 2] = */
+      {wt.strides()[0], wt.strides()[1], wt.strides()[2]},
+      /* const size_t out_strides[NDIM + 2] = */
+      {out.strides()[0], out.strides()[1], out.strides()[2]},
+      /* const int groups = */ groups,
+      /* const bool flip = */ flip};
+
+  // Direct to explicit gemm conv
+  if (groups > 1) {
+    return explicit_gemm_conv_group_ND_gpu(s, d, in, wt, out, conv_params);
+  } else {
+    return explicit_gemm_conv_ND_gpu(s, d, in, wt, out, conv_params);
+  }
+}
+
+void conv_2D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    const int groups,
+    bool flip,
+    std::vector<array>& copies) {
+  // Make conv params
+  MLXConvParams<2> conv_params{
+      /* const int  N = */ static_cast<int>(in.shape(0)),
+      /* const int  C = */ static_cast<int>(in.shape(3)),
+      /* const int  O = */ static_cast<int>(wt.shape(0)),
+      /* const int iS[NDIM] = */
+      {static_cast<int>(in.shape(1)), static_cast<int>(in.shape(2))},
+      /* const int wS[NDIM] = */
+      {static_cast<int>(wt.shape(1)), static_cast<int>(wt.shape(2))},
+      /* const int oS[NDIM] = */
+      {static_cast<int>(out.shape(1)), static_cast<int>(out.shape(2))},
+      /* const int str[NDIM] = */ {wt_strides[0], wt_strides[1]},
+      /* const int pad[NDIM] = */ {padding[0], padding[1]},
+      /* const int kdil[NDIM] = */ {wt_dilation[0], wt_dilation[1]},
+      /* const int idil[NDIM] = */ {in_dilation[0], in_dilation[1]},
+      /* const size_t in_strides[NDIM + 2] = */
+      {in.strides(0), in.strides(1), in.strides(2), in.strides(3)},
+      /* const size_t wt_strides[NDIM + 2] = */
+      {wt.strides(0), wt.strides(1), wt.strides(2), wt.strides(3)},
+      /* const size_t out_strides[NDIM + 2] = */
+      {out.strides(0), out.strides(1), out.strides(2), out.strides(3)},
+      /* const int groups = */ groups,
+      /* const bool flip = */ flip,
+  };
+  dispatch_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
+}
+
+void conv_3D_gpu(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    const array& wt,
+    array out,
+    const std::vector<int>& padding,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
+    std::vector<array>& copies) {
+  // Make conv params
+  MLXConvParams<3> conv_params{
+      /* const int  N = */ static_cast<int>(in.shape(0)),
+      /* const int  C = */ static_cast<int>(in.shape(4)),
+      /* const int  O = */ static_cast<int>(wt.shape(0)),
+      /* const int iS[NDIM] = */
+      {static_cast<int>(in.shape(1)),
+       static_cast<int>(in.shape(2)),
+       static_cast<int>(in.shape(3))},
+      /* const int wS[NDIM] = */
+      {static_cast<int>(wt.shape(1)),
+       static_cast<int>(wt.shape(2)),
+       static_cast<int>(wt.shape(3))},
+      /* const int oS[NDIM] = */
+      {static_cast<int>(out.shape(1)),
+       static_cast<int>(out.shape(2)),
+       static_cast<int>(out.shape(3))},
+      /* const int str[NDIM] = */ {wt_strides[0], wt_strides[1], wt_strides[2]},
+      /* const int pad[NDIM] = */ {padding[0], padding[1], padding[2]},
+      /* const int kdil[NDIM] = */
+      {wt_dilation[0], wt_dilation[1], wt_dilation[2]},
+      /* const int idil[NDIM] = */
+      {in_dilation[0], in_dilation[1], in_dilation[2]},
+      /* const size_t in_strides[NDIM + 2] = */
+      {in.strides()[0],
+       in.strides()[1],
+       in.strides()[2],
+       in.strides()[3],
+       in.strides()[4]},
+      /* const size_t wt_strides[NDIM + 2] = */
+      {wt.strides()[0],
+       wt.strides()[1],
+       wt.strides()[2],
+       wt.strides()[3],
+       wt.strides()[4]},
+      /* const size_t out_strides[NDIM + 2] = */
+      {out.strides()[0],
+       out.strides()[1],
+       out.strides()[2],
+       out.strides()[3],
+       out.strides()[4]},
+      /* const int groups = */ 1,
+      /* const bool flip = */ flip,
+  };
+  return explicit_gemm_conv_ND_gpu(s, d, in, wt, out, conv_params);
+}
+
+} // namespace
+
+void Convolution::eval_gpu(const std::vector<array>& inputs, array& out) {
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  // Ensure contiguity
+  std::vector<array> copies;
+  auto in = inputs[0];
+  auto wt = inputs[1];
+  if (!in.flags().row_contiguous) {
+    in = contiguous_copy_gpu(in, s);
+    copies.push_back(in);
+  }
+  if (!wt.flags().row_contiguous) {
+    wt = contiguous_copy_gpu(wt, s);
+    copies.push_back(wt);
+  }
+
+  // 3D conv
+  if (out.ndim() == 5) {
+    conv_3D_gpu(
+        s,
+        d,
+        in,
+        wt,
+        out,
+        padding_lo_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_,
+        copies);
+  }
+  // 2D conv
+  else if (out.ndim() == 4) {
+    conv_2D_gpu(
+        s,
+        d,
+        in,
+        wt,
+        out,
+        padding_lo_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        groups_,
+        flip_,
+        copies);
+  }
+  // 1D conv
+  else if (out.ndim() == 3) {
+    conv_1D_gpu(
+        s,
+        d,
+        in,
+        wt,
+        out,
+        padding_lo_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        groups_,
+        flip_,
+        copies);
+  }
+  // Throw error
+  else {
+    throw std::invalid_argument(
+        "[Convolution::eval_gpu] Only supports 1D, 2D or 3D convolutions.");
+  }
+
+  // Record copies
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/copy.cpp b/macos/third_party/mlx/mlx/backend/metal/copy.cpp
new file mode 100644
index 0000000..6b79128
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/copy.cpp
@@ -0,0 +1,236 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+
+namespace mlx::core {
+
+constexpr int MAX_COPY_SPECIALIZED_DIMS = 3;
+
+void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
+  bool donated = set_copy_output_data(in, out, ctype);
+  if (donated && in.dtype() == out.dtype()) {
+    // If the output has the same type as the input then there is nothing to
+    // copy, just use the buffer.
+    return;
+  }
+  if (ctype == CopyType::GeneralGeneral) {
+    ctype = CopyType::General;
+  }
+  copy_gpu_inplace(in, out, ctype, s);
+}
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const Shape& data_shape,
+    const Strides& strides_in_pre,
+    const Strides& strides_out_pre,
+    int64_t inp_offset,
+    int64_t out_offset,
+    CopyType ctype,
+    const Stream& s,
+    std::optional<array> dynamic_i_offset /* = std::nullopt */,
+    std::optional<array> dynamic_o_offset /* = std::nullopt */) {
+  if (out.size() == 0) {
+    return;
+  }
+
+  // Try to collapse contiguous dims
+  auto maybe_collapse =
+      [ctype, &data_shape, &strides_in_pre, &strides_out_pre]() {
+        if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+          auto [shape, strides] = collapse_contiguous_dims(
+              data_shape,
+              std::vector{strides_in_pre, strides_out_pre},
+              /* size_cap = */ INT32_MAX);
+          return std::make_tuple(shape, strides[0], strides[1]);
+        } else {
+          Strides e{};
+          return std::make_tuple(Shape{}, e, e);
+        }
+      };
+  auto [shape, strides_in_, strides_out_] = maybe_collapse();
+  int ndim = shape.size();
+  bool large;
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    // Allow for negative strides
+    large = in.data_size() > INT32_MAX || out.data_size() > INT32_MAX;
+  } else {
+    large = out.data_size() > UINT32_MAX;
+  }
+  bool dynamic = dynamic_i_offset || dynamic_o_offset;
+  auto& d = metal::device(s.device);
+  int work_per_thread = 1;
+  std::string kernel_name;
+  switch (ctype) {
+    case CopyType::Scalar:
+      kernel_name = large ? "s2" : "s";
+      break;
+    case CopyType::Vector:
+      kernel_name = large ? "v2" : "v";
+      break;
+    case CopyType::General:
+      kernel_name = "g";
+      break;
+    case CopyType::GeneralGeneral:
+      kernel_name = "gg";
+      break;
+  }
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    if (shape.size() <= MAX_COPY_SPECIALIZED_DIMS) {
+      kernel_name += std::to_string(shape.size());
+    } else {
+      work_per_thread = large ? 4 : 2;
+      concatenate(kernel_name, "n", std::to_string(work_per_thread));
+    }
+    if (large) {
+      kernel_name += "large";
+    }
+    if (dynamic) {
+      kernel_name += "_dynamic";
+      if (ctype != CopyType::GeneralGeneral) {
+        throw std::runtime_error(
+            "[Copy::eval_gpu] Dynamic output offset requires GeneralGeneral copy");
+      }
+    }
+  } else {
+    work_per_thread = get_work_per_thread(out.dtype(), out.data_size());
+    if (!large && work_per_thread > 1) {
+      kernel_name += "n";
+    }
+  }
+  concatenate(kernel_name, "_copy", type_to_name(in), type_to_name(out));
+  auto kernel = dynamic ? get_dynamic_copy_kernel(d, kernel_name, in, out)
+                        : get_copy_kernel(d, kernel_name, in, out);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  bool donate_in = in.data_shared_ptr() == nullptr;
+
+  inp_offset *= size_of(in.dtype());
+  out_offset *= size_of(out.dtype());
+
+  compute_encoder.set_input_array(donate_in ? out : in, 0, inp_offset);
+  compute_encoder.set_output_array(out, 1, out_offset);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    Strides strides_in{strides_in_.begin(), strides_in_.end()};
+    Strides strides_out{strides_out_.begin(), strides_out_.end()};
+    if (ndim > 3) {
+      compute_encoder.set_vector_bytes(shape, ndim, 2);
+    }
+    compute_encoder.set_vector_bytes(strides_in, ndim, 3);
+    if (ctype == CopyType::GeneralGeneral) {
+      compute_encoder.set_vector_bytes(strides_out, ndim, 4);
+    }
+
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+
+    size_t data_size = 1;
+    for (auto& s : shape)
+      data_size *= s;
+    size_t rest = data_size / (dim0 * dim1);
+
+    if (ndim > MAX_COPY_SPECIALIZED_DIMS) {
+      compute_encoder.set_bytes(ndim, 5);
+      dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    }
+    if (dynamic) {
+      if (dynamic_i_offset) {
+        compute_encoder.set_input_array(*dynamic_i_offset, 6);
+      } else {
+        compute_encoder.set_bytes(0ll, 6);
+      }
+      if (dynamic_o_offset) {
+        compute_encoder.set_input_array(*dynamic_o_offset, 7);
+      } else {
+        compute_encoder.set_bytes(0ll, 7);
+      }
+    }
+
+    // NB assuming thread_group_size is a power of 2 larger than 32 x 32
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::copy] Must use 1024 sized block");
+    }
+
+    auto group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    size_t nthreads = ceildiv(out.data_size(), work_per_thread);
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    MTL::Size grid_dims;
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(out.data_size(), 2);
+      grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+    } else {
+      compute_encoder.set_bytes<int>(out.data_size(), 2);
+      grid_dims = MTL::Size(nthreads, 1, 1);
+    }
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void fill_gpu(const array& val, array& out, const Stream& s) {
+  if (out.size() == 0) {
+    return;
+  }
+  out.set_data(allocator::malloc(out.nbytes()));
+  bool large = out.data_size() > UINT32_MAX;
+  int work_per_thread = get_work_per_thread(out.dtype(), out.data_size());
+  auto& d = metal::device(s.device);
+  std::string kernel_name = large ? "s2" : (work_per_thread > 1 ? "sn" : "s");
+  concatenate(kernel_name, "_copy", type_to_name(val), type_to_name(out));
+  auto kernel = get_copy_kernel(d, kernel_name, val, out);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  compute_encoder.set_input_array(val, 0);
+  compute_encoder.set_output_array(out, 1);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  size_t nthreads = ceildiv(out.data_size(), work_per_thread);
+  if (thread_group_size > nthreads) {
+    thread_group_size = nthreads;
+  }
+  MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+  MTL::Size grid_dims;
+  if (large) {
+    compute_encoder.set_bytes<int64_t>(out.data_size(), 2);
+    grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+  } else {
+    compute_encoder.set_bytes<int>(out.data_size(), 2);
+    grid_dims = MTL::Size(nthreads, 1, 1);
+  }
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void reshape_gpu(const array& in, array& out, Stream s) {
+  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
+  if (copy_necessary) {
+    out.set_data(allocator::malloc(out.nbytes()));
+    copy_gpu_inplace(
+        in,
+        out,
+        in.shape(),
+        in.strides(),
+        make_contiguous_strides(in.shape()),
+        0,
+        0,
+        CopyType::General,
+        s);
+  } else {
+    shared_buffer_reshape(in, out_strides, out);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/custom_kernel.cpp b/macos/third_party/mlx/mlx/backend/metal/custom_kernel.cpp
new file mode 100644
index 0000000..eec6645
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/custom_kernel.cpp
@@ -0,0 +1,430 @@
+// Copyright © 2024 Apple Inc.
+
+#include <iostream>
+#include <regex>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/jit/includes.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fast {
+
+struct CustomKernelCache {
+  std::unordered_map<std::string, std::string> libraries;
+};
+
+static CustomKernelCache& cache() {
+  static CustomKernelCache cache_;
+  return cache_;
+};
+
+std::string write_signature(
+    std::string func_name,
+    const std::string& header,
+    const std::string& source,
+    const std::vector<std::string>& input_names,
+    const std::vector<array>& inputs,
+    const std::vector<std::string>& output_names,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args,
+    const std::vector<std::string>& attributes,
+    const std::vector<std::tuple<bool, bool, bool>>& shape_infos,
+    bool atomic_outputs) {
+  std::string kernel_source;
+  kernel_source.reserve(header.size() + source.size() + 16384);
+  kernel_source += header;
+  // Auto-generate a function signature based on `template_args`
+  // and the dtype/shape of the arrays passed as `inputs`.
+  if (!template_args.empty()) {
+    kernel_source += "template <";
+    int i = 0;
+    for (const auto& [name, arg] : template_args) {
+      std::string param_type;
+      if (std::holds_alternative<int>(arg)) {
+        param_type = "int";
+      } else if (std::holds_alternative<bool>(arg)) {
+        param_type = "bool";
+      } else if (std::holds_alternative<Dtype>(arg)) {
+        param_type = "typename";
+      }
+      if (i > 0) {
+        kernel_source += ", ";
+      }
+      kernel_source += param_type;
+      kernel_source += " ";
+      kernel_source += name;
+      i++;
+    }
+    kernel_source += ">\n";
+  }
+  kernel_source += "[[kernel]] void ";
+  kernel_source += func_name;
+  kernel_source += "(\n";
+
+  int index = 0;
+  constexpr int max_constant_array_size = 8;
+  // Add inputs
+  for (int i = 0; i < inputs.size(); ++i) {
+    const auto& name = input_names[i];
+    const auto& arr = inputs[i];
+    auto dtype = get_type_string(arr.dtype());
+    std::string location =
+        arr.size() < max_constant_array_size ? "constant" : "device";
+    std::string ref = arr.ndim() == 0 ? "&" : "*";
+    kernel_source += "  const ";
+    kernel_source += location;
+    kernel_source += " ";
+    kernel_source += dtype;
+    kernel_source += ref;
+    kernel_source += " ";
+    kernel_source += name;
+    kernel_source += " [[buffer(";
+    kernel_source += std::to_string(index);
+    kernel_source += ")]],\n";
+    index++;
+    // Add input shape, strides and ndim if present in the source
+    if (arr.ndim() > 0) {
+      if (std::get<0>(shape_infos[i])) {
+        kernel_source +=
+            ("  const constant int* " + name + "_shape [[buffer(" +
+             std::to_string(index) + ")]],\n");
+        index++;
+      }
+      if (std::get<1>(shape_infos[i])) {
+        kernel_source +=
+            ("  const constant int64_t* " + name + "_strides [[buffer(" +
+             std::to_string(index) + ")]],\n");
+        index++;
+      }
+      if (std::get<2>(shape_infos[i])) {
+        kernel_source +=
+            ("  const constant int& " + name + "_ndim [[buffer(" +
+             std::to_string(index) + ")]],\n");
+        index++;
+      }
+    }
+  }
+  // Add outputs
+  for (int i = 0; i < output_names.size(); ++i) {
+    const auto& name = output_names[i];
+    const auto& dtype = output_dtypes[i];
+    kernel_source += "  device ";
+    auto type_string = get_type_string(dtype);
+    if (atomic_outputs) {
+      kernel_source += "atomic<";
+    }
+    kernel_source += type_string;
+    if (atomic_outputs) {
+      kernel_source += ">";
+    }
+    kernel_source += "* ";
+    kernel_source += name;
+    kernel_source += " [[buffer(";
+    kernel_source += std::to_string(index);
+    kernel_source += ")]]";
+    if (index < inputs.size() + output_names.size() - 1 ||
+        attributes.size() > 0) {
+      kernel_source += ",\n";
+    } else {
+      kernel_source += ") {\n";
+    }
+    index++;
+  }
+
+  index = 0;
+  for (const auto& attr : attributes) {
+    kernel_source += attr;
+    if (index < attributes.size() - 1) {
+      kernel_source += ",\n";
+    } else {
+      kernel_source += ") {\n";
+    }
+    index++;
+  }
+  kernel_source += source;
+  kernel_source += "\n}\n";
+  return kernel_source;
+}
+
+std::string write_template(
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args) {
+  std::ostringstream template_def;
+  template_def << "<";
+  int i = 0;
+  for (const auto& [name, arg] : template_args) {
+    if (i > 0) {
+      template_def << ", ";
+    }
+    if (std::holds_alternative<int>(arg)) {
+      template_def << std::get<int>(arg);
+    } else if (std::holds_alternative<bool>(arg)) {
+      template_def << std::get<bool>(arg);
+    } else if (std::holds_alternative<Dtype>(arg)) {
+      template_def << get_type_string(std::get<Dtype>(arg));
+    }
+    i++;
+  }
+  template_def << ">";
+  return template_def.str();
+}
+
+CustomKernelFunction metal_kernel(
+    const std::string& name,
+    const std::vector<std::string>& input_names,
+    const std::vector<std::string>& output_names,
+    const std::string& source,
+    const std::string& header /* = "" */,
+    bool ensure_row_contiguous /* = true */,
+    bool atomic_outputs /* = false */) {
+  if (output_names.empty()) {
+    throw std::invalid_argument(
+        "[metal_kernel] Must specify at least one output.");
+  }
+  std::vector<std::tuple<bool, bool, bool>> shape_infos;
+  for (auto& n : input_names) {
+    std::tuple<bool, bool, bool> shape_info;
+    std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos;
+    std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos;
+    std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos;
+    shape_infos.push_back(shape_info);
+  }
+  const std::vector<std::pair<std::string, std::string>> metal_attributes = {
+      {"dispatch_quadgroups_per_threadgroup", "uint"},
+      {"dispatch_simdgroups_per_threadgroup", "uint"},
+      {"dispatch_threads_per_threadgroup", "uint3"},
+      {"grid_origin", "uint3"},
+      {"grid_size", "uint3"},
+      {"quadgroup_index_in_threadgroup", "uint"},
+      {"quadgroups_per_threadgroup", "uint"},
+      {"simdgroup_index_in_threadgroup", "uint"},
+      {"simdgroups_per_threadgroup", "uint"},
+      {"thread_execution_width", "uint"},
+      {"thread_index_in_quadgroup", "uint"},
+      {"thread_index_in_simdgroup", "uint"},
+      {"thread_index_in_threadgroup", "uint"},
+      {"thread_position_in_grid", "uint3"},
+      {"thread_position_in_threadgroup", "uint3"},
+      {"threadgroup_position_in_grid", "uint3"},
+      {"threadgroups_per_grid", "uint3"},
+      {"threads_per_grid", "uint3"},
+      {"threads_per_simdgroup", "uint"},
+      {"threads_per_threadgroup", "uint3"},
+  };
+
+  std::vector<std::string> attributes;
+  for (const auto& [attr, dtype] : metal_attributes) {
+    if (source.find(attr) != std::string::npos) {
+      attributes.push_back("  " + dtype + " " + attr + " [[" + attr + "]]");
+    }
+  }
+
+  return [=,
+          shape_infos = std::move(shape_infos),
+          attributes = std::move(attributes)](
+             const std::vector<array>& inputs,
+             const std::vector<Shape>& output_shapes,
+             const std::vector<Dtype>& output_dtypes,
+             std::tuple<int, int, int> grid,
+             std::tuple<int, int, int> threadgroup,
+             const std::vector<std::pair<std::string, TemplateArg>>&
+                 template_args = {},
+             std::optional<float> init_value = std::nullopt,
+             bool verbose = false,
+             StreamOrDevice s_ = {}) {
+    if (inputs.size() != input_names.size()) {
+      std::ostringstream msg;
+      msg << "[metal_kernel] Expected `inputs` to have size "
+          << input_names.size() << " but got size " << inputs.size() << "."
+          << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_shapes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[metal_kernel] Expected `output_shapes` to have size "
+          << output_names.size() << " but got size " << output_shapes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_dtypes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[metal_kernel] Expected `output_dtypes` to have size "
+          << output_names.size() << " but got size " << output_dtypes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+
+    auto s = to_stream(s_);
+    if (s.device != Device::gpu) {
+      throw std::invalid_argument("[metal_kernel] Only supports the GPU.");
+    }
+
+    std::string kernel_name = "custom_kernel_" + name;
+    std::string template_def = "";
+    if (!template_args.empty()) {
+      std::regex disallowed_chars("\\<|\\>|(, )");
+      template_def = write_template(template_args);
+      auto template_hash =
+          std::regex_replace(template_def, disallowed_chars, "_");
+      template_hash.pop_back();
+      kernel_name += "_";
+      kernel_name += template_hash;
+    }
+
+    std::string kernel_source = write_signature(
+        kernel_name,
+        header,
+        source,
+        input_names,
+        inputs,
+        output_names,
+        output_dtypes,
+        template_args,
+        attributes,
+        shape_infos,
+        atomic_outputs);
+
+    if (!template_args.empty()) {
+      template_def = kernel_name + template_def;
+      kernel_source += "\ntemplate [[host_name(\"";
+      kernel_source += kernel_name;
+      kernel_source += "\")]] [[kernel]] decltype(";
+      kernel_source += template_def;
+      kernel_source += ") ";
+      kernel_source += template_def;
+      kernel_source += ";\n";
+    }
+
+    if (verbose) {
+      std::cout << "Generated source code for `" << name << "`:" << std::endl
+                << "```" << std::endl
+                << kernel_source << std::endl
+                << "```" << std::endl;
+    }
+
+    return array::make_arrays(
+        std::move(output_shapes),
+        std::move(output_dtypes),
+        std::make_shared<CustomKernel>(
+            s,
+            std::move(kernel_name),
+            std::move(kernel_source),
+            grid,
+            threadgroup,
+            shape_infos,
+            ensure_row_contiguous,
+            init_value,
+            std::vector<ScalarArg>{},
+            false,
+            0),
+        std::move(inputs));
+  };
+}
+
+void CustomKernel::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  // silence some warnings
+  (void)is_precompiled_;
+  (void)shared_memory_;
+
+  auto& s = stream();
+
+  std::vector<array> copies;
+
+  for (auto& out : outputs) {
+    if (init_value_) {
+      copies.emplace_back(init_value_.value(), out.dtype());
+      fill_gpu(copies.back(), out, s);
+    } else {
+      out.set_data(allocator::malloc(out.nbytes()));
+    }
+  }
+
+  auto check_input = [&copies, &s, this](const array& x) -> const array {
+    bool no_copy = x.flags().row_contiguous;
+    if (!ensure_row_contiguous_ || no_copy) {
+      return x;
+    } else {
+      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
+      copy_gpu(x, copies.back(), CopyType::General, s);
+      return copies.back();
+    }
+  };
+  std::vector<array> checked_inputs;
+  for (const array& in : inputs) {
+    checked_inputs.push_back(check_input(in));
+  }
+
+  auto& d = metal::device(s.device);
+
+  {
+    // Clear kernels from the device library cache if needed
+    auto& kernel_cache = cache();
+    if (auto it = kernel_cache.libraries.find(name_);
+        it != kernel_cache.libraries.end()) {
+      if (it->second != source_) {
+        auto& d = metal::device(s.device);
+        d.clear_library(name_);
+        it->second = source_;
+      }
+    } else {
+      kernel_cache.libraries.emplace(name_, source_);
+    }
+  }
+
+  auto lib = d.get_library(name_, [this] { return metal::utils() + source_; });
+  auto kernel = d.get_kernel(name_, lib);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  int index = 0;
+  for (int i = 0; i < checked_inputs.size(); i++) {
+    const array& in = checked_inputs[i];
+    auto& shape_info = shape_infos_[i];
+    compute_encoder.set_input_array(in, index);
+    index++;
+    if (in.ndim() > 0) {
+      int ndim = in.ndim();
+      if (std::get<0>(shape_info)) {
+        compute_encoder.set_vector_bytes(in.shape(), ndim, index);
+        index++;
+      }
+      if (std::get<1>(shape_info)) {
+        compute_encoder.set_vector_bytes(in.strides(), ndim, index);
+        index++;
+      }
+      if (std::get<2>(shape_info)) {
+        compute_encoder.set_bytes(ndim, index);
+        index++;
+      }
+    }
+  }
+  for (auto& out : outputs) {
+    compute_encoder.set_output_array(out, index);
+    index++;
+  }
+
+  const auto [tx, ty, tz] = threadgroup_;
+  auto tg_size = tx * ty * tz;
+  auto max_tg_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (tg_size > max_tg_size) {
+    std::ostringstream msg;
+    msg << "Thread group size (" << tg_size << ") is greater than "
+        << " the maximum allowed threads per threadgroup (" << max_tg_size
+        << ").";
+    throw std::invalid_argument(msg.str());
+  }
+
+  const auto [gx, gy, gz] = grid_;
+  MTL::Size group_dims =
+      MTL::Size(std::min(tx, gx), std::min(ty, gy), std::min(tz, gz));
+  MTL::Size grid_dims = MTL::Size(gx, gy, gz);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/backend/metal/device.cpp b/macos/third_party/mlx/mlx/backend/metal/device.cpp
new file mode 100644
index 0000000..de61b0f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/device.cpp
@@ -0,0 +1,816 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstdlib>
+#include <sstream>
+
+#define NS_PRIVATE_IMPLEMENTATION
+#define CA_PRIVATE_IMPLEMENTATION
+#define MTL_PRIVATE_IMPLEMENTATION
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::metal {
+
+namespace {
+
+constexpr const char* default_mtllib_path = METAL_PATH;
+
+auto get_metal_version() {
+  auto get_metal_version_ = []() {
+    if (__builtin_available(macOS 26, iOS 26, tvOS 26, visionOS 26, *)) {
+      return MTL::LanguageVersion4_0;
+    } else if (__builtin_available(macOS 15, iOS 18, tvOS 18, visionOS 2, *)) {
+      return MTL::LanguageVersion3_2;
+    } else {
+      return MTL::LanguageVersion3_1;
+    }
+  };
+  static auto metal_version_ = get_metal_version_();
+  return metal_version_;
+}
+
+auto load_device() {
+  auto devices = MTL::CopyAllDevices();
+  auto device = static_cast<MTL::Device*>(devices->object(0))
+      ?: MTL::CreateSystemDefaultDevice();
+  if (!device) {
+    throw std::runtime_error("Failed to load device");
+  }
+  return device;
+}
+std::pair<MTL::Library*, NS::Error*> load_library_from_path(
+    MTL::Device* device,
+    const char* path) {
+  auto library = NS::String::string(path, NS::UTF8StringEncoding);
+  NS::Error* error;
+  auto lib = device->newLibrary(library, &error);
+
+  return std::make_pair(lib, error);
+}
+
+#ifdef SWIFTPM_BUNDLE
+MTL::Library* try_load_bundle(
+    MTL::Device* device,
+    NS::URL* url,
+    const std::string& lib_name) {
+  std::string bundle_path = std::string(url->fileSystemRepresentation()) + "/" +
+      SWIFTPM_BUNDLE + ".bundle";
+  auto bundle = NS::Bundle::alloc()->init(
+      NS::String::string(bundle_path.c_str(), NS::UTF8StringEncoding));
+  if (bundle != nullptr) {
+    std::string resource_path =
+        std::string(bundle->resourceURL()->fileSystemRepresentation()) + "/" +
+        lib_name + ".metallib";
+    auto [lib, error] = load_library_from_path(device, resource_path.c_str());
+    if (lib) {
+      return lib;
+    }
+  }
+  return nullptr;
+}
+
+MTL::Library* try_load_framework(
+    MTL::Device* device,
+    NS::URL* url,
+    const std::string& lib_name) {
+  std::string resource_path = std::string(url->fileSystemRepresentation()) +
+      "/" + lib_name + ".metallib";
+  auto [lib, error] = load_library_from_path(device, resource_path.c_str());
+  if (lib) {
+    return lib;
+  }
+  return nullptr;
+}
+#endif
+
+// Firstly, search for the metallib in the same path as this binary
+std::pair<MTL::Library*, NS::Error*> load_colocated_library(
+    MTL::Device* device,
+    const std::string& relative_path) {
+  auto path = current_binary_dir() / relative_path;
+  if (!path.has_extension()) {
+    path.replace_extension(".metallib");
+  }
+
+  return load_library_from_path(device, path.c_str());
+}
+
+std::pair<MTL::Library*, NS::Error*> load_swiftpm_library(
+    MTL::Device* device,
+    const std::string& lib_name) {
+#ifdef SWIFTPM_BUNDLE
+  MTL::Library* library =
+      try_load_bundle(device, NS::Bundle::mainBundle()->bundleURL(), lib_name);
+  if (library != nullptr) {
+    return {library, nullptr};
+  }
+  auto bundles = NS::Bundle::allBundles();
+  for (int i = 0, c = (int)bundles->count(); i < c; i++) {
+    auto bundle = reinterpret_cast<NS::Bundle*>(bundles->object(i));
+    library = try_load_bundle(device, bundle->resourceURL(), lib_name);
+    if (library != nullptr) {
+      return {library, nullptr};
+    }
+  }
+  // if SWIFTPM_BUNDLE is a framework identifier, try loading from that
+  auto frameworks = NS::Bundle::allFrameworks();
+  for (int i = 0, c = (int)frameworks->count(); i < c; i++) {
+    const auto bundle = reinterpret_cast<NS::Bundle*>(frameworks->object(i));
+    const auto identifier = bundle->bundleIdentifier();
+    if (identifier != nullptr &&
+        !strcmp(identifier->utf8String(), SWIFTPM_BUNDLE)) {
+      library = try_load_framework(device, bundle->resourceURL(), lib_name);
+      if (library != nullptr) {
+        return {library, nullptr};
+      }
+    }
+  }
+#endif
+  return {nullptr, nullptr};
+}
+
+MTL::Library* load_default_library(MTL::Device* device) {
+  NS::Error* error[5];
+  MTL::Library* lib;
+  // First try the colocated mlx.metallib
+  std::tie(lib, error[0]) = load_colocated_library(device, "mlx");
+  if (lib) {
+    return lib;
+  }
+
+  std::tie(lib, error[1]) = load_colocated_library(device, "Resources/mlx");
+  if (lib) {
+    return lib;
+  }
+
+  // Then try default.metallib in a SwiftPM bundle if we have one
+  std::tie(lib, error[2]) = load_swiftpm_library(device, "default");
+  if (lib) {
+    return lib;
+  }
+
+  // Try lo load resources from Framework resources if SwiftPM wrapped as a
+  // dynamic framework.
+  std::tie(lib, error[3]) = load_colocated_library(device, "Resources/default");
+  if (lib) {
+    return lib;
+  }
+
+  // Finally try default_mtllib_path
+  std::tie(lib, error[4]) = load_library_from_path(device, default_mtllib_path);
+  if (!lib) {
+    std::ostringstream msg;
+    msg << "Failed to load the default metallib. ";
+    for (int i = 0; i < 5; i++) {
+      if (error[i] != nullptr) {
+        msg << error[i]->localizedDescription()->utf8String() << " ";
+      }
+    }
+    throw std::runtime_error(msg.str());
+  }
+  return lib;
+}
+
+MTL::Library* load_library(
+    MTL::Device* device,
+    const std::string& lib_name,
+    const std::string& lib_path) {
+  // We have been given a path that ends in metallib so try to load it
+  if (lib_path.size() > 9 &&
+      std::equal(lib_path.end() - 9, lib_path.end(), ".metallib")) {
+    auto [lib, error] = load_library_from_path(device, lib_path.c_str());
+    if (!lib) {
+      std::ostringstream msg;
+      msg << "Failed to load the metallib from <" << lib_path << "> with error "
+          << error->localizedDescription()->utf8String();
+      throw std::runtime_error(msg.str());
+    }
+    return lib;
+  }
+
+  // We have been given a path so try to load from lib_path / lib_name.metallib
+  if (lib_path.size() > 0) {
+    std::string full_path = lib_path + "/" + lib_name + ".metallib";
+    auto [lib, error] = load_library_from_path(device, full_path.c_str());
+    if (!lib) {
+      std::ostringstream msg;
+      msg << "Failed to load the metallib from <" << full_path
+          << "> with error " << error->localizedDescription()->utf8String();
+      throw std::runtime_error(msg.str());
+    }
+    return lib;
+  }
+
+  // Try to load the colocated library
+  {
+    auto [lib, error] = load_colocated_library(device, lib_name);
+    if (lib) {
+      return lib;
+    }
+  }
+
+  // Try to load the library from swiftpm
+  {
+    auto [lib, error] = load_swiftpm_library(device, lib_name);
+    if (lib) {
+      return lib;
+    }
+  }
+
+  std::ostringstream msg;
+  msg << "Failed to load the metallib " << lib_name << ".metallib. "
+      << "We attempted to load it from <" << current_binary_dir() << "/"
+      << lib_name << ".metallib>";
+#ifdef SWIFTPM_BUNDLE
+  msg << " and from the Swift PM bundle.";
+#endif
+  throw std::runtime_error(msg.str());
+}
+
+} // namespace
+
+CommandEncoder::CommandEncoder(DeviceStream& stream) : stream_(stream) {
+  enc_ = stream_.buffer->computeCommandEncoder(MTL::DispatchTypeConcurrent);
+  enc_->retain();
+}
+
+CommandEncoder::~CommandEncoder() {
+  enc_->endEncoding();
+  enc_->release();
+}
+
+void CommandEncoder::set_buffer(
+    const MTL::Buffer* buf,
+    int idx,
+    int64_t offset /* = 0 */) {
+  enc_->setBuffer(buf, offset, idx);
+}
+
+void CommandEncoder::set_input_array(
+    const array& a,
+    int idx,
+    int64_t offset /* = 0 */) {
+  if (all_inputs_.insert(a.buffer().ptr()).second) {
+    stream_.buffer_sizes += a.data_size();
+  }
+  auto r_buf = static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
+  needs_barrier_ =
+      needs_barrier_ | (prev_outputs_.find(r_buf) != prev_outputs_.end());
+  auto a_buf = static_cast<const MTL::Buffer*>(a.buffer().ptr());
+  enc_->setBuffer(a_buf, a.offset() + offset, idx);
+}
+
+void CommandEncoder::set_output_array(
+    array& a,
+    int idx,
+    int64_t offset /* = 0 */) {
+  // Add barriers before adding the output to the output set
+  set_input_array(a, idx, offset);
+  register_output_array(a);
+}
+
+void CommandEncoder::register_output_array(const array& a) {
+  all_outputs_.insert(a.buffer().ptr());
+
+  auto buf = static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
+  if (concurrent_) {
+    concurrent_outputs_.insert(buf);
+  } else {
+    next_outputs_.insert(buf);
+  }
+}
+
+void CommandEncoder::maybeInsertBarrier() {
+  if (needs_barrier_) {
+    enc_->memoryBarrier(MTL::BarrierScopeBuffers);
+    needs_barrier_ = false;
+    prev_outputs_ = std::move(next_outputs_);
+  } else {
+    prev_outputs_.insert(next_outputs_.begin(), next_outputs_.end());
+  }
+  next_outputs_.clear();
+}
+
+void CommandEncoder::dispatch_threadgroups(
+    MTL::Size grid_dims,
+    MTL::Size group_dims) {
+  maybeInsertBarrier();
+  stream_.buffer_ops++;
+  enc_->dispatchThreadgroups(grid_dims, group_dims);
+}
+
+void CommandEncoder::dispatch_threads(
+    MTL::Size grid_dims,
+    MTL::Size group_dims) {
+  maybeInsertBarrier();
+  stream_.buffer_ops++;
+  enc_->dispatchThreads(grid_dims, group_dims);
+}
+
+void CommandEncoder::barrier() {
+  enc_->memoryBarrier(MTL::BarrierScopeBuffers);
+}
+
+Device::Device() {
+  auto pool = new_scoped_memory_pool();
+  device_ = load_device();
+  default_library_ = load_default_library(device_);
+  arch_ = std::string(device_->architecture()->name()->utf8String());
+  int ag_tens = arch_[arch_.size() - 3] - '0';
+  int ag_ones = arch_[arch_.size() - 2] - '0';
+  arch_gen_ = ag_tens * 10 + ag_ones;
+  auto arch = arch_.back();
+  switch (arch) {
+    case 'p': // phone
+      max_ops_per_buffer_ = 20;
+      max_mb_per_buffer_ = 40;
+      break;
+    case 'g': // base, pro
+      max_ops_per_buffer_ = 40;
+      max_mb_per_buffer_ = 40;
+      break;
+    case 's': // max
+      max_ops_per_buffer_ = 50;
+      max_mb_per_buffer_ = 50;
+      break;
+    case 'd': // ultra
+      max_ops_per_buffer_ = 50;
+      max_mb_per_buffer_ = 50;
+      break;
+    default: // default to medium
+      max_ops_per_buffer_ = 40;
+      max_mb_per_buffer_ = 40;
+      break;
+  }
+  max_ops_per_buffer_ = env::max_ops_per_buffer(max_ops_per_buffer_);
+  max_mb_per_buffer_ = env::max_mb_per_buffer(max_mb_per_buffer_);
+}
+
+Device::~Device() {
+  auto pool = new_scoped_memory_pool();
+  for (auto& [l, kernel_map] : library_kernels_) {
+    l->release();
+    for (auto& [_, k] : kernel_map) {
+      k->release();
+    }
+  }
+  stream_map_.clear();
+  device_->release();
+}
+
+void Device::new_queue(int index) {
+  auto thread_pool = metal::new_scoped_memory_pool();
+  auto q = device_->newCommandQueue();
+  debug_set_stream_queue_label(q, index);
+  if (!q) {
+    throw std::runtime_error(
+        "[metal::Device] Failed to make new command queue.");
+  }
+  stream_map_.emplace(index, q);
+  if (residency_set_ != nullptr) {
+    q->addResidencySet(residency_set_);
+  }
+}
+
+MTL::CommandQueue* Device::get_queue(Stream stream) {
+  return get_stream_(stream.index).queue;
+}
+
+bool Device::command_buffer_needs_commit(int index) {
+  auto& stream = get_stream_(index);
+  return (stream.buffer_ops > max_ops_per_buffer_) ||
+      ((stream.buffer_sizes >> 20) > max_mb_per_buffer_);
+}
+
+MTL::CommandBuffer* Device::get_command_buffer(int index) {
+  auto& stream = get_stream_(index);
+  if (stream.buffer == nullptr) {
+    stream.buffer = stream.queue->commandBufferWithUnretainedReferences();
+    if (!stream.buffer) {
+      throw std::runtime_error(
+          "[metal::Device] Unable to create new command buffer");
+    }
+    // Increment ref count so the buffer is not garbage collected
+    stream.buffer->retain();
+  }
+  return stream.buffer;
+}
+
+void Device::commit_command_buffer(int index) {
+  auto& stream = get_stream_(index);
+  stream.buffer->commit();
+  stream.buffer->release();
+  stream.buffer = nullptr;
+  stream.buffer_ops = 0;
+  stream.buffer_sizes = 0;
+}
+
+void Device::add_temporary(array arr, int index) {
+  get_stream_(index).temporaries.push_back(std::move(arr));
+}
+
+void Device::add_temporaries(std::vector<array> arrays, int index) {
+  if (arrays.empty()) {
+    return;
+  }
+  auto& stream = get_stream_(index);
+  stream.temporaries.insert(
+      stream.temporaries.end(),
+      std::make_move_iterator(arrays.begin()),
+      std::make_move_iterator(arrays.end()));
+}
+
+void Device::end_encoding(int index) {
+  auto& stream = get_stream_(index);
+  if (stream.encoder != nullptr) {
+    // Each command encoder has a unique fence. We also store a map of
+    // all previous outputs of command encoders to their corresponding fence.
+    // - The command encoder records its inputs and outputs.
+    // - Wait on a fence if any inputs in the encoder are outputs of a previous
+    //   encoder.
+    // - Update the map of outputs to include this command encoder's outputs.
+    // - Always signal this command encoders fence.
+    // - Add a completion handler for this command encoder that removes outputs
+    //   from the map to limit the growth of the map and avoid unnecessary waits
+    // - Temporaries are a special case as they do not cross command encoder
+    //   boundaries. These can be removed early from the encoders inputs and
+    //   outputs since they don't need synchronization.
+    auto& enc = *stream.encoder;
+    // Remove temporaries from inputs and outputs
+    for (auto& t : stream.temporaries) {
+      enc.outputs().erase(t.buffer().ptr());
+      enc.inputs().erase(t.buffer().ptr());
+    }
+
+    // Keep references to the fences we waited on and put them
+    // in the completion handler so they are not prematurely released
+    std::unordered_set<std::shared_ptr<Fence>> waiting_on;
+    {
+      std::lock_guard<std::mutex> lk(stream.fence_mtx);
+      for (auto in : enc.inputs()) {
+        if (auto it = stream.outputs.find(in); it != stream.outputs.end()) {
+          // If we've already waited on a fence, don't wait on it again.
+          if (waiting_on.find(it->second) == waiting_on.end()) {
+            enc.wait_for_fence(it->second->fence);
+            waiting_on.insert(it->second);
+          }
+        }
+      }
+      for (auto out : enc.outputs()) {
+        stream.outputs[out] = stream.fence;
+      }
+    }
+    enc.update_fence(stream.fence->fence);
+    stream.buffer->addCompletedHandler(
+        [&stream,
+         waiting_on = std::move(waiting_on),
+         fence = std::move(stream.fence),
+         outputs = std::move(enc.outputs()),
+         temporaries =
+             std::move(stream.temporaries)](MTL::CommandBuffer*) mutable {
+          temporaries.clear();
+          std::lock_guard<std::mutex> lk(stream.fence_mtx);
+          for (auto o : outputs) {
+            if (auto it = stream.outputs.find(o); it != stream.outputs.end()) {
+              if (it->second == fence) {
+                stream.outputs.erase(it);
+              }
+            }
+          }
+        });
+  }
+  stream.encoder = nullptr;
+}
+
+CommandEncoder& Device::get_command_encoder(int index) {
+  auto& stream = get_stream_(index);
+  if (stream.encoder == nullptr) {
+    // Ensure there is an active command buffer
+    if (stream.buffer == nullptr) {
+      get_command_buffer(index);
+    }
+    stream.encoder = std::make_unique<CommandEncoder>(stream);
+    stream.fence = std::make_shared<Fence>(device_->newFence());
+  }
+  return *stream.encoder;
+}
+
+MTL::Library* Device::get_library(
+    const std::string& name,
+    const std::string& path /* = "" */) {
+  {
+    std::shared_lock rlock(library_mtx_);
+    if (auto it = library_map_.find(name); it != library_map_.end()) {
+      return it->second;
+    }
+  }
+
+  std::unique_lock wlock(library_mtx_);
+  if (auto it = library_map_.find(name); it != library_map_.end()) {
+    return it->second;
+  }
+
+  auto new_lib = load_library(device_, name, path.c_str());
+  library_map_.insert({name, new_lib});
+  return new_lib;
+}
+
+MTL::Library* Device::build_library_(const std::string& source_string) {
+  auto pool = new_scoped_memory_pool();
+
+  auto ns_code =
+      NS::String::string(source_string.c_str(), NS::ASCIIStringEncoding);
+
+  NS::Error* error = nullptr;
+  auto options = MTL::CompileOptions::alloc()->init();
+  options->setFastMathEnabled(false);
+  options->setLanguageVersion(get_metal_version());
+#ifndef NDEBUG
+  if (options->languageVersion() >= MTL::LanguageVersion3_2) {
+    options->setEnableLogging(true);
+  }
+#endif
+  auto mtl_lib = device_->newLibrary(ns_code, options, &error);
+  options->release();
+
+  // Throw error if unable to compile library
+  if (!mtl_lib) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to build metal library from source\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  return mtl_lib;
+}
+
+MTL::Function* Device::get_function_(
+    const std::string& name,
+    MTL::Library* mtl_lib) {
+  // Pull kernel from library
+  auto ns_name = NS::String::string(name.c_str(), NS::ASCIIStringEncoding);
+  auto mtl_function = mtl_lib->newFunction(ns_name);
+
+  return mtl_function;
+}
+
+MTL::Function* Device::get_function_(
+    const std::string& name,
+    const std::string& specialized_name,
+    const MTLFCList& func_consts,
+    MTL::Library* mtl_lib) {
+  if (func_consts.empty() && (specialized_name == name)) {
+    return get_function_(name, mtl_lib);
+  }
+
+  // Prepare function constants
+  auto mtl_func_consts = MTL::FunctionConstantValues::alloc()->init();
+
+  for (auto [value, type, index] : func_consts) {
+    mtl_func_consts->setConstantValue(value, type, index);
+  }
+
+  // Prepare function desc
+  auto desc = MTL::FunctionDescriptor::functionDescriptor();
+  desc->setName(NS::String::string(name.c_str(), NS::ASCIIStringEncoding));
+  desc->setSpecializedName(
+      NS::String::string(specialized_name.c_str(), NS::ASCIIStringEncoding));
+  desc->setConstantValues(mtl_func_consts);
+
+  // Pull kernel from library
+  NS::Error* error = nullptr;
+  auto mtl_function = mtl_lib->newFunction(desc, &error);
+
+  // Throw error if unable to build metal function
+  if (!mtl_function) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load function " << name << "\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  mtl_func_consts->release();
+
+  return mtl_function;
+}
+
+MTL::ComputePipelineState* Device::get_kernel_(
+    const std::string& name,
+    const MTL::Function* mtl_function) {
+  // Compile kernel to compute pipeline
+  NS::Error* error = nullptr;
+  MTL::ComputePipelineState* kernel;
+
+  if (mtl_function) {
+    kernel = device_->newComputePipelineState(mtl_function, &error);
+  }
+
+  // Throw error if unable to compile metal function
+  if (!mtl_function || !kernel) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load kernel " << name << "\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  return kernel;
+}
+
+MTL::ComputePipelineState* Device::get_kernel_(
+    const std::string& name,
+    const MTL::Function* mtl_function,
+    const MTL::LinkedFunctions* linked_functions) {
+  // Check inputs
+  if (!linked_functions) {
+    return get_kernel_(name, mtl_function);
+  }
+
+  if (!mtl_function) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load kernel " << name << "\n";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Prepare compute pipeline state descriptor
+  auto desc = MTL::ComputePipelineDescriptor::alloc()->init();
+  desc->setComputeFunction(mtl_function);
+  desc->setLinkedFunctions(linked_functions);
+
+  // Compile kernel to compute pipeline
+  NS::Error* error = nullptr;
+  auto kernel = device_->newComputePipelineState(
+      desc, MTL::PipelineOptionNone, nullptr, &error);
+
+  // Throw error if unable to compile metal function
+  if (!kernel) {
+    std::ostringstream msg;
+    msg << "[metal::Device] Unable to load kernel " << name << "\n";
+    if (error) {
+      msg << error->localizedDescription()->utf8String() << "\n";
+    }
+    throw std::runtime_error(msg.str());
+  }
+
+  return kernel;
+}
+
+MTL::Library* Device::get_library_(const std::string& name) {
+  std::shared_lock lock(library_mtx_);
+  auto it = library_map_.find(name);
+  return (it != library_map_.end()) ? it->second : nullptr;
+}
+
+MTL::Library* Device::get_library(
+    const std::string& name,
+    const std::function<std::string(void)>& builder) {
+  {
+    std::shared_lock rlock(library_mtx_);
+    if (auto it = library_map_.find(name); it != library_map_.end()) {
+      return it->second;
+    }
+  }
+
+  std::unique_lock wlock(library_mtx_);
+  if (auto it = library_map_.find(name); it != library_map_.end()) {
+    return it->second;
+  }
+
+  auto mtl_lib = build_library_(builder());
+  library_map_.insert({name, mtl_lib});
+  return mtl_lib;
+}
+
+void Device::clear_library(const std::string& name) {
+  std::unique_lock wlock(library_mtx_);
+  if (auto it = library_map_.find(name); it != library_map_.end()) {
+    auto kernel_map_it = library_kernels_.find(it->second);
+    for (auto& [_, kernel] : kernel_map_it->second) {
+      kernel->release();
+    }
+    library_kernels_.erase(kernel_map_it);
+    it->second->release();
+    library_map_.erase(it);
+  }
+}
+
+MTL::LinkedFunctions* Device::get_linked_functions_(
+    const std::vector<MTL::Function*>& funcs) {
+  if (funcs.empty()) {
+    return nullptr;
+  }
+
+  auto lfuncs = MTL::LinkedFunctions::linkedFunctions();
+
+  std::vector<NS::Object*> objs(funcs.size());
+  for (int i = 0; i < funcs.size(); i++) {
+    objs[i] = funcs[i];
+  }
+
+  NS::Array* funcs_arr = NS::Array::array(objs.data(), funcs.size());
+
+  lfuncs->setPrivateFunctions(funcs_arr);
+
+  return lfuncs;
+}
+
+MTL::ComputePipelineState* Device::get_kernel_(
+    const std::string& base_name,
+    MTL::Library* mtl_lib,
+    const std::string& hash_name,
+    const MTLFCList& func_consts /* = {} */,
+    const std::vector<MTL::Function*>& linked_functions /* = {} */) {
+  // Single writer allowed
+  std::unique_lock wlock(kernel_mtx_);
+
+  // Try loading again to avoid loading twice
+  auto& kernel_map_ = library_kernels_[mtl_lib];
+  if (auto it = kernel_map_.find(hash_name); it != kernel_map_.end()) {
+    return it->second;
+  }
+
+  auto pool = new_scoped_memory_pool();
+
+  // Pull kernel from library
+  auto mtl_function = get_function_(base_name, hash_name, func_consts, mtl_lib);
+
+  // Compile kernel to compute pipeline
+  auto mtl_linked_funcs = get_linked_functions_(linked_functions);
+  auto kernel = get_kernel_(hash_name, mtl_function, mtl_linked_funcs);
+
+  mtl_function->release();
+  mtl_linked_funcs->release();
+
+  // Add kernel to cache
+  kernel_map_.insert({hash_name, kernel});
+
+  return kernel;
+}
+
+MTL::ComputePipelineState* Device::get_kernel(
+    const std::string& base_name,
+    MTL::Library* mtl_lib,
+    const std::string& hash_name /* = "" */,
+    const MTLFCList& func_consts /* = {} */,
+    const std::vector<MTL::Function*>& linked_functions /* = {} */) {
+  const auto& kname = hash_name.empty() ? base_name : hash_name;
+  {
+    // Multiple readers allowed
+    std::shared_lock lock(kernel_mtx_);
+
+    // Look for cached kernel
+    auto& kernel_map_ = library_kernels_[mtl_lib];
+    if (auto it = kernel_map_.find(kname); it != kernel_map_.end()) {
+      return it->second;
+    }
+  }
+  return get_kernel_(base_name, mtl_lib, kname, func_consts, linked_functions);
+}
+
+MTL::ComputePipelineState* Device::get_kernel(
+    const std::string& base_name,
+    const std::string& hash_name /*  = "" */,
+    const MTLFCList& func_consts /*  = {} */,
+    const std::vector<MTL::Function*>& linked_functions /*  = {} */) {
+  return get_kernel(
+      base_name, default_library_, hash_name, func_consts, linked_functions);
+}
+
+void Device::set_residency_set(const MTL::ResidencySet* residency_set) {
+  if (residency_set_ != nullptr) {
+    throw std::runtime_error(
+        "[Device::set_residency_set] Can only be set once.");
+  }
+  if (residency_set == nullptr) {
+    return;
+  }
+  residency_set_ = residency_set;
+  // Attach residency set to existing command queues
+  for (auto& [_, stream] : stream_map_) {
+    stream.queue->addResidencySet(residency_set_);
+  }
+}
+
+Device& device(mlx::core::Device) {
+  static Device metal_device;
+  return metal_device;
+}
+
+std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool() {
+  auto dtor = [](void* ptr) {
+    static_cast<NS::AutoreleasePool*>(ptr)->release();
+  };
+  return std::unique_ptr<void, std::function<void(void*)>>(
+      NS::AutoreleasePool::alloc()->init(), dtor);
+}
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/device.h b/macos/third_party/mlx/mlx/backend/metal/device.h
new file mode 100644
index 0000000..e6162d7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/device.h
@@ -0,0 +1,289 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <Metal/Metal.hpp>
+#include <functional>
+#include <mutex>
+#include <shared_mutex>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/array.h"
+#include "mlx/device.h"
+
+namespace mlx::core::metal {
+
+using MTLFCList =
+    std::vector<std::tuple<const void*, MTL::DataType, NS::UInteger>>;
+
+struct DeviceStream;
+
+struct MLX_API CommandEncoder {
+  explicit CommandEncoder(DeviceStream& stream);
+  CommandEncoder(const CommandEncoder&) = delete;
+  CommandEncoder& operator=(const CommandEncoder&) = delete;
+
+  struct ConcurrentContext {
+    ConcurrentContext(CommandEncoder& enc) : enc(enc) {
+      enc.concurrent_ = true;
+    }
+    ~ConcurrentContext() {
+      enc.concurrent_ = false;
+      enc.prev_outputs_.insert(
+          enc.concurrent_outputs_.begin(), enc.concurrent_outputs_.end());
+      enc.concurrent_outputs_.clear();
+    }
+
+   private:
+    CommandEncoder& enc;
+  };
+
+  void set_input_array(const array& a, int idx, int64_t offset = 0);
+  void set_output_array(array& a, int idx, int64_t offset = 0);
+  void register_output_array(const array& a);
+  void dispatch_threadgroups(MTL::Size grid_dims, MTL::Size group_dims);
+  void dispatch_threads(MTL::Size grid_dims, MTL::Size group_dims);
+  void maybeInsertBarrier();
+  void set_buffer(const MTL::Buffer* buf, int idx, int64_t offset = 0);
+
+  void set_compute_pipeline_state(MTL::ComputePipelineState* kernel) {
+    enc_->setComputePipelineState(kernel);
+  }
+
+  void wait_for_fence(MTL::Fence* fence) {
+    enc_->waitForFence(fence);
+  }
+
+  void update_fence(MTL::Fence* fence) {
+    enc_->updateFence(fence);
+  }
+
+  template <typename Vec, typename = std::enable_if_t<is_vector_v<Vec>>>
+  void set_vector_bytes(const Vec& vec, size_t nelems, int idx) {
+    enc_->setBytes(vec.data(), nelems * sizeof(typename Vec::value_type), idx);
+  }
+  template <typename Vec, typename = std::enable_if_t<is_vector_v<Vec>>>
+  void set_vector_bytes(const Vec& vec, int idx) {
+    return set_vector_bytes(vec, vec.size(), idx);
+  }
+
+  template <typename T>
+  void set_bytes(const T* v, int n, int idx) {
+    return enc_->setBytes(v, n * sizeof(T), idx);
+  }
+
+  template <typename T>
+  void set_bytes(const T& v, int idx) {
+    return enc_->setBytes(&v, sizeof(T), idx);
+  }
+
+  void set_threadgroup_memory_length(size_t length, int idx) {
+    enc_->setThreadgroupMemoryLength(length, idx);
+  }
+
+  ConcurrentContext start_concurrent() {
+    return ConcurrentContext(*this);
+  }
+  ~CommandEncoder();
+
+  // Inputs to all kernels in the encoder including temporaries
+  std::unordered_set<const void*>& inputs() {
+    return all_inputs_;
+  };
+
+  // Outputs of all kernels in the encoder including temporaries
+  std::unordered_set<const void*>& outputs() {
+    return all_outputs_;
+  };
+
+  void barrier();
+
+ private:
+  DeviceStream& stream_;
+  MTL::ComputeCommandEncoder* enc_;
+  bool needs_barrier_{false};
+  bool concurrent_{false};
+  std::unordered_set<MTL::Resource*> prev_outputs_;
+  std::unordered_set<MTL::Resource*> next_outputs_;
+  std::unordered_set<MTL::Resource*> concurrent_outputs_;
+  std::unordered_set<const void*> all_inputs_;
+  std::unordered_set<const void*> all_outputs_;
+};
+
+struct Fence {
+  Fence(MTL::Fence* fence) : fence(fence) {}
+  ~Fence() {
+    fence->release();
+  }
+  MTL::Fence* fence;
+};
+
+struct DeviceStream {
+  DeviceStream(MTL::CommandQueue* queue) : queue(queue) {};
+  ~DeviceStream() {
+    queue->release();
+    if (buffer != nullptr) {
+      buffer->release();
+    }
+  };
+  MTL::CommandQueue* queue;
+  // A map of prior command encoder outputs to their corresponding fence
+  std::unordered_map<const void*, std::shared_ptr<Fence>> outputs;
+  // Used to allow thread-safe access to the outputs map
+  std::mutex fence_mtx;
+
+  // Data updated between command buffers
+  MTL::CommandBuffer* buffer{nullptr};
+  int buffer_ops{0};
+  size_t buffer_sizes{0};
+
+  // The command encoder, fence, and temporaries are updated between command
+  // encoders
+  std::unique_ptr<CommandEncoder> encoder{nullptr};
+  std::shared_ptr<Fence> fence;
+  std::vector<array> temporaries;
+};
+
+class MLX_API Device {
+ public:
+  Device();
+  Device(const Device&) = delete;
+  Device& operator=(const Device&) = delete;
+  ~Device();
+
+  MTL::Device* mtl_device() {
+    return device_;
+  };
+
+  const std::string& get_architecture() {
+    return arch_;
+  }
+
+  int get_architecture_gen() const {
+    return arch_gen_;
+  }
+
+  void new_queue(int index);
+
+  MTL::CommandQueue* get_queue(Stream stream);
+
+  MTL::CommandBuffer* get_command_buffer(int index);
+  bool command_buffer_needs_commit(int index);
+  void commit_command_buffer(int index);
+  CommandEncoder& get_command_encoder(int index);
+  void end_encoding(int index);
+
+  MTL::Library* get_library(
+      const std::string& name,
+      const std::string& path = "");
+
+  MTL::Library* get_library(
+      const std::string& name,
+      const std::function<std::string(void)>& builder);
+
+  void clear_library(const std::string& name);
+
+  MTL::ComputePipelineState* get_kernel(
+      const std::string& base_name,
+      MTL::Library* mtl_lib,
+      const std::string& hash_name = "",
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
+
+  MTL::ComputePipelineState* get_kernel(
+      const std::string& base_name,
+      const std::string& hash_name = "",
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
+
+  MTL::ArgumentEncoder* argument_encoder(
+      const std::vector<MTL::ArgumentDescriptor*>& arg_descs) const;
+
+  // Record temporary arrays for the given stream index
+  void add_temporary(array arr, int index);
+  void add_temporaries(std::vector<array> arrays, int index);
+
+  void set_residency_set(const MTL::ResidencySet* residency_set);
+
+ private:
+  DeviceStream& get_stream_(int index) {
+    return stream_map_.find(index)->second;
+  }
+  MTL::Library* get_library_cache_(const std::string& name);
+
+  MTL::Library* get_library_(const std::string& name);
+  MTL::Library* build_library_(const std::string& source_string);
+
+  MTL::Function* get_function_(const std::string& name, MTL::Library* mtl_lib);
+
+  MTL::Function* get_function_(
+      const std::string& name,
+      const std::string& specialized_name,
+      const MTLFCList& func_consts,
+      MTL::Library* mtl_lib);
+
+  MTL::LinkedFunctions* get_linked_functions_(
+      const std::vector<MTL::Function*>& funcs);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& name,
+      const MTL::Function* mtl_function);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& name,
+      const MTL::Function* mtl_function,
+      const MTL::LinkedFunctions* linked_functions);
+
+  MTL::ComputePipelineState* get_kernel_(
+      const std::string& base_name,
+      MTL::Library* mtl_lib,
+      const std::string& hash_name,
+      const MTLFCList& func_consts = {},
+      const std::vector<MTL::Function*>& linked_functions = {});
+
+  MTL::Device* device_;
+  std::unordered_map<int32_t, DeviceStream> stream_map_;
+
+  std::shared_mutex kernel_mtx_;
+  std::shared_mutex library_mtx_;
+  std::unordered_map<std::string, MTL::Library*> library_map_;
+  MTL::Library* default_library_;
+  std::unordered_map<
+      MTL::Library*,
+      std::unordered_map<std::string, MTL::ComputePipelineState*>>
+      library_kernels_;
+  const MTL::ResidencySet* residency_set_{nullptr};
+  std::string arch_;
+  int arch_gen_;
+  int max_ops_per_buffer_;
+  int max_mb_per_buffer_;
+};
+
+MLX_API Device& device(mlx::core::Device);
+
+std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
+
+inline bool is_nax_available() {
+#ifdef MLX_METAL_NO_NAX
+  return false;
+#else
+  auto _check_nax = []() {
+    bool can_use_nax = false;
+    if (__builtin_available(
+            macOS 26.2, iOS 26.2, tvOS 26.2, visionOS 26.2, *)) {
+      can_use_nax = true;
+    }
+    auto& d = metal::device(mlx::core::Device::gpu);
+    auto arch = d.get_architecture().back();
+    auto gen = d.get_architecture_gen();
+    can_use_nax &= gen >= (arch == 'p' ? 18 : 17);
+    return can_use_nax;
+  };
+  static bool is_nax_available_ = _check_nax();
+  return is_nax_available_;
+#endif
+}
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/device_info.cpp b/macos/third_party/mlx/mlx/backend/metal/device_info.cpp
new file mode 100644
index 0000000..dd18fc6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/device_info.cpp
@@ -0,0 +1,58 @@
+// Copyright © 2026 Apple Inc.
+
+#include <sys/sysctl.h>
+
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/metal.h"
+
+namespace mlx::core::gpu {
+
+bool is_available() {
+  return metal::is_available();
+}
+
+int device_count() {
+  return 1;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index) {
+  auto init_device_info = []()
+      -> std::unordered_map<std::string, std::variant<std::string, size_t>> {
+    auto pool = metal::new_scoped_memory_pool();
+    auto raw_device = metal::device(mlx::core::Device::gpu).mtl_device();
+    auto name = std::string(raw_device->name()->utf8String());
+    auto arch = std::string(raw_device->architecture()->name()->utf8String());
+
+    size_t memsize = 0;
+    size_t length = sizeof(memsize);
+    sysctlbyname("hw.memsize", &memsize, &length, NULL, 0);
+
+    size_t rsrc_limit = 0;
+    sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0);
+    if (rsrc_limit == 0) {
+      rsrc_limit = 499000;
+    }
+
+    return {
+        {"device_name", name},
+        {"architecture", arch},
+        {"max_buffer_length", raw_device->maxBufferLength()},
+        {"max_recommended_working_set_size",
+         raw_device->recommendedMaxWorkingSetSize()},
+        {"memory_size", memsize},
+        {"resource_limit", rsrc_limit}};
+  };
+  static auto device_info_ = init_device_info();
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+
+  if (device_index == 0) {
+    return device_info_;
+  } else {
+    return empty;
+  }
+}
+
+} // namespace mlx::core::gpu
diff --git a/macos/third_party/mlx/mlx/backend/metal/distributed.cpp b/macos/third_party/mlx/mlx/backend/metal/distributed.cpp
new file mode 100644
index 0000000..217ee3c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/distributed.cpp
@@ -0,0 +1,38 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/distributed/ops.h"
+#include "mlx/distributed/primitives.h"
+#include "mlx/fence.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core::distributed {
+
+void AllReduce::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[AllReduce::eval_gpu] has no GPU implementation.");
+}
+
+void AllGather::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[AllGather::eval_gpu] has no GPU implementation.");
+}
+
+void Send::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[Send::eval_gpu] has no GPU implementation.");
+}
+
+void Recv::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error("[Recv::eval_gpu] has no GPU implementation.");
+}
+
+void ReduceScatter::eval_gpu(const std::vector<array>&, std::vector<array>&) {
+  throw std::runtime_error(
+      "[ReduceScatter::eval_gpu] has no GPU implementation.");
+}
+
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/backend/metal/eval.cpp b/macos/third_party/mlx/mlx/backend/metal/eval.cpp
new file mode 100644
index 0000000..bd58a69
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/eval.cpp
@@ -0,0 +1,97 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <memory>
+
+#include "mlx/backend/gpu/eval.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core::gpu {
+
+void new_stream(Stream stream) {
+  if (stream.device == mlx::core::Device::gpu) {
+    metal::device(stream.device).new_queue(stream.index);
+  }
+}
+
+inline void check_error(MTL::CommandBuffer* cbuf) {
+  if (cbuf->status() == MTL::CommandBufferStatusError) {
+    std::ostringstream msg;
+    msg << "[METAL] Command buffer execution failed: "
+        << cbuf->error()->localizedDescription()->utf8String();
+    throw std::runtime_error(msg.str());
+  }
+}
+
+void eval(array& arr) {
+  auto pool = metal::new_scoped_memory_pool();
+  auto s = arr.primitive().stream();
+  auto& d = metal::device(s.device);
+  auto command_buffer = d.get_command_buffer(s.index);
+
+  auto outputs = arr.outputs();
+  {
+    // If the array is a tracer hold a reference
+    // to its inputs so they don't get donated
+    std::vector<array> inputs;
+    if (arr.is_tracer()) {
+      inputs = arr.inputs();
+    }
+
+    debug_set_primitive_buffer_label(command_buffer, arr.primitive());
+    arr.primitive().eval_gpu(arr.inputs(), outputs);
+  }
+  std::unordered_set<std::shared_ptr<array::Data>> buffers;
+  for (auto& in : arr.inputs()) {
+    buffers.insert(in.data_shared_ptr());
+  }
+  for (auto& s : arr.siblings()) {
+    buffers.insert(s.data_shared_ptr());
+  }
+  // Remove the output if it was donated to by an input
+  if (auto it = buffers.find(arr.data_shared_ptr()); it != buffers.end()) {
+    buffers.erase(it);
+  }
+
+  if (d.command_buffer_needs_commit(s.index)) {
+    d.end_encoding(s.index);
+    scheduler::notify_new_task(s);
+    command_buffer->addCompletedHandler(
+        [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+          scheduler::notify_task_completion(s);
+          check_error(cbuf);
+        });
+    d.commit_command_buffer(s.index);
+    d.get_command_buffer(s.index);
+  } else {
+    command_buffer->addCompletedHandler(
+        [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {
+          check_error(cbuf);
+        });
+  }
+}
+
+void finalize(Stream s) {
+  auto pool = metal::new_scoped_memory_pool();
+  auto& d = metal::device(s.device);
+  auto cb = d.get_command_buffer(s.index);
+  d.end_encoding(s.index);
+  cb->addCompletedHandler([](MTL::CommandBuffer* cbuf) { check_error(cbuf); });
+  d.commit_command_buffer(s.index);
+  d.get_command_buffer(s.index);
+}
+
+void synchronize(Stream s) {
+  auto pool = metal::new_scoped_memory_pool();
+  auto& d = metal::device(s.device);
+  auto cb = d.get_command_buffer(s.index);
+  cb->retain();
+  d.end_encoding(s.index);
+  d.commit_command_buffer(s.index);
+  cb->waitUntilCompleted();
+  check_error(cb);
+  cb->release();
+}
+
+} // namespace mlx::core::gpu
diff --git a/macos/third_party/mlx/mlx/backend/metal/event.cpp b/macos/third_party/mlx/mlx/backend/metal/event.cpp
new file mode 100644
index 0000000..39a4923
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/event.cpp
@@ -0,0 +1,62 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/event.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core {
+
+Event::Event(Stream stream) : stream_(stream) {
+  auto dtor = [](void* ptr) {
+    auto p = metal::new_scoped_memory_pool();
+    static_cast<MTL::SharedEvent*>(ptr)->release();
+  };
+  auto p = metal::new_scoped_memory_pool();
+  event_ = std::shared_ptr<void>(
+      metal::device(Device::gpu).mtl_device()->newSharedEvent(), dtor);
+  if (event_ == nullptr) {
+    throw std::runtime_error(
+        "[Event::Event] Failed to create Metal shared event.");
+  }
+}
+
+void Event::wait() {
+  if (!static_cast<MTL::SharedEvent*>(event_.get())
+           ->waitUntilSignaledValue(value(), -1)) {
+    throw std::runtime_error("[Event::wait] Timed out");
+  }
+}
+
+void Event::wait(Stream stream) {
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [*this]() mutable { wait(); });
+  } else {
+    auto& d = metal::device(stream.device);
+    d.end_encoding(stream.index);
+    auto command_buffer = d.get_command_buffer(stream.index);
+    command_buffer->encodeWait(static_cast<MTL::Event*>(event_.get()), value());
+    command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {});
+  }
+}
+
+void Event::signal(Stream stream) {
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [*this]() mutable {
+      static_cast<MTL::SharedEvent*>(event_.get())->setSignaledValue(value());
+    });
+  } else {
+    auto& d = metal::device(stream.device);
+    d.end_encoding(stream.index);
+    auto command_buffer = d.get_command_buffer(stream.index);
+    command_buffer->encodeSignalEvent(
+        static_cast<MTL::Event*>(event_.get()), value());
+    command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {});
+  }
+}
+
+bool Event::is_signaled() const {
+  return static_cast<MTL::SharedEvent*>(event_.get())->signaledValue() >=
+      value();
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/fence.cpp b/macos/third_party/mlx/mlx/backend/metal/fence.cpp
new file mode 100644
index 0000000..da73d5d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/fence.cpp
@@ -0,0 +1,162 @@
+// Copyright © 2024 Apple Inc.
+#include "mlx/fence.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/scheduler.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+struct FenceImpl {
+  FenceImpl() {
+    auto d = metal::device(Device::gpu).mtl_device();
+    if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
+      use_fast = false;
+    } else if (__builtin_available(macOS 15, iOS 18, *)) {
+      use_fast = env::metal_fast_synch();
+    }
+
+    if (!use_fast) {
+      auto p = metal::new_scoped_memory_pool();
+      fence = static_cast<void*>(d->newSharedEvent());
+    } else {
+      auto buf = allocator::malloc(sizeof(uint32_t)).ptr();
+      fence = static_cast<void*>(buf);
+      cpu_value()[0] = 0;
+    }
+  }
+
+  ~FenceImpl() {
+    if (!use_fast) {
+      // Wraps Metal SharedEvent
+      auto p = metal::new_scoped_memory_pool();
+      static_cast<MTL::SharedEvent*>(fence)->release();
+    } else {
+      allocator::free(allocator::Buffer{static_cast<MTL::Buffer*>(fence)});
+    }
+  }
+  bool use_fast{false};
+  uint32_t count{0};
+  void* fence;
+
+  std::atomic_uint* cpu_value() {
+    return static_cast<std::atomic_uint*>(
+        static_cast<MTL::Buffer*>(fence)->contents());
+  }
+};
+
+Fence::Fence(Stream) {
+  auto dtor = [](void* ptr) { delete static_cast<FenceImpl*>(ptr); };
+  fence_ = std::shared_ptr<void>(new FenceImpl{}, dtor);
+}
+
+void Fence::wait(Stream stream, const array& x) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      if (!f.use_fast) {
+        if (!static_cast<MTL::SharedEvent*>(f.fence)->waitUntilSignaledValue(
+                count, -1)) {
+          throw std::runtime_error("[Fence::wait] Timed out");
+        }
+        return;
+      }
+      while (f.cpu_value()[0] < count) {
+      }
+    });
+    return;
+  }
+
+  auto& d = metal::device(stream.device);
+  auto idx = stream.index;
+
+  if (!f.use_fast) {
+    d.end_encoding(idx);
+    auto command_buffer = d.get_command_buffer(idx);
+    command_buffer->encodeWait(static_cast<MTL::Event*>(f.fence), f.count);
+    command_buffer->addCompletedHandler(
+        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+    return;
+  }
+
+  auto& compute_encoder = d.get_command_encoder(idx);
+
+  // Register outputs to ensure that no kernels which depends on the
+  // output starts before this one is done
+  compute_encoder.register_output_array(x);
+
+  auto kernel = d.get_kernel("fence_wait");
+  MTL::Size kernel_dims = MTL::Size(1, 1, 1);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  auto buf = static_cast<MTL::Buffer*>(f.fence);
+  compute_encoder.set_buffer(buf, 0);
+  compute_encoder.set_bytes(f.count, 1);
+  compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
+
+  d.get_command_buffer(idx)->addCompletedHandler(
+      [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+}
+
+void Fence::update(Stream stream, const array& x, bool cross_device) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+  f.count++;
+
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      if (!f.use_fast) {
+        static_cast<MTL::SharedEvent*>(f.fence)->setSignaledValue(count);
+        return;
+      }
+
+      f.cpu_value()[0] = count;
+    });
+    return;
+  }
+
+  auto& d = metal::device(stream.device);
+  auto idx = stream.index;
+  if (!f.use_fast) {
+    d.end_encoding(idx);
+    auto command_buffer = d.get_command_buffer(idx);
+    command_buffer->encodeSignalEvent(
+        static_cast<MTL::Event*>(f.fence), f.count);
+    command_buffer->addCompletedHandler(
+        [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+    return;
+  }
+
+  // Launch input visibility kernels
+  auto& compute_encoder = d.get_command_encoder(idx);
+  if (cross_device) {
+    auto kernel = d.get_kernel("input_coherent");
+    uint32_t nthreads = (x.data_size() * x.itemsize() + sizeof(uint32_t) - 1) /
+        sizeof(uint32_t);
+    MTL::Size group_dims = MTL::Size(1024, 1, 1);
+    MTL::Size grid_dims = MTL::Size((nthreads + 1024 - 1) / 1024, 1, 1);
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(x, 0);
+    compute_encoder.set_bytes(nthreads, 1);
+    compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+  }
+
+  // Barrier on previous kernels
+  compute_encoder.barrier();
+
+  // Launch value update kernel
+  auto kernel = d.get_kernel("fence_update");
+  MTL::Size kernel_dims = MTL::Size(1, 1, 1);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  auto buf = static_cast<MTL::Buffer*>(f.fence);
+  compute_encoder.set_buffer(buf, 0);
+  compute_encoder.set_bytes(f.count, 1);
+  compute_encoder.dispatch_threads(kernel_dims, kernel_dims);
+
+  d.get_command_buffer(idx)->addCompletedHandler(
+      [fence_ = fence_](MTL::CommandBuffer* cbuf) {});
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/fft.cpp b/macos/third_party/mlx/mlx/backend/metal/fft.cpp
new file mode 100644
index 0000000..d99e1ba
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/fft.cpp
@@ -0,0 +1,807 @@
+// Copyright © 2024 Apple Inc.
+#include <cassert>
+#include <complex>
+#include <map>
+#include <numeric>
+#include <set>
+
+#include "mlx/3rdparty/pocketfft.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/gpu/slicing.h"
+#include "mlx/backend/metal/binary.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/unary.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+using MTLFC = std::tuple<const void*, MTL::DataType, NS::UInteger>;
+
+#define MAX_STOCKHAM_FFT_SIZE 4096
+#define MAX_RADER_FFT_SIZE 2048
+#define MAX_BLUESTEIN_FFT_SIZE 2048
+// Threadgroup memory batching improves throughput for small n
+#define MIN_THREADGROUP_MEM_SIZE 256
+// For strided reads/writes, coalesce at least this many complex64s
+#define MIN_COALESCE_WIDTH 4
+
+inline const std::vector<int> supported_radices() {
+  // Ordered by preference in decomposition.
+  return {13, 11, 8, 7, 6, 5, 4, 3, 2};
+}
+
+std::vector<int> prime_factors(int n) {
+  int z = 2;
+  std::vector<int> factors;
+  while (z * z <= n) {
+    if (n % z == 0) {
+      factors.push_back(z);
+      n /= z;
+    } else {
+      z++;
+    }
+  }
+  if (n > 1) {
+    factors.push_back(n);
+  }
+  return factors;
+}
+
+struct FourStepParams {
+  bool required = false;
+  bool first_step = true;
+  int n1 = 0;
+  int n2 = 0;
+};
+
+// Forward Declaration
+void fft_op(
+    const array& in,
+    array& out,
+    size_t axis,
+    bool inverse,
+    bool real,
+    const FourStepParams four_step_params,
+    bool inplace,
+    const Stream& s);
+
+struct FFTPlan {
+  int n = 0;
+  // Number of steps for each radix in the Stockham decomposition
+  std::vector<int> stockham;
+  // Number of steps for each radix in the Rader decomposition
+  std::vector<int> rader;
+  // Rader factor, 1 if no rader factors
+  int rader_n = 1;
+  int bluestein_n = -1;
+  // Four step FFT
+  bool four_step = false;
+  int n1 = 0;
+  int n2 = 0;
+};
+
+int next_fast_n(int n) {
+  return next_power_of_2(n);
+}
+
+std::vector<int> plan_stockham_fft(int n) {
+  auto radices = supported_radices();
+  std::vector<int> plan(radices.size(), 0);
+  int orig_n = n;
+  if (n == 1) {
+    return plan;
+  }
+  for (int i = 0; i < radices.size(); i++) {
+    int radix = radices[i];
+    // Manually tuned radices for powers of 2
+    if (is_power_of_2(orig_n) && orig_n < 512 && radix > 4) {
+      continue;
+    }
+    while (n % radix == 0) {
+      plan[i] += 1;
+      n /= radix;
+      if (n == 1) {
+        return plan;
+      }
+    }
+  }
+  throw std::runtime_error("Unplannable");
+}
+
+FFTPlan plan_fft(int n) {
+  auto radices = supported_radices();
+  std::set<int> radices_set(radices.begin(), radices.end());
+
+  FFTPlan plan;
+  plan.n = n;
+  plan.rader = std::vector<int>(radices.size(), 0);
+  auto factors = prime_factors(n);
+  int remaining_n = n;
+
+  // Four Step FFT when N is too large for shared mem.
+  if (n > MAX_STOCKHAM_FFT_SIZE && is_power_of_2(n)) {
+    // For power's of two we have a fast, no transpose four step implementation.
+    plan.four_step = true;
+    // Rough heuristic for choosing faster powers of two when we can
+    plan.n2 = n > 65536 ? 1024 : 64;
+    plan.n1 = n / plan.n2;
+    return plan;
+  } else if (n > MAX_STOCKHAM_FFT_SIZE) {
+    // Otherwise we use a multi-upload Bluestein's
+    plan.four_step = true;
+    plan.bluestein_n = next_fast_n(2 * n - 1);
+    return plan;
+  }
+
+  for (int factor : factors) {
+    // Make sure the factor is a supported radix
+    if (radices_set.find(factor) == radices_set.end()) {
+      // We only support a single Rader factor currently
+      // TODO(alexbarron) investigate weirdness with large
+      // Rader sizes -- possibly a compiler issue?
+      if (plan.rader_n > 1 || n > MAX_RADER_FFT_SIZE) {
+        plan.four_step = n > MAX_BLUESTEIN_FFT_SIZE;
+        plan.bluestein_n = next_fast_n(2 * n - 1);
+        plan.stockham = plan_stockham_fft(plan.bluestein_n);
+        plan.rader = std::vector<int>(radices.size(), 0);
+        return plan;
+      }
+      // See if we can use Rader's algorithm to Stockham decompose n - 1
+      auto rader_factors = prime_factors(factor - 1);
+      for (int rf : rader_factors) {
+        // We don't nest Rader's algorithm so if `factor - 1`
+        // isn't Stockham decomposable we give up and do Bluestein's.
+        if (radices_set.find(rf) == radices_set.end()) {
+          plan.four_step = n > MAX_BLUESTEIN_FFT_SIZE;
+          plan.bluestein_n = next_fast_n(2 * n - 1);
+          plan.stockham = plan_stockham_fft(plan.bluestein_n);
+          plan.rader = std::vector<int>(radices.size(), 0);
+          return plan;
+        }
+      }
+      plan.rader = plan_stockham_fft(factor - 1);
+      plan.rader_n = factor;
+      remaining_n /= factor;
+    }
+  }
+
+  plan.stockham = plan_stockham_fft(remaining_n);
+  return plan;
+}
+
+int compute_elems_per_thread(FFTPlan plan) {
+  // Heuristics for selecting an efficient number
+  // of threads to use for a particular mixed-radix FFT.
+  auto n = plan.n;
+
+  std::vector<int> steps;
+  auto radices = supported_radices();
+  steps.insert(steps.end(), plan.stockham.begin(), plan.stockham.end());
+  steps.insert(steps.end(), plan.rader.begin(), plan.rader.end());
+  std::set<int> used_radices;
+  for (int i = 0; i < steps.size(); i++) {
+    int radix = radices[i % radices.size()];
+    if (steps[i] > 0) {
+      used_radices.insert(radix);
+    }
+  }
+
+  // Manual tuning for 7/11/13
+  if (used_radices.find(7) != used_radices.end() &&
+      (used_radices.find(11) != used_radices.end() ||
+       used_radices.find(13) != used_radices.end())) {
+    return 7;
+  } else if (
+      used_radices.find(11) != used_radices.end() &&
+      used_radices.find(13) != used_radices.end()) {
+    return 11;
+  }
+
+  // TODO(alexbarron) Some really weird stuff is going on
+  // for certain `elems_per_thread` on large composite n.
+  // Possibly a compiler issue?
+  if (n == 3159)
+    return 13;
+  if (n == 3645)
+    return 5;
+  if (n == 3969)
+    return 7;
+  if (n == 1982)
+    return 5;
+
+  if (used_radices.size() == 1) {
+    return *(used_radices.begin());
+  }
+  if (used_radices.size() == 2) {
+    if (used_radices.find(11) != used_radices.end() ||
+        used_radices.find(13) != used_radices.end()) {
+      return std::accumulate(used_radices.begin(), used_radices.end(), 0) / 2;
+    }
+    std::vector<int> radix_vec(used_radices.begin(), used_radices.end());
+    return radix_vec[1];
+  }
+  // In all other cases use the second smallest radix.
+  std::vector<int> radix_vec(used_radices.begin(), used_radices.end());
+  return radix_vec[1];
+}
+
+// Rader
+int mod_exp(int x, int y, int n) {
+  int out = 1;
+  while (y) {
+    if (y & 1) {
+      out = out * x % n;
+    }
+    y >>= 1;
+    x = x * x % n;
+  }
+  return out;
+}
+
+int primitive_root(int n) {
+  auto factors = prime_factors(n - 1);
+
+  for (int r = 2; r < n - 1; r++) {
+    bool found = true;
+    for (int factor : factors) {
+      if (mod_exp(r, (n - 1) / factor, n) == 1) {
+        found = false;
+        break;
+      }
+    }
+    if (found) {
+      return r;
+    }
+  }
+  return -1;
+}
+
+std::tuple<array, array, array> compute_raders_constants(
+    int rader_n,
+    const Stream& s) {
+  int proot = primitive_root(rader_n);
+  // Fermat's little theorem
+  int inv = mod_exp(proot, rader_n - 2, rader_n);
+  std::vector<short> g_q(rader_n - 1);
+  std::vector<short> g_minus_q(rader_n - 1);
+  for (int i = 0; i < rader_n - 1; i++) {
+    g_q[i] = mod_exp(proot, i, rader_n);
+    g_minus_q[i] = mod_exp(inv, i, rader_n);
+  }
+  array g_q_arr(g_q.begin(), {rader_n - 1});
+  array g_minus_q_arr(g_minus_q.begin(), {rader_n - 1});
+
+  std::vector<std::complex<float>> b_q(rader_n - 1);
+  for (int i = 0; i < rader_n - 1; i++) {
+    float pi_i = (float)g_minus_q[i] * -2.0 * M_PI / rader_n;
+    b_q[i] = std::exp(std::complex<float>(0, pi_i));
+  }
+
+  array b_q_fft({rader_n - 1}, complex64, nullptr, {});
+  b_q_fft.set_data(allocator::malloc(b_q_fft.nbytes()));
+  auto b_q_fft_ptr =
+      reinterpret_cast<std::complex<float>*>(b_q_fft.data<complex64_t>());
+  std::ptrdiff_t item_size = b_q_fft.itemsize();
+  size_t fft_size = rader_n - 1;
+  // This FFT is always small (<4096, batch 1) so save some overhead
+  // and do it on the CPU
+  pocketfft::c2c(
+      /* shape= */ {fft_size},
+      /* stride_in= */ {item_size},
+      /* stride_out= */ {item_size},
+      /* axes= */ {0},
+      /* forward= */ true,
+      /* data_in= */ b_q.data(),
+      /* data_out= */ b_q_fft_ptr,
+      /* scale= */ 1.0f);
+  return std::make_tuple(b_q_fft, g_q_arr, g_minus_q_arr);
+}
+
+// Bluestein
+std::pair<array, array> compute_bluestein_constants(int n, int bluestein_n) {
+  // We need to calculate the Bluestein twiddle factors
+  // in double precision for the overall numerical stability
+  // of Bluestein's FFT algorithm to be acceptable.
+  //
+  // Metal doesn't support float64, so instead we
+  // manually implement the required operations on cpu.
+  //
+  // In numpy:
+  // w_k = np.exp(-1j * np.pi / N * (np.arange(-N + 1, N) ** 2))
+  // w_q = np.fft.fft(1/w_k)
+  // return w_k, w_q
+  std::vector<std::complex<float>> w_k_vec(n);
+  std::vector<std::complex<float>> w_q_vec(bluestein_n, 0);
+
+  for (int i = -n + 1; i < n; i++) {
+    double theta = pow(i, 2) * M_PI / (double)n;
+    w_q_vec[i + n - 1] = std::exp(std::complex<double>(0, theta));
+    if (i >= 0) {
+      w_k_vec[i] = std::exp(std::complex<double>(0, -theta));
+    }
+  }
+
+  array w_k({n}, complex64, nullptr, {});
+  w_k.set_data(allocator::malloc(w_k.nbytes()));
+  std::copy(w_k_vec.begin(), w_k_vec.end(), w_k.data<complex64_t>());
+
+  array w_q({bluestein_n}, complex64, nullptr, {});
+  w_q.set_data(allocator::malloc(w_q.nbytes()));
+  auto w_q_ptr =
+      reinterpret_cast<std::complex<float>*>(w_q.data<complex64_t>());
+
+  std::ptrdiff_t item_size = w_q.itemsize();
+  size_t fft_size = bluestein_n;
+  pocketfft::c2c(
+      /* shape= */ {fft_size},
+      /* stride_in= */ {item_size},
+      /* stride_out= */ {item_size},
+      /* axes= */ {0},
+      /* forward= */ true,
+      /* data_in= */ w_q_vec.data(),
+      /* data_out= */ w_q_ptr,
+      /* scale= */ 1.0f);
+  return std::make_tuple(w_k, w_q);
+}
+
+void multi_upload_bluestein_fft(
+    const array& in,
+    array& out,
+    size_t axis,
+    bool inverse,
+    bool real,
+    FFTPlan& plan,
+    std::vector<array>& copies,
+    const Stream& s) {
+  // TODO(alexbarron) Implement fused kernels for mutli upload bluestein's
+  // algorithm
+  int n = inverse ? out.shape(axis) : in.shape(axis);
+  auto [w_k, w_q] = compute_bluestein_constants(n, plan.bluestein_n);
+  copies.push_back(w_k);
+  copies.push_back(w_q);
+
+  auto temp_shape = inverse ? out.shape() : in.shape();
+  array temp(temp_shape, complex64, nullptr, {});
+  array temp1(temp_shape, complex64, nullptr, {});
+
+  if (real && !inverse) {
+    // Convert float32->complex64
+    copy_gpu(in, temp, CopyType::General, s);
+    copies.push_back(temp);
+  } else if (real && inverse) {
+    int back_offset = n % 2 == 0 ? 2 : 1;
+    auto slice_shape = in.shape();
+    slice_shape[axis] -= back_offset;
+    array slice_temp(slice_shape, complex64, nullptr, {});
+    array conj_temp(in.shape(), complex64, nullptr, {});
+    copies.push_back(conj_temp);
+
+    Shape rstarts(in.ndim(), 0);
+    Shape rstrides(in.ndim(), 1);
+    rstarts[axis] = in.shape(axis) - back_offset;
+    rstrides[axis] = -1;
+    unary_op_gpu({in}, conj_temp, "Conjugate", s);
+    slice_gpu(in, slice_temp, rstarts, rstrides, s);
+    concatenate_gpu({conj_temp, slice_temp}, temp, (int)axis, s);
+    copies.push_back(temp);
+  } else if (inverse) {
+    unary_op_gpu({in}, temp, "Conjugate", s);
+    copies.push_back(temp);
+  } else {
+    temp.copy_shared_buffer(in);
+  }
+
+  Strides b_strides(in.ndim(), 0);
+  b_strides[axis] = 1;
+  array w_k_broadcast(temp.shape(), complex64, nullptr, {});
+  w_k_broadcast.copy_shared_buffer(w_k, b_strides, {}, w_k.data_size());
+  binary_op_gpu({temp, w_k_broadcast}, temp1, "Multiply", s);
+
+  std::vector<std::pair<int, int>> pads;
+  auto padded_shape = out.shape();
+  padded_shape[axis] = plan.bluestein_n;
+  array pad_temp(padded_shape, complex64, nullptr, {});
+  auto zero = array(complex64_t{0.0f, 0.0f});
+  copies.push_back(zero);
+  pad_gpu(temp1, zero, pad_temp, {(int)axis}, {0}, s);
+  copies.push_back(pad_temp);
+
+  array pad_temp1(padded_shape, complex64, nullptr, {});
+  fft_op(
+      pad_temp,
+      pad_temp1,
+      axis,
+      /*inverse=*/false,
+      /*real=*/false,
+      FourStepParams(),
+      /*inplace=*/false,
+      s);
+  copies.push_back(pad_temp1);
+
+  array w_q_broadcast(pad_temp1.shape(), complex64, nullptr, {});
+  w_q_broadcast.copy_shared_buffer(w_q, b_strides, {}, w_q.data_size());
+  binary_op_gpu_inplace({pad_temp1, w_q_broadcast}, pad_temp, "Multiply", s);
+
+  fft_op(
+      pad_temp,
+      pad_temp1,
+      axis,
+      /* inverse= */ true,
+      /* real= */ false,
+      FourStepParams(),
+      /*inplace=*/true,
+      s);
+
+  int offset = plan.bluestein_n - (2 * n - 1);
+  Shape starts(in.ndim(), 0);
+  Shape strides(in.ndim(), 1);
+  starts[axis] = plan.bluestein_n - offset - n;
+
+  array temp2(temp_shape, complex64, nullptr, {});
+  slice_gpu(pad_temp1, temp2, starts, strides, s);
+
+  binary_op_gpu_inplace({temp2, w_k_broadcast}, temp1, "Multiply", s);
+
+  if (real && !inverse) {
+    Shape rstarts(in.ndim(), 0);
+    Shape rstrides(in.ndim(), 1);
+    slice_gpu(temp1, out, rstarts, strides, s);
+  } else if (real && inverse) {
+    Strides b_strides(in.ndim(), 0);
+    auto inv_n = array({1.0f / n}, {1}, float32);
+    array temp_float(out.shape(), out.dtype(), nullptr, {});
+    copies.push_back(temp_float);
+    copies.push_back(inv_n);
+    copies.push_back(temp1);
+
+    copy_gpu(temp1, temp_float, CopyType::General, s);
+    binary_op_gpu({temp_float, inv_n}, out, "Multiply", s);
+  } else if (inverse) {
+    auto inv_n = array({1.0f / n}, {1}, complex64);
+    array temp3(temp_shape, complex64, nullptr, {});
+    unary_op_gpu({temp1}, temp3, "Conjugate", s);
+    binary_op_gpu({temp3, inv_n}, out, "Multiply", s);
+    copies.push_back(inv_n);
+    copies.push_back(temp1);
+    copies.push_back(temp3);
+  } else {
+    out.copy_shared_buffer(temp1);
+  }
+}
+
+void four_step_fft(
+    const array& in,
+    array& out,
+    size_t axis,
+    bool inverse,
+    bool real,
+    FFTPlan& plan,
+    std::vector<array>& copies,
+    const Stream& s,
+    bool in_place) {
+  if (plan.bluestein_n == -1) {
+    // Fast no transpose implementation for powers of 2.
+    FourStepParams four_step_params = {
+        /* required= */ true, /* first_step= */ true, plan.n1, plan.n2};
+    auto temp_shape = (real && inverse) ? out.shape() : in.shape();
+    array temp(temp_shape, complex64, nullptr, {});
+    fft_op(
+        in, temp, axis, inverse, real, four_step_params, /*inplace=*/false, s);
+    four_step_params.first_step = false;
+    fft_op(
+        temp,
+        out,
+        axis,
+        inverse,
+        real,
+        four_step_params,
+        /*inplace=*/in_place,
+        s);
+    copies.push_back(temp);
+  } else {
+    multi_upload_bluestein_fft(in, out, axis, inverse, real, plan, copies, s);
+  }
+}
+
+void fft_op(
+    const array& in,
+    array& out,
+    size_t axis,
+    bool inverse,
+    bool real,
+    const FourStepParams four_step_params,
+    bool inplace,
+    const Stream& s) {
+  auto& d = metal::device(s.device);
+
+  size_t n = out.dtype() == float32 ? out.shape(axis) : in.shape(axis);
+  if (n == 1) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  if (four_step_params.required) {
+    // Four Step FFT decomposes into two FFTs: n1 on columns, n2 on rows
+    n = four_step_params.first_step ? four_step_params.n1 : four_step_params.n2;
+  }
+
+  // Make sure that the array is contiguous and has stride 1 in the FFT dim
+  std::vector<array> copies;
+  auto check_input = [&axis, &copies, &s](const array& x) {
+    // TODO: Pass the strides to the kernel so
+    // we can avoid the copy when x is not contiguous.
+    bool no_copy = x.strides()[axis] == 1 &&
+        (x.flags().row_contiguous || x.flags().col_contiguous);
+    if (no_copy) {
+      return x;
+    } else {
+      array x_copy(x.shape(), x.dtype(), nullptr, {});
+      Strides strides;
+      int64_t cur_stride = x.shape(axis);
+      for (int a = 0; a < x.ndim(); a++) {
+        if (a == axis) {
+          strides.push_back(1);
+        } else {
+          strides.push_back(cur_stride);
+          cur_stride *= x.shape(a);
+        }
+      }
+
+      auto flags = x.flags();
+      auto [data_size, is_row_contiguous, is_col_contiguous] =
+          check_contiguity(x.shape(), strides);
+
+      flags.col_contiguous = is_col_contiguous;
+      flags.row_contiguous = is_row_contiguous;
+      flags.contiguous = data_size == x_copy.size();
+
+      x_copy.set_data(allocator::malloc(x.nbytes()), data_size, strides, flags);
+      copy_gpu_inplace(x, x_copy, CopyType::GeneralGeneral, s);
+      copies.push_back(x_copy);
+      return x_copy;
+    }
+  };
+  const array& in_contiguous = check_input(in);
+
+  // real to complex: n -> (n/2)+1
+  // complex to real: (n/2)+1 -> n
+  auto out_strides = in_contiguous.strides();
+  size_t out_data_size = in_contiguous.data_size();
+  if (in.shape(axis) != out.shape(axis)) {
+    for (int i = 0; i < out_strides.size(); i++) {
+      if (out_strides[i] != 1) {
+        out_strides[i] = out_strides[i] / in.shape(axis) * out.shape(axis);
+      }
+    }
+    out_data_size = out_data_size / in.shape(axis) * out.shape(axis);
+  }
+
+  auto plan = plan_fft(n);
+  if (plan.four_step) {
+    four_step_fft(in, out, axis, inverse, real, plan, copies, s, inplace);
+    d.add_temporaries(std::move(copies), s.index);
+    return;
+  }
+
+  // TODO: allow donation here
+  if (!inplace) {
+    out.set_data(
+        allocator::malloc(out.nbytes()),
+        out_data_size,
+        out_strides,
+        in_contiguous.flags());
+  }
+
+  auto radices = supported_radices();
+  int fft_size = plan.bluestein_n > 0 ? plan.bluestein_n : n;
+
+  // Setup function constants
+  bool power_of_2 = is_power_of_2(fft_size);
+
+  auto make_int = [](int* a, int i) {
+    return std::make_tuple(a, MTL::DataType::DataTypeInt, i);
+  };
+  auto make_bool = [](bool* a, int i) {
+    return std::make_tuple(a, MTL::DataType::DataTypeBool, i);
+  };
+
+  std::vector<MTLFC> func_consts = {
+      make_bool(&inverse, 0), make_bool(&power_of_2, 1)};
+
+  // Start of radix/rader step constants
+  int index = 4;
+  for (int i = 0; i < plan.stockham.size(); i++) {
+    func_consts.push_back(make_int(&plan.stockham[i], index));
+    index += 1;
+  }
+  for (int i = 0; i < plan.rader.size(); i++) {
+    func_consts.push_back(make_int(&plan.rader[i], index));
+    index += 1;
+  }
+  int elems_per_thread = compute_elems_per_thread(plan);
+  func_consts.push_back(make_int(&elems_per_thread, 2));
+
+  int rader_m = n / plan.rader_n;
+  func_consts.push_back(make_int(&rader_m, 3));
+
+  // The overall number of FFTs we're going to compute for this input
+  size_t size = out.dtype() == float32 ? out.size() : in.size();
+  if (real && inverse && four_step_params.required) {
+    size = out.size();
+  }
+  int total_batch_size = size / n;
+  int threads_per_fft = (fft_size + elems_per_thread - 1) / elems_per_thread;
+
+  // We batch among threadgroups for improved efficiency when n is small
+  int threadgroup_batch_size = std::max(MIN_THREADGROUP_MEM_SIZE / fft_size, 1);
+  if (four_step_params.required) {
+    // Require a threadgroup batch size of at least 4 for four step FFT
+    // so we can coalesce the memory accesses.
+    threadgroup_batch_size =
+        std::max(threadgroup_batch_size, MIN_COALESCE_WIDTH);
+  }
+  int threadgroup_mem_size = next_power_of_2(threadgroup_batch_size * fft_size);
+  // FFTs up to 2^20 are currently supported
+  assert(threadgroup_mem_size <= MAX_STOCKHAM_FFT_SIZE);
+
+  // ceil divide
+  int batch_size =
+      (total_batch_size + threadgroup_batch_size - 1) / threadgroup_batch_size;
+
+  if (real && !four_step_params.required) {
+    // We can perform 2 RFFTs at once so the batch size is halved.
+    batch_size = (batch_size + 2 - 1) / 2;
+  }
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto in_type_str = in.dtype() == float32 ? "float" : "float2";
+  auto out_type_str = out.dtype() == float32 ? "float" : "float2";
+  // Only required by four step
+  int step = -1;
+  {
+    std::ostringstream kname;
+    std::string inv_string = inverse ? "true" : "false";
+    std::string real_string = real ? "true" : "false";
+    std::string func_name;
+    if (plan.bluestein_n > 0) {
+      kname << "bluestein_fft_mem_" << threadgroup_mem_size << "_"
+            << in_type_str << "_" << out_type_str;
+      func_name = "bluestein_fft";
+    } else if (plan.rader_n > 1) {
+      kname << "rader_fft_mem_" << threadgroup_mem_size << "_" << in_type_str
+            << "_" << out_type_str;
+      func_name = "rader_fft";
+    } else if (four_step_params.required) {
+      step = four_step_params.first_step ? 0 : 1;
+      kname << "four_step_mem_" << threadgroup_mem_size << "_" << in_type_str
+            << "_" << out_type_str << "_" << step << "_" << real_string;
+      func_name = "four_step_fft";
+    } else {
+      kname << "fft_mem_" << threadgroup_mem_size << "_" << in_type_str << "_"
+            << out_type_str;
+      func_name = "fft";
+    }
+    std::string base_name = kname.str();
+    // We use a specialized kernel for each FFT size
+    kname << "_n" << fft_size << "_inv_" << inverse;
+    std::string hash_name = kname.str();
+    auto template_def = func_name == "four_step_fft" ? get_template_definition(
+                                                           base_name,
+                                                           func_name,
+                                                           threadgroup_mem_size,
+                                                           in_type_str,
+                                                           out_type_str,
+                                                           step,
+                                                           real)
+                                                     : get_template_definition(
+                                                           base_name,
+                                                           func_name,
+                                                           threadgroup_mem_size,
+                                                           in_type_str,
+                                                           out_type_str);
+    auto kernel =
+        get_fft_kernel(d, base_name, hash_name, func_consts, template_def);
+
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(in_contiguous, 0);
+    compute_encoder.set_output_array(out, 1);
+
+    if (plan.bluestein_n > 0) {
+      // Precomputed twiddle factors for Bluestein's
+      auto [w_k, w_q] = compute_bluestein_constants(n, plan.bluestein_n);
+      copies.push_back(w_q);
+      copies.push_back(w_k);
+
+      compute_encoder.set_input_array(w_q, 2); // w_q
+      compute_encoder.set_input_array(w_k, 3); // w_k
+      compute_encoder.set_bytes(n, 4);
+      compute_encoder.set_bytes(plan.bluestein_n, 5);
+      compute_encoder.set_bytes(total_batch_size, 6);
+    } else if (plan.rader_n > 1) {
+      auto [b_q, g_q, g_minus_q] = compute_raders_constants(plan.rader_n, s);
+      copies.push_back(b_q);
+      copies.push_back(g_q);
+      copies.push_back(g_minus_q);
+
+      compute_encoder.set_input_array(b_q, 2);
+      compute_encoder.set_input_array(g_q, 3);
+      compute_encoder.set_input_array(g_minus_q, 4);
+      compute_encoder.set_bytes(n, 5);
+      compute_encoder.set_bytes(total_batch_size, 6);
+      compute_encoder.set_bytes(plan.rader_n, 7);
+    } else if (four_step_params.required) {
+      compute_encoder.set_bytes(four_step_params.n1, 2);
+      compute_encoder.set_bytes(four_step_params.n2, 3);
+      compute_encoder.set_bytes(total_batch_size, 4);
+    } else {
+      compute_encoder.set_bytes(n, 2);
+      compute_encoder.set_bytes(total_batch_size, 3);
+    }
+
+    auto group_dims = MTL::Size(1, threadgroup_batch_size, threads_per_fft);
+    auto grid_dims =
+        MTL::Size(batch_size, threadgroup_batch_size, threads_per_fft);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+void fft_op(
+    const array& in,
+    array& out,
+    size_t axis,
+    bool inverse,
+    bool real,
+    bool inplace,
+    const Stream& s) {
+  fft_op(in, out, axis, inverse, real, FourStepParams(), inplace, s);
+}
+
+void nd_fft_op(
+    const array& in,
+    array& out,
+    const std::vector<size_t>& axes,
+    bool inverse,
+    bool real,
+    const Stream& s) {
+  // Perform ND FFT on GPU as a series of 1D FFTs
+  auto temp_shape = inverse ? in.shape() : out.shape();
+  std::vector<array> temp_arrs;
+  temp_arrs.emplace_back(temp_shape, complex64, nullptr, std::vector<array>{});
+  if (axes.size() > 2) {
+    temp_arrs.emplace_back(
+        temp_shape, complex64, nullptr, std::vector<array>{});
+  }
+  for (int i = axes.size() - 1; i >= 0; i--) {
+    int reverse_index = axes.size() - i - 1;
+    // For 5D and above, we don't want to reallocate our two temporary arrays
+    bool inplace = reverse_index >= 3 && i != 0;
+    // Opposite order for fft vs ifft
+    int index = inverse ? reverse_index : i;
+    size_t axis = axes[index];
+    // Mirror np.fft.(i)rfftn and perform a real transform
+    // only on the final axis.
+    bool step_real = (real && index == axes.size() - 1);
+    const array& in_arr = i == axes.size() - 1 ? in : temp_arrs[i % 2];
+    array& out_arr = i == 0 ? out : temp_arrs[1 - i % 2];
+    fft_op(in_arr, out_arr, axis, inverse, step_real, inplace, s);
+  }
+
+  auto& d = metal::device(s.device);
+  d.add_temporaries(std::move(temp_arrs), s.index);
+}
+
+void FFT::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& in = inputs[0];
+
+  if (axes_.size() > 1) {
+    nd_fft_op(in, out, axes_, inverse_, real_, s);
+  } else {
+    fft_op(in, out, axes_[0], inverse_, real_, /*inplace=*/false, s);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/hadamard.cpp b/macos/third_party/mlx/mlx/backend/metal/hadamard.cpp
new file mode 100644
index 0000000..65a8771
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/hadamard.cpp
@@ -0,0 +1,198 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/hadamard.h"
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/jit/includes.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+constexpr int MAX_HADAMARD_THREADS_PER_GROUP = 256;
+
+std::string gen_hadamard_codelet(int m) {
+  // Generate a O(m^2) hadamard codelet for a given M
+  // using the hadamard matrices above
+  //
+  // e.g. m = 2
+  // METAL_FUNC void hadamard_m(thread float *x) {
+  //   float tmp[2];
+  //   tmp[0] = + x[0] + x[1];
+  //   tmp[1] = + x[0] - x[1];
+  //   for (int i = 0; i < 2; i++) { x[i] = tmp[i]; }
+  // }
+  //
+  auto h_matrices = hadamard_matrices();
+  auto& matrix = h_matrices[m];
+
+  std::ostringstream source;
+  source << "METAL_FUNC void hadamard_radix_m(thread float *x) {" << std::endl;
+  if (m == 1) {
+    source << "}" << std::endl;
+    return source.str();
+  }
+  source << "  float tmp[" << m << "];" << std::endl;
+  auto start = 1;
+  auto end = matrix.find('\n', start);
+
+  int index = 0;
+  while (end != std::string_view::npos) {
+    source << "  tmp[" << index << "] = ";
+    auto row = matrix.substr(start, end - start);
+    for (int i = 0; i < row.length(); i++) {
+      source << " " << row[i] << " x[" << i << "]";
+    }
+    source << ";" << std::endl;
+    start = end + 1;
+    end = matrix.find('\n', start);
+    index++;
+  }
+  source << "  for (int i = 0; i < " << m << "; i++) { x[i] = tmp[i]; }"
+         << std::endl;
+  source << "}" << std::endl;
+  return source.str();
+}
+
+void hadamard_mn_contiguous(
+    const array& x,
+    array& y,
+    int m,
+    int n1,
+    int n2,
+    float scale,
+    metal::Device& d,
+    const Stream& s) {
+  int n = n1 * n2;
+  int read_width_n1 = n1 == 2 ? 2 : 4;
+  int read_width_n2 = n2 == 2 ? 2 : 4;
+  int read_width_m = (n == 2 || m == 28) ? 2 : 4;
+  int max_radix_1 = std::min(n1, 16);
+  int max_radix_2 = std::min(n2, 16);
+  float scale_n1 = 1.0;
+  float scale_n2 = (m == 1) ? scale : 1.0;
+  float scale_m = scale;
+
+  // n2 is a row contiguous power of 2 hadamard transform
+  MTL::Size group_dims_n2(n2 / max_radix_2, 1, 1);
+  MTL::Size grid_dims_n2(n2 / max_radix_2, x.size() / n2, 1);
+
+  // n1 is a strided power of 2 hadamard transform with stride n2
+  MTL::Size group_dims_n1(n1 / max_radix_1, 1, 1);
+  MTL::Size grid_dims_n1(n1 / max_radix_1, x.size() / n, n2);
+
+  // m is a strided hadamard transform with stride n = n1 * n2
+  MTL::Size group_dims_m(
+      std::min(n / read_width_m, MAX_HADAMARD_THREADS_PER_GROUP), 1, 1);
+  MTL::Size grid_dims_m(
+      group_dims_m.width, x.size() / m / read_width_m / group_dims_m.width, 1);
+
+  // Make the kernel
+  std::string kname;
+  kname.reserve(32);
+  concatenate(kname, "hadamard_", n * m, "_", type_to_name(x));
+  auto lib = d.get_library(kname, [&]() {
+    std::string kernel;
+    concatenate(
+        kernel,
+        metal::utils(),
+        gen_hadamard_codelet(m),
+        metal::hadamard(),
+        get_template_definition(
+            "n2" + kname,
+            "hadamard_n",
+            get_type_string(x.dtype()),
+            n2,
+            max_radix_2,
+            read_width_n2));
+    if (n1 > 1) {
+      kernel += get_template_definition(
+          "n1" + kname,
+          "hadamard_n",
+          get_type_string(x.dtype()),
+          n1,
+          max_radix_1,
+          read_width_n1,
+          n2);
+    }
+    if (m > 1) {
+      kernel += get_template_definition(
+          "m" + kname,
+          "hadamard_m",
+          get_type_string(x.dtype()),
+          n,
+          m,
+          read_width_m);
+    }
+    return kernel;
+  });
+
+  // Launch the strided transform for n1
+  if (n1 > 1) {
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel("n1" + kname, lib);
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(x, 0);
+    compute_encoder.set_output_array(y, 1);
+    compute_encoder.set_bytes(scale_n1, 2);
+    compute_encoder.dispatch_threads(grid_dims_n1, group_dims_n1);
+  }
+
+  // Launch the transform for n2
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel("n2" + kname, lib);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(n1 > 1 ? y : x, 0);
+  compute_encoder.set_output_array(y, 1);
+  compute_encoder.set_bytes(scale_n2, 2);
+  compute_encoder.dispatch_threads(grid_dims_n2, group_dims_n2);
+
+  // Launch the strided transform for m
+  if (m > 1) {
+    auto kernel = d.get_kernel("m" + kname, lib);
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(y, 0);
+    compute_encoder.set_output_array(y, 1);
+    compute_encoder.set_bytes(scale_m, 2);
+    compute_encoder.dispatch_threads(grid_dims_m, group_dims_m);
+  }
+}
+
+void Hadamard::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  // Split the hadamard transform so that all of them work on vectors smaller
+  // than 8192 elements.
+  //
+  // We decompose it in the following way:
+  //
+  // n = m * n1 * n2 = m * 2^k1 * 2^k2
+  //
+  // where m is in (1, 12, 20, 28) and n1 and n2 <= 8192
+  auto [n, m] = decompose_hadamard(in.shape().back());
+  int n1 = 1, n2 = n;
+  if (n > 8192) {
+    for (n2 = 2; n2 * n2 < n; n2 *= 2) {
+    }
+    n1 = n / n2;
+  }
+
+  if (in.flags().row_contiguous) {
+    if (in.is_donatable()) {
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(allocator::malloc(out.nbytes()));
+    }
+    hadamard_mn_contiguous(in, out, m, n1, n2, scale_, d, s);
+  } else {
+    copy_gpu(in, out, CopyType::General, s);
+    hadamard_mn_contiguous(out, out, m, n1, n2, scale_, d, s);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/indexing.cpp b/macos/third_party/mlx/mlx/backend/metal/indexing.cpp
new file mode 100644
index 0000000..48311dc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/indexing.cpp
@@ -0,0 +1,727 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <fmt/format.h>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/jit/includes.h"
+#include "mlx/backend/metal/jit/indexing.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/scan.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/dtype.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+constexpr int METAL_MAX_INDEX_ARRAYS = 20;
+
+std::pair<std::string, std::string> make_index_args(
+    const std::string& idx_type,
+    int nidx) {
+  std::ostringstream idx_args;
+  std::ostringstream idx_arr;
+  for (int i = 0; i < nidx; ++i) {
+    idx_args << fmt::format(
+        "const device {0} *idx{1} [[buffer({2})]],", idx_type, i, 20 + i);
+    idx_arr << fmt::format("idx{0}", i);
+    if (i < nidx - 1) {
+      idx_args << "\n";
+      idx_arr << ",";
+    }
+  }
+  return {idx_args.str(), idx_arr.str()};
+}
+
+void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& src = inputs[0];
+  int nidx = inputs.size() - 1;
+
+  if (nidx > METAL_MAX_INDEX_ARRAYS) {
+    std::ostringstream msg;
+    msg << "[Gather::eval_gpu] Gathering with more than "
+        << METAL_MAX_INDEX_ARRAYS << " index arrays not yet supported.";
+    throw std::runtime_error(msg.str());
+  }
+
+  out.set_data(allocator::malloc(out.nbytes()));
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  size_t slice_size = 1;
+  for (auto s : slice_sizes_) {
+    slice_size *= s;
+  }
+
+  bool large_index = nidx && inputs[1].size() > INT32_MAX;
+  bool large_src = src.size() > INT32_MAX;
+  bool large_out = out.size() > INT32_MAX;
+  bool large = large_index || large_src || large_out;
+
+  std::string idx_type_name = nidx ? type_to_name(inputs[1]) : "";
+
+  if (src.flags().row_contiguous && nidx == 1 && axes_[0] == 0 &&
+      inputs[1].flags().row_contiguous && slice_size == src.strides()[0]) {
+    int work_per_thread = (slice_size > 8 && src.dtype().size() < 4) ? 2 : 1;
+    auto& indices = inputs[1];
+    std::string kernel_name = fmt::format(
+        "gather_front{0}_{1}_{2}_{3}",
+        type_to_name(out),
+        idx_type_name,
+        large ? "int64_t" : "int",
+        work_per_thread);
+    std::string lib_name = kernel_name;
+
+    auto lib = d.get_library(lib_name, [&]() {
+      std::string kernel_source = metal::utils();
+      kernel_source += metal::gather_front();
+      kernel_source += get_template_definition(
+          kernel_name,
+          "gather_front",
+          get_type_string(out.dtype()),
+          get_type_string(indices.dtype()),
+          large ? "int64_t" : "int",
+          work_per_thread);
+
+      return kernel_source;
+    });
+
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kernel_name, lib);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    size_t dim_x = (slice_size + work_per_thread - 1) / work_per_thread;
+    size_t dim_y = indices.size();
+    auto group_dims = get_block_dims(dim_x, dim_y, 1);
+    MTL::Size grid_dims = MTL::Size(dim_x, dim_y, 1);
+
+    compute_encoder.set_input_array(src, 0);
+    compute_encoder.set_input_array(indices, 1);
+    compute_encoder.set_output_array(out, 2);
+    compute_encoder.set_bytes(slice_size, 3);
+    compute_encoder.set_bytes(src.shape(0), 4);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+    return;
+  }
+
+  int idx_ndim = nidx ? inputs[1].ndim() : 0;
+  size_t ndim = src.ndim();
+
+  std::string kernel_name = fmt::format(
+      "gather{0}{1}_{2}_{3}_{4}",
+      type_to_name(out),
+      idx_type_name,
+      nidx,
+      idx_ndim,
+      large ? "int64_t" : "int");
+  std::string lib_name = kernel_name;
+
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source = metal::utils();
+    kernel_source += metal::gather();
+    std::string out_type_str = get_type_string(out.dtype());
+    std::string idx_type_str =
+        nidx ? get_type_string(inputs[1].dtype()) : "bool";
+    auto [idx_args, idx_arr] = make_index_args(idx_type_str, nidx);
+
+    // Index dimension specializations
+    kernel_source += fmt::format(
+        gather_kernels,
+        type_to_name(out) + idx_type_name,
+        out_type_str,
+        idx_type_str,
+        nidx,
+        idx_args,
+        idx_arr,
+        idx_ndim,
+        large ? "int64_t" : "int");
+    return kernel_source;
+  });
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kernel_name, lib);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Launch 3D grid of threads
+  // First two dimensions for the indices, the last one for the slice
+  size_t dim0 = 1;
+  size_t dim1 = 1;
+  if (nidx) {
+    if (inputs[1].ndim() >= 1) {
+      dim0 = inputs[1].shape(0);
+    }
+    if (inputs[1].ndim() >= 2) {
+      dim1 = inputs[1].size() / dim0;
+    }
+  }
+  size_t dim2 = slice_size;
+  auto group_dims = get_block_dims(dim0, dim1, dim2);
+  MTL::Size grid_dims = MTL::Size(dim0, dim1, dim2);
+
+  // Collect all idx shapes and strides into one place
+  std::vector<int> idx_shapes;
+  std::vector<size_t> idx_strides;
+  std::vector<char> idx_contigs;
+  for (int i = 0; i < nidx; ++i) {
+    idx_shapes.insert(
+        idx_shapes.end(),
+        inputs[i + 1].shape().begin(),
+        inputs[i + 1].shape().end());
+    idx_strides.insert(
+        idx_strides.end(),
+        inputs[i + 1].strides().begin(),
+        inputs[i + 1].strides().end());
+    idx_contigs.push_back(inputs[i + 1].flags().row_contiguous);
+  }
+
+  // Set all the buffers
+  compute_encoder.set_input_array(src, 0);
+  compute_encoder.set_output_array(out, 1);
+
+  // Set source info
+  compute_encoder.set_vector_bytes(src.shape(), 2);
+  compute_encoder.set_vector_bytes(src.strides(), 3);
+  compute_encoder.set_bytes(ndim, 4);
+  compute_encoder.set_vector_bytes(slice_sizes_, 5);
+  compute_encoder.set_vector_bytes(axes_, 6);
+
+  // Set index info
+  //
+  // We don't need to check for empty idx_shapes because gather has a
+  // idx_ndim == 0 specialization
+  compute_encoder.set_vector_bytes(idx_shapes, 7);
+  compute_encoder.set_vector_bytes(idx_strides, 8);
+  compute_encoder.set_vector_bytes(idx_contigs, 9);
+  compute_encoder.set_bytes(idx_ndim, 10);
+
+  // Set index buffers
+  for (int i = 0; i < nidx; ++i) {
+    compute_encoder.set_input_array(inputs[i + 1], 20 + i);
+  }
+
+  // Launch grid
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
+  if (size_of(out.dtype()) == 8) {
+    std::ostringstream msg;
+    msg << "[Scatter::eval_gpu] Does not support " << out.dtype();
+    throw std::invalid_argument(msg.str());
+  }
+
+  int nidx = axes_.size();
+  if (nidx > METAL_MAX_INDEX_ARRAYS) {
+    std::ostringstream msg;
+    msg << "[Scatter::eval_gpu] Gathering with more than "
+        << METAL_MAX_INDEX_ARRAYS << " index arrays not yet supported.";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Copy src into out
+  CopyType copy_type;
+  if (inputs[0].data_size() == 1) {
+    copy_type = CopyType::Scalar;
+  } else if (inputs[0].flags().row_contiguous) {
+    copy_type = CopyType::Vector;
+  } else {
+    copy_type = CopyType::General;
+  }
+  copy_gpu(inputs[0], out, copy_type);
+
+  auto& upd = inputs.back();
+
+  // Empty update
+  if (upd.size() == 0) {
+    return;
+  }
+
+  // Get stream
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  int idx_ndim = nidx ? inputs[1].ndim() : 0;
+  size_t idx_size = nidx ? inputs[1].size() : 1;
+
+  auto idx_to_out = idx_size / out.size();
+  int nwork;
+  if (idx_ndim <= 1 || idx_to_out < 1) {
+    nwork = 1;
+  } else if (idx_to_out <= 4) {
+    nwork = 4;
+  } else if (idx_to_out < 16) {
+    nwork = 8;
+  } else if (idx_to_out < 32) {
+    nwork = 16;
+  } else {
+    nwork = 32;
+  }
+
+  std::string idx_type_name = nidx ? type_to_name(inputs[1]) : "";
+  std::string op_name;
+  switch (reduce_type_) {
+    case Scatter::None:
+      op_name = "none";
+      break;
+    case Scatter::Sum:
+      op_name = "sum";
+      break;
+    case Scatter::Prod:
+      op_name = "prod";
+      break;
+    case Scatter::Max:
+      op_name = "max";
+      break;
+    case Scatter::Min:
+      op_name = "min";
+      break;
+  }
+  auto upd_contig = upd.flags().row_contiguous;
+  bool large_out = out.size() > INT32_MAX;
+  bool large_idx = nidx && (inputs[1].size() > INT32_MAX);
+  bool large_upd = upd.size() > INT32_MAX;
+  bool large = large_out || large_idx || large_upd;
+  std::string kernel_name = fmt::format(
+      "scatter{0}{1}_{2}_{3}_{4}_nwork{5}_{6}",
+      type_to_name(out),
+      idx_type_name,
+      op_name,
+      nidx,
+      upd_contig ? "updc_true" : "updc_false",
+      nwork,
+      large ? "int64_t" : "int");
+  std::string lib_name = kernel_name;
+
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source = metal::utils();
+    concatenate(kernel_source, metal::reduce_utils(), metal::scatter());
+
+    std::string out_type_str = get_type_string(out.dtype());
+    std::string idx_type_str =
+        nidx ? get_type_string(inputs[1].dtype()) : "bool";
+    std::string op_type;
+    switch (reduce_type_) {
+      case Scatter::None:
+        op_type = "None";
+        break;
+      case Scatter::Sum:
+        op_type = "Sum<{0}>";
+        break;
+      case Scatter::Prod:
+        op_type = "Prod<{0}>";
+        break;
+      case Scatter::Max:
+        op_type = "Max<{0}>";
+        break;
+      case Scatter::Min:
+        op_type = "Min<{0}>";
+        break;
+    }
+    if (reduce_type_ != Scatter::None) {
+      op_type = fmt::format(fmt::runtime(op_type), out_type_str);
+    }
+    auto [idx_args, idx_arr] = make_index_args(idx_type_str, nidx);
+
+    kernel_source += fmt::format(
+        scatter_kernels,
+        type_to_name(out) + idx_type_name + "_" + op_name,
+        out_type_str,
+        idx_type_str,
+        op_type,
+        nidx,
+        idx_args,
+        idx_arr,
+        upd_contig,
+        nwork,
+        large ? "int64_t" : "int");
+    return kernel_source;
+  });
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kernel_name, lib);
+
+  size_t nthreads = upd.size();
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Set all the buffers
+  compute_encoder.set_input_array(upd, 1);
+  compute_encoder.set_output_array(out, 2);
+
+  // Set update info
+  size_t upd_ndim = upd.ndim();
+  size_t upd_size = 1;
+  for (int i = idx_ndim; i < upd.ndim(); ++i) {
+    upd_size *= upd.shape(i);
+  }
+  // Collect all idx shapes and strides into one place
+  Shape idx_shapes;
+  Strides idx_strides;
+  // To access .data() use char instead of bool
+  // bool is 1 byte in Metal so this is safe
+  std::vector<char> idx_contigs;
+  for (int i = 0; i < nidx; ++i) {
+    idx_shapes.insert(
+        idx_shapes.end(),
+        inputs[i + 1].shape().begin(),
+        inputs[i + 1].shape().end());
+    idx_strides.insert(
+        idx_strides.end(),
+        inputs[i + 1].strides().begin(),
+        inputs[i + 1].strides().end());
+    idx_contigs.push_back(inputs[i + 1].flags().row_contiguous);
+  }
+
+  if (upd_ndim == 0) {
+    // Need placeholders so Metal doesn't complain
+    int shape_ = 0;
+    int64_t stride_ = 0;
+    compute_encoder.set_bytes(shape_, 3);
+    compute_encoder.set_bytes(stride_, 4);
+  } else {
+    compute_encoder.set_vector_bytes(upd.shape(), 3);
+    compute_encoder.set_vector_bytes(upd.strides(), 4);
+  }
+  compute_encoder.set_bytes(upd_ndim, 5);
+  compute_encoder.set_bytes(upd_size, 6);
+
+  // Set output info
+  size_t out_ndim = out.ndim();
+  if (out_ndim == 0) {
+    // Need placeholders so Metal doesn't complain
+    int shape_ = 0;
+    int64_t stride_ = 0;
+    compute_encoder.set_bytes(shape_, 7);
+    compute_encoder.set_bytes(stride_, 8);
+  } else {
+    compute_encoder.set_vector_bytes(out.shape(), 7);
+    compute_encoder.set_vector_bytes(out.strides(), 8);
+  }
+  compute_encoder.set_bytes(out_ndim, 9);
+  compute_encoder.set_vector_bytes(axes_, 10);
+
+  // Set index info
+  if (idx_ndim == 0) {
+    // Add a 0 in idx_shapes and strides to avoid the missing buffer binding
+    // error in the metal API.
+    idx_shapes.push_back(0);
+    idx_strides.push_back(0);
+    idx_contigs.push_back(false);
+  }
+  compute_encoder.set_vector_bytes(idx_shapes, 11);
+  compute_encoder.set_vector_bytes(idx_strides, 12);
+  compute_encoder.set_vector_bytes(idx_contigs, 13);
+  compute_encoder.set_bytes(idx_ndim, 14);
+  compute_encoder.set_bytes(idx_size, 15);
+
+  // Set index buffers
+  for (int i = 0; i < nidx; ++i) {
+    compute_encoder.set_input_array(inputs[i + 1], 20 + i);
+  }
+
+  // Launch grid
+  auto grid_y = (nthreads / upd_size);
+  grid_y = (grid_y + nwork - 1) / nwork;
+  MTL::Size grid_dims = MTL::Size(upd_size, grid_y, 1);
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (thread_group_size != 1024) {
+    throw std::runtime_error("[Scatter::eval_gpu] Invalid number of threads");
+  }
+  MTL::Size group_dims = get_block_dims(upd_size, grid_y, 1);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& src = inputs[0];
+  auto& idx = inputs[1];
+
+  out.set_data(allocator::malloc(out.nbytes()));
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  size_t ndim = src.ndim();
+
+  bool large = idx.size() > INT32_MAX || src.size() > INT32_MAX;
+
+  std::string kernel_name = fmt::format(
+      "gather_axis{0}{1}_{2}",
+      type_to_name(out),
+      type_to_name(idx),
+      large ? "int64_t" : "int");
+  std::string lib_name = kernel_name;
+  kernel_name += src.flags().row_contiguous ? "c" : "nc";
+  kernel_name += idx.flags().row_contiguous ? "c" : "nc";
+
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source = metal::utils();
+    kernel_source += metal::gather_axis();
+    std::string out_type_str = get_type_string(out.dtype());
+    std::string idx_type_str = get_type_string(idx.dtype());
+    for (int i = 0; i < 4; ++i) {
+      bool sc = i & 1;
+      bool ic = i & 2;
+      kernel_source += get_template_definition(
+          lib_name + (sc ? "c" : "nc") + (ic ? "c" : "nc"),
+          "gather_axis",
+          out_type_str,
+          idx_type_str,
+          large ? "int64_t" : "int",
+          sc ? "true" : "false",
+          ic ? "true" : "false");
+    }
+    return kernel_source;
+  });
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kernel_name, lib);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Grid [size post, index size, size pre]
+  size_t size_pre = 1;
+  size_t size_post = 1;
+  for (int i = 0; i < axis_; ++i) {
+    size_pre *= idx.shape(i);
+  }
+  for (int i = axis_ + 1; i < idx.ndim(); ++i) {
+    size_post *= idx.shape(i);
+  }
+
+  int idx_ax_size = idx.shape(axis_);
+  auto group_dims = get_block_dims(size_post, idx_ax_size, size_pre);
+  MTL::Size grid_dims = MTL::Size(size_post, idx_ax_size, size_pre);
+
+  // Set all the buffers
+  compute_encoder.set_input_array(src, 0);
+  compute_encoder.set_input_array(idx, 1);
+  compute_encoder.set_output_array(out, 2);
+
+  // Set source info
+  compute_encoder.set_vector_bytes(remove_index(idx.shape(), axis_), 3);
+  compute_encoder.set_vector_bytes(remove_index(src.strides(), axis_), 4);
+  compute_encoder.set_vector_bytes(remove_index(idx.strides(), axis_), 5);
+  compute_encoder.set_bytes(ndim - 1, 6);
+  compute_encoder.set_bytes(axis_, 7);
+  compute_encoder.set_bytes(src.shape(axis_), 8);
+  compute_encoder.set_bytes(src.strides(axis_), 9);
+  compute_encoder.set_bytes(idx.strides(axis_), 10);
+
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& src = inputs[0];
+  auto& idx = inputs[1];
+  auto& upd = inputs[2];
+
+  // Copy src into out
+  CopyType copy_type;
+  if (src.data_size() == 1) {
+    copy_type = CopyType::Scalar;
+  } else if (src.flags().row_contiguous) {
+    copy_type = CopyType::Vector;
+  } else {
+    copy_type = CopyType::General;
+  }
+  copy_gpu(src, out, copy_type);
+
+  // Empty update
+  if (upd.size() == 0) {
+    return;
+  }
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  size_t ndim = src.ndim();
+
+  bool large = idx.size() > INT32_MAX || src.size() > INT32_MAX;
+
+  std::string op_name;
+  switch (reduce_type_) {
+    case ScatterAxis::None:
+      op_name = "none";
+      break;
+    case ScatterAxis::Sum:
+      op_name = "sum";
+      break;
+  }
+
+  std::string kernel_name = fmt::format(
+      "scatter_axis{0}{1}_{2}_{3}",
+      type_to_name(out),
+      type_to_name(idx),
+      op_name,
+      large ? "int64_t" : "int");
+  std::string lib_name = kernel_name;
+  kernel_name += upd.flags().row_contiguous ? "c" : "nc";
+  kernel_name += idx.flags().row_contiguous ? "c" : "nc";
+
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source = metal::utils();
+    kernel_source += metal::reduce_utils();
+    kernel_source += metal::scatter_axis();
+    std::string out_type_str = get_type_string(out.dtype());
+    std::string idx_type_str = get_type_string(idx.dtype());
+    std::string op_type;
+    switch (reduce_type_) {
+      case ScatterAxis::None:
+        op_type = "None";
+        break;
+      case ScatterAxis::Sum:
+        op_type = "Sum<" + out_type_str + ">";
+        break;
+    }
+
+    for (int i = 0; i < 4; ++i) {
+      bool uc = i & 1;
+      bool ic = i & 2;
+      kernel_source += get_template_definition(
+          lib_name + (uc ? "c" : "nc") + (ic ? "c" : "nc"),
+          "scatter_axis",
+          out_type_str,
+          idx_type_str,
+          large ? "int64_t" : "int",
+          op_type,
+          uc ? "true" : "false",
+          ic ? "true" : "false");
+    }
+    return kernel_source;
+  });
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kernel_name, lib);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Grid [size post, index size, size pre]
+  size_t size_pre = 1;
+  size_t size_post = 1;
+  for (int i = 0; i < axis_; ++i) {
+    size_pre *= idx.shape(i);
+  }
+  for (int i = axis_ + 1; i < idx.ndim(); ++i) {
+    size_post *= idx.shape(i);
+  }
+
+  int idx_ax_size = idx.shape(axis_);
+  auto group_dims = get_block_dims(size_post, idx_ax_size, size_pre);
+  MTL::Size grid_dims = MTL::Size(size_post, idx_ax_size, size_pre);
+
+  // Set all the buffers
+  compute_encoder.set_input_array(upd, 0);
+  compute_encoder.set_input_array(idx, 1);
+  compute_encoder.set_output_array(out, 2);
+
+  // Set source info
+  if (ndim > 1) {
+    compute_encoder.set_vector_bytes(remove_index(idx.shape(), axis_), 3);
+    compute_encoder.set_vector_bytes(remove_index(upd.strides(), axis_), 4);
+    compute_encoder.set_vector_bytes(remove_index(idx.strides(), axis_), 5);
+  } else {
+    // The following will be ignored in the kernel but we still have to set
+    // some value so that metal validation passes.
+    compute_encoder.set_vector_bytes(idx.shape(), 3);
+    compute_encoder.set_vector_bytes(upd.strides(), 4);
+    compute_encoder.set_vector_bytes(idx.strides(), 5);
+  }
+  compute_encoder.set_bytes(ndim - 1, 6);
+  compute_encoder.set_bytes(axis_, 7);
+  compute_encoder.set_bytes(out.shape(axis_), 8);
+  compute_encoder.set_bytes(upd.strides(axis_), 9);
+  compute_encoder.set_bytes(idx.strides(axis_), 10);
+
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void MaskedScatter::eval_gpu(const std::vector<array>& inputs, array& out) {
+  const array& dst = inputs[0];
+  const array& mask = inputs[1];
+  const array& src = inputs[2];
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  const size_t total = mask.size();
+  const CopyType ct = (total == 1)
+      ? CopyType::Scalar
+      : (dst.flags().row_contiguous ? CopyType::Vector : CopyType::General);
+  copy_gpu(dst, out, ct, s);
+  if (total == 0) {
+    return;
+  }
+
+  array mask_flat = flatten_in_eval(mask, 1, -1, s);
+  if (mask_flat.data<void>() != mask.data<void>()) {
+    d.add_temporary(mask_flat, s.index);
+  }
+
+  if (!mask_flat.flags().row_contiguous) {
+    mask_flat = contiguous_copy_gpu(mask_flat, s);
+    d.add_temporary(mask_flat, s.index);
+  }
+
+  // Prefix (exclusive) of mask → scatter_offsets
+  array scatter_offsets(mask_flat.shape(), uint32, nullptr, {});
+  scatter_offsets.set_data(allocator::malloc(scatter_offsets.nbytes()));
+  d.add_temporary(scatter_offsets, s.index);
+
+  scan_gpu_inplace(
+      mask_flat,
+      scatter_offsets,
+      Scan::Sum,
+      /*axis=*/1,
+      /*reverse=*/false,
+      /*inclusive=*/false,
+      s);
+
+  // Kernel selection/build
+  static constexpr std::string_view kBaseName = "masked_assign";
+  const std::string dtype_tag = type_to_name(out.dtype());
+  const std::string value_type = get_type_string(out.dtype());
+  const std::string contiguous =
+      (src.flags().row_contiguous) ? "true" : "false";
+  const std::string kernel_name =
+      fmt::format("{}_{}_{}", kBaseName, dtype_tag, contiguous);
+
+  auto lib = d.get_library(kernel_name, [&]() {
+    std::string source = metal::utils();
+    source += metal::masked_scatter();
+    source +=
+        fmt::format(masked_assign_kernel, kernel_name, value_type, contiguous);
+    return source;
+  });
+  auto kernel = d.get_kernel(kernel_name, lib);
+
+  // Binding
+  int bind_idx = 0;
+  const int ndim = static_cast<int>(src.ndim());
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(mask_flat, bind_idx++);
+  compute_encoder.set_input_array(scatter_offsets, bind_idx++);
+  compute_encoder.set_input_array(src, bind_idx++);
+  compute_encoder.set_output_array(out, bind_idx++);
+  compute_encoder.set_vector_bytes(src.shape(), bind_idx++);
+  compute_encoder.set_vector_bytes(src.strides(), bind_idx++);
+  compute_encoder.set_bytes(ndim, bind_idx++);
+  compute_encoder.set_bytes(src.size() / src.shape(0), bind_idx++);
+  compute_encoder.set_bytes(mask_flat.size() / mask.shape(0), bind_idx++);
+
+  // Dispatch
+  auto group_dims = get_block_dims(total, 1, 1);
+  MTL::Size grid_dims(total, 1, 1);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/jit/includes.h b/macos/third_party/mlx/mlx/backend/metal/jit/includes.h
new file mode 100644
index 0000000..a6ef0f1
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/jit/includes.h
@@ -0,0 +1,58 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+namespace mlx::core::metal {
+
+const char* utils();
+const char* binary_ops();
+const char* unary_ops();
+const char* ternary_ops();
+const char* reduce_utils();
+const char* gather();
+const char* scatter();
+const char* masked_scatter();
+
+const char* arange();
+const char* unary();
+const char* binary();
+const char* binary_two();
+const char* copy();
+const char* fft();
+const char* gather_axis();
+const char* gather_front();
+const char* hadamard();
+const char* logsumexp();
+const char* quantized_utils();
+const char* quantized();
+const char* fp_quantized();
+const char* ternary();
+const char* scan();
+const char* scatter_axis();
+const char* softmax();
+const char* sort();
+const char* reduce();
+
+const char* gemm();
+const char* steel_gemm_fused();
+const char* steel_gemm_masked();
+const char* steel_gemm_splitk();
+const char* steel_gemm_gather();
+const char* steel_gemm_segmented();
+const char* conv();
+const char* steel_conv();
+const char* steel_conv_general();
+const char* gemv_masked();
+const char* steel_attention();
+
+const char* gemm_nax();
+const char* steel_gemm_fused_nax();
+const char* steel_gemm_gather_nax();
+const char* steel_gemm_splitk_nax();
+
+const char* quantized_nax();
+const char* fp_quantized_nax();
+
+const char* steel_attention_nax();
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/jit/indexing.h b/macos/third_party/mlx/mlx/backend/metal/jit/indexing.h
new file mode 100644
index 0000000..fa141fc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/jit/indexing.h
@@ -0,0 +1,76 @@
+// Copyright © 2023-2024 Apple Inc.
+
+constexpr std::string_view gather_kernels = R"(
+[[kernel]] void gather{0}_{3}_{6}_{7}(
+    const device {1}* src [[buffer(0)]],
+    device {1}* out [[buffer(1)]],
+    const constant int* src_shape [[buffer(2)]],
+    const constant int64_t* src_strides [[buffer(3)]],
+    const constant size_t& src_ndim [[buffer(4)]],
+    const constant int* slice_sizes [[buffer(5)]],
+    const constant int* axes [[buffer(6)]],
+    const constant int* idx_shapes [[buffer(7)]],
+    const constant int64_t* idx_strides [[buffer(8)]],
+    const constant bool* idx_contigs [[buffer(9)]],
+    const constant int& idx_ndim [[buffer(10)]],
+    {4}
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {{
+  Indices<{2}, {3}> idxs{{
+    {{ {5} }}, idx_shapes, idx_strides, idx_contigs, idx_ndim}};
+
+  return gather_impl<{1}, {2}, {3}, {6}, {7}>(
+      src,
+      out,
+      src_shape,
+      src_strides,
+      src_ndim,
+      slice_sizes,
+      axes,
+      idxs,
+      index,
+      grid_dim);
+}}
+)";
+
+constexpr std::string_view scatter_kernels = R"(
+[[kernel]] void scatter{0}_{4}_updc_{7}_nwork{8}_{9}(
+    const device {1}* updates [[buffer(1)]],
+    device mlx_atomic<{1}>* out [[buffer(2)]],
+    const constant int* upd_shape [[buffer(3)]],
+    const constant int64_t* upd_strides [[buffer(4)]],
+    const constant size_t& upd_ndim [[buffer(5)]],
+    const constant size_t& upd_size [[buffer(6)]],
+    const constant int* out_shape [[buffer(7)]],
+    const constant int64_t* out_strides [[buffer(8)]],
+    const constant size_t& out_ndim [[buffer(9)]],
+    const constant int* axes [[buffer(10)]],
+    const constant int* idx_shapes [[buffer(11)]],
+    const constant int64_t* idx_strides [[buffer(12)]],
+    const constant bool* idx_contigs [[buffer(13)]],
+    const constant int& idx_ndim [[buffer(14)]],
+    const constant size_t& idx_size [[buffer(15)]],
+    {5}
+    uint2 gid [[thread_position_in_grid]]) {{
+  Indices<{2}, {4}> idxs{{ {{ {6} }}, idx_shapes, idx_strides, idx_contigs, idx_ndim}};
+
+  return scatter_impl<{1}, {2}, {3}, {4}, {7}, {8}, {9}>(
+      updates,
+      out,
+      upd_shape,
+      upd_strides,
+      upd_ndim,
+      upd_size,
+      out_shape,
+      out_strides,
+      out_ndim,
+      axes,
+      idx_size,
+      idxs,
+      gid);
+}}
+)";
+
+constexpr std::string_view masked_assign_kernel = R"(
+template [[host_name("{0}")]] [[kernel]] decltype(masked_assign_impl<{1}, {2}>) masked_assign_impl<{1}, {2}>;
+)";
diff --git a/macos/third_party/mlx/mlx/backend/metal/jit_kernels.cpp b/macos/third_party/mlx/mlx/backend/metal/jit_kernels.cpp
new file mode 100644
index 0000000..b657457
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/jit_kernels.cpp
@@ -0,0 +1,1118 @@
+// Copyright © 2024 Apple Inc.
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/metal/jit/includes.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+
+using namespace fmt::literals;
+
+namespace mlx::core {
+
+MTL::ComputePipelineState* get_arange_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out) {
+  auto lib = d.get_library(kernel_name, [&]() {
+    std::string kernel_source = metal::utils();
+    kernel_source += metal::arange();
+    kernel_source += get_template_definition(
+        kernel_name, "arange", get_type_string(out.dtype()));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_unary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    auto in_t = get_type_string(in_type);
+    auto out_t = get_type_string(out_type);
+    std::string kernel_source = metal::utils();
+    concatenate(kernel_source, metal::unary_ops(), metal::unary());
+    kernel_source +=
+        get_template_definition("v_" + lib_name, "unary_v", in_t, out_t, op, 1);
+    if (get_work_per_thread(in_type) > 1) {
+      kernel_source +=
+          get_template_definition("vn_" + lib_name, "unary_v", in_t, out_t, op);
+    }
+    kernel_source +=
+        get_template_definition("v2_" + lib_name, "unary_v2", in_t, out_t, op);
+    kernel_source += get_template_definition(
+        "gn1_" + lib_name, "unary_g", in_t, out_t, op, 1, "int");
+    kernel_source += get_template_definition(
+        "gn4large_" + lib_name, "unary_g", in_t, out_t, op, 4);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+void append_binary_kernels(
+    const std::string& lib_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op,
+    std::string& kernel_source) {
+  const std::array<std::pair<std::string, std::string>, 7> kernel_types = {{
+      {"ss", "binary_ss"},
+      {"vs2", "binary_vs2"},
+      {"sv2", "binary_sv2"},
+      {"vv2", "binary_vv2"},
+      {"g1large", "binary_g_nd1"},
+      {"g2large", "binary_g_nd2"},
+      {"g3large", "binary_g_nd3"},
+  }};
+  auto in_t = get_type_string(in_type);
+  auto out_t = get_type_string(out_type);
+
+  for (auto& [name, func] : kernel_types) {
+    kernel_source +=
+        get_template_definition(name + "_" + lib_name, func, in_t, out_t, op);
+  }
+  kernel_source += get_template_definition(
+      "vs_" + lib_name, "binary_vs", in_t, out_t, op, 1);
+  kernel_source += get_template_definition(
+      "sv_" + lib_name, "binary_sv", in_t, out_t, op, 1);
+  kernel_source += get_template_definition(
+      "vv_" + lib_name, "binary_vv", in_t, out_t, op, 1);
+
+  if (get_work_per_thread(in_type) > 1) {
+    kernel_source += get_template_definition(
+        "vsn_" + lib_name, "binary_vs", in_t, out_t, op);
+    kernel_source += get_template_definition(
+        "svn_" + lib_name, "binary_sv", in_t, out_t, op);
+    kernel_source += get_template_definition(
+        "vvn_" + lib_name, "binary_vv", in_t, out_t, op);
+  }
+
+  kernel_source += get_template_definition(
+      "g1_" + lib_name, "binary_g_nd1", in_t, out_t, op, "int");
+  kernel_source += get_template_definition(
+      "g2_" + lib_name, "binary_g_nd2", in_t, out_t, op, "int");
+  kernel_source += get_template_definition(
+      "g3_" + lib_name, "binary_g_nd3", in_t, out_t, op, "int");
+  kernel_source += get_template_definition(
+      "gn2_" + lib_name, "binary_g", in_t, out_t, op, 2, "int");
+  kernel_source += get_template_definition(
+      "gn4large_" + lib_name, "binary_g", in_t, out_t, op, 4);
+}
+
+MTL::ComputePipelineState* get_binary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    kernel_source = metal::utils();
+    concatenate(kernel_source, metal::binary_ops(), metal::binary());
+    append_binary_kernels(lib_name, in_type, out_type, op, kernel_source);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_binary_two_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source = metal::utils();
+    concatenate(kernel_source, metal::binary_ops(), metal::binary_two());
+    append_binary_kernels(lib_name, in_type, out_type, op, kernel_source);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_ternary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype type,
+    const char* op) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    auto t_str = get_type_string(type);
+    std::string kernel_source = metal::utils();
+    concatenate(kernel_source, metal::ternary_ops(), metal::ternary());
+    const std::array<std::pair<std::string, std::string>, 3> kernel_types = {{
+        {"g1large", "ternary_g_nd1"},
+        {"g2large", "ternary_g_nd2"},
+        {"g3large", "ternary_g_nd3"},
+    }};
+    for (auto& [name, func] : kernel_types) {
+      kernel_source +=
+          get_template_definition(name + "_" + lib_name, func, t_str, op);
+    }
+
+    kernel_source += get_template_definition(
+        "v2_" + lib_name, "ternary_v2", t_str, op, false, false);
+    kernel_source += get_template_definition(
+        "sv2_" + lib_name, "ternary_v2", t_str, op, true, false);
+    kernel_source += get_template_definition(
+        "vs2_" + lib_name, "ternary_v2", t_str, op, false, true);
+
+    if (get_work_per_thread(type) > 1) {
+      kernel_source += get_template_definition(
+          "vn_" + lib_name, "ternary_v", t_str, op, false, false);
+      kernel_source += get_template_definition(
+          "svn_" + lib_name, "ternary_v", t_str, op, true, false);
+      kernel_source += get_template_definition(
+          "vsn_" + lib_name, "ternary_v", t_str, op, false, true);
+    }
+
+    kernel_source += get_template_definition(
+        "v_" + lib_name, "ternary_v", t_str, op, false, false, 1);
+    kernel_source += get_template_definition(
+        "sv_" + lib_name, "ternary_v", t_str, op, true, false, 1);
+    kernel_source += get_template_definition(
+        "vs_" + lib_name, "ternary_v", t_str, op, false, true, 1);
+    kernel_source += get_template_definition(
+        "g1_" + lib_name, "ternary_g_nd1", t_str, op, "int");
+    kernel_source += get_template_definition(
+        "g2_" + lib_name, "ternary_g_nd2", t_str, op, "int");
+    kernel_source += get_template_definition(
+        "g3_" + lib_name, "ternary_g_nd3", t_str, op, "int");
+    kernel_source += get_template_definition(
+        "gn2_" + lib_name, "ternary_g", t_str, op, 2, "int");
+    kernel_source += get_template_definition(
+        "gn4large_" + lib_name, "ternary_g", t_str, op, 4);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source = metal::utils();
+    kernel_source += metal::copy();
+    auto in_type = get_type_string(in.dtype());
+    auto out_type = get_type_string(out.dtype());
+    kernel_source += get_template_definition(
+        "s_" + lib_name, "copy_s", in_type, out_type, 1);
+    kernel_source +=
+        get_template_definition("s2_" + lib_name, "copy_s2", in_type, out_type);
+    kernel_source += get_template_definition(
+        "v_" + lib_name, "copy_v", in_type, out_type, 1);
+    kernel_source +=
+        get_template_definition("v2_" + lib_name, "copy_v2", in_type, out_type);
+
+    if (get_work_per_thread(out.dtype()) > 1) {
+      kernel_source += get_template_definition(
+          "sn_" + lib_name, "copy_s", in_type, out_type);
+      kernel_source += get_template_definition(
+          "vn_" + lib_name, "copy_v", in_type, out_type);
+    }
+
+    kernel_source += get_template_definition(
+        "g1_" + lib_name, "copy_g_nd1", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "g2_" + lib_name, "copy_g_nd2", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "g3_" + lib_name, "copy_g_nd3", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "gn2_" + lib_name, "copy_g", in_type, out_type, 2, "int");
+    kernel_source += get_template_definition(
+        "gg1_" + lib_name, "copy_gg_nd1", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "gg2_" + lib_name, "copy_gg_nd2", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "gg3_" + lib_name, "copy_gg_nd3", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "ggn2_" + lib_name, "copy_gg", in_type, out_type, 2, "int");
+    kernel_source += get_template_definition(
+        "g1large_" + lib_name, "copy_g_nd1", in_type, out_type);
+    kernel_source += get_template_definition(
+        "g2large_" + lib_name, "copy_g_nd2", in_type, out_type);
+    kernel_source += get_template_definition(
+        "g3large_" + lib_name, "copy_g_nd3", in_type, out_type);
+    kernel_source += get_template_definition(
+        "gn4large_" + lib_name, "copy_g", in_type, out_type, 4);
+    kernel_source += get_template_definition(
+        "gg1large_" + lib_name, "copy_gg_nd1", in_type, out_type);
+    kernel_source += get_template_definition(
+        "gg2large_" + lib_name, "copy_gg_nd2", in_type, out_type);
+    kernel_source += get_template_definition(
+        "gg3large_" + lib_name, "copy_gg_nd3", in_type, out_type);
+    kernel_source += get_template_definition(
+        "ggn4large_" + lib_name, "copy_gg", in_type, out_type, 4);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_dynamic_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source = metal::utils();
+    kernel_source += metal::copy();
+    auto in_type = get_type_string(in.dtype());
+    auto out_type = get_type_string(out.dtype());
+    kernel_source += get_template_definition(
+        "gg1_" + lib_name, "copy_gg_dynamic_nd1", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "gg2_" + lib_name, "copy_gg_dynamic_nd2", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "gg3_" + lib_name, "copy_gg_dynamic_nd3", in_type, out_type, "int");
+    kernel_source += get_template_definition(
+        "ggn2_" + lib_name, "copy_gg_dynamic", in_type, out_type, 2, "int");
+    kernel_source += get_template_definition(
+        "gg1large_" + lib_name, "copy_gg_dynamic_nd1", in_type, out_type);
+    kernel_source += get_template_definition(
+        "gg2large_" + lib_name, "copy_gg_dynamic_nd2", in_type, out_type);
+    kernel_source += get_template_definition(
+        "gg3large_" + lib_name, "copy_gg_dynamic_nd3", in_type, out_type);
+    kernel_source += get_template_definition(
+        "ggn4large_" + lib_name, "copy_gg_dynamic", in_type, out_type, 4);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_softmax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool precise,
+    const array& out) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&] {
+    std::string kernel_source = metal::utils();
+    auto in_type = get_type_string(out.dtype());
+    auto acc_type = get_type_string(precise ? float32 : out.dtype());
+    kernel_source += metal::softmax();
+    kernel_source += get_template_definition(
+        "block_" + lib_name, "softmax_single_row", in_type, acc_type);
+    kernel_source += get_template_definition(
+        "looped_" + lib_name, "softmax_looped", in_type, acc_type);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_logsumexp_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&] {
+    auto t_str = get_type_string(out.dtype());
+    std::string kernel_source;
+    kernel_source = metal::utils();
+    kernel_source += metal::logsumexp();
+    kernel_source +=
+        get_template_definition("block_" + lib_name, "logsumexp", t_str);
+    kernel_source += get_template_definition(
+        "looped_" + lib_name, "logsumexp_looped", t_str);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_scan_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool reverse,
+    bool inclusive,
+    const std::string& reduce_type,
+    const array& in,
+    const array& out) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    auto out_type = get_type_string(out.dtype());
+    std::string op = "Cum" + reduce_type + "<" + out_type + ">";
+    op[3] = toupper(op[3]);
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::scan();
+    const std::array<std::pair<std::string, std::string>, 2> scan_kernels = {{
+        {"contig_", "contiguous_scan"},
+        {"strided_", "strided_scan"},
+    }};
+    for (auto& [prefix, kernel] : scan_kernels) {
+      kernel_source << get_template_definition(
+          prefix + lib_name,
+          kernel,
+          get_type_string(in.dtype()),
+          get_type_string(out.dtype()),
+          op,
+          in.itemsize() <= 4 ? 4 : 2,
+          inclusive,
+          reverse);
+    }
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    int bn,
+    int tn) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    auto in_type = get_type_string(in.dtype());
+    auto out_type = get_type_string(out.dtype());
+    kernel_source << metal::utils() << metal::sort();
+    for (bool is_argsort : {true, false}) {
+      std::string bool_string = is_argsort ? "true" : "false";
+      std::string func_string = is_argsort ? "carg_" : "c_";
+      kernel_source << get_template_definition(
+          func_string + lib_name,
+          "block_sort",
+          in_type,
+          out_type,
+          bool_string,
+          bn,
+          tn);
+      kernel_source << get_template_definition(
+          "n" + func_string + lib_name,
+          "block_sort_nc",
+          in_type,
+          out_type,
+          bool_string,
+          bn,
+          tn);
+    }
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_mb_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& idx,
+    int bn,
+    int tn) {
+  std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::sort();
+    std::array<std::pair<std::string, std::string>, 3> kernel_types = {
+        {{"sort_", "mb_block_sort"},
+         {"partition_", "mb_block_partition"},
+         {"merge_", "mb_block_merge"}}};
+    for (auto& [name, func] : kernel_types) {
+      kernel_source << get_template_definition(
+          name + lib_name,
+          func,
+          get_type_string(in.dtype()),
+          get_type_string(idx.dtype()),
+          "true",
+          bn,
+          tn);
+    }
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_reduce_init_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& func_name,
+    const std::string& op_name,
+    const Dtype& out_type) {
+  auto lib = d.get_library(kernel_name, [&]() {
+    std::string op_type = op_name;
+    op_type[0] = std::toupper(op_name[0]);
+    auto out_t = get_type_string(out_type);
+    std::string op = op_type + "<" + out_t + ">";
+    std::string kernel_source = metal::utils();
+    kernel_source += metal::reduce_utils();
+    kernel_source += metal::reduce();
+    kernel_source += get_template_definition(kernel_name, func_name, out_t, op);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_reduce_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& func_name,
+    const std::string& op_name,
+    const Dtype& in_type,
+    const Dtype& out_type,
+    const std::string& idx_t,
+    int ndim /* = -1 */,
+    int bm /* = -1 */,
+    int bn /* = -1 */) {
+  auto lib = d.get_library(kernel_name, [&]() {
+    std::string op_type = op_name;
+    op_type[0] = std::toupper(op_name[0]);
+    auto in_t = get_type_string(in_type);
+    auto out_t = get_type_string(out_type);
+    std::string op = op_type + "<" + out_t + ">";
+    std::string kernel_source = metal::utils();
+    concatenate(kernel_source, metal::reduce_utils(), metal::reduce());
+    if (bm >= 0) {
+      kernel_source += get_template_definition(
+          kernel_name, func_name, in_t, out_t, op, idx_t, ndim, bm, bn);
+    } else if (ndim >= 0) {
+      kernel_source += get_template_definition(
+          kernel_name, func_name, in_t, out_t, op, idx_t, ndim);
+    } else {
+      kernel_source += get_template_definition(
+          kernel_name, func_name, in_t, out_t, op, idx_t);
+    }
+    return kernel_source;
+  });
+  auto st = d.get_kernel(kernel_name, lib);
+  return st;
+}
+
+MTL::ComputePipelineState* get_steel_gemm_fused_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm()
+                  << metal::steel_gemm_fused()
+                  << get_template_definition(
+                         lib_name,
+                         "gemm",
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         transpose_a,
+                         transpose_b);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool mn_aligned,
+    bool k_aligned) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm()
+                  << metal::steel_gemm_splitk()
+                  << get_template_definition(
+                         lib_name,
+                         "gemm_splitk",
+                         get_type_string(in.dtype()),
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         transpose_a,
+                         transpose_b,
+                         mn_aligned,
+                         k_aligned);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_accum_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    bool axbpy) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm()
+                  << metal::steel_gemm_splitk()
+                  << get_template_definition(
+                         lib_name,
+                         axbpy ? "gemm_splitk_accum_axpby"
+                               : "gemm_splitk_accum",
+                         get_type_string(in.dtype()),
+                         get_type_string(out.dtype()));
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool mn_aligned,
+    bool k_aligned) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    auto out_mask_type = mask_out.has_value()
+        ? get_type_string((*mask_out).dtype())
+        : "nomask_t";
+    auto op_mask_type =
+        mask_op.has_value() ? get_type_string((*mask_op).dtype()) : "nomask_t";
+    kernel_source << metal::utils() << metal::gemm()
+                  << metal::steel_gemm_masked()
+                  << get_template_definition(
+                         lib_name,
+                         "block_masked_gemm",
+                         get_type_string(out.dtype()),
+                         out_mask_type,
+                         op_mask_type,
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         transpose_a,
+                         transpose_b,
+                         mn_aligned,
+                         k_aligned);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_gather_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm(),
+        metal::steel_gemm_gather(),
+        get_template_definition(
+            lib_name,
+            rhs ? "gather_mm_rhs" : "gather_mm",
+            get_type_string(out.dtype()),
+            bm,
+            bn,
+            bk,
+            wm,
+            wn,
+            transpose_a,
+            transpose_b));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_segmented_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm(),
+        metal::steel_gemm_segmented(),
+        get_template_definition(
+            lib_name,
+            "segmented_mm",
+            get_type_string(out.dtype()),
+            bm,
+            bn,
+            bk,
+            wm,
+            wn,
+            transpose_a,
+            transpose_b));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_gemv_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_mat,
+    int bm,
+    int bn,
+    int sm,
+    int sn,
+    int tm,
+    int tn,
+    bool contiguous) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    auto out_mask_type = mask_out.has_value()
+        ? get_type_string((*mask_out).dtype())
+        : "nomask_t";
+    auto op_mask_type =
+        mask_op.has_value() ? get_type_string((*mask_op).dtype()) : "nomask_t";
+    kernel_source << metal::utils() << metal::gemv_masked()
+                  << get_template_definition(
+                         lib_name,
+                         (transpose_mat) ? "gemv_t_masked" : "gemv_masked",
+                         get_type_string(out.dtype()),
+                         out_mask_type,
+                         op_mask_type,
+                         bm,
+                         bn,
+                         sm,
+                         sn,
+                         tm,
+                         tn,
+                         contiguous ? 0 : 1);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_steel_conv_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    int n_channel_specialization,
+    bool small_filter) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::conv() << metal::steel_conv()
+                  << get_template_definition(
+                         lib_name,
+                         "implicit_gemm_conv_2d",
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         n_channel_specialization,
+                         small_filter);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_steel_conv_general_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::conv()
+                  << metal::steel_conv_general()
+                  << get_template_definition(
+                         lib_name,
+                         "implicit_gemm_conv_2d_general",
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_fft_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const std::string& template_def) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    std::string kernel_string;
+    kernel_source << metal::fft() << template_def;
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_quantized_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm(),
+        metal::quantized_utils(),
+        (mode == "affine") ? metal::quantized() : metal::fp_quantized(),
+        template_def);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_gather_qmm_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source, metal::utils(), metal::quantized_utils(), metal::gemm());
+    bool is_affine = mode == "affine";
+    concatenate(
+        kernel_source,
+        is_affine ? metal::quantized() : metal::fp_quantized(),
+        get_template_definition(
+            lib_name,
+            (is_affine ? "affine" : "fp") + std::string("_gather_qmm_rhs"),
+            get_type_string(x.dtype()),
+            group_size,
+            bits,
+            bm,
+            bn,
+            bk,
+            wm,
+            wn,
+            transpose));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_fused_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm_nax()
+                  << metal::steel_gemm_fused_nax()
+                  << get_template_definition(
+                         lib_name,
+                         "gemm",
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         transpose_a,
+                         transpose_b);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm_nax(),
+        metal::steel_gemm_gather_nax(),
+        get_template_definition(
+            lib_name,
+            rhs ? "gather_mm_rhs_nax" : "gather_mm_nax",
+            get_type_string(out.dtype()),
+            bm,
+            bn,
+            bk,
+            wm,
+            wn,
+            transpose_a,
+            transpose_b));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm_nax()
+                  << metal::steel_gemm_splitk_nax()
+                  << get_template_definition(
+                         lib_name,
+                         "gemm_splitk_nax",
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         transpose_a,
+                         transpose_b);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm_nax(),
+        metal::quantized_utils(),
+        (mode == "affine") ? metal::quantized_nax() : metal::fp_quantized_nax(),
+        template_def);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_gather_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm_nax(),
+        metal::quantized_utils());
+    bool is_affine = mode == "affine";
+    concatenate(
+        kernel_source,
+        is_affine ? metal::quantized_nax() : metal::fp_quantized_nax(),
+        get_template_definition(
+            lib_name,
+            (is_affine ? "affine" : "fp") + std::string("_gather_qmm_rhs_nax"),
+            get_type_string(x.dtype()),
+            group_size,
+            bits,
+            bm,
+            bn,
+            bk,
+            wm,
+            wn,
+            transpose));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_attention_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::steel_attention(),
+        get_template_definition(
+            lib_name,
+            "attention",
+            get_type_string(q.dtype()),
+            bq,
+            bk,
+            bd,
+            wm,
+            wn,
+            get_type_string(m.dtype())));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_attention_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::steel_attention_nax(),
+        get_template_definition(
+            lib_name,
+            "attention_nax",
+            get_type_string(q.dtype()),
+            bq,
+            bk,
+            bd,
+            wm,
+            wn,
+            get_type_string(m.dtype())));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels.h b/macos/third_party/mlx/mlx/backend/metal/kernels.h
new file mode 100644
index 0000000..82aa4f9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels.h
@@ -0,0 +1,375 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <fmt/format.h>
+
+#include "mlx/array.h"
+#include "mlx/backend/metal/device.h"
+
+namespace mlx::core {
+
+MTL::ComputePipelineState* get_arange_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out);
+
+MTL::ComputePipelineState* get_unary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op);
+
+MTL::ComputePipelineState* get_binary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op);
+
+MTL::ComputePipelineState* get_binary_two_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype in_type,
+    Dtype out_type,
+    const char* op);
+
+MTL::ComputePipelineState* get_ternary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype type,
+    const char* op);
+
+MTL::ComputePipelineState* get_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out);
+
+MTL::ComputePipelineState* get_dynamic_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out);
+
+MTL::ComputePipelineState* get_softmax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool precise,
+    const array& out);
+
+MTL::ComputePipelineState* get_logsumexp_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out);
+
+MTL::ComputePipelineState* get_scan_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool reverse,
+    bool inclusive,
+    const std::string& reduce_type,
+    const array& in,
+    const array& out);
+
+MTL::ComputePipelineState* get_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    int bn,
+    int tn);
+
+MTL::ComputePipelineState* get_mb_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& idx,
+    int bn,
+    int tn);
+
+MTL::ComputePipelineState* get_reduce_init_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& func_name,
+    const std::string& op_name,
+    const Dtype& out_type);
+
+MTL::ComputePipelineState* get_reduce_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& func_name,
+    const std::string& op_name,
+    const Dtype& in_type,
+    const Dtype& out_type,
+    const std::string& idx_t,
+    int ndim = -1,
+    int bm = -1,
+    int bn = -1);
+
+MTL::ComputePipelineState* get_steel_gemm_fused_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool mn_aligned,
+    bool k_aligned);
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_accum_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& in,
+    const array& out,
+    bool axbpy);
+
+MTL::ComputePipelineState* get_steel_gemm_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool mn_aligned,
+    bool k_aligned);
+
+MTL::ComputePipelineState* get_steel_gemm_gather_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs);
+
+MTL::ComputePipelineState* get_steel_gemm_segmented_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_steel_conv_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    int n_channel_specialization,
+    bool small_filter);
+
+MTL::ComputePipelineState* get_gemv_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_mat,
+    int bm,
+    int bn,
+    int sm,
+    int sn,
+    int tm,
+    int tn,
+    bool contiguous);
+
+MTL::ComputePipelineState* get_steel_conv_general_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_fft_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const std::string& template_def);
+
+MTL::ComputePipelineState* get_quantized_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode);
+
+MTL::ComputePipelineState* get_gather_qmm_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose);
+
+MTL::ComputePipelineState* get_steel_gemm_fused_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs);
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode);
+
+MTL::ComputePipelineState* get_gather_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose);
+
+MTL::ComputePipelineState* get_steel_attention_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m);
+
+MTL::ComputePipelineState* get_steel_attention_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m);
+
+// Create a GPU kernel template definition for JIT compilation
+template <typename... Args>
+std::string get_template_definition(
+    std::string_view name,
+    std::string_view func,
+    Args... args) {
+  std::ostringstream s;
+  s << func << "<";
+  bool first = true;
+  auto add_arg = [&s, &first](const auto& arg) {
+    if (!first) {
+      s << ", ";
+    }
+    first = false;
+    s << arg;
+  };
+  (add_arg(args), ...);
+  s << ">";
+  return fmt::format(
+      "\ntemplate [[host_name(\"{0}\")]] [[kernel]] decltype({1}) {1};\n",
+      name,
+      s.str());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/metal/kernels/CMakeLists.txt
new file mode 100644
index 0000000..7d010e6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/CMakeLists.txt
@@ -0,0 +1,193 @@
+set(BASE_HEADERS
+    bf16.h
+    bf16_math.h
+    complex.h
+    defines.h
+    erf.h
+    expm1f.h
+    fp8.h
+    logging.h
+    utils.h)
+
+function(build_kernel_base TARGET SRCFILE DEPS)
+  set(METAL_FLAGS
+      -x
+      metal
+      -Wall
+      -Wextra
+      -fno-fast-math
+      -Wno-c++17-extensions
+      -Wno-c++20-extensions)
+  if(MLX_METAL_DEBUG)
+    set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
+  endif()
+  if(CMAKE_BUILD_TYPE STREQUAL "Debug" AND MLX_METAL_VERSION GREATER_EQUAL 320)
+    set(METAL_FLAGS ${METAL_FLAGS} -fmetal-enable-logging)
+  endif()
+  if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
+    set(METAL_FLAGS ${METAL_FLAGS}
+                    "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+  endif()
+  add_custom_command(
+    COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE}
+            -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air
+    DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS}
+    OUTPUT ${TARGET}.air
+    COMMENT "Building ${TARGET}.air"
+    VERBATIM)
+endfunction(build_kernel_base)
+
+function(build_kernel KERNEL)
+  set(SRCFILE ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL}.metal)
+  cmake_path(GET KERNEL STEM TARGET)
+  build_kernel_base(${TARGET} ${SRCFILE} "${ARGN}")
+  set(KERNEL_AIR
+      ${TARGET}.air ${KERNEL_AIR}
+      PARENT_SCOPE)
+endfunction(build_kernel)
+
+build_kernel(arg_reduce)
+build_kernel(conv steel/conv/params.h)
+build_kernel(gemv steel/utils.h)
+build_kernel(layer_norm)
+build_kernel(random)
+build_kernel(rms_norm)
+build_kernel(rope)
+build_kernel(scaled_dot_product_attention sdpa_vector.h)
+if(MLX_METAL_VERSION GREATER_EQUAL 320)
+  build_kernel(fence)
+endif()
+
+set(STEEL_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/conv/conv.h
+    steel/conv/loader.h
+    steel/conv/loaders/loader_channel_l.h
+    steel/conv/loaders/loader_channel_n.h
+    steel/conv/loaders/loader_general.h
+    steel/conv/kernels/steel_conv.h
+    steel/conv/kernels/steel_conv_general.h
+    steel/gemm/gemm.h
+    steel/gemm/mma.h
+    steel/gemm/loader.h
+    steel/gemm/params.h
+    steel/gemm/transforms.h
+    steel/gemm/kernels/steel_gemm_fused.h
+    steel/gemm/kernels/steel_gemm_gather.h
+    steel/gemm/kernels/steel_gemm_masked.h
+    steel/gemm/kernels/steel_gemm_segmented.h
+    steel/gemm/kernels/steel_gemm_splitk.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h)
+
+set(STEEL_ATTN_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/gemm/gemm.h
+    steel/gemm/mma.h
+    steel/gemm/loader.h
+    steel/gemm/transforms.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/attn/attn.h
+    steel/attn/loader.h
+    steel/attn/mma.h
+    steel/attn/params.h
+    steel/attn/transforms.h
+    steel/attn/kernels/steel_attention.h)
+
+set(STEEL_NAX_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/gemm/params.h
+    steel/gemm/transforms.h
+    steel/gemm/nax.h
+    steel/gemm/gemm_nax.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/gemm/kernels/steel_gemm_fused_nax.h
+    steel/gemm/kernels/steel_gemm_gather_nax.h
+    steel/gemm/kernels/steel_gemm_splitk_nax.h)
+
+set(STEEL_NAX_ATTN_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/attn/nax.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/attn/params.h
+    steel/attn/kernels/steel_attention_nax.h)
+
+if(NOT MLX_METAL_JIT)
+  build_kernel(arange arange.h)
+  build_kernel(binary binary.h binary_ops.h)
+  build_kernel(binary_two binary_two.h)
+  build_kernel(copy copy.h)
+  build_kernel(fft fft.h fft/radix.h fft/readwrite.h)
+  build_kernel(
+    reduce
+    atomic.h
+    reduction/ops.h
+    reduction/reduce_init.h
+    reduction/reduce_all.h
+    reduction/reduce_col.h
+    reduction/reduce_row.h)
+  build_kernel(quantized quantized.h quantized_utils.h ${STEEL_HEADERS})
+  build_kernel(fp_quantized fp4.h fp8.h fp_quantized.h quantized_utils.h
+               ${STEEL_HEADERS})
+  build_kernel(scan scan.h)
+  build_kernel(softmax softmax.h)
+  build_kernel(logsumexp logsumexp.h)
+  build_kernel(sort sort.h)
+  build_kernel(ternary ternary.h ternary_ops.h)
+  build_kernel(unary unary.h unary_ops.h)
+  build_kernel(steel/conv/kernels/steel_conv ${STEEL_HEADERS})
+  build_kernel(steel/conv/kernels/steel_conv_general ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_fused ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_gather ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_masked ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_splitk ${STEEL_HEADERS})
+  build_kernel(steel/gemm/kernels/steel_gemm_segmented ${STEEL_HEADERS})
+  build_kernel(gemv_masked steel/utils.h)
+  build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS})
+
+  if((MLX_METAL_VERSION GREATER_EQUAL 400) AND (MACOS_SDK_VERSION GREATER_EQUAL
+                                                26.2))
+
+    build_kernel(steel/gemm/kernels/steel_gemm_fused_nax ${STEEL_NAX_HEADERS})
+    build_kernel(steel/gemm/kernels/steel_gemm_gather_nax ${STEEL_NAX_HEADERS})
+    build_kernel(steel/gemm/kernels/steel_gemm_splitk_nax ${STEEL_NAX_HEADERS})
+
+    build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS})
+    build_kernel(fp_quantized_nax fp4.h fp8.h fp_quantized_nax.h
+                 ${STEEL_NAX_HEADERS})
+
+    build_kernel(steel/attn/kernels/steel_attention_nax
+                 ${STEEL_NAX_ATTN_HEADERS})
+
+  else()
+    target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX)
+  endif()
+
+endif()
+
+add_custom_command(
+  OUTPUT ${MLX_METAL_PATH}/mlx.metallib
+  COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
+          ${MLX_METAL_PATH}/mlx.metallib
+  DEPENDS ${KERNEL_AIR}
+  COMMENT "Building mlx.metallib"
+  VERBATIM)
+
+add_custom_target(mlx-metallib DEPENDS ${MLX_METAL_PATH}/mlx.metallib)
+
+add_dependencies(mlx mlx-metallib)
+
+# Install metallib
+include(GNUInstallDirs)
+
+install(
+  FILES ${MLX_METAL_PATH}/mlx.metallib
+  DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  COMPONENT metallib)
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/arange.h b/macos/third_party/mlx/mlx/backend/metal/kernels/arange.h
new file mode 100644
index 0000000..5448fe9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/arange.h
@@ -0,0 +1,9 @@
+// Copyright © 2023-2024 Apple Inc.
+template <typename T>
+[[kernel]] void arange(
+    constant const T& start,
+    constant const T& step,
+    device T* out,
+    uint index [[thread_position_in_grid]]) {
+  out[index] = start + index * step;
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/arange.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/arange.metal
new file mode 100644
index 0000000..fb56c1c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/arange.metal
@@ -0,0 +1,20 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/arange.h"
+
+#define instantiate_arange(tname, type)                                 \
+  instantiate_kernel("arange" #tname, arange, type)
+
+instantiate_arange(uint8, uint8_t)
+instantiate_arange(uint16, uint16_t)
+instantiate_arange(uint32, uint32_t)
+instantiate_arange(uint64, uint64_t)
+instantiate_arange(int8, int8_t)
+instantiate_arange(int16, int16_t)
+instantiate_arange(int32, int32_t)
+instantiate_arange(int64, int64_t)
+instantiate_arange(float16, half)
+instantiate_arange(float32, float)
+instantiate_arange(bfloat16, bfloat16_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/arg_reduce.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/arg_reduce.metal
new file mode 100644
index 0000000..4a83d8e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/arg_reduce.metal
@@ -0,0 +1,182 @@
+// Copyright © 2023 Apple Inc.
+
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+template <typename U>
+struct IndexValPair {
+  uint32_t index;
+  U val;
+};
+
+template <typename U>
+struct ArgMin {
+  static constexpr constant U init = Limits<U>::max;
+
+  IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> current) {
+    if (best.val > current.val ||
+        (best.val == current.val && best.index > current.index)) {
+      return current;
+    } else {
+      return best;
+    }
+  }
+
+  template <int N>
+  IndexValPair<U>
+  reduce_many(IndexValPair<U> best, thread U* vals, uint32_t offset) {
+    for (int i = 0; i < N; i++) {
+      if (vals[i] < best.val) {
+        best.val = vals[i];
+        best.index = offset + i;
+      }
+    }
+    return best;
+  }
+};
+
+template <typename U>
+struct ArgMax {
+  static constexpr constant U init = Limits<U>::min;
+
+  IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> current) {
+    if (best.val < current.val ||
+        (best.val == current.val && best.index > current.index)) {
+      return current;
+    } else {
+      return best;
+    }
+  }
+
+  template <int N>
+  IndexValPair<U>
+  reduce_many(IndexValPair<U> best, thread U* vals, uint32_t offset) {
+    for (int i = 0; i < N; i++) {
+      if (vals[i] > best.val) {
+        best.val = vals[i];
+        best.index = offset + i;
+      }
+    }
+    return best;
+  }
+};
+
+template <typename U>
+IndexValPair<U> simd_shuffle_down(IndexValPair<U> data, uint16_t delta) {
+  return IndexValPair<U>{
+      simd_shuffle_down(data.index, delta), simd_shuffle_down(data.val, delta)};
+}
+
+template <typename T, typename Op, int N_READS = 4>
+[[kernel]] void arg_reduce_general(
+    const device T* in [[buffer(0)]],
+    device uint32_t* out [[buffer(1)]],
+    const constant int* shape [[buffer(2)]],
+    const constant int64_t* in_strides [[buffer(3)]],
+    const constant int64_t* out_strides [[buffer(4)]],
+    const constant size_t& ndim [[buffer(5)]],
+    const constant int64_t& axis_stride [[buffer(6)]],
+    const constant size_t& axis_size [[buffer(7)]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 gsize [[threads_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint simd_size [[threads_per_simdgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  // Shapes and strides *do not* contain the reduction axis. The reduction size
+  // and stride are provided in axis_stride and axis_size.
+  //
+  // Note: in shape == out shape with this convention.
+  //
+  // The sketch of the kernel is as follows.
+  //    1. Launch prod(shape) * thread_group_size threads.
+  //    2. Loop ceildiv(axis_size / lsize) times
+  //    3. Read input values
+  //    4. Reduce among them and go to 3
+  //    4. Reduce in each simd_group
+  //    6. Write in the thread local memory
+  //    6. Reduce them across thread group
+  //    7. Write the output without need for atomic
+  Op op;
+
+  // Compute the input/output index. There is one beginning and one output for
+  // the whole threadgroup.
+  int64_t row_idx = gid.y + static_cast<int64_t>(gsize.y) * gid.z;
+  auto in_idx = elem_to_loc(row_idx, shape, in_strides, ndim);
+  auto out_idx = elem_to_loc(row_idx, shape, out_strides, ndim);
+
+  IndexValPair<T> best{0, Op::init};
+
+  threadgroup IndexValPair<T> local_data[32];
+
+  // Loop over the reduction axis in lsize*N_READS buckets
+  for (uint r = 0; r < ceildiv(axis_size, N_READS * lsize.x); r++) {
+    // Read the current value
+    uint32_t current_index = r * lsize.x * N_READS + lid.x * N_READS;
+    uint32_t offset = current_index;
+    const device T* current_in = in + in_idx + current_index * axis_stride;
+    T vals[N_READS];
+    for (int i = 0; i < N_READS; i++) {
+      vals[i] = (current_index < axis_size) ? *current_in : T(Op::init);
+      current_index++;
+      current_in += axis_stride;
+    }
+    best = op.template reduce_many<N_READS>(best, vals, offset);
+  }
+  // At this point we have reduced the axis into thread group best values so we
+  // need to reduce across the thread group.
+
+  // First per simd reduction.
+  for (uint offset = simd_size / 2; offset > 0; offset /= 2) {
+    IndexValPair<T> neighbor = simd_shuffle_down(best, offset);
+    best = op.reduce(best, neighbor);
+  }
+
+  // Write to the threadgroup memory
+  if (simd_lane_id == 0) {
+    local_data[simd_group_id] = best;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id != 0) {
+    return;
+  }
+
+  // Read the appropriate value from local data and perform one simd reduction
+  uint simd_groups = ceildiv(lsize.x, simd_size);
+  if (simd_lane_id < simd_groups) {
+    best = local_data[simd_lane_id];
+  }
+  for (uint offset = simd_size / 2; offset > 0; offset /= 2) {
+    IndexValPair<T> neighbor = simd_shuffle_down(best, offset);
+    best = op.reduce(best, neighbor);
+  }
+
+  // Finally write the output
+  if (lid.x == 0) {
+    out[out_idx] = best.index;
+  }
+}
+
+// clang-format off
+#define instantiate_arg_reduce(name, itype)                      \
+  instantiate_kernel(                                            \
+      "argmin_" #name, arg_reduce_general, itype, ArgMin<itype>) \
+  instantiate_kernel(                                            \
+      "argmax_" #name, arg_reduce_general, itype, ArgMax<itype>)
+
+instantiate_arg_reduce(bool_, bool)
+instantiate_arg_reduce(uint8, uint8_t)
+instantiate_arg_reduce(uint16, uint16_t)
+instantiate_arg_reduce(uint32, uint32_t)
+instantiate_arg_reduce(uint64, uint64_t)
+instantiate_arg_reduce(int8, int8_t)
+instantiate_arg_reduce(int16, int16_t)
+instantiate_arg_reduce(int32, int32_t)
+instantiate_arg_reduce(int64, int64_t)
+instantiate_arg_reduce(float16, half)
+instantiate_arg_reduce(float32, float)
+instantiate_arg_reduce(bfloat16, bfloat16_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/atomic.h b/macos/third_party/mlx/mlx/backend/metal/kernels/atomic.h
new file mode 100644
index 0000000..93952c2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/atomic.h
@@ -0,0 +1,345 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <metal_atomic>
+#include <metal_stdlib>
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// Atomic utils
+///////////////////////////////////////////////////////////////////////////////
+
+#pragma METAL internals : enable
+template <typename T>
+constexpr constant bool is_metal_atomic = _disjunction<
+    is_same<T, int>,
+    is_same<T, uint>,
+    is_same<T, ulong>,
+    is_same<T, float>>::value;
+
+#pragma METAL internals : disable
+
+template <typename T, typename = void>
+struct mlx_atomic {
+  atomic<uint> val;
+};
+
+template <typename T>
+struct mlx_atomic<T, enable_if_t<is_metal_atomic<T>>> {
+  atomic<T> val;
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Native metal atomics
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC T
+mlx_atomic_load_explicit(device mlx_atomic<T>* object, size_t offset) {
+  return atomic_load_explicit(&(object[offset].val), memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC void
+mlx_atomic_store_explicit(device mlx_atomic<T>* object, T val, size_t offset) {
+  atomic_store_explicit(&(object[offset].val), val, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_and_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  atomic_fetch_and_explicit(&(object[offset].val), val, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_or_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  atomic_fetch_or_explicit(&(object[offset].val), val, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_min_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  atomic_fetch_min_explicit(&(object[offset].val), val, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_max_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  atomic_fetch_max_explicit(&(object[offset].val), val, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_add_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  atomic_fetch_add_explicit(&(object[offset].val), val, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_mul_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  T expected = mlx_atomic_load_explicit(object, offset);
+  while (!mlx_atomic_compare_exchange_weak_explicit(
+      object, &expected, val * expected, offset)) {
+  }
+}
+
+template <typename T, enable_if_t<is_metal_atomic<T>, bool> = true>
+METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit(
+    device mlx_atomic<T>* object,
+    thread T* expected,
+    T val,
+    size_t offset) {
+  return atomic_compare_exchange_weak_explicit(
+      &(object[offset].val),
+      expected,
+      val,
+      memory_order_relaxed,
+      memory_order_relaxed);
+}
+
+// Specialization for float since it does not atomic_fetch_min_explicit
+template <>
+METAL_FUNC void mlx_atomic_fetch_min_explicit<float>(
+    device mlx_atomic<float>* object,
+    float val,
+    size_t offset) {
+  float expected = mlx_atomic_load_explicit(object, offset);
+  while (val < expected) {
+    if (mlx_atomic_compare_exchange_weak_explicit(
+            object, &expected, val, offset)) {
+      return;
+    }
+  }
+}
+
+// Specialization for float since it does not atomic_fetch_max_explicit
+template <>
+METAL_FUNC void mlx_atomic_fetch_max_explicit<float>(
+    device mlx_atomic<float>* object,
+    float val,
+    size_t offset) {
+  float expected = mlx_atomic_load_explicit(object, offset);
+  while (val > expected) {
+    if (mlx_atomic_compare_exchange_weak_explicit(
+            object, &expected, val, offset)) {
+      return;
+    }
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Custom atomics
+///////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename T>
+constexpr constant uint packing_size = sizeof(uint) / sizeof(T);
+
+template <typename T>
+union uint_or_packed {
+  T val[packing_size<T>];
+  uint bits;
+};
+
+template <typename T, typename Op>
+struct mlx_atomic_update_helper {
+  uint operator()(uint_or_packed<T> init, T update, size_t elem_offset) {
+    Op op;
+    init.val[elem_offset] = op(update, init.val[elem_offset]);
+    return init.bits;
+  }
+};
+
+template <typename T, typename Op>
+METAL_FUNC void mlx_atomic_update_and_store(
+    device mlx_atomic<T>* object,
+    T update,
+    size_t offset) {
+  size_t pack_offset = offset / packing_size<T>;
+  size_t elem_offset = offset % packing_size<T>;
+
+  mlx_atomic_update_helper<T, Op> helper;
+  uint_or_packed<T> expected;
+  expected.bits =
+      atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed);
+
+  while (Op::condition(update, expected.val[elem_offset]) &&
+         !mlx_atomic_compare_exchange_weak_explicit(
+             object,
+             &(expected.bits),
+             helper(expected, update, elem_offset),
+             pack_offset)) {
+  }
+}
+
+template <typename T>
+struct __None {
+  static bool condition(T a, T b) {
+#pragma unused(a)
+#pragma unused(b)
+    return true;
+  }
+
+  T operator()(T a, T b) {
+#pragma unused(b)
+    return a;
+  }
+};
+
+template <typename T>
+struct __Add {
+  static bool condition(T a, T b) {
+#pragma unused(a)
+#pragma unused(b)
+    return true;
+  }
+
+  T operator()(T a, T b) {
+    return a + b;
+  }
+};
+
+template <typename T>
+struct __Mul {
+  static bool condition(T a, T b) {
+#pragma unused(a)
+    return b != 0;
+  }
+
+  T operator()(T a, T b) {
+    return a * b;
+  }
+};
+
+template <typename T>
+struct __Max {
+  static bool condition(T a, T b) {
+    return a > b;
+  }
+
+  T operator()(T a, T b) {
+    return max(a, b);
+  }
+};
+
+template <typename T>
+struct __Min {
+  static bool condition(T a, T b) {
+    return a < b;
+  }
+
+  T operator()(T a, T b) {
+    return min(a, b);
+  }
+};
+
+} // namespace
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC T
+mlx_atomic_load_explicit(device mlx_atomic<T>* object, size_t offset) {
+  size_t pack_offset = offset / sizeof(T);
+  size_t elem_offset = offset % sizeof(T);
+  uint_or_packed<T> packed_val;
+  packed_val.bits =
+      atomic_load_explicit(&(object[pack_offset].val), memory_order_relaxed);
+  return packed_val.val[elem_offset];
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC void
+mlx_atomic_store_explicit(device mlx_atomic<T>* object, T val, size_t offset) {
+  mlx_atomic_update_and_store<T, __None<T>>(object, val, offset);
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_and_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  size_t pack_offset = offset / packing_size<T>;
+  size_t elem_offset = offset % packing_size<T>;
+  uint_or_packed<T> identity;
+  identity.bits = __UINT32_MAX__;
+  identity.val[elem_offset] = val;
+
+  atomic_fetch_and_explicit(
+      &(object[pack_offset].val), identity.bits, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_or_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  size_t pack_offset = offset / packing_size<T>;
+  size_t elem_offset = offset % packing_size<T>;
+  uint_or_packed<T> identity;
+  identity.bits = 0;
+  identity.val[elem_offset] = val;
+
+  atomic_fetch_or_explicit(
+      &(object[pack_offset].val), identity.bits, memory_order_relaxed);
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_min_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  mlx_atomic_update_and_store<T, __Min<T>>(object, val, offset);
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_max_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  mlx_atomic_update_and_store<T, __Max<T>>(object, val, offset);
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_add_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  mlx_atomic_update_and_store<T, __Add<T>>(object, val, offset);
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC void mlx_atomic_fetch_mul_explicit(
+    device mlx_atomic<T>* object,
+    T val,
+    size_t offset) {
+  mlx_atomic_update_and_store<T, __Mul<T>>(object, val, offset);
+}
+
+template <typename T, enable_if_t<!is_metal_atomic<T>, bool> = true>
+METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit(
+    device mlx_atomic<T>* object,
+    thread uint* expected,
+    uint val,
+    size_t offset) {
+  return atomic_compare_exchange_weak_explicit(
+      &(object[offset].val),
+      expected,
+      val,
+      memory_order_relaxed,
+      memory_order_relaxed);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/bf16.h b/macos/third_party/mlx/mlx/backend/metal/kernels/bf16.h
new file mode 100644
index 0000000..aa3c3c7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/bf16.h
@@ -0,0 +1,16 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+using namespace metal;
+
+typedef bfloat bfloat16_t;
+inline uint16_t bfloat16_to_uint16(const bfloat16_t x) {
+  return as_type<uint16_t>(x);
+}
+
+inline bfloat16_t uint16_to_bfloat16(const uint16_t x) {
+  return as_type<bfloat16_t>(x);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/bf16_math.h b/macos/third_party/mlx/mlx/backend/metal/kernels/bf16_math.h
new file mode 100644
index 0000000..0643fb3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/bf16_math.h
@@ -0,0 +1,380 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+///////////////////////////////////////////////////////////////////////////////
+// Metal math for bfloat16
+///////////////////////////////////////////////////////////////////////////////
+
+/*
+
+Following the Metal Shading Language Specification (Metal 3.1)
+
+"bfloat is an extended itypeing point type that only allows implicit conversion
+ to a type of greater itypeing point rank. While bfloat can be implicitly
+ converted to itype, it cannot be implicitly converted to half, and neither
+ itype nor half can be implicitly converted to bfloat."
+
+Further, as far as I can tell, the stdlib math/simd functions are not defined
+for bfloat and calling with an argument of type bfloat will result in that
+argument getting implicitly converted to itype which then returns an output
+that is (likely) a itype which cannot be implicitly converted into a bfloat
+
+This leads to situations where
+bfloat a = 5.0bf;
+bfloat b = metal::abs(a); // this will throw an error since abs return itype
+bfloat c = static_cast<bfloat>(metal::abs(a)); // this is fine
+
+For the moment, I will be adding overloaded instantiations of the math
+functions to accordingly automatically handle the casting
+
+*/
+
+#define instantiate_metal_math_funcs(itype, otype, ctype, mfast)               \
+                                                                               \
+  METAL_FUNC otype abs(itype x) {                                              \
+    return static_cast<otype>(__metal_fabs(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype acos(itype x) {                                             \
+    return static_cast<otype>(__metal_acos(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype acosh(itype x) {                                            \
+    return static_cast<otype>(__metal_acosh(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype asin(itype x) {                                             \
+    return static_cast<otype>(__metal_asin(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype asinh(itype x) {                                            \
+    return static_cast<otype>(__metal_asinh(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype atan(itype y_over_x) {                                      \
+    return static_cast<otype>(                                                 \
+        __metal_atan(static_cast<ctype>(y_over_x), mfast));                    \
+  }                                                                            \
+  METAL_FUNC otype atan2(itype y, itype x) {                                   \
+    return static_cast<otype>(                                                 \
+        __metal_atan2(static_cast<ctype>(y), static_cast<ctype>(x), mfast));   \
+  }                                                                            \
+  METAL_FUNC otype atanh(itype x) {                                            \
+    return static_cast<otype>(__metal_atanh(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype ceil(itype x) {                                             \
+    return static_cast<otype>(__metal_ceil(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype cos(itype x) {                                              \
+    return static_cast<otype>(__metal_cos(static_cast<ctype>(x), mfast));      \
+  }                                                                            \
+  METAL_FUNC otype cosh(itype x) {                                             \
+    return static_cast<otype>(__metal_cosh(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype cospi(itype x) {                                            \
+    return static_cast<otype>(__metal_cospi(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype divide(itype x, itype y) {                                  \
+    return static_cast<otype>(                                                 \
+        __metal_divide(static_cast<ctype>(x), static_cast<ctype>(y), mfast));  \
+  }                                                                            \
+  METAL_FUNC otype exp(itype x) {                                              \
+    return static_cast<otype>(__metal_exp(static_cast<ctype>(x), mfast));      \
+  }                                                                            \
+  METAL_FUNC otype exp10(itype x) {                                            \
+    return static_cast<otype>(__metal_exp10(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype exp2(itype x) {                                             \
+    return static_cast<otype>(__metal_exp2(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype fabs(itype x) {                                             \
+    return static_cast<otype>(__metal_fabs(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype fdim(itype x, itype y) {                                    \
+    ctype t = static_cast<ctype>(x - y);                                       \
+    return static_cast<otype>(select(t, ctype(0), t < ctype(0) || x == y));    \
+  }                                                                            \
+  METAL_FUNC otype floor(itype x) {                                            \
+    return static_cast<otype>(__metal_floor(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype fma(itype x, itype y, itype z) {                            \
+    return static_cast<otype>(__metal_fma(                                     \
+        static_cast<ctype>(x), static_cast<ctype>(y), static_cast<ctype>(z))); \
+  }                                                                            \
+  METAL_FUNC otype fmax(itype x, itype y) {                                    \
+    return static_cast<otype>(                                                 \
+        __metal_fmax(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype fmax3(itype x, itype y, itype z) {                          \
+    return static_cast<otype>(__metal_fmax3(                                   \
+        static_cast<ctype>(x),                                                 \
+        static_cast<ctype>(y),                                                 \
+        static_cast<ctype>(z),                                                 \
+        mfast));                                                               \
+  }                                                                            \
+  METAL_FUNC otype fmedian3(itype x, itype y, itype z) {                       \
+    return static_cast<otype>(__metal_fmedian3(                                \
+        static_cast<ctype>(x),                                                 \
+        static_cast<ctype>(y),                                                 \
+        static_cast<ctype>(z),                                                 \
+        mfast));                                                               \
+  }                                                                            \
+  METAL_FUNC otype fmin(itype x, itype y) {                                    \
+    return static_cast<otype>(                                                 \
+        __metal_fmin(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype fmin3(itype x, itype y, itype z) {                          \
+    return static_cast<otype>(__metal_fmin3(                                   \
+        static_cast<ctype>(x),                                                 \
+        static_cast<ctype>(y),                                                 \
+        static_cast<ctype>(z),                                                 \
+        mfast));                                                               \
+  }                                                                            \
+  METAL_FUNC otype fmod(itype x, itype y) {                                    \
+    return static_cast<otype>(                                                 \
+        __metal_fmod(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype fract(itype x) {                                            \
+    return static_cast<otype>(__metal_fract(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype frexp(itype x, thread int& exp) {                           \
+    return static_cast<otype>(__metal_frexp(static_cast<ctype>(x), &exp));     \
+  }                                                                            \
+  METAL_FUNC otype ldexp(itype x, int k) {                                     \
+    return static_cast<otype>(__metal_ldexp(static_cast<ctype>(x), k, mfast)); \
+  }                                                                            \
+  METAL_FUNC otype log(itype x) {                                              \
+    return static_cast<otype>(__metal_log(static_cast<ctype>(x), mfast));      \
+  }                                                                            \
+  METAL_FUNC otype log10(itype x) {                                            \
+    return static_cast<otype>(__metal_log10(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype log2(itype x) {                                             \
+    return static_cast<otype>(__metal_log2(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype max(itype x, itype y) {                                     \
+    return static_cast<otype>(                                                 \
+        __metal_fmax(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype max3(itype x, itype y, itype z) {                           \
+    return static_cast<otype>(__metal_fmax3(                                   \
+        static_cast<ctype>(x),                                                 \
+        static_cast<ctype>(y),                                                 \
+        static_cast<ctype>(z),                                                 \
+        mfast));                                                               \
+  }                                                                            \
+  METAL_FUNC otype median3(itype x, itype y, itype z) {                        \
+    return static_cast<otype>(__metal_fmedian3(                                \
+        static_cast<ctype>(x),                                                 \
+        static_cast<ctype>(y),                                                 \
+        static_cast<ctype>(z),                                                 \
+        mfast));                                                               \
+  }                                                                            \
+  METAL_FUNC otype min(itype x, itype y) {                                     \
+    return static_cast<otype>(                                                 \
+        __metal_fmin(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype min3(itype x, itype y, itype z) {                           \
+    return static_cast<otype>(__metal_fmin3(                                   \
+        static_cast<ctype>(x),                                                 \
+        static_cast<ctype>(y),                                                 \
+        static_cast<ctype>(z),                                                 \
+        mfast));                                                               \
+  }                                                                            \
+  METAL_FUNC otype nextafter(itype x, itype y) {                               \
+    return static_cast<otype>(                                                 \
+        __metal_nextafter(static_cast<ctype>(x), static_cast<ctype>(y)));      \
+  }                                                                            \
+  METAL_FUNC otype pow(itype x, itype y) {                                     \
+    return static_cast<otype>(                                                 \
+        __metal_pow(static_cast<ctype>(x), static_cast<ctype>(y), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype powr(itype x, itype y) {                                    \
+    return static_cast<otype>(                                                 \
+        __metal_powr(static_cast<ctype>(x), static_cast<ctype>(y), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype rint(itype x) {                                             \
+    return static_cast<otype>(__metal_rint(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype round(itype x) {                                            \
+    return static_cast<otype>(__metal_round(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype rsqrt(itype x) {                                            \
+    return static_cast<otype>(__metal_rsqrt(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype sin(itype x) {                                              \
+    return static_cast<otype>(__metal_sin(static_cast<ctype>(x), mfast));      \
+  }                                                                            \
+  METAL_FUNC otype sinh(itype x) {                                             \
+    return static_cast<otype>(__metal_sinh(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype sinpi(itype x) {                                            \
+    return static_cast<otype>(__metal_sinpi(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype sqrt(itype x) {                                             \
+    return static_cast<otype>(__metal_sqrt(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype tan(itype x) {                                              \
+    return static_cast<otype>(__metal_tan(static_cast<ctype>(x), mfast));      \
+  }                                                                            \
+  METAL_FUNC otype tanh(itype x) {                                             \
+    return static_cast<otype>(__metal_tanh(static_cast<ctype>(x), mfast));     \
+  }                                                                            \
+  METAL_FUNC otype tanpi(itype x) {                                            \
+    return static_cast<otype>(__metal_tanpi(static_cast<ctype>(x), mfast));    \
+  }                                                                            \
+  METAL_FUNC otype trunc(itype x) {                                            \
+    return static_cast<otype>(__metal_trunc(static_cast<ctype>(x), mfast));    \
+  }
+
+namespace metal {
+
+instantiate_metal_math_funcs(
+    bfloat16_t,
+    bfloat16_t,
+    float,
+    __METAL_MAYBE_FAST_MATH__);
+
+namespace fast {
+
+instantiate_metal_math_funcs(
+    bfloat16_t,
+    bfloat16_t,
+    float,
+    __METAL_FAST_MATH__);
+
+} // namespace fast
+
+namespace precise {
+
+instantiate_metal_math_funcs(
+    bfloat16_t,
+    bfloat16_t,
+    float,
+    __METAL_PRECISE_MATH__);
+
+} // namespace precise
+
+} // namespace metal
+
+///////////////////////////////////////////////////////////////////////////////
+// Metal simd for bfloat16
+///////////////////////////////////////////////////////////////////////////////
+
+#define instantiate_metal_simd_comm_funcs(                                   \
+    itype, otype, ctype, itype_to_ctype, ctype_to_otype)                     \
+                                                                             \
+  METAL_FUNC otype simd_broadcast(itype data, ushort broadcast_lane_id) {    \
+    return ctype_to_otype(                                                   \
+        __metal_simd_broadcast(itype_to_ctype(data), broadcast_lane_id));    \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle(itype data, ushort simd_lane_id) {           \
+    return ctype_to_otype(                                                   \
+        __metal_simd_shuffle(itype_to_ctype(data), simd_lane_id));           \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_and_fill_down(                               \
+      itype data, itype filling_data, ushort delta, ushort modulo) {         \
+    return ctype_to_otype(__metal_simd_shuffle_and_fill_down(                \
+        itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_and_fill_down(                               \
+      itype data, itype filling_data, ushort delta) {                        \
+    return ctype_to_otype(__metal_simd_shuffle_and_fill_down(                \
+        itype_to_ctype(data),                                                \
+        itype_to_ctype(filling_data),                                        \
+        delta,                                                               \
+        __metal_get_simdgroup_size(ushort())));                              \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_and_fill_up(                                 \
+      itype data, itype filling_data, ushort delta, ushort modulo) {         \
+    return ctype_to_otype(__metal_simd_shuffle_and_fill_up(                  \
+        itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_and_fill_up(                                 \
+      itype data, itype filling_data, ushort delta) {                        \
+    return ctype_to_otype(__metal_simd_shuffle_and_fill_up(                  \
+        itype_to_ctype(data),                                                \
+        itype_to_ctype(filling_data),                                        \
+        delta,                                                               \
+        __metal_get_simdgroup_size(ushort())));                              \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_down(itype data, ushort delta) {             \
+    return ctype_to_otype(                                                   \
+        __metal_simd_shuffle_down(itype_to_ctype(data), delta));             \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_rotate_down(itype data, ushort delta) {      \
+    return ctype_to_otype(                                                   \
+        __metal_simd_shuffle_rotate_down(itype_to_ctype(data), delta));      \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_rotate_up(itype data, ushort delta) {        \
+    return ctype_to_otype(                                                   \
+        __metal_simd_shuffle_rotate_up(itype_to_ctype(data), delta));        \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_up(itype data, ushort delta) {               \
+    return ctype_to_otype(                                                   \
+        __metal_simd_shuffle_up(itype_to_ctype(data), delta));               \
+  }                                                                          \
+                                                                             \
+  METAL_FUNC otype simd_shuffle_xor(itype data, ushort mask) {               \
+    return ctype_to_otype(                                                   \
+        __metal_simd_shuffle_xor(itype_to_ctype(data), mask));               \
+  }
+
+#define instantiate_metal_simd_reduction_funcs(itype, otype, ctype)            \
+                                                                               \
+  METAL_FUNC otype simd_max(itype data) {                                      \
+    return static_cast<otype>(__metal_simd_max(static_cast<ctype>(data)));     \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_min(itype data) {                                      \
+    return static_cast<otype>(__metal_simd_min(static_cast<ctype>(data)));     \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_prefix_exclusive_product(itype data) {                 \
+    return static_cast<otype>(                                                 \
+        __metal_simd_prefix_exclusive_product(static_cast<ctype>(data)));      \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_prefix_exclusive_sum(itype data) {                     \
+    return static_cast<otype>(                                                 \
+        __metal_simd_prefix_exclusive_sum(static_cast<ctype>(data)));          \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_prefix_inclusive_product(itype data) {                 \
+    return static_cast<otype>(                                                 \
+        __metal_simd_prefix_inclusive_product(static_cast<ctype>(data)));      \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_prefix_inclusive_sum(itype data) {                     \
+    return static_cast<otype>(                                                 \
+        __metal_simd_prefix_inclusive_sum(static_cast<ctype>(data)));          \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_product(itype data) {                                  \
+    return static_cast<otype>(__metal_simd_product(static_cast<ctype>(data))); \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_sum(itype data) {                                      \
+    return static_cast<otype>(__metal_simd_sum(static_cast<ctype>(data)));     \
+  }                                                                            \
+                                                                               \
+  METAL_FUNC otype simd_xor(itype data) {                                      \
+    return static_cast<otype>(__metal_simd_xor(static_cast<ctype>(data)));     \
+  }
+
+namespace metal {
+
+instantiate_metal_simd_comm_funcs(
+    bfloat16_t,
+    bfloat16_t,
+    uint16_t,
+    bfloat16_to_uint16,
+    uint16_to_bfloat16);
+instantiate_metal_simd_reduction_funcs(bfloat16_t, bfloat16_t, float);
+
+} // namespace metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/binary.h b/macos/third_party/mlx/mlx/backend/metal/kernels/binary.h
new file mode 100644
index 0000000..f1df885
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/binary.h
@@ -0,0 +1,199 @@
+// Copyright © 2024 Apple Inc.
+
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_ss(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    uint index [[thread_position_in_grid]]) {
+  c[index] = Op()(a[0], b[0]);
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_sv(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      c[index + i] = Op()(a[0], b[index + i]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      c[index + i] = Op()(a[0], b[index + i]);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vs(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      c[index + i] = Op()(a[index + i], b[0]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      c[index + i] = Op()(a[index + i], b[0]);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vv(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      c[index + i] = Op()(a[index + i], b[index + i]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      c[index + i] = Op()(a[index + i], b[index + i]);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_sv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      c[offset + i] = Op()(a[0], b[offset + i]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      c[offset + i] = Op()(a[0], b[offset + i]);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vs2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      c[offset + i] = Op()(a[offset + i], b[0]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      c[offset + i] = Op()(a[offset + i], b[0]);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      c[offset + i] = Op()(a[offset + i], b[offset + i]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      c[offset + i] = Op()(a[offset + i], b[offset + i]);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, typename IdxT = int64_t>
+[[kernel]] void binary_g_nd1(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant const int64_t& a_stride,
+    constant const int64_t& b_stride,
+    uint index [[thread_position_in_grid]]) {
+  auto a_idx = elem_to_loc_1<IdxT>(index, a_stride);
+  auto b_idx = elem_to_loc_1<IdxT>(index, b_stride);
+  c[index] = Op()(a[a_idx], b[b_idx]);
+}
+
+template <typename T, typename U, typename Op, typename IdxT = int64_t>
+[[kernel]] void binary_g_nd2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant const int64_t a_strides[2],
+    constant const int64_t b_strides[2],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_2<IdxT>(index, a_strides);
+  auto b_idx = elem_to_loc_2<IdxT>(index, b_strides);
+  IdxT out_idx = index.x + IdxT(grid_dim.x) * index.y;
+  c[out_idx] = Op()(a[a_idx], b[b_idx]);
+}
+
+template <typename T, typename U, typename Op, typename IdxT = int64_t>
+[[kernel]] void binary_g_nd3(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant const int64_t a_strides[3],
+    constant const int64_t b_strides[3],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_3<IdxT>(index, a_strides);
+  auto b_idx = elem_to_loc_3<IdxT>(index, b_strides);
+  IdxT out_idx = index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z);
+  c[out_idx] = Op()(a[a_idx], b[b_idx]);
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N = 1,
+    typename IdxT = int64_t>
+[[kernel]] void binary_g(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    constant const int& ndim,
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto idx = elem_to_loc_2_nd<IdxT>(
+      {N * index.x, index.y, index.z}, shape, a_strides, b_strides, ndim);
+  auto xshape = shape[ndim - 1];
+  IdxT out_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z);
+  IdxT a_xstride = a_strides[ndim - 1];
+  IdxT b_xstride = b_strides[ndim - 1];
+  for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) {
+    c[out_idx++] = Op()(a[idx.x], b[idx.y]);
+    idx.x += a_xstride;
+    idx.y += b_xstride;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/binary.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/binary.metal
new file mode 100644
index 0000000..17ed13c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/binary.metal
@@ -0,0 +1,109 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_integer>
+#include <metal_math>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/binary_ops.h"
+#include "mlx/backend/metal/kernels/binary.h"
+
+#define instantiate_binary_work_per_thread(op, tname, itype, otype)     \
+  instantiate_kernel("svn_" #op #tname, binary_sv, itype, otype, op)    \
+  instantiate_kernel("vsn_" #op #tname, binary_vs, itype, otype, op)    \
+  instantiate_kernel("vvn_" #op #tname, binary_vv, itype, otype, op)    \
+
+#define instantiate_binary_base(op, tname, itype, otype)                    \
+  instantiate_kernel("ss_" #op #tname, binary_ss, itype, otype, op)         \
+  instantiate_kernel("sv_" #op #tname, binary_sv, itype, otype, op, 1)      \
+  instantiate_kernel("vs_" #op #tname, binary_vs, itype, otype, op, 1)      \
+  instantiate_kernel("vv_" #op #tname, binary_vv, itype, otype, op, 1)      \
+  instantiate_kernel("sv2_" #op #tname, binary_sv2, itype, otype, op)       \
+  instantiate_kernel("vs2_" #op #tname, binary_vs2, itype, otype, op)       \
+  instantiate_kernel("vv2_" #op #tname, binary_vv2, itype, otype, op)       \
+  instantiate_kernel("gn2_" #op #tname, binary_g, itype, otype, op, 2, int) \
+  instantiate_kernel("gn4large_" #op #tname, binary_g, itype, otype, op, 4) \
+  instantiate_kernel("g1_" #op #tname, binary_g_nd1, itype, otype, op, int) \
+  instantiate_kernel("g1large_" #op #tname, binary_g_nd1, itype, otype, op) \
+  instantiate_kernel("g2_" #op #tname, binary_g_nd2, itype, otype, op, int) \
+  instantiate_kernel("g2large_" #op #tname, binary_g_nd2, itype, otype, op) \
+  instantiate_kernel("g3_" #op #tname, binary_g_nd3, itype, otype, op, int) \
+  instantiate_kernel("g3large_" #op #tname, binary_g_nd3, itype, otype, op)
+
+#define instantiate_binary_all(op, tname, itype, otype)       \
+  instantiate_binary_base(op, tname, itype, otype)            \
+  instantiate_binary_work_per_thread(op, tname, itype, otype)
+
+#define instantiate_binary_integer(op)                    \
+  instantiate_binary_all(op, uint8, uint8_t, uint8_t)     \
+  instantiate_binary_all(op, uint16, uint16_t, uint16_t)  \
+  instantiate_binary_all(op, uint32, uint32_t, uint32_t)  \
+  instantiate_binary_base(op, uint64, uint64_t, uint64_t) \
+  instantiate_binary_all(op, int8, int8_t, int8_t)        \
+  instantiate_binary_all(op, int16, int16_t, int16_t)     \
+  instantiate_binary_all(op, int32, int32_t, int32_t)     \
+  instantiate_binary_base(op, int64, int64_t, int64_t)
+
+#define instantiate_binary_float(op)                \
+  instantiate_binary_all(op, float16, half, half)   \
+  instantiate_binary_all(op, float32, float, float) \
+  instantiate_binary_all(op, bfloat16, bfloat16_t, bfloat16_t)
+
+#define instantiate_binary_types(op)                              \
+  instantiate_binary_all(op, bool_, bool, bool)                   \
+  instantiate_binary_integer(op)                                  \
+  instantiate_binary_base(op, complex64, complex64_t, complex64_t)\
+  instantiate_binary_float(op)
+
+#define instantiate_binary_types_bool(op)                \
+  instantiate_binary_all(op, bool_, bool, bool)          \
+  instantiate_binary_all(op, uint8, uint8_t, bool)       \
+  instantiate_binary_all(op, uint16, uint16_t, bool)     \
+  instantiate_binary_all(op, uint32, uint32_t, bool)     \
+  instantiate_binary_base(op, uint64, uint64_t, bool)    \
+  instantiate_binary_all(op, int8, int8_t, bool)         \
+  instantiate_binary_all(op, int16, int16_t, bool)       \
+  instantiate_binary_all(op, int32, int32_t, bool)       \
+  instantiate_binary_base(op, int64, int64_t, bool)      \
+  instantiate_binary_all(op, float16, half, bool)        \
+  instantiate_binary_all(op, float32, float, bool)       \
+  instantiate_binary_all(op, bfloat16, bfloat16_t, bool) \
+  instantiate_binary_base(op, complex64, complex64_t, bool)
+
+instantiate_binary_types(Add)
+instantiate_binary_types(Divide)
+instantiate_binary_types_bool(Equal)
+instantiate_binary_types_bool(Greater)
+instantiate_binary_types_bool(GreaterEqual)
+instantiate_binary_types_bool(Less)
+instantiate_binary_types_bool(LessEqual)
+instantiate_binary_types_bool(NotEqual)
+instantiate_binary_float(LogAddExp)
+instantiate_binary_base(LogAddExp, complex64, complex64_t, complex64_t)
+instantiate_binary_types(Maximum)
+instantiate_binary_types(Minimum)
+instantiate_binary_types(Multiply)
+instantiate_binary_types(Subtract)
+instantiate_binary_types(Power)
+instantiate_binary_types(Remainder)
+instantiate_binary_float(ArcTan2)
+
+// NaNEqual only needed for floating point types with boolean output
+instantiate_binary_all(NaNEqual, float16, half, bool)
+instantiate_binary_all(NaNEqual, float32, float, bool)
+instantiate_binary_all(NaNEqual, bfloat16, bfloat16_t, bool)
+instantiate_binary_base(NaNEqual, complex64, complex64_t, bool)
+
+instantiate_binary_all(LogicalOr, bool_, bool, bool)
+instantiate_binary_all(LogicalAnd, bool_, bool, bool)
+
+// Bitwise ops only need integer types and bool (except for l/r shift)
+instantiate_binary_integer(BitwiseAnd)
+instantiate_binary_all(BitwiseAnd, bool_, bool, bool)
+instantiate_binary_integer(BitwiseOr)
+instantiate_binary_all(BitwiseOr, bool_, bool, bool)
+instantiate_binary_integer(BitwiseXor)
+instantiate_binary_all(BitwiseXor, bool_, bool, bool)
+instantiate_binary_integer(LeftShift)
+instantiate_binary_integer(RightShift) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/binary_ops.h b/macos/third_party/mlx/mlx/backend/metal/kernels/binary_ops.h
new file mode 100644
index 0000000..4e3d881
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/binary_ops.h
@@ -0,0 +1,330 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <metal_integer>
+#include <metal_math>
+
+constant mlx::os_log logger("mlx", "binary_ops");
+
+struct Add {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x + y;
+  }
+};
+
+struct FloorDivide {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x / y;
+  }
+  template <>
+  float operator()(float x, float y) {
+    return trunc(x / y);
+  }
+  template <>
+  half operator()(half x, half y) {
+    return trunc(x / y);
+  }
+  template <>
+  bfloat16_t operator()(bfloat16_t x, bfloat16_t y) {
+    return trunc(x / y);
+  }
+};
+
+struct Divide {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x / y;
+  }
+};
+
+struct Remainder {
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T> & !metal::is_signed_v<T>, T>
+  operator()(T x, T y) {
+    return x % y;
+  }
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T> & metal::is_signed_v<T>, T>
+  operator()(T x, T y) {
+    auto r = x % y;
+    if (r != 0 && (r < 0 != y < 0)) {
+      r += y;
+    }
+    return r;
+  }
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) {
+    T r = fmod(x, y);
+    if (r != 0 && (r < 0 != y < 0)) {
+      r += y;
+    }
+    return r;
+  }
+  template <>
+  complex64_t operator()(complex64_t x, complex64_t y) {
+    return x % y;
+  }
+};
+
+struct Equal {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x == y;
+  }
+};
+
+struct NaNEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x == y || (metal::isnan(x) && metal::isnan(y));
+  }
+  template <>
+  bool operator()(complex64_t x, complex64_t y) {
+    return x == y ||
+        (metal::isnan(x.real) && metal::isnan(y.real) && metal::isnan(x.imag) &&
+         metal::isnan(y.imag)) ||
+        (x.real == y.real && metal::isnan(x.imag) && metal::isnan(y.imag)) ||
+        (metal::isnan(x.real) && metal::isnan(y.real) && x.imag == y.imag);
+  }
+};
+
+struct Greater {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x > y;
+  }
+};
+
+struct GreaterEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x >= y;
+  }
+};
+
+struct Less {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x < y;
+  }
+};
+
+struct LessEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x <= y;
+  }
+};
+
+struct LogAddExp {
+  template <typename T>
+  T operator()(T x, T y) {
+    if (metal::isnan(x) || metal::isnan(y)) {
+      return metal::numeric_limits<T>::quiet_NaN();
+    }
+    constexpr T inf = metal::numeric_limits<T>::infinity();
+    T maxval = metal::max(x, y);
+    T minval = metal::min(x, y);
+    return (minval == -inf || maxval == inf)
+        ? maxval
+        : (maxval + log1p(metal::exp(minval - maxval)));
+  };
+
+  complex64_t operator()(complex64_t x, complex64_t y) {
+    if (metal::isnan(x.real) || metal::isnan(x.imag) || metal::isnan(y.real) ||
+        metal::isnan(y.imag)) {
+      return metal::numeric_limits<float>::quiet_NaN();
+    }
+    constexpr float inf = metal::numeric_limits<float>::infinity();
+    complex64_t maxval = x > y ? x : y;
+    complex64_t minval = x < y ? x : y;
+    if (minval.real == -inf || maxval.real == inf)
+      return maxval;
+    float m = metal::exp(minval.real - maxval.real);
+    complex64_t dexp{
+        m * metal::cos(minval.imag - maxval.imag),
+        m * metal::sin(minval.imag - maxval.imag),
+    };
+    return maxval + log1p(dexp);
+  }
+};
+
+struct Maximum {
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T x, T y) {
+    return metal::max(x, y);
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) {
+    if (metal::isnan(x)) {
+      return x;
+    }
+    return x > y ? x : y;
+  }
+
+  template <>
+  complex64_t operator()(complex64_t x, complex64_t y) {
+    if (metal::isnan(x.real) || metal::isnan(x.imag)) {
+      return x;
+    }
+    return x > y ? x : y;
+  }
+};
+
+struct Minimum {
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T x, T y) {
+    return metal::min(x, y);
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) {
+    if (metal::isnan(x)) {
+      return x;
+    }
+    return x < y ? x : y;
+  }
+
+  template <>
+  complex64_t operator()(complex64_t x, complex64_t y) {
+    if (metal::isnan(x.real) || metal::isnan(x.imag)) {
+      return x;
+    }
+    return x < y ? x : y;
+  }
+};
+
+struct Multiply {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x * y;
+  }
+};
+
+struct NotEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x != y;
+  }
+  template <>
+  bool operator()(complex64_t x, complex64_t y) {
+    return x.real != y.real || x.imag != y.imag;
+  }
+};
+
+struct Power {
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T base, T exp) {
+    return metal::pow(base, exp);
+  }
+
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T base, T exp) {
+    T res = 1;
+    // Undefined to raise integer to negative power
+    if (exp < 0) {
+      logger.log_debug(
+          "int pow exp<0 (base=%ld exp=%ld)", (long)base, (long)exp);
+      return 0;
+    }
+
+    while (exp) {
+      if (exp & 1) {
+        res *= base;
+      }
+      exp >>= 1;
+      base *= base;
+    }
+    return res;
+  }
+
+  template <>
+  complex64_t operator()(complex64_t x, complex64_t y) {
+    if (x.real == 0 && x.imag == 0) {
+      if (metal::isnan(y.real) || metal::isnan(y.imag)) {
+        auto nan = metal::numeric_limits<float>::quiet_NaN();
+        return {nan, nan};
+      }
+      return {0.0, 0.0};
+    }
+    auto x_theta = metal::atan2(x.imag, x.real);
+    auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag);
+    auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta);
+    auto phase = y.imag * x_ln_r + y.real * x_theta;
+    return {mag * metal::cos(phase), mag * metal::sin(phase)};
+  }
+};
+
+struct Subtract {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x - y;
+  }
+};
+
+struct LogicalAnd {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x && y;
+  };
+};
+
+struct LogicalOr {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x || y;
+  };
+};
+
+struct BitwiseAnd {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x & y;
+  };
+};
+
+struct BitwiseOr {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x | y;
+  };
+};
+
+struct BitwiseXor {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x ^ y;
+  };
+};
+
+struct LeftShift {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x << y;
+  };
+};
+
+struct RightShift {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x >> y;
+  };
+};
+
+struct ArcTan2 {
+  template <typename T>
+  T operator()(T y, T x) {
+    return metal::precise::atan2(y, x);
+  }
+};
+
+struct DivMod {
+  template <typename T>
+  metal::array<T, 2> operator()(T x, T y) {
+    return {FloorDivide{}(x, y), Remainder{}(x, y)};
+  };
+};
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/binary_two.h b/macos/third_party/mlx/mlx/backend/metal/kernels/binary_two.h
new file mode 100644
index 0000000..4455e4c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/binary_two.h
@@ -0,0 +1,244 @@
+// Copyright © 2024 Apple Inc.
+
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_ss(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    uint index [[thread_position_in_grid]]) {
+  auto out = Op()(a[0], b[0]);
+  c[index] = out[0];
+  d[index] = out[1];
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_sv(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      auto out = Op()(a[0], b[index + i]);
+      c[index + i] = out[0];
+      d[index + i] = out[1];
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto out = Op()(a[0], b[index + i]);
+      c[index + i] = out[0];
+      d[index + i] = out[1];
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vs(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      auto out = Op()(a[index + i], b[0]);
+      c[index + i] = out[0];
+      d[index + i] = out[1];
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto out = Op()(a[index + i], b[0]);
+      c[index + i] = out[0];
+      d[index + i] = out[1];
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vv(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      auto out = Op()(a[index + i], b[index + i]);
+      c[index + i] = out[0];
+      d[index + i] = out[1];
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto out = Op()(a[index + i], b[index + i]);
+      c[index + i] = out[0];
+      d[index + i] = out[1];
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_sv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      auto out = Op()(a[0], b[offset + i]);
+      c[offset + i] = out[0];
+      d[offset + i] = out[1];
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto out = Op()(a[0], b[offset + i]);
+      c[offset + i] = out[0];
+      d[offset + i] = out[1];
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vs2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      auto out = Op()(a[offset + i], b[0]);
+      c[offset + i] = out[0];
+      d[offset + i] = out[1];
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto out = Op()(a[offset + i], b[0]);
+      c[offset + i] = out[0];
+      d[offset + i] = out[1];
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void binary_vv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      auto out = Op()(a[offset + i], b[offset + i]);
+      c[offset + i] = out[0];
+      d[offset + i] = out[1];
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto out = Op()(a[offset + i], b[offset + i]);
+      c[offset + i] = out[0];
+      d[offset + i] = out[1];
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, typename IdxT = int64_t>
+[[kernel]] void binary_g_nd1(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant const int64_t& a_stride,
+    constant const int64_t& b_stride,
+    uint index [[thread_position_in_grid]]) {
+  auto a_idx = elem_to_loc_1<IdxT>(index, a_stride);
+  auto b_idx = elem_to_loc_1<IdxT>(index, b_stride);
+  auto out = Op()(a[a_idx], b[b_idx]);
+  c[index] = out[0];
+  d[index] = out[1];
+}
+
+template <typename T, typename U, typename Op, typename IdxT = int64_t>
+[[kernel]] void binary_g_nd2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant const int64_t a_strides[2],
+    constant const int64_t b_strides[2],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_2<IdxT>(index, a_strides);
+  auto b_idx = elem_to_loc_2<IdxT>(index, b_strides);
+  IdxT out_idx = index.x + IdxT(grid_dim.x) * index.y;
+  auto out = Op()(a[a_idx], b[b_idx]);
+  c[out_idx] = out[0];
+  d[out_idx] = out[1];
+}
+
+template <typename T, typename U, typename Op, typename IdxT = int64_t>
+[[kernel]] void binary_g_nd3(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant const int64_t a_strides[3],
+    constant const int64_t b_strides[3],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_3<IdxT>(index, a_strides);
+  auto b_idx = elem_to_loc_3<IdxT>(index, b_strides);
+  IdxT out_idx = index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z);
+  auto out = Op()(a[a_idx], b[b_idx]);
+  c[out_idx] = out[0];
+  d[out_idx] = out[1];
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N = 1,
+    typename IdxT = int64_t>
+[[kernel]] void binary_g(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    constant const int& ndim,
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto idx = elem_to_loc_2_nd<IdxT>(
+      {N * index.x, index.y, index.z}, shape, a_strides, b_strides, ndim);
+  auto xshape = shape[ndim - 1];
+  IdxT out_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z);
+  IdxT a_xstride = a_strides[ndim - 1];
+  IdxT b_xstride = b_strides[ndim - 1];
+  for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) {
+    auto out = Op()(a[idx.x], b[idx.y]);
+    c[out_idx] = out[0];
+    d[out_idx++] = out[1];
+    idx.x += a_xstride;
+    idx.y += b_xstride;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/binary_two.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/binary_two.metal
new file mode 100644
index 0000000..c7d3ecd
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/binary_two.metal
@@ -0,0 +1,54 @@
+// Copyright © 2024 Apple Inc.
+#include <metal_integer>
+#include <metal_math>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/binary_ops.h"
+#include "mlx/backend/metal/kernels/binary_two.h"
+
+#define instantiate_binary_work_per_thread(op, tname, itype, otype)     \
+  instantiate_kernel("svn_" #op #tname, binary_sv, itype, otype, op)    \
+  instantiate_kernel("vsn_" #op #tname, binary_vs, itype, otype, op)    \
+  instantiate_kernel("vvn_" #op #tname, binary_vv, itype, otype, op)
+
+#define instantiate_binary_base(op, tname, itype, otype)                    \
+  instantiate_kernel("ss_" #op #tname, binary_ss, itype, otype, op)         \
+  instantiate_kernel("sv_" #op #tname, binary_sv, itype, otype, op, 1)      \
+  instantiate_kernel("vs_" #op #tname, binary_vs, itype, otype, op, 1)      \
+  instantiate_kernel("vv_" #op #tname, binary_vv, itype, otype, op, 1)      \
+  instantiate_kernel("sv2_" #op #tname, binary_sv2, itype, otype, op)       \
+  instantiate_kernel("vs2_" #op #tname, binary_vs2, itype, otype, op)       \
+  instantiate_kernel("vv2_" #op #tname, binary_vv2, itype, otype, op)       \
+  instantiate_kernel("gn2_" #op #tname, binary_g, itype, otype, op, 2, int) \
+  instantiate_kernel("gn4large_" #op #tname, binary_g, itype, otype, op, 4) \
+  instantiate_kernel("g1_" #op #tname, binary_g_nd1, itype, otype, op, int) \
+  instantiate_kernel("g2_" #op #tname, binary_g_nd2, itype, otype, op, int) \
+  instantiate_kernel("g3_" #op #tname, binary_g_nd3, itype, otype, op, int) \
+  instantiate_kernel("g1large_" #op #tname, binary_g_nd1, itype, otype, op) \
+  instantiate_kernel("g2large_" #op #tname, binary_g_nd2, itype, otype, op) \
+  instantiate_kernel("g3large_" #op #tname, binary_g_nd3, itype, otype, op)
+
+#define instantiate_binary_all(op, tname, itype, otype)       \
+  instantiate_binary_base(op, tname, itype, otype)            \
+  instantiate_binary_work_per_thread(op, tname, itype, otype)
+
+#define instantiate_binary_float(op)                \
+  instantiate_binary_all(op, float16, half, half)   \
+  instantiate_binary_all(op, float32, float, float) \
+  instantiate_binary_all(op, bfloat16, bfloat16_t, bfloat16_t)
+
+#define instantiate_binary_types(op)                               \
+  instantiate_binary_all(op, bool_, bool, bool)                    \
+  instantiate_binary_all(op, uint8, uint8_t, uint8_t)              \
+  instantiate_binary_all(op, uint16, uint16_t, uint16_t)           \
+  instantiate_binary_all(op, uint32, uint32_t, uint32_t)           \
+  instantiate_binary_base(op, uint64, uint64_t, uint64_t)          \
+  instantiate_binary_all(op, int8, int8_t, int8_t)                 \
+  instantiate_binary_all(op, int16, int16_t, int16_t)              \
+  instantiate_binary_all(op, int32, int32_t, int32_t)              \
+  instantiate_binary_base(op, int64, int64_t, int64_t)             \
+  instantiate_binary_base(op, complex64, complex64_t, complex64_t) \
+  instantiate_binary_float(op)
+
+instantiate_binary_types(DivMod) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/cexpf.h b/macos/third_party/mlx/mlx/backend/metal/kernels/cexpf.h
new file mode 100644
index 0000000..b45fe6a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/cexpf.h
@@ -0,0 +1,134 @@
+// Copyright © 2025 Apple Inc.
+// Copyright © 2008-2013 NVIDIA Corporation
+// Copyright © 2013 Filipe RNC Maia
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Forked from
+// https://github.com/NVIDIA/cccl/blob/main/thrust/thrust/detail/complex/cexpf.h
+
+// TODO: We should use thrust::exp but the thrust header in old CUDA versions
+// can not be used in JIT.
+
+#pragma once
+
+#include <metal_math>
+
+using ieee_float_shape_type = union {
+  float value;
+  uint32_t word;
+};
+
+inline void get_float_word(thread uint32_t& i, float d) {
+  ieee_float_shape_type gf_u;
+  gf_u.value = (d);
+  (i) = gf_u.word;
+}
+
+inline void get_float_word(thread int32_t& i, float d) {
+  ieee_float_shape_type gf_u;
+  gf_u.value = (d);
+  (i) = gf_u.word;
+}
+
+inline void set_float_word(thread float& d, uint32_t i) {
+  ieee_float_shape_type sf_u;
+  sf_u.word = (i);
+  (d) = sf_u.value;
+}
+
+inline float frexp_expf(float x, thread int* expt) {
+  const uint32_t k = 235;
+  const float kln2 = 162.88958740F;
+
+  float exp_x;
+  uint32_t hx;
+
+  exp_x = metal::exp(x - kln2);
+  get_float_word(hx, exp_x);
+  *expt = (hx >> 23) - (0x7f + 127) + k;
+  set_float_word(exp_x, (hx & 0x7fffff) | ((0x7f + 127) << 23));
+  return exp_x;
+}
+
+inline complex64_t ldexp_cexpf(complex64_t z, int expt) {
+  float x, y, exp_x, scale1, scale2;
+  int ex_expt, half_expt;
+
+  x = z.real;
+  y = z.imag;
+  exp_x = frexp_expf(x, &ex_expt);
+  expt += ex_expt;
+
+  half_expt = expt / 2;
+  set_float_word(scale1, (0x7f + half_expt) << 23);
+  half_expt = expt - half_expt;
+  set_float_word(scale2, (0x7f + half_expt) << 23);
+
+  return complex64_t{
+      metal::cos(y) * exp_x * scale1 * scale2,
+      metal::sin(y) * exp_x * scale1 * scale2};
+}
+
+inline complex64_t cexpf(const thread complex64_t& z) {
+  float x, y, exp_x;
+  uint32_t hx, hy;
+
+  const uint32_t exp_ovfl = 0x42b17218, cexp_ovfl = 0x43400074;
+
+  x = z.real;
+  y = z.imag;
+
+  get_float_word(hy, y);
+  hy &= 0x7fffffff;
+
+  /* cexp(x + I 0) = exp(x) + I 0 */
+  if (hy == 0) {
+    return complex64_t{metal::exp(x), y};
+  }
+  get_float_word(hx, x);
+  /* cexp(0 + I y) = cos(y) + I sin(y) */
+  if ((hx & 0x7fffffff) == 0) {
+    return complex64_t{metal::cos(y), metal::sin(y)};
+  }
+  if (hy >= 0x7f800000) {
+    if ((hx & 0x7fffffff) != 0x7f800000) {
+      /* cexp(finite|NaN +- I Inf|NaN) = NaN + I NaN */
+      return complex64_t{y - y, y - y};
+    } else if (hx & 0x80000000) {
+      /* cexp(-Inf +- I Inf|NaN) = 0 + I 0 */
+      return complex64_t{0.0, 0.0};
+    } else {
+      /* cexp(+Inf +- I Inf|NaN) = Inf + I NaN */
+      return complex64_t{x, y - y};
+    }
+  }
+
+  if (hx >= exp_ovfl && hx <= cexp_ovfl) {
+    /*
+     * x is between 88.7 and 192, so we must scale to avoid
+     * overflow in expf(x).
+     */
+    return ldexp_cexpf(z, 0);
+  } else {
+    /*
+     * Cases covered here:
+     *  -  x < exp_ovfl and exp(x) won't overflow (common case)
+     *  -  x > cexp_ovfl, so exp(x) * s overflows for all s > 0
+     *  -  x = +-Inf (generated by exp())
+     *  -  x = NaN (spurious inexact exception from y)
+     */
+    exp_x = metal::exp(x);
+    return complex64_t{exp_x * metal::cos(y), exp_x * metal::sin(y)};
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/complex.h b/macos/third_party/mlx/mlx/backend/metal/kernels/complex.h
new file mode 100644
index 0000000..6e39148
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/complex.h
@@ -0,0 +1,173 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+using namespace metal;
+
+struct complex64_t;
+
+template <typename T>
+static constexpr constant bool can_convert_to_complex64 =
+    !is_same_v<T, complex64_t> && is_convertible_v<T, float>;
+
+template <typename T>
+static constexpr constant bool can_convert_from_complex64 =
+    !is_same_v<T, complex64_t> &&
+    (is_convertible_v<float, T> || is_convertible_v<bfloat16_t, T>);
+
+struct complex64_t {
+  float real;
+  float imag;
+
+  // Constructors
+  constexpr complex64_t(float real, float imag) : real(real), imag(imag) {};
+  constexpr complex64_t() : real(0), imag(0) {};
+  constexpr complex64_t() threadgroup : real(0), imag(0) {};
+
+  // Conversions to complex64_t
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_complex64<T>>::type>
+  constexpr complex64_t(T x) thread : real(x), imag(0) {}
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_complex64<T>>::type>
+  constexpr complex64_t(T x) threadgroup : real(x), imag(0) {}
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_complex64<T>>::type>
+  constexpr complex64_t(T x) device : real(x), imag(0) {}
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_to_complex64<T>>::type>
+  constexpr complex64_t(T x) constant : real(x), imag(0) {}
+
+  // Conversions from complex64_t
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_complex64<T>>::type>
+  constexpr operator T() const thread {
+    return static_cast<T>(real);
+  }
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_complex64<T>>::type>
+  constexpr operator T() const threadgroup {
+    return static_cast<T>(real);
+  }
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_complex64<T>>::type>
+  constexpr operator T() const device {
+    return static_cast<T>(real);
+  }
+
+  template <
+      typename T,
+      typename = typename enable_if<can_convert_from_complex64<T>>::type>
+  constexpr operator T() const constant {
+    return static_cast<T>(real);
+  }
+};
+
+constexpr complex64_t operator-(complex64_t x) {
+  return {-x.real, -x.imag};
+}
+
+constexpr bool operator>=(complex64_t a, complex64_t b) {
+  return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag);
+}
+
+constexpr bool operator>(complex64_t a, complex64_t b) {
+  return (a.real > b.real) || (a.real == b.real && a.imag > b.imag);
+}
+
+constexpr bool operator<=(complex64_t a, complex64_t b) {
+  return operator>=(b, a);
+}
+
+constexpr bool operator<(complex64_t a, complex64_t b) {
+  return operator>(b, a);
+}
+
+constexpr bool operator==(complex64_t a, complex64_t b) {
+  return a.real == b.real && a.imag == b.imag;
+}
+
+constexpr complex64_t operator+(complex64_t a, complex64_t b) {
+  return {a.real + b.real, a.imag + b.imag};
+}
+
+constexpr thread complex64_t& operator+=(thread complex64_t& a, complex64_t b) {
+  a.real += b.real;
+  a.imag += b.imag;
+  return a;
+}
+
+constexpr threadgroup complex64_t& operator+=(
+    threadgroup complex64_t& a,
+    complex64_t b) {
+  a.real += b.real;
+  a.imag += b.imag;
+  return a;
+}
+
+constexpr device complex64_t& operator+=(device complex64_t& a, complex64_t b) {
+  a.real += b.real;
+  a.imag += b.imag;
+  return a;
+}
+
+constexpr complex64_t operator+(float a, complex64_t b) {
+  return {a + b.real, b.imag};
+}
+constexpr complex64_t operator+(complex64_t a, float b) {
+  return {a.real + b, a.imag};
+}
+
+constexpr complex64_t operator-(complex64_t a, complex64_t b) {
+  return {a.real - b.real, a.imag - b.imag};
+}
+constexpr complex64_t operator-(float a, complex64_t b) {
+  return {a - b.real, -b.imag};
+}
+constexpr complex64_t operator-(complex64_t a, float b) {
+  return {a.real - b, a.imag};
+}
+
+constexpr complex64_t operator*(complex64_t a, complex64_t b) {
+  return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real};
+}
+
+constexpr complex64_t operator/(complex64_t a, complex64_t b) {
+  auto denom = b.real * b.real + b.imag * b.imag;
+  auto x = a.real * b.real + a.imag * b.imag;
+  auto y = a.imag * b.real - a.real * b.imag;
+  return {x / denom, y / denom};
+}
+
+constexpr complex64_t operator/(float a, complex64_t b) {
+  auto denom = b.real * b.real + b.imag * b.imag;
+  auto x = a * b.real;
+  auto y = -a * b.imag;
+  return {x / denom, y / denom};
+}
+
+constexpr complex64_t operator%(complex64_t a, complex64_t b) {
+  auto real = a.real - (b.real * static_cast<int64_t>(a.real / b.real));
+  auto imag = a.imag - (b.imag * static_cast<int64_t>(a.imag / b.imag));
+  if (real != 0 && (real < 0 != b.real < 0)) {
+    real += b.real;
+  }
+  if (imag != 0 && (imag < 0 != b.imag < 0)) {
+    imag += b.imag;
+  }
+  return {real, imag};
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/conv.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/conv.metal
new file mode 100644
index 0000000..50a14ad
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/conv.metal
@@ -0,0 +1,701 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+#include "mlx/backend/metal/kernels/utils.h"
+
+#define MLX_MTL_CONST static constant constexpr const
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+/// Naive unfold with dilation
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, int N>
+[[kernel]] void naive_unfold_Nd(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const constant MLXConvParams<N>* params [[buffer(2)]],
+    uint3 gid [[thread_position_in_grid]]) {
+  int filter_size = params->C;
+  for (short i = 0; i < N; i++)
+    filter_size *= params->wS[i];
+
+  int out_pixels = 1;
+  for (short i = 0; i < N; i++)
+    out_pixels *= params->oS[i];
+
+  // Set out
+  out += gid.z * filter_size + gid.y * (params->C);
+
+  // Coordinates in input
+  int is[N] = {0};
+
+  // gid.z: N oS (Batch and row in unfolded output)
+  // gid.y: wS (Filter location to unfold input)
+  // gid.x: C (channel)
+
+  int n = (gid.z) / out_pixels;
+  int oS = (gid.z) % out_pixels;
+  int wS = gid.y;
+
+  bool valid = n < params->N;
+
+  // Unroll dimensions
+  for (int i = N - 1; i >= 0; --i) {
+    int os_ = (oS % params->oS[i]);
+    int ws_ = (wS % params->wS[i]);
+
+    ws_ = params->flip ? params->wS[i] - ws_ - 1 : ws_;
+
+    int is_ = os_ * params->str[i] - params->pad[i] + ws_ * params->kdil[i];
+    int is_max = 1 + params->idil[i] * (params->iS[i] - 1);
+
+    valid &= is_ >= 0 && is_ < is_max && (is_ % params->idil[i] == 0);
+
+    is[i] = is_ / params->idil[i];
+
+    oS /= params->oS[i];
+    wS /= params->wS[i];
+  }
+
+  if (valid) {
+    size_t in_offset = n * params->in_strides[0];
+
+    for (int i = 0; i < N; ++i) {
+      in_offset += is[i] * params->in_strides[i + 1];
+    }
+
+    out[gid.x] = in[in_offset + gid.x];
+  } else {
+    out[gid.x] = T(0);
+  }
+}
+
+// This kernel unfolds the input array of size (N, *spatial_dims, C)
+// into an array of size (N x *spatial_dims, C x *kernel_dims).
+template <typename T, int N>
+[[kernel]] void naive_unfold_transpose_Nd(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const constant MLXConvParams<N>* params [[buffer(2)]],
+    uint3 gid [[thread_position_in_grid]]) {
+  int filter_size = params->C;
+  for (short i = 0; i < N; i++)
+    filter_size *= params->wS[i];
+
+  int out_pixels = 1;
+  for (short i = 0; i < N; i++)
+    out_pixels *= params->oS[i];
+
+  // Set out
+  out += gid.z * filter_size + gid.x * (filter_size / params->C);
+
+  // Coordinates in input
+  int is[N] = {0};
+
+  // gid.z: N oS (Batch and row in unfolded output)
+  // gid.y: wS (Filter location to unfold input)
+  // gid.x: C (channel)
+
+  int n = (gid.z) / out_pixels;
+  int oS = (gid.z) % out_pixels;
+  int wS = gid.y;
+
+  bool valid = n < params->N;
+
+  // Unroll dimensions
+  int kernel_stride = 1;
+  for (int i = N - 1; i >= 0; --i) {
+    int os_ = (oS % params->oS[i]);
+    int ws_ = (wS % params->wS[i]);
+    out += ws_ * kernel_stride;
+
+    ws_ = params->flip ? params->wS[i] - ws_ - 1 : ws_;
+
+    int is_ = os_ * params->str[i] - params->pad[i] + ws_ * params->kdil[i];
+    int is_max = 1 + params->idil[i] * (params->iS[i] - 1);
+
+    valid &= is_ >= 0 && is_ < is_max && (is_ % params->idil[i] == 0);
+
+    is[i] = is_ / params->idil[i];
+
+    oS /= params->oS[i];
+    wS /= params->wS[i];
+
+    kernel_stride *= params->wS[i];
+  }
+
+  if (valid) {
+    size_t in_offset = n * params->in_strides[0];
+
+    for (int i = 0; i < N; ++i) {
+      in_offset += is[i] * params->in_strides[i + 1];
+    }
+
+    out[0] = in[in_offset + gid.x];
+  } else {
+    out[0] = T(0);
+  }
+}
+
+#define instantiate_naive_unfold_nd(name, itype, n)                            \
+  template [[host_name("naive_unfold_nd_" #name "_" #n)]] [[kernel]] void      \
+  naive_unfold_Nd(                                                             \
+      const device itype* in [[buffer(0)]],                                    \
+      device itype* out [[buffer(1)]],                                         \
+      const constant MLXConvParams<n>* params [[buffer(2)]],                   \
+      uint3 gid [[thread_position_in_grid]]);                                  \
+  template                                                                     \
+      [[host_name("naive_unfold_transpose_nd_" #name "_" #n)]] [[kernel]] void \
+      naive_unfold_transpose_Nd(                                               \
+          const device itype* in [[buffer(0)]],                                \
+          device itype* out [[buffer(1)]],                                     \
+          const constant MLXConvParams<n>* params [[buffer(2)]],               \
+          uint3 gid [[thread_position_in_grid]]);
+
+#define instantiate_naive_unfold_nd_dims(name, itype)                      \
+  instantiate_naive_unfold_nd(name, itype, 1) instantiate_naive_unfold_nd( \
+      name, itype, 2) instantiate_naive_unfold_nd(name, itype, 3)
+
+instantiate_naive_unfold_nd_dims(float32, float);
+instantiate_naive_unfold_nd_dims(float16, half);
+instantiate_naive_unfold_nd_dims(bfloat16, bfloat16_t);
+
+///////////////////////////////////////////////////////////////////////////////
+/// Depthwise convolution kernels
+///////////////////////////////////////////////////////////////////////////////
+
+constant int ker_h [[function_constant(00)]];
+constant int ker_w [[function_constant(01)]];
+constant int str_h [[function_constant(10)]];
+constant int str_w [[function_constant(11)]];
+constant int tgp_h [[function_constant(100)]];
+constant int tgp_w [[function_constant(101)]];
+constant bool do_flip [[function_constant(200)]];
+
+constant int span_h = tgp_h * str_h + ker_h - 1;
+constant int span_w = tgp_w * str_w + ker_w - 1;
+constant int span_hw = span_h * span_w;
+
+template <typename T>
+[[kernel]] void depthwise_conv_2d(
+    const device T* in [[buffer(0)]],
+    const device T* wt [[buffer(1)]],
+    device T* out [[buffer(2)]],
+    const constant MLXConvParams<2>& params [[buffer(3)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int tc = 8;
+  constexpr int tw = 8;
+  constexpr int th = 4;
+
+  constexpr int c_per_thr = 8;
+
+  constexpr int TGH = th * 2 + 6;
+  constexpr int TGW = tw * 2 + 6;
+  constexpr int TGC = tc;
+
+  threadgroup T ins[TGH * TGW * TGC];
+
+  const int n_tgblocks_h = params.oS[0] / th;
+  const int n = tid.z / n_tgblocks_h;
+  const int tghid = tid.z % n_tgblocks_h;
+  const int oh = tghid * th + lid.z;
+  const int ow = gid.y;
+  const int c = gid.x;
+
+  in += n * params.in_strides[0];
+
+  // Load in
+  {
+    constexpr int n_threads = th * tw * tc;
+    const int tg_oh = (tghid * th) * str_h - params.pad[0];
+    const int tg_ow = (tid.y * tw) * str_w - params.pad[1];
+    const int tg_c = tid.x * tc;
+
+    const int thread_idx = simd_gid * 32 + simd_lid;
+    constexpr int thr_per_hw = tc / c_per_thr;
+    constexpr int hw_per_group = n_threads / thr_per_hw;
+
+    const int thr_c = thread_idx % thr_per_hw;
+    const int thr_hw = thread_idx / thr_per_hw;
+
+    for (int hw = thr_hw; hw < span_hw; hw += hw_per_group) {
+      const int h = hw / span_w;
+      const int w = hw % span_w;
+
+      const int ih = tg_oh + h;
+      const int iw = tg_ow + w;
+
+      const int in_s_offset = h * span_w * TGC + w * TGC;
+
+      if (ih >= 0 && ih < params.iS[0] && iw >= 0 && iw < params.iS[1]) {
+        const auto in_load =
+            in + ih * params.in_strides[1] + iw * params.in_strides[2] + tg_c;
+
+        MLX_MTL_PRAGMA_UNROLL
+        for (int cc = 0; cc < c_per_thr; ++cc) {
+          ins[in_s_offset + c_per_thr * thr_c + cc] =
+              in_load[c_per_thr * thr_c + cc];
+        }
+      } else {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int cc = 0; cc < c_per_thr; ++cc) {
+          ins[in_s_offset + c_per_thr * thr_c + cc] = T(0);
+        }
+      }
+    }
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  wt += c * params.wt_strides[0];
+
+  const auto ins_ptr =
+      &ins[lid.z * str_h * span_w * TGC + lid.y * str_w * TGC + lid.x];
+  float o = 0.;
+  for (int h = 0; h < ker_h; ++h) {
+    for (int w = 0; w < ker_w; ++w) {
+      int wt_h = h;
+      int wt_w = w;
+      if (do_flip) {
+        wt_h = ker_h - h - 1;
+        wt_w = ker_w - w - 1;
+      }
+      auto inv = ins_ptr[h * span_w * TGC + w * TGC];
+      auto wtv = wt[wt_h * ker_w + wt_w];
+      o += inv * wtv;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_none);
+
+  out += n * params.out_strides[0] + oh * params.out_strides[1] +
+      ow * params.out_strides[2];
+  out[c] = static_cast<T>(o);
+}
+
+#define instantiate_depthconv2d(iname, itype) \
+  instantiate_kernel("depthwise_conv_2d_" #iname, depthwise_conv_2d, itype)
+
+instantiate_depthconv2d(float32, float);
+instantiate_depthconv2d(float16, half);
+instantiate_depthconv2d(bfloat16, bfloat16_t);
+
+template <typename T, typename IdxT>
+[[kernel]] void depthwise_conv_1d(
+    const device T* in [[buffer(0)]],
+    const device T* w [[buffer(1)]],
+    device T* out [[buffer(2)]],
+    constant const IdxT strides[3],
+    constant const int& kernel_size,
+    uint3 tid [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  out += (tid.z * static_cast<IdxT>(grid_dim.y) + tid.y) * grid_dim.x + tid.x;
+  in += tid.z * strides[0] + tid.y * strides[1] + tid.x * strides[2];
+  w += tid.x * kernel_size;
+
+  float acc = 0.0;
+  for (int i = 0; i < kernel_size; ++i) {
+    acc += static_cast<float>(in[0]) * w[i];
+    in += strides[1];
+  }
+  *out = static_cast<T>(acc);
+}
+
+#define instantiate_depthconv1d(iname, itype)                         \
+  instantiate_kernel(                                                 \
+      "depthwise_conv_1d_" #iname, depthwise_conv_1d, itype, int32_t) \
+      instantiate_kernel(                                             \
+          "depthwise_conv_1d_" #iname "_large",                       \
+          depthwise_conv_1d,                                          \
+          itype,                                                      \
+          int64_t)
+
+instantiate_depthconv1d(float32, float);
+instantiate_depthconv1d(float16, half);
+instantiate_depthconv1d(bfloat16, bfloat16_t);
+
+///////////////////////////////////////////////////////////////////////////////
+/// Winograd kernels
+///////////////////////////////////////////////////////////////////////////////
+
+template <int M, int R, int S>
+struct WinogradTransforms {};
+
+template <>
+struct WinogradTransforms<6, 3, 8> {
+  MLX_MTL_CONST int OUT_TILE_SIZE = 6;
+  MLX_MTL_CONST int FILTER_SIZE = 3;
+  MLX_MTL_CONST int IN_TILE_SIZE = OUT_TILE_SIZE + FILTER_SIZE - 1;
+  MLX_MTL_CONST int SIMD_MATRIX_SIZE = 8;
+  MLX_MTL_CONST float in_transform[SIMD_MATRIX_SIZE][SIMD_MATRIX_SIZE] = {
+      {1.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f},
+      {0.00f, 1.00f, -1.00f, 0.50f, -0.50f, 2.00f, -2.00f, -1.00f},
+      {-5.25f, 1.00f, 1.00f, 0.25f, 0.25f, 4.00f, 4.00f, 0.00f},
+      {0.00f, -4.25f, 4.25f, -2.50f, 2.50f, -2.50f, 2.50f, 5.25f},
+      {5.25f, -4.25f, -4.25f, -1.25f, -1.25f, -5.00f, -5.00f, 0.00f},
+      {0.00f, 1.00f, -1.00f, 2.00f, -2.00f, 0.50f, -0.50f, -5.25f},
+      {-1.00f, 1.00f, 1.00f, 1.00f, 1.00f, 1.00f, 1.00f, 0.00f},
+      {0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 1.00f},
+  };
+
+  MLX_MTL_CONST float out_transform[SIMD_MATRIX_SIZE][SIMD_MATRIX_SIZE] = {
+      {1.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f},
+      {1.00f, 1.00f, 1.00f, 1.00f, 1.00f, 1.00f},
+      {1.00f, -1.00f, 1.00f, -1.00f, 1.00f, -1.00f},
+      {1.00f, 2.00f, 4.00f, 8.00f, 16.00f, 32.00f},
+      {1.00f, -2.00f, 4.00f, -8.00f, 16.00f, -32.00f},
+      {1.00f, 0.50f, 0.25f, 0.125f, 0.0625f, 0.03125f},
+      {1.00f, -0.50f, 0.25f, -0.125f, 0.0625f, -0.03125f},
+      {0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 1.00f},
+  };
+
+  MLX_MTL_CONST float wt_transform[SIMD_MATRIX_SIZE][SIMD_MATRIX_SIZE] = {
+      {1.00, 0.00, 0.00},
+      {-2.0 / 9.00, -2.0 / 9.00, -2.0 / 9.00},
+      {-2.0 / 9.00, 2.0 / 9.00, -2.0 / 9.00},
+      {1.0 / 90.0, 1.0 / 45.0, 2.0 / 45.0},
+      {1.0 / 90.0, -1.0 / 45.0, 2.0 / 45.0},
+      {32.0 / 45.0, 16.0 / 45.0, 8.0 / 45.0},
+      {32.0 / 45.0, -16.0 / 45.0, 8.0 / 45.0},
+      {0.00, 0.00, 1.00},
+  };
+};
+
+constant constexpr const float WinogradTransforms<6, 3, 8>::wt_transform[8][8];
+constant constexpr const float WinogradTransforms<6, 3, 8>::in_transform[8][8];
+constant constexpr const float WinogradTransforms<6, 3, 8>::out_transform[8][8];
+
+template <typename T, int BC = 32, int BO = 4, int M = 6, int R = 3>
+[[kernel, max_total_threads_per_threadgroup(BO * 32)]] void
+winograd_conv_2d_weight_transform(
+    const device T* wt_in [[buffer(0)]],
+    device T* wt_out [[buffer(1)]],
+    const constant int& C [[buffer(2)]],
+    const constant int& O [[buffer(3)]],
+    uint tid [[threadgroup_position_in_grid]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]]) {
+  using WGT = WinogradTransforms<M, R, 8>;
+
+  // Get lane position in simdgroup
+  const short qid = simd_lane_id / 4;
+  const short sm = (qid & 4) + (simd_lane_id / 2) % 4;
+  const short sn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
+
+  // Initialize G matrix
+  simdgroup_matrix<float, 8, 8> G;
+  G.thread_elements()[0] = WGT::wt_transform[sm][sn];
+  G.thread_elements()[1] = WGT::wt_transform[sm][sn + 1];
+
+  // Initialize Gt matrix
+  simdgroup_matrix<float, 8, 8> Gt;
+  Gt.thread_elements()[0] = WGT::wt_transform[sn][sm];
+  Gt.thread_elements()[1] = WGT::wt_transform[sn + 1][sm];
+
+  // Move to the correct output filter
+  size_t ko = BO * tid + simd_group_id;
+  wt_in += ko * R * R * C;
+
+  // wt_out is stored transposed (A x A x C x O)
+  short ohw_0 = sm * 8 + sn;
+  short ohw_1 = sm * 8 + sn + 1;
+  device T* wt_out_0 = wt_out + ohw_0 * C * O + ko;
+  device T* wt_out_1 = wt_out + ohw_1 * C * O + ko;
+
+  // Prepare shared memory
+  threadgroup T Ws[BO][R][R][BC];
+
+  // Loop over C
+  for (int bc = 0; bc < C; bc += BC) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Read into shared memory
+    for (int kh = 0; kh < R; ++kh) {
+      for (int kw = 0; kw < R; ++kw) {
+        for (int kc = simd_lane_id; kc < BC; kc += 32) {
+          Ws[simd_group_id][kh][kw][kc] = wt_in[kh * R * C + kw * C + kc];
+        }
+      }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Do transform and store the result
+    for (int c = 0; c < BC; ++c) {
+      simdgroup_matrix<float, 8, 8> g;
+      g.thread_elements()[0] =
+          sm < R && sn < R ? Ws[simd_group_id][sm][sn][c] : T(0);
+      g.thread_elements()[1] =
+          sm < R && sn + 1 < R ? Ws[simd_group_id][sm][sn + 1][c] : T(0);
+
+      simdgroup_matrix<float, 8, 8> g_out = (G * g) * Gt;
+      wt_out_0[c * O] = static_cast<T>(g_out.thread_elements()[0]);
+      wt_out_1[c * O] = static_cast<T>(g_out.thread_elements()[1]);
+    }
+
+    wt_in += BC;
+    wt_out_0 += BC * O;
+    wt_out_1 += BC * O;
+  }
+}
+
+#define instantiate_winograd_conv_2d_weight_transform_base(name, itype, bc)   \
+  template [[host_name(                                                       \
+      "winograd_conv_2d_weight_transform_" #name "_bc" #bc)]] [[kernel]] void \
+  winograd_conv_2d_weight_transform<itype, bc>(                               \
+      const device itype* wt_in [[buffer(0)]],                                \
+      device itype* wt_out [[buffer(1)]],                                     \
+      const constant int& C [[buffer(2)]],                                    \
+      const constant int& O [[buffer(3)]],                                    \
+      uint tid [[threadgroup_position_in_grid]],                              \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],                  \
+      uint simd_lane_id [[thread_index_in_simdgroup]]);
+
+template <typename T, int BC, int WM, int WN, int M = 6, int R = 3>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
+winograd_conv_2d_input_transform(
+    const device T* inp_in [[buffer(0)]],
+    device T* inp_out [[buffer(1)]],
+    const constant MLXConvParams<2>& params [[buffer(2)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 tgp_per_grid [[threadgroups_per_grid]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  using WGT = WinogradTransforms<M, R, 8>;
+  constexpr int A = WGT::IN_TILE_SIZE;
+  constexpr int N_SIMD_GROUPS = WM * WN;
+
+  // Get lane position in simdgroup
+  const short qid = simd_lane_id / 4;
+  const short sm = (qid & 4) + (simd_lane_id / 2) % 4;
+  const short sn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
+
+  // Initialize B matrix
+  simdgroup_matrix<float, 8, 8> B;
+  B.thread_elements()[0] = WGT::in_transform[sm][sn];
+  B.thread_elements()[1] = WGT::in_transform[sm][sn + 1];
+
+  // Initialize Bt matrix
+  simdgroup_matrix<float, 8, 8> Bt;
+  Bt.thread_elements()[0] = WGT::in_transform[sn][sm];
+  Bt.thread_elements()[1] = WGT::in_transform[sn + 1][sm];
+
+  // Resolve input tile
+  constexpr int TH = (A / WM);
+  constexpr int TW = (A / WN);
+  int kh = TH * (simd_group_id / WN);
+  int kw = TW * (simd_group_id % WN);
+  int bh = M * tid.y + kh;
+  int bw = M * tid.x + kw;
+
+  // Move to the correct input tile
+  inp_in += tid.z * params.in_strides[0] + bh * params.in_strides[1] +
+      bw * params.in_strides[2];
+
+  // Pre compute strides
+  int jump_in[TH][TW];
+
+  for (int h = 0; h < TH; h++) {
+    for (int w = 0; w < TW; w++) {
+      jump_in[h][w] = h * params.in_strides[1] + w * params.in_strides[2];
+    }
+  }
+
+  // inp_out is stored interleaved (A x A x tiles x C)
+  size_t N_TILES = tgp_per_grid.x * tgp_per_grid.y * tgp_per_grid.z;
+  size_t tile_id =
+      tid.z * tgp_per_grid.x * tgp_per_grid.y + tid.y * tgp_per_grid.x + tid.x;
+  size_t ohw_0 = sm * 8 + sn;
+  size_t ohw_1 = sm * 8 + sn + 1;
+  device T* inp_out_0 =
+      inp_out + ohw_0 * N_TILES * params.C + tile_id * params.C;
+  device T* inp_out_1 =
+      inp_out + ohw_1 * N_TILES * params.C + tile_id * params.C;
+
+  // Prepare shared memory
+  threadgroup T Is[A][A][BC];
+
+  // Loop over C
+  for (int bc = 0; bc < params.C; bc += BC) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Read into shared memory
+    for (int h = 0; h < TH; h++) {
+      for (int w = 0; w < TW; w++) {
+        const device T* in_ptr = inp_in + jump_in[h][w];
+        for (int c = simd_lane_id; c < BC; c += 32) {
+          Is[kh + h][kw + w][c] = in_ptr[c];
+        }
+      }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Do transform and store the result
+    for (int c = simd_group_id; c < BC; c += N_SIMD_GROUPS) {
+      simdgroup_matrix<float, 8, 8> I;
+      I.thread_elements()[0] = Is[sm][sn][c];
+      I.thread_elements()[1] = Is[sm][sn + 1][c];
+
+      simdgroup_matrix<float, 8, 8> I_out = (Bt * I) * B;
+      inp_out_0[c] = static_cast<T>(I_out.thread_elements()[0]);
+      inp_out_1[c] = static_cast<T>(I_out.thread_elements()[1]);
+    }
+
+    inp_in += BC;
+    inp_out_0 += BC;
+    inp_out_1 += BC;
+  }
+}
+
+#define instantiate_winograd_conv_2d_input_transform(name, itype, bc)        \
+  template [[host_name(                                                      \
+      "winograd_conv_2d_input_transform_" #name "_bc" #bc)]] [[kernel]] void \
+  winograd_conv_2d_input_transform<itype, bc, 2, 2>(                         \
+      const device itype* inp_in [[buffer(0)]],                              \
+      device itype* inp_out [[buffer(1)]],                                   \
+      const constant MLXConvParams<2>& params [[buffer(2)]],                 \
+      uint3 tid [[threadgroup_position_in_grid]],                            \
+      uint3 lid [[thread_position_in_threadgroup]],                          \
+      uint3 tgp_per_grid [[threadgroups_per_grid]],                          \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],                 \
+      uint simd_lane_id [[thread_index_in_simdgroup]]);
+
+template <typename T, int BO, int WM, int WN, int M = 6, int R = 3>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
+winograd_conv_2d_output_transform(
+    const device T* out_in [[buffer(0)]],
+    device T* out_out [[buffer(1)]],
+    const constant MLXConvParams<2>& params [[buffer(2)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 tgp_per_grid [[threadgroups_per_grid]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  using WGT = WinogradTransforms<M, R, 8>;
+  constexpr int N_SIMD_GROUPS = WM * WN;
+
+  // Get lane position in simdgroup
+  const short qid = simd_lane_id / 4;
+  const short sm = (qid & 4) + (simd_lane_id / 2) % 4;
+  const short sn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
+
+  // Initialize A matrix
+  simdgroup_matrix<float, 8, 8> B;
+  B.thread_elements()[0] = WGT::out_transform[sm][sn];
+  B.thread_elements()[1] = WGT::out_transform[sm][sn + 1];
+
+  // Initialize At matrix
+  simdgroup_matrix<float, 8, 8> Bt;
+  Bt.thread_elements()[0] = WGT::out_transform[sn][sm];
+  Bt.thread_elements()[1] = WGT::out_transform[sn + 1][sm];
+
+  // Out_in comes in shape (A x A x tiles x O)
+  // We do transform and then write out to out_out in shape (N, H, W, O)
+
+  // Resolve output tile
+  constexpr int TH = (M / WM);
+  constexpr int TW = (M / WN);
+  int kh = TH * (simd_group_id / WN);
+  int kw = TW * (simd_group_id % WN);
+  int bh = M * tid.y + kh;
+  int bw = M * tid.x + kw;
+
+  // Move to the correct input tile
+  out_out += tid.z * params.out_strides[0] + bh * params.out_strides[1] +
+      bw * params.out_strides[2];
+
+  // Pre compute strides
+  int jump_in[TH][TW];
+
+  for (int h = 0; h < TH; h++) {
+    for (int w = 0; w < TW; w++) {
+      bool valid = ((bh + h) < params.oS[0]) && ((bw + w) < params.oS[1]);
+      jump_in[h][w] =
+          valid ? h * params.out_strides[1] + w * params.out_strides[2] : -1;
+    }
+  }
+
+  // out_in is stored interleaved (A x A x tiles x O)
+  size_t N_TILES = tgp_per_grid.x * tgp_per_grid.y * tgp_per_grid.z;
+  size_t tile_id =
+      tid.z * tgp_per_grid.x * tgp_per_grid.y + tid.y * tgp_per_grid.x + tid.x;
+  size_t ohw_0 = sm * 8 + sn;
+  size_t ohw_1 = sm * 8 + sn + 1;
+  const device T* out_in_0 =
+      out_in + ohw_0 * N_TILES * params.O + tile_id * params.O;
+  const device T* out_in_1 =
+      out_in + ohw_1 * N_TILES * params.O + tile_id * params.O;
+
+  // Prepare shared memory
+  threadgroup T Os[M][M][BO];
+
+  // Loop over O
+  for (int bo = 0; bo < params.O; bo += BO) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Do transform and store the result
+    for (int c = simd_group_id; c < BO; c += N_SIMD_GROUPS) {
+      simdgroup_matrix<float, 8, 8> O_mat;
+      O_mat.thread_elements()[0] = out_in_0[c];
+      O_mat.thread_elements()[1] = out_in_1[c];
+
+      simdgroup_matrix<float, 8, 8> O_out = (Bt * (O_mat * B));
+      if ((sm < M) && (sn < M)) {
+        Os[sm][sn][c] = static_cast<T>(O_out.thread_elements()[0]);
+      }
+      if ((sm < M) && ((sn + 1) < M)) {
+        Os[sm][sn + 1][c] = static_cast<T>(O_out.thread_elements()[1]);
+      }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Read out from shared memory
+    for (int h = 0; h < TH; h++) {
+      for (int w = 0; w < TW; w++) {
+        if (jump_in[h][w] >= 0) {
+          device T* out_ptr = out_out + jump_in[h][w];
+          for (int c = simd_lane_id; c < BO; c += 32) {
+            out_ptr[c] = Os[kh + h][kw + w][c];
+          }
+        }
+      }
+    }
+
+    out_out += BO;
+    out_in_0 += BO;
+    out_in_1 += BO;
+  }
+}
+
+#define instantiate_winograd_conv_2d_output_transform(name, itype, bo)        \
+  template [[host_name(                                                       \
+      "winograd_conv_2d_output_transform_" #name "_bo" #bo)]] [[kernel]] void \
+  winograd_conv_2d_output_transform<itype, bo, 2, 2>(                         \
+      const device itype* out_in [[buffer(0)]],                               \
+      device itype* out_out [[buffer(1)]],                                    \
+      const constant MLXConvParams<2>& params [[buffer(2)]],                  \
+      uint3 tid [[threadgroup_position_in_grid]],                             \
+      uint3 lid [[thread_position_in_threadgroup]],                           \
+      uint3 tgp_per_grid [[threadgroups_per_grid]],                           \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],                  \
+      uint simd_lane_id [[thread_index_in_simdgroup]]);
+
+// clang-format off
+#define instantiate_winograd_conv_2d(name, itype)                     \
+  instantiate_winograd_conv_2d_weight_transform_base(name, itype, 32) \
+  instantiate_winograd_conv_2d_input_transform(name, itype, 32)       \
+  instantiate_winograd_conv_2d_output_transform(name, itype, 32) // clang-format on
+
+// clang-format off
+instantiate_winograd_conv_2d(float32, float);
+instantiate_winograd_conv_2d(bfloat16, bfloat16_t);
+instantiate_winograd_conv_2d(float16, half); // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/copy.h b/macos/third_party/mlx/mlx/backend/metal/kernels/copy.h
new file mode 100644
index 0000000..cf22347
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/copy.h
@@ -0,0 +1,276 @@
+// Copyright © 2024 Apple Inc.
+
+template <typename T, typename U, int N = WorkPerThread<U>::n>
+[[kernel]] void copy_s(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      dst[index + i] = static_cast<U>(src[0]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      dst[index + i] = static_cast<U>(src[0]);
+    }
+  }
+}
+
+template <typename T, typename U, int N = WorkPerThread<U>::n>
+[[kernel]] void copy_v(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      dst[index + i] = static_cast<U>(src[index + i]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      dst[index + i] = static_cast<U>(src[index + i]);
+    }
+  }
+}
+
+template <typename T, typename U, int N = WorkPerThread<U>::n>
+[[kernel]] void copy_s2(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      dst[offset + i] = static_cast<U>(src[0]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      dst[offset + i] = static_cast<U>(src[0]);
+    }
+  }
+}
+
+template <typename T, typename U, int N = WorkPerThread<U>::n>
+[[kernel]] void copy_v2(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      dst[offset + i] = static_cast<U>(src[offset + i]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      dst[offset + i] = static_cast<U>(src[offset + i]);
+    }
+  }
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_g_nd1(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t& src_stride [[buffer(3)]],
+    uint index [[thread_position_in_grid]]) {
+  auto src_idx = elem_to_loc_1<IdxT>(index, src_stride);
+  dst[index] = static_cast<U>(src[src_idx]);
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_g_nd2(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto src_idx = elem_to_loc_2<IdxT>(index, src_strides);
+  IdxT dst_idx = index.x + IdxT(grid_dim.x) * index.y;
+  dst[dst_idx] = static_cast<U>(src[src_idx]);
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_g_nd3(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto src_idx = elem_to_loc_3<IdxT>(index, src_strides);
+  IdxT dst_idx =
+      index.x + IdxT(grid_dim.x) * (index.y + IdxT(grid_dim.y) * index.z);
+  dst[dst_idx] = static_cast<U>(src[src_idx]);
+}
+
+template <typename T, typename U, int N = 1, typename IdxT = int64_t>
+[[kernel]] void copy_g(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int* src_shape [[buffer(2)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int& ndim [[buffer(5)]],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto src_idx = elem_to_loc<IdxT>(
+      {N * index.x, index.y, index.z}, src_shape, src_strides, ndim);
+  if (N == 1) {
+    IdxT dst_idx =
+        index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z);
+    dst[dst_idx] = static_cast<U>(src[src_idx]);
+    return;
+  }
+  auto xshape = src_shape[ndim - 1];
+  IdxT dst_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z);
+  auto src_xstride = src_strides[ndim - 1];
+  for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) {
+    dst[dst_idx + i] = static_cast<U>(src[src_idx]);
+    src_idx += src_xstride;
+  }
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_gg_nd1(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t& src_stride [[buffer(3)]],
+    constant const int64_t& dst_stride [[buffer(4)]],
+    uint index [[thread_position_in_grid]]) {
+  auto src_idx = elem_to_loc_1<IdxT>(index, src_stride);
+  auto dst_idx = elem_to_loc_1<IdxT>(index, dst_stride);
+  dst[dst_idx] = static_cast<U>(src[src_idx]);
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_gg_nd2(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
+    uint2 index [[thread_position_in_grid]]) {
+  auto src_idx = elem_to_loc_2<IdxT>(index, src_strides);
+  auto dst_idx = elem_to_loc_2<IdxT>(index, dst_strides);
+  dst[dst_idx] = static_cast<U>(src[src_idx]);
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_gg_nd3(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
+    uint3 index [[thread_position_in_grid]]) {
+  auto src_idx = elem_to_loc_3<IdxT>(index, src_strides);
+  auto dst_idx = elem_to_loc_3<IdxT>(index, dst_strides);
+  dst[dst_idx] = static_cast<U>(src[src_idx]);
+}
+
+template <typename T, typename U, int N = 1, typename IdxT = int64_t>
+[[kernel]] void copy_gg(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int* src_shape [[buffer(2)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
+    constant const int& ndim [[buffer(5)]],
+    uint3 index [[thread_position_in_grid]]) {
+  auto idx = elem_to_loc_2_nd<IdxT>(
+      {N * index.x, index.y, index.z},
+      src_shape,
+      src_strides,
+      dst_strides,
+      ndim);
+  if (N == 1) {
+    dst[idx.y] = static_cast<U>(src[idx.x]);
+    return;
+  }
+  IdxT src_xstride = src_strides[ndim - 1];
+  IdxT dst_xstride = dst_strides[ndim - 1];
+  auto xshape = src_shape[ndim - 1];
+  for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) {
+    dst[idx.y] = static_cast<U>(src[idx.x]);
+    idx.x += src_xstride;
+    idx.y += dst_xstride;
+  }
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_gg_dynamic_nd1(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t& src_stride [[buffer(3)]],
+    constant const int64_t& dst_stride [[buffer(4)]],
+    constant const int64_t& src_offset [[buffer(6)]],
+    constant const int64_t& dst_offset [[buffer(7)]],
+    uint index [[thread_position_in_grid]]) {
+  auto src_idx = elem_to_loc_1<IdxT>(index, src_stride);
+  auto dst_idx = elem_to_loc_1<IdxT>(index, dst_stride);
+  dst[dst_idx + dst_offset] = src[src_idx + src_offset];
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_gg_dynamic_nd2(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
+    constant const int64_t& src_offset [[buffer(6)]],
+    constant const int64_t& dst_offset [[buffer(7)]],
+    uint2 index [[thread_position_in_grid]]) {
+  auto src_idx = elem_to_loc_2<IdxT>(index, src_strides);
+  auto dst_idx = elem_to_loc_2<IdxT>(index, dst_strides);
+  dst[dst_idx + dst_offset] = src[src_idx + src_offset];
+}
+
+template <typename T, typename U, typename IdxT = int64_t>
+[[kernel]] void copy_gg_dynamic_nd3(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
+    constant const int64_t& src_offset [[buffer(6)]],
+    constant const int64_t& dst_offset [[buffer(7)]],
+    uint3 index [[thread_position_in_grid]]) {
+  auto src_idx = elem_to_loc_3<IdxT>(index, src_strides);
+  auto dst_idx = elem_to_loc_3<IdxT>(index, dst_strides);
+  dst[dst_idx + dst_offset] = src[src_idx + src_offset];
+}
+
+template <typename T, typename U, int N = 1, typename IdxT = int64_t>
+[[kernel]] void copy_gg_dynamic(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    constant const int* src_shape [[buffer(2)]],
+    constant const int64_t* src_strides [[buffer(3)]],
+    constant const int64_t* dst_strides [[buffer(4)]],
+    constant const int& ndim [[buffer(5)]],
+    constant const int64_t& src_offset [[buffer(6)]],
+    constant const int64_t& dst_offset [[buffer(7)]],
+    uint3 index [[thread_position_in_grid]]) {
+  src += src_offset;
+  dst += dst_offset;
+  auto idx = elem_to_loc_2_nd<IdxT>(
+      {N * index.x, index.y, index.z},
+      src_shape,
+      src_strides,
+      dst_strides,
+      ndim);
+  if (N == 1) {
+    dst[idx.y] = src[idx.x];
+    return;
+  }
+  IdxT src_xstride = src_strides[ndim - 1];
+  IdxT dst_xstride = dst_strides[ndim - 1];
+  auto xshape = src_shape[ndim - 1];
+  for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) {
+    dst[idx.y] = src[idx.x];
+    idx.x += src_xstride;
+    idx.y += dst_xstride;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/copy.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/copy.metal
new file mode 100644
index 0000000..fcf8884
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/copy.metal
@@ -0,0 +1,75 @@
+// Copyright © 2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/copy.h"
+
+#define instantiate_copy_work_per_thread(tname, itype, otype)         \
+  instantiate_kernel("sn_copy" #tname, copy_s, itype, otype)          \
+  instantiate_kernel("vn_copy" #tname, copy_v, itype, otype)
+
+#define instantiate_copy_base(tname, itype, otype)                    \
+  instantiate_kernel("s_copy" #tname, copy_s, itype, otype, 1)        \
+  instantiate_kernel("v_copy" #tname, copy_v, itype, otype, 1)        \
+  instantiate_kernel("s2_copy" #tname, copy_s2, itype, otype)         \
+  instantiate_kernel("v2_copy" #tname, copy_v2, itype, otype)         \
+  instantiate_kernel("g1_copy" #tname, copy_g_nd1, itype, otype, int) \
+  instantiate_kernel("g2_copy" #tname, copy_g_nd2, itype, otype, int) \
+  instantiate_kernel("g3_copy" #tname, copy_g_nd3, itype, otype, int) \
+  instantiate_kernel("gn2_copy" #tname, copy_g, itype, otype, 2, int) \
+  instantiate_kernel("g1large_copy" #tname, copy_g_nd1, itype, otype) \
+  instantiate_kernel("g2large_copy" #tname, copy_g_nd2, itype, otype) \
+  instantiate_kernel("g3large_copy" #tname, copy_g_nd3, itype, otype) \
+  instantiate_kernel("gn4large_copy" #tname, copy_g, itype, otype, 4)
+
+#define instantiate_copy_all(tname, itype, otype) \
+  instantiate_copy_base(tname, itype, otype)      \
+  instantiate_copy_work_per_thread(tname, itype, otype)
+
+#define instantiate_copy_same(tname, type)                                            \
+  instantiate_kernel("gg1_copy" #tname, copy_gg_nd1, type, type, int)                 \
+  instantiate_kernel("gg2_copy" #tname, copy_gg_nd2, type, type, int)                 \
+  instantiate_kernel("gg3_copy" #tname, copy_gg_nd3, type, type, int)                 \
+  instantiate_kernel("ggn2_copy" #tname, copy_gg, type, type, 2, int)                 \
+  instantiate_kernel("gg1large_copy" #tname, copy_gg_nd1, type, type)                 \
+  instantiate_kernel("gg2large_copy" #tname, copy_gg_nd2, type, type)                 \
+  instantiate_kernel("gg3large_copy" #tname, copy_gg_nd3, type, type)                 \
+  instantiate_kernel("ggn4large_copy" #tname, copy_gg, type, type, 4)                 \
+  instantiate_kernel("gg1_dynamic_copy" #tname, copy_gg_dynamic_nd1, type, type, int) \
+  instantiate_kernel("gg2_dynamic_copy" #tname, copy_gg_dynamic_nd2, type, type, int) \
+  instantiate_kernel("gg3_dynamic_copy" #tname, copy_gg_dynamic_nd3, type, type, int) \
+  instantiate_kernel("ggn2_dynamic_copy" #tname, copy_gg_dynamic, type, type, 2, int) \
+  instantiate_kernel("gg1large_dynamic_copy" #tname, copy_gg_dynamic_nd1, type, type) \
+  instantiate_kernel("gg2large_dynamic_copy" #tname, copy_gg_dynamic_nd2, type, type) \
+  instantiate_kernel("gg3large_dynamic_copy" #tname, copy_gg_dynamic_nd3, type, type) \
+  instantiate_kernel("ggn4large_dynamic_copy" #tname, copy_gg_dynamic, type, type, 4)
+
+#define instantiate_copy_itype(itname, itype)                \
+  instantiate_copy_same(itname ##itname, itype)              \
+  instantiate_copy_all(itname ##bool_, itype, bool)          \
+  instantiate_copy_all(itname ##uint8, itype, uint8_t)       \
+  instantiate_copy_all(itname ##uint16, itype, uint16_t)     \
+  instantiate_copy_all(itname ##uint32, itype, uint32_t)     \
+  instantiate_copy_base(itname ##uint64, itype, uint64_t)    \
+  instantiate_copy_all(itname ##int8, itype, int8_t)         \
+  instantiate_copy_all(itname ##int16, itype, int16_t)       \
+  instantiate_copy_all(itname ##int32, itype, int32_t)       \
+  instantiate_copy_base(itname ##int64, itype, int64_t)      \
+  instantiate_copy_all(itname ##float16, itype, half)        \
+  instantiate_copy_all(itname ##float32, itype, float)       \
+  instantiate_copy_all(itname ##bfloat16, itype, bfloat16_t) \
+  instantiate_copy_base(itname ##complex64, itype, complex64_t)
+
+instantiate_copy_itype(bool_, bool)
+instantiate_copy_itype(uint8, uint8_t)
+instantiate_copy_itype(uint16, uint16_t)
+instantiate_copy_itype(uint32, uint32_t)
+instantiate_copy_itype(uint64, uint64_t)
+instantiate_copy_itype(int8, int8_t)
+instantiate_copy_itype(int16, int16_t)
+instantiate_copy_itype(int32, int32_t)
+instantiate_copy_itype(int64, int64_t)
+instantiate_copy_itype(float16, half)
+instantiate_copy_itype(float32, float)
+instantiate_copy_itype(bfloat16, bfloat16_t)
+instantiate_copy_itype(complex64, complex64_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/defines.h b/macos/third_party/mlx/mlx/backend/metal/kernels/defines.h
new file mode 100644
index 0000000..c369adb
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/defines.h
@@ -0,0 +1,24 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#if defined __METAL__ || defined MLX_METAL_JIT
+#define MTL_CONST constant
+#else
+#define MTL_CONST
+#endif
+
+static MTL_CONST constexpr int MAX_REDUCE_SPECIALIZED_DIMS = 4;
+static MTL_CONST constexpr int REDUCE_N_READS = 4;
+static MTL_CONST constexpr int REDUCE_N_WRITES = 4;
+static MTL_CONST constexpr int SOFTMAX_N_READS = 4;
+static MTL_CONST constexpr int RMS_N_READS = 4;
+static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096;
+
+// Instantiate a templated kernel.
+// Extra args are used as template parameters:
+// e.g. instantiate_kernel(binary_int, binary, a, b) ->
+// [[host_name(binary_int)]] [kernel] binary<a, b>
+#define instantiate_kernel(name, func, ...) \
+  template [[host_name(                     \
+      name)]] [[kernel]] decltype(func<__VA_ARGS__>) func<__VA_ARGS__>;
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/erf.h b/macos/third_party/mlx/mlx/backend/metal/kernels/erf.h
new file mode 100644
index 0000000..d367eef
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/erf.h
@@ -0,0 +1,69 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+#include <metal_math>
+#include "mlx/backend/metal/kernels/expm1f.h"
+
+/*
+ * Approximation to the error function.
+ * Based on code from:
+ * https://stackoverflow.com/questions/35148198/efficient-faithfully-rounded-implementation-of-error-function-erff#answer-35148199
+ */
+float erf(float a) {
+  float r, s, t, u;
+  t = metal::abs(a);
+  s = a * a;
+  if (t > 0.927734375f) {
+    // maximum error 0.99527 ulp
+    r = metal::fma(
+        -1.72853470e-5f, t, 3.83197126e-4f); // -0x1.220000p-16,0x1.91cfb2p-12
+    u = metal::fma(
+        -3.88396438e-3f, t, 2.42546219e-2f); // -0x1.fd1438p-9, 0x1.8d6342p-6
+    r = metal::fma(r, s, u);
+    r = metal::fma(r, t, -1.06777877e-1f); // -0x1.b55cb8p-4
+    r = metal::fma(r, t, -6.34846687e-1f); // -0x1.450aa0p-1
+    r = metal::fma(r, t, -1.28717512e-1f); // -0x1.079d0cp-3
+    r = metal::fma(r, t, -t);
+    r = -expm1f(r);
+    r = metal::copysign(r, a);
+  } else {
+    // maximum error 0.98929 ulp
+    r = -5.96761703e-4f; // -0x1.38e000p-11
+    r = metal::fma(r, s, 4.99119423e-3f); //  0x1.471a58p-8
+    r = metal::fma(r, s, -2.67681349e-2f); // -0x1.b691b2p-6
+    r = metal::fma(r, s, 1.12819925e-1f); //  0x1.ce1c44p-4
+    r = metal::fma(r, s, -3.76125336e-1f); // -0x1.812700p-2
+    r = metal::fma(r, s, 1.28379166e-1f); //  0x1.06eba8p-3
+    r = metal::fma(r, a, a);
+  }
+  return r;
+}
+
+float erfinv(float a) {
+  auto t = metal::fma(a, 0.0f - a, 1.0f);
+  t = metal::log(t);
+  float p;
+  if (metal::abs(t) > 6.125f) { // maximum ulp error = 2.35793
+    p = 3.03697567e-10f; //  0x1.4deb44p-32
+    p = metal::fma(p, t, 2.93243101e-8f); //  0x1.f7c9aep-26
+    p = metal::fma(p, t, 1.22150334e-6f); //  0x1.47e512p-20
+    p = metal::fma(p, t, 2.84108955e-5f); //  0x1.dca7dep-16
+    p = metal::fma(p, t, 3.93552968e-4f); //  0x1.9cab92p-12
+    p = metal::fma(p, t, 3.02698812e-3f); //  0x1.8cc0dep-9
+    p = metal::fma(p, t, 4.83185798e-3f); //  0x1.3ca920p-8
+    p = metal::fma(p, t, -2.64646143e-1f); // -0x1.0eff66p-2
+    p = metal::fma(p, t, 8.40016484e-1f); //  0x1.ae16a4p-1
+  } else { // maximum ulp error = 2.35002
+    p = 5.43877832e-9f; //  0x1.75c000p-28
+    p = metal::fma(p, t, 1.43285448e-7f); //  0x1.33b402p-23
+    p = metal::fma(p, t, 1.22774793e-6f); //  0x1.499232p-20
+    p = metal::fma(p, t, 1.12963626e-7f); //  0x1.e52cd2p-24
+    p = metal::fma(p, t, -5.61530760e-5f); // -0x1.d70bd0p-15
+    p = metal::fma(p, t, -1.47697632e-4f); // -0x1.35be90p-13
+    p = metal::fma(p, t, 2.31468678e-3f); //  0x1.2f6400p-9
+    p = metal::fma(p, t, 1.15392581e-2f); //  0x1.7a1e50p-7
+    p = metal::fma(p, t, -2.32015476e-1f); // -0x1.db2aeep-3
+    p = metal::fma(p, t, 8.86226892e-1f); //  0x1.c5bf88p-1
+  }
+  return a * p;
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/expm1f.h b/macos/third_party/mlx/mlx/backend/metal/kernels/expm1f.h
new file mode 100644
index 0000000..68224e1
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/expm1f.h
@@ -0,0 +1,90 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <metal_math>
+
+// Original license copied below:
+//  Copyright (c) 2015-2023 Norbert Juffa
+//  All rights reserved.
+//
+//  Redistribution and use in source and binary forms, with or without
+//  modification, are permitted provided that the following conditions
+//  are met:
+//
+//  1. Redistributions of source code must retain the above copyright
+//     notice, this list of conditions and the following disclaimer.
+//
+//  2. Redistributions in binary form must reproduce the above copyright
+//     notice, this list of conditions and the following disclaimer in the
+//     documentation and/or other materials provided with the distribution.
+//
+//  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+//  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+//  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+//  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+//  HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+//  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+//  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+//  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+//  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+//  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+//  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+/* Compute exponential base e minus 1. Maximum ulp error = 0.997458
+
+   i = rint(a/log(2)), f = a-i*log(2). Then expm1(a) = 2**i * (expm1(f)+1) - 1.
+   Compute r = expm1(f). Then expm1(a)= 2 * (0.5 * 2**i * r + 0.5 * 2**i - 0.5).
+   With t = 0.5*2**i, expm1(a) = 2*(r * t + t-0.5). However, for best accuracy,
+   when i == 1, expm1(a)= 2*(r + 0.5), and when i == 0, expm1(a) = r.
+
+   NOTE: Scale factor b is only applied if i < 0 or i > 1 (should be power of 2)
+*/
+float expm1f_scaled_unchecked(float a, float b) {
+  float f, j, r, s, t, u, v, x, y;
+  int i;
+
+  // exp(a) = 2**i * exp(f); i = rintf (a / log(2))
+  j = fma(1.442695f, a, 12582912.f); // 0x1.715476p0, 0x1.8p23
+  j = j - 12582912.0f; // 0x1.8p23
+  i = (int)j;
+  f = fma(j, -6.93145752e-1f, a);
+
+  // approximate r = exp(f)-1 on interval [-log(2)/2, +log(2)/2]
+  s = f * f;
+  if (a == 0.0f)
+    s = a; // ensure -0 is passed through
+  // err = 0.997458  ulp1 = 11081805
+  r = 1.97350979e-4f; // 0x1.9de000p-13
+  r = fma(r, f, 1.39309070e-3f); // 0x1.6d30bcp-10
+  r = fma(r, f, 8.33343994e-3f); // 0x1.1111f6p-7
+  r = fma(r, f, 4.16668020e-2f); // 0x1.55559ep-5
+  r = fma(r, f, 1.66666716e-1f); // 0x1.55555cp-3
+  r = fma(r, f, 4.99999970e-1f); // 0x1.fffffep-2
+  u = (j == 1) ? (f + 0.5f) : f;
+  v = fma(r, s, u);
+  s = 0.5f * b;
+  t = ldexp(s, i);
+  y = t - s;
+  x = (t - y) - s; // double-float canonicalization of difference
+  r = fma(v, t, x) + y;
+  r = r + r;
+  if (j == 0)
+    r = v;
+  if (j == 1)
+    r = v + v;
+  return r;
+}
+
+/* Compute exponential base e minus 1. max ulp err = 0.99746 */
+float expm1f(float a) {
+  float r;
+
+  r = expm1f_scaled_unchecked(a, 1.0f);
+  /* handle severe overflow and underflow */
+  if (abs(a - 1.0f) > 88.0f) {
+    r = pow(2, a);
+    r = fma(r, r, -1.0f);
+  }
+  return r;
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fence.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/fence.metal
new file mode 100644
index 0000000..cd5f901
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fence.metal
@@ -0,0 +1,52 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma METAL internals : enable
+
+#ifndef __METAL_MEMORY_SCOPE_SYSTEM__
+#define __METAL_MEMORY_SCOPE_SYSTEM__ 3
+#endif
+namespace metal {
+constexpr constant metal::thread_scope thread_scope_system =
+    static_cast<thread_scope>(__METAL_MEMORY_SCOPE_SYSTEM__);
+}
+
+#include <metal_atomic>
+
+[[kernel]] void input_coherent(
+    volatile coherent(system) device uint* input [[buffer(0)]],
+    const constant uint& size [[buffer(1)]],
+    uint index [[thread_position_in_grid]]) {
+  if (index < size) {
+    input[index] = input[index];
+  }
+  metal::atomic_thread_fence(
+      metal::mem_flags::mem_device,
+      metal::memory_order_seq_cst,
+      metal::thread_scope_system);
+}
+
+// single thread kernel to update timestamp
+[[kernel]] void fence_update(
+    volatile coherent(system) device uint* timestamp [[buffer(0)]],
+    constant uint& value [[buffer(1)]]) {
+  timestamp[0] = value;
+  metal::atomic_thread_fence(
+      metal::mem_flags::mem_device,
+      metal::memory_order_seq_cst,
+      metal::thread_scope_system);
+}
+
+// single thread kernel to spin wait for timestamp value
+[[kernel]] void fence_wait(
+    volatile coherent(system) device uint* timestamp [[buffer(0)]],
+    constant uint& value [[buffer(1)]]) {
+  while (1) {
+    metal::atomic_thread_fence(
+        metal::mem_flags::mem_device,
+        metal::memory_order_seq_cst,
+        metal::thread_scope_system);
+    if (timestamp[0] >= value) {
+      break;
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fft.h b/macos/third_party/mlx/mlx/backend/metal/kernels/fft.h
new file mode 100644
index 0000000..e478a85
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fft.h
@@ -0,0 +1,486 @@
+// Copyright © 2024 Apple Inc.
+
+// Metal FFT using Stockham's algorithm
+//
+// References:
+// - VkFFT (https://github.com/DTolm/VkFFT)
+// - Eric Bainville's excellent page (http://www.bealto.com/gpu-fft.html)
+
+#include <metal_common>
+
+#include "mlx/backend/metal/kernels/fft/radix.h"
+#include "mlx/backend/metal/kernels/fft/readwrite.h"
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+using namespace metal;
+
+#define MAX_RADIX 13
+// Reached when elems_per_thread_ = 6, max_radix = 13
+// and some threads have to do 3 radix 6s requiring 18 float2s.
+#define MAX_OUTPUT_SIZE 18
+
+// Specialize for a particular value of N at runtime
+STEEL_CONST bool inv_ [[function_constant(0)]];
+STEEL_CONST bool is_power_of_2_ [[function_constant(1)]];
+STEEL_CONST int elems_per_thread_ [[function_constant(2)]];
+// rader_m = n / rader_n
+STEEL_CONST int rader_m_ [[function_constant(3)]];
+// Stockham steps
+STEEL_CONST int radix_13_steps_ [[function_constant(4)]];
+STEEL_CONST int radix_11_steps_ [[function_constant(5)]];
+STEEL_CONST int radix_8_steps_ [[function_constant(6)]];
+STEEL_CONST int radix_7_steps_ [[function_constant(7)]];
+STEEL_CONST int radix_6_steps_ [[function_constant(8)]];
+STEEL_CONST int radix_5_steps_ [[function_constant(9)]];
+STEEL_CONST int radix_4_steps_ [[function_constant(10)]];
+STEEL_CONST int radix_3_steps_ [[function_constant(11)]];
+STEEL_CONST int radix_2_steps_ [[function_constant(12)]];
+// Rader steps
+STEEL_CONST int rader_13_steps_ [[function_constant(13)]];
+STEEL_CONST int rader_11_steps_ [[function_constant(14)]];
+STEEL_CONST int rader_8_steps_ [[function_constant(15)]];
+STEEL_CONST int rader_7_steps_ [[function_constant(16)]];
+STEEL_CONST int rader_6_steps_ [[function_constant(17)]];
+STEEL_CONST int rader_5_steps_ [[function_constant(18)]];
+STEEL_CONST int rader_4_steps_ [[function_constant(19)]];
+STEEL_CONST int rader_3_steps_ [[function_constant(20)]];
+STEEL_CONST int rader_2_steps_ [[function_constant(21)]];
+
+// See "radix.h" for radix codelets
+typedef void (*RadixFunc)(thread float2*, thread float2*);
+
+// Perform a single radix n butterfly with appropriate twiddles
+template <int radix, RadixFunc radix_func>
+METAL_FUNC void radix_butterfly(
+    int i,
+    int p,
+    thread float2* x,
+    thread short* indices,
+    thread float2* y) {
+  // i: the index in the overall DFT that we're processing.
+  // p: the size of the DFTs we're merging at this step.
+  // m: how many threads are working on this DFT.
+  int k, j;
+
+  // Use faster bitwise operations when working with powers of two
+  constexpr bool radix_p_2 = (radix & (radix - 1)) == 0;
+  if (radix_p_2 && is_power_of_2_) {
+    constexpr short power = __builtin_ctz(radix);
+    k = i & (p - 1);
+    j = ((i - k) << power) + k;
+  } else {
+    k = i % p;
+    j = (i / p) * radix * p + k;
+  }
+
+  // Apply twiddles
+  if (p > 1) {
+    float2 twiddle_1 = get_twiddle(k, radix * p);
+    float2 twiddle = twiddle_1;
+    x[1] = complex_mul(x[1], twiddle);
+
+    STEEL_PRAGMA_UNROLL
+    for (int t = 2; t < radix; t++) {
+      twiddle = complex_mul(twiddle, twiddle_1);
+      x[t] = complex_mul(x[t], twiddle);
+    }
+  }
+
+  radix_func(x, y);
+
+  STEEL_PRAGMA_UNROLL
+  for (int t = 0; t < radix; t++) {
+    indices[t] = j + t * p;
+  }
+}
+
+// Perform all the radix steps required for a
+// particular radix size n.
+template <int radix, RadixFunc radix_func>
+METAL_FUNC void radix_n_steps(
+    int i,
+    thread int* p,
+    int m,
+    int n,
+    int num_steps,
+    thread float2* inputs,
+    thread short* indices,
+    thread float2* values,
+    threadgroup float2* buf) {
+  int m_r = n / radix;
+  // When combining different sized radices, we have to do
+  // multiple butterflies in a single thread.
+  // E.g. n = 28 = 4 * 7
+  // 4 threads, 7 elems_per_thread
+  // All threads do 1 radix7 butterfly.
+  // 3 threads do 2 radix4 butterflies.
+  // 1 thread does 1 radix4 butterfly.
+  int max_radices_per_thread = (elems_per_thread_ + radix - 1) / radix;
+
+  int index = 0;
+  int r_index = 0;
+  for (int s = 0; s < num_steps; s++) {
+    for (int t = 0; t < max_radices_per_thread; t++) {
+      index = i + t * m;
+      if (index < m_r) {
+        for (int r = 0; r < radix; r++) {
+          inputs[r] = buf[index + r * m_r];
+        }
+        radix_butterfly<radix, radix_func>(
+            index, *p, inputs, indices + t * radix, values + t * radix);
+      }
+    }
+
+    // Wait until all threads have read their inputs into thread local mem
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    for (int t = 0; t < max_radices_per_thread; t++) {
+      index = i + t * m;
+      if (index < m_r) {
+        for (int r = 0; r < radix; r++) {
+          r_index = t * radix + r;
+          buf[indices[r_index]] = values[r_index];
+        }
+      }
+    }
+
+    // Wait until all threads have written back to threadgroup mem
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    *p *= radix;
+  }
+}
+
+#define RADIX_STEP(radix, radix_func, num_steps) \
+  radix_n_steps<radix, radix_func>(              \
+      fft_idx, p, m, n, num_steps, inputs, indices, values, buf);
+
+template <bool rader = false>
+METAL_FUNC void
+perform_fft(int fft_idx, thread int* p, int m, int n, threadgroup float2* buf) {
+  float2 inputs[MAX_RADIX];
+  short indices[MAX_OUTPUT_SIZE];
+  float2 values[MAX_OUTPUT_SIZE];
+
+  RADIX_STEP(2, radix2, rader ? rader_2_steps_ : radix_2_steps_);
+  RADIX_STEP(3, radix3, rader ? rader_3_steps_ : radix_3_steps_);
+  RADIX_STEP(4, radix4, rader ? rader_4_steps_ : radix_4_steps_);
+  RADIX_STEP(5, radix5, rader ? rader_5_steps_ : radix_5_steps_);
+  RADIX_STEP(6, radix6, rader ? rader_6_steps_ : radix_6_steps_);
+  RADIX_STEP(7, radix7, rader ? rader_7_steps_ : radix_7_steps_);
+  RADIX_STEP(8, radix8, rader ? rader_8_steps_ : radix_8_steps_);
+  RADIX_STEP(11, radix11, rader ? rader_11_steps_ : radix_11_steps_);
+  RADIX_STEP(13, radix13, rader ? rader_13_steps_ : radix_13_steps_);
+}
+
+// Each FFT is computed entirely in shared GPU memory.
+//
+// N is decomposed into radix-n DFTs:
+// e.g. 128 = 2 * 4 * 4 * 4
+template <int tg_mem_size, typename in_T, typename out_T>
+[[kernel]] void fft(
+    const device in_T* in [[buffer(0)]],
+    device out_T* out [[buffer(1)]],
+    constant const int& n,
+    constant const int& batch_size,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  threadgroup float2 shared_in[tg_mem_size];
+
+  thread ReadWriter<in_T, out_T> read_writer = ReadWriter<in_T, out_T>(
+      in,
+      &shared_in[0],
+      out,
+      n,
+      batch_size,
+      elems_per_thread_,
+      elem,
+      grid,
+      inv_);
+
+  if (read_writer.out_of_bounds()) {
+    return;
+  };
+  read_writer.load();
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  int p = 1;
+  int fft_idx = elem.z; // Thread index in DFT
+  int m = grid.z; // Threads per DFT
+  int tg_idx = elem.y * n; // Index of this DFT in threadgroup
+  threadgroup float2* buf = &shared_in[tg_idx];
+
+  perform_fft(fft_idx, &p, m, n, buf);
+
+  read_writer.write();
+}
+
+template <int tg_mem_size, typename in_T, typename out_T>
+[[kernel]] void rader_fft(
+    const device in_T* in [[buffer(0)]],
+    device out_T* out [[buffer(1)]],
+    const device float2* raders_b_q [[buffer(2)]],
+    const device short* raders_g_q [[buffer(3)]],
+    const device short* raders_g_minus_q [[buffer(4)]],
+    constant const int& n,
+    constant const int& batch_size,
+    constant const int& rader_n,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Use Rader's algorithm to compute fast FFTs
+  // when a prime factor `p` of `n` is greater than 13 but
+  // has `p - 1` Stockham decomposable into to prime factors <= 13.
+  //
+  // E.g. n = 102
+  //        = 2 * 3 * 17
+  // .      = 2 * 3 * RADER(16)
+  // .      = 2 * 3 * RADER(4 * 4)
+  //
+  // In numpy:
+  //   x_perm = x[g_q]
+  //   y = np.fft.fft(x_perm) * b_q
+  //   z = np.fft.ifft(y) + x[0]
+  //   out = z[g_minus_q]
+  //   out[0]  = x[1:].sum()
+  //
+  // Where the g_q and g_minus_q are permutations formed
+  // by the group under multiplicative modulo N using the
+  // primitive root of N and b_q is a constant.
+  // See https://en.wikipedia.org/wiki/Rader%27s_FFT_algorithm
+  //
+  // Rader's uses fewer operations than Bluestein's and so
+  // is more accurate. It's also faster in most cases.
+  threadgroup float2 shared_in[tg_mem_size];
+
+  thread ReadWriter<in_T, out_T> read_writer = ReadWriter<in_T, out_T>(
+      in,
+      &shared_in[0],
+      out,
+      n,
+      batch_size,
+      elems_per_thread_,
+      elem,
+      grid,
+      inv_);
+
+  if (read_writer.out_of_bounds()) {
+    return;
+  };
+  read_writer.load();
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // The number of the threads we're using for each DFT
+  int m = grid.z;
+
+  int fft_idx = elem.z;
+  int tg_idx = elem.y * n;
+  threadgroup float2* buf = &shared_in[tg_idx];
+
+  // rader_m = n / rader_n;
+  int rader_m = rader_m_;
+
+  // We have to load two x_0s for each thread since sometimes
+  // elems_per_thread_ crosses a boundary.
+  // E.g. with n = 34, rader_n = 17, elems_per_thread_ = 4
+  // 0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 7 7 7 8 8
+  // 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
+  short x_0_index =
+      metal::min(fft_idx * elems_per_thread_ / (rader_n - 1), rader_m - 1);
+  float2 x_0[2] = {buf[x_0_index], buf[x_0_index + 1]};
+
+  // Do the Rader permutation in shared memory
+  float2 temp[MAX_RADIX];
+  int max_index = n - rader_m - 1;
+  for (int e = 0; e < elems_per_thread_; e++) {
+    short index = metal::min(fft_idx * elems_per_thread_ + e, max_index);
+    short g_q = raders_g_q[index / rader_m];
+    temp[e] = buf[rader_m + (g_q - 1) * rader_m + index % rader_m];
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  for (int e = 0; e < elems_per_thread_; e++) {
+    short index = metal::min(fft_idx * elems_per_thread_ + e, max_index);
+    buf[index + rader_m] = temp[e];
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Rader FFT on x[rader_m:]
+  int p = 1;
+  perform_fft</*rader=*/true>(fft_idx, &p, m, n - rader_m, buf + rader_m);
+
+  // x_1 + ... + x_n is computed for us in the first FFT step so
+  // we save it in the first rader_m indices of the array for later.
+  int x_sum_index = metal::min(fft_idx, rader_m - 1);
+  buf[x_sum_index] = buf[rader_m + x_sum_index * (rader_n - 1)];
+
+  float2 inv = {1.0f, -1.0f};
+  for (int e = 0; e < elems_per_thread_; e++) {
+    short index = metal::min(fft_idx * elems_per_thread_ + e, max_index);
+    short interleaved_index =
+        index / rader_m + (index % rader_m) * (rader_n - 1);
+    temp[e] = complex_mul(
+        buf[rader_m + interleaved_index],
+        raders_b_q[interleaved_index % (rader_n - 1)]);
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  for (int e = 0; e < elems_per_thread_; e++) {
+    short index = metal::min(fft_idx * elems_per_thread_ + e, max_index);
+    buf[rader_m + index] = temp[e] * inv;
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Rader IFFT on x[rader_m:]
+  p = 1;
+  perform_fft</*rader=*/true>(fft_idx, &p, m, n - rader_m, buf + rader_m);
+
+  float2 rader_inv_factor = {1.0f / (rader_n - 1), -1.0f / (rader_n - 1)};
+
+  for (int e = 0; e < elems_per_thread_; e++) {
+    short index = metal::min(fft_idx * elems_per_thread_ + e, n - rader_m - 1);
+    short diff_index = index / (rader_n - 1) - x_0_index;
+    temp[e] = buf[rader_m + index] * rader_inv_factor + x_0[diff_index];
+  }
+
+  // Use the sum of elements that was computed in the first FFT
+  float2 x_sum = buf[x_0_index] + x_0[0];
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  for (int e = 0; e < elems_per_thread_; e++) {
+    short index = metal::min(fft_idx * elems_per_thread_ + e, max_index);
+    short g_q_index = index % (rader_n - 1);
+    short g_q = raders_g_minus_q[g_q_index];
+    short out_index = index - g_q_index + g_q + (index / (rader_n - 1));
+    buf[out_index] = temp[e];
+  }
+
+  buf[x_0_index * rader_n] = x_sum;
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  p = rader_n;
+  perform_fft(fft_idx, &p, m, n, buf);
+
+  read_writer.write();
+}
+
+template <int tg_mem_size, typename in_T, typename out_T>
+[[kernel]] void bluestein_fft(
+    const device in_T* in [[buffer(0)]],
+    device out_T* out [[buffer(1)]],
+    const device float2* w_q [[buffer(2)]],
+    const device float2* w_k [[buffer(3)]],
+    constant const int& length,
+    constant const int& n,
+    constant const int& batch_size,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Computes arbitrary length FFTs with Bluestein's algorithm
+  //
+  // In numpy:
+  //   bluestein_n = next_power_of_2(2*n - 1)
+  //   out = w_k * np.fft.ifft(np.fft.fft(w_k * in, bluestein_n) * w_q)
+  //
+  // Where w_k and w_q are precomputed on CPU in high precision as:
+  //   w_k = np.exp(-1j * np.pi / n * (np.arange(-n + 1, n) ** 2))
+  //   w_q = np.fft.fft(1/w_k[-n:])
+  threadgroup float2 shared_in[tg_mem_size];
+
+  thread ReadWriter<in_T, out_T> read_writer = ReadWriter<in_T, out_T>(
+      in,
+      &shared_in[0],
+      out,
+      n,
+      batch_size,
+      elems_per_thread_,
+      elem,
+      grid,
+      inv_);
+
+  if (read_writer.out_of_bounds()) {
+    return;
+  };
+  read_writer.load_padded(length, w_k);
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  int p = 1;
+  int fft_idx = elem.z; // Thread index in DFT
+  int m = grid.z; // Threads per DFT
+  int tg_idx = elem.y * n; // Index of this DFT in threadgroup
+  threadgroup float2* buf = &shared_in[tg_idx];
+
+  // fft
+  perform_fft(fft_idx, &p, m, n, buf);
+
+  float2 inv = float2(1.0f, -1.0f);
+  for (int t = 0; t < elems_per_thread_; t++) {
+    int index = fft_idx + t * m;
+    buf[index] = complex_mul(buf[index], w_q[index]) * inv;
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // ifft
+  p = 1;
+  perform_fft(fft_idx, &p, m, n, buf);
+
+  read_writer.write_padded(length, w_k);
+}
+
+template <
+    int tg_mem_size,
+    typename in_T,
+    typename out_T,
+    int step,
+    bool real = false>
+[[kernel]] void four_step_fft(
+    const device in_T* in [[buffer(0)]],
+    device out_T* out [[buffer(1)]],
+    constant const int& n1,
+    constant const int& n2,
+    constant const int& batch_size,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Fast four step FFT implementation for powers of 2.
+  int overall_n = n1 * n2;
+  int n = step == 0 ? n1 : n2;
+  int stride = step == 0 ? n2 : n1;
+
+  // The number of the threads we're using for each DFT
+  int m = grid.z;
+  int fft_idx = elem.z;
+
+  threadgroup float2 shared_in[tg_mem_size];
+  threadgroup float2* buf = &shared_in[elem.y * n];
+
+  using read_writer_t = ReadWriter<in_T, out_T, step, real>;
+  read_writer_t read_writer = read_writer_t(
+      in,
+      &shared_in[0],
+      out,
+      n,
+      batch_size,
+      elems_per_thread_,
+      elem,
+      grid,
+      inv_);
+
+  if (read_writer.out_of_bounds()) {
+    return;
+  };
+  read_writer.load_strided(stride, overall_n);
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  int p = 1;
+  perform_fft(fft_idx, &p, m, n, buf);
+
+  read_writer.write_strided(stride, overall_n);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fft.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/fft.metal
new file mode 100644
index 0000000..590b558
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fft.metal
@@ -0,0 +1,67 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/fft.h"
+
+#define instantiate_fft(tg_mem_size, in_T, out_T)   \
+  instantiate_kernel(                               \
+      "fft_mem_" #tg_mem_size "_" #in_T "_" #out_T, \
+      fft,                                          \
+      tg_mem_size,                                  \
+      in_T,                                         \
+      out_T)
+
+#define instantiate_rader(tg_mem_size, in_T, out_T)       \
+  instantiate_kernel(                                     \
+      "rader_fft_mem_" #tg_mem_size "_" #in_T "_" #out_T, \
+      rader_fft,                                          \
+      tg_mem_size,                                        \
+      in_T,                                               \
+      out_T)
+
+#define instantiate_bluestein(tg_mem_size, in_T, out_T)       \
+  instantiate_kernel(                                         \
+      "bluestein_fft_mem_" #tg_mem_size "_" #in_T "_" #out_T, \
+      bluestein_fft,                                          \
+      tg_mem_size,                                            \
+      in_T,                                                   \
+      out_T)
+
+#define instantiate_four_step(tg_mem_size, in_T, out_T, step, real)           \
+  instantiate_kernel(                                                         \
+      "four_step_mem_" #tg_mem_size "_" #in_T "_" #out_T "_" #step "_" #real, \
+      four_step_fft,                                                          \
+      tg_mem_size,                                                            \
+      in_T,                                                                   \
+      out_T,                                                                  \
+      step,                                                                   \
+      real)
+
+// clang-format off
+#define instantiate_ffts(tg_mem_size)                        \
+  instantiate_fft(tg_mem_size, float2, float2) \
+  instantiate_fft(tg_mem_size, float, float2) \
+  instantiate_fft(tg_mem_size, float2, float) \
+  instantiate_rader(tg_mem_size, float2, float2) \
+  instantiate_rader(tg_mem_size, float, float2) \
+  instantiate_rader(tg_mem_size, float2, float) \
+  instantiate_bluestein(tg_mem_size, float2, float2) \
+  instantiate_bluestein(tg_mem_size, float, float2) \
+  instantiate_bluestein(tg_mem_size, float2, float) \
+  instantiate_four_step(tg_mem_size, float2, float2, 0, /*real=*/false) \
+  instantiate_four_step(tg_mem_size, float2, float2, 1, /*real=*/false) \
+  instantiate_four_step(tg_mem_size, float, float2, 0, /*real=*/true) \
+  instantiate_four_step(tg_mem_size, float2, float2, 1, /*real=*/true) \
+  instantiate_four_step(tg_mem_size, float2, float2, 0, /*real=*/true) \
+  instantiate_four_step(tg_mem_size, float2, float, 1, /*real=*/true)
+
+// It's substantially faster to statically define the
+// threadgroup memory size rather than using
+// `setThreadgroupMemoryLength` on the compute encoder.
+// For non-power of 2 sizes we round up the shared memory.
+instantiate_ffts(256)
+instantiate_ffts(512)
+instantiate_ffts(1024)
+instantiate_ffts(2048)
+// 4096 is the max that will fit into 32KB of threadgroup memory.
+instantiate_ffts(4096) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fft/radix.h b/macos/third_party/mlx/mlx/backend/metal/kernels/fft/radix.h
new file mode 100644
index 0000000..bd61eef
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fft/radix.h
@@ -0,0 +1,328 @@
+// Copyright © 2024 Apple Inc.
+
+/* Radix kernels
+
+We provide optimized, single threaded Radix codelets
+for n=2,3,4,5,6,7,8,10,11,12,13.
+
+For n=2,3,4,5,6 we hand write the codelets.
+For n=8,10,12 we combine smaller codelets.
+For n=7,11,13 we use Rader's algorithm which decomposes
+them into (n-1)=6,10,12 codelets. */
+
+#pragma once
+
+#include <metal_common>
+#include <metal_math>
+#include <metal_stdlib>
+
+METAL_FUNC float2 complex_mul(float2 a, float2 b) {
+  return float2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
+}
+
+// Complex mul followed by conjugate
+METAL_FUNC float2 complex_mul_conj(float2 a, float2 b) {
+  return float2(a.x * b.x - a.y * b.y, -a.x * b.y - a.y * b.x);
+}
+
+// Compute an FFT twiddle factor
+METAL_FUNC float2 get_twiddle(int k, int p) {
+  float theta = -2.0f * k * M_PI_F / p;
+
+  float2 twiddle = {metal::fast::cos(theta), metal::fast::sin(theta)};
+  return twiddle;
+}
+
+METAL_FUNC void radix2(thread float2* x, thread float2* y) {
+  y[0] = x[0] + x[1];
+  y[1] = x[0] - x[1];
+}
+
+METAL_FUNC void radix3(thread float2* x, thread float2* y) {
+  float pi_2_3 = -0.8660254037844387;
+
+  float2 a_1 = x[1] + x[2];
+  float2 a_2 = x[1] - x[2];
+
+  y[0] = x[0] + a_1;
+  float2 b_1 = x[0] - 0.5 * a_1;
+  float2 b_2 = pi_2_3 * a_2;
+
+  float2 b_2_j = {-b_2.y, b_2.x};
+  y[1] = b_1 + b_2_j;
+  y[2] = b_1 - b_2_j;
+}
+
+METAL_FUNC void radix4(thread float2* x, thread float2* y) {
+  float2 z_0 = x[0] + x[2];
+  float2 z_1 = x[0] - x[2];
+  float2 z_2 = x[1] + x[3];
+  float2 z_3 = x[1] - x[3];
+  float2 z_3_i = {z_3.y, -z_3.x};
+
+  y[0] = z_0 + z_2;
+  y[1] = z_1 + z_3_i;
+  y[2] = z_0 - z_2;
+  y[3] = z_1 - z_3_i;
+}
+
+METAL_FUNC void radix5(thread float2* x, thread float2* y) {
+  float2 root_5_4 = 0.5590169943749475;
+  float2 sin_2pi_5 = 0.9510565162951535;
+  float2 sin_1pi_5 = 0.5877852522924731;
+
+  float2 a_1 = x[1] + x[4];
+  float2 a_2 = x[2] + x[3];
+  float2 a_3 = x[1] - x[4];
+  float2 a_4 = x[2] - x[3];
+
+  float2 a_5 = a_1 + a_2;
+  float2 a_6 = root_5_4 * (a_1 - a_2);
+  float2 a_7 = x[0] - a_5 / 4;
+  float2 a_8 = a_7 + a_6;
+  float2 a_9 = a_7 - a_6;
+  float2 a_10 = sin_2pi_5 * a_3 + sin_1pi_5 * a_4;
+  float2 a_11 = sin_1pi_5 * a_3 - sin_2pi_5 * a_4;
+  float2 a_10_j = {a_10.y, -a_10.x};
+  float2 a_11_j = {a_11.y, -a_11.x};
+
+  y[0] = x[0] + a_5;
+  y[1] = a_8 + a_10_j;
+  y[2] = a_9 + a_11_j;
+  y[3] = a_9 - a_11_j;
+  y[4] = a_8 - a_10_j;
+}
+
+METAL_FUNC void radix6(thread float2* x, thread float2* y) {
+  float sin_pi_3 = 0.8660254037844387;
+  float2 a_1 = x[2] + x[4];
+  float2 a_2 = x[0] - a_1 / 2;
+  float2 a_3 = sin_pi_3 * (x[2] - x[4]);
+  float2 a_4 = x[5] + x[1];
+  float2 a_5 = x[3] - a_4 / 2;
+  float2 a_6 = sin_pi_3 * (x[5] - x[1]);
+  float2 a_7 = x[0] + a_1;
+
+  float2 a_3_i = {a_3.y, -a_3.x};
+  float2 a_6_i = {a_6.y, -a_6.x};
+  float2 a_8 = a_2 + a_3_i;
+  float2 a_9 = a_2 - a_3_i;
+  float2 a_10 = x[3] + a_4;
+  float2 a_11 = a_5 + a_6_i;
+  float2 a_12 = a_5 - a_6_i;
+
+  y[0] = a_7 + a_10;
+  y[1] = a_8 - a_11;
+  y[2] = a_9 + a_12;
+  y[3] = a_7 - a_10;
+  y[4] = a_8 + a_11;
+  y[5] = a_9 - a_12;
+}
+
+METAL_FUNC void radix7(thread float2* x, thread float2* y) {
+  // Rader's algorithm
+  float2 inv = {1 / 6.0, -1 / 6.0};
+
+  // fft
+  float2 in1[6] = {x[1], x[3], x[2], x[6], x[4], x[5]};
+  radix6(in1, y + 1);
+
+  y[0] = y[1] + x[0];
+
+  // b_q
+  y[1] = complex_mul_conj(y[1], float2(-1, 0));
+  y[2] = complex_mul_conj(y[2], float2(2.44013336, -1.02261879));
+  y[3] = complex_mul_conj(y[3], float2(2.37046941, -1.17510629));
+  y[4] = complex_mul_conj(y[4], float2(0, -2.64575131));
+  y[5] = complex_mul_conj(y[5], float2(2.37046941, 1.17510629));
+  y[6] = complex_mul_conj(y[6], float2(-2.44013336, -1.02261879));
+
+  // ifft
+  radix6(y + 1, x + 1);
+
+  y[1] = x[1] * inv + x[0];
+  y[5] = x[2] * inv + x[0];
+  y[4] = x[3] * inv + x[0];
+  y[6] = x[4] * inv + x[0];
+  y[2] = x[5] * inv + x[0];
+  y[3] = x[6] * inv + x[0];
+}
+
+METAL_FUNC void radix8(thread float2* x, thread float2* y) {
+  float cos_pi_4 = 0.7071067811865476;
+  float2 w_0 = {cos_pi_4, -cos_pi_4};
+  float2 w_1 = {-cos_pi_4, -cos_pi_4};
+  float2 temp[8] = {x[0], x[2], x[4], x[6], x[1], x[3], x[5], x[7]};
+  radix4(temp, x);
+  radix4(temp + 4, x + 4);
+
+  y[0] = x[0] + x[4];
+  y[4] = x[0] - x[4];
+  float2 x_5 = complex_mul(x[5], w_0);
+  y[1] = x[1] + x_5;
+  y[5] = x[1] - x_5;
+  float2 x_6 = {x[6].y, -x[6].x};
+  y[2] = x[2] + x_6;
+  y[6] = x[2] - x_6;
+  float2 x_7 = complex_mul(x[7], w_1);
+  y[3] = x[3] + x_7;
+  y[7] = x[3] - x_7;
+}
+
+template <bool raders_perm>
+METAL_FUNC void radix10(thread float2* x, thread float2* y) {
+  float2 w[4];
+  w[0] = {0.8090169943749475, -0.5877852522924731};
+  w[1] = {0.30901699437494745, -0.9510565162951535};
+  w[2] = {-w[1].x, w[1].y};
+  w[3] = {-w[0].x, w[0].y};
+
+  if (raders_perm) {
+    float2 temp[10] = {
+        x[0], x[3], x[4], x[8], x[2], x[1], x[7], x[9], x[6], x[5]};
+    radix5(temp, x);
+    radix5(temp + 5, x + 5);
+  } else {
+    float2 temp[10] = {
+        x[0], x[2], x[4], x[6], x[8], x[1], x[3], x[5], x[7], x[9]};
+    radix5(temp, x);
+    radix5(temp + 5, x + 5);
+  }
+
+  y[0] = x[0] + x[5];
+  y[5] = x[0] - x[5];
+  for (int t = 1; t < 5; t++) {
+    float2 a = complex_mul(x[t + 5], w[t - 1]);
+    y[t] = x[t] + a;
+    y[t + 5] = x[t] - a;
+  }
+}
+
+METAL_FUNC void radix11(thread float2* x, thread float2* y) {
+  // Raders Algorithm
+  float2 inv = {1 / 10.0, -1 / 10.0};
+
+  // fft
+  radix10<true>(x + 1, y + 1);
+
+  y[0] = y[1] + x[0];
+
+  // b_q
+  y[1] = complex_mul_conj(y[1], float2(-1, 0));
+  y[2] = complex_mul_conj(y[2], float2(0.955301878, -3.17606649));
+  y[3] = complex_mul_conj(y[3], float2(2.63610556, 2.01269656));
+  y[4] = complex_mul_conj(y[4], float2(2.54127802, 2.13117479));
+  y[5] = complex_mul_conj(y[5], float2(2.07016210, 2.59122150));
+  y[6] = complex_mul_conj(y[6], float2(0, -3.31662479));
+  y[7] = complex_mul_conj(y[7], float2(2.07016210, -2.59122150));
+  y[8] = complex_mul_conj(y[8], float2(-2.54127802, 2.13117479));
+  y[9] = complex_mul_conj(y[9], float2(2.63610556, -2.01269656));
+  y[10] = complex_mul_conj(y[10], float2(-0.955301878, -3.17606649));
+
+  // ifft
+  radix10<false>(y + 1, x + 1);
+
+  y[1] = x[1] * inv + x[0];
+  y[6] = x[2] * inv + x[0];
+  y[3] = x[3] * inv + x[0];
+  y[7] = x[4] * inv + x[0];
+  y[9] = x[5] * inv + x[0];
+  y[10] = x[6] * inv + x[0];
+  y[5] = x[7] * inv + x[0];
+  y[8] = x[8] * inv + x[0];
+  y[4] = x[9] * inv + x[0];
+  y[2] = x[10] * inv + x[0];
+}
+
+template <bool raders_perm>
+METAL_FUNC void radix12(thread float2* x, thread float2* y) {
+  float2 w[6];
+  float sin_pi_3 = 0.8660254037844387;
+  w[0] = {sin_pi_3, -0.5};
+  w[1] = {0.5, -sin_pi_3};
+  w[2] = {0, -1};
+  w[3] = {-0.5, -sin_pi_3};
+  w[4] = {-sin_pi_3, -0.5};
+
+  if (raders_perm) {
+    float2 temp[12] = {
+        x[0],
+        x[3],
+        x[2],
+        x[11],
+        x[8],
+        x[9],
+        x[1],
+        x[7],
+        x[5],
+        x[10],
+        x[4],
+        x[6]};
+    radix6(temp, x);
+    radix6(temp + 6, x + 6);
+  } else {
+    float2 temp[12] = {
+        x[0],
+        x[2],
+        x[4],
+        x[6],
+        x[8],
+        x[10],
+        x[1],
+        x[3],
+        x[5],
+        x[7],
+        x[9],
+        x[11]};
+    radix6(temp, x);
+    radix6(temp + 6, x + 6);
+  }
+
+  y[0] = x[0] + x[6];
+  y[6] = x[0] - x[6];
+  for (int t = 1; t < 6; t++) {
+    float2 a = complex_mul(x[t + 6], w[t - 1]);
+    y[t] = x[t] + a;
+    y[t + 6] = x[t] - a;
+  }
+}
+
+METAL_FUNC void radix13(thread float2* x, thread float2* y) {
+  // Raders Algorithm
+  float2 inv = {1 / 12.0, -1 / 12.0};
+
+  // fft
+  radix12<true>(x + 1, y + 1);
+
+  y[0] = y[1] + x[0];
+
+  // b_q
+  y[1] = complex_mul_conj(y[1], float2(-1, 0));
+  y[2] = complex_mul_conj(y[2], float2(3.07497206, -1.88269669));
+  y[3] = complex_mul_conj(y[3], float2(3.09912468, 1.84266823));
+  y[4] = complex_mul_conj(y[4], float2(3.45084438, -1.04483161));
+  y[5] = complex_mul_conj(y[5], float2(0.91083583, 3.48860690));
+  y[6] = complex_mul_conj(y[6], float2(-3.60286363, 0.139189267));
+  y[7] = complex_mul_conj(y[7], float2(3.60555128, 0));
+  y[8] = complex_mul_conj(y[8], float2(3.60286363, 0.139189267));
+  y[9] = complex_mul_conj(y[9], float2(0.91083583, -3.48860690));
+  y[10] = complex_mul_conj(y[10], float2(-3.45084438, -1.04483161));
+  y[11] = complex_mul_conj(y[11], float2(3.09912468, -1.84266823));
+  y[12] = complex_mul_conj(y[12], float2(-3.07497206, -1.88269669));
+
+  // ifft
+  radix12<false>(y + 1, x + 1);
+
+  y[1] = x[1] * inv + x[0];
+  y[7] = x[2] * inv + x[0];
+  y[10] = x[3] * inv + x[0];
+  y[5] = x[4] * inv + x[0];
+  y[9] = x[5] * inv + x[0];
+  y[11] = x[6] * inv + x[0];
+  y[12] = x[7] * inv + x[0];
+  y[6] = x[8] * inv + x[0];
+  y[3] = x[9] * inv + x[0];
+  y[8] = x[10] * inv + x[0];
+  y[4] = x[11] * inv + x[0];
+  y[2] = x[12] * inv + x[0];
+}
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fft/readwrite.h b/macos/third_party/mlx/mlx/backend/metal/kernels/fft/readwrite.h
new file mode 100644
index 0000000..0dc6299
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fft/readwrite.h
@@ -0,0 +1,624 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_common>
+
+#include "mlx/backend/metal/kernels/fft/radix.h"
+
+/* FFT helpers for reading and writing from/to device memory.
+
+For many sizes, GPU FFTs are memory bandwidth bound so
+read/write performance is important.
+
+Where possible, we read 128 bits sequentially in each thread,
+coalesced with accesses from adjacent threads for optimal performance.
+
+We implement specialized reading/writing for:
+  - FFT
+  - RFFT
+  - IRFFT
+
+Each with support for:
+  - Contiguous reads
+  - Padded reads
+  - Strided reads
+*/
+
+#define MAX_RADIX 13
+
+using namespace metal;
+
+template <
+    typename in_T,
+    typename out_T,
+    int step = 0,
+    bool four_step_real = false>
+struct ReadWriter {
+  const device in_T* in;
+  threadgroup float2* buf;
+  device out_T* out;
+  int n;
+  int batch_size;
+  int elems_per_thread;
+  uint3 elem;
+  uint3 grid;
+  int threads_per_tg;
+  bool inv;
+
+  // Used for strided access
+  int strided_device_idx = 0;
+  int strided_shared_idx = 0;
+
+  METAL_FUNC ReadWriter(
+      const device in_T* in_,
+      threadgroup float2* buf_,
+      device out_T* out_,
+      const short n_,
+      const int batch_size_,
+      const short elems_per_thread_,
+      const uint3 elem_,
+      const uint3 grid_,
+      const bool inv_)
+      : in(in_),
+        buf(buf_),
+        out(out_),
+        n(n_),
+        batch_size(batch_size_),
+        elems_per_thread(elems_per_thread_),
+        elem(elem_),
+        grid(grid_),
+        inv(inv_) {
+    // Account for padding on last threadgroup
+    threads_per_tg = elem.x == grid.x - 1
+        ? (batch_size - (grid.x - 1) * grid.y) * grid.z
+        : grid.y * grid.z;
+  }
+
+  // ifft(x) = 1/n * conj(fft(conj(x)))
+  METAL_FUNC float2 post_in(float2 elem) const {
+    return inv ? float2(elem.x, -elem.y) : elem;
+  }
+
+  // Handle float case for generic RFFT alg
+  METAL_FUNC float2 post_in(float elem) const {
+    return float2(elem, 0);
+  }
+
+  METAL_FUNC float2 pre_out(float2 elem) const {
+    return inv ? float2(elem.x / n, -elem.y / n) : elem;
+  }
+
+  METAL_FUNC float2 pre_out(float2 elem, int length) const {
+    return inv ? float2(elem.x / length, -elem.y / length) : elem;
+  }
+
+  METAL_FUNC bool out_of_bounds() const {
+    // Account for possible extra threadgroups
+    int grid_index = elem.x * grid.y + elem.y;
+    return grid_index >= batch_size;
+  }
+
+  METAL_FUNC void load() const {
+    size_t batch_idx = size_t(elem.x * grid.y) * n;
+    short tg_idx = elem.y * grid.z + elem.z;
+    short max_index = grid.y * n - 2;
+
+    // 2 complex64s = 128 bits
+    constexpr int read_width = 2;
+    for (short e = 0; e < (elems_per_thread / read_width); e++) {
+      short index = read_width * tg_idx + read_width * threads_per_tg * e;
+      index = metal::min(index, max_index);
+      // vectorized reads
+      buf[index] = post_in(in[batch_idx + index]);
+      buf[index + 1] = post_in(in[batch_idx + index + 1]);
+    }
+    max_index += 1;
+    if (elems_per_thread % 2 != 0) {
+      short index = tg_idx +
+          read_width * threads_per_tg * (elems_per_thread / read_width);
+      index = metal::min(index, max_index);
+      buf[index] = post_in(in[batch_idx + index]);
+    }
+  }
+
+  METAL_FUNC void write() const {
+    size_t batch_idx = size_t(elem.x * grid.y) * n;
+    short tg_idx = elem.y * grid.z + elem.z;
+    short max_index = grid.y * n - 2;
+
+    constexpr int read_width = 2;
+    for (short e = 0; e < (elems_per_thread / read_width); e++) {
+      short index = read_width * tg_idx + read_width * threads_per_tg * e;
+      index = metal::min(index, max_index);
+      // vectorized reads
+      out[batch_idx + index] = pre_out(buf[index]);
+      out[batch_idx + index + 1] = pre_out(buf[index + 1]);
+    }
+    max_index += 1;
+    if (elems_per_thread % 2 != 0) {
+      short index = tg_idx +
+          read_width * threads_per_tg * (elems_per_thread / read_width);
+      index = metal::min(index, max_index);
+      out[batch_idx + index] = pre_out(buf[index]);
+    }
+  }
+
+  // Padded IO for Bluestein's algorithm
+  METAL_FUNC void load_padded(int length, const device float2* w_k) const {
+    size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length;
+    int fft_idx = elem.z;
+    int m = grid.z;
+
+    threadgroup float2* seq_buf = buf + elem.y * n;
+    for (int e = 0; e < elems_per_thread; e++) {
+      int index = metal::min(fft_idx + e * m, n - 1);
+      if (index < length) {
+        float2 elem = post_in(in[batch_idx + index]);
+        seq_buf[index] = complex_mul(elem, w_k[index]);
+      } else {
+        seq_buf[index] = 0.0;
+      }
+    }
+  }
+
+  METAL_FUNC void write_padded(int length, const device float2* w_k) const {
+    size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length;
+    int fft_idx = elem.z;
+    int m = grid.z;
+    float2 inv_factor = {1.0f / n, -1.0f / n};
+
+    threadgroup float2* seq_buf = buf + elem.y * n;
+    for (int e = 0; e < elems_per_thread; e++) {
+      int index = metal::min(fft_idx + e * m, n - 1);
+      if (index < length) {
+        float2 elem = seq_buf[index + length - 1] * inv_factor;
+        out[batch_idx + index] = pre_out(complex_mul(elem, w_k[index]), length);
+      }
+    }
+  }
+
+  // Strided IO for four step FFT
+  METAL_FUNC void compute_strided_indices(int stride, int overall_n) {
+    // Use the batch threadgroup dimension to coalesce memory accesses:
+    // e.g. stride = 12
+    // device      | shared mem
+    // 0  1  2  3  |  0 12 - -
+    // -  -  -  -  |  1 13 - -
+    // -  -  -  -  |  2 14 - -
+    // 12 13 14 15 |  3 15 - -
+    int coalesce_width = grid.y;
+    int tg_idx = elem.y * grid.z + elem.z;
+    int outer_batch_size = stride / coalesce_width;
+
+    int strided_batch_idx = (elem.x % outer_batch_size) * coalesce_width +
+        overall_n * (elem.x / outer_batch_size);
+    strided_device_idx = strided_batch_idx +
+        tg_idx / coalesce_width * elems_per_thread * stride +
+        tg_idx % coalesce_width;
+    strided_shared_idx = (tg_idx % coalesce_width) * n +
+        tg_idx / coalesce_width * elems_per_thread;
+  }
+
+  // Four Step FFT First Step
+  METAL_FUNC void load_strided(int stride, int overall_n) {
+    compute_strided_indices(stride, overall_n);
+    for (int e = 0; e < elems_per_thread; e++) {
+      buf[strided_shared_idx + e] =
+          post_in(in[strided_device_idx + e * stride]);
+    }
+  }
+
+  METAL_FUNC void write_strided(int stride, int overall_n) {
+    for (int e = 0; e < elems_per_thread; e++) {
+      float2 output = buf[strided_shared_idx + e];
+      int combined_idx = (strided_device_idx + e * stride) % overall_n;
+      int ij = (combined_idx / stride) * (combined_idx % stride);
+      // Apply four step twiddles at end of first step
+      float2 twiddle = get_twiddle(ij, overall_n);
+      out[strided_device_idx + e * stride] = complex_mul(output, twiddle);
+    }
+  }
+};
+
+// Four Step FFT Second Step
+template <>
+METAL_FUNC void ReadWriter<float2, float2, /*step=*/1>::load_strided(
+    int stride,
+    int overall_n) {
+  // Silence compiler warnings
+  (void)stride;
+  (void)overall_n;
+  // Don't invert between steps
+  bool default_inv = inv;
+  inv = false;
+  load();
+  inv = default_inv;
+}
+
+template <>
+METAL_FUNC void ReadWriter<float2, float2, /*step=*/1>::write_strided(
+    int stride,
+    int overall_n) {
+  compute_strided_indices(stride, overall_n);
+  for (int e = 0; e < elems_per_thread; e++) {
+    float2 output = buf[strided_shared_idx + e];
+    out[strided_device_idx + e * stride] = pre_out(output, overall_n);
+  }
+}
+
+// For RFFT, we interleave batches of two real sequences into one complex one:
+//
+// z_k = x_k + j.y_k
+// X_k = (Z_k + Z_(N-k)*) / 2
+// Y_k = -j * ((Z_k - Z_(N-k)*) / 2)
+//
+// This roughly doubles the throughput over the regular FFT.
+template <>
+METAL_FUNC bool ReadWriter<float, float2>::out_of_bounds() const {
+  int grid_index = elem.x * grid.y + elem.y;
+  // We pack two sequences into one for RFFTs
+  return grid_index * 2 >= batch_size;
+}
+
+template <>
+METAL_FUNC void ReadWriter<float, float2>::load() const {
+  size_t batch_idx = size_t(elem.x * grid.y) * n * 2 + elem.y * n * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n;
+
+  // No out of bounds accesses on odd batch sizes
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_in =
+      batch_size % 2 == 1 && grid_index * 2 == batch_size - 1 ? 0 : n;
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  for (int e = 0; e < elems_per_thread; e++) {
+    int index = metal::min(fft_idx + e * m, n - 1);
+    seq_buf[index].x = in[batch_idx + index];
+    seq_buf[index].y = in[batch_idx + index + next_in];
+  }
+}
+
+template <>
+METAL_FUNC void ReadWriter<float, float2>::write() const {
+  short n_over_2 = (n / 2) + 1;
+
+  size_t batch_idx =
+      size_t(elem.x * grid.y) * n_over_2 * 2 + elem.y * n_over_2 * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n;
+
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_out =
+      batch_size % 2 == 1 && grid_index * 2 == batch_size - 1 ? 0 : n_over_2;
+
+  float2 conj = {1, -1};
+  float2 minus_j = {0, -1};
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  for (int e = 0; e < elems_per_thread / 2 + 1; e++) {
+    int index = metal::min(fft_idx + e * m, n_over_2 - 1);
+    // x_0 = z_0.real
+    // y_0 = z_0.imag
+    if (index == 0) {
+      out[batch_idx + index] = {seq_buf[index].x, 0};
+      out[batch_idx + index + next_out] = {seq_buf[index].y, 0};
+    } else {
+      float2 x_k = seq_buf[index];
+      float2 x_n_minus_k = seq_buf[n - index] * conj;
+      out[batch_idx + index] = (x_k + x_n_minus_k) / 2;
+      out[batch_idx + index + next_out] =
+          complex_mul(((x_k - x_n_minus_k) / 2), minus_j);
+    }
+  }
+}
+
+template <>
+METAL_FUNC void ReadWriter<float, float2>::load_padded(
+    int length,
+    const device float2* w_k) const {
+  size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n;
+
+  // No out of bounds accesses on odd batch sizes
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_in =
+      batch_size % 2 == 1 && grid_index * 2 == batch_size - 1 ? 0 : length;
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  for (int e = 0; e < elems_per_thread; e++) {
+    int index = metal::min(fft_idx + e * m, n - 1);
+    if (index < length) {
+      float2 elem =
+          float2(in[batch_idx + index], in[batch_idx + index + next_in]);
+      seq_buf[index] = complex_mul(elem, w_k[index]);
+    } else {
+      seq_buf[index] = 0;
+    }
+  }
+}
+
+template <>
+METAL_FUNC void ReadWriter<float, float2>::write_padded(
+    int length,
+    const device float2* w_k) const {
+  int length_over_2 = (length / 2) + 1;
+  size_t batch_idx =
+      size_t(elem.x * grid.y) * length_over_2 * 2 + elem.y * length_over_2 * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n + length - 1;
+
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_out = batch_size % 2 == 1 && grid_index * 2 == batch_size - 1
+      ? 0
+      : length_over_2;
+
+  float2 conj = {1, -1};
+  float2 inv_factor = {1.0f / n, -1.0f / n};
+  float2 minus_j = {0, -1};
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  for (int e = 0; e < elems_per_thread / 2 + 1; e++) {
+    int index = metal::min(fft_idx + e * m, length_over_2 - 1);
+    // x_0 = z_0.real
+    // y_0 = z_0.imag
+    if (index == 0) {
+      float2 elem = complex_mul(w_k[index], seq_buf[index] * inv_factor);
+      out[batch_idx + index] = float2(elem.x, 0);
+      out[batch_idx + index + next_out] = float2(elem.y, 0);
+    } else {
+      float2 x_k = complex_mul(w_k[index], seq_buf[index] * inv_factor);
+      float2 x_n_minus_k = complex_mul(
+          w_k[length - index], seq_buf[length - index] * inv_factor);
+      x_n_minus_k *= conj;
+      // w_k should happen before this extraction
+      out[batch_idx + index] = (x_k + x_n_minus_k) / 2;
+      out[batch_idx + index + next_out] =
+          complex_mul(((x_k - x_n_minus_k) / 2), minus_j);
+    }
+  }
+}
+
+// For IRFFT, we do the opposite
+//
+// Z_k = X_k + j.Y_k
+// x_k = Re(Z_k)
+// Y_k = Imag(Z_k)
+template <>
+METAL_FUNC bool ReadWriter<float2, float>::out_of_bounds() const {
+  int grid_index = elem.x * grid.y + elem.y;
+  // We pack two sequences into one for IRFFTs
+  return grid_index * 2 >= batch_size;
+}
+
+template <>
+METAL_FUNC void ReadWriter<float2, float>::load() const {
+  short n_over_2 = (n / 2) + 1;
+  size_t batch_idx =
+      size_t(elem.x * grid.y) * n_over_2 * 2 + elem.y * n_over_2 * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n;
+
+  // No out of bounds accesses on odd batch sizes
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_in =
+      batch_size % 2 == 1 && grid_index * 2 == batch_size - 1 ? 0 : n_over_2;
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  float2 conj = {1, -1};
+  float2 plus_j = {0, 1};
+
+  for (int t = 0; t < elems_per_thread / 2 + 1; t++) {
+    int index = metal::min(fft_idx + t * m, n_over_2 - 1);
+    float2 x = in[batch_idx + index];
+    float2 y = in[batch_idx + index + next_in];
+    // NumPy forces first input to be real
+    bool first_val = index == 0;
+    // NumPy forces last input on even irffts to be real
+    bool last_val = n % 2 == 0 && index == n_over_2 - 1;
+    if (first_val || last_val) {
+      x = float2(x.x, 0);
+      y = float2(y.x, 0);
+    }
+    seq_buf[index] = x + complex_mul(y, plus_j);
+    seq_buf[index].y = -seq_buf[index].y;
+    if (index > 0 && !last_val) {
+      seq_buf[n - index] = (x * conj) + complex_mul(y * conj, plus_j);
+      seq_buf[n - index].y = -seq_buf[n - index].y;
+    }
+  }
+}
+
+template <>
+METAL_FUNC void ReadWriter<float2, float>::write() const {
+  int batch_idx = elem.x * grid.y * n * 2 + elem.y * n * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n;
+
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_out =
+      batch_size % 2 == 1 && grid_index * 2 == batch_size - 1 ? 0 : n;
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  for (int e = 0; e < elems_per_thread; e++) {
+    int index = metal::min(fft_idx + e * m, n - 1);
+    out[batch_idx + index] = seq_buf[index].x / n;
+    out[batch_idx + index + next_out] = seq_buf[index].y / -n;
+  }
+}
+
+template <>
+METAL_FUNC void ReadWriter<float2, float>::load_padded(
+    int length,
+    const device float2* w_k) const {
+  int n_over_2 = (n / 2) + 1;
+  int length_over_2 = (length / 2) + 1;
+
+  size_t batch_idx =
+      size_t(elem.x * grid.y) * length_over_2 * 2 + elem.y * length_over_2 * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n;
+
+  // No out of bounds accesses on odd batch sizes
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_in = batch_size % 2 == 1 && grid_index * 2 == batch_size - 1
+      ? 0
+      : length_over_2;
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  float2 conj = {1, -1};
+  float2 plus_j = {0, 1};
+
+  for (int t = 0; t < elems_per_thread / 2 + 1; t++) {
+    int index = metal::min(fft_idx + t * m, n_over_2 - 1);
+    float2 x = in[batch_idx + index];
+    float2 y = in[batch_idx + index + next_in];
+    if (index < length_over_2) {
+      bool last_val = length % 2 == 0 && index == length_over_2 - 1;
+      if (last_val) {
+        x = float2(x.x, 0);
+        y = float2(y.x, 0);
+      }
+      float2 elem1 = x + complex_mul(y, plus_j);
+      seq_buf[index] = complex_mul(elem1 * conj, w_k[index]);
+      if (index > 0 && !last_val) {
+        float2 elem2 = (x * conj) + complex_mul(y * conj, plus_j);
+        seq_buf[length - index] =
+            complex_mul(elem2 * conj, w_k[length - index]);
+      }
+    } else {
+      short pad_index = metal::min(length + (index - length_over_2) * 2, n - 2);
+      seq_buf[pad_index] = 0;
+      seq_buf[pad_index + 1] = 0;
+    }
+  }
+}
+
+template <>
+METAL_FUNC void ReadWriter<float2, float>::write_padded(
+    int length,
+    const device float2* w_k) const {
+  size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2;
+  threadgroup float2* seq_buf = buf + elem.y * n + length - 1;
+
+  int grid_index = elem.x * grid.y + elem.y;
+  short next_out =
+      batch_size % 2 == 1 && grid_index * 2 == batch_size - 1 ? 0 : length;
+
+  short m = grid.z;
+  short fft_idx = elem.z;
+
+  float2 inv_factor = {1.0f / n, -1.0f / n};
+  for (int e = 0; e < elems_per_thread; e++) {
+    int index = fft_idx + e * m;
+    if (index < length) {
+      float2 output = complex_mul(seq_buf[index] * inv_factor, w_k[index]);
+      out[batch_idx + index] = output.x / length;
+      out[batch_idx + index + next_out] = output.y / -length;
+    }
+  }
+}
+
+// Four Step RFFT
+template <>
+METAL_FUNC void
+ReadWriter<float2, float2, /*step=*/1, /*real=*/true>::load_strided(
+    int stride,
+    int overall_n) {
+  // Silence compiler warnings
+  (void)stride;
+  (void)overall_n;
+  // Don't invert between steps
+  bool default_inv = inv;
+  inv = false;
+  load();
+  inv = default_inv;
+}
+
+template <>
+METAL_FUNC void
+ReadWriter<float2, float2, /*step=*/1, /*real=*/true>::write_strided(
+    int stride,
+    int overall_n) {
+  int overall_n_over_2 = overall_n / 2 + 1;
+  int coalesce_width = grid.y;
+  int tg_idx = elem.y * grid.z + elem.z;
+  int outer_batch_size = stride / coalesce_width;
+
+  int strided_batch_idx = (elem.x % outer_batch_size) * coalesce_width +
+      overall_n_over_2 * (elem.x / outer_batch_size);
+  strided_device_idx = strided_batch_idx +
+      tg_idx / coalesce_width * elems_per_thread / 2 * stride +
+      tg_idx % coalesce_width;
+  strided_shared_idx = (tg_idx % coalesce_width) * n +
+      tg_idx / coalesce_width * elems_per_thread / 2;
+  for (int e = 0; e < elems_per_thread / 2; e++) {
+    float2 output = buf[strided_shared_idx + e];
+    out[strided_device_idx + e * stride] = output;
+  }
+
+  // Add on n/2 + 1 element
+  if (tg_idx == 0 && elem.x % outer_batch_size == 0) {
+    out[strided_batch_idx + overall_n / 2] = buf[n / 2];
+  }
+}
+
+// Four Step IRFFT
+template <>
+METAL_FUNC void
+ReadWriter<float2, float2, /*step=*/0, /*real=*/true>::load_strided(
+    int stride,
+    int overall_n) {
+  int overall_n_over_2 = overall_n / 2 + 1;
+  auto conj = float2(1, -1);
+
+  compute_strided_indices(stride, overall_n);
+  // Translate indices in terms of N - k
+  for (int e = 0; e < elems_per_thread; e++) {
+    int device_idx = strided_device_idx + e * stride;
+    int overall_batch = device_idx / overall_n;
+    int overall_index = device_idx % overall_n;
+    if (overall_index < overall_n_over_2) {
+      device_idx -= overall_batch * (overall_n - overall_n_over_2);
+      buf[strided_shared_idx + e] = in[device_idx] * conj;
+    } else {
+      int conj_idx = overall_n - overall_index;
+      device_idx = overall_batch * overall_n_over_2 + conj_idx;
+      buf[strided_shared_idx + e] = in[device_idx];
+    }
+  }
+}
+
+template <>
+METAL_FUNC void
+ReadWriter<float2, float, /*step=*/1, /*real=*/true>::load_strided(
+    int stride,
+    int overall_n) {
+  // Silence compiler warnings
+  (void)stride;
+  (void)overall_n;
+  bool default_inv = inv;
+  inv = false;
+  load();
+  inv = default_inv;
+}
+
+template <>
+METAL_FUNC void
+ReadWriter<float2, float, /*step=*/1, /*real=*/true>::write_strided(
+    int stride,
+    int overall_n) {
+  compute_strided_indices(stride, overall_n);
+
+  for (int e = 0; e < elems_per_thread; e++) {
+    out[strided_device_idx + e * stride] =
+        pre_out(buf[strided_shared_idx + e], overall_n).x;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fp4.h b/macos/third_party/mlx/mlx/backend/metal/kernels/fp4.h
new file mode 100644
index 0000000..25642f2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fp4.h
@@ -0,0 +1,48 @@
+#pragma once
+
+struct fp4_e2m1 {
+  fp4_e2m1(float x) {
+    if (metal::isnan(x)) {
+      bits = 0x7;
+      return;
+    }
+
+    const uint8_t sign_bit = (metal::signbit(x)) ? 0x8 : 0x0;
+    x = metal::abs(x);
+
+    if (x > 5.0f) {
+      bits = 0x7;
+    } else if (x >= 3.5f) {
+      bits = 0x6;
+    } else if (x > 2.5f) {
+      bits = 0x5;
+    } else if (x >= 1.75f) {
+      bits = 0x4;
+    } else if (x > 1.25f) {
+      bits = 0x3;
+    } else if (x >= 0.75f) {
+      bits = 0x2;
+    } else if (x > 0.25f) {
+      bits = 0x1;
+    } else {
+      bits = 0x0;
+    }
+    bits |= sign_bit;
+  }
+
+  operator float16_t() {
+    half converted = as_type<half>(ushort((bits & 7) << 9));
+    converted *= 16384.0;
+    return bits & 8 ? -converted : converted;
+  }
+
+  operator float() {
+    return static_cast<float>(this->operator float16_t());
+  }
+
+  operator bfloat16_t() {
+    return static_cast<bfloat16_t>(this->operator float16_t());
+  }
+
+  uint8_t bits;
+};
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fp8.h b/macos/third_party/mlx/mlx/backend/metal/kernels/fp8.h
new file mode 100644
index 0000000..60d34be
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fp8.h
@@ -0,0 +1,80 @@
+#pragma once
+
+struct fp8_e4m3 {
+  template <typename T>
+  fp8_e4m3(T f) {
+    // From PyTorch
+    // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L148
+    uint32_t fp8_max = 543 << 21;
+    uint32_t denorm_mask = 141 << 23;
+    uint32_t f_bits = as_type<uint32_t>(static_cast<float>(f));
+    uint32_t sign = f_bits & 0x80000000;
+    f_bits ^= sign;
+    if (f_bits >= fp8_max) {
+      // Default behavior saturates to min/max
+      bits = 0x7E;
+    } else {
+      if (f_bits < (121 << 23)) {
+        f_bits = as_type<uint32_t>(
+            as_type<float>(f_bits) + as_type<float>(denorm_mask));
+        bits = static_cast<uint8_t>(f_bits - denorm_mask);
+      } else {
+        // resulting mantissa is odd
+        uint8_t mant_odd = (f_bits >> 20) & 1;
+        f_bits += ((uint32_t)(7 - 127) << 23) + 0x7FFFF;
+        f_bits += mant_odd;
+        bits = static_cast<uint8_t>(f_bits >> 20);
+      }
+    }
+    bits |= static_cast<uint8_t>(sign >> 24);
+  }
+
+  operator float16_t() {
+    uint16_t v = (bits & 127) << 7;
+    half converted = as_type<half>(v);
+    converted *= 256.0;
+    auto sign = bits & 128;
+    return (sign ? -converted : converted);
+  }
+
+  operator bfloat16_t() {
+    return static_cast<bfloat16_t>(this->operator float16_t());
+  }
+
+  operator float() {
+    return static_cast<float>(this->operator float16_t());
+  }
+
+  uint8_t bits;
+};
+
+struct fp8_e8m0 {
+  fp8_e8m0(float x) {
+    if (!metal::isfinite(x)) {
+      bits = 0xFF;
+      return;
+    }
+    if (x < 0.0f) {
+      bits = 0x00;
+      return;
+    }
+    float le = metal::log2(x);
+    int n = int(metal::round(le));
+
+    n = n < -127 ? -127 : n;
+    n = n > 127 ? 127 : n;
+    bits = static_cast<uint8_t>(n + 127);
+  }
+
+  operator bfloat16_t() {
+    uint16_t out = (bits == 0 ? 0x40 : (static_cast<uint16_t>(bits) << 7));
+    return as_type<bfloat16_t>(out);
+  }
+
+  operator float() {
+    uint32_t out = (bits == 0 ? 0x400000 : (static_cast<uint16_t>(bits) << 23));
+    return as_type<float>(out);
+  }
+
+  uint8_t bits;
+};
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized.h b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized.h
new file mode 100644
index 0000000..5c5e4b2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized.h
@@ -0,0 +1,1850 @@
+// Copyright © 2025 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/fp4.h"
+#include "mlx/backend/metal/kernels/fp8.h"
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+using namespace metal;
+
+#define MLX_MTL_CONST static constant constexpr const
+
+MLX_MTL_CONST int SIMD_SIZE = 32;
+MLX_MTL_CONST int QUAD_SIZE = 4;
+
+template <int wsize = 8, int bits = 4>
+inline constexpr short get_pack_factor() {
+  return wsize / bits;
+}
+
+template <int wsize = 8>
+inline constexpr short get_bytes_per_pack() {
+  return wsize / 8;
+}
+
+template <typename T, int group_size>
+static inline T dequantize_scale(uint8_t s) {
+  if constexpr (group_size == 16) {
+    // Use nv scale
+    return T(*(thread fp8_e4m3*)(&s));
+  } else {
+    return T(*(thread fp8_e8m0*)(&s));
+  }
+}
+
+template <int bits>
+struct Quantize {
+  uint8_t operator()(float x) {
+    if (bits == 8) {
+      return fp8_e4m3(x).bits;
+    } else {
+      return fp4_e2m1(x).bits;
+    }
+  }
+};
+
+template <int bits, typename U = float>
+struct Dequantize {
+  U operator()(uint8_t x) {
+    if constexpr (bits == 8) {
+      return U(*(thread fp8_e4m3*)(&x));
+    } else {
+      return U(*(thread fp4_e2m1*)(&x));
+    }
+  }
+};
+
+template <typename T, typename U, int values_per_thread>
+inline void load_vector(const device T* x, thread U* x_thread) {
+#pragma unroll
+  for (int i = 0; i < values_per_thread; i++) {
+    x_thread[i] = x[i];
+  }
+}
+
+template <typename T, typename U, int values_per_thread>
+inline void load_vector_safe(const device T* x, thread U* x_thread, int N) {
+  for (int i = 0; i < N; i++) {
+    x_thread[i] = x[i];
+  }
+
+  for (int i = N; i < values_per_thread; i++) {
+    x_thread[i] = 0;
+  }
+}
+
+template <typename U, int values_per_thread, int bits>
+inline U qdot(const device uint8_t* w, const thread U* x_thread, U scale) {
+  U accum = 0;
+  if constexpr (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * Dequantize<4>{}(ws[i]) +
+           x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) +
+           x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) +
+           x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12));
+    }
+  } else {
+    for (int i = 0; i < values_per_thread; i++) {
+      accum += x_thread[i] * Dequantize<8>{}(w[i]);
+    }
+  }
+
+  return scale * accum;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline U
+qdot_safe(const device uint8_t* w, const thread U* x_thread, U scale, int N) {
+  U accum = 0;
+
+  if constexpr (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (N / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * Dequantize<4>{}(ws[i]) +
+           x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) +
+           x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) +
+           x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12));
+    }
+  } else {
+    for (int i = 0; i < N; i++) {
+      accum += x_thread[i] * Dequantize<8>{}(w[i]);
+    }
+  }
+  return scale * accum;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) {
+  if constexpr (bits == 4) {
+    for (int i = 0; i < (values_per_thread / 2); i++) {
+      result[2 * i] += x * scale * Dequantize<4>{}(w[i]);
+      result[2 * i + 1] += x * scale * Dequantize<4>{}(w[i] >> 4);
+    }
+  } else {
+    for (int i = 0; i < values_per_thread; i++) {
+      result[i] += x * scale * Dequantize<8>{}(w[i]);
+    }
+  }
+}
+
+template <typename U, int bits>
+inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) {
+  if constexpr (bits == 4) {
+    w_local[0] = scale * Dequantize<4, U>{}(w);
+    w_local[1] = scale * Dequantize<4, U>{}(w >> 4);
+  } else {
+    w_local[0] = scale * Dequantize<8, U>{}(w);
+  }
+}
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short group_size,
+    short bits>
+struct QuantizedBlockLoader {
+  MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>();
+  MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack();
+  MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
+  MLX_MTL_CONST short n_reads =
+      (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
+  MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS;
+  MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1;
+
+  static_assert(
+      (n_reads * pack_factor) <= group_size,
+      "The number of reads per thread must be less than the group size.");
+
+  const int src_ld;
+  const int tile_stride;
+  short group_step_cnt;
+  const int group_stride;
+
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  threadgroup T* dst;
+  const device uint8_t* src;
+  const device uint8_t* scales;
+
+  QuantizedBlockLoader(
+      const device uint8_t* src_,
+      const device uint8_t* scales_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(
+            reduction_dim ? BCOLS_PACKED * bytes_per_pack
+                          : BROWS * src_ld * bytes_per_pack / pack_factor),
+        group_step_cnt(0),
+        group_stride(BROWS * src_ld / group_size),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(n_reads * thread_idx / BCOLS_PACKED),
+        bj((n_reads * thread_idx) % BCOLS_PACKED),
+        dst(dst_ + bi * dst_ld + bj * pack_factor),
+        src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
+            bj * bytes_per_pack),
+        scales(
+            scales_ + bi * src_ld / group_size +
+            (bj * pack_factor) / group_size) {}
+
+  void load_unsafe() const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    T scale = dequantize_scale<T, group_size>(*scales);
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, bits>(
+          src[i * bytes_per_pack], scale, dst + i * pack_factor);
+    }
+  }
+
+  void load_safe(short2 src_tile_dim) const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    T scale = dequantize_scale<T, group_size>(*scales);
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, bits>(
+          src[i * bytes_per_pack], scale, dst + i * pack_factor);
+    }
+  }
+
+  void next() {
+    src += tile_stride;
+    if (reduction_dim == 1) {
+      if (group_steps > 1) {
+        group_step_cnt++;
+        if (group_step_cnt == group_steps) {
+          group_step_cnt = 0;
+          scales++;
+        }
+      } else {
+        scales += scale_step;
+      }
+    } else {
+      scales += group_stride;
+    }
+  }
+};
+
+template <typename T, int group_size, int bits, int D>
+METAL_FUNC void fp_qmv_quad_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    constant int& in_vec_size,
+    const constant int& out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint quad_gid [[quadgroup_index_in_threadgroup]],
+    uint quad_lid [[thread_index_in_quadgroup]]) {
+  constexpr int quads_per_simd = SIMD_SIZE / QUAD_SIZE;
+  constexpr int pack_factor = get_pack_factor<32, bits>();
+  constexpr int values_per_thread = D / QUAD_SIZE;
+  constexpr int steps_per_thread =
+      values_per_thread < group_size ? 1 : values_per_thread / group_size;
+  constexpr int values_per_step = values_per_thread / steps_per_thread;
+  constexpr int packs_per_thread = values_per_thread / pack_factor;
+  constexpr int packs_per_step = values_per_step / pack_factor;
+  constexpr int results_per_quadgroup = 8;
+
+  typedef float U;
+
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_quadgroup] = {0};
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const int out_row = tid.y * quads_per_simd * results_per_quadgroup + quad_gid;
+
+  w += out_row * in_vec_size_w + quad_lid * packs_per_thread;
+  scales +=
+      out_row * in_vec_size_g + (quad_lid * values_per_thread) / group_size;
+  x += tid.x * in_vec_size + quad_lid * values_per_thread;
+  y += tid.x * out_vec_size + out_row;
+
+  load_vector<T, U, values_per_thread>(x, x_thread);
+
+  for (int row = 0; row < results_per_quadgroup; row++) {
+    auto wl = (const device uint8_t*)(w + row * in_vec_size_w * quads_per_simd);
+    const device uint8_t* sl = scales + row * in_vec_size_g * quads_per_simd;
+#pragma unroll
+    for (int k = 0; k < steps_per_thread; ++k) {
+      U s = dequantize_scale<U, group_size>(sl[0]);
+      if (row * quads_per_simd + out_row < out_vec_size) {
+        result[row] += qdot<U, values_per_step, bits>(
+            wl, x_thread + k * values_per_step, s);
+      }
+      sl++;
+      wl += (sizeof(uint32_t) / sizeof(uint8_t)) * packs_per_step;
+    }
+  }
+
+  for (int row = 0; row < results_per_quadgroup; row++) {
+    result[row] = quad_sum(result[row]);
+    if (quad_lid == 0 && row * quads_per_simd + out_row < out_vec_size) {
+      y[row * quads_per_simd] = static_cast<T>(result[row]);
+    }
+  }
+}
+
+template <typename T, int group_size, int bits>
+METAL_FUNC void fp_qmv_fast_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int packs_per_thread = 2;
+  constexpr int num_simdgroups = 2;
+  constexpr int results_per_simdgroup = 4;
+  constexpr int pack_factor = get_pack_factor<32, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<32>();
+  constexpr int values_per_thread = pack_factor * packs_per_thread;
+  constexpr int block_size = values_per_thread * SIMD_SIZE;
+  constexpr int scale_step_per_thread = group_size / values_per_thread;
+
+  const device uint8_t* ws = (const device uint8_t*)w;
+
+  typedef float U;
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_simdgroup] = {0};
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) +
+      simd_gid * results_per_simdgroup;
+
+  ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack;
+  scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+  x += tid.x * in_vec_size + simd_lid * values_per_thread;
+  y += tid.x * out_vec_size + out_row;
+
+  for (int k = 0; k < in_vec_size; k += block_size) {
+    load_vector<T, U, values_per_thread>(x, x_thread);
+
+    for (int row = 0; row < results_per_simdgroup; row++) {
+      auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+      const device auto* sl = scales + row * in_vec_size_g;
+
+      U s = dequantize_scale<U, group_size>(sl[0]);
+      result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
+    }
+
+    ws += block_size * bytes_per_pack / pack_factor;
+    scales += block_size / group_size;
+    x += block_size;
+  }
+
+  for (int row = 0; row < results_per_simdgroup; row++) {
+    result[row] = simd_sum(result[row]);
+    if (simd_lid == 0) {
+      y[row] = static_cast<T>(result[row]);
+    }
+  }
+}
+
+template <typename T, int group_size, int bits>
+METAL_FUNC void fp_qmv_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int num_simdgroups = 2;
+  constexpr int results_per_simdgroup = 4;
+  constexpr int packs_per_thread = 1;
+  constexpr int pack_factor = get_pack_factor<32, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<32>();
+
+  constexpr int values_per_thread = pack_factor * packs_per_thread;
+  constexpr int block_size = values_per_thread * SIMD_SIZE;
+  constexpr int scale_step_per_thread = group_size / values_per_thread;
+
+  const device uint8_t* ws = (const device uint8_t*)w;
+
+  typedef float U;
+
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_simdgroup] = {0};
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) +
+      simd_gid * results_per_simdgroup;
+  const int used_out_row = min(out_vec_size - results_per_simdgroup, out_row);
+
+  if (out_row >= out_vec_size) {
+    return;
+  }
+
+  // In this case we need to properly guard all our reads because there isn't
+  // even 1 tile in the matrix
+  if (out_vec_size < (num_simdgroups * results_per_simdgroup)) {
+    ws +=
+        out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack;
+    scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    x += tid.x * in_vec_size + simd_lid * values_per_thread;
+    y += tid.x * out_vec_size + out_row;
+
+    int k = 0;
+    for (; k < in_vec_size - block_size; k += block_size) {
+      load_vector<T, U, values_per_thread>(x, x_thread);
+
+      for (int row = 0;
+           row < results_per_simdgroup && out_row + row < out_vec_size;
+           row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device auto* sl = scales + row * in_vec_size_g;
+
+        uint8_t s = sl[0];
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
+      }
+
+      ws += block_size * bytes_per_pack / pack_factor;
+      scales += block_size / group_size;
+      x += block_size;
+    }
+    const int remaining = clamp(
+        static_cast<int>(in_vec_size - k - simd_lid * values_per_thread),
+        0,
+        values_per_thread);
+    if (remaining > 0) {
+      load_vector_safe<T, U, values_per_thread>(x, x_thread, remaining);
+
+      for (int row = 0;
+           row < results_per_simdgroup && out_row + row < out_vec_size;
+           row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device auto* sl = scales + row * in_vec_size_g;
+
+        U s = dequantize_scale<U, group_size>(sl[0]);
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
+      }
+    }
+
+    for (int row = 0;
+         row < results_per_simdgroup && out_row + row < out_vec_size;
+         row++) {
+      result[row] = simd_sum(result[row]);
+      if (simd_lid == 0) {
+        y[row] = static_cast<T>(result[row]);
+      }
+    }
+  }
+
+  // In this case the last tile is moved back to redo some output values
+  else {
+    ws += used_out_row * in_vec_size_w +
+        simd_lid * packs_per_thread * bytes_per_pack;
+    scales += used_out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    x += tid.x * in_vec_size + simd_lid * values_per_thread;
+    y += tid.x * out_vec_size + used_out_row;
+
+    int k = 0;
+    for (; k < in_vec_size - block_size; k += block_size) {
+      load_vector<T, U, values_per_thread>(x, x_thread);
+
+      for (int row = 0; row < results_per_simdgroup; row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device auto* sl = scales + row * in_vec_size_g;
+
+        U s = dequantize_scale<U, group_size>(sl[0]);
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
+      }
+
+      ws += block_size * bytes_per_pack / pack_factor;
+      scales += block_size / group_size;
+      x += block_size;
+    }
+    const int remaining = clamp(
+        static_cast<int>(in_vec_size - k - simd_lid * values_per_thread),
+        0,
+        values_per_thread);
+    if (remaining > 0) {
+      load_vector_safe<T, U, values_per_thread>(x, x_thread, remaining);
+
+      for (int row = 0; row < results_per_simdgroup; row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device auto* sl = scales + row * in_vec_size_g;
+
+        U s = dequantize_scale<U, group_size>(sl[0]);
+        result[row] +=
+            qdot_safe<U, values_per_thread, bits>(wl, x_thread, s, remaining);
+      }
+    }
+    for (int row = 0; row < results_per_simdgroup; row++) {
+      result[row] = simd_sum(result[row]);
+      if (simd_lid == 0) {
+        y[row] = static_cast<T>(result[row]);
+      }
+    }
+  }
+}
+
+template <typename T, const int group_size, int bits>
+METAL_FUNC void fp_qvm_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const int in_vec_size,
+    const int out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int num_simdgroups = 2;
+  constexpr int pack_factor = get_pack_factor<32, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack();
+
+  constexpr int tn = group_size / pack_factor;
+  constexpr int block_size = SIMD_SIZE;
+
+  using W_T = uint32_t;
+  const device W_T* ws = (const device W_T*)w;
+
+  typedef float U;
+  typedef struct {
+    W_T wi[tn * bytes_per_pack];
+  } vec_w;
+
+  thread vec_w w_local;
+  thread U result[tn * pack_factor] = {0};
+  thread U scale = 0;
+  thread U x_local = 0;
+
+  // Adjust positions
+  const int out_vec_size_w = out_vec_size * bytes_per_pack / pack_factor;
+  const int out_vec_size_g = out_vec_size / group_size;
+  // 32 * (tid.y * 2 + simd_gid)
+  int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid);
+  ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w;
+  scales += out_col / group_size + simd_lid * out_vec_size_g;
+  x += tid.x * in_vec_size + simd_lid;
+  y += tid.x * out_vec_size + out_col;
+
+  if (out_col >= out_vec_size) {
+    return;
+  }
+
+  // Loop over in_vec in blocks of block_size
+  int remaining = in_vec_size % block_size;
+  if (remaining == 0) {
+    for (int i = 0; i < in_vec_size; i += block_size) {
+      x_local = *x;
+      scale = dequantize_scale<U, group_size>(*scales);
+      w_local = *((device vec_w*)ws);
+      qouter<U, tn * pack_factor, bits>(
+          (thread uint8_t*)&w_local, x_local, scale, result);
+
+      x += block_size;
+      scales += block_size * out_vec_size_g;
+      ws += block_size * out_vec_size_w;
+    }
+  } else {
+    for (int i = block_size; i < in_vec_size; i += block_size) {
+      x_local = *x;
+      scale = dequantize_scale<U, group_size>(*scales);
+      w_local = *((device vec_w*)ws);
+
+      qouter<U, tn * pack_factor, bits>(
+          (thread uint8_t*)&w_local, x_local, scale, result);
+
+      x += block_size;
+      scales += block_size * out_vec_size_g;
+      ws += block_size * out_vec_size_w;
+    }
+    if (static_cast<int>(simd_lid) < remaining) {
+      x_local = *x;
+      scale = dequantize_scale<U, group_size>(*scales);
+      w_local = *((device vec_w*)ws);
+    } else {
+      x_local = 0;
+      scale = 0;
+    }
+    qouter<U, tn * pack_factor, bits>(
+        (thread uint8_t*)&w_local, x_local, scale, result);
+  }
+
+// Accumulate in the simdgroup
+#pragma clang loop unroll(full)
+  for (int k = 0; k < tn * pack_factor; k++) {
+    result[k] = simd_sum(result[k]);
+  }
+
+  // Store the result
+  if (simd_lid == 0) {
+#pragma clang loop unroll(full)
+    for (int k = 0; k < tn * pack_factor; k++) {
+      y[k] = static_cast<T>(result[k]);
+    }
+  }
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+METAL_FUNC void fp_qmm_t_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    threadgroup T* Xs,
+    threadgroup T* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  (void)lid;
+
+  constexpr int WM = 2;
+  constexpr int WN = 2;
+  constexpr int pack_factor = get_pack_factor<8, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack();
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  // Instantiate the appropriate BlockMMA and Loader
+  using mma_t = mlx::steel::
+      BlockMMA<T, T, BM, BN, BK, WM, WN, false, true, BK_padded, BK_padded>;
+  using loader_x_t =
+      mlx::steel::BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE>;
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      BN,
+      BK,
+      BK_padded,
+      1,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  // Set the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+
+  auto wl = (const device uint8_t*)w;
+
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * K_w;
+  scales += y_col * K_g;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the x loader and mma operation
+  const short num_els = min(BM, M - y_row);
+  const short num_outs = min(BN, N - y_col);
+  loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
+  loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid);
+  mma_t mma_op(simd_gid, simd_lid);
+
+  if (num_els < BM) {
+    if (!aligned_N && num_outs < BN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_safe(short2(BK, num_outs));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  } else {
+    if (!aligned_N && num_outs < BN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_safe(short2(BK, num_outs));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  }
+
+  // Store results to device memory
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (num_els < BM || num_outs < BN) {
+    mma_op.store_result_safe(y, N, short2(num_outs, num_els));
+  } else {
+    mma_op.store_result(y, N);
+  }
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+METAL_FUNC void fp_qmm_n_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    threadgroup T* Xs,
+    threadgroup T* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  (void)lid;
+
+  constexpr int WM = 2;
+  constexpr int WN = 2;
+  constexpr int pack_factor = get_pack_factor<8, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack();
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  // Instantiate the appropriate BlockMMA and Loader
+  using mma_t = mlx::steel::
+      BlockMMA<T, T, BM, BN, BK, WM, WN, false, false, BK_padded, BN_padded>;
+  using loader_x_t = mlx::steel::
+      BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE, 1, 4>;
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      BK,
+      BN,
+      BN_padded,
+      0,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  auto wl = (const device uint8_t*)w;
+
+  // Set the block
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * bytes_per_pack / pack_factor;
+  scales += y_col / group_size;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the x loader and mma operation
+  const short num_els = min(BM, M - y_row);
+  loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
+  loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid);
+  mma_t mma_op(simd_gid, simd_lid);
+
+  if (num_els < BM) {
+    if ((K % BK) != 0) {
+      const int k_blocks = K / BK;
+      for (int k = 0; k < k_blocks; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+      const short num_k = K - k_blocks * BK;
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      loader_x.load_safe(short2(num_k, num_els));
+      loader_w.load_safe(short2(BN, num_k));
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      mma_op.mma(Xs, Ws);
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  } else {
+    if ((K % BK) != 0) {
+      const int k_blocks = K / BK;
+      for (int k = 0; k < k_blocks; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+      const short num_k = K - k_blocks * BK;
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      loader_x.load_safe(short2(num_k, BM));
+      loader_w.load_safe(short2(BN, num_k));
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      mma_op.mma(Xs, Ws);
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  }
+
+  // Store results to device memory
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (num_els < BM) {
+    mma_op.store_result_safe(y, N, short2(BN, num_els));
+  } else {
+    mma_op.store_result(y, N);
+  }
+}
+
+template <typename T>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device uint8_t*& scales,
+    device T*& y,
+    int output_stride,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx = tid.z;
+  uint32_t w_idx = tid.z;
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+  }
+  y += tid.z * output_stride;
+}
+
+template <typename T>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device uint8_t*& scales,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T*& y,
+    int output_stride,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx;
+  uint32_t w_idx;
+  if (batch_ndims == 1) {
+    x_idx = lhs_indices[tid.z * lhs_strides[0]];
+    w_idx = rhs_indices[tid.z * rhs_strides[0]];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        tid.z, batch_shape, lhs_strides, rhs_strides, batch_ndims);
+    x_idx = lhs_indices[idx.x];
+    w_idx = rhs_indices[idx.y];
+  }
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+  }
+  y += tid.z * output_stride;
+}
+
+template <typename T, int group_size, int bits, int D, bool batched>
+[[kernel]] void fp_qmv_quad(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint quad_gid [[quadgroup_index_in_threadgroup]],
+    uint quad_lid [[thread_index_in_quadgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+  fp_qmv_quad_impl<T, group_size, bits, D>(
+      w, scales, x, y, in_vec_size, out_vec_size, tid, quad_gid, quad_lid);
+}
+
+template <typename T, int group_size, int bits, bool batched>
+[[kernel]] void fp_qmv_fast(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+  fp_qmv_fast_impl<T, group_size, bits>(
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+template <typename T, const int group_size, int bits, bool batched>
+[[kernel]] void fp_qmv(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+  fp_qmv_impl<T, group_size, bits>(
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+template <typename T, const int group_size, int bits, bool batched>
+[[kernel]] void fp_qvm(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+  fp_qvm_impl<T, group_size, bits>(
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+template <typename T, const int group_size, int bits, int split_k = 32>
+[[kernel]] void fp_qvm_split_k(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& final_block_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      y,
+      out_vec_size * M,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+
+  // When (in_vec_size % split_k != 0) the final block needs to be smaller
+  int in_vec_size_adj =
+      tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size;
+
+  fp_qvm_impl<T, group_size, bits>(
+      w, scales, x, y, in_vec_size_adj, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const bool batched,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void fp_qmm_t(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BN * BK_padded];
+
+  if (batched) {
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+  fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN>(
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool batched,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void fp_qmm_n(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BK * BN_padded];
+
+  if (batched) {
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+
+  fp_qmm_n_impl<T, group_size, bits, BM, BK, BN>(
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <typename T, int group_size, int bits>
+[[kernel]] void fp_gather_qmv_fast(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      lhs_indices,
+      rhs_indices,
+      y,
+      out_vec_size * M,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+  fp_qmv_fast_impl<T, group_size, bits>(
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+template <typename T, int group_size, int bits>
+[[kernel]] void fp_gather_qmv(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      lhs_indices,
+      rhs_indices,
+      y,
+      out_vec_size * M,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+  fp_qmv_impl<T, group_size, bits>(
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+template <typename T, int group_size, int bits>
+[[kernel]] void fp_gather_qvm(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      lhs_indices,
+      rhs_indices,
+      y,
+      out_vec_size * M,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+  fp_qvm_impl<T, group_size, bits>(
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void fp_gather_qmm_t(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BN * BK_padded];
+
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+  fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN>(
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void fp_gather_qmm_n(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BK * BN_padded];
+
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+  fp_qmm_n_impl<T, group_size, bits, BM, BK, BN>(
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    int group_size,
+    int bits,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose>
+[[kernel]] void fp_gather_qmm_rhs(
+    const device T* x,
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device uint32_t* indices,
+    device T* y,
+    const constant int& M,
+    const constant int& N,
+    const constant int& K,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]]) {
+  constexpr int pack_factor = get_pack_factor<8, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack();
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  using mma_t = mlx::steel::BlockMMA<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      false,
+      transpose,
+      BK_padded,
+      transpose ? BK_padded : BN_padded>;
+  using loader_x_t =
+      mlx::steel::BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE>;
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      transpose ? BN : BK,
+      transpose ? BK : BN,
+      transpose ? BK_padded : BN_padded,
+      transpose,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded];
+
+  // Compute the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int N_w = N * bytes_per_pack / pack_factor;
+  const int N_g = N / group_size;
+  const int K_it = K / BK;
+  const size_t stride_w = transpose ? N * K_w : K * N_w;
+  const size_t stride_s = transpose ? N * K_g : K * N_g;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+  const size_t y_row_long = size_t(y_row);
+  const size_t y_col_long = size_t(y_col);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, M - y_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, N - y_col));
+
+  // Calculate the final tiles in the case that K is not aligned
+  const int k_remain = K - K_it * BK;
+  const short2 tile_x = short2(k_remain, tgp_bm);
+  const short2 tile_w =
+      transpose ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+  // Move x and output to the correct block
+  auto wl = (const device uint8_t*)w;
+  x += y_row_long * K;
+  y += y_row_long * N + y_col_long;
+  wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor;
+  scales += transpose ? y_col_long * K_g : y_col / group_size;
+
+  // Do as many matmuls as necessary
+  uint32_t index;
+  short offset;
+  uint32_t index_next = indices[y_row];
+  short offset_next = 0;
+  int n = 0;
+  while (n < tgp_bm) {
+    n++;
+    offset = offset_next;
+    index = index_next;
+    offset_next = tgp_bm;
+    for (; n < tgp_bm; n++) {
+      if (indices[y_row + n] != index) {
+        offset_next = n;
+        index_next = indices[y_row + n];
+        break;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Prepare threadgroup mma operation
+    thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+    // Prepare threadgroup loading operations
+    thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id);
+    thread loader_w_t loader_w(
+        wl + index * stride_w,
+        scales + index * stride_s,
+        transpose ? K : N,
+        Ws,
+        simd_group_id,
+        simd_lane_id);
+
+    // Matrices are all aligned check nothing
+    if (align_M && align_N) {
+      gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it);
+      if (!align_K) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+      }
+
+      // Store results to device memory
+      if (offset_next - offset == BM) {
+        mma_op.store_result(y, N);
+      } else {
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(BN, offset_next));
+      }
+    } else {
+      // Tile aligned so check outside of the hot loop
+      if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) {
+        gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+
+        // Store results to device memory
+        if (offset_next - offset == BM) {
+          mma_op.store_result(y, N);
+        } else {
+          mma_op.store_result_slice(
+              y, N, short2(0, offset), short2(BN, offset_next));
+        }
+      }
+
+      // Tile partially aligned check rows
+      else if (align_N || tgp_bn == BN) {
+        gemm_loop_unaligned<false, true, transpose>(
+            Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(BN, offset_next));
+      }
+
+      // Tile partially aligned check cols
+      else if (align_M || tgp_bm == BM) {
+        gemm_loop_unaligned<true, false, transpose>(
+            Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(tgp_bn, offset_next));
+      }
+
+      // Nothing aligned so check both rows and cols
+      else {
+        gemm_loop_unaligned<false, false, transpose>(
+            Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(tgp_bn, offset_next));
+      }
+    }
+  }
+}
+
+template <typename T, const int group_size, const int bits>
+[[kernel]] void fp_quantize(
+    const device T* w [[buffer(0)]],
+    device uint8_t* out [[buffer(1)]],
+    device uint8_t* scales [[buffer(2)]],
+    uint2 tidx [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  constexpr bool use_mx_scale = group_size == 32;
+  size_t index = tidx.x + grid_dim.x * size_t(tidx.y);
+
+  float scale;
+  float w_thread = w[index];
+  if (use_mx_scale) {
+    scale = simd_max(abs(w_thread));
+  } else {
+    float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0);
+    float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0);
+    scale = tidx.x < 16 ? w_max_l : w_max_r;
+  }
+  scale /= bits == 4 ? 6.0f : 448.0f;
+
+  using ScaleType = metal::conditional_t<use_mx_scale, fp8_e8m0, fp8_e4m3>;
+  auto s = ScaleType(scale);
+  uint8_t q_scale = s.bits;
+  scale = float(s);
+
+  size_t gindex = index / group_size;
+  if (index % group_size == 0) {
+    scales[gindex] = q_scale;
+  }
+
+  uint8_t output = Quantize<bits>{}(scale == 0 ? 0.0f : w_thread / scale);
+  if (bits == 4) {
+    uint8_t sval = simd_shuffle_down(output, 1);
+    output |= sval << bits;
+  }
+  constexpr int pack_factor = bits == 8 ? 1 : 2;
+  if (index % pack_factor == 0) {
+    out[index / pack_factor] = output;
+  }
+}
+
+template <typename T, const int group_size, const int bits>
+[[kernel]] void fp_dequantize(
+    const device uint8_t* w [[buffer(0)]],
+    const device uint8_t* scales [[buffer(1)]],
+    device T* out [[buffer(3)]],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  constexpr bool use_mx_scale = group_size == 32;
+  constexpr int pack_factor = bits == 8 ? 1 : 2;
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  size_t oindex = offset * pack_factor;
+  size_t gindex = oindex / group_size;
+
+  out += oindex;
+
+  using ScaleType = metal::conditional_t<use_mx_scale, fp8_e8m0, fp8_e4m3>;
+  auto q_scale = ((device ScaleType*)(scales))[gindex];
+  auto scale = float(q_scale);
+
+  uint val = w[offset];
+#pragma clang loop unroll(full)
+  for (int i = 0; i < pack_factor; i++) {
+    uint8_t d;
+    if (bits == 4) {
+      d = (val >> (bits * i)) & 0x0f;
+    } else if (bits == 8) {
+      d = val;
+    }
+    out[i] = static_cast<T>(scale * Dequantize<bits>{}(d));
+  }
+}
+
+template <typename T, const int group_size, const int bits>
+[[kernel]] void fp_quantize_dequantize(
+    const device T* w [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    uint2 tidx [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  constexpr bool use_mx_scale = group_size == 32;
+  size_t index = tidx.x + grid_dim.x * size_t(tidx.y);
+
+  float scale;
+  float w_thread = w[index];
+  if (use_mx_scale) {
+    scale = simd_max(abs(w_thread));
+  } else {
+    float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0);
+    float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0);
+    scale = tidx.x < 16 ? w_max_l : w_max_r;
+  }
+  scale /= bits == 4 ? 6.0f : 448.0f;
+
+  using ScaleType = metal::conditional_t<use_mx_scale, fp8_e8m0, fp8_e4m3>;
+  auto s = ScaleType(scale);
+  scale = float(s);
+
+  uint8_t output = Quantize<bits>{}(scale == 0 ? 0.0f : w_thread / scale);
+
+  out[index] = static_cast<T>(scale * Dequantize<bits>{}(output));
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized.metal
new file mode 100644
index 0000000..01a2df0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized.metal
@@ -0,0 +1,153 @@
+// Copyright © 2025 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/quantized_utils.h"
+#include "mlx/backend/metal/kernels/fp_quantized.h"
+
+#define instantiate_quantized(mode, name, type, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits, \
+      fp_ ## name, \
+      type, \
+      group_size,   \
+      bits)
+
+#define instantiate_quantized_batched(mode, name, type, batched, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \
+      fp_ ## name,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      batched)
+
+#define instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \
+      fp_ ## name,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      aligned)
+
+#define instantiate_quantized_aligned_batched(mode, name, type, aligned, batched, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_batch_" #batched, \
+      fp_ ## name,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      aligned, \
+      batched)
+
+#define instantiate_quantized_quad(mode, name, type, D, batched, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \
+      fp_ ## name,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      D,       \
+      batched)
+
+#define instantiate_quantized_split_k(mode, name, type, split_k, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \
+      fp_ ## name,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      split_k)
+
+#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \
+      func,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      bm,      \
+      bn,      \
+      bk,      \
+      wm,      \
+      wn,      \
+      transpose)
+
+#define instantiate_quantized_batched_wrap(name, type, mode, group_size, bits) \
+  instantiate_quantized_batched(mode, name, type, 1, group_size, bits)         \
+  instantiate_quantized_batched(mode, name, type, 0, group_size, bits)
+
+#define instantiate_quantized_all_batched(type, mode, group_size, bits) \
+  instantiate_quantized_batched_wrap(qmv_fast, type, mode, group_size, bits) \
+  instantiate_quantized_batched_wrap(qmv, type, mode, group_size, bits)      \
+  instantiate_quantized_batched_wrap(qvm, type, mode, group_size, bits) \
+  instantiate_quantized_batched_wrap(qmm_n, type, mode, group_size, bits)
+
+#define instantiate_quantized_all_single(type, mode, group_size, bits) \
+  instantiate_quantized(mode, gather_qmv_fast, type, group_size, bits) \
+  instantiate_quantized(mode, gather_qmv, type, group_size, bits)      \
+  instantiate_quantized(mode, gather_qvm, type, group_size, bits) \
+  instantiate_quantized(mode, gather_qmm_n, type, group_size, bits)
+
+#define instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_aligned(mode, gather_qmm_t, type, true, group_size, bits)      \
+  instantiate_quantized_aligned(mode, gather_qmm_t, type, false, group_size, bits)     \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 1, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 0, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 1, group_size, bits) \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 0, group_size, bits)
+
+#define instantiate_quantized_all_quad(type, mode, group_size, bits) \
+  instantiate_quantized_quad(mode, qmv_quad, type, 64, 1, group_size, bits)  \
+  instantiate_quantized_quad(mode, qmv_quad, type, 64, 0, group_size, bits)  \
+  instantiate_quantized_quad(mode, qmv_quad, type, 128, 1, group_size, bits) \
+  instantiate_quantized_quad(mode, qmv_quad, type, 128, 0, group_size, bits)
+
+#define instantiate_quantized_all_splitk(type, mode, group_size, bits) \
+  instantiate_quantized_split_k(mode, qvm_split_k, type, 8, group_size, bits) \
+  instantiate_quantized_split_k(mode, qvm_split_k, type, 32, group_size, bits)
+
+#define instantiate_quantized_all_rhs(type, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nt, type, 16, 32, 32, 1, 2, true, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false, mode, group_size, bits)
+
+#define instantiate_quantize_dequantize(type, mode, group_size, bits) \
+  instantiate_kernel( \
+    #mode "_quantize_dequantize_" #type "_gs_" #group_size "_b_" #bits, \
+    fp_quantize_dequantize, \
+    type, \
+    group_size,  \
+    bits) \
+  instantiate_kernel( \
+    #mode "_quantize_" #type "_gs_" #group_size "_b_" #bits, \
+    fp_quantize, \
+    type, \
+    group_size,  \
+    bits) \
+  instantiate_kernel( \
+    #mode "_dequantize_" #type "_gs_" #group_size "_b_" #bits, \
+    fp_dequantize, \
+    type, \
+    group_size,  \
+    bits)
+
+#define instantiate_quantized_modes(type, mode, group_size, bits) \
+  instantiate_quantized_all_batched(type, mode, group_size, bits) \
+  instantiate_quantized_all_single(type, mode, group_size, bits)  \
+  instantiate_quantized_all_quad(type, mode, group_size, bits)    \
+  instantiate_quantized_all_splitk(type, mode, group_size, bits)  \
+  instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_all_rhs(type, mode, group_size, bits)     \
+  instantiate_quantize_dequantize(type, mode, group_size, bits)
+
+#define instantiate_quantized_types(type) \
+  instantiate_quantized_modes(type, nvfp4, 16, 4) \
+  instantiate_quantized_modes(type, mxfp8, 32, 8) \
+  instantiate_quantized_modes(type, mxfp4, 32, 4)
+
+instantiate_quantized_types(float)
+instantiate_quantized_types(bfloat16_t)
+instantiate_quantized_types(float16_t)
+    // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h
new file mode 100644
index 0000000..80e1c4c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h
@@ -0,0 +1,1044 @@
+// Copyright © 2025 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/fp4.h"
+#include "mlx/backend/metal/kernels/fp8.h"
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+using namespace metal;
+
+#define MLX_MTL_CONST static constant constexpr const
+
+MLX_MTL_CONST int SIMD_SIZE = 32;
+MLX_MTL_CONST int QUAD_SIZE = 4;
+
+template <int wsize = 8, int bits>
+inline constexpr short get_pack_factor() {
+  return wsize / bits;
+}
+
+template <int wsize = 8>
+inline constexpr short get_bytes_per_pack() {
+  return wsize / 8;
+}
+
+template <typename T, int group_size>
+static inline T dequantize_scale(uint8_t s) {
+  if constexpr (group_size == 16) {
+    // Use nv scale
+    return T(*(thread fp8_e4m3*)(&s));
+  } else {
+    return T(*(thread fp8_e8m0*)(&s));
+  }
+}
+
+template <int bits>
+struct Quantize {
+  uint8_t operator()(float x) {
+    if (bits == 8) {
+      return fp8_e4m3(x).bits;
+    } else {
+      return fp4_e2m1(x).bits;
+    }
+  }
+};
+
+template <int bits, typename U = float>
+struct Dequantize {
+  U operator()(uint8_t x) {
+    if constexpr (bits == 8) {
+      return U(*(thread fp8_e4m3*)(&x));
+    } else {
+      return U(*(thread fp4_e2m1*)(&x));
+    }
+  }
+};
+
+template <typename U, int bits>
+inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) {
+  if constexpr (bits == 4) {
+    w_local[0] = scale * Dequantize<4, U>{}(w);
+    w_local[1] = scale * Dequantize<4, U>{}(w >> 4);
+  } else {
+    w_local[0] = scale * Dequantize<8, U>{}(w);
+  }
+}
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short group_size,
+    short bits>
+struct QuantizedBlockLoader {
+  MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>();
+  MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack();
+  MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
+  MLX_MTL_CONST short n_reads =
+      (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
+
+  MLX_MTL_CONST short n_reads_per_scale = (n_reads * pack_factor) <= group_size
+      ? n_reads
+      : (group_size / pack_factor);
+  MLX_MTL_CONST short n_steps_per_read = n_reads / n_reads_per_scale;
+
+  MLX_MTL_CONST short n_groups = BCOLS / group_size;
+
+  const int src_ld;
+  const int tile_stride;
+  const int group_stride;
+
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  const short group_id;
+
+  threadgroup T* dst;
+  const device uint8_t* src;
+  const device uint8_t* scales;
+
+  QuantizedBlockLoader(
+      const device uint8_t* src_,
+      const device uint8_t* scales_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(
+            reduction_dim ? BCOLS_PACKED * bytes_per_pack
+                          : BROWS * src_ld * bytes_per_pack / pack_factor),
+        group_stride(BROWS * src_ld / group_size),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(n_reads * thread_idx / BCOLS_PACKED),
+        bj((n_reads * thread_idx) % BCOLS_PACKED),
+        group_id((bj * pack_factor) / group_size),
+        dst(dst_ + bi * dst_ld + bj * pack_factor),
+        src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
+            bj * bytes_per_pack),
+        scales(scales_ + bi * src_ld / group_size + group_id) {}
+
+  void load_unsafe() const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    int k = 0;
+    for (int i = 0; i < n_steps_per_read; i++) {
+      T scale = dequantize_scale<T, group_size>(scales[i]);
+      for (int j = 0; j < n_reads_per_scale; j++) {
+        dequantize<T, bits>(
+            src[k * bytes_per_pack], scale, dst + k * pack_factor);
+        k++;
+      }
+    }
+  }
+
+  void load_safe(short2 src_tile_dim) const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    int k = 0;
+    for (int i = 0; i < n_steps_per_read; i++) {
+      T scale = dequantize_scale<T, group_size>(scales[i]);
+      for (int j = 0; j < n_reads_per_scale; j++) {
+        dequantize<T, bits>(
+            src[k * bytes_per_pack], scale, dst + k * pack_factor);
+        k++;
+      }
+    }
+  }
+
+  void next() {
+    src += tile_stride;
+    if (reduction_dim == 1) {
+      scales += n_groups;
+    } else {
+      scales += n_groups * group_stride;
+    }
+  }
+};
+
+using namespace mlx::steel;
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2,
+    typename Wtype = bfloat>
+METAL_FUNC void fp_qmm_t_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    threadgroup Wtype* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  (void)lid;
+
+  constexpr int pack_factor = get_pack_factor<8, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack();
+
+  constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
+
+  // Instantiate Loader
+  using loader_w_t = QuantizedBlockLoader<
+      Wtype,
+      BN,
+      BK,
+      BK_padded,
+      1,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  // Set the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+
+  auto wl = (const device uint8_t*)w;
+
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * K_w;
+  scales += y_col * K_g;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the weight loader
+  loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid);
+
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  const short tm = SM * (simd_gid / WN);
+  const short tn = SN * (simd_gid % WN);
+
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = true;
+
+  const short sgp_sm = min(SM, short(M - (y_row + tm)));
+  const bool is_unaligned_sm = (sgp_sm != SM);
+
+  const short sgp_sn = aligned_N ? SN : min(SN, short(N - (y_col + tn)));
+
+  const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col)));
+  const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN);
+
+  using AccumType = float;
+
+  using ASubTile = NAXSubTile<T, UM, UK>;
+  using BSubTile = NAXSubTile<Wtype, UN, UK>;
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  Dtile.clear();
+
+  x += tm * K;
+
+  dispatch_bool(!is_unaligned_sm, [&](auto kAlignedM) {
+    dispatch_bool(aligned_N || !is_unaligned_bn, [&](auto kAlignedN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if constexpr (kAlignedN.value) {
+          loader_w.load_unsafe();
+        } else {
+          loader_w.load_safe(short2(BK, tgp_bn));
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        STEEL_PRAGMA_NO_UNROLL
+        for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+          NAXTile<T, TM, TK, ASubTile> Atile;
+          NAXTile<Wtype, TN, TK, BSubTile> Btile;
+
+          volatile int compiler_barrier;
+
+          if constexpr (kAlignedM.value) {
+            Atile.load(x + kk1, K);
+          } else {
+            Atile.load_safe(x + kk1, K, short2(SK, sgp_sm));
+          }
+
+          Btile.template load<Wtype, BK_padded, 1>(Ws + tn * BK_padded + kk1);
+
+          tile_matmad_nax(
+              Dtile,
+              Atile,
+              metal::bool_constant<transpose_a>{},
+              Btile,
+              metal::bool_constant<transpose_b>{});
+
+          (void)compiler_barrier;
+        }
+
+        x += BK;
+        loader_w.next();
+      }
+
+      // Store results to device memory
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      if constexpr (kAlignedM.value && kAlignedN.value) {
+        Dtile.store(y + tm * N + tn, N);
+      } else if (kAlignedM.value && sgp_sn == SN) {
+        Dtile.store(y + tm * N + tn, N);
+      } else {
+        Dtile.store_safe(y + tm * N + tn, N, short2(sgp_sn, sgp_sm));
+      }
+    });
+  });
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2,
+    typename Wtype = bfloat>
+METAL_FUNC void fp_qmm_n_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    threadgroup T* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  (void)lid;
+  (void)M;
+
+  constexpr int pack_factor = get_pack_factor<8, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack();
+
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      BK,
+      BN,
+      BN_padded,
+      0,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  // Set the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+
+  auto wl = (const device uint8_t*)w;
+
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * K_w;
+  scales += y_col * K_g;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the x loader and mma operation
+  // const short num_els = min(BM, M - y_row);
+  // const short num_outs = min(BN, N - y_col);
+  loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid);
+
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  const short tm = SM * (simd_gid / WN);
+  const short tn = SN * (simd_gid % WN);
+
+  const short ldb_tgp = BN_padded;
+
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = false;
+
+  using AccumType = float;
+
+  using ASubTile = NAXSubTile<T, UM, UK>;
+  using BSubTile = NAXSubTile<T, UK, UN>;
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  Dtile.clear();
+
+  x += tm * K;
+
+  for (int k = 0; k < K; k += BK) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    loader_w.load_unsafe();
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_NO_UNROLL
+    for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+      NAXTile<T, TM, TK, ASubTile> Atile;
+      NAXTile<Wtype, TK, TN, BSubTile> Btile;
+
+      volatile int compiler_barrier;
+
+      Atile.load(x + kk1, K);
+      Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * ldb_tgp);
+
+      tile_matmad_nax(
+          Dtile,
+          Atile,
+          metal::bool_constant<transpose_a>{},
+          Btile,
+          metal::bool_constant<transpose_b>{});
+
+      (void)compiler_barrier;
+    }
+
+    x += BK;
+    loader_w.next();
+  }
+
+  // Store results to device memory
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  Dtile.store(y + tm * N + tn, N);
+}
+
+template <typename T, typename S>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device S*& scales,
+    device T*& y,
+    int output_stride,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx = tid.z;
+  uint32_t w_idx = tid.z;
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+  }
+  y += tid.z * output_stride;
+}
+
+template <typename T, typename S>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device S*& scales,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T*& y,
+    int output_stride,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx;
+  uint32_t w_idx;
+  if (batch_ndims == 1) {
+    x_idx = lhs_indices[tid.z * lhs_strides[0]];
+    w_idx = rhs_indices[tid.z * rhs_strides[0]];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        tid.z, batch_shape, lhs_strides, rhs_strides, batch_ndims);
+    x_idx = lhs_indices[idx.x];
+    w_idx = rhs_indices[idx.y];
+  }
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+  }
+  y += tid.z * output_stride;
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const bool batched,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2,
+    typename Wtype = bfloat>
+[[kernel]] void fp_qmm_t_nax(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
+
+  threadgroup Wtype Ws[BN * BK_padded];
+
+  if (batched) {
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+  fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN, WM, WN, Wtype>(
+      w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool batched,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2,
+    typename Wtype = bfloat>
+[[kernel]] void fp_qmm_n_nax(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BK * BN_padded];
+
+  if (batched) {
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+
+  fp_qmm_n_impl<T, group_size, bits, BM, BK, BN, WM, WN, Wtype>(
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2,
+    typename Wtype = bfloat>
+[[kernel]] void fp_gather_qmm_t_nax(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
+
+  threadgroup Wtype Ws[BN * BK_padded];
+
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+  fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN, WM, WN, Wtype>(
+      w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2,
+    typename Wtype = bfloat>
+[[kernel]] void fp_gather_qmm_n_nax(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BK * BN_padded];
+
+  adjust_matrix_offsets(
+      x,
+      w,
+      scales,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      tid);
+  fp_qmm_n_impl<T, group_size, bits, BM, BK, BN, WM, WN, Wtype>(
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    int group_size,
+    const int bits,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose,
+    typename Wtype = bfloat>
+[[kernel]] void fp_gather_qmm_rhs_nax(
+    const device T* x,
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device uint32_t* indices,
+    device T* y,
+    const constant int& M,
+    const constant int& N,
+    const constant int& K,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]]) {
+  constexpr int pack_factor = get_pack_factor<8, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack();
+  constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
+  constexpr int BN_padded = (BN + 16 / sizeof(Wtype));
+
+  using loader_w_t = QuantizedBlockLoader<
+      Wtype,
+      transpose ? BN : BK,
+      transpose ? BK : BN,
+      transpose ? BK_padded : BN_padded,
+      transpose,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded];
+
+  // Compute the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int N_w = N * bytes_per_pack / pack_factor;
+  const int N_g = N / group_size;
+  const int K_it = K / BK;
+  const size_t stride_w = transpose ? N * K_w : K * N_w;
+  const size_t stride_s = transpose ? N * K_g : K * N_g;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+  const size_t y_row_long = size_t(y_row);
+  const size_t y_col_long = size_t(y_col);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, M - y_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, N - y_col));
+
+  // Calculate the final tiles in the case that K is not aligned
+  const int k_remain = K - K_it * BK;
+  const short2 tile_w =
+      transpose ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+  // Move x and output to the correct block
+  auto wl = (const device uint8_t*)w;
+  x += y_row_long * K;
+  y += y_row_long * N + y_col_long;
+  wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor;
+  scales += transpose ? y_col_long * K_g : y_col / group_size;
+
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  const short tm = SM * (simd_group_id / WN);
+  const short tn = SN * (simd_group_id % WN);
+
+  const short sgp_sm =
+      align_M ? SM : min(SM, short(max(0, (M - (y_row + tm)))));
+  const short sgp_sn =
+      align_N ? SN : min(SN, short(max(0, (N - (y_col + tn)))));
+
+  const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
+  const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN);
+
+  constexpr short BR = transpose ? TN : TK;
+  constexpr short BC = transpose ? TK : TN;
+
+  using AccumType = float;
+
+  using ASubTile = NAXSubTile<T, UM, UK>;
+  using BSubTile = NAXSubTile<Wtype, transpose ? UN : UK, transpose ? UK : UN>;
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+
+  // Do as many matmuls as necessary
+  uint32_t index;
+  short offset;
+  uint32_t index_next = indices[y_row];
+  short offset_next = 0;
+  int n = 0;
+  while (n < tgp_bm) {
+    n++;
+    offset = offset_next;
+    index = index_next;
+    offset_next = tgp_bm;
+    for (; n < tgp_bm; n++) {
+      if (indices[y_row + n] != index) {
+        offset_next = n;
+        index_next = indices[y_row + n];
+        break;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Prepare threadgroup mma operation
+    NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+    Dtile.clear();
+
+    const device T* xn = x + tm * K;
+
+    // Prepare threadgroup loading operations
+    thread loader_w_t loader_w(
+        wl + index * stride_w,
+        scales + index * stride_s,
+        transpose ? K : N,
+        Ws,
+        simd_group_id,
+        simd_lane_id);
+
+    dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+      dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) {
+        for (int k = 0; k < K_it; k++) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          if constexpr (kAlignedN.value) {
+            loader_w.load_unsafe();
+          } else {
+            loader_w.load_safe(
+                transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK));
+          }
+
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+
+          STEEL_PRAGMA_NO_UNROLL
+          for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+            NAXTile<T, TM, TK, ASubTile> Atile;
+            NAXTile<Wtype, BR, BC, BSubTile> Btile;
+
+            volatile int compiler_barrier;
+
+            if constexpr (kAlignedM.value) {
+              Atile.load(xn + kk1, K);
+            } else {
+              Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm));
+            }
+
+            if constexpr (transpose) {
+              Btile.template load<Wtype, BK_padded, 1>(
+                  Ws + tn * BK_padded + kk1);
+            } else {
+              Btile.template load<Wtype, BN_padded, 1>(
+                  Ws + tn + kk1 * BN_padded);
+            }
+
+            tile_matmad_nax(
+                Dtile,
+                Atile,
+                metal::bool_constant<false>{},
+                Btile,
+                metal::bool_constant<transpose>{});
+
+            (void)compiler_barrier;
+          }
+
+          xn += BK;
+          loader_w.next();
+        }
+
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          loader_w.load_safe(tile_w);
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+
+          STEEL_PRAGMA_NO_UNROLL
+          for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+            NAXTile<T, TM, TK, ASubTile> Atile;
+            NAXTile<Wtype, BR, BC, BSubTile> Btile;
+
+            volatile int compiler_barrier;
+
+            const short psk = min(int(SK), max(0, (BK - kk1)));
+            Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm));
+
+            if constexpr (transpose) {
+              Btile.template load<Wtype, BK_padded, 1>(
+                  Ws + tn * BK_padded + kk1);
+            } else {
+              Btile.template load<Wtype, BN_padded, 1>(
+                  Ws + tn + kk1 * BN_padded);
+            }
+
+            tile_matmad_nax(
+                Dtile,
+                Atile,
+                metal::bool_constant<false>{},
+                Btile,
+                metal::bool_constant<transpose>{});
+
+            (void)compiler_barrier;
+          }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm));
+        const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm));
+
+        // Store results to device memory
+        if constexpr (kAlignedN.value) {
+          if (m_lo_lim == 0 && m_hi_lim == SM) {
+            Dtile.store(y + tm * N + tn, N);
+          } else {
+            Dtile.store_slice(
+                y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim));
+          }
+        } else {
+          Dtile.store_slice(
+              y + tm * N + tn,
+              N,
+              short2(0, m_lo_lim),
+              short2(sgp_sn, m_hi_lim));
+        }
+      });
+    });
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal
new file mode 100644
index 0000000..4d65a38
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal
@@ -0,0 +1,79 @@
+// Copyright © 2025 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/quantized_utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/nax.h"
+#include "mlx/backend/metal/kernels/fp_quantized_nax.h"
+
+
+#define instantiate_quantized_batched(mode, name, type, bm, bn, bk, wm, wn, batched, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched, \
+      fp_ ## name,  \
+      type,         \
+      group_size,           \
+      bits,            \
+      batched)
+
+#define instantiate_quantized_aligned(mode, name, type, bm, bn, bk, wm, wn, aligned, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned, \
+      fp_ ## name, \
+      type,        \
+      group_size,          \
+      bits,           \
+      aligned)
+
+#define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched, \
+      fp_ ## name,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      aligned, \
+      batched)
+
+#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \
+      func,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      bm,      \
+      bn,      \
+      bk,      \
+      wm,      \
+      wn,      \
+      transpose)
+
+
+#define instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_aligned(mode, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, true, group_size, bits)      \
+  instantiate_quantized_aligned(mode, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, false, group_size, bits)     \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1, group_size, bits) \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits)
+
+
+#define instantiate_quantized_all_rhs(type, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false, mode, group_size, bits)
+
+#define instantiate_quantized_modes(type, mode, group_size, bits) \
+  instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_all_rhs(type, mode, group_size, bits)
+
+#define instantiate_quantized_types(type) \
+  instantiate_quantized_modes(type, nvfp4, 16, 4) \
+  instantiate_quantized_modes(type, mxfp8, 32, 8) \
+  instantiate_quantized_modes(type, mxfp4, 32, 4)
+
+instantiate_quantized_types(float)
+instantiate_quantized_types(bfloat16_t)
+instantiate_quantized_types(float16_t)
+    // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/gemv.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/gemv.metal
new file mode 100644
index 0000000..277b68a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/gemv.metal
@@ -0,0 +1,868 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+/// Matrix vector multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+#define MLX_MTL_CONST static constant constexpr const
+
+template <typename U>
+struct DefaultAccT {
+  using type = float;
+};
+template <>
+struct DefaultAccT<complex64_t> {
+  using type = complex64_t;
+};
+
+template <
+    typename T,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoAxpby, /* Do out = alpha * out + beta * bias */
+    typename AccT = typename DefaultAccT<T>::type>
+struct GEMVKernel {
+  using acc_type = AccT;
+
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
+
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  static_assert(
+      SN == 4 || SN == 8 || SN == 16 || SN == 32,
+      "gemv block must have a width of 4, 8, 16, or 32");
+
+  // - The matrix of size (M = out_vec_size, K = in_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM rows
+  //    and the corresponding scalar from the vector
+  // 2. The thread then multiplies and adds to accumulate its local result for
+  //    the block
+  // 3. At the end, each thread has accumulated results over all blocks across
+  //    the rows. These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated blockM outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid has blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results
+  //     remain zero)
+  //   * The last thread that partially overlaps with the matrix is shifted
+  //     inwards such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BN > 1 ? BN*(blockM + TM) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BN > 1;
+
+  template <typename U = T>
+  static METAL_FUNC void
+  load_unsafe(const device T* src, thread U dst[TN], const int src_offset = 0) {
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tn = 0; tn < TN; tn++) {
+      dst[tn] = static_cast<U>(src[src_offset + tn]);
+    }
+  }
+
+  template <typename U = T>
+  static METAL_FUNC void load_safe(
+      const device T* src,
+      thread U dst[TN],
+      const int src_offset = 0,
+      const int src_size = TN) {
+    if (src_offset + TN <= src_size) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = static_cast<U>(src[src_offset + tn]);
+      }
+    } else { // Edgecase
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = src_offset + tn < src_size
+            ? static_cast<U>(src[src_offset + tn])
+            : U(0);
+      }
+    }
+  }
+
+  static METAL_FUNC void run(
+      const device T* mat [[buffer(0)]],
+      const device T* in_vec [[buffer(1)]],
+      const device T* bias [[buffer(2)]],
+      device T* out_vec [[buffer(3)]],
+      const constant int& in_vec_size [[buffer(4)]],
+      const constant int& out_vec_size [[buffer(5)]],
+      const constant int& matrix_ld [[buffer(6)]],
+      const constant float& alpha [[buffer(7)]],
+      const constant float& beta [[buffer(8)]],
+      const constant int& bias_stride [[buffer(14)]],
+      threadgroup AccT* tgp_memory [[threadgroup(0)]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]],
+      uint simd_gid [[simdgroup_index_in_threadgroup]],
+      uint simd_lid [[thread_index_in_simdgroup]]) {
+    // Appease compiler
+    (void)lid;
+
+    // Thread local accumulation results
+    thread AccT result[TM] = {0};
+    thread T inter[TN];
+    thread AccT v_coeff[TN];
+
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
+
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = BN != 1 ? SM * (simd_gid / BN) : int(SM * simd_gid);
+    const int simdN = BN != 1 ? SN * (simd_gid % BN) : 0;
+
+    int bm = (simdM + thrM) * TM;
+    int bn = (simdN + thrN) * TN;
+
+    // Block position
+    int out_row = tid.x * blockM + bm;
+
+    // Exit simdgroup if rows out of bound
+    if (out_row >= out_vec_size)
+      return;
+
+    // Adjust tail simdgroup to ensure in bound reads
+    out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM;
+
+    // Advance matrix
+    mat += out_row * matrix_ld;
+
+    constexpr const uniform<int> loop_stride = make_uniform(blockN);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
+
+    // Loop over in_vec in blocks of blockN
+    for (int i = 0; i < n_iter; ++i) {
+      load_unsafe<AccT>(in_vec, v_coeff, bn);
+
+      // Per thread work loop
+      int mat_offset = 0;
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        // Load for the row
+        load_unsafe(mat, inter, mat_offset + bn);
+
+        // Accumulate results
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          result[tm] += inter[tn] * v_coeff[tn];
+        }
+
+        mat_offset += matrix_ld;
+      }
+
+      bn += blockN;
+    }
+
+    if (leftover > 0) {
+      load_safe<AccT>(in_vec, v_coeff, bn, in_size);
+
+      // Per thread work loop
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        // Load for the row
+        load_safe(&mat[tm * matrix_ld], inter, bn, in_size);
+
+        // Accumulate results
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          result[tm] += inter[tn] * v_coeff[tn];
+        }
+      }
+    }
+
+    // Simdgroup accumulations
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tm = 0; tm < TM; tm++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sn = (SN / 2); sn >= 1; sn >>= 1) {
+        result[tm] += simd_shuffle_down(result[tm], sn);
+      }
+    }
+
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup AccT* tgp_results = tgp_memory + sgN * (blockM + TM) + bm;
+      if (thrN == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          tgp_results[tm] = result[tm];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgN == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgn = 1; sgn < BN; sgn++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tm = 0; tm < TM; tm++) {
+              result[tm] += tgp_results[sgn * (blockM + TM) + tm];
+            }
+          }
+        }
+      }
+    }
+
+    // Write outputs
+    if (simdN == 0 && thrN == 0) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        if (kDoAxpby) {
+          out_vec[out_row + tm] =
+              static_cast<T>(alpha) * static_cast<T>(result[tm]) +
+              static_cast<T>(beta) * bias[(out_row + tm) * bias_stride];
+        } else {
+          out_vec[out_row + tm] = static_cast<T>(result[tm]);
+        }
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoAxpby, /* Do out = alpha * out + beta * bias */
+    typename AccT = typename DefaultAccT<T>::type>
+struct GEMVTKernel {
+  using acc_type = AccT;
+
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
+
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM contiguous rows
+  //    and the corresponding scalar from the vector
+  // 2. The thread then accumulates its local result for the block
+  // 3. At the end, each thread has accumulated results over all blocks across
+  //    the rows. These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated BN * TN outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid has blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results
+  //     remain zero)
+  //   * The last thread that partially overlaps with the matrix is shifted
+  //     inwards such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BM > 1 ? BM*(blockN + TN) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BM > 1;
+
+  static METAL_FUNC void run(
+      const device T* mat [[buffer(0)]],
+      const device T* in_vec [[buffer(1)]],
+      const device T* bias [[buffer(2)]],
+      device T* out_vec [[buffer(3)]],
+      const constant int& in_vec_size [[buffer(4)]],
+      const constant int& out_vec_size [[buffer(5)]],
+      const constant int& marix_ld [[buffer(6)]],
+      const constant float& alpha [[buffer(7)]],
+      const constant float& beta [[buffer(8)]],
+      const constant int& bias_stride [[buffer(14)]],
+      threadgroup AccT* tgp_memory [[threadgroup(0)]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]],
+      uint simd_gid [[simdgroup_index_in_threadgroup]],
+      uint simd_lid [[thread_index_in_simdgroup]]) {
+    // Appease compiler
+    (void)lid;
+
+    // Thread local accumulation results
+    AccT result[TN] = {0};
+    T inter[TN];
+    AccT v_coeff[TM];
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
+
+    const int sgM = BN != 1 ? (simd_gid / BN) : int(simd_gid);
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = SM * sgM;
+    const int simdN = SN * sgN;
+
+    int cm = (simdM + thrM);
+    int cn = (simdN + thrN);
+
+    int bm = cm * TM;
+    int bn = cn * TN;
+
+    int out_col = tid.x * blockN + bn;
+
+    constexpr const uniform<int> loop_stride = make_uniform(blockM);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
+
+    // Edgecase handling
+    if (out_col < out_vec_size) {
+      out_col = out_col + TN < out_vec_size ? out_col : out_vec_size - TN;
+
+      // Per thread accumulation main loop
+      for (int i = 0; i < n_iter; ++i) {
+        // Adding a threadgroup_barrier improves performance slightly
+        // This is possibly it may help exploit cache better
+        threadgroup_barrier(mem_flags::mem_none);
+
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          v_coeff[tm] = static_cast<AccT>(in_vec[bm + tm]);
+        }
+
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          auto vc = static_cast<AccT>(v_coeff[tm]);
+          for (int tn = 0; tn < TN; tn++) {
+            inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
+          }
+          for (int tn = 0; tn < TN; tn++) {
+            result[tn] += vc * inter[tn];
+          }
+        }
+
+        bm += blockM;
+      }
+
+      if (leftover > 0) {
+        for (int tm = 0; tm < TM && bm + tm < in_vec_size; tm++) {
+          v_coeff[tm] = static_cast<AccT>(in_vec[bm + tm]);
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
+          }
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            result[tn] += v_coeff[tm] * inter[tn];
+          }
+        }
+      }
+    }
+
+    // Simdgroup accumulations
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tn = 0; tn < TN; tn++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sm = (SM / 2); sm >= 1; sm >>= 1) {
+        result[tn] += simd_shuffle_down(result[tn], SN * sm);
+      }
+    }
+
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup AccT* tgp_results = tgp_memory + sgM * (blockN + TN) + bn;
+      if (thrM == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          tgp_results[tn] = result[tn];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgM == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgm = 1; sgm < BM; sgm++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              result[tn] += tgp_results[sgm * (blockN + TN) + tn];
+            }
+          }
+        }
+      }
+    }
+
+    // Threadgroup accumulation and writing out results
+    if (cm == 0 && out_col < out_vec_size) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int j = 0; j < TN; j++) {
+        if (kDoAxpby) {
+          out_vec[out_col + j] =
+              static_cast<T>(alpha) * static_cast<T>(result[j]) +
+              static_cast<T>(beta) * bias[(out_col + j) * bias_stride];
+        } else {
+          out_vec[out_col + j] = static_cast<T>(result[j]);
+        }
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Matrix vector multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoNCBatch, /* Batch ndim > 1 */
+    const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
+[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    const device T* bias [[buffer(2)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant float& alpha [[buffer(7)]],
+    const constant float& beta [[buffer(8)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant int64_t* vector_batch_stride [[buffer(11)]],
+    const constant int64_t* matrix_batch_stride [[buffer(12)]],
+    const constant int64_t* bias_batch_stride [[buffer(13)]],
+    const constant int& bias_stride [[buffer(14)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  using gemv_kernel = GEMVKernel<T, BM, BN, SM, SN, TM, TN, kDoAxpby>;
+  threadgroup typename gemv_kernel::acc_type tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  // Update batch offsets
+  if (kDoNCBatch) {
+    in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim);
+    mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim);
+
+    if (kDoAxpby) {
+      bias += elem_to_loc(tid.z, batch_shape, bias_batch_stride, batch_ndim);
+    }
+
+  } else {
+    in_vec += tid.z * vector_batch_stride[0];
+    mat += tid.z * matrix_batch_stride[0];
+
+    if (kDoAxpby) {
+      bias += tid.z * bias_batch_stride[0];
+    }
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      bias,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      alpha,
+      beta,
+      bias_stride,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}
+
+#define instantiate_gemv_helper(                                      \
+    name, itype, bm, bn, sm, sn, tm, tn, nc, axpby)                   \
+  instantiate_kernel(                                                 \
+      "gemv_" #name "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn "_tm" #tm \
+      "_tn" #tn "_nc" #nc "_axpby" #axpby,                            \
+      gemv,                                                           \
+      itype,                                                          \
+      bm,                                                             \
+      bn,                                                             \
+      sm,                                                             \
+      sn,                                                             \
+      tm,                                                             \
+      tn,                                                             \
+      nc,                                                             \
+      axpby)
+
+// clang-format off
+#define instantiate_gemv(name, itype, bm, bn, sm, sn, tm, tn)        \
+  instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 0) \
+  instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 1) \
+  instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 0) \
+  instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 1) // clang-format on
+
+// clang-format off
+#define instantiate_gemv_blocks(name, itype) \
+  instantiate_gemv(name, itype, 1,  8, 1, 32, 4, 4) \
+  instantiate_gemv(name, itype, 1,  8, 1, 32, 1, 4) \
+  instantiate_gemv(name, itype, 1,  1, 8,  4, 4, 4) \
+  instantiate_gemv(name, itype, 1,  1, 8,  4, 1, 4) \
+  instantiate_gemv(name, itype, 4,  1, 1, 32, 1, 4) \
+  instantiate_gemv(name, itype, 4,  1, 1, 32, 4, 4) \
+  instantiate_gemv(name, itype, 8,  1, 1, 32, 4, 4) // clang-format on
+
+instantiate_gemv_blocks(float32, float);
+instantiate_gemv_blocks(float16, half);
+instantiate_gemv_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_blocks(complex64, complex64_t);
+
+template <
+    typename T,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_gather(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    const device T* bias [[buffer(2)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant float& alpha [[buffer(7)]],
+    const constant float& beta [[buffer(8)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant int64_t* index_batch_strides [[buffer(11)]],
+    const constant int& vector_batch_ndim [[buffer(12)]],
+    const constant int* vector_batch_shape [[buffer(13)]],
+    const constant int64_t* vector_batch_stride [[buffer(14)]],
+    const constant int& matrix_batch_ndim [[buffer(15)]],
+    const constant int* matrix_batch_shape [[buffer(16)]],
+    const constant int64_t* matrix_batch_stride [[buffer(17)]],
+    const constant uint32_t* vec_indices [[buffer(18)]],
+    const constant uint32_t* mat_indices [[buffer(19)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  using gemv_kernel = GEMVKernel<T, BM, BN, SM, SN, TM, TN, false>;
+  threadgroup typename gemv_kernel::acc_type tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  uint32_t indx_vec;
+  uint32_t indx_mat;
+
+  // Update batch offsets
+  if (batch_ndim > 1) {
+    const constant auto* veci_bstrides = index_batch_strides;
+    const constant auto* mati_bstrides = index_batch_strides + batch_ndim;
+
+    ulong2 batch_offsets = elem_to_loc_broadcast(
+        tid.z, batch_shape, veci_bstrides, mati_bstrides, batch_ndim);
+
+    indx_vec = vec_indices[batch_offsets.x];
+    indx_mat = mat_indices[batch_offsets.y];
+
+  } else {
+    indx_vec = vec_indices[index_batch_strides[0] * tid.z];
+    indx_mat = mat_indices[index_batch_strides[batch_ndim] * tid.z];
+  }
+
+  if (vector_batch_ndim > 1) {
+    in_vec += elem_to_loc(
+        indx_vec, vector_batch_shape, vector_batch_stride, vector_batch_ndim);
+  } else {
+    in_vec += indx_vec * vector_batch_stride[0];
+  }
+
+  if (matrix_batch_ndim > 1) {
+    mat += elem_to_loc(
+        indx_mat, matrix_batch_shape, matrix_batch_stride, matrix_batch_ndim);
+  } else {
+    mat += indx_mat * matrix_batch_stride[0];
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      bias,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      alpha,
+      beta,
+      batch_ndim, // Not used
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}
+
+// clang-format off
+#define instantiate_gemv_bs_helper(nm, itype, bm, bn, sm, sn, tm, tn) \
+  instantiate_kernel(                                                 \
+    "gemv_gather_" #nm "_bm" #bm "_bn" #bn "_sm" #sm                  \
+                       "_sn" #sn "_tm" #tm "_tn" #tn,                 \
+    gemv_gather, itype, bm, bn, sm, sn, tm, tn)
+
+#define instantiate_gemv_bs_blocks(name, itype)              \
+  instantiate_gemv_bs_helper(name, itype, 4, 1, 1, 32, 1, 4) \
+  instantiate_gemv_bs_helper(name, itype, 4, 1, 1, 32, 4, 4) \
+  instantiate_gemv_bs_helper(name, itype, 8, 1, 1, 32, 4, 4) // clang-format on
+
+instantiate_gemv_bs_blocks(float32, float);
+instantiate_gemv_bs_blocks(float16, half);
+instantiate_gemv_bs_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_bs_blocks(complex64, complex64_t);
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoNCBatch, /* Batch ndim > 1 */
+    const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
+[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    const device T* bias [[buffer(2)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant float& alpha [[buffer(7)]],
+    const constant float& beta [[buffer(8)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant int64_t* vector_batch_stride [[buffer(11)]],
+    const constant int64_t* matrix_batch_stride [[buffer(12)]],
+    const constant int64_t* bias_batch_stride [[buffer(13)]],
+    const constant int& bias_stride [[buffer(14)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  using gemv_kernel = GEMVTKernel<T, BM, BN, SM, SN, TM, TN, kDoAxpby>;
+  threadgroup typename gemv_kernel::acc_type tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  // Update batch offsets
+  if (kDoNCBatch) {
+    in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim);
+    mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim);
+
+    if (kDoAxpby) {
+      bias += elem_to_loc(tid.z, batch_shape, bias_batch_stride, batch_ndim);
+    }
+
+  } else {
+    in_vec += tid.z * vector_batch_stride[0];
+    mat += tid.z * matrix_batch_stride[0];
+
+    if (kDoAxpby) {
+      bias += tid.z * bias_batch_stride[0];
+    }
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      bias,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      alpha,
+      beta,
+      bias_stride,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}
+
+// clang-format off
+#define instantiate_gemv_t_helper(                          \
+    name, itype, bm, bn, sm, sn, tm, tn, nc, axpby)         \
+  instantiate_kernel(                                       \
+    "gemv_t_" #name "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn \
+       "_tm" #tm "_tn" #tn "_nc" #nc "_axpby" #axpby,       \
+  gemv_t, itype, bm, bn, sm, sn, tm, tn, nc, axpby)
+
+#define instantiate_gemv_t(name, itype, bm, bn, sm, sn, tm, tn)        \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 0) \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 1) \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 0) \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 1) // clang-format on
+
+// clang-format off
+#define instantiate_gemv_t_blocks(name, itype) \
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 4, 1) \
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 4,  8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 16, 8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 16, 4, 8, 4, 4) // clang-format on
+
+// clang-format off
+instantiate_gemv_t_blocks(float32, float);
+instantiate_gemv_t_blocks(float16, half);
+instantiate_gemv_t_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_t_blocks(complex64, complex64_t); // clang-format on
+
+template <
+    typename T,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t_gather(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    const device T* bias [[buffer(2)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant float& alpha [[buffer(7)]],
+    const constant float& beta [[buffer(8)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant int64_t* index_batch_strides [[buffer(11)]],
+    const constant int& vector_batch_ndim [[buffer(12)]],
+    const constant int* vector_batch_shape [[buffer(13)]],
+    const constant int64_t* vector_batch_stride [[buffer(14)]],
+    const constant int& matrix_batch_ndim [[buffer(15)]],
+    const constant int* matrix_batch_shape [[buffer(16)]],
+    const constant int64_t* matrix_batch_stride [[buffer(17)]],
+    const constant uint32_t* vec_indices [[buffer(18)]],
+    const constant uint32_t* mat_indices [[buffer(19)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  using gemv_kernel = GEMVTKernel<T, BM, BN, SM, SN, TM, TN, false>;
+  threadgroup typename gemv_kernel::acc_type tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  uint32_t indx_vec;
+  uint32_t indx_mat;
+
+  // Update batch offsets
+  if (batch_ndim > 1) {
+    const constant auto* veci_bstrides = index_batch_strides;
+    const constant auto* mati_bstrides = index_batch_strides + batch_ndim;
+
+    ulong2 batch_offsets = elem_to_loc_broadcast(
+        tid.z, batch_shape, veci_bstrides, mati_bstrides, batch_ndim);
+
+    indx_vec = vec_indices[batch_offsets.x];
+    indx_mat = mat_indices[batch_offsets.y];
+
+  } else {
+    indx_vec = vec_indices[index_batch_strides[0] * tid.z];
+    indx_mat = mat_indices[index_batch_strides[batch_ndim] * tid.z];
+  }
+
+  if (vector_batch_ndim > 1) {
+    in_vec += elem_to_loc(
+        indx_vec, vector_batch_shape, vector_batch_stride, vector_batch_ndim);
+  } else {
+    in_vec += indx_vec * vector_batch_stride[0];
+  }
+
+  if (matrix_batch_ndim > 1) {
+    mat += elem_to_loc(
+        indx_mat, matrix_batch_shape, matrix_batch_stride, matrix_batch_ndim);
+  } else {
+    mat += indx_mat * matrix_batch_stride[0];
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      bias,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      alpha,
+      beta,
+      batch_ndim, // Not used,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}
+
+// clang-format off
+#define instantiate_gemv_t_bs_helper(                  \
+    nm, itype, bm, bn, sm, sn, tm, tn)                 \
+  instantiate_kernel(                                  \
+    "gemv_t_gather_" #nm "_bm" #bm "_bn" #bn "_sm" #sm \
+       "_sn" #sn "_tm" #tm "_tn" #tn,                  \
+  gemv_t_gather, itype, bm, bn, sm, sn, tm, tn)
+
+#define instantiate_gemv_t_bs_blocks(name, itype)              \
+  instantiate_gemv_t_bs_helper(name, itype, 1,  2, 8, 4, 4, 1) \
+  instantiate_gemv_t_bs_helper(name, itype, 1,  2, 8, 4, 4, 4) \
+  instantiate_gemv_t_bs_helper(name, itype, 1,  4, 8, 4, 4, 4) \
+  instantiate_gemv_t_bs_helper(name, itype, 1, 16, 8, 4, 4, 4) \
+  instantiate_gemv_t_bs_helper(name, itype, 1, 16, 4, 8, 4, 4) // clang-format on
+
+// clang-format off
+instantiate_gemv_t_bs_blocks(float32, float);
+instantiate_gemv_t_bs_blocks(float16, half);
+instantiate_gemv_t_bs_blocks(bfloat16, bfloat16_t);
+instantiate_gemv_t_bs_blocks(complex64, complex64_t); // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/gemv_masked.h b/macos/third_party/mlx/mlx/backend/metal/kernels/gemv_masked.h
new file mode 100644
index 0000000..407d14b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/gemv_masked.h
@@ -0,0 +1,827 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace metal;
+
+#define MLX_MTL_CONST static constant constexpr const
+#define MLX_MTL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
+
+struct _NoMask {
+  char x;
+
+  constexpr METAL_FUNC operator bool() {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const threadgroup {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const device {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const constant {
+    return true;
+  }
+};
+
+typedef struct _NoMask nomask_t;
+
+template <typename OutT, typename InT = OutT>
+struct ScaleOp {
+  OutT scale;
+
+  METAL_FUNC OutT apply(InT x) const {
+    return static_cast<OutT>(x) * scale;
+  }
+};
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    typename AccT = float>
+struct GEMVKernel {
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
+
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  static_assert(
+      SN == 8 || SN == 16 || SN == 32,
+      "gemv block must have a width of 8, 16, or 32");
+
+  static_assert(blockN >= blockM, "Masked gemv must have blockN >= blockM");
+
+  MLX_MTL_CONST bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  MLX_MTL_CONST bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  MLX_MTL_CONST bool has_mul_operand_mask =
+      has_operand_mask && !metal::is_same_v<op_mask_t, bool>;
+  MLX_MTL_CONST bool has_mul_output_mask =
+      has_output_mask && !metal::is_same_v<out_mask_t, bool>;
+
+  // - The matrix of size (M = out_vec_size, K = in_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM rows
+  //    and the corresponding scalar from the vector
+  // 2. The thread then multiplies and adds to accumulate its local result for
+  //    the block
+  // 3. At the end, each thread has accumulated results over all blocks across
+  //    the rows. These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated blockM outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid has blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results
+  //     remain zero)
+  //   * The last thread that partially overlaps with the matrix is shifted
+  //     inwards such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BN > 1 ? BN*(blockM + TM) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BN > 1;
+
+  template <typename U = T>
+  static METAL_FUNC void
+  load_unsafe(const device T* src, thread U dst[TN], const int src_offset = 0) {
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tn = 0; tn < TN; tn++) {
+      dst[tn] = static_cast<U>(src[src_offset + tn]);
+    }
+  }
+
+  template <typename U = T>
+  static METAL_FUNC void load_safe(
+      const device T* src,
+      thread U dst[TN],
+      const int src_offset = 0,
+      const int src_size = TN) {
+    if (src_offset + TN <= src_size) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = static_cast<U>(src[src_offset + tn]);
+      }
+    } else { // Edgecase
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = src_offset + tn < src_size
+            ? static_cast<U>(src[src_offset + tn])
+            : U(0);
+      }
+    }
+  }
+
+  static METAL_FUNC void run(
+      const device T* mat [[buffer(0)]],
+      const device T* in_vec [[buffer(1)]],
+      device T* out_vec [[buffer(3)]],
+      const constant int& in_vec_size [[buffer(4)]],
+      const constant int& out_vec_size [[buffer(5)]],
+      const constant int& matrix_ld [[buffer(6)]],
+      const device out_mask_t* out_mask [[buffer(20)]],
+      const device op_mask_t* mat_mask [[buffer(21)]],
+      const device op_mask_t* vec_mask [[buffer(22)]],
+      const constant int* mask_strides [[buffer(23)]],
+      threadgroup AccT* tgp_memory [[threadgroup(0)]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]],
+      uint simd_gid [[simdgroup_index_in_threadgroup]],
+      uint simd_lid [[thread_index_in_simdgroup]]) {
+    // Appease compiler
+    (void)lid;
+
+    // Thread local accumulation results
+    thread AccT result[TM] = {0};
+    thread T inter[TN];
+    thread AccT v_coeff[TN];
+
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
+
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = BN != 1 ? SM * (simd_gid / BN) : int(SM * simd_gid);
+    const int simdN = BN != 1 ? SN * (simd_gid % BN) : 0;
+
+    int bm = (simdM + thrM) * TM;
+    int bn = (simdN + thrN) * TN;
+
+    // Block position
+    int out_row = tid.x * blockM + bm;
+
+    // Exit simdgroup if rows out of bound
+    if (out_row >= out_vec_size)
+      return;
+
+    // Adjust tail simdgroup to ensure in bound reads
+    out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM;
+
+    // Prepare mask offsets
+    const constant int* out_mask_strides = mask_strides;
+    const constant int* mat_mask_strides =
+        mask_strides + (has_output_mask ? 2 : 0);
+    const constant int* vec_mask_strides =
+        mat_mask_strides + (has_operand_mask ? 2 : 0);
+
+    const int m_block_idx = blockN > blockM ? out_row / blockN : int(tid.x);
+
+    const int out_mask_offset =
+        !has_output_mask ? 0 : m_block_idx * out_mask_strides[1];
+
+    int mat_mask_offset =
+        !has_operand_mask ? 0 : m_block_idx * mat_mask_strides[1];
+    int vec_mask_offset = 0;
+    const int mat_mask_step = !has_operand_mask ? 0 : mat_mask_strides[0];
+    const int vec_mask_step = !has_operand_mask ? 0 : vec_mask_strides[1];
+
+    T out_scale{1};
+
+    // Check output mask
+    if (has_output_mask) {
+      auto mask_out = out_mask[out_mask_offset];
+
+      // Write zeros and return if mask is 0
+      if (!mask_out) {
+        if (simdN == 0 && thrN == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tm = 0; tm < TM; tm++) {
+            out_vec[out_row + tm] = T(0.);
+          }
+        }
+
+        return;
+      }
+
+      // Store scalar if multiplicative mask
+      if (has_mul_output_mask) {
+        out_scale = T(mask_out);
+      }
+    }
+
+    // Advance matrix
+    mat += out_row * matrix_ld;
+
+    // Prepare for loop
+    constexpr const uniform<int> loop_stride = make_uniform(blockN);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
+
+    // Loop over in_vec in blocks of blockN
+    for (int i = 0; i < n_iter; ++i) {
+      if (!has_operand_mask ||
+          (bool(mat_mask[mat_mask_offset]) &&
+           bool(vec_mask[vec_mask_offset]))) {
+        T block_scale{1};
+        if (has_mul_operand_mask) {
+          block_scale =
+              T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+        }
+
+        load_unsafe<AccT>(in_vec, v_coeff, bn);
+
+        // Apply scale
+        if (has_mul_operand_mask) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            v_coeff[tn] *= block_scale;
+          }
+        }
+
+        // Per thread work loop
+        int mat_offset = 0;
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          // Load for the row
+          load_unsafe(mat, inter, mat_offset + bn);
+
+          // Accumulate results
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            result[tm] += inter[tn] * v_coeff[tn];
+          }
+
+          mat_offset += matrix_ld;
+        }
+      }
+
+      bn += blockN;
+      mat_mask_offset += mat_mask_step;
+      vec_mask_offset += vec_mask_step;
+    }
+
+    if (leftover > 0) {
+      if (!has_operand_mask ||
+          (bool(mat_mask[mat_mask_offset]) &&
+           bool(vec_mask[vec_mask_offset]))) {
+        T block_scale{1};
+        if (has_mul_operand_mask) {
+          block_scale =
+              T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+        }
+
+        load_safe<AccT>(in_vec, v_coeff, bn, in_size);
+
+        // Apply scale
+        if (has_mul_operand_mask) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            v_coeff[tn] *= block_scale;
+          }
+        }
+
+        // Per thread work loop
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          // Load for the row
+          load_safe(&mat[tm * matrix_ld], inter, bn, in_size);
+
+          // Accumulate results
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            result[tm] += inter[tn] * v_coeff[tn];
+          }
+        }
+      }
+    }
+
+    // Apply out scale
+    if (has_mul_output_mask) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        result[tm] *= out_scale;
+      }
+    }
+
+    // Simdgroup accumulations
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tm = 0; tm < TM; tm++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sn = (SN / 2); sn >= 1; sn >>= 1) {
+        result[tm] += simd_shuffle_down(result[tm], sn);
+      }
+    }
+
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup AccT* tgp_results = tgp_memory + sgN * (blockM + TM) + bm;
+      if (thrN == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          tgp_results[tm] = result[tm];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgN == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgn = 1; sgn < BN; sgn++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tm = 0; tm < TM; tm++) {
+              result[tm] += tgp_results[sgn * (blockM + TM) + tm];
+            }
+          }
+        }
+      }
+    }
+
+    // Write outputs
+    if (simdN == 0 && thrN == 0) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        out_vec[out_row + tm] = static_cast<T>(result[tm]);
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    typename AccT = float>
+struct GEMVTKernel {
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
+
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  MLX_MTL_CONST bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  MLX_MTL_CONST bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  MLX_MTL_CONST bool has_mul_operand_mask =
+      has_operand_mask && !metal::is_same_v<op_mask_t, bool>;
+  MLX_MTL_CONST bool has_mul_output_mask =
+      has_output_mask && !metal::is_same_v<out_mask_t, bool>;
+
+  // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM contiguous rows
+  //    and the corresponding scalar from the vector
+  // 2. The thread then accumulates its local result for the block
+  // 3. At the end, each thread has accumulated results over all blocks across
+  //    the rows. These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated BN * TN outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid has blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results
+  //     remain zero)
+  //   * The last thread that partially overlaps with the matrix is shifted
+  //     inwards such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BM > 1 ? BM*(blockN + TN) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BM > 1;
+
+  static METAL_FUNC void run(
+      const device T* mat [[buffer(0)]],
+      const device T* in_vec [[buffer(1)]],
+      device T* out_vec [[buffer(3)]],
+      const constant int& in_vec_size [[buffer(4)]],
+      const constant int& out_vec_size [[buffer(5)]],
+      const constant int& marix_ld [[buffer(6)]],
+      const device out_mask_t* out_mask [[buffer(20)]],
+      const device op_mask_t* mat_mask [[buffer(21)]],
+      const device op_mask_t* vec_mask [[buffer(22)]],
+      const constant int* mask_strides [[buffer(23)]],
+      threadgroup AccT* tgp_memory [[threadgroup(0)]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]],
+      uint simd_gid [[simdgroup_index_in_threadgroup]],
+      uint simd_lid [[thread_index_in_simdgroup]]) {
+    // Appease compiler
+    (void)lid;
+
+    // Thread local accumulation results
+    AccT result[TN] = {0};
+    T inter[TN];
+    AccT v_coeff[TM];
+
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
+
+    const int sgM = BN != 1 ? (simd_gid / BN) : int(simd_gid);
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = SM * sgM;
+    const int simdN = SN * sgN;
+
+    int cm = (simdM + thrM);
+    int cn = (simdN + thrN);
+
+    int bm = cm * TM;
+    int bn = cn * TN;
+
+    int out_col = tid.x * blockN + bn;
+
+    // Prepare mask offsets
+    const constant int* out_mask_strides = mask_strides;
+    const constant int* mat_mask_strides =
+        out_mask_strides + (has_output_mask ? 2 : 0);
+    const constant int* vec_mask_strides =
+        mat_mask_strides + (has_operand_mask ? 2 : 0);
+
+    const int n_block_idx = blockM > blockN ? out_col / blockM : int(tid.x);
+
+    const int out_mask_offset =
+        !has_output_mask ? 0 : n_block_idx; // * out_mask_strides[0];
+
+    int mat_mask_offset =
+        !has_operand_mask ? 0 : n_block_idx * mat_mask_strides[0];
+    int vec_mask_offset = 0;
+    const int mat_mask_step = !has_operand_mask ? 0 : mat_mask_strides[1];
+    const int vec_mask_step = !has_operand_mask ? 0 : vec_mask_strides[0];
+
+    T out_scale{1};
+
+    // Check output mask
+    if (has_output_mask) {
+      auto mask_out = out_mask[out_mask_offset];
+
+      // Write zeros and return if mask is 0
+      if (!mask_out) {
+        if (cm == 0 && out_col < out_vec_size) {
+          if (out_col + TN <= out_vec_size) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              out_vec[out_col + tn] = T(0.);
+            }
+          } else {
+            for (int tn = 0; tn < TN && (out_col + tn) < out_vec_size; tn++) {
+              out_vec[out_col + tn] = T(0.);
+            }
+          }
+        }
+
+        return;
+      }
+
+      // Store scalar if multiplicative mask
+      if (has_mul_output_mask) {
+        out_scale = T(mask_out);
+      }
+    }
+
+    // Prepare for loop
+    constexpr const uniform<int> loop_stride = make_uniform(blockM);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
+
+    // Edgecase handling
+    if (out_col < out_vec_size) {
+      out_col = (out_col + TN) <= out_vec_size ? out_col : out_vec_size - TN;
+
+      // Per thread accumulation main loop
+      for (int i = 0; i < n_iter; ++i) {
+        // Adding a threadgroup_barrier improves performance slightly
+        // This is possibly it may help exploit cache better
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (!has_operand_mask ||
+            (bool(mat_mask[mat_mask_offset]) &&
+             bool(vec_mask[vec_mask_offset]))) {
+          T block_scale{1};
+          if (has_mul_operand_mask) {
+            block_scale =
+                T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+          }
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tm = 0; tm < TM; tm++) {
+            v_coeff[tm] = static_cast<AccT>(in_vec[bm + tm]);
+          }
+
+          // Apply scale
+          if (has_mul_operand_mask) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tm = 0; tm < TM; tm++) {
+              v_coeff[tm] *= block_scale;
+            }
+          }
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tm = 0; tm < TM; tm++) {
+            for (int tn = 0; tn < TN; tn++) {
+              inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
+            }
+            for (int tn = 0; tn < TN; tn++) {
+              result[tn] += v_coeff[tm] * inter[tn];
+            }
+          }
+        }
+
+        bm += blockM;
+        mat_mask_offset += mat_mask_step;
+        vec_mask_offset += vec_mask_step;
+      }
+
+      if (leftover > 0) {
+        if (!has_operand_mask ||
+            (bool(mat_mask[mat_mask_offset]) &&
+             bool(vec_mask[vec_mask_offset]))) {
+          T block_scale{1};
+          if (has_mul_operand_mask) {
+            block_scale =
+                T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+          }
+
+          for (int tm = 0; tm < TM && bm + tm < in_vec_size; tm++) {
+            v_coeff[tm] = static_cast<AccT>(in_vec[bm + tm]);
+
+            if (has_mul_operand_mask) {
+              v_coeff[tm] *= block_scale;
+            }
+
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
+            }
+
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              result[tn] += v_coeff[tm] * inter[tn];
+            }
+          }
+        }
+      }
+    }
+
+    // Apply out scale
+    if (has_mul_output_mask) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        result[tn] *= out_scale;
+      }
+    }
+
+    // Simdgroup accumulations
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tn = 0; tn < TN; tn++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sm = (SM / 2); sm >= 1; sm >>= 1) {
+        result[tn] += simd_shuffle_down(result[tn], SN * sm);
+      }
+    }
+
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup AccT* tgp_results = tgp_memory + sgM * (blockN + TN) + bn;
+      if (thrM == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          tgp_results[tn] = result[tn];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgM == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgm = 1; sgm < BM; sgm++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              result[tn] += tgp_results[sgm * (blockN + TN) + tn];
+            }
+          }
+        }
+      }
+    }
+
+    // Threadgroup accumulation and writing out results
+    if (cm == 0 && out_col < out_vec_size) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int j = 0; j < TN; j++) {
+        out_vec[out_col + j] = static_cast<T>(result[j]);
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Matrix vector multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoNCBatch> /* Batch ndim > 1 */
+[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_masked(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant int64_t* vector_batch_stride [[buffer(11)]],
+    const constant int64_t* matrix_batch_stride [[buffer(12)]],
+    const device out_mask_t* out_mask [[buffer(20)]],
+    const device op_mask_t* mat_mask [[buffer(21)]],
+    const device op_mask_t* vec_mask [[buffer(22)]],
+    const constant int* mask_strides [[buffer(23)]],
+    const constant int64_t* mask_batch_strides [[buffer(24)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  using gemv_kernel =
+      GEMVKernel<T, out_mask_t, op_mask_t, BM, BN, SM, SN, TM, TN>;
+  threadgroup float tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  constexpr bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  constexpr bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  // Update batch offsets
+  if (kDoNCBatch) {
+    in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim);
+    mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim);
+
+    if (has_output_mask) {
+      out_mask +=
+          elem_to_loc(tid.z, batch_shape, mask_batch_strides, batch_ndim);
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      const constant auto* mask_strides_mat = mask_batch_strides;
+      const constant auto* mask_strides_vec = mask_strides_mat + batch_ndim;
+
+      ulong2 batch_offsets = elem_to_loc_broadcast(
+          tid.z, batch_shape, mask_strides_mat, mask_strides_vec, batch_ndim);
+
+      mat_mask += batch_offsets.x;
+      vec_mask += batch_offsets.y;
+    }
+
+  } else {
+    in_vec += tid.z * vector_batch_stride[0];
+    mat += tid.z * matrix_batch_stride[0];
+
+    if (has_output_mask) {
+      out_mask += tid.z * mask_batch_strides[0];
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      mat_mask += tid.z * mask_batch_strides[0];
+      vec_mask += tid.z * mask_batch_strides[batch_ndim];
+    }
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      out_mask,
+      mat_mask,
+      vec_mask,
+      mask_strides,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoNCBatch> /* Batch ndim > 1 */
+[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t_masked(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant int64_t* vector_batch_stride [[buffer(11)]],
+    const constant int64_t* matrix_batch_stride [[buffer(12)]],
+    const device out_mask_t* out_mask [[buffer(20)]],
+    const device op_mask_t* mat_mask [[buffer(21)]],
+    const device op_mask_t* vec_mask [[buffer(22)]],
+    const constant int* mask_strides [[buffer(23)]],
+    const constant int64_t* mask_batch_strides [[buffer(24)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  using gemv_kernel =
+      GEMVTKernel<T, out_mask_t, op_mask_t, BM, BN, SM, SN, TM, TN>;
+  threadgroup float tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  constexpr bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  constexpr bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  // Update batch offsets
+  if (kDoNCBatch) {
+    in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim);
+    mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim);
+
+    if (has_output_mask) {
+      out_mask +=
+          elem_to_loc(tid.z, batch_shape, mask_batch_strides, batch_ndim);
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      const constant auto* mask_strides_mat = mask_batch_strides;
+      const constant auto* mask_strides_vec = mask_strides_mat + batch_ndim;
+
+      ulong2 batch_offsets = elem_to_loc_broadcast(
+          tid.z, batch_shape, mask_strides_mat, mask_strides_vec, batch_ndim);
+
+      mat_mask += batch_offsets.x;
+      vec_mask += batch_offsets.y;
+    }
+
+  } else {
+    in_vec += tid.z * vector_batch_stride[0];
+    mat += tid.z * matrix_batch_stride[0];
+
+    if (has_output_mask) {
+      out_mask += tid.z * mask_batch_strides[0];
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      mat_mask += tid.z * mask_batch_strides[0];
+      vec_mask += tid.z * mask_batch_strides[batch_ndim];
+    }
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      out_mask,
+      mat_mask,
+      vec_mask,
+      mask_strides,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/gemv_masked.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/gemv_masked.metal
new file mode 100644
index 0000000..394250e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/gemv_masked.metal
@@ -0,0 +1,76 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// clang-format off
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/gemv_masked.h"
+
+#define instantiate_gemv_helper(                                           \
+    outm_n, outm_t, opm_n, opm_t, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_kernel(                                                      \
+    "gemv_outmask_" #outm_n "_opmask_" #opm_n "_" #name                    \
+      "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn "_tm" #tm                    \
+      "_tn" #tn "_nc" #nc,                                                 \
+  gemv_masked, itype, outm_t, opm_t, bm, bn, sm, sn, tm, tn, nc)
+
+#define instantiate_gemv_base(name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(bool_, bool, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(name, itype, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(bool_, bool, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(name, itype, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(nomask, nomask_t, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(nomask, nomask_t, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(bool_, bool, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(name, itype, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc)
+
+#define instantiate_gemv(name, itype, bm, bn, sm, sn, tm, tn)   \
+  instantiate_gemv_base(name, itype, bm, bn, sm, sn, tm, tn, 0) \
+  instantiate_gemv_base(name, itype, bm, bn, sm, sn, tm, tn, 1)
+
+#define instantiate_gemv_blocks(name, itype) \
+  instantiate_gemv(name, itype, 2, 1, 4,  8, 1, 4) \
+  instantiate_gemv(name, itype, 2, 1, 4,  8, 4, 4) \
+  instantiate_gemv(name, itype, 2, 1, 2, 16, 1, 4) \
+  instantiate_gemv(name, itype, 2, 1, 2, 16, 4, 4) \
+  instantiate_gemv(name, itype, 4, 1, 2, 16, 4, 4)
+
+instantiate_gemv_blocks(float32, float);
+instantiate_gemv_blocks(float16, half);
+instantiate_gemv_blocks(bfloat16, bfloat16_t);
+
+#define instantiate_gemv_t_helper(                                           \
+    outm_n, outm_t, opm_n, opm_t, name, itype, bm, bn, sm, sn, tm, tn, nc)   \
+  instantiate_kernel(                                                        \
+    "gemv_t_outmask_" #outm_n "_opmask_" #opm_n "_" #name                    \
+      "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn "_tm" #tm                      \
+      "_tn" #tn "_nc" #nc,                                                   \
+  gemv_t_masked, itype, outm_t, opm_t, bm, bn, sm, sn, tm, tn, nc)
+
+#define instantiate_gemv_t_base(name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(bool_, bool, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(name, itype, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(bool_, bool, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(name, itype, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(nomask, nomask_t, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(nomask, nomask_t, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(bool_, bool, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(name, itype, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc)
+
+#define instantiate_gemv_t(name, itype, bm, bn, sm, sn, tm, tn)   \
+  instantiate_gemv_t_base(name, itype, bm, bn, sm, sn, tm, tn, 0) \
+  instantiate_gemv_t_base(name, itype, bm, bn, sm, sn, tm, tn, 1)
+
+#define instantiate_gemv_t_blocks(name, itype) \
+  instantiate_gemv_t(name, itype, 1, 1,  8, 4, 4, 1) \
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 1,  8, 4, 8, 1) \
+  instantiate_gemv_t(name, itype, 1, 1,  8, 4, 8, 4) \
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 8, 4) \
+  instantiate_gemv_t(name, itype, 1, 4,  8, 4, 8, 4)
+
+instantiate_gemv_t_blocks(float32, float);
+instantiate_gemv_t_blocks(float16, half);
+instantiate_gemv_t_blocks(bfloat16, bfloat16_t); // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/hadamard.h b/macos/third_party/mlx/mlx/backend/metal/kernels/hadamard.h
new file mode 100644
index 0000000..9f2311c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/hadamard.h
@@ -0,0 +1,182 @@
+// Copyright © 2024 Apple Inc.
+#include <metal_common>
+#include <metal_compute>
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+using namespace metal;
+
+// Thread local Hadamard transform for 2^R
+template <short R>
+METAL_FUNC void radix_func(thread float* x) {
+  constexpr short logR = __builtin_ctz(R);
+  short h = 1;
+  STEEL_PRAGMA_UNROLL
+  for (short s = 0; s < logR; s++) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < R / 2; i++) {
+      short k = i & (h - 1);
+      short j = ((i - k) << 1) + k;
+      float a = x[j];
+      float b = x[j + h];
+      x[j] = a + b;
+      x[j + h] = a - b;
+    }
+    h <<= 1;
+  }
+}
+
+template <typename T, int N, int max_radix, int read_width, int stride = 1>
+[[kernel]] void hadamard_n(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const float& scale,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Compute a Hadamard transform of size N = 2^k
+  //
+  // Equivalent to:
+  //    from scipy.linalg import hadamard
+  //    y = hadamard(len(x)) @ x
+
+  constexpr short num_threads = N / max_radix;
+  constexpr short logN = __builtin_ctz(N);
+  constexpr short logR = __builtin_ctz(max_radix);
+  constexpr short num_steps = logN / logR;
+  constexpr short logFinal = logN % logR;
+  constexpr short final_radix = 1 << (logFinal);
+
+  int batch_idx = elem.y * N * stride + elem.z;
+  short i = elem.x;
+
+  threadgroup T buf[N];
+
+  // Read values from device
+  if (stride == 1) {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix / read_width; j++) {
+      short index = j * read_width * num_threads + i * read_width;
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < read_width; r++) {
+        buf[index + r] = in[batch_idx + index + r];
+      }
+    }
+  } else {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix; j++) {
+      buf[j * num_threads + i] = in[batch_idx + (j * num_threads + i) * stride];
+    }
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  float x[max_radix];
+  short h = 1;
+
+  STEEL_PRAGMA_UNROLL
+  for (short s = 0; s < num_steps; s++) {
+    short k = i & (h - 1);
+    short j = ((i - k) << logR) + k;
+
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < max_radix; r++) {
+      x[r] = buf[j + h * r];
+    }
+
+    radix_func<max_radix>(x);
+
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < max_radix; r++) {
+      buf[j + h * r] = T(x[r]);
+    }
+
+    h <<= logR;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // Do the final radix
+  // e.g. max_radix = 16
+  //      N = 1024 = 16 * 16 * 4
+  if (final_radix > 1) {
+    // Each thread does multiple butterflies
+    STEEL_PRAGMA_UNROLL
+    for (int t = 0; t < max_radix / final_radix; t++) {
+      short index = i + t * num_threads;
+      short k = index & (h - 1);
+      short j = ((index - k) << logFinal) + k;
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < final_radix; r++) {
+        x[r] = buf[j + h * r];
+      }
+
+      radix_func<final_radix>(x);
+
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < final_radix; r++) {
+        buf[j + h * r] = T(x[r]);
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // Write values to device
+  if (stride == 1) {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix / read_width; j++) {
+      short index = j * read_width * num_threads + i * read_width;
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < read_width; r++) {
+        out[batch_idx + index + r] = T(buf[index + r] * scale);
+      }
+    }
+  } else {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < max_radix; j++) {
+      out[batch_idx + (j * num_threads + i) * stride] =
+          buf[j * num_threads + i];
+    }
+  }
+}
+
+template <typename T, int N, int M, int read_width>
+[[kernel]] void hadamard_m(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const float& scale,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Compute a Hadamard transform of size M
+  // using a naive O(M^2) codelet.
+  //
+  // This kernel is the second stage in the computation
+  // of a Hadamard transform of size M*N where N = 2^k.
+
+  int index = elem.x * grid.y + elem.y;
+  short i = index % (N / read_width);
+  int batch_idx = index / (N / read_width) * M * N;
+
+  float x[read_width][M];
+  STEEL_PRAGMA_UNROLL
+  for (short c = 0; c < M; c++) {
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      x[r][c] = in[batch_idx + c * N + i * read_width + r];
+    }
+  }
+
+  STEEL_PRAGMA_UNROLL
+  for (short r = 0; r < read_width; r++) {
+    // This function is JIT compiled for M
+    // using the Hadamard matrix strings in `metal/hadamard.cpp`
+    hadamard_radix_m(x[r]);
+  }
+
+  // Write back to device
+  STEEL_PRAGMA_UNROLL
+  for (short c = 0; c < M; c++) {
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      out[batch_idx + c * N + i * read_width + r] = T(x[r][c] * scale);
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather.h b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather.h
new file mode 100644
index 0000000..8b93c01
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather.h
@@ -0,0 +1,51 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <typename T, typename IdxT, int NIDX, int IDX_NDIM, typename LocT>
+METAL_FUNC void gather_impl(
+    const device T* src [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const constant int* src_shape [[buffer(2)]],
+    const constant int64_t* src_strides [[buffer(3)]],
+    const constant size_t& src_ndim [[buffer(4)]],
+    const constant int* slice_sizes [[buffer(5)]],
+    const constant int* axes [[buffer(6)]],
+    const thread Indices<IdxT, NIDX>& indices,
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  LocT src_idx = 0;
+  for (int i = 0; i < NIDX; ++i) {
+    LocT idx_loc;
+    if (IDX_NDIM == 0) {
+      idx_loc = 0;
+    } else if (IDX_NDIM == 1) {
+      idx_loc = index.x * static_cast<LocT>(indices.strides[indices.ndim * i]);
+    } else {
+      idx_loc = index.x * static_cast<LocT>(indices.strides[indices.ndim * i]);
+      idx_loc += indices.row_contiguous[i]
+          ? index.y
+          : elem_to_loc<LocT>(
+                index.y,
+                &indices.shapes[indices.ndim * i + 1],
+                &indices.strides[indices.ndim * i + 1],
+                indices.ndim - 1);
+    }
+    auto ax = axes[i];
+    auto idx_val = offset_neg_idx(indices.buffers[i][idx_loc], src_shape[ax]);
+    src_idx += static_cast<LocT>(idx_val) * static_cast<LocT>(src_strides[ax]);
+  }
+
+  auto src_offset =
+      elem_to_loc<LocT>(index.z, slice_sizes, src_strides, src_ndim);
+
+  LocT out_idx = index.z;
+  if (IDX_NDIM == 1) {
+    out_idx += static_cast<LocT>(grid_dim.z) * index.x;
+  } else if (IDX_NDIM >= 2) {
+    out_idx += grid_dim.z * (index.x * static_cast<LocT>(grid_dim.y) + index.y);
+  }
+  out[out_idx] = src[src_offset + src_idx];
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather_axis.h b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather_axis.h
new file mode 100644
index 0000000..bf490ad
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather_axis.h
@@ -0,0 +1,44 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+template <typename T, typename IdxT, typename LocT, bool SrcC, bool IdxC>
+[[kernel]] void gather_axis(
+    const device T* src [[buffer(0)]],
+    const device IdxT* indices [[buffer(1)]],
+    device T* out [[buffer(2)]],
+    const constant int* shape [[buffer(3)]],
+    const constant int64_t* src_strides [[buffer(4)]],
+    const constant int64_t* idx_strides [[buffer(5)]],
+    const constant size_t& ndim [[buffer(6)]],
+    const constant int& axis [[buffer(7)]],
+    const constant int& axis_size [[buffer(8)]],
+    const constant size_t& src_ax_stride [[buffer(9)]],
+    const constant size_t& idx_ax_stride [[buffer(10)]],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  LocT elem_idx = index.z * static_cast<LocT>(grid_dim.x);
+  LocT out_idx = elem_idx * grid_dim.y + index.x;
+
+  LocT idx_loc = index.y * static_cast<LocT>(idx_ax_stride);
+  if (IdxC) {
+    idx_loc += out_idx;
+  } else {
+    idx_loc += elem_to_loc<LocT>(elem_idx + index.x, shape, idx_strides, ndim);
+  }
+
+  auto idx_val = indices[idx_loc];
+  if (is_signed_v<IdxT>) {
+    idx_val = (idx_val < 0) ? idx_val + axis_size : idx_val;
+  }
+
+  LocT src_idx = idx_val * static_cast<LocT>(src_ax_stride);
+  if (SrcC) {
+    src_idx += elem_idx * axis_size + index.x;
+  } else {
+    src_idx += elem_to_loc<LocT>(elem_idx + index.x, shape, src_strides, ndim);
+  }
+
+  out_idx += index.y * static_cast<LocT>(grid_dim.x);
+  out[out_idx] = src[src_idx];
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather_front.h b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather_front.h
new file mode 100644
index 0000000..1389e4c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/gather_front.h
@@ -0,0 +1,24 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <typename T, typename IdxT, typename LocT, int N>
+[[kernel]] void gather_front(
+    const device T* src,
+    const device IdxT* indices,
+    device T* out,
+    const constant int64_t& stride,
+    const constant int& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto idx = offset_neg_idx(indices[index.y], size);
+  LocT src_idx = static_cast<LocT>(stride) * idx;
+  LocT out_idx = static_cast<LocT>(stride) * index.y;
+
+  int s_idx = N * index.x;
+  for (int i = 0; i < N && s_idx < stride; ++i, ++s_idx) {
+    out[out_idx + s_idx] = src[src_idx + s_idx];
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/indexing.h b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/indexing.h
new file mode 100644
index 0000000..2a4b4f9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/indexing.h
@@ -0,0 +1,23 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+template <typename IdxT, int NIDX>
+struct Indices {
+  const array<const device IdxT*, NIDX> buffers;
+  const constant int* shapes;
+  const constant int64_t* strides;
+  const constant bool* row_contiguous;
+  const int ndim;
+};
+
+template <typename IdxT>
+METAL_FUNC size_t offset_neg_idx(IdxT idx, int size) {
+  if (is_unsigned_v<IdxT>) {
+    return idx;
+  } else {
+    return (idx < 0) ? idx + size : idx;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/masked_scatter.h b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/masked_scatter.h
new file mode 100644
index 0000000..2ba5474
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/masked_scatter.h
@@ -0,0 +1,41 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+constant mlx::os_log logger("mlx", "masked_assign");
+
+template <typename T, bool src_contiguous>
+[[kernel]] void masked_assign_impl(
+    const device bool* mask [[buffer(0)]],
+    const device uint* scatter_offsets [[buffer(1)]],
+    const device T* src [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    const constant int* src_shapes [[buffer(4)]],
+    const constant int64_t* src_strides [[buffer(5)]],
+    const constant int& src_ndim [[buffer(6)]],
+    const constant int64_t& src_batch_size [[buffer(7)]],
+    const constant int64_t& mask_batch_size [[buffer(8)]],
+    uint idx [[thread_position_in_grid]]) {
+  const bool mask_value = mask[idx];
+  if (!mask_value) {
+    return;
+  }
+
+  const uint src_index = scatter_offsets[idx];
+  if (src_index >= src_batch_size) {
+    logger.log_debug("Out of bound read from src");
+    return;
+  }
+
+  const uint batch_idx = idx / mask_batch_size;
+
+  if (src_contiguous) {
+    out[idx] = src[batch_idx * src_batch_size + src_index];
+  } else {
+    out[idx] = src[elem_to_loc<uint>(
+        batch_idx * src_batch_size + src_index,
+        src_shapes,
+        src_strides,
+        src_ndim)];
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/scatter.h b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/scatter.h
new file mode 100644
index 0000000..f0217b3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/scatter.h
@@ -0,0 +1,59 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <
+    typename T,
+    typename IdxT,
+    typename Op,
+    int NIDX,
+    bool UPD_ROW_CONTIG,
+    int NWORK,
+    typename LocT>
+METAL_FUNC void scatter_impl(
+    const device T* updates,
+    device mlx_atomic<T>* out,
+    const constant int* upd_shape,
+    const constant int64_t* upd_strides,
+    const constant size_t& upd_ndim,
+    const constant size_t& upd_size,
+    const constant int* out_shape,
+    const constant int64_t* out_strides,
+    const constant size_t& out_ndim,
+    const constant int* axes,
+    const constant size_t& idx_size,
+    const thread Indices<IdxT, NIDX>& indices,
+    uint2 gid [[thread_position_in_grid]]) {
+  Op op;
+
+  auto ind_idx = gid.y * NWORK;
+  LocT out_offset = 0;
+  if (upd_size > 1) {
+    out_offset = elem_to_loc<LocT>(
+        gid.x, upd_shape + indices.ndim, out_strides, out_ndim);
+  }
+
+  for (int j = 0; j < NWORK && ind_idx < idx_size; ++j, ind_idx++) {
+    LocT out_idx = out_offset;
+    for (int i = 0; i < NIDX; ++i) {
+      auto idx_loc = indices.row_contiguous[i]
+          ? ind_idx
+          : elem_to_loc<LocT>(
+                ind_idx,
+                &indices.shapes[indices.ndim * i],
+                &indices.strides[indices.ndim * i],
+                indices.ndim);
+      auto ax = axes[i];
+      auto idx_val = offset_neg_idx(indices.buffers[i][idx_loc], out_shape[ax]);
+      out_idx +=
+          static_cast<LocT>(idx_val) * static_cast<LocT>(out_strides[ax]);
+    }
+    auto upd_idx = ind_idx * static_cast<LocT>(upd_size) + gid.x;
+    if constexpr (!UPD_ROW_CONTIG) {
+      upd_idx = elem_to_loc<LocT>(upd_idx, upd_shape, upd_strides, upd_ndim);
+    }
+    op.atomic_update(out, updates[upd_idx], out_idx);
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/scatter_axis.h b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/scatter_axis.h
new file mode 100644
index 0000000..73fd7ab
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/indexing/scatter_axis.h
@@ -0,0 +1,52 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+template <
+    typename T,
+    typename IdxT,
+    typename LocT,
+    typename Op,
+    bool UpdC,
+    bool IdxC>
+[[kernel]] void scatter_axis(
+    const device T* upd [[buffer(0)]],
+    const device IdxT* indices [[buffer(1)]],
+    device mlx_atomic<T>* out [[buffer(2)]],
+    const constant int* shape [[buffer(3)]],
+    const constant int64_t* upd_strides [[buffer(4)]],
+    const constant int64_t* idx_strides [[buffer(5)]],
+    const constant size_t& ndim [[buffer(6)]],
+    const constant int& axis [[buffer(7)]],
+    const constant int& out_axis_size [[buffer(8)]],
+    const constant size_t& upd_ax_stride [[buffer(9)]],
+    const constant size_t& idx_ax_stride [[buffer(10)]],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  Op op;
+
+  LocT elem_idx = index.z * static_cast<LocT>(grid_dim.x);
+
+  LocT idx_loc = index.y * static_cast<LocT>(idx_ax_stride);
+  if (IdxC) {
+    idx_loc += elem_idx * grid_dim.y + index.x;
+  } else {
+    idx_loc += elem_to_loc<LocT>(elem_idx + index.x, shape, idx_strides, ndim);
+  }
+
+  auto idx_val = indices[idx_loc];
+  if (is_signed_v<IdxT>) {
+    idx_val = (idx_val < 0) ? idx_val + out_axis_size : idx_val;
+  }
+
+  LocT upd_idx = index.y * static_cast<LocT>(upd_ax_stride);
+  if (UpdC) {
+    upd_idx += elem_idx * grid_dim.y + index.x;
+  } else {
+    upd_idx += elem_to_loc<LocT>(elem_idx + index.x, shape, upd_strides, ndim);
+  }
+
+  LocT out_idx = elem_idx * static_cast<LocT>(out_axis_size) +
+      idx_val * grid_dim.x + index.x;
+  op.atomic_update(out, upd[upd_idx], out_idx);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/layer_norm.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/layer_norm.metal
new file mode 100644
index 0000000..ea77b53
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/layer_norm.metal
@@ -0,0 +1,433 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_common>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+constant bool has_w [[function_constant(20)]];
+
+template <int N = 1>
+inline void initialize_buffer(
+    threadgroup float* xs,
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  if (simd_group_id == 0) {
+    for (int i = 0; i < N; i++) {
+      xs[N * simd_lane_id + i] = 0;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+}
+
+template <int N = 1>
+inline void threadgroup_sum(
+    thread float* x,
+    threadgroup float* xs,
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  for (int i = 0; i < N; i++) {
+    x[i] = simd_sum(x[i]);
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_lane_id == 0) {
+    for (int i = 0; i < N; i++) {
+      xs[N * simd_group_id + i] = x[i];
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  for (int i = 0; i < N; i++) {
+    x[i] = xs[N * simd_lane_id + i];
+    x[i] = simd_sum(x[i]);
+  }
+}
+
+template <typename T, int N_READS = 8>
+[[kernel]] void layer_norm_single_row(
+    const device T* x,
+    const device T* w,
+    const device T* b,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    constant uint& b_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+
+  // Initialize the registers and threadgroup memory
+  float thread_x[N_READS] = {0};
+  threadgroup float local_buffer[SIMD_SIZE] = {0};
+  initialize_buffer(local_buffer, simd_lane_id, simd_group_id);
+
+  // Advance the pointers
+  x += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  b += b_stride * lid * N_READS;
+  out += gid * size_t(axis_size) + lid * N_READS;
+
+  // Compute some variables for reading writing etc
+  const bool safe = lid * N_READS + N_READS <= axis_size;
+  const int n = axis_size - lid * N_READS;
+
+  // Read the inputs
+  if (safe) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] = x[i];
+    }
+  } else {
+    for (int i = 0; i < n; i++) {
+      thread_x[i] = x[i];
+    }
+  }
+
+  // Compute the mean
+  float mean = 0;
+  for (int i = 0; i < N_READS; i++) {
+    mean += thread_x[i];
+  }
+  threadgroup_sum(&mean, local_buffer, simd_lane_id, simd_group_id);
+  mean /= axis_size;
+
+  // Compute the normalizer
+  float normalizer = 0;
+  if (!safe) {
+    for (int i = n; i < N_READS; i++) {
+      thread_x[i] = mean;
+    }
+  }
+  for (int i = 0; i < N_READS; i++) {
+    thread_x[i] -= mean;
+    normalizer += thread_x[i] * thread_x[i];
+  }
+  threadgroup_sum(&normalizer, local_buffer, simd_lane_id, simd_group_id);
+  normalizer = metal::precise::rsqrt(normalizer / axis_size + eps);
+
+  // Write the outputs
+  if (safe) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] *= normalizer;
+      out[i] = w[w_stride * i] * static_cast<T>(thread_x[i]) + b[b_stride * i];
+    }
+  } else {
+    for (int i = 0; i < n; i++) {
+      thread_x[i] *= normalizer;
+      out[i] = w[w_stride * i] * static_cast<T>(thread_x[i]) + b[b_stride * i];
+    }
+  }
+}
+
+template <typename T, int N_READS = 4>
+[[kernel]] void layer_norm_looped(
+    const device T* x,
+    const device T* w,
+    const device T* b,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    constant uint& b_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup float local_buffer[SIMD_SIZE];
+  initialize_buffer(local_buffer, simd_lane_id, simd_group_id);
+
+  x += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  b += b_stride * lid * N_READS;
+
+  // Compute the mean
+  float mean = 0;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        mean += x[i + r];
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          mean += x[i + r];
+        }
+      }
+    }
+  }
+  threadgroup_sum(&mean, local_buffer, simd_lane_id, simd_group_id);
+  mean /= axis_size;
+
+  // Compute the normalizer
+  float normalizer = 0;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float t = x[i + r] - mean;
+        normalizer += t * t;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float t = x[i + r] - mean;
+          normalizer += t * t;
+        }
+      }
+    }
+  }
+  threadgroup_sum(&normalizer, local_buffer, simd_lane_id, simd_group_id);
+  normalizer = metal::precise::rsqrt(normalizer / axis_size + eps);
+
+  // Write the outputs
+  out += gid * size_t(axis_size) + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = (x[r + i] - mean) * normalizer;
+        out[r + i] =
+            w[w_stride * (i + r)] * static_cast<T>(xi) + b[b_stride * (i + r)];
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = (x[r + i] - mean) * normalizer;
+          out[r + i] = w[w_stride * (i + r)] * static_cast<T>(xi) +
+              b[b_stride * (i + r)];
+        }
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = 8>
+[[kernel]] void vjp_layer_norm_single_row(
+    const device T* x,
+    const device T* w,
+    const device T* g,
+    device T* gx,
+    device T* gw,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+
+  // Advance the input pointers
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+
+  // Initialize the registers and threadgroup memory
+  float thread_x[N_READS] = {0};
+  float thread_w[N_READS] = {0};
+  float thread_g[N_READS] = {0};
+  threadgroup float local_buffer[3 * SIMD_SIZE];
+  initialize_buffer<3>(local_buffer, simd_lane_id, simd_group_id);
+
+  // Compute some variables for reading writing etc
+  const bool safe = lid * N_READS + N_READS <= axis_size;
+  const int n = axis_size - lid * N_READS;
+
+  // Read the inputs
+  if (safe) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] = x[i];
+      thread_g[i] = g[i];
+      thread_w[i] = w[i * w_stride];
+    }
+  } else {
+    for (int i = 0; i < n; i++) {
+      thread_x[i] = x[i];
+      thread_g[i] = g[i];
+      thread_w[i] = w[i * w_stride];
+    }
+  }
+
+  // Compute the mean
+  float mean = 0;
+  for (int i = 0; i < N_READS; i++) {
+    mean += thread_x[i];
+  }
+  threadgroup_sum(&mean, local_buffer, simd_lane_id, simd_group_id);
+  mean /= axis_size;
+
+  // Compute the neccesary scaling factors using the mean
+  if (!safe) {
+    for (int i = n; i < N_READS; i++) {
+      thread_x[i] = mean;
+    }
+  }
+  float factors[3] = {0};
+  constexpr int meanwg = 0;
+  constexpr int meanwgxc = 1;
+  constexpr int normalizer2 = 2;
+  for (int i = 0; i < N_READS; i++) {
+    thread_x[i] -= mean;
+    factors[meanwg] += thread_w[i] * thread_g[i];
+    factors[meanwgxc] += thread_w[i] * thread_g[i] * thread_x[i];
+    factors[normalizer2] += thread_x[i] * thread_x[i];
+  }
+  threadgroup_sum<3>(factors, local_buffer, simd_lane_id, simd_group_id);
+  factors[meanwg] /= axis_size;
+  factors[meanwgxc] /= axis_size;
+  factors[normalizer2] = 1 / (factors[normalizer2] / axis_size + eps);
+  float normalizer = metal::precise::sqrt(factors[normalizer2]);
+
+  // Write the outputs
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
+  if (safe) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] *= normalizer;
+      gx[i] = static_cast<T>(
+          normalizer * (thread_w[i] * thread_g[i] - factors[meanwg]) -
+          thread_x[i] * factors[meanwgxc] * factors[normalizer2]);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+      }
+    }
+  } else {
+    for (int i = 0; i < n; i++) {
+      thread_x[i] *= normalizer;
+      gx[i] = static_cast<T>(
+          normalizer * (thread_w[i] * thread_g[i] - factors[meanwg]) -
+          thread_x[i] * factors[meanwgxc] * factors[normalizer2]);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i]);
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = 4>
+[[kernel]] void vjp_layer_norm_looped(
+    const device T* x,
+    const device T* w,
+    const device T* g,
+    device T* gx,
+    device T* gw,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+
+  // Advance the input pointers
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+
+  threadgroup float local_buffer[3 * SIMD_SIZE];
+  initialize_buffer<3>(local_buffer, simd_lane_id, simd_group_id);
+
+  // Compute the mean
+  float mean = 0;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        mean += x[i + r];
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          mean += x[i + r];
+        }
+      }
+    }
+  }
+  threadgroup_sum(&mean, local_buffer, simd_lane_id, simd_group_id);
+  mean /= axis_size;
+
+  // Compute the neccesary scaling factors using the mean
+  float factors[3] = {0};
+  constexpr int meanwg = 0;
+  constexpr int meanwgxc = 1;
+  constexpr int normalizer2 = 2;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float t = x[i + r] - mean;
+        float wi = w[(i + r) * w_stride];
+        float gi = g[i + r];
+        float wg = wi * gi;
+        factors[meanwg] += wg;
+        factors[meanwgxc] += wg * t;
+        factors[normalizer2] += t * t;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float t = x[i + r] - mean;
+          float wi = w[(i + r) * w_stride];
+          float gi = g[i + r];
+          float wg = wi * gi;
+          factors[meanwg] += wg;
+          factors[meanwgxc] += wg * t;
+          factors[normalizer2] += t * t;
+        }
+      }
+    }
+  }
+  threadgroup_sum<3>(factors, local_buffer, simd_lane_id, simd_group_id);
+  factors[meanwg] /= axis_size;
+  factors[meanwgxc] /= axis_size;
+  factors[normalizer2] = 1 / (factors[normalizer2] / axis_size + eps);
+  float normalizer = metal::precise::sqrt(factors[normalizer2]);
+
+  // Write the outputs
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = (x[i + r] - mean) * normalizer;
+        float wi = w[(i + r) * w_stride];
+        float gi = g[i + r];
+        gx[i + r] = static_cast<T>(
+            normalizer * (wi * gi - factors[meanwg]) -
+            xi * factors[meanwgxc] * factors[normalizer2]);
+        if (has_w) {
+          gw[i + r] = static_cast<T>(gi * xi);
+        }
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = (x[i + r] - mean) * normalizer;
+          float wi = w[(i + r) * w_stride];
+          float gi = g[i + r];
+          gx[i + r] = static_cast<T>(
+              normalizer * (wi * gi - factors[meanwg]) -
+              xi * factors[meanwgxc] * factors[normalizer2]);
+          if (has_w) {
+            gw[i + r] = static_cast<T>(gi * xi);
+          }
+        }
+      }
+    }
+  }
+}
+
+// clang-format off
+#define instantiate_layer_norm(name, itype)                                       \
+  instantiate_kernel("layer_norm" #name, layer_norm_single_row, itype)            \
+  instantiate_kernel("vjp_layer_norm" #name, vjp_layer_norm_single_row, itype)    \
+  instantiate_kernel("layer_norm_looped" #name, layer_norm_looped, itype)         \
+  instantiate_kernel("vjp_layer_norm_looped" #name, vjp_layer_norm_looped, itype)
+
+instantiate_layer_norm(float32, float)
+instantiate_layer_norm(float16, half)
+instantiate_layer_norm(bfloat16, bfloat16_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/logging.h b/macos/third_party/mlx/mlx/backend/metal/kernels/logging.h
new file mode 100644
index 0000000..7b3ee04
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/logging.h
@@ -0,0 +1,26 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#if defined(__METAL_VERSION__) && (__METAL_VERSION__ >= 320)
+#include <metal_logging>
+
+namespace mlx {
+using os_log = metal::os_log;
+} // namespace mlx
+
+#else
+
+namespace mlx {
+struct os_log {
+  constexpr os_log(constant char*, constant char*) constant {}
+
+  template <typename... Args>
+  void log_debug(constant char*, Args...) const {}
+
+  template <typename... Args>
+  void log_debug(constant char*, Args...) const constant {}
+};
+} // namespace mlx
+
+#endif
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/logsumexp.h b/macos/third_party/mlx/mlx/backend/metal/kernels/logsumexp.h
new file mode 100644
index 0000000..c746050
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/logsumexp.h
@@ -0,0 +1,140 @@
+// Copyright © 2025 Apple Inc.
+
+template <typename T, typename AccT = float, int N_READS = 4>
+[[kernel]] void logsumexp(
+    const device T* in,
+    device T* out,
+    constant int& axis_size,
+    uint gid [[threadgroup_position_in_grid]],
+    uint _lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  int lid = _lid;
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup AccT local_max[SIMD_SIZE];
+  threadgroup AccT local_normalizer[SIMD_SIZE];
+
+  AccT ld[N_READS];
+
+  in += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      ld[i] = AccT(in[i]);
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      ld[i] =
+          ((lid * N_READS + i) < axis_size) ? AccT(in[i]) : Limits<AccT>::min;
+    }
+  }
+  if (simd_group_id == 0) {
+    local_max[simd_lane_id] = Limits<AccT>::min;
+    local_normalizer[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Get the max
+  AccT maxval = Limits<AccT>::finite_min;
+  for (int i = 0; i < N_READS; i++) {
+    maxval = (maxval < ld[i]) ? ld[i] : maxval;
+  }
+  maxval = simd_max(maxval);
+  if (simd_lane_id == 0) {
+    local_max[simd_group_id] = maxval;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    maxval = simd_max(local_max[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_max[0] = maxval;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  maxval = local_max[0];
+
+  // Compute exp(x_i - maxval) and store the partial sums in local_normalizer
+  AccT normalizer = 0;
+  for (int i = 0; i < N_READS; i++) {
+    normalizer += fast::exp(ld[i] - maxval);
+  }
+  normalizer = simd_sum(normalizer);
+  if (simd_lane_id == 0) {
+    local_normalizer[simd_group_id] = normalizer;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    normalizer = simd_sum(local_normalizer[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      out[gid] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
+    }
+  }
+}
+
+template <typename T, typename AccT = float, int N_READS = 4>
+[[kernel]] void logsumexp_looped(
+    const device T* in,
+    device T* out,
+    constant int& axis_size,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  in += gid * size_t(axis_size);
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup AccT local_max[SIMD_SIZE];
+  threadgroup AccT local_normalizer[SIMD_SIZE];
+
+  // Get the max and the normalizer in one go
+  AccT prevmax;
+  AccT maxval = Limits<AccT>::finite_min;
+  AccT normalizer = 0;
+  for (int r = 0; r < static_cast<int>(ceildiv(axis_size, N_READS * lsize));
+       r++) {
+    int offset = r * lsize * N_READS + lid * N_READS;
+    AccT vals[N_READS];
+    if (offset + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        vals[i] = AccT(in[offset + i]);
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        vals[i] =
+            (offset + i < axis_size) ? AccT(in[offset + i]) : Limits<AccT>::min;
+      }
+    }
+    prevmax = maxval;
+    for (int i = 0; i < N_READS; i++) {
+      maxval = (maxval < vals[i]) ? vals[i] : maxval;
+    }
+    normalizer *= fast::exp(prevmax - maxval);
+    for (int i = 0; i < N_READS; i++) {
+      normalizer += fast::exp(vals[i] - maxval);
+    }
+  }
+  prevmax = maxval;
+  maxval = simd_max(maxval);
+  normalizer *= fast::exp(prevmax - maxval);
+  normalizer = simd_sum(normalizer);
+
+  prevmax = maxval;
+  if (simd_lane_id == 0) {
+    local_max[simd_group_id] = maxval;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  maxval = simd_max(local_max[simd_lane_id]);
+  normalizer *= fast::exp(prevmax - maxval);
+  if (simd_lane_id == 0) {
+    local_normalizer[simd_group_id] = normalizer;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  normalizer = simd_sum(local_normalizer[simd_lane_id]);
+
+  if (lid == 0) {
+    out[gid] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/logsumexp.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/logsumexp.metal
new file mode 100644
index 0000000..eb76436
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/logsumexp.metal
@@ -0,0 +1,18 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_common>
+#include <metal_simdgroup>
+
+using namespace metal;
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/logsumexp.h"
+
+#define instantiate_logsumexp(name, itype)                               \
+  instantiate_kernel("block_logsumexp_" #name, logsumexp, itype)         \
+  instantiate_kernel("looped_logsumexp_" #name, logsumexp_looped, itype) \
+
+instantiate_logsumexp(float32, float)
+instantiate_logsumexp(float16, half)
+instantiate_logsumexp(bfloat16, bfloat16_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/quantized.h b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized.h
new file mode 100644
index 0000000..5ac4c6e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized.h
@@ -0,0 +1,2508 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+using namespace metal;
+
+#define MLX_MTL_CONST static constant constexpr const
+
+MLX_MTL_CONST int SIMD_SIZE = 32;
+MLX_MTL_CONST int QUAD_SIZE = 4;
+
+template <int bits, int wsize = 8>
+inline constexpr short get_pack_factor() {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+template <int bits, int wsize = 8>
+inline constexpr short get_bytes_per_pack() {
+  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
+template <typename T, typename U, int values_per_thread, int bits>
+inline U load_vector(const device T* x, thread U* x_thread) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U sum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 4.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 64.0f;
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < values_per_thread; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 8.0f;
+      x_thread[i + 2] = x[i + 2] / 64.0f;
+      x_thread[i + 3] = x[i + 3] / 2.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 128.0f;
+      x_thread[i + 6] = x[i + 6] / 4.0f;
+      x_thread[i + 7] = x[i + 7] / 32.0f;
+    }
+  }
+
+  else if (bits == 4) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 16.0f;
+      x_thread[i + 2] = x[i + 2] / 256.0f;
+      x_thread[i + 3] = x[i + 3] / 4096.0f;
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < values_per_thread; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 32.0f;
+      x_thread[i + 2] = x[i + 2] / 4.0f;
+      x_thread[i + 3] = x[i + 3] / 128.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 2.0f;
+      x_thread[i + 6] = x[i + 6] / 64.0f;
+      x_thread[i + 7] = x[i + 7] / 8.0f;
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 64.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 4.0f;
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      sum += x[i];
+      x_thread[i] = x[i];
+    }
+  }
+
+  return sum;
+}
+
+template <typename T, typename U, int values_per_thread, int bits>
+inline U load_vector_safe(const device T* x, thread U* x_thread, int N) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U sum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < N; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 4.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 64.0f;
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < N; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 8.0f;
+      x_thread[i + 2] = x[i + 2] / 64.0f;
+      x_thread[i + 3] = x[i + 3] / 2.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 128.0f;
+      x_thread[i + 6] = x[i + 6] / 4.0f;
+      x_thread[i + 7] = x[i + 7] / 32.0f;
+    }
+  }
+
+  else if (bits == 4) {
+    for (int i = 0; i < N; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 16.0f;
+      x_thread[i + 2] = x[i + 2] / 256.0f;
+      x_thread[i + 3] = x[i + 3] / 4096.0f;
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < N; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 32.0f;
+      x_thread[i + 2] = x[i + 2] / 4.0f;
+      x_thread[i + 3] = x[i + 3] / 128.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 2.0f;
+      x_thread[i + 6] = x[i + 6] / 64.0f;
+      x_thread[i + 7] = x[i + 7] / 8.0f;
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < N; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 64.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 4.0f;
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < N; i++) {
+      sum += x[i];
+      x_thread[i] = x[i];
+    }
+  }
+
+  for (int i = N; i < values_per_thread; i++) {
+    x_thread[i] = 0;
+  }
+
+  return sum;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline U qdot(
+    const device uint8_t* w,
+    const thread U* x_thread,
+    U scale,
+    U bias,
+    U sum) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U accum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (w[i] & 0x03) +
+           x_thread[4 * i + 1] * (w[i] & 0x0c) +
+           x_thread[4 * i + 2] * (w[i] & 0x30) +
+           x_thread[4 * i + 3] * (w[i] & 0xc0));
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      x_thread += 8 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x07) * x_thread[0];
+      accum += (w[0] & 0x38) * x_thread[1];
+      accum += (w[0] & 0xc0) * x_thread[2];
+      accum += (w[1] & 0x01) * (x_thread[2] * 256.0f);
+
+      accum += (w[1] & 0x0e) * x_thread[3];
+      accum += (w[1] & 0x70) * x_thread[4];
+      accum += (w[1] & 0x80) * x_thread[5];
+      accum += (w[2] & 0x03) * (x_thread[5] * 256.0f);
+
+      accum += (w[2] & 0x1c) * x_thread[6];
+      accum += (w[2] & 0xe0) * x_thread[7];
+    }
+  }
+
+  else if (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (ws[i] & 0x000f) +
+           x_thread[4 * i + 1] * (ws[i] & 0x00f0) +
+           x_thread[4 * i + 2] * (ws[i] & 0x0f00) +
+           x_thread[4 * i + 3] * (ws[i] & 0xf000));
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      x_thread += 8 * i;
+      w += 5 * i;
+
+      accum += (w[0] & 0x1f) * x_thread[0];
+      accum += (w[0] & 0xe0) * x_thread[1];
+      accum += (w[1] & 0x3) * (x_thread[1] * 256.0f);
+      accum += (w[1] & 0x7c) * x_thread[2];
+      accum += (w[1] & 0x80) * x_thread[3];
+      accum += (w[2] & 0xf) * (x_thread[3] * 256.0f);
+      accum += (w[2] & 0xf0) * x_thread[4];
+      accum += (w[3] & 0x1) * (x_thread[4] * 256.0f);
+      accum += (w[3] & 0x3e) * x_thread[5];
+      accum += (w[3] & 0xc0) * x_thread[6];
+      accum += (w[4] & 0x7) * (x_thread[6] * 256.0f);
+      accum += (w[4] & 0xf8) * x_thread[7];
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      x_thread += 4 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x3f) * x_thread[0];
+
+      accum += (w[0] & 0xc0) * x_thread[1];
+      accum += (w[1] & 0x0f) * (x_thread[1] * 256.0f);
+
+      accum += (w[1] & 0xf0) * x_thread[2];
+      accum += (w[2] & 0x03) * (x_thread[2] * 256.0f);
+
+      accum += (w[2] & 0xfc) * x_thread[3];
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      accum += x_thread[i] * w[i];
+    }
+  }
+
+  return scale * accum + sum * bias;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline U qdot_safe(
+    const device uint8_t* w,
+    const thread U* x_thread,
+    U scale,
+    U bias,
+    U sum,
+    int N) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U accum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < (N / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (w[i] & 0x03) +
+           x_thread[4 * i + 1] * (w[i] & 0x0c) +
+           x_thread[4 * i + 2] * (w[i] & 0x30) +
+           x_thread[4 * i + 3] * (w[i] & 0xc0));
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (N / 8); i++) {
+      x_thread += 8 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x07) * x_thread[0];
+      accum += (w[0] & 0x38) * x_thread[1];
+      accum += (w[0] & 0xc0) * x_thread[2];
+      accum += (w[1] & 0x01) * (x_thread[2] * 256.0f);
+
+      accum += (w[1] & 0x0e) * x_thread[3];
+      accum += (w[1] & 0x70) * x_thread[4];
+      accum += (w[1] & 0x80) * x_thread[5];
+      accum += (w[2] & 0x03) * (x_thread[5] * 256.0f);
+
+      accum += (w[2] & 0x1c) * x_thread[6];
+      accum += (w[2] & 0xe0) * x_thread[7];
+    }
+  }
+
+  else if (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (N / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (ws[i] & 0x000f) +
+           x_thread[4 * i + 1] * (ws[i] & 0x00f0) +
+           x_thread[4 * i + 2] * (ws[i] & 0x0f00) +
+           x_thread[4 * i + 3] * (ws[i] & 0xf000));
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (N / 8); i++) {
+      x_thread += 8 * i;
+      w += 5 * i;
+
+      accum += (w[0] & 0x1f) * x_thread[0];
+      accum += (w[0] & 0xe0) * x_thread[1];
+      accum += (w[1] & 0x3) * (x_thread[1] * 256.0f);
+      accum += (w[1] & 0x7c) * x_thread[2];
+      accum += (w[1] & 0x80) * x_thread[3];
+      accum += (w[2] & 0xf) * (x_thread[3] * 256.0f);
+      accum += (w[2] & 0xf0) * x_thread[4];
+      accum += (w[3] & 0x1) * (x_thread[4] * 256.0f);
+      accum += (w[3] & 0x3e) * x_thread[5];
+      accum += (w[3] & 0xc0) * x_thread[6];
+      accum += (w[4] & 0x7) * (x_thread[6] * 256.0f);
+      accum += (w[4] & 0xf8) * x_thread[7];
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (N / 4); i++) {
+      x_thread += 4 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x3f) * x_thread[0];
+
+      accum += (w[0] & 0xc0) * x_thread[1];
+      accum += (w[1] & 0x0f) * (x_thread[1] * 256.0f);
+
+      accum += (w[1] & 0xf0) * x_thread[2];
+      accum += (w[2] & 0x03) * (x_thread[2] * 256.0f);
+
+      accum += (w[2] & 0xfc) * x_thread[3];
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < N; i++) {
+      accum += x_thread[i] * w[i];
+    }
+  }
+
+  return scale * accum + sum * bias;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline void
+qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  if (bits == 2) {
+    U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f};
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias);
+      result[4 * i + 1] += x * (s[1] * (w[i] & 0x0c) + bias);
+      result[4 * i + 2] += x * (s[2] * (w[i] & 0x30) + bias);
+      result[4 * i + 3] += x * (s[3] * (w[i] & 0xc0) + bias);
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      uint8_t w0 = w[3 * i];
+      uint8_t w1 = w[3 * i + 1];
+      uint8_t w2 = w[3 * i + 2];
+
+      result[8 * i] += x * ((w0 & 0x7) * scale + bias);
+      result[8 * i + 1] += x * (((w0 & 0x38) >> 3) * scale + bias);
+      result[8 * i + 2] +=
+          x * ((((w0 & 0xc0) >> 6) + ((w1 & 0x1) << 2)) * scale + bias);
+      result[8 * i + 3] += x * (((w1 & 0xe) >> 1) * scale + bias);
+      result[8 * i + 4] += x * (((w1 & 0x70) >> 4) * scale + bias);
+      result[8 * i + 5] +=
+          x * ((((w1 & 0x80) >> 7) + ((w2 & 0x3) << 1)) * scale + bias);
+      result[8 * i + 6] += x * (((w2 & 0x1c) >> 2) * scale + bias);
+      result[8 * i + 7] += x * (((w2 & 0xe0) >> 5) * scale + bias);
+    }
+  }
+
+  else if (bits == 4) {
+    U s[2] = {scale, scale / 16.0f};
+    for (int i = 0; i < (values_per_thread / 2); i++) {
+      result[2 * i] += x * (s[0] * (w[i] & 0x0f) + bias);
+      result[2 * i + 1] += x * (s[1] * (w[i] & 0xf0) + bias);
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      uint8_t w0 = w[5 * i];
+      uint8_t w1 = w[5 * i + 1];
+      uint8_t w2 = w[5 * i + 2];
+      uint8_t w3 = w[5 * i + 3];
+      uint8_t w4 = w[5 * i + 4];
+      result[8 * i] += x * ((w0 & 0x1f) * scale + bias);
+      result[8 * i + 1] +=
+          x * ((((w0 & 0xe0) >> 5) + ((w1 & 0x3) << 3)) * scale + bias);
+      result[8 * i + 2] += x * (((w1 & 0x7c) >> 2) * scale + bias);
+      result[8 * i + 3] +=
+          x * ((((w1 & 0x80) >> 7) + ((w2 & 0xf) << 1)) * scale + bias);
+      result[8 * i + 4] +=
+          x * ((((w2 & 0xf0) >> 4) + ((w3 & 0x1) << 4)) * scale + bias);
+      result[8 * i + 5] += x * (((w3 & 0x3e) >> 1) * scale + bias);
+      result[8 * i + 6] +=
+          x * ((((w3 & 0xc0) >> 6) + ((w4 & 0x7) << 2)) * scale + bias);
+      result[8 * i + 7] += x * (((w4 & 0xf8) >> 3) * scale + bias);
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      uint8_t w0 = w[3 * i];
+      uint8_t w1 = w[3 * i + 1];
+      uint8_t w2 = w[3 * i + 2];
+
+      result[4 * i] += x * ((w0 & 0x3f) * scale + bias);
+      result[4 * i + 1] +=
+          x * ((((w0 >> 6) & 0x03) + ((w1 & 0x0f) << 2)) * scale + bias);
+      result[4 * i + 2] +=
+          x * ((((w1 >> 4) & 0x0f) + ((w2 & 0x03) << 4)) * scale + bias);
+      result[4 * i + 3] += x * (((w2 >> 2) & 0x3f) * scale + bias);
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      result[i] += x * (scale * w[i] + bias);
+    }
+  }
+}
+
+template <typename U, int N, int bits>
+inline void
+dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  if (bits == 2) {
+    U s[4] = {
+        scale,
+        scale / static_cast<U>(4.0f),
+        scale / static_cast<U>(16.0f),
+        scale / static_cast<U>(64.0f)};
+    for (int i = 0; i < (N / 4); i++) {
+      w_local[4 * i] = s[0] * (w[i] & 0x03) + bias;
+      w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias;
+      w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias;
+      w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias;
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (N / 8); i++) {
+      w_local += 8 * i;
+      w += 3 * i;
+
+      w_local[0] = (w[0] & 0x7) * scale + bias;
+      w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias;
+      w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias;
+      w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias;
+      w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias;
+      w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias;
+      w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias;
+      w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias;
+    }
+  }
+
+  else if (bits == 4) {
+    U s[2] = {scale, scale / static_cast<U>(16.0f)};
+    for (int i = 0; i < (N / 2); i++) {
+      w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias;
+      w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias;
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (N / 8); i++) {
+      w_local += 8 * i;
+      w += 5 * i;
+
+      w_local[0] = (w[0] & 0x1f) * scale + bias;
+      w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias;
+      w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias;
+      w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias;
+      w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias;
+      w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias;
+      w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias;
+      w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias;
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (N / 4); i++) {
+      w_local += 4 * i;
+      w += 3 * i;
+      w_local[0] = (w[0] & 0x3f) * scale + bias;
+      w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias;
+      w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias;
+      w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias;
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < N; i++) {
+      w_local[i] = scale * w[i] + bias;
+    }
+  }
+}
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short group_size,
+    short bits>
+struct QuantizedBlockLoader {
+  static_assert(
+      BCOLS <= group_size,
+      "The group size should be larger than the columns");
+  static_assert(
+      group_size % BCOLS == 0,
+      "The group size should be divisible by the columns");
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  MLX_MTL_CONST short pack_factor = get_pack_factor<bits, 8>();
+  MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack<bits>();
+  MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
+  MLX_MTL_CONST short n_reads =
+      (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
+  MLX_MTL_CONST short group_steps = group_size / BCOLS;
+
+  const int src_ld;
+  const int tile_stride;
+  short group_step_cnt;
+  const int group_stride;
+
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  threadgroup T* dst;
+  const device uint8_t* src;
+  const device T* scales;
+  const device T* biases;
+
+  QuantizedBlockLoader(
+      const device uint8_t* src_,
+      const device T* scales_,
+      const device T* biases_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(
+            reduction_dim ? BCOLS_PACKED * bytes_per_pack
+                          : BROWS * src_ld * bytes_per_pack / pack_factor),
+        group_step_cnt(0),
+        group_stride(BROWS * src_ld / group_size),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(n_reads * thread_idx / BCOLS_PACKED),
+        bj((n_reads * thread_idx) % BCOLS_PACKED),
+        dst(dst_ + bi * dst_ld + bj * pack_factor),
+        src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
+            bj * bytes_per_pack),
+        scales(scales_ + bi * src_ld / group_size),
+        biases(biases_ + bi * src_ld / group_size) {}
+
+  void load_unsafe() const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    T scale = *scales;
+    T bias = *biases;
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, pack_factor, bits>(
+          src + i * bytes_per_pack, scale, bias, dst + i * pack_factor);
+    }
+  }
+
+  void load_safe(short2 src_tile_dim) const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    T scale = *scales;
+    T bias = *biases;
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, pack_factor, bits>(
+          (device uint8_t*)(src + i * bytes_per_pack),
+          scale,
+          bias,
+          dst + i * pack_factor);
+    }
+  }
+
+  void next() {
+    src += tile_stride;
+    if (reduction_dim == 1) {
+      if (group_steps > 1) {
+        group_step_cnt++;
+        if (group_step_cnt == group_steps) {
+          group_step_cnt = 0;
+          scales++;
+          biases++;
+        }
+      } else {
+        scales++;
+        biases++;
+      }
+    } else {
+      scales += group_stride;
+      biases += group_stride;
+    }
+  }
+};
+
+template <typename T, int group_size, int bits, int D>
+METAL_FUNC void qmv_quad_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    constant int& in_vec_size,
+    const constant int& out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint quad_gid [[quadgroup_index_in_threadgroup]],
+    uint quad_lid [[thread_index_in_quadgroup]]) {
+  constexpr int quads_per_simd = SIMD_SIZE / QUAD_SIZE;
+  constexpr int pack_factor = 32 / bits;
+  constexpr int values_per_thread = D / QUAD_SIZE;
+  constexpr int packs_per_thread = values_per_thread / pack_factor;
+  constexpr int scale_step_per_thread = group_size / values_per_thread;
+  constexpr int results_per_quadgroup = 8;
+
+  typedef float U;
+
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_quadgroup] = {0};
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const int out_row = tid.y * quads_per_simd * results_per_quadgroup + quad_gid;
+
+  w += out_row * in_vec_size_w + quad_lid * packs_per_thread;
+  scales += out_row * in_vec_size_g + quad_lid / scale_step_per_thread;
+  biases += out_row * in_vec_size_g + quad_lid / scale_step_per_thread;
+  x += tid.x * in_vec_size + quad_lid * values_per_thread;
+  y += tid.x * out_vec_size + out_row;
+
+  U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
+
+  for (int row = 0; row < results_per_quadgroup; row++) {
+    auto wl = (const device uint8_t*)(w + row * in_vec_size_w * quads_per_simd);
+    const device T* sl = scales + row * in_vec_size_g * quads_per_simd;
+    const device T* bl = biases + row * in_vec_size_g * quads_per_simd;
+
+    U s = sl[0];
+    U b = bl[0];
+    if (row * quads_per_simd + out_row < out_vec_size) {
+      result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
+    }
+  }
+
+  for (int row = 0; row < results_per_quadgroup; row++) {
+    result[row] = quad_sum(result[row]);
+    if (quad_lid == 0 && row * quads_per_simd + out_row < out_vec_size) {
+      y[row * quads_per_simd] = static_cast<T>(result[row]);
+    }
+  }
+}
+
+template <typename T, int group_size, int bits>
+METAL_FUNC void qmv_fast_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int packs_per_thread = bits == 2 ? 1 : 2;
+  constexpr int num_simdgroups = 2;
+  constexpr int results_per_simdgroup = 4;
+  constexpr int pack_factor = get_pack_factor<bits, 32>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits, 32>();
+  constexpr int values_per_thread = pack_factor * packs_per_thread;
+  constexpr int block_size = values_per_thread * SIMD_SIZE;
+  constexpr int scale_step_per_thread = group_size / values_per_thread;
+
+  const device uint8_t* ws = (const device uint8_t*)w;
+
+  typedef float U;
+
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_simdgroup] = {0};
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) +
+      simd_gid * results_per_simdgroup;
+
+  ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack;
+  scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+  biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+  x += tid.x * in_vec_size + simd_lid * values_per_thread;
+  y += tid.x * out_vec_size + out_row;
+
+  for (int k = 0; k < in_vec_size; k += block_size) {
+    U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
+
+    for (int row = 0; row < results_per_simdgroup; row++) {
+      auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+      const device T* sl = scales + row * in_vec_size_g;
+      const device T* bl = biases + row * in_vec_size_g;
+
+      U s = sl[0];
+      U b = bl[0];
+      result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
+    }
+
+    ws += block_size * bytes_per_pack / pack_factor;
+    scales += block_size / group_size;
+    biases += block_size / group_size;
+    x += block_size;
+  }
+
+  for (int row = 0; row < results_per_simdgroup; row++) {
+    result[row] = simd_sum(result[row]);
+    if (simd_lid == 0) {
+      y[row] = static_cast<T>(result[row]);
+    }
+  }
+}
+
+template <typename T, int group_size, int bits>
+METAL_FUNC void qmv_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int num_simdgroups = 2;
+  constexpr int results_per_simdgroup = 4;
+  constexpr int packs_per_thread = 1;
+  constexpr int pack_factor = get_pack_factor<bits, 32>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits, 32>();
+
+  constexpr int values_per_thread = pack_factor * packs_per_thread;
+  constexpr int block_size = values_per_thread * SIMD_SIZE;
+  constexpr int scale_step_per_thread = group_size / values_per_thread;
+
+  const device uint8_t* ws = (const device uint8_t*)w;
+
+  typedef float U;
+
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_simdgroup] = {0};
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) +
+      simd_gid * results_per_simdgroup;
+  const int used_out_row = min(out_vec_size - results_per_simdgroup, out_row);
+
+  if (out_row >= out_vec_size) {
+    return;
+  }
+
+  // In this case we need to properly guard all our reads because there isn't
+  // even 1 tile in the matrix
+  if (out_vec_size < (num_simdgroups * results_per_simdgroup)) {
+    ws +=
+        out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack;
+    scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    x += tid.x * in_vec_size + simd_lid * values_per_thread;
+    y += tid.x * out_vec_size + out_row;
+
+    int k = 0;
+    for (; k < in_vec_size - block_size; k += block_size) {
+      U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
+
+      for (int row = 0;
+           row < results_per_simdgroup && out_row + row < out_vec_size;
+           row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device T* sl = scales + row * in_vec_size_g;
+        const device T* bl = biases + row * in_vec_size_g;
+
+        U s = sl[0];
+        U b = bl[0];
+        result[row] +=
+            qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
+      }
+
+      ws += block_size * bytes_per_pack / pack_factor;
+      scales += block_size / group_size;
+      biases += block_size / group_size;
+      x += block_size;
+    }
+    const int remaining = clamp(
+        static_cast<int>(in_vec_size - k - simd_lid * values_per_thread),
+        0,
+        values_per_thread);
+    if (remaining > 0) {
+      U sum = load_vector_safe<T, U, values_per_thread, bits>(
+          x, x_thread, remaining);
+
+      for (int row = 0;
+           row < results_per_simdgroup && out_row + row < out_vec_size;
+           row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device T* sl = scales + row * in_vec_size_g;
+        const device T* bl = biases + row * in_vec_size_g;
+
+        U s = sl[0];
+        U b = bl[0];
+        result[row] += qdot_safe<U, values_per_thread, bits>(
+            wl, x_thread, s, b, sum, remaining);
+      }
+    }
+
+    for (int row = 0;
+         row < results_per_simdgroup && out_row + row < out_vec_size;
+         row++) {
+      result[row] = simd_sum(result[row]);
+      if (simd_lid == 0) {
+        y[row] = static_cast<T>(result[row]);
+      }
+    }
+  }
+
+  // In this case the last tile is moved back to redo some output values
+  else {
+    ws += used_out_row * in_vec_size_w +
+        simd_lid * packs_per_thread * bytes_per_pack;
+    scales += used_out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    biases += used_out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    x += tid.x * in_vec_size + simd_lid * values_per_thread;
+    y += tid.x * out_vec_size + used_out_row;
+
+    int k = 0;
+    for (; k < in_vec_size - block_size; k += block_size) {
+      U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
+
+      for (int row = 0; row < results_per_simdgroup; row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device T* sl = scales + row * in_vec_size_g;
+        const device T* bl = biases + row * in_vec_size_g;
+
+        U s = sl[0];
+        U b = bl[0];
+        result[row] +=
+            qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
+      }
+
+      ws += block_size * bytes_per_pack / pack_factor;
+      scales += block_size / group_size;
+      biases += block_size / group_size;
+      x += block_size;
+    }
+    const int remaining = clamp(
+        static_cast<int>(in_vec_size - k - simd_lid * values_per_thread),
+        0,
+        values_per_thread);
+    if (remaining > 0) {
+      U sum = load_vector_safe<T, U, values_per_thread, bits>(
+          x, x_thread, remaining);
+
+      for (int row = 0; row < results_per_simdgroup; row++) {
+        auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
+        const device T* sl = scales + row * in_vec_size_g;
+        const device T* bl = biases + row * in_vec_size_g;
+
+        U s = sl[0];
+        U b = bl[0];
+        result[row] += qdot_safe<U, values_per_thread, bits>(
+            wl, x_thread, s, b, sum, remaining);
+      }
+    }
+    for (int row = 0; row < results_per_simdgroup; row++) {
+      result[row] = simd_sum(result[row]);
+      if (simd_lid == 0) {
+        y[row] = static_cast<T>(result[row]);
+      }
+    }
+  }
+}
+
+template <typename T, const int group_size, const int bits>
+METAL_FUNC void qvm_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    const int in_vec_size,
+    const int out_vec_size,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
+  constexpr int num_simdgroups = 2;
+  constexpr int pack_factor = get_pack_factor<bits, 32>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+
+  constexpr int tn = 32 / pack_factor;
+  constexpr int block_size = SIMD_SIZE;
+
+  using W_T =
+      typename ConditionalType<power_of_2_bits, uint32_t, uint8_t>::type;
+  const device W_T* ws = (const device W_T*)w;
+
+  typedef float U;
+  typedef struct {
+    W_T wi[tn * bytes_per_pack];
+  } vec_w;
+
+  thread vec_w w_local;
+  thread U result[tn * pack_factor] = {0};
+  thread U scale = 1;
+  thread U bias = 0;
+  thread U x_local = 0;
+
+  // Adjust positions
+  const int out_vec_size_w = out_vec_size * bytes_per_pack / pack_factor;
+  const int out_vec_size_g = out_vec_size / group_size;
+  int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid);
+  ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w;
+  scales += out_col / group_size + simd_lid * out_vec_size_g;
+  biases += out_col / group_size + simd_lid * out_vec_size_g;
+  x += tid.x * in_vec_size + simd_lid;
+  y += tid.x * out_vec_size + out_col;
+
+  if (out_col >= out_vec_size) {
+    return;
+  }
+
+  // Loop over in_vec in blocks of block_size
+  int remaining = in_vec_size % block_size;
+  if (remaining == 0) {
+    for (int i = 0; i < in_vec_size; i += block_size) {
+      x_local = *x;
+      scale = *scales;
+      bias = *biases;
+      w_local = *((device vec_w*)ws);
+      qouter<U, tn * pack_factor, bits>(
+          (thread uint8_t*)&w_local, x_local, scale, bias, result);
+
+      x += block_size;
+      scales += block_size * out_vec_size_g;
+      biases += block_size * out_vec_size_g;
+      ws += block_size * out_vec_size_w;
+    }
+  } else {
+    for (int i = block_size; i < in_vec_size; i += block_size) {
+      x_local = *x;
+      scale = *scales;
+      bias = *biases;
+      w_local = *((device vec_w*)ws);
+
+      qouter<U, tn * pack_factor, bits>(
+          (thread uint8_t*)&w_local, x_local, scale, bias, result);
+
+      x += block_size;
+      scales += block_size * out_vec_size_g;
+      biases += block_size * out_vec_size_g;
+      ws += block_size * out_vec_size_w;
+    }
+    if (static_cast<int>(simd_lid) < remaining) {
+      x_local = *x;
+      scale = *scales;
+      bias = *biases;
+      w_local = *((device vec_w*)ws);
+    } else {
+      x_local = 0;
+      scale = 0;
+      bias = 0;
+    }
+    qouter<U, tn * pack_factor, bits>(
+        (thread uint8_t*)&w_local, x_local, scale, bias, result);
+  }
+
+// Accumulate in the simdgroup
+#pragma clang loop unroll(full)
+  for (int k = 0; k < tn * pack_factor; k++) {
+    result[k] = simd_sum(result[k]);
+  }
+
+  // Store the result
+  if (simd_lid == 0) {
+#pragma clang loop unroll(full)
+    for (int k = 0; k < tn * pack_factor; k++) {
+      y[k] = static_cast<T>(result[k]);
+    }
+  }
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+METAL_FUNC void qmm_t_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    threadgroup T* Xs,
+    threadgroup T* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  (void)lid;
+
+  constexpr int WM = 2;
+  constexpr int WN = 2;
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  // Instantiate the appropriate BlockMMA and Loader
+  using mma_t = mlx::steel::
+      BlockMMA<T, T, BM, BN, BK, WM, WN, false, true, BK_padded, BK_padded>;
+  using loader_x_t =
+      mlx::steel::BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE>;
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      BN,
+      BK,
+      BK_padded,
+      1,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  // Set the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+
+  auto wl = (const device uint8_t*)w;
+
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * K_w;
+  scales += y_col * K_g;
+  biases += y_col * K_g;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the x loader and mma operation
+  const short num_els = min(BM, M - y_row);
+  const short num_outs = min(BN, N - y_col);
+  loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
+  loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid);
+  mma_t mma_op(simd_gid, simd_lid);
+
+  if (num_els < BM) {
+    if (!aligned_N && num_outs < BN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_safe(short2(BK, num_outs));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  } else {
+    if (!aligned_N && num_outs < BN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_safe(short2(BK, num_outs));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  }
+
+  // Store results to device memory
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (num_els < BM || num_outs < BN) {
+    mma_op.store_result_safe(y, N, short2(num_outs, num_els));
+  } else {
+    mma_op.store_result(y, N);
+  }
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+METAL_FUNC void qmm_n_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    threadgroup T* Xs,
+    threadgroup T* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  (void)lid;
+
+  constexpr int WM = 2;
+  constexpr int WN = 2;
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  // Instantiate the appropriate BlockMMA and Loader
+  using mma_t = mlx::steel::
+      BlockMMA<T, T, BM, BN, BK, WM, WN, false, false, BK_padded, BN_padded>;
+  using loader_x_t = mlx::steel::
+      BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE, 1, 4>;
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      BK,
+      BN,
+      BN_padded,
+      0,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  auto wl = (const device uint8_t*)w;
+
+  // Set the block
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * bytes_per_pack / pack_factor;
+  scales += y_col / group_size;
+  biases += y_col / group_size;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the x loader and mma operation
+  const short num_els = min(BM, M - y_row);
+  loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
+  loader_w_t loader_w(wl, scales, biases, N, Ws, simd_gid, simd_lid);
+  mma_t mma_op(simd_gid, simd_lid);
+
+  if (num_els < BM) {
+    if ((K % BK) != 0) {
+      const int k_blocks = K / BK;
+      for (int k = 0; k < k_blocks; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+      const short num_k = K - k_blocks * BK;
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      loader_x.load_safe(short2(num_k, num_els));
+      loader_w.load_safe(short2(BN, num_k));
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      mma_op.mma(Xs, Ws);
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  } else {
+    if ((K % BK) != 0) {
+      const int k_blocks = K / BK;
+      for (int k = 0; k < k_blocks; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+      const short num_k = K - k_blocks * BK;
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      loader_x.load_safe(short2(num_k, BM));
+      loader_w.load_safe(short2(BN, num_k));
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      mma_op.mma(Xs, Ws);
+    } else {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_unsafe();
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
+    }
+  }
+
+  // Store results to device memory
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (num_els < BM) {
+    mma_op.store_result_safe(y, N, short2(BN, num_els));
+  } else {
+    mma_op.store_result(y, N);
+  }
+}
+
+template <typename T>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device T*& scales,
+    const device T*& biases,
+    device T*& y,
+    int output_stride,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int64_t* b_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx = tid.z;
+  uint32_t w_idx = tid.z;
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+    biases += w_idx * b_strides[0];
+  } else {
+    ulong3 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, b_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+    biases += idx.z;
+  }
+  y += tid.z * output_stride;
+}
+
+template <typename T>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device T*& scales,
+    const device T*& biases,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T*& y,
+    int output_stride,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int64_t* b_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx;
+  uint32_t w_idx;
+  if (batch_ndims == 1) {
+    x_idx = lhs_indices[tid.z * lhs_strides[0]];
+    w_idx = rhs_indices[tid.z * rhs_strides[0]];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        tid.z, batch_shape, lhs_strides, rhs_strides, batch_ndims);
+    x_idx = lhs_indices[idx.x];
+    w_idx = rhs_indices[idx.y];
+  }
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+    biases += w_idx * b_strides[0];
+  } else {
+    ulong3 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, b_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+    biases += idx.z;
+  }
+  y += tid.z * output_stride;
+}
+
+template <typename T, int group_size, int bits, int D, bool batched>
+[[kernel]] void affine_qmv_quad(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    const constant int& x_batch_ndims [[buffer(7)]],
+    const constant int* x_shape [[buffer(8)]],
+    const constant int64_t* x_strides [[buffer(9)]],
+    const constant int& w_batch_ndims [[buffer(10)]],
+    const constant int* w_shape [[buffer(11)]],
+    const constant int64_t* w_strides [[buffer(12)]],
+    const constant int64_t* s_strides [[buffer(13)]],
+    const constant int64_t* b_strides [[buffer(14)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint quad_gid [[quadgroup_index_in_threadgroup]],
+    uint quad_lid [[thread_index_in_quadgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+  qmv_quad_impl<T, group_size, bits, D>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      quad_gid,
+      quad_lid);
+}
+
+template <typename T, int group_size, int bits, bool batched>
+[[kernel]] void affine_qmv_fast(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    const constant int& x_batch_ndims [[buffer(7)]],
+    const constant int* x_shape [[buffer(8)]],
+    const constant int64_t* x_strides [[buffer(9)]],
+    const constant int& w_batch_ndims [[buffer(10)]],
+    const constant int* w_shape [[buffer(11)]],
+    const constant int64_t* w_strides [[buffer(12)]],
+    const constant int64_t* s_strides [[buffer(13)]],
+    const constant int64_t* b_strides [[buffer(14)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+  qmv_fast_impl<T, group_size, bits>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
+template <typename T, const int group_size, const int bits, bool batched>
+[[kernel]] void affine_qmv(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    const constant int& x_batch_ndims [[buffer(7)]],
+    const constant int* x_shape [[buffer(8)]],
+    const constant int64_t* x_strides [[buffer(9)]],
+    const constant int& w_batch_ndims [[buffer(10)]],
+    const constant int* w_shape [[buffer(11)]],
+    const constant int64_t* w_strides [[buffer(12)]],
+    const constant int64_t* s_strides [[buffer(13)]],
+    const constant int64_t* b_strides [[buffer(14)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+  qmv_impl<T, group_size, bits>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
+template <typename T, const int group_size, const int bits, bool batched>
+[[kernel]] void affine_qvm(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    const constant int& x_batch_ndims [[buffer(7)]],
+    const constant int* x_shape [[buffer(8)]],
+    const constant int64_t* x_strides [[buffer(9)]],
+    const constant int& w_batch_ndims [[buffer(10)]],
+    const constant int* w_shape [[buffer(11)]],
+    const constant int64_t* w_strides [[buffer(12)]],
+    const constant int64_t* s_strides [[buffer(13)]],
+    const constant int64_t* b_strides [[buffer(14)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  if (batched) {
+    int M = x_shape[x_batch_ndims];
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+  qvm_impl<T, group_size, bits>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
+template <typename T, const int group_size, const int bits, int split_k = 32>
+[[kernel]] void affine_qvm_split_k(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    const constant int& x_batch_ndims [[buffer(7)]],
+    const constant int* x_shape [[buffer(8)]],
+    const constant int64_t* x_strides [[buffer(9)]],
+    const constant int& w_batch_ndims [[buffer(10)]],
+    const constant int* w_shape [[buffer(11)]],
+    const constant int64_t* w_strides [[buffer(12)]],
+    const constant int64_t* s_strides [[buffer(13)]],
+    const constant int64_t* b_strides [[buffer(14)]],
+    const constant int& final_block_size [[buffer(15)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      y,
+      out_vec_size * M,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+
+  // When (in_vec_size % split_k != 0) the final block needs to be smaller
+  int in_vec_size_adj =
+      tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size;
+
+  qvm_impl<T, group_size, bits>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size_adj,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const bool batched,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void affine_qmm_t(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& K [[buffer(5)]],
+    const constant int& N [[buffer(6)]],
+    const constant int& M [[buffer(7)]],
+    const constant int& x_batch_ndims [[buffer(8)]],
+    const constant int* x_shape [[buffer(9)]],
+    const constant int64_t* x_strides [[buffer(10)]],
+    const constant int& w_batch_ndims [[buffer(11)]],
+    const constant int* w_shape [[buffer(12)]],
+    const constant int64_t* w_strides [[buffer(13)]],
+    const constant int64_t* s_strides [[buffer(14)]],
+    const constant int64_t* b_strides [[buffer(15)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BN * BK_padded];
+
+  if (batched) {
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+  qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN>(
+      w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool batched,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void affine_qmm_n(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& K [[buffer(5)]],
+    const constant int& N [[buffer(6)]],
+    const constant int& M [[buffer(7)]],
+    const constant int& x_batch_ndims [[buffer(8)]],
+    const constant int* x_shape [[buffer(9)]],
+    const constant int64_t* x_strides [[buffer(10)]],
+    const constant int& w_batch_ndims [[buffer(11)]],
+    const constant int* w_shape [[buffer(12)]],
+    const constant int64_t* w_strides [[buffer(13)]],
+    const constant int64_t* s_strides [[buffer(14)]],
+    const constant int64_t* b_strides [[buffer(15)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BK * BN_padded];
+
+  if (batched) {
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+
+  qmm_n_impl<T, group_size, bits, BM, BK, BN>(
+      w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <typename T, int group_size, int bits>
+[[kernel]] void affine_gather_qmv_fast(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    const device uint32_t* lhs_indices [[buffer(4)]],
+    const device uint32_t* rhs_indices [[buffer(5)]],
+    device T* y [[buffer(6)]],
+    const constant int& in_vec_size [[buffer(7)]],
+    const constant int& out_vec_size [[buffer(8)]],
+    const constant int& x_batch_ndims [[buffer(9)]],
+    const constant int* x_shape [[buffer(10)]],
+    const constant int64_t* x_strides [[buffer(11)]],
+    const constant int& w_batch_ndims [[buffer(12)]],
+    const constant int* w_shape [[buffer(13)]],
+    const constant int64_t* w_strides [[buffer(14)]],
+    const constant int64_t* s_strides [[buffer(15)]],
+    const constant int64_t* b_strides [[buffer(16)]],
+    const constant int& batch_ndims [[buffer(17)]],
+    const constant int* batch_shape [[buffer(18)]],
+    const constant int64_t* lhs_strides [[buffer(19)]],
+    const constant int64_t* rhs_strides [[buffer(20)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      y,
+      out_vec_size * M,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+  qmv_fast_impl<T, group_size, bits>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
+template <typename T, int group_size, int bits>
+[[kernel]] void affine_gather_qmv(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    const device uint32_t* lhs_indices [[buffer(4)]],
+    const device uint32_t* rhs_indices [[buffer(5)]],
+    device T* y [[buffer(6)]],
+    const constant int& in_vec_size [[buffer(7)]],
+    const constant int& out_vec_size [[buffer(8)]],
+    const constant int& x_batch_ndims [[buffer(9)]],
+    const constant int* x_shape [[buffer(10)]],
+    const constant int64_t* x_strides [[buffer(11)]],
+    const constant int& w_batch_ndims [[buffer(12)]],
+    const constant int* w_shape [[buffer(13)]],
+    const constant int64_t* w_strides [[buffer(14)]],
+    const constant int64_t* s_strides [[buffer(15)]],
+    const constant int64_t* b_strides [[buffer(16)]],
+    const constant int& batch_ndims [[buffer(17)]],
+    const constant int* batch_shape [[buffer(18)]],
+    const constant int64_t* lhs_strides [[buffer(19)]],
+    const constant int64_t* rhs_strides [[buffer(20)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      y,
+      out_vec_size * M,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+  qmv_impl<T, group_size, bits>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
+template <typename T, int group_size, int bits>
+[[kernel]] void affine_gather_qvm(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    const device uint32_t* lhs_indices [[buffer(4)]],
+    const device uint32_t* rhs_indices [[buffer(5)]],
+    device T* y [[buffer(6)]],
+    const constant int& in_vec_size [[buffer(7)]],
+    const constant int& out_vec_size [[buffer(8)]],
+    const constant int& x_batch_ndims [[buffer(9)]],
+    const constant int* x_shape [[buffer(10)]],
+    const constant int64_t* x_strides [[buffer(11)]],
+    const constant int& w_batch_ndims [[buffer(12)]],
+    const constant int* w_shape [[buffer(13)]],
+    const constant int64_t* w_strides [[buffer(14)]],
+    const constant int64_t* s_strides [[buffer(15)]],
+    const constant int64_t* b_strides [[buffer(16)]],
+    const constant int& batch_ndims [[buffer(17)]],
+    const constant int* batch_shape [[buffer(18)]],
+    const constant int64_t* lhs_strides [[buffer(19)]],
+    const constant int64_t* rhs_strides [[buffer(20)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  int M = x_shape[x_batch_ndims];
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      y,
+      out_vec_size * M,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+  qvm_impl<T, group_size, bits>(
+      w,
+      scales,
+      biases,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void affine_gather_qmm_t(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    const device uint32_t* lhs_indices [[buffer(4)]],
+    const device uint32_t* rhs_indices [[buffer(5)]],
+    device T* y [[buffer(6)]],
+    const constant int& K [[buffer(7)]],
+    const constant int& N [[buffer(8)]],
+    const constant int& M [[buffer(9)]],
+    const constant int& x_batch_ndims [[buffer(10)]],
+    const constant int* x_shape [[buffer(11)]],
+    const constant int64_t* x_strides [[buffer(12)]],
+    const constant int& w_batch_ndims [[buffer(13)]],
+    const constant int* w_shape [[buffer(14)]],
+    const constant int64_t* w_strides [[buffer(15)]],
+    const constant int64_t* s_strides [[buffer(16)]],
+    const constant int64_t* b_strides [[buffer(17)]],
+    const constant int& batch_ndims [[buffer(18)]],
+    const constant int* batch_shape [[buffer(19)]],
+    const constant int64_t* lhs_strides [[buffer(20)]],
+    const constant int64_t* rhs_strides [[buffer(21)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BN * BK_padded];
+
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+  qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN>(
+      w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 32,
+    const int BK = 32,
+    const int BN = 32>
+[[kernel]] void affine_gather_qmm_n(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    const device uint32_t* lhs_indices [[buffer(4)]],
+    const device uint32_t* rhs_indices [[buffer(5)]],
+    device T* y [[buffer(6)]],
+    const constant int& K [[buffer(7)]],
+    const constant int& N [[buffer(8)]],
+    const constant int& M [[buffer(9)]],
+    const constant int& x_batch_ndims [[buffer(10)]],
+    const constant int* x_shape [[buffer(11)]],
+    const constant int64_t* x_strides [[buffer(12)]],
+    const constant int& w_batch_ndims [[buffer(13)]],
+    const constant int* w_shape [[buffer(14)]],
+    const constant int64_t* w_strides [[buffer(15)]],
+    const constant int64_t* s_strides [[buffer(16)]],
+    const constant int64_t* b_strides [[buffer(17)]],
+    const constant int& batch_ndims [[buffer(18)]],
+    const constant int* batch_shape [[buffer(19)]],
+    const constant int64_t* lhs_strides [[buffer(20)]],
+    const constant int64_t* rhs_strides [[buffer(21)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[BK * BN_padded];
+
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+  qmm_n_impl<T, group_size, bits, BM, BK, BN>(
+      w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    int group_size,
+    int bits,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose>
+[[kernel]] void affine_gather_qmm_rhs(
+    const device T* x [[buffer(0)]],
+    const device uint32_t* w [[buffer(1)]],
+    const device T* scales [[buffer(2)]],
+    const device T* biases [[buffer(3)]],
+    const device uint32_t* indices [[buffer(4)]],
+    device T* y [[buffer(5)]],
+    const constant int& M [[buffer(6)]],
+    const constant int& N [[buffer(7)]],
+    const constant int& K [[buffer(8)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]]) {
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  using mma_t = mlx::steel::BlockMMA<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      false,
+      transpose,
+      BK_padded,
+      transpose ? BK_padded : BN_padded>;
+  using loader_x_t =
+      mlx::steel::BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE>;
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      transpose ? BN : BK,
+      transpose ? BK : BN,
+      transpose ? BK_padded : BN_padded,
+      transpose,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  threadgroup T Xs[BM * BK_padded];
+  threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded];
+
+  // Compute the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int N_w = N * bytes_per_pack / pack_factor;
+  const int N_g = N / group_size;
+  const int K_it = K / BK;
+  const size_t stride_w = transpose ? N * K_w : K * N_w;
+  const size_t stride_s = transpose ? N * K_g : K * N_g;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+  const size_t y_row_long = size_t(y_row);
+  const size_t y_col_long = size_t(y_col);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, M - y_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, N - y_col));
+
+  // Calculate the final tiles in the case that K is not aligned
+  const int k_remain = K - K_it * BK;
+  const short2 tile_x = short2(k_remain, tgp_bm);
+  const short2 tile_w =
+      transpose ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+  // Move x and output to the correct block
+  auto wl = (const device uint8_t*)w;
+  x += y_row_long * K;
+  y += y_row_long * N + y_col_long;
+  wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor;
+  scales += transpose ? y_col_long * K_g : y_col / group_size;
+  biases += transpose ? y_col_long * K_g : y_col / group_size;
+
+  // Do as many matmuls as necessary
+  uint32_t index;
+  short offset;
+  uint32_t index_next = indices[y_row];
+  short offset_next = 0;
+  int n = 0;
+  while (n < tgp_bm) {
+    n++;
+    offset = offset_next;
+    index = index_next;
+    offset_next = tgp_bm;
+    for (; n < tgp_bm; n++) {
+      if (indices[y_row + n] != index) {
+        offset_next = n;
+        index_next = indices[y_row + n];
+        break;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Prepare threadgroup mma operation
+    thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+    // Prepare threadgroup loading operations
+    thread loader_x_t loader_x(x, K, Xs, simd_group_id, simd_lane_id);
+    thread loader_w_t loader_w(
+        wl + index * stride_w,
+        scales + index * stride_s,
+        biases + index * stride_s,
+        transpose ? K : N,
+        Ws,
+        simd_group_id,
+        simd_lane_id);
+
+    // Matrices are all aligned check nothing
+    if (align_M && align_N) {
+      gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it);
+      if (!align_K) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        gemm_loop_finalize(Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+      }
+
+      // Store results to device memory
+      if (offset_next - offset == BM) {
+        mma_op.store_result(y, N);
+      } else {
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(BN, offset_next));
+      }
+    } else {
+      // Tile aligned so check outside of the hot loop
+      if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) {
+        gemm_loop_aligned(Xs, Ws, mma_op, loader_x, loader_w, K_it);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+
+        // Store results to device memory
+        if (offset_next - offset == BM) {
+          mma_op.store_result(y, N);
+        } else {
+          mma_op.store_result_slice(
+              y, N, short2(0, offset), short2(BN, offset_next));
+        }
+      }
+
+      // Tile partially aligned check rows
+      else if (align_N || tgp_bn == BN) {
+        gemm_loop_unaligned<false, true, transpose>(
+            Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(BN, offset_next));
+      }
+
+      // Tile partially aligned check cols
+      else if (align_M || tgp_bm == BM) {
+        gemm_loop_unaligned<true, false, transpose>(
+            Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(tgp_bn, offset_next));
+      }
+
+      // Nothing aligned so check both rows and cols
+      else {
+        gemm_loop_unaligned<false, false, transpose>(
+            Xs, Ws, mma_op, loader_x, loader_w, K_it, tgp_bm, tgp_bn, BK);
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          gemm_loop_finalize(
+              Xs, Ws, mma_op, loader_x, loader_w, tile_x, tile_w);
+        }
+        mma_op.store_result_slice(
+            y, N, short2(0, offset), short2(tgp_bn, offset_next));
+      }
+    }
+  }
+}
+
+template <typename T, const int group_size, const int bits>
+[[kernel]] void affine_quantize(
+    const device T* w [[buffer(0)]],
+    device uint8_t* out [[buffer(1)]],
+    device T* scales [[buffer(2)]],
+    device T* biases [[buffer(3)]],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  constexpr float eps = 1e-7;
+  constexpr int simd_size = 32;
+  constexpr float n_bins = (1 << bits) - 1;
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+  constexpr int values_per_reduce = group_size / simd_size;
+  constexpr int writes_per_reduce = pack_factor / values_per_reduce;
+  constexpr int writes_per_pack =
+      writes_per_reduce > 1 ? 1 : values_per_reduce / pack_factor;
+  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
+
+  static_assert(
+      group_size % simd_size == 0,
+      "Group size must be divisible by simd size.");
+
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  size_t in_index = offset * values_per_reduce;
+  size_t out_index = power_of_2_bits
+      ? offset * writes_per_pack
+      : offset * bytes_per_pack / writes_per_reduce;
+
+  float w_thread[values_per_reduce];
+  float w_min = Limits<T>::max;
+  float w_max = 0;
+
+#pragma clang loop unroll(full)
+  for (int i = 0; i < values_per_reduce; i++) {
+    float val = w[in_index + i];
+    w_thread[i] = val;
+    w_min = min(w_min, val);
+    w_max = max(w_max, val);
+  }
+
+  w_min = simd_min(w_min);
+  w_max = simd_max(w_max);
+
+  float scale = max((w_max - w_min) / n_bins, eps);
+  bool side = abs(w_min) > abs(w_max);
+  scale = side ? scale : -scale;
+  float edge = side ? w_min : w_max;
+  float q0 = round(edge / scale);
+  bool at_zero = q0 == 0.0f;
+  scale = at_zero ? scale : edge / q0;
+  float bias = at_zero ? 0 : edge;
+
+  // Write out the scales and biases
+  size_t gindex = in_index / group_size;
+  if (in_index % group_size == 0) {
+    scales[gindex] = static_cast<T>(scale);
+    biases[gindex] = static_cast<T>(bias);
+  }
+
+  using OutType = metal::conditional_t<bits == 5, uint64_t, uint32_t>;
+  OutType output = 0;
+
+#pragma clang loop unroll(full)
+  for (int i = 0; i < values_per_reduce; i++) {
+    uint8_t val = min(round((w_thread[i] - bias) / scale), n_bins);
+    if (bits == 8) {
+      output = val;
+    } else {
+      output |= val << (bits * (i % pack_factor));
+    }
+
+    if (pack_factor < values_per_reduce && i % pack_factor == pack_factor - 1) {
+      out[out_index + i / pack_factor] = output;
+      output = 0;
+    } else {
+#pragma clang loop unroll(full)
+      for (int j = 1; j < writes_per_reduce; j++) {
+        uint8_t sval = simd_shuffle_down(val, j);
+        output |= static_cast<OutType>(sval)
+            << (bits * (j * values_per_reduce + i));
+      }
+    }
+  }
+  if (bits == 3 || bits == 6) {
+    if (in_index % pack_factor == 0 && out_index % bytes_per_pack == 0) {
+      out[out_index] = output & 0xff;
+      out[out_index + 1] = (output & 0xff00) >> 8;
+      out[out_index + 2] = (output & 0xff0000) >> 16;
+    }
+  } else if (bits == 5) {
+    if (in_index % pack_factor == 0 && out_index % bytes_per_pack == 0) {
+      out[out_index] = output & 0xff;
+      out[out_index + 1] = (output & 0xff00) >> 8;
+      out[out_index + 2] = (output & 0xff0000) >> 16;
+      out[out_index + 3] = (output & 0xff000000) >> 24;
+      out[out_index + 4] = (output & 0xff00000000) >> 32;
+    }
+  } else {
+    if (writes_per_reduce > 0 && out_index % writes_per_reduce == 0) {
+      out[out_index / writes_per_reduce] = output;
+    }
+  }
+}
+
+template <typename T, const int group_size, const int bits>
+[[kernel]] void affine_dequantize(
+    const device uint8_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  size_t oindex = offset * pack_factor;
+  size_t gindex = oindex / group_size;
+  T scale = scales[gindex];
+  T bias = biases[gindex];
+
+  out += oindex;
+
+  if (bits == 3) {
+    w += offset * bytes_per_pack;
+    out[0] = (w[0] & 0x7) * scale + bias;
+    out[1] = ((w[0] & 0x38) >> 3) * scale + bias;
+    out[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias;
+    out[3] = ((w[1] & 0xe) >> 1) * scale + bias;
+    out[4] = ((w[1] & 0x70) >> 4) * scale + bias;
+    out[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias;
+    out[6] = ((w[2] & 0x1c) >> 2) * scale + bias;
+    out[7] = ((w[2] & 0xe0) >> 5) * scale + bias;
+  } else if (bits == 5) {
+    w += offset * bytes_per_pack;
+    out[0] = (w[0] & 0x1f) * scale + bias;
+    out[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias;
+    out[2] = ((w[1] & 0x7c) >> 2) * scale + bias;
+    out[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias;
+    out[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias;
+    out[5] = ((w[3] & 0x3e) >> 1) * scale + bias;
+    out[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias;
+    out[7] = ((w[4] & 0xf8) >> 3) * scale + bias;
+  } else if (bits == 6) {
+    w += offset * bytes_per_pack;
+    out[0] = (w[0] & 0x3f) * scale + bias;
+    out[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias;
+    out[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias;
+    out[3] = ((w[2] >> 2) & 0x3f) * scale + bias;
+  } else {
+    uint val = w[offset];
+#pragma clang loop unroll(full)
+    for (int i = 0; i < pack_factor; i++) {
+      uint8_t d;
+      if (bits == 2) {
+        d = (val >> (bits * i)) & 0x03;
+      } else if (bits == 4) {
+        d = (val >> (bits * i)) & 0x0f;
+      } else if (bits == 8) {
+        d = val;
+      }
+      out[i] = scale * d + bias;
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/quantized.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized.metal
new file mode 100644
index 0000000..f734b9b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized.metal
@@ -0,0 +1,144 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/quantized_utils.h"
+#include "mlx/backend/metal/kernels/quantized.h"
+
+#define instantiate_quantized(name, type, group_size, bits)     \
+  instantiate_kernel(                                                    \
+      #name "_" #type "_gs_" #group_size "_b_" #bits,                    \
+      name,                                                              \
+      type,                                                              \
+      group_size,                                                        \
+      bits)
+
+#define instantiate_quantized_batched(name, type, group_size, bits, batched)     \
+  instantiate_kernel(                                                    \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \
+      name,                                                              \
+      type,                                                              \
+      group_size,                                                        \
+      bits,                                                              \
+      batched)
+
+#define instantiate_quantized_aligned(name, type, group_size, bits, aligned)     \
+  instantiate_kernel(                                                                     \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \
+      name,                                                                  \
+      type,                                                                  \
+      group_size,                                                            \
+      bits,                                                                  \
+      aligned)
+
+#define instantiate_quantized_aligned_batched(name, type, group_size, bits, aligned, batched)     \
+  instantiate_kernel(                                                                     \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_batch_" #batched, \
+      name,                                                                  \
+      type,                                                                  \
+      group_size,                                                            \
+      bits,                                                                  \
+      aligned,                                                               \
+      batched)
+
+#define instantiate_quantized_quad(name, type, group_size, bits, D, batched)     \
+  instantiate_kernel(                                                            \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \
+      name,                                                         \
+      type,                                                         \
+      group_size,                                                   \
+      bits,                                                         \
+      D,                                                            \
+      batched)
+
+#define instantiate_quantized_split_k(name, type, group_size, bits, split_k)     \
+  instantiate_kernel(                                                            \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \
+      name,                                                         \
+      type,                                                         \
+      group_size,                                                   \
+      bits,                                                         \
+      split_k)
+
+#define instantiate_gather_qmm_rhs(func, name, type, group_size, bits, bm, bn, bk, wm, wn, transpose)        \
+  instantiate_kernel(                                                                                        \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \
+      func,                                                         \
+      type,                                                         \
+      group_size,                                                   \
+      bits,                                                         \
+      bm,                                                           \
+      bn,                                                           \
+      bk,                                                           \
+      wm,                                                           \
+      wn,                                                           \
+      transpose)
+
+#define instantiate_quantized_batched_wrap(name, type, group_size, bits) \
+  instantiate_quantized_batched(name, type, group_size, bits, 1)      \
+  instantiate_quantized_batched(name, type, group_size, bits, 0)
+
+#define instantiate_quantized_all_batched(type, group_size, bits) \
+  instantiate_quantized_batched_wrap(affine_qmv_fast, type, group_size, bits)     \
+  instantiate_quantized_batched_wrap(affine_qmv, type, group_size, bits)     \
+  instantiate_quantized_batched_wrap(affine_qvm, type, group_size, bits)     \
+  instantiate_quantized_batched_wrap(affine_qmm_n, type, group_size, bits)
+
+#define instantiate_quantized_all_single(type, group_size, bits) \
+  instantiate_quantized(affine_quantize, type, group_size, bits) \
+  instantiate_quantized(affine_dequantize, type, group_size, bits)     \
+  instantiate_quantized(affine_gather_qmv_fast, type, group_size, bits)     \
+  instantiate_quantized(affine_gather_qmv, type, group_size, bits)     \
+  instantiate_quantized(affine_gather_qvm, type, group_size, bits)     \
+  instantiate_quantized(affine_gather_qmm_n, type, group_size, bits)
+
+#define instantiate_quantized_all_aligned(type, group_size, bits)   \
+  instantiate_quantized_aligned(affine_gather_qmm_t, type, group_size, bits, true) \
+  instantiate_quantized_aligned(affine_gather_qmm_t, type, group_size, bits, false) \
+  instantiate_quantized_aligned_batched(affine_qmm_t, type, group_size, bits, true, 1) \
+  instantiate_quantized_aligned_batched(affine_qmm_t, type, group_size, bits, true, 0) \
+  instantiate_quantized_aligned_batched(affine_qmm_t, type, group_size, bits, false, 1) \
+  instantiate_quantized_aligned_batched(affine_qmm_t, type, group_size, bits, false, 0)
+
+#define instantiate_quantized_all_quad(type, group_size, bits)   \
+  instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 64, 1)   \
+  instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 64, 0)   \
+  instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 1)  \
+  instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 0)
+
+#define instantiate_quantized_all_splitk(type, group_size, bits)   \
+  instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 8)   \
+  instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 32)
+
+#define instantiate_quantized_all_rhs(type, group_size, bits) \
+  instantiate_gather_qmm_rhs(affine_gather_qmm_rhs, affine_gather_qmm_rhs_nt, type, group_size, bits, 16, 32, 32, 1, 2, true) \
+  instantiate_gather_qmm_rhs(affine_gather_qmm_rhs, affine_gather_qmm_rhs_nn, type, group_size, bits, 16, 32, 32, 1, 2, false)
+
+#define instantiate_quantized_funcs(type, group_size, bits) \
+  instantiate_quantized_all_single(type, group_size, bits)  \
+  instantiate_quantized_all_batched(type, group_size, bits) \
+  instantiate_quantized_all_aligned(type, group_size, bits) \
+  instantiate_quantized_all_quad(type, group_size, bits)    \
+  instantiate_quantized_all_splitk(type, group_size, bits)  \
+  instantiate_quantized_all_rhs(type, group_size, bits)
+
+#define instantiate_quantized_types(group_size, bits)       \
+  instantiate_quantized_funcs(float, group_size, bits)      \
+  instantiate_quantized_funcs(float16_t, group_size, bits)  \
+  instantiate_quantized_funcs(bfloat16_t, group_size, bits)
+
+#define instantiate_quantized_groups(bits) \
+  instantiate_quantized_types(128, bits)   \
+  instantiate_quantized_types(64, bits)    \
+  instantiate_quantized_types(32, bits)
+
+#define instantiate_quantized_all() \
+  instantiate_quantized_groups(2) \
+  instantiate_quantized_groups(3) \
+  instantiate_quantized_groups(4) \
+  instantiate_quantized_groups(5) \
+  instantiate_quantized_groups(6) \
+  instantiate_quantized_groups(8)
+
+instantiate_quantized_all() // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_nax.h
new file mode 100644
index 0000000..c26ff64
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_nax.h
@@ -0,0 +1,1705 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+using namespace metal;
+using namespace mlx::steel;
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+using namespace metal;
+
+#define MLX_MTL_CONST static constant constexpr const
+
+MLX_MTL_CONST int SIMD_SIZE = 32;
+MLX_MTL_CONST int QUAD_SIZE = 4;
+
+template <int bits, int wsize = 8>
+inline constexpr short get_pack_factor() {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+template <int bits, int wsize = 8>
+inline constexpr short get_bytes_per_pack() {
+  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
+template <typename T, typename U, int values_per_thread, int bits>
+inline U load_vector(const device T* x, thread U* x_thread) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U sum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 4.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 64.0f;
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < values_per_thread; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 8.0f;
+      x_thread[i + 2] = x[i + 2] / 64.0f;
+      x_thread[i + 3] = x[i + 3] / 2.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 128.0f;
+      x_thread[i + 6] = x[i + 6] / 4.0f;
+      x_thread[i + 7] = x[i + 7] / 32.0f;
+    }
+  }
+
+  else if (bits == 4) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 16.0f;
+      x_thread[i + 2] = x[i + 2] / 256.0f;
+      x_thread[i + 3] = x[i + 3] / 4096.0f;
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < values_per_thread; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 32.0f;
+      x_thread[i + 2] = x[i + 2] / 4.0f;
+      x_thread[i + 3] = x[i + 3] / 128.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 2.0f;
+      x_thread[i + 6] = x[i + 6] / 64.0f;
+      x_thread[i + 7] = x[i + 7] / 8.0f;
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 64.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 4.0f;
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      sum += x[i];
+      x_thread[i] = x[i];
+    }
+  }
+
+  return sum;
+}
+
+template <typename T, typename U, int values_per_thread, int bits>
+inline U load_vector_safe(const device T* x, thread U* x_thread, int N) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U sum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < N; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 4.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 64.0f;
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < N; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 8.0f;
+      x_thread[i + 2] = x[i + 2] / 64.0f;
+      x_thread[i + 3] = x[i + 3] / 2.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 128.0f;
+      x_thread[i + 6] = x[i + 6] / 4.0f;
+      x_thread[i + 7] = x[i + 7] / 32.0f;
+    }
+  }
+
+  else if (bits == 4) {
+    for (int i = 0; i < N; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 16.0f;
+      x_thread[i + 2] = x[i + 2] / 256.0f;
+      x_thread[i + 3] = x[i + 3] / 4096.0f;
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < N; i += 8) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
+          x[i + 6] + x[i + 7];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 32.0f;
+      x_thread[i + 2] = x[i + 2] / 4.0f;
+      x_thread[i + 3] = x[i + 3] / 128.0f;
+      x_thread[i + 4] = x[i + 4] / 16.0f;
+      x_thread[i + 5] = x[i + 5] / 2.0f;
+      x_thread[i + 6] = x[i + 6] / 64.0f;
+      x_thread[i + 7] = x[i + 7] / 8.0f;
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < N; i += 4) {
+      sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
+      x_thread[i] = x[i];
+      x_thread[i + 1] = x[i + 1] / 64.0f;
+      x_thread[i + 2] = x[i + 2] / 16.0f;
+      x_thread[i + 3] = x[i + 3] / 4.0f;
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < N; i++) {
+      sum += x[i];
+      x_thread[i] = x[i];
+    }
+  }
+
+  for (int i = N; i < values_per_thread; i++) {
+    x_thread[i] = 0;
+  }
+
+  return sum;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline U qdot(
+    const device uint8_t* w,
+    const thread U* x_thread,
+    U scale,
+    U bias,
+    U sum) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U accum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (w[i] & 0x03) +
+           x_thread[4 * i + 1] * (w[i] & 0x0c) +
+           x_thread[4 * i + 2] * (w[i] & 0x30) +
+           x_thread[4 * i + 3] * (w[i] & 0xc0));
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      x_thread += 8 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x07) * x_thread[0];
+      accum += (w[0] & 0x38) * x_thread[1];
+      accum += (w[0] & 0xc0) * x_thread[2];
+      accum += (w[1] & 0x01) * (x_thread[2] * 256.0f);
+
+      accum += (w[1] & 0x0e) * x_thread[3];
+      accum += (w[1] & 0x70) * x_thread[4];
+      accum += (w[1] & 0x80) * x_thread[5];
+      accum += (w[2] & 0x03) * (x_thread[5] * 256.0f);
+
+      accum += (w[2] & 0x1c) * x_thread[6];
+      accum += (w[2] & 0xe0) * x_thread[7];
+    }
+  }
+
+  else if (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (ws[i] & 0x000f) +
+           x_thread[4 * i + 1] * (ws[i] & 0x00f0) +
+           x_thread[4 * i + 2] * (ws[i] & 0x0f00) +
+           x_thread[4 * i + 3] * (ws[i] & 0xf000));
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      x_thread += 8 * i;
+      w += 5 * i;
+
+      accum += (w[0] & 0x1f) * x_thread[0];
+      accum += (w[0] & 0xe0) * x_thread[1];
+      accum += (w[1] & 0x3) * (x_thread[1] * 256.0f);
+      accum += (w[1] & 0x7c) * x_thread[2];
+      accum += (w[1] & 0x80) * x_thread[3];
+      accum += (w[2] & 0xf) * (x_thread[3] * 256.0f);
+      accum += (w[2] & 0xf0) * x_thread[4];
+      accum += (w[3] & 0x1) * (x_thread[4] * 256.0f);
+      accum += (w[3] & 0x3e) * x_thread[5];
+      accum += (w[3] & 0xc0) * x_thread[6];
+      accum += (w[4] & 0x7) * (x_thread[6] * 256.0f);
+      accum += (w[4] & 0xf8) * x_thread[7];
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      x_thread += 4 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x3f) * x_thread[0];
+
+      accum += (w[0] & 0xc0) * x_thread[1];
+      accum += (w[1] & 0x0f) * (x_thread[1] * 256.0f);
+
+      accum += (w[1] & 0xf0) * x_thread[2];
+      accum += (w[2] & 0x03) * (x_thread[2] * 256.0f);
+
+      accum += (w[2] & 0xfc) * x_thread[3];
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      accum += x_thread[i] * w[i];
+    }
+  }
+
+  return scale * accum + sum * bias;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline U qdot_safe(
+    const device uint8_t* w,
+    const thread U* x_thread,
+    U scale,
+    U bias,
+    U sum,
+    int N) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  U accum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < (N / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (w[i] & 0x03) +
+           x_thread[4 * i + 1] * (w[i] & 0x0c) +
+           x_thread[4 * i + 2] * (w[i] & 0x30) +
+           x_thread[4 * i + 3] * (w[i] & 0xc0));
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (N / 8); i++) {
+      x_thread += 8 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x07) * x_thread[0];
+      accum += (w[0] & 0x38) * x_thread[1];
+      accum += (w[0] & 0xc0) * x_thread[2];
+      accum += (w[1] & 0x01) * (x_thread[2] * 256.0f);
+
+      accum += (w[1] & 0x0e) * x_thread[3];
+      accum += (w[1] & 0x70) * x_thread[4];
+      accum += (w[1] & 0x80) * x_thread[5];
+      accum += (w[2] & 0x03) * (x_thread[5] * 256.0f);
+
+      accum += (w[2] & 0x1c) * x_thread[6];
+      accum += (w[2] & 0xe0) * x_thread[7];
+    }
+  }
+
+  else if (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (N / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * (ws[i] & 0x000f) +
+           x_thread[4 * i + 1] * (ws[i] & 0x00f0) +
+           x_thread[4 * i + 2] * (ws[i] & 0x0f00) +
+           x_thread[4 * i + 3] * (ws[i] & 0xf000));
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (N / 8); i++) {
+      x_thread += 8 * i;
+      w += 5 * i;
+
+      accum += (w[0] & 0x1f) * x_thread[0];
+      accum += (w[0] & 0xe0) * x_thread[1];
+      accum += (w[1] & 0x3) * (x_thread[1] * 256.0f);
+      accum += (w[1] & 0x7c) * x_thread[2];
+      accum += (w[1] & 0x80) * x_thread[3];
+      accum += (w[2] & 0xf) * (x_thread[3] * 256.0f);
+      accum += (w[2] & 0xf0) * x_thread[4];
+      accum += (w[3] & 0x1) * (x_thread[4] * 256.0f);
+      accum += (w[3] & 0x3e) * x_thread[5];
+      accum += (w[3] & 0xc0) * x_thread[6];
+      accum += (w[4] & 0x7) * (x_thread[6] * 256.0f);
+      accum += (w[4] & 0xf8) * x_thread[7];
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (N / 4); i++) {
+      x_thread += 4 * i;
+      w += 3 * i;
+
+      accum += (w[0] & 0x3f) * x_thread[0];
+
+      accum += (w[0] & 0xc0) * x_thread[1];
+      accum += (w[1] & 0x0f) * (x_thread[1] * 256.0f);
+
+      accum += (w[1] & 0xf0) * x_thread[2];
+      accum += (w[2] & 0x03) * (x_thread[2] * 256.0f);
+
+      accum += (w[2] & 0xfc) * x_thread[3];
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < N; i++) {
+      accum += x_thread[i] * w[i];
+    }
+  }
+
+  return scale * accum + sum * bias;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline void
+qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  if (bits == 2) {
+    U s[4] = {scale, scale / 4.0f, scale / 16.0f, scale / 64.0f};
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      result[4 * i] += x * (s[0] * (w[i] & 0x03) + bias);
+      result[4 * i + 1] += x * (s[1] * (w[i] & 0x0c) + bias);
+      result[4 * i + 2] += x * (s[2] * (w[i] & 0x30) + bias);
+      result[4 * i + 3] += x * (s[3] * (w[i] & 0xc0) + bias);
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      uint8_t w0 = w[3 * i];
+      uint8_t w1 = w[3 * i + 1];
+      uint8_t w2 = w[3 * i + 2];
+
+      result[8 * i] += x * ((w0 & 0x7) * scale + bias);
+      result[8 * i + 1] += x * (((w0 & 0x38) >> 3) * scale + bias);
+      result[8 * i + 2] +=
+          x * ((((w0 & 0xc0) >> 6) + ((w1 & 0x1) << 2)) * scale + bias);
+      result[8 * i + 3] += x * (((w1 & 0xe) >> 1) * scale + bias);
+      result[8 * i + 4] += x * (((w1 & 0x70) >> 4) * scale + bias);
+      result[8 * i + 5] +=
+          x * ((((w1 & 0x80) >> 7) + ((w2 & 0x3) << 1)) * scale + bias);
+      result[8 * i + 6] += x * (((w2 & 0x1c) >> 2) * scale + bias);
+      result[8 * i + 7] += x * (((w2 & 0xe0) >> 5) * scale + bias);
+    }
+  }
+
+  else if (bits == 4) {
+    U s[2] = {scale, scale / 16.0f};
+    for (int i = 0; i < (values_per_thread / 2); i++) {
+      result[2 * i] += x * (s[0] * (w[i] & 0x0f) + bias);
+      result[2 * i + 1] += x * (s[1] * (w[i] & 0xf0) + bias);
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (values_per_thread / 8); i++) {
+      uint8_t w0 = w[5 * i];
+      uint8_t w1 = w[5 * i + 1];
+      uint8_t w2 = w[5 * i + 2];
+      uint8_t w3 = w[5 * i + 3];
+      uint8_t w4 = w[5 * i + 4];
+      result[8 * i] += x * ((w0 & 0x1f) * scale + bias);
+      result[8 * i + 1] +=
+          x * ((((w0 & 0xe0) >> 5) + ((w1 & 0x3) << 3)) * scale + bias);
+      result[8 * i + 2] += x * (((w1 & 0x7c) >> 2) * scale + bias);
+      result[8 * i + 3] +=
+          x * ((((w1 & 0x80) >> 7) + ((w2 & 0xf) << 1)) * scale + bias);
+      result[8 * i + 4] +=
+          x * ((((w2 & 0xf0) >> 4) + ((w3 & 0x1) << 4)) * scale + bias);
+      result[8 * i + 5] += x * (((w3 & 0x3e) >> 1) * scale + bias);
+      result[8 * i + 6] +=
+          x * ((((w3 & 0xc0) >> 6) + ((w4 & 0x7) << 2)) * scale + bias);
+      result[8 * i + 7] += x * (((w4 & 0xf8) >> 3) * scale + bias);
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      uint8_t w0 = w[3 * i];
+      uint8_t w1 = w[3 * i + 1];
+      uint8_t w2 = w[3 * i + 2];
+
+      result[4 * i] += x * ((w0 & 0x3f) * scale + bias);
+      result[4 * i + 1] +=
+          x * ((((w0 >> 6) & 0x03) + ((w1 & 0x0f) << 2)) * scale + bias);
+      result[4 * i + 2] +=
+          x * ((((w1 >> 4) & 0x0f) + ((w2 & 0x03) << 4)) * scale + bias);
+      result[4 * i + 3] += x * (((w2 >> 2) & 0x3f) * scale + bias);
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      result[i] += x * (scale * w[i] + bias);
+    }
+  }
+}
+
+template <typename U, int N, int bits>
+inline void
+dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) {
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  if (bits == 2) {
+    U s[4] = {
+        scale,
+        scale / static_cast<U>(4.0f),
+        scale / static_cast<U>(16.0f),
+        scale / static_cast<U>(64.0f)};
+    for (int i = 0; i < (N / 4); i++) {
+      w_local[4 * i] = s[0] * (w[i] & 0x03) + bias;
+      w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias;
+      w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias;
+      w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias;
+    }
+  }
+
+  else if (bits == 3) {
+    for (int i = 0; i < (N / 8); i++) {
+      w_local += 8 * i;
+      w += 3 * i;
+
+      w_local[0] = (w[0] & 0x7) * scale + bias;
+      w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias;
+      w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias;
+      w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias;
+      w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias;
+      w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias;
+      w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias;
+      w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias;
+    }
+  }
+
+  else if (bits == 4) {
+    U s[2] = {scale, scale / static_cast<U>(16.0f)};
+    for (int i = 0; i < (N / 2); i++) {
+      w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias;
+      w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias;
+    }
+  }
+
+  else if (bits == 5) {
+    for (int i = 0; i < (N / 8); i++) {
+      w_local += 8 * i;
+      w += 5 * i;
+
+      w_local[0] = (w[0] & 0x1f) * scale + bias;
+      w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias;
+      w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias;
+      w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias;
+      w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias;
+      w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias;
+      w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias;
+      w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias;
+    }
+  }
+
+  else if (bits == 6) {
+    for (int i = 0; i < (N / 4); i++) {
+      w_local += 4 * i;
+      w += 3 * i;
+      w_local[0] = (w[0] & 0x3f) * scale + bias;
+      w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias;
+      w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias;
+      w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias;
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < N; i++) {
+      w_local[i] = scale * w[i] + bias;
+    }
+  }
+}
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short group_size,
+    short bits>
+struct QuantizedBlockLoader {
+  static_assert(
+      BCOLS <= group_size,
+      "The group size should be larger than the columns");
+  static_assert(
+      group_size % BCOLS == 0,
+      "The group size should be divisible by the columns");
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  MLX_MTL_CONST short pack_factor = get_pack_factor<bits, 8>();
+  MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack<bits>();
+  MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
+  MLX_MTL_CONST short n_reads =
+      (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
+  MLX_MTL_CONST short group_steps = group_size / BCOLS;
+
+  const int src_ld;
+  const int tile_stride;
+  short group_step_cnt;
+  const int group_stride;
+
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  threadgroup T* dst;
+  const device uint8_t* src;
+  const device T* scales;
+  const device T* biases;
+
+  QuantizedBlockLoader(
+      const device uint8_t* src_,
+      const device T* scales_,
+      const device T* biases_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(
+            reduction_dim ? BCOLS_PACKED * bytes_per_pack
+                          : BROWS * src_ld * bytes_per_pack / pack_factor),
+        group_step_cnt(0),
+        group_stride(BROWS * src_ld / group_size),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(n_reads * thread_idx / BCOLS_PACKED),
+        bj((n_reads * thread_idx) % BCOLS_PACKED),
+        dst(dst_ + bi * dst_ld + bj * pack_factor),
+        src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
+            bj * bytes_per_pack),
+        scales(scales_ + bi * src_ld / group_size),
+        biases(biases_ + bi * src_ld / group_size) {}
+
+  void load_unsafe() const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    T scale = *scales;
+    T bias = *biases;
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, pack_factor, bits>(
+          src + i * bytes_per_pack, scale, bias, dst + i * pack_factor);
+    }
+  }
+
+  void load_safe(short2 src_tile_dim) const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    T scale = *scales;
+    T bias = *biases;
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, pack_factor, bits>(
+          (device uint8_t*)(src + i * bytes_per_pack),
+          scale,
+          bias,
+          dst + i * pack_factor);
+    }
+  }
+
+  void next() {
+    src += tile_stride;
+    if (reduction_dim == 1) {
+      if (group_steps > 1) {
+        group_step_cnt++;
+        if (group_step_cnt == group_steps) {
+          group_step_cnt = 0;
+          scales++;
+          biases++;
+        }
+      } else {
+        scales++;
+        biases++;
+      }
+    } else {
+      scales += group_stride;
+      biases += group_stride;
+    }
+  }
+};
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short bits>
+struct QuantizedBlockLoader<
+    T,
+    BROWS,
+    BCOLS,
+    dst_ld,
+    reduction_dim,
+    tgp_size,
+    32,
+    bits> {
+  MLX_MTL_CONST short group_size = 32;
+
+  static_assert(
+      BCOLS % group_size == 0,
+      "The group size should be divisible by the columns");
+  static_assert(
+      bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
+          bits == 8,
+      "Template undefined for bits not in {2, 3, 4, 5, 6, 8}");
+
+  MLX_MTL_CONST short pack_factor = get_pack_factor<bits, 8>();
+  MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack<bits>();
+  MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
+  MLX_MTL_CONST short n_reads =
+      (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
+  MLX_MTL_CONST short n_groups = BCOLS / group_size;
+
+  static_assert(
+      (BCOLS_PACKED / n_reads) == n_groups,
+      "Other configurations are not yet supported");
+
+  const int src_ld;
+  const int tile_stride;
+  const int group_stride;
+
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  const short group_id;
+
+  threadgroup T* dst;
+  const device uint8_t* src;
+  const device T* scales;
+  const device T* biases;
+
+  QuantizedBlockLoader(
+      const device uint8_t* src_,
+      const device T* scales_,
+      const device T* biases_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(
+            reduction_dim ? BCOLS_PACKED * bytes_per_pack
+                          : BROWS * src_ld * bytes_per_pack / pack_factor),
+        group_stride(BROWS * src_ld / group_size),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(n_reads * thread_idx / BCOLS_PACKED),
+        bj((n_reads * thread_idx) % BCOLS_PACKED),
+        group_id((bj * pack_factor) / group_size),
+        dst(dst_ + bi * dst_ld + bj * pack_factor),
+        src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
+            bj * bytes_per_pack),
+        scales(scales_ + bi * src_ld / group_size + group_id),
+        biases(biases_ + bi * src_ld / group_size + group_id) {}
+
+  void load_unsafe() const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    T scale = *scales;
+    T bias = *biases;
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, pack_factor, bits>(
+          src + i * bytes_per_pack, scale, bias, dst + i * pack_factor);
+    }
+  }
+
+  void load_safe(short2 src_tile_dim) const {
+    if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
+      return;
+    }
+
+    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
+      for (int i = 0; i < n_reads * pack_factor; i++) {
+        dst[i] = T(0);
+      }
+      return;
+    }
+
+    T scale = *scales;
+    T bias = *biases;
+    for (int i = 0; i < n_reads; i++) {
+      dequantize<T, pack_factor, bits>(
+          (device uint8_t*)(src + i * bytes_per_pack),
+          scale,
+          bias,
+          dst + i * pack_factor);
+    }
+  }
+
+  void next() {
+    src += tile_stride;
+    if (reduction_dim == 1) {
+      // if (group_steps > 1) {
+      //   group_step_cnt++;
+      //   if (group_step_cnt == group_steps) {
+      //     group_step_cnt = 0;
+      //     scales++;
+      //     biases++;
+      //   }
+      // } else {
+      scales += n_groups;
+      biases += n_groups;
+      // }
+    } else {
+      scales += n_groups * group_stride;
+      biases += n_groups * group_stride;
+    }
+  }
+};
+
+template <typename T>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device T*& scales,
+    const device T*& biases,
+    device T*& y,
+    int output_stride,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int64_t* b_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx = tid.z;
+  uint32_t w_idx = tid.z;
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+    biases += w_idx * b_strides[0];
+  } else {
+    ulong3 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, b_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+    biases += idx.z;
+  }
+  y += tid.z * output_stride;
+}
+
+template <typename T>
+METAL_FUNC void adjust_matrix_offsets(
+    const device T*& x,
+    const device uint32_t*& w,
+    const device T*& scales,
+    const device T*& biases,
+    const device uint32_t* lhs_indices,
+    const device uint32_t* rhs_indices,
+    device T*& y,
+    int output_stride,
+    const constant int& batch_ndims,
+    const constant int* batch_shape,
+    const constant int64_t* lhs_strides,
+    const constant int64_t* rhs_strides,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    const constant int64_t* b_strides,
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  // Set the input/output matrices
+  uint32_t x_idx;
+  uint32_t w_idx;
+  if (batch_ndims == 1) {
+    x_idx = lhs_indices[tid.z * lhs_strides[0]];
+    w_idx = rhs_indices[tid.z * rhs_strides[0]];
+  } else {
+    ulong2 idx = elem_to_loc_broadcast(
+        tid.z, batch_shape, lhs_strides, rhs_strides, batch_ndims);
+    x_idx = lhs_indices[idx.x];
+    w_idx = rhs_indices[idx.y];
+  }
+  if (x_batch_ndims == 1) {
+    x += x_idx * x_strides[0];
+  } else {
+    x += elem_to_loc(x_idx, x_shape, x_strides, x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += w_idx * w_strides[0];
+    scales += w_idx * s_strides[0];
+    biases += w_idx * b_strides[0];
+  } else {
+    ulong3 idx = elem_to_loc_broadcast(
+        w_idx, w_shape, w_strides, s_strides, b_strides, w_batch_ndims);
+    w += idx.x;
+    scales += idx.y;
+    biases += idx.z;
+  }
+  y += tid.z * output_stride;
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2>
+METAL_FUNC void qmm_t_nax_tgp_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    threadgroup T* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  (void)lid;
+
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      BN,
+      BK,
+      BK_padded,
+      1,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  // Set the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+
+  auto wl = (const device uint8_t*)w;
+
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * K_w;
+  scales += y_col * K_g;
+  biases += y_col * K_g;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the weight loader
+  loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid);
+
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  const short tm = SM * (simd_gid / WN);
+  const short tn = SN * (simd_gid % WN);
+
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = true;
+
+  const short sgp_sm = min(SM, short(M - (y_row + tm)));
+  const bool is_unaligned_sm = (sgp_sm != SM);
+
+  const short sgp_sn = aligned_N ? SN : min(SN, short(N - (y_col + tn)));
+
+  const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col)));
+  const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN);
+
+  using AccumType = float;
+
+  using ASubTile = NAXSubTile<T, UM, UK>;
+  using BSubTile = NAXSubTile<T, UN, UK>;
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  Dtile.clear();
+
+  x += tm * K;
+
+  dispatch_bool(!is_unaligned_sm, [&](auto kAlignedM) {
+    dispatch_bool(aligned_N || !is_unaligned_bn, [&](auto kAlignedN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if constexpr (kAlignedN.value) {
+          loader_w.load_unsafe();
+        } else {
+          loader_w.load_safe(short2(BK, tgp_bn));
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        STEEL_PRAGMA_NO_UNROLL
+        for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+          NAXTile<T, TM, TK, ASubTile> Atile;
+          NAXTile<T, TN, TK, BSubTile> Btile;
+
+          volatile int compiler_barrier;
+
+          if constexpr (kAlignedM.value) {
+            Atile.load(x + kk1, K);
+          } else {
+            Atile.load_safe(x + kk1, K, short2(SK, sgp_sm));
+          }
+
+          Btile.template load<T, BK_padded, 1>(Ws + tn * BK_padded + kk1);
+
+          tile_matmad_nax(
+              Dtile,
+              Atile,
+              metal::bool_constant<transpose_a>{},
+              Btile,
+              metal::bool_constant<transpose_b>{});
+
+          (void)compiler_barrier;
+        }
+
+        x += BK;
+        loader_w.next();
+      }
+
+      // Store results to device memory
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      if constexpr (kAlignedM.value && kAlignedN.value) {
+        Dtile.store(y + tm * N + tn, N);
+      } else if (kAlignedM.value && sgp_sn == SN) {
+        Dtile.store(y + tm * N + tn, N);
+      } else {
+        Dtile.store_safe(y + tm * N + tn, N, short2(sgp_sn, sgp_sm));
+      }
+    });
+  });
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2>
+METAL_FUNC void qmm_n_nax_tgp_impl(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    threadgroup T* Ws,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+  (void)M;
+
+  static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
+  static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
+
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      BK,
+      BN,
+      BN_padded,
+      0,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  // Set the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+
+  auto wl = (const device uint8_t*)w;
+
+  x += y_row * static_cast<int64_t>(K);
+  wl += y_col * K_w;
+  scales += y_col * K_g;
+  biases += y_col * K_g;
+  y += y_row * static_cast<int64_t>(N) + y_col;
+
+  // Make the x loader and mma operation
+  // const short num_els = min(BM, M - y_row);
+  // const short num_outs = min(BN, N - y_col);
+  loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid);
+
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  const short tm = SM * (simd_gid / WN);
+  const short tn = SN * (simd_gid % WN);
+
+  const short ldb_tgp = BN_padded;
+
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = false;
+
+  using AccumType = float;
+
+  using ASubTile = NAXSubTile<T, UM, UK>;
+  using BSubTile = NAXSubTile<T, UK, UN>;
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  Dtile.clear();
+
+  x += tm * K;
+
+  for (int k = 0; k < K; k += BK) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    loader_w.load_unsafe();
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_NO_UNROLL
+    for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+      NAXTile<T, TM, TK, ASubTile> Atile;
+      NAXTile<T, TK, TN, BSubTile> Btile;
+
+      volatile int compiler_barrier;
+
+      Atile.load(x + kk1, K);
+      Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * ldb_tgp);
+
+      tile_matmad_nax(
+          Dtile,
+          Atile,
+          metal::bool_constant<transpose_a>{},
+          Btile,
+          metal::bool_constant<transpose_b>{});
+
+      (void)compiler_barrier;
+    }
+
+    x += BK;
+    loader_w.next();
+  }
+
+  // Store results to device memory
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  Dtile.store(y + tm * N + tn, N);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const bool batched,
+    const int BM = 64,
+    const int BK = 32,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2>
+[[kernel]] void affine_qmm_t_nax(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& K [[buffer(5)]],
+    const constant int& N [[buffer(6)]],
+    const constant int& M [[buffer(7)]],
+    const constant int& x_batch_ndims [[buffer(8)]],
+    const constant int* x_shape [[buffer(9)]],
+    const constant int64_t* x_strides [[buffer(10)]],
+    const constant int& w_batch_ndims [[buffer(11)]],
+    const constant int* w_shape [[buffer(12)]],
+    const constant int64_t* w_strides [[buffer(13)]],
+    const constant int64_t* s_strides [[buffer(14)]],
+    const constant int64_t* b_strides [[buffer(15)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  threadgroup T Ws[BN * BK_padded];
+
+  if (batched) {
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+  qmm_t_nax_tgp_impl<T, group_size, bits, aligned_N, BM, BK, BN, WM, WN>(
+      w, scales, biases, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool batched,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2>
+[[kernel]] void affine_qmm_n_nax(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& K [[buffer(5)]],
+    const constant int& N [[buffer(6)]],
+    const constant int& M [[buffer(7)]],
+    const constant int& x_batch_ndims [[buffer(8)]],
+    const constant int* x_shape [[buffer(9)]],
+    const constant int64_t* x_strides [[buffer(10)]],
+    const constant int& w_batch_ndims [[buffer(11)]],
+    const constant int* w_shape [[buffer(12)]],
+    const constant int64_t* w_strides [[buffer(13)]],
+    const constant int64_t* s_strides [[buffer(14)]],
+    const constant int64_t* b_strides [[buffer(15)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Ws[BK * BN_padded];
+
+  if (batched) {
+    adjust_matrix_offsets<T>(
+        x,
+        w,
+        scales,
+        biases,
+        y,
+        M * N,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        b_strides,
+        tid);
+  }
+
+  qmm_n_nax_tgp_impl<T, group_size, bits, BM, BK, BN, WM, WN>(
+      w, scales, biases, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const bool aligned_N,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2>
+[[kernel]] void affine_gather_qmm_t_nax(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    const device uint32_t* lhs_indices [[buffer(4)]],
+    const device uint32_t* rhs_indices [[buffer(5)]],
+    device T* y [[buffer(6)]],
+    const constant int& K [[buffer(7)]],
+    const constant int& N [[buffer(8)]],
+    const constant int& M [[buffer(9)]],
+    const constant int& x_batch_ndims [[buffer(10)]],
+    const constant int* x_shape [[buffer(11)]],
+    const constant int64_t* x_strides [[buffer(12)]],
+    const constant int& w_batch_ndims [[buffer(13)]],
+    const constant int* w_shape [[buffer(14)]],
+    const constant int64_t* w_strides [[buffer(15)]],
+    const constant int64_t* s_strides [[buffer(16)]],
+    const constant int64_t* b_strides [[buffer(17)]],
+    const constant int& batch_ndims [[buffer(18)]],
+    const constant int* batch_shape [[buffer(19)]],
+    const constant int64_t* lhs_strides [[buffer(20)]],
+    const constant int64_t* rhs_strides [[buffer(21)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+
+  threadgroup T Ws[BN * BK_padded];
+
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+  qmm_t_nax_tgp_impl<T, group_size, bits, aligned_N, BM, BK, BN, WM, WN>(
+      w, scales, biases, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    const int group_size,
+    const int bits,
+    const int BM = 64,
+    const int BK = 64,
+    const int BN = 64,
+    const int WM = 2,
+    const int WN = 2>
+[[kernel]] void affine_gather_qmm_n_nax(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    const device uint32_t* lhs_indices [[buffer(4)]],
+    const device uint32_t* rhs_indices [[buffer(5)]],
+    device T* y [[buffer(6)]],
+    const constant int& K [[buffer(7)]],
+    const constant int& N [[buffer(8)]],
+    const constant int& M [[buffer(9)]],
+    const constant int& x_batch_ndims [[buffer(10)]],
+    const constant int* x_shape [[buffer(11)]],
+    const constant int64_t* x_strides [[buffer(12)]],
+    const constant int& w_batch_ndims [[buffer(13)]],
+    const constant int* w_shape [[buffer(14)]],
+    const constant int64_t* w_strides [[buffer(15)]],
+    const constant int64_t* s_strides [[buffer(16)]],
+    const constant int64_t* b_strides [[buffer(17)]],
+    const constant int& batch_ndims [[buffer(18)]],
+    const constant int* batch_shape [[buffer(19)]],
+    const constant int64_t* lhs_strides [[buffer(20)]],
+    const constant int64_t* rhs_strides [[buffer(21)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint lid [[thread_index_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  threadgroup T Ws[BK * BN_padded];
+
+  adjust_matrix_offsets<T>(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      y,
+      M * N,
+      batch_ndims,
+      batch_shape,
+      lhs_strides,
+      rhs_strides,
+      x_batch_ndims,
+      x_shape,
+      x_strides,
+      w_batch_ndims,
+      w_shape,
+      w_strides,
+      s_strides,
+      b_strides,
+      tid);
+  qmm_n_nax_tgp_impl<T, group_size, bits, BM, BK, BN, WM, WN>(
+      w, scales, biases, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
+}
+
+template <
+    typename T,
+    int group_size,
+    int bits,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose>
+[[kernel]] void affine_gather_qmm_rhs_nax(
+    const device T* x [[buffer(0)]],
+    const device uint32_t* w [[buffer(1)]],
+    const device T* scales [[buffer(2)]],
+    const device T* biases [[buffer(3)]],
+    const device uint32_t* indices [[buffer(4)]],
+    device T* y [[buffer(5)]],
+    const constant int& M [[buffer(6)]],
+    const constant int& N [[buffer(7)]],
+    const constant int& K [[buffer(8)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]]) {
+  constexpr int pack_factor = get_pack_factor<bits, 8>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+  constexpr int BK_padded = (BK + 16 / sizeof(T));
+  constexpr int BN_padded = (BN + 16 / sizeof(T));
+
+  using loader_w_t = QuantizedBlockLoader<
+      T,
+      transpose ? BN : BK,
+      transpose ? BK : BN,
+      transpose ? BK_padded : BN_padded,
+      transpose,
+      WM * WN * SIMD_SIZE,
+      group_size,
+      bits>;
+
+  threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded];
+
+  // Compute the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int N_w = N * bytes_per_pack / pack_factor;
+  const int N_g = N / group_size;
+  const int K_it = K / BK;
+  const size_t stride_w = transpose ? N * K_w : K * N_w;
+  const size_t stride_s = transpose ? N * K_g : K * N_g;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
+  const size_t y_row_long = size_t(y_row);
+  const size_t y_col_long = size_t(y_col);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, M - y_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, N - y_col));
+
+  // Calculate the final tiles in the case that K is not aligned
+  const int k_remain = K - K_it * BK;
+  const short2 tile_w =
+      transpose ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+  // Move x and output to the correct block
+  auto wl = (const device uint8_t*)w;
+  x += y_row_long * K;
+  y += y_row_long * N + y_col_long;
+  wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor;
+  scales += transpose ? y_col_long * K_g : y_col / group_size;
+  biases += transpose ? y_col_long * K_g : y_col / group_size;
+
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  const short tm = SM * (simd_group_id / WN);
+  const short tn = SN * (simd_group_id % WN);
+
+  const short sgp_sm =
+      align_M ? SM : min(SM, short(max(0, (M - (y_row + tm)))));
+  const short sgp_sn =
+      align_N ? SN : min(SN, short(max(0, (N - (y_col + tn)))));
+
+  const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
+  const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN);
+
+  constexpr short BR = transpose ? TN : TK;
+  constexpr short BC = transpose ? TK : TN;
+
+  using AccumType = float;
+
+  using ASubTile = NAXSubTile<T, UM, UK>;
+  using BSubTile = NAXSubTile<T, transpose ? UN : UK, transpose ? UK : UN>;
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+
+  // Do as many matmuls as necessary
+  uint32_t index;
+  short offset;
+  uint32_t index_next = indices[y_row];
+  short offset_next = 0;
+  int n = 0;
+  while (n < tgp_bm) {
+    n++;
+    offset = offset_next;
+    index = index_next;
+    offset_next = tgp_bm;
+    for (; n < tgp_bm; n++) {
+      if (indices[y_row + n] != index) {
+        offset_next = n;
+        index_next = indices[y_row + n];
+        break;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_none);
+
+    NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+    Dtile.clear();
+
+    const device T* xn = x + tm * K;
+
+    // Prepare threadgroup loading operations
+    thread loader_w_t loader_w(
+        wl + index * stride_w,
+        scales + index * stride_s,
+        biases + index * stride_s,
+        transpose ? K : N,
+        Ws,
+        simd_group_id,
+        simd_lane_id);
+
+    dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+      dispatch_bool(align_N || !is_unaligned_bn, [&](auto kAlignedN) {
+        for (int k = 0; k < K_it; k++) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          if constexpr (kAlignedN.value) {
+            loader_w.load_unsafe();
+          } else {
+            loader_w.load_safe(
+                transpose ? short2(BK, tgp_bn) : short2(tgp_bn, BK));
+          }
+
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+
+          STEEL_PRAGMA_NO_UNROLL
+          for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+            NAXTile<T, TM, TK, ASubTile> Atile;
+            NAXTile<T, BR, BC, BSubTile> Btile;
+
+            volatile int compiler_barrier;
+
+            if constexpr (kAlignedM.value) {
+              Atile.load(xn + kk1, K);
+            } else {
+              Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm));
+            }
+
+            if constexpr (transpose) {
+              Btile.template load<T, BK_padded, 1>(Ws + tn * BK_padded + kk1);
+            } else {
+              Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * BN_padded);
+            }
+
+            tile_matmad_nax(
+                Dtile,
+                Atile,
+                metal::bool_constant<false>{},
+                Btile,
+                metal::bool_constant<transpose>{});
+
+            (void)compiler_barrier;
+          }
+
+          xn += BK;
+          loader_w.next();
+        }
+
+        if (!align_K) {
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+          loader_w.load_safe(tile_w);
+          threadgroup_barrier(mem_flags::mem_threadgroup);
+
+          STEEL_PRAGMA_NO_UNROLL
+          for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+            NAXTile<T, TM, TK, ASubTile> Atile;
+            NAXTile<T, BR, BC, BSubTile> Btile;
+
+            volatile int compiler_barrier;
+
+            const short psk = min(int(SK), max(0, (BK - kk1)));
+            Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm));
+
+            if constexpr (transpose) {
+              Btile.template load<T, BK_padded, 1>(Ws + tn * BK_padded + kk1);
+            } else {
+              Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * BN_padded);
+            }
+
+            tile_matmad_nax(
+                Dtile,
+                Atile,
+                metal::bool_constant<false>{},
+                Btile,
+                metal::bool_constant<transpose>{});
+
+            (void)compiler_barrier;
+          }
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm));
+        const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm));
+
+        // Store results to device memory
+        if constexpr (kAlignedN.value) {
+          if (m_lo_lim == 0 && m_hi_lim == SM) {
+            Dtile.store(y + tm * N + tn, N);
+          } else {
+            Dtile.store_slice(
+                y + tm * N + tn, N, short2(0, m_lo_lim), short2(SN, m_hi_lim));
+          }
+        } else {
+          Dtile.store_slice(
+              y + tm * N + tn,
+              N,
+              short2(0, m_lo_lim),
+              short2(sgp_sn, m_hi_lim));
+        }
+      });
+    });
+  }
+}
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_nax.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_nax.metal
new file mode 100644
index 0000000..5a9c9fb
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_nax.metal
@@ -0,0 +1,106 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/steel/gemm/nax.h"
+#include "mlx/backend/metal/kernels/steel/gemm/loader.h"
+#include "mlx/backend/metal/kernels/quantized_nax.h"
+
+#define instantiate_quantized(name, type, group_size, bits, bm, bn, bk, wm, wn)  \
+  instantiate_kernel(                                                    \
+      #name "_" #type "_gs_" #group_size "_b_" #bits,                    \
+      name,                                                              \
+      type,                                                              \
+      group_size,                                                        \
+      bits, bm, bk, bn, wm, wn)
+
+#define instantiate_quantized_batched(name, type, group_size, bits, bm, bn, bk, wm, wn, batched)     \
+  instantiate_kernel(                                                    \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched, \
+      name,                                                              \
+      type,                                                              \
+      group_size,                                                        \
+      bits,                                                              \
+      batched, bm, bk, bn, wm, wn)
+
+#define instantiate_quantized_aligned(name, type, group_size, bits, bm, bn, bk, wm, wn, aligned)     \
+  instantiate_kernel(                                                                     \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned, \
+      name,                                                                  \
+      type,                                                                  \
+      group_size,                                                            \
+      bits,                                                                  \
+      aligned, bm, bk, bn, wm, wn)
+
+#define instantiate_quantized_aligned_batched(name, type, group_size, bits, bm, bn, bk, wm, wn, aligned, batched)     \
+  instantiate_kernel(                                                                     \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched, \
+      name,                                                                  \
+      type,                                                                  \
+      group_size,                                                            \
+      bits,                                                                  \
+      aligned,                                                               \
+      batched, bm, bk, bn, wm, wn)
+
+#define instantiate_gather_qmm_rhs(func, name, type, group_size, bits, bm, bn, bk, wm, wn, transpose)        \
+  instantiate_kernel(                                                                                        \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \
+      func,                                                         \
+      type,                                                         \
+      group_size,                                                   \
+      bits,                                                         \
+      bm,                                                           \
+      bn,                                                           \
+      bk,                                                           \
+      wm,                                                           \
+      wn,                                                           \
+      transpose)
+
+#define instantiate_quantized_batched_wrap(name, type, group_size, bits) \
+  instantiate_quantized_batched(name, type, group_size, bits, 64, 64, 64, 2, 2, 1)      \
+  instantiate_quantized_batched(name, type, group_size, bits, 64, 64, 64, 2, 2, 0)
+
+#define instantiate_quantized_all_batched(type, group_size, bits) \
+  instantiate_quantized_batched_wrap(affine_qmm_n_nax, type, group_size, bits)
+
+
+#define instantiate_quantized_all_single(type, group_size, bits) \
+  instantiate_quantized(affine_gather_qmm_n_nax, type, group_size, bits, 64, 64, 64, 2, 2)
+
+#define instantiate_quantized_all_aligned(type, group_size, bits)   \
+  instantiate_quantized_aligned(affine_gather_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true) \
+  instantiate_quantized_aligned(affine_gather_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 1) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 0) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 1) \
+  instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0)
+
+#define instantiate_quantized_all_rhs(type, group_size, bits) \
+  instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 64, 64, 64, 2, 2, true) \
+  instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false)
+
+#define instantiate_quantized_funcs(type, group_size, bits) \
+  instantiate_quantized_all_batched(type, group_size, bits) \
+  instantiate_quantized_all_aligned(type, group_size, bits) \
+  instantiate_quantized_all_rhs(type, group_size, bits)
+
+#define instantiate_quantized_types(group_size, bits)       \
+  instantiate_quantized_funcs(float, group_size, bits)      \
+  instantiate_quantized_funcs(float16_t, group_size, bits)  \
+  instantiate_quantized_funcs(bfloat16_t, group_size, bits)  
+
+#define instantiate_quantized_groups(bits) \
+  instantiate_quantized_types(128, bits)   \
+  instantiate_quantized_types(64, bits)    \
+  instantiate_quantized_types(32, bits)
+
+#define instantiate_quantized_all() \
+  instantiate_quantized_groups(2) \
+  instantiate_quantized_groups(3) \
+  instantiate_quantized_groups(4) \
+  instantiate_quantized_groups(5) \
+  instantiate_quantized_groups(6) \
+  instantiate_quantized_groups(8)
+
+instantiate_quantized_all() // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_utils.h b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_utils.h
new file mode 100644
index 0000000..38253f8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/quantized_utils.h
@@ -0,0 +1,90 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+template <typename T, typename mma_t, typename loader_a_t, typename loader_b_t>
+METAL_FUNC void gemm_loop_aligned(
+    threadgroup T* As,
+    threadgroup T* Bs,
+    thread mma_t& mma_op,
+    thread loader_a_t& loader_a,
+    thread loader_b_t& loader_b,
+    const int k_iterations) {
+  for (int k = 0; k < k_iterations; k++) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load elements into threadgroup memory
+    loader_a.load_unsafe();
+    loader_b.load_unsafe();
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(As, Bs);
+
+    // Prepare for next iteration
+    loader_a.next();
+    loader_b.next();
+  }
+}
+
+template <
+    bool rows_aligned,
+    bool cols_aligned,
+    bool transpose,
+    typename T,
+    typename mma_t,
+    typename loader_a_t,
+    typename loader_b_t>
+METAL_FUNC void gemm_loop_unaligned(
+    threadgroup T* As,
+    threadgroup T* Bs,
+    thread mma_t& mma_op,
+    thread loader_a_t& loader_a,
+    thread loader_b_t& loader_b,
+    const int k_iterations,
+    const short tgp_bm,
+    const short tgp_bn,
+    const short tgp_bk) {
+  for (int k = 0; k < k_iterations; k++) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load elements into threadgroup memory
+    if (rows_aligned) {
+      loader_a.load_unsafe();
+    } else {
+      loader_a.load_safe(short2(tgp_bk, tgp_bm));
+    }
+    if (cols_aligned) {
+      loader_b.load_unsafe();
+    } else {
+      loader_b.load_safe(
+          transpose ? short2(tgp_bk, tgp_bn) : short2(tgp_bn, tgp_bk));
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(As, Bs);
+
+    // Prepare for next iteration
+    loader_a.next();
+    loader_b.next();
+  }
+}
+
+template <typename T, typename mma_t, typename loader_a_t, typename loader_b_t>
+METAL_FUNC void gemm_loop_finalize(
+    threadgroup T* As,
+    threadgroup T* Bs,
+    thread mma_t& mma_op,
+    thread loader_a_t& loader_a,
+    thread loader_b_t& loader_b,
+    const short2 tile_a,
+    const short2 tile_b) {
+  loader_a.load_safe(tile_a);
+  loader_b.load_safe(tile_b);
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  mma_op.mma(As, Bs);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/random.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/random.metal
new file mode 100644
index 0000000..ccbd464
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/random.metal
@@ -0,0 +1,103 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+static constexpr constant uint32_t rotations[2][4] = {
+    {13, 15, 26, 6},
+    {17, 29, 16, 24}};
+
+union rbits {
+  uint2 val;
+  uchar4 bytes[2];
+};
+
+rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
+  uint4 ks = {key.x, key.y, key.x ^ key.y ^ 0x1BD11BDA};
+
+  rbits v;
+  v.val.x = count.x + ks[0];
+  v.val.y = count.y + ks[1];
+
+  for (int i = 0; i < 5; ++i) {
+    for (auto r : rotations[i % 2]) {
+      v.val.x += v.val.y;
+      v.val.y = (v.val.y << r) | (v.val.y >> (32 - r));
+      v.val.y ^= v.val.x;
+    }
+    v.val.x += ks[(i + 1) % 3];
+    v.val.y += ks[(i + 2) % 3] + i + 1;
+  }
+
+  return v;
+}
+
+[[kernel]] void rbitsc(
+    device const uint32_t* keys,
+    device char* out,
+    constant const bool& odd,
+    constant const uint& bytes_per_key,
+    uint2 grid_dim [[threads_per_grid]],
+    uint2 index [[thread_position_in_grid]]) {
+  auto kidx = 2 * index.x;
+  auto key = uint2(keys[kidx], keys[kidx + 1]);
+  auto half_size = grid_dim.y - odd;
+  out += index.x * bytes_per_key;
+  bool drop_last = odd && (index.y == half_size);
+  auto bits = threefry2x32_hash(
+      key, uint2(index.y, drop_last ? 0 : index.y + grid_dim.y));
+  size_t idx = size_t(index.y) << 2;
+  for (int i = 0; i < 4; ++i) {
+    out[idx + i] = bits.bytes[0][i];
+  }
+  if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index.y) + grid_dim.y) << 2;
+    if ((index.y + 1) == half_size && (bytes_per_key % 4) > 0) {
+      int edge_bytes = (bytes_per_key % 4);
+      for (int i = 0; i < edge_bytes; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    } else {
+      for (int i = 0; i < 4; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    }
+  }
+}
+
+[[kernel]] void rbits(
+    device const uint32_t* keys,
+    device char* out,
+    constant const bool& odd,
+    constant const uint& bytes_per_key,
+    constant const int& ndim,
+    constant const int* key_shape,
+    constant const int64_t* key_strides,
+    uint2 grid_dim [[threads_per_grid]],
+    uint2 index [[thread_position_in_grid]]) {
+  auto kidx = 2 * index.x;
+  auto k1_elem = elem_to_loc(kidx, key_shape, key_strides, ndim);
+  auto k2_elem = elem_to_loc(kidx + 1, key_shape, key_strides, ndim);
+  auto key = uint2(keys[k1_elem], keys[k2_elem]);
+  auto half_size = grid_dim.y - odd;
+  out += size_t(index.x) * bytes_per_key;
+  bool drop_last = odd && (index.y == half_size);
+  auto bits = threefry2x32_hash(
+      key, uint2(index.y, drop_last ? 0 : index.y + grid_dim.y));
+  size_t idx = size_t(index.y) << 2;
+  for (int i = 0; i < 4; ++i) {
+    out[idx + i] = bits.bytes[0][i];
+  }
+  if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index.y) + grid_dim.y) << 2;
+    if ((index.y + 1) == half_size && (bytes_per_key % 4) > 0) {
+      int edge_bytes = (bytes_per_key % 4);
+      for (int i = 0; i < edge_bytes; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    } else {
+      for (int i = 0; i < 4; ++i) {
+        out[idx + i] = bits.bytes[1][i];
+      }
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduce.h b/macos/third_party/mlx/mlx/backend/metal/kernels/reduce.h
new file mode 100644
index 0000000..ee5c3d5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduce.h
@@ -0,0 +1,5 @@
+#pragma once
+#include "mlx/backend/metal/kernels/reduction/reduce_all.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_col.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_init.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_row.h"
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduce.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/reduce.metal
new file mode 100644
index 0000000..de5dfba
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduce.metal
@@ -0,0 +1,169 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_atomic>
+#include <metal_simdgroup>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/atomic.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+#include "mlx/backend/metal/kernels/reduce.h"
+
+#define instantiate_init_reduce(name, tname, type, op) \
+  instantiate_kernel("init_reduce_" #name #tname, init_reduce, type, op<type>)
+
+instantiate_init_reduce(and, bool_, bool, And)
+instantiate_init_reduce(or, bool_, bool, Or)
+
+#define instantiate_init_sum_prod(name, op)                 \
+  instantiate_init_reduce(name, int32, int32_t, op)         \
+  instantiate_init_reduce(name, int64, int64_t, op)         \
+  instantiate_init_reduce(name, float16, float16_t, op)     \
+  instantiate_init_reduce(name, bfloat16, bfloat16_t, op)   \
+  instantiate_init_reduce(name, float32, float, op)         \
+  instantiate_init_reduce(name, complex64, complex64_t, op)
+
+instantiate_init_sum_prod(sum, Sum)
+instantiate_init_sum_prod(prod, Prod)
+
+#define instantiate_init_min_max(name, op)                   \
+  instantiate_init_reduce(name, bool_, bool, op)             \
+  instantiate_init_reduce(name, int8, int8_t, op)            \
+  instantiate_init_reduce(name, int16, int16_t, op)          \
+  instantiate_init_reduce(name, int32, int32_t, op)          \
+  instantiate_init_reduce(name, int64, int64_t, op)          \
+  instantiate_init_reduce(name, uint8, uint8_t, op)          \
+  instantiate_init_reduce(name, uint16, uint16_t, op)        \
+  instantiate_init_reduce(name, uint32, uint32_t, op)        \
+  instantiate_init_reduce(name, uint64, uint64_t, op)        \
+  instantiate_init_reduce(name, float16, float16_t, op)      \
+  instantiate_init_reduce(name, bfloat16, bfloat16_t, op)    \
+  instantiate_init_reduce(name, float32, float, op)          \
+  instantiate_init_reduce(name, complex64, complex64_t, op)
+
+instantiate_init_min_max(min, Min)
+instantiate_init_min_max(max, Max)
+
+#define instantiate_all_reduce(name, itype, otype, op) \
+  instantiate_kernel("all_reduce_" #name,              \
+                     all_reduce,                       \
+                     itype, otype, op)
+
+#define instantiate_col_reduce_small(name, itype, otype, op, dim)          \
+  instantiate_kernel("col_reduce_small_" #dim "_reduce_" #name,            \
+                     col_reduce_small,                                     \
+                     itype, otype, op, int, dim)                           \
+  instantiate_kernel("col_reduce_longcolumn_" #dim "_reduce_" #name,       \
+                     col_reduce_longcolumn,                                \
+                     itype, otype, op, int, dim)                           \
+  instantiate_kernel("col_reduce_small_large_" #dim "_reduce_" #name,      \
+                     col_reduce_small,                                     \
+                     itype, otype, op, int64_t, dim)                       \
+  instantiate_kernel("col_reduce_longcolumn_large_" #dim "_reduce_" #name, \
+                     col_reduce_longcolumn,                                \
+                     itype, otype, op, int64_t, dim)
+
+#define instantiate_col_reduce_looped_tile(name, itype, otype, op, dim, bm, bn)        \
+  instantiate_kernel("col_reduce_looped_" #dim "_" #bm "_" #bn "_reduce_" #name,       \
+                     col_reduce_looped,                                                \
+                     itype, otype, op, int, dim, bm, bn)                               \
+  instantiate_kernel("col_reduce_looped_large_" #dim "_" #bm "_" #bn "_reduce_" #name, \
+                     col_reduce_looped,                                                \
+                     itype, otype, op, int64_t, dim, bm, bn)
+
+#define instantiate_col_reduce_2pass_tile(name, itype, otype, op, dim, bm, bn)        \
+  instantiate_kernel("col_reduce_2pass_" #dim "_" #bm "_" #bn "_reduce_" #name,       \
+                     col_reduce_2pass,                                                \
+                     itype, otype, op, int, dim, bm, bn)                              \
+  instantiate_kernel("col_reduce_2pass_large_" #dim "_" #bm "_" #bn "_reduce_" #name, \
+                     col_reduce_2pass,                                                \
+                     itype, otype, op, int64_t, dim, bm, bn)
+
+#define instantiate_col_reduce_looped(name, itype, otype, op, dim)        \
+  instantiate_col_reduce_looped_tile(name, itype, otype, op, dim, 32, 32) \
+  instantiate_col_reduce_2pass_tile(name, itype, otype, op, dim, 32, 32)
+
+#define instantiate_col_reduce_general(name, itype, otype, op) \
+  instantiate_col_reduce_small(name, itype, otype, op, 1)      \
+  instantiate_col_reduce_small(name, itype, otype, op, 2)      \
+  instantiate_col_reduce_small(name, itype, otype, op, 5)      \
+  instantiate_col_reduce_looped(name, itype, otype, op, 1)     \
+  instantiate_col_reduce_looped(name, itype, otype, op, 2)     \
+  instantiate_col_reduce_looped(name, itype, otype, op, 5)
+
+#define instantiate_row_reduce_small(name, itype, otype, op, dim)     \
+  instantiate_kernel("row_reduce_small_" #dim "_reduce_" #name,       \
+                     row_reduce_small,                                \
+                     itype, otype, op, int, dim)                      \
+  instantiate_kernel("row_reduce_small_large_" #dim "_reduce_" #name, \
+                     row_reduce_small,                                \
+                     itype, otype, op, int64_t, dim)
+
+#define instantiate_row_reduce_looped(name, itype, otype, op, dim)       \
+  instantiate_kernel("row_reduce_looped_" #dim "_reduce_" #name,         \
+                     row_reduce_looped,                                  \
+                     itype, otype, op, int, dim)                         \
+  instantiate_kernel("row_reduce_looped_large_" #dim "_reduce_" #name,   \
+                     row_reduce_looped,                                  \
+                     itype, otype, op, int64_t, dim)
+
+#define instantiate_row_reduce_general(name, itype, otype, op) \
+  instantiate_row_reduce_small(name, itype, otype, op, 1)      \
+  instantiate_row_reduce_small(name, itype, otype, op, 2)      \
+  instantiate_row_reduce_small(name, itype, otype, op, 5)      \
+  instantiate_row_reduce_looped(name, itype, otype, op, 1)     \
+  instantiate_row_reduce_looped(name, itype, otype, op, 2)     \
+  instantiate_row_reduce_looped(name, itype, otype, op, 5)     \
+  instantiate_kernel("row_reduce_simple_" #name,               \
+                     row_reduce_simple,                        \
+                     itype, otype, op)
+
+#define instantiate_reduce_functions(name, tname, itype, otype, op)    \
+  instantiate_all_reduce(name##tname, itype, otype, op<otype>)         \
+  instantiate_row_reduce_general(name##tname, itype, otype, op<otype>) \
+  instantiate_col_reduce_general(name##tname, itype, otype, op<otype>)
+
+#define instantiate_and_or(name, op)                           \
+  instantiate_reduce_functions(name, bool_, bool, bool, op)    \
+  instantiate_reduce_functions(name, int16, int16_t, bool, op) \
+  instantiate_reduce_functions(name, int32, int32_t, bool, op) \
+  instantiate_reduce_functions(name, int64, int64_t, bool, op)
+
+instantiate_and_or(and, And)
+instantiate_and_or(or, Or)
+
+#define instantiate_sum_prod(name, op)                                       \
+  instantiate_reduce_functions(name, uint8, uint8_t, int32_t, op)            \
+  instantiate_reduce_functions(name, uint16, uint16_t, uint32_t, op)         \
+  instantiate_reduce_functions(name, uint32, uint32_t, uint32_t, op)         \
+  instantiate_reduce_functions(name, uint64, uint64_t, uint64_t, op)         \
+  instantiate_reduce_functions(name, int8, int8_t, int32_t, op)              \
+  instantiate_reduce_functions(name, int16, int16_t, int32_t, op)            \
+  instantiate_reduce_functions(name, int32, int32_t, int32_t, op)            \
+  instantiate_reduce_functions(name, int64, int64_t, int64_t, op)            \
+  instantiate_reduce_functions(name, float16, float16_t, float16_t, op)      \
+  instantiate_reduce_functions(name, bfloat16, bfloat16_t, bfloat16_t, op)   \
+  instantiate_reduce_functions(name, float32, float, float, op)              \
+  instantiate_reduce_functions(name, complex64, complex64_t, complex64_t, op)
+
+instantiate_sum_prod(sum, Sum)
+instantiate_sum_prod(prod, Prod)
+
+#define instantiate_min_max(name, op)                                        \
+  instantiate_reduce_functions(name, int8, int8_t, int8_t, op)               \
+  instantiate_reduce_functions(name, int16, int16_t, int16_t, op)            \
+  instantiate_reduce_functions(name, int32, int32_t, int32_t, op)            \
+  instantiate_reduce_functions(name, int64, int64_t, int64_t, op)            \
+  instantiate_reduce_functions(name, uint8, uint8_t, uint8_t, op)            \
+  instantiate_reduce_functions(name, uint16, uint16_t, uint16_t, op)         \
+  instantiate_reduce_functions(name, uint32, uint32_t, uint32_t, op)         \
+  instantiate_reduce_functions(name, uint64, uint64_t, uint64_t, op)         \
+  instantiate_reduce_functions(name, float16, float16_t, float16_t, op)      \
+  instantiate_reduce_functions(name, bfloat16, bfloat16_t, bfloat16_t, op)   \
+  instantiate_reduce_functions(name, float32, float, float, op)              \
+  instantiate_reduce_functions(name, complex64, complex64_t, complex64_t, op)
+
+instantiate_min_max(min, Min)
+instantiate_min_max(max, Max)
+    // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduce_utils.h b/macos/third_party/mlx/mlx/backend/metal/kernels/reduce_utils.h
new file mode 100644
index 0000000..279a7af
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduce_utils.h
@@ -0,0 +1,6 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/atomic.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/ops.h b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/ops.h
new file mode 100644
index 0000000..11d8e83
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/ops.h
@@ -0,0 +1,275 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <metal_atomic>
+#include <metal_simdgroup>
+
+#define DEFINE_SIMD_REDUCE()                                             \
+  template <typename T, metal::enable_if_t<sizeof(T) < 8, bool> = true>  \
+  T simd_reduce(T val) {                                                 \
+    return simd_reduce_impl(val);                                        \
+  }                                                                      \
+                                                                         \
+  template <typename T, metal::enable_if_t<sizeof(T) == 8, bool> = true> \
+  T simd_reduce(T val) {                                                 \
+    for (short i = simd_size / 2; i > 0; i /= 2) {                       \
+      val = operator()(val, simd_shuffle_down(val, i));                  \
+    }                                                                    \
+    return val;                                                          \
+  }
+
+static constant constexpr const uint8_t simd_size = 32;
+
+union bool4_or_uint {
+  bool4 b;
+  unsigned int i;
+};
+
+struct None {
+  template <typename T>
+  void atomic_update(device mlx_atomic<T>* out, T val, size_t offset = 0) {
+    mlx_atomic_store_explicit(out, val, offset);
+  }
+};
+
+template <typename U = bool>
+struct And {
+  DEFINE_SIMD_REDUCE()
+
+  bool simd_reduce_impl(bool val) {
+    return simd_all(val);
+  }
+
+  static constexpr constant bool init = true;
+
+  void atomic_update(
+      device mlx_atomic<unsigned int>* out,
+      bool val,
+      int elem_idx,
+      size_t offset = 0) {
+    if (!val) {
+      bool4_or_uint update;
+      update.b = {true, true, true, true};
+      update.b[elem_idx] = false;
+      mlx_atomic_fetch_and_explicit(out, update.i, offset);
+    }
+  }
+
+  void
+  atomic_update(device mlx_atomic<bool>* out, bool val, size_t offset = 0) {
+    if (!val) {
+      mlx_atomic_store_explicit(out, val, offset);
+    }
+  }
+
+  // Non atomic update
+  void update(device bool* out, bool val) {
+    *out &= val;
+  }
+
+  // Operator
+  bool operator()(bool a, bool b) {
+    return a && b;
+  }
+};
+
+template <typename U = bool>
+struct Or {
+  DEFINE_SIMD_REDUCE()
+
+  bool simd_reduce_impl(bool val) {
+    return simd_any(val);
+  }
+
+  static constexpr constant bool init = false;
+
+  void atomic_update(
+      device mlx_atomic<unsigned int>* out,
+      bool val,
+      int elem_idx,
+      size_t offset = 0) {
+    if (val) {
+      bool4_or_uint update;
+      update.b = {false, false, false, false};
+      update.b[elem_idx] = true;
+      mlx_atomic_fetch_or_explicit(out, update.i, offset);
+    }
+  }
+
+  void
+  atomic_update(device mlx_atomic<bool>* out, bool val, size_t offset = 0) {
+    if (val) {
+      mlx_atomic_store_explicit(out, val, offset);
+    }
+  }
+
+  // Non atomic update
+  void update(device bool* out, bool val) {
+    *out |= val;
+  }
+
+  // Operator
+  bool operator()(bool a, bool b) {
+    return a || b;
+  }
+};
+
+template <typename U>
+struct Sum {
+  DEFINE_SIMD_REDUCE()
+
+  template <typename T>
+  T simd_reduce_impl(T val) {
+    return simd_sum(val);
+  }
+
+  static constexpr constant U init = U(0);
+
+  template <typename T>
+  void atomic_update(device mlx_atomic<T>* out, T val, size_t offset = 0) {
+    mlx_atomic_fetch_add_explicit(out, val, offset);
+  }
+
+  // Operator
+  U operator()(U a, U b) {
+    return a + b;
+  }
+};
+
+template <typename U>
+struct Prod {
+  DEFINE_SIMD_REDUCE()
+
+  template <typename T>
+  T simd_reduce_impl(T val) {
+    return simd_product(val);
+  }
+
+  static constexpr constant U init = U(1);
+
+  template <typename T>
+  void atomic_update(device mlx_atomic<T>* out, T val, size_t offset = 0) {
+    mlx_atomic_fetch_mul_explicit(out, val, offset);
+  }
+
+  // Operator
+  U operator()(U a, U b) {
+    return a * b;
+  }
+};
+
+template <typename U>
+struct Min {
+  DEFINE_SIMD_REDUCE()
+
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
+    return simd_min(val);
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
+    if (simd_any(val != val)) {
+      return static_cast<T>(NAN);
+    }
+    return simd_min(val);
+  }
+
+  static constexpr constant U init = Limits<U>::max;
+
+  template <typename T>
+  void atomic_update(device mlx_atomic<T>* out, T val, size_t offset = 0) {
+    mlx_atomic_fetch_min_explicit(out, val, offset);
+  }
+
+  // Operator
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T a, T b) {
+    return a < b ? a : b;
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T a, T b) {
+    if (metal::isnan(a) || metal::isnan(b)) {
+      return static_cast<T>(NAN);
+    } else {
+      return a < b ? a : b;
+    }
+  }
+
+  template <>
+  complex64_t operator()(complex64_t a, complex64_t b) {
+    bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real);
+    bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag);
+
+    if (!real_is_nan && !imag_is_nan) {
+      return a < b ? a : b;
+    } else if (real_is_nan && !imag_is_nan) {
+      return complex64_t(
+          static_cast<float>(NAN), a.imag < b.imag ? a.imag : b.imag);
+    } else if (!real_is_nan && imag_is_nan) {
+      return complex64_t(
+          a.real < b.real ? a.real : b.real, static_cast<float>(NAN));
+    } else {
+      return complex64_t(static_cast<float>(NAN), static_cast<float>(NAN));
+    }
+  };
+};
+template <typename U>
+struct Max {
+  DEFINE_SIMD_REDUCE()
+
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
+    return simd_max(val);
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
+    if (simd_any(val != val)) {
+      return static_cast<T>(NAN);
+    }
+    return simd_max(val);
+  }
+
+  static constexpr constant U init = Limits<U>::min;
+
+  template <typename T>
+  void atomic_update(device mlx_atomic<T>* out, T val, size_t offset = 0) {
+    mlx_atomic_fetch_max_explicit(out, val, offset);
+  }
+
+  // Operator
+  template <typename T>
+  metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T a, T b) {
+    return a > b ? a : b;
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T a, T b) {
+    if (metal::isnan(a) || metal::isnan(b)) {
+      return static_cast<T>(NAN);
+    } else {
+      return a > b ? a : b;
+    }
+  }
+
+  template <>
+  complex64_t operator()(complex64_t a, complex64_t b) {
+    bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real);
+    bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag);
+
+    if (!real_is_nan && !imag_is_nan) {
+      return a > b ? a : b;
+    } else if (real_is_nan && !imag_is_nan) {
+      return complex64_t(
+          static_cast<float>(NAN), a.imag > b.imag ? a.imag : b.imag);
+    } else if (!real_is_nan && imag_is_nan) {
+      return complex64_t(
+          a.real > b.real ? a.real : b.real, static_cast<float>(NAN));
+    } else {
+      return complex64_t(static_cast<float>(NAN), static_cast<float>(NAN));
+    }
+  }
+};
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_all.h b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_all.h
new file mode 100644
index 0000000..e0d0839
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_all.h
@@ -0,0 +1,66 @@
+// Copyright © 2023-2024 Apple Inc.
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    typename IdxT = int64_t,
+    int N_READS = REDUCE_N_READS>
+[[kernel]] void all_reduce(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& in_size [[buffer(2)]],
+    const constant size_t& row_size [[buffer(3)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  Op op;
+  threadgroup U shared_vals[simd_size];
+
+  U total = Op::init;
+  IdxT start_idx = gid.y * IdxT(row_size);
+  IdxT actual_row =
+      (start_idx + row_size <= in_size) ? row_size : in_size - start_idx;
+  IdxT blocks = actual_row / (lsize.x * N_READS);
+  int extra = actual_row - blocks * (lsize.x * N_READS);
+  extra -= lid.x * N_READS;
+  start_idx += lid.x * N_READS;
+  in += start_idx;
+
+  if (extra >= N_READS) {
+    blocks++;
+    extra = 0;
+  }
+
+  for (IdxT b = 0; b < blocks; b++) {
+    for (int i = 0; i < N_READS; i++) {
+      total = op(static_cast<U>(in[i]), total);
+    }
+    in += lsize.x * N_READS;
+  }
+  if (extra > 0) {
+    for (int i = 0; i < extra; i++) {
+      total = op(static_cast<U>(in[i]), total);
+    }
+  }
+
+  // Reduction within simd group
+  total = op.simd_reduce(total);
+  if (simd_per_group > 1) {
+    if (simd_lane_id == 0) {
+      shared_vals[simd_group_id] = total;
+    }
+
+    // Reduction within thread group
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    total = lid.x < simd_per_group ? shared_vals[lid.x] : op.init;
+    total = op.simd_reduce(total);
+  }
+
+  if (lid.x == 0) {
+    out[gid.y] = total;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h
new file mode 100644
index 0000000..c109faf
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h
@@ -0,0 +1,398 @@
+// Copyright © 2023-2024 Apple Inc.
+
+template <typename T, typename U, typename Op, typename IdxT, int NDIMS>
+[[kernel]] void col_reduce_small(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant int64_t& reduction_stride [[buffer(3)]],
+    const constant int* shape [[buffer(4)]],
+    const constant int64_t* strides [[buffer(5)]],
+    const constant int& ndim [[buffer(6)]],
+    const constant int* reduce_shape [[buffer(7)]],
+    const constant int64_t* reduce_strides [[buffer(8)]],
+    const constant int& reduce_ndim [[buffer(9)]],
+    const constant size_t& non_col_reductions [[buffer(10)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]]) {
+  constexpr int n_reads = 4;
+  Op op;
+  LoopedElemToLoc<NDIMS, IdxT, (NDIMS > 2)> loop(reduce_ndim);
+  const device T* row;
+
+  U totals[n_reads];
+  for (int i = 0; i < n_reads; i++) {
+    totals[i] = Op::init;
+  }
+
+  IdxT column = IdxT(gid.x) * lsize.x * n_reads + lid.x * n_reads;
+  if (column >= reduction_stride) {
+    return;
+  }
+  bool safe = column + n_reads <= reduction_stride;
+
+  IdxT out_idx = gid.y + gsize.y * IdxT(gid.z);
+  IdxT in_idx = elem_to_loc<IdxT>(out_idx, shape, strides, ndim);
+  in += in_idx + column;
+
+  IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size);
+  loop.next(lid.y, reduce_shape, reduce_strides);
+  for (IdxT r = lid.y; r < total_rows; r += lsize.y) {
+    row = in + loop.location();
+    if (safe) {
+      for (int i = 0; i < n_reads; i++) {
+        totals[i] = op(static_cast<U>(row[i]), totals[i]);
+      }
+    } else {
+      U vals[n_reads];
+      for (int i = 0; i < n_reads; i++) {
+        vals[i] =
+            (column + i < reduction_stride) ? static_cast<U>(row[i]) : op.init;
+      }
+      for (int i = 0; i < n_reads; i++) {
+        totals[i] = op(vals[i], totals[i]);
+      }
+    }
+    loop.next(lsize.y, reduce_shape, reduce_strides);
+  }
+
+  if (lsize.y > 1) {
+    // lsize.y should be <= 8
+    threadgroup U shared_vals[32 * 8 * n_reads];
+    for (int i = 0; i < n_reads; i++) {
+      shared_vals[lid.y * lsize.x * n_reads + lid.x * n_reads + i] = totals[i];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (lid.y == 0) {
+      for (int i = 0; i < n_reads; i++) {
+        totals[i] = shared_vals[lid.x * n_reads + i];
+      }
+      for (uint j = 1; j < lsize.y; j++) {
+        for (int i = 0; i < n_reads; i++) {
+          totals[i] =
+              op(shared_vals[j * lsize.x * n_reads + lid.x * n_reads + i],
+                 totals[i]);
+        }
+      }
+    }
+  }
+
+  if (lid.y == 0) {
+    out += out_idx * IdxT(reduction_stride) + column;
+    if (safe) {
+      for (int i = 0; i < n_reads; i++) {
+        out[i] = totals[i];
+      }
+    } else {
+      for (int i = 0; column + i < reduction_stride; i++) {
+        out[i] = totals[i];
+      }
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, typename IdxT, int NDIMS>
+[[kernel]] void col_reduce_longcolumn(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& reduction_stride [[buffer(3)]],
+    const constant int* shape [[buffer(4)]],
+    const constant int64_t* strides [[buffer(5)]],
+    const constant int& ndim [[buffer(6)]],
+    const constant int* reduce_shape [[buffer(7)]],
+    const constant int64_t* reduce_strides [[buffer(8)]],
+    const constant int& reduce_ndim [[buffer(9)]],
+    const constant size_t& non_col_reductions [[buffer(10)]],
+    const constant size_t& out_size [[buffer(11)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]]) {
+  Op op;
+  LoopedElemToLoc<NDIMS, IdxT, (NDIMS > 2)> loop(reduce_ndim);
+  const device T* row;
+
+  IdxT out_idx = gid.x + gsize.x * IdxT(gid.y);
+  IdxT in_idx = elem_to_loc<IdxT>(out_idx, shape, strides, ndim);
+  in += in_idx + lid.x;
+
+  U total = Op::init;
+  IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size);
+  loop.next(gid.z * lsize.y + lid.y, reduce_shape, reduce_strides);
+  for (IdxT r = gid.z * lsize.y + lid.y; r < total_rows;
+       r += lsize.y * gsize.z) {
+    row = in + loop.location();
+    total = op(static_cast<U>(*row), total);
+    loop.next(lsize.y * gsize.z, reduce_shape, reduce_strides);
+  }
+
+  threadgroup U shared_vals[32 * 32];
+  shared_vals[lid.y * lsize.x + lid.x] = total;
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (lid.y == 0) {
+    for (uint i = 1; i < lsize.y; i++) {
+      total = op(total, shared_vals[i * lsize.x + lid.x]);
+    }
+    out[gid.z * IdxT(out_size) + out_idx * IdxT(reduction_stride) + lid.x] =
+        total;
+  }
+}
+
+/**
+ * Our approach is the following simple looped approach:
+ *  1. Each thread keeps running totals for BN / n_simdgroups outputs.
+ *  2. Load a tile BM, BN in registers and accumulate in the running totals
+ *  3. Move ahead by BM steps until the column axis and the non column
+ *     reductions are exhausted.
+ *  6. If BM == 32 then transpose in SM and simd reduce the running totals.
+ *     Otherwise write in shared memory and BN threads accumulate the running
+ *     totals with a loop.
+ *  7. Write them to the output
+ */
+template <
+    typename T,
+    typename U,
+    typename Op,
+    typename IdxT,
+    int NDIMS,
+    int BM,
+    int BN>
+[[kernel]] void col_reduce_looped(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant int64_t& reduction_stride [[buffer(3)]],
+    const constant int* shape [[buffer(4)]],
+    const constant int64_t* strides [[buffer(5)]],
+    const constant int& ndim [[buffer(6)]],
+    const constant int* reduce_shape [[buffer(7)]],
+    const constant int64_t* reduce_strides [[buffer(8)]],
+    const constant int& reduce_ndim [[buffer(9)]],
+    const constant size_t& non_col_reductions [[buffer(10)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  Op op;
+  constexpr int n_simdgroups = 8;
+  constexpr short tgp_size = n_simdgroups * simd_size;
+  constexpr short n_reads = (BM * BN) / tgp_size;
+  constexpr short n_read_blocks = BN / n_reads;
+
+  threadgroup U shared_vals[BN * BM];
+  U totals[n_reads];
+  LoopedElemToLoc<NDIMS, IdxT, (NDIMS > 2)> loop(reduce_ndim);
+  const device T* row;
+
+  for (int i = 0; i < n_reads; i++) {
+    totals[i] = Op::init;
+  }
+
+  short lid = simd_group_id * simd_size + simd_lane_id;
+  short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks);
+  IdxT column = BN * gid.x + offset.x;
+  bool safe = column + n_reads <= reduction_stride;
+
+  IdxT out_idx = gid.y + gsize.y * IdxT(gid.z);
+  IdxT in_idx = elem_to_loc<IdxT>(out_idx, shape, strides, ndim);
+  in += in_idx + column;
+
+  IdxT total = IdxT(non_col_reductions) * IdxT(reduction_size);
+  loop.next(offset.y, reduce_shape, reduce_strides);
+  for (IdxT r = offset.y; r < total; r += BM) {
+    row = in + loop.location();
+
+    if (safe) {
+      for (int i = 0; i < n_reads; i++) {
+        totals[i] = op(static_cast<U>(row[i]), totals[i]);
+      }
+    } else {
+      U vals[n_reads];
+      for (int i = 0; i < n_reads; i++) {
+        vals[i] =
+            (column + i < reduction_stride) ? static_cast<U>(row[i]) : op.init;
+      }
+      for (int i = 0; i < n_reads; i++) {
+        totals[i] = op(vals[i], totals[i]);
+      }
+    }
+
+    loop.next(BM, reduce_shape, reduce_strides);
+  }
+
+  // We can use a simd reduction to accumulate across BM so each thread writes
+  // the partial output to SM and then each simdgroup does BN / n_simdgroups
+  // accumulations.
+  if (BM == 32) {
+    constexpr int n_outputs = BN / n_simdgroups;
+    static_assert(
+        BM != 32 || n_outputs == n_reads,
+        "The tile should be selected such that n_outputs == n_reads");
+    for (int i = 0; i < n_reads; i++) {
+      shared_vals[offset.y * BN + offset.x + i] = totals[i];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    short2 out_offset(simd_group_id * n_outputs, simd_lane_id);
+    for (int i = 0; i < n_outputs; i++) {
+      totals[i] =
+          op.simd_reduce(shared_vals[out_offset.y * BN + out_offset.x + i]);
+    }
+
+    // Write the output.
+    if (simd_lane_id == 0) {
+      IdxT out_column = BN * gid.x + out_offset.x;
+      out += out_idx * IdxT(reduction_stride) + out_column;
+      if (out_column + n_outputs <= reduction_stride) {
+        for (int i = 0; i < n_outputs; i++) {
+          out[i] = totals[i];
+        }
+      } else {
+        for (int i = 0; out_column + i < reduction_stride; i++) {
+          out[i] = totals[i];
+        }
+      }
+    }
+  }
+
+  // Each thread holds n_reads partial results. We write them all out to shared
+  // memory and threads with offset.y == 0 aggregate the columns and write the
+  // outputs.
+  else {
+    short x_block = offset.x / n_reads;
+    for (int i = 0; i < n_reads; i++) {
+      shared_vals[x_block * BM * n_reads + i * BM + offset.y] = totals[i];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (offset.y == 0) {
+      for (int i = 0; i < n_reads; i++) {
+        for (int j = 1; j < BM; j++) {
+          totals[i] =
+              op(shared_vals[x_block * BM * n_reads + i * BM + j], totals[i]);
+        }
+      }
+    }
+
+    // Write the output.
+    if (offset.y == 0) {
+      out += out_idx * IdxT(reduction_stride) + column;
+      if (safe) {
+        for (int i = 0; i < n_reads; i++) {
+          out[i] = totals[i];
+        }
+      } else {
+        for (int i = 0; column + i < reduction_stride; i++) {
+          out[i] = totals[i];
+        }
+      }
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    typename IdxT,
+    int NDIMS,
+    int BM,
+    int BN>
+[[kernel]] void col_reduce_2pass(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant int64_t& reduction_stride [[buffer(3)]],
+    const constant int* shape [[buffer(4)]],
+    const constant int64_t* strides [[buffer(5)]],
+    const constant int& ndim [[buffer(6)]],
+    const constant int* reduce_shape [[buffer(7)]],
+    const constant int64_t* reduce_strides [[buffer(8)]],
+    const constant int& reduce_ndim [[buffer(9)]],
+    const constant size_t& non_col_reductions [[buffer(10)]],
+    const constant size_t& out_size [[buffer(11)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  Op op;
+  constexpr int n_simdgroups = 8;
+  constexpr short tgp_size = n_simdgroups * simd_size;
+  constexpr short n_reads = (BM * BN) / tgp_size;
+  constexpr short n_read_blocks = BN / n_reads;
+  constexpr int n_outputs = BN / n_simdgroups;
+  constexpr short outer_blocks = 32;
+  static_assert(BM == 32, "BM should be equal to 32");
+
+  threadgroup U shared_vals[BN * BM];
+  U totals[n_reads];
+  LoopedElemToLoc<NDIMS, IdxT, (NDIMS > 2)> loop(reduce_ndim);
+  const device T* row;
+
+  for (int i = 0; i < n_reads; i++) {
+    totals[i] = Op::init;
+  }
+
+  short lid = simd_group_id * simd_size + simd_lane_id;
+  short2 offset((lid % n_read_blocks) * n_reads, lid / n_read_blocks);
+  IdxT column = BN * gid.x + offset.x;
+  bool safe = column + n_reads <= reduction_stride;
+
+  IdxT full_idx = gid.y + gsize.y * IdxT(gid.z);
+  IdxT block_idx = full_idx / IdxT(out_size);
+  IdxT out_idx = full_idx % IdxT(out_size);
+  IdxT in_idx = elem_to_loc<IdxT>(out_idx, shape, strides, ndim);
+  in += in_idx + column;
+
+  IdxT total = IdxT(non_col_reductions) * IdxT(reduction_size);
+  loop.next(offset.y + block_idx * BM, reduce_shape, reduce_strides);
+  for (IdxT r = offset.y + block_idx * BM; r < total; r += outer_blocks * BM) {
+    row = in + loop.location();
+
+    if (safe) {
+      for (int i = 0; i < n_reads; i++) {
+        totals[i] = op(static_cast<U>(row[i]), totals[i]);
+      }
+    } else {
+      U vals[n_reads];
+      for (int i = 0; i < n_reads; i++) {
+        vals[i] =
+            (column + i < reduction_stride) ? static_cast<U>(row[i]) : op.init;
+      }
+      for (int i = 0; i < n_reads; i++) {
+        totals[i] = op(vals[i], totals[i]);
+      }
+    }
+
+    loop.next(outer_blocks * BM, reduce_shape, reduce_strides);
+  }
+
+  // We can use a simd reduction to accumulate across BM so each thread writes
+  // the partial output to SM and then each simdgroup does BN / n_simdgroups
+  // accumulations.
+  for (int i = 0; i < n_reads; i++) {
+    shared_vals[offset.y * BN + offset.x + i] = totals[i];
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  short2 out_offset(simd_group_id * n_outputs, simd_lane_id);
+  for (int i = 0; i < n_outputs; i++) {
+    totals[i] =
+        op.simd_reduce(shared_vals[out_offset.y * BN + out_offset.x + i]);
+  }
+
+  // Write the output.
+  if (simd_lane_id == 0) {
+    IdxT out_column = BN * gid.x + out_offset.x;
+    out += full_idx * IdxT(reduction_stride) + out_column;
+    if (out_column + n_outputs <= reduction_stride) {
+      for (int i = 0; i < n_outputs; i++) {
+        out[i] = totals[i];
+      }
+    } else {
+      for (int i = 0; out_column + i < reduction_stride; i++) {
+        out[i] = totals[i];
+      }
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_init.h b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_init.h
new file mode 100644
index 0000000..604efa7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_init.h
@@ -0,0 +1,8 @@
+// Copyright © 2023-2024 Apple Inc.
+
+template <typename T, typename Op>
+[[kernel]] void init_reduce(
+    device T* out [[buffer(0)]],
+    uint tid [[thread_position_in_grid]]) {
+  out[tid] = Op::init;
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_row.h b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_row.h
new file mode 100644
index 0000000..936d75b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/reduction/reduce_row.h
@@ -0,0 +1,369 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// Row reduction utilities
+// - `per_thread_row_reduce` collaborative partial reduction in the threadgroup
+// - `threadgroup_reduce` collaborative reduction in the threadgroup such that
+//   lid.x == 0 holds the reduced value
+// - `thread_reduce` simple loop and reduce the row
+
+/**
+ * The thread group collaboratively reduces across the rows with bounds
+ * checking. In the end each thread holds a part of the reduction.
+ */
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS = REDUCE_N_READS,
+    int N_WRITES = REDUCE_N_WRITES>
+METAL_FUNC void per_thread_row_reduce(
+    thread U totals[N_WRITES],
+    const device T* inputs[N_WRITES],
+    int blocks,
+    int extra,
+    uint lsize_x,
+    uint lid_x) {
+  Op op;
+
+  // Set up the accumulator registers
+  for (int i = 0; i < N_WRITES; i++) {
+    totals[i] = Op::init;
+  }
+
+  // Loop over the reduction size within thread group
+  for (int i = 0; i < blocks; i++) {
+    for (int j = 0; j < N_WRITES; j++) {
+      for (int i = 0; i < N_READS; i++) {
+        totals[j] = op(static_cast<U>(inputs[j][i]), totals[j]);
+      }
+
+      inputs[j] += lsize_x * N_READS;
+    }
+  }
+
+  // Separate case for the last set as we close the reduction size
+  int index = lid_x * N_READS;
+  if (index + N_READS <= extra) {
+    for (int j = 0; j < N_WRITES; j++) {
+      for (int i = 0; i < N_READS; i++) {
+        totals[j] = op(static_cast<U>(inputs[j][i]), totals[j]);
+      }
+    }
+  } else {
+    for (int j = 0; j < N_WRITES; j++) {
+      for (int i = 0; index + i < extra; i++) {
+        totals[j] = op(static_cast<U>(inputs[j][i]), totals[j]);
+      }
+    }
+  }
+}
+
+/**
+ * Consecutive rows in a contiguous array.
+ */
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS = REDUCE_N_READS,
+    int N_WRITES = REDUCE_N_WRITES>
+METAL_FUNC void per_thread_row_reduce(
+    thread U totals[N_WRITES],
+    const device T* in,
+    const constant size_t& reduction_size,
+    int blocks,
+    int extra,
+    uint lsize_x,
+    uint lid_x) {
+  // Set up the input pointers
+  const device T* inputs[N_WRITES];
+  inputs[0] = in + lid_x * N_READS;
+  for (int i = 1; i < N_READS; i++) {
+    inputs[i] = inputs[i - 1] + reduction_size;
+  }
+
+  per_thread_row_reduce<T, U, Op, N_READS, N_WRITES>(
+      totals, inputs, blocks, extra, lsize_x, lid_x);
+}
+
+/**
+ * Consecutive rows in an arbitrarily ordered array.
+ */
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS = REDUCE_N_READS,
+    int N_WRITES = REDUCE_N_WRITES>
+METAL_FUNC void per_thread_row_reduce(
+    thread U totals[N_WRITES],
+    const device T* in,
+    const int64_t row_idx,
+    int blocks,
+    int extra,
+    const constant int* shape,
+    const constant int64_t* strides,
+    const constant int& ndim,
+    uint lsize_x,
+    uint lid_x) {
+  // Set up the input pointers
+  const device T* inputs[N_WRITES];
+  in += lid_x * N_READS;
+  for (int i = 0; i < N_READS; i++) {
+    inputs[i] = in + elem_to_loc(row_idx + i, shape, strides, ndim);
+  }
+
+  per_thread_row_reduce<T, U, Op, N_READS, N_WRITES>(
+      totals, inputs, blocks, extra, lsize_x, lid_x);
+}
+
+/**
+ * Reduce within the threadgroup.
+ */
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS = REDUCE_N_READS,
+    int N_WRITES = REDUCE_N_WRITES>
+METAL_FUNC void threadgroup_reduce(
+    thread U totals[N_WRITES],
+    threadgroup U* shared_vals,
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  Op op;
+
+  // Simdgroup first
+  for (int i = 0; i < N_WRITES; i++) {
+    totals[i] = op.simd_reduce(totals[i]);
+  }
+
+  // Across simdgroups
+  if (simd_per_group > 1) {
+    if (simd_lane_id == 0) {
+      for (int i = 0; i < N_WRITES; i++) {
+        shared_vals[simd_group_id * N_WRITES + i] = totals[i];
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    U values[N_WRITES];
+    for (int i = 0; i < N_WRITES; i++) {
+      values[i] = (lid.x < simd_per_group) ? shared_vals[lid.x * N_WRITES + i]
+                                           : op.init;
+    }
+
+    for (int i = 0; i < N_WRITES; i++) {
+      totals[i] = op.simd_reduce(values[i]);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
+METAL_FUNC void
+thread_reduce(thread U& total, const device T* row, int blocks, int extra) {
+  Op op;
+  for (int i = 0; i < blocks; i++) {
+    U vals[N_READS];
+    for (int j = 0; j < N_READS; j++) {
+      vals[j] = row[j];
+    }
+    for (int j = 0; j < N_READS; j++) {
+      total = op(vals[j], total);
+    }
+    row += N_READS;
+  }
+  for (int i = 0; i < extra; i++) {
+    total = op(*row++, total);
+  }
+}
+
+// Reduction kernels
+// - `row_reduce_small` depending on the non-row reductions and row size it
+//   either just loops over everything or a simd collaboratively reduces the
+//   non_row reductions. In the first case one thread is responsible for one
+//   output on the 2nd one simd is responsible for one output.
+// - `row_reduce_simple` simple contiguous row reduction
+// - `row_reduce_looped` simply loop and reduce each row for each non-row
+//   reduction. One threadgroup is responsible for one output.
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    typename IdxT,
+    int NDIMS,
+    int N_READS = REDUCE_N_READS>
+[[kernel]] void row_reduce_small(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant int64_t& row_size [[buffer(2)]],
+    const constant int64_t& non_row_reductions [[buffer(3)]],
+    const constant int* shape [[buffer(4)]],
+    const constant int64_t* strides [[buffer(5)]],
+    const constant int& ndim [[buffer(6)]],
+    const constant int* reduce_shape [[buffer(7)]],
+    const constant int64_t* reduce_strides [[buffer(8)]],
+    const constant int& reduce_ndim [[buffer(9)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint3 tid [[thread_position_in_grid]],
+    uint3 tsize [[threads_per_grid]]) {
+  Op op;
+
+  U total_val = Op::init;
+  LoopedElemToLoc<NDIMS, IdxT, (NDIMS > 2)> loop(reduce_ndim);
+
+  // Precompute some row reduction numbers
+  const device T* row;
+  int blocks = IdxT(row_size) / N_READS;
+  int extra = IdxT(row_size) % N_READS;
+
+  if ((non_row_reductions < 32 && row_size <= 8) || non_row_reductions <= 8) {
+    // Simple loop over non_row_reductions and reduce the row in the thread.
+    IdxT out_idx = tid.x + tsize.x * IdxT(tid.y);
+    in += elem_to_loc<IdxT>(out_idx, shape, strides, ndim);
+
+    for (uint r = 0; r < non_row_reductions; r++) {
+      row = in + loop.location();
+      thread_reduce<T, U, Op, N_READS>(total_val, row, blocks, extra);
+      loop.next(reduce_shape, reduce_strides);
+    }
+
+    out[out_idx] = total_val;
+  } else {
+    // Collaboratively reduce over non_row_reductions in the simdgroup. Each
+    // thread reduces every 32nd row and then a simple simd reduce.
+    IdxT out_idx = gid.y + gsize.y * IdxT(gid.z);
+    in += elem_to_loc<IdxT>(out_idx, shape, strides, ndim);
+
+    loop.next(simd_lane_id, reduce_shape, reduce_strides);
+
+    for (uint r = simd_lane_id; r < non_row_reductions; r += simd_size) {
+      row = in + loop.location();
+      thread_reduce<T, U, Op, N_READS>(total_val, row, blocks, extra);
+      loop.next(simd_size, reduce_shape, reduce_strides);
+    }
+
+    total_val = op.simd_reduce(total_val);
+
+    if (simd_lane_id == 0) {
+      out[out_idx] = total_val;
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    typename IdxT = int64_t,
+    int N_READS = REDUCE_N_READS,
+    int N_WRITES = REDUCE_N_WRITES>
+[[kernel]] void row_reduce_simple(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant int64_t& out_size [[buffer(3)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  threadgroup U shared_vals[simd_size * N_WRITES];
+  U totals[N_WRITES];
+
+  // Move to the row
+  IdxT out_idx = N_WRITES * (gid.y + gsize.y * IdxT(gid.z));
+  if (out_idx + N_WRITES > out_size) {
+    out_idx = out_size - N_WRITES;
+  }
+  in += out_idx * IdxT(reduction_size);
+  out += out_idx;
+
+  // Each thread reduces across the row
+  int blocks = IdxT(reduction_size) / (lsize.x * N_READS);
+  int extra = reduction_size - blocks * (lsize.x * N_READS);
+  per_thread_row_reduce<T, U, Op, N_READS, N_WRITES>(
+      totals, in, reduction_size, blocks, extra, lsize.x, lid.x);
+
+  // Reduce across the threadgroup
+  threadgroup_reduce<T, U, Op, N_READS, N_WRITES>(
+      totals, shared_vals, lid, simd_lane_id, simd_per_group, simd_group_id);
+
+  // Write the output
+  if (lid.x == 0) {
+    for (int i = 0; i < N_WRITES; i++) {
+      out[i] = totals[i];
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    typename IdxT,
+    int NDIMS,
+    int N_READS = REDUCE_N_READS>
+[[kernel]] void row_reduce_looped(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant int64_t& row_size [[buffer(2)]],
+    const constant int64_t& non_row_reductions [[buffer(3)]],
+    const constant int* shape [[buffer(4)]],
+    const constant int64_t* strides [[buffer(5)]],
+    const constant int& ndim [[buffer(6)]],
+    const constant int* reduce_shape [[buffer(7)]],
+    const constant int64_t* reduce_strides [[buffer(8)]],
+    const constant int& reduce_ndim [[buffer(9)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  Op op;
+  threadgroup U shared_vals[simd_size];
+  U total = Op::init;
+
+  IdxT out_idx = gid.y + gsize.y * IdxT(gid.z);
+
+  // lid.x * N_READS breaks the per_thread_row_reduce interface a bit. Maybe it
+  // needs a small refactor.
+  in += elem_to_loc<IdxT>(out_idx, shape, strides, ndim) + lid.x * N_READS;
+
+  LoopedElemToLoc<NDIMS, IdxT, (NDIMS > 2)> loop(reduce_ndim);
+  const device T* row;
+  int blocks = IdxT(row_size) / (lsize.x * N_READS);
+  int extra = row_size - blocks * (lsize.x * N_READS);
+
+  for (IdxT i = 0; i < non_row_reductions; i++) {
+    row = in + loop.location();
+
+    // Each thread reduces across the row
+    U row_total;
+    per_thread_row_reduce<T, U, Op, N_READS, 1>(
+        &row_total, &row, blocks, extra, lsize.x, lid.x);
+
+    // Aggregate across rows
+    total = op(total, row_total);
+
+    loop.next(reduce_shape, reduce_strides);
+  }
+
+  // Reduce across the threadgroup
+  threadgroup_reduce<T, U, Op, N_READS, 1>(
+      &total, shared_vals, lid, simd_lane_id, simd_per_group, simd_group_id);
+
+  // Write the output
+  if (lid.x == 0) {
+    out[out_idx] = total;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/rms_norm.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/rms_norm.metal
new file mode 100644
index 0000000..62f2457
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/rms_norm.metal
@@ -0,0 +1,391 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_common>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+using namespace metal;
+
+constant bool has_w [[function_constant(20)]];
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void rms_single_row(
+    const device T* x,
+    const device T* w,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup float local_inv_mean[1];
+  threadgroup float local_sums[SIMD_SIZE];
+
+  float acc = 0;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      float xi = x[i];
+      acc += xi * xi;
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        float xi = x[i];
+        acc += xi * xi;
+      }
+    }
+  }
+  acc = simd_sum(acc);
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sums[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sums[simd_group_id] = acc;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    acc = simd_sum(local_sums[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_inv_mean[0] = metal::precise::rsqrt(acc / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write the outputs
+  out += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      out[i] = w[w_stride * i] * static_cast<T>(x[i] * local_inv_mean[0]);
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        out[i] = w[w_stride * i] * static_cast<T>(x[i] * local_inv_mean[0]);
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void rms_looped(
+    const device T* x,
+    const device T* w,
+    device T* out,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int SIMD_SIZE = 32;
+  threadgroup float local_inv_mean[1];
+  threadgroup float local_sums[SIMD_SIZE];
+
+  float acc = 0;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        acc += xi * xi;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          acc += xi * xi;
+        }
+      }
+    }
+  }
+  acc = simd_sum(acc);
+  //  Initialize shared memory
+  if (simd_group_id == 0) {
+    local_sums[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write simd accumulations into shared memory
+  if (simd_lane_id == 0) {
+    local_sums[simd_group_id] = acc;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Accumulate over simd groups
+  if (simd_group_id == 0) {
+    acc = simd_sum(local_sums[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_inv_mean[0] = metal::precise::rsqrt(acc / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Write the outputs
+  out += gid * size_t(axis_size) + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        out[r + i] = w[w_stride * (i + r)] *
+            static_cast<T>(x[r + i] * local_inv_mean[0]);
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          out[r + i] = w[w_stride * (i + r)] *
+              static_cast<T>(x[r + i] * local_inv_mean[0]);
+        }
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void vjp_rms_single_row(
+    const device T* x,
+    const device T* w,
+    const device T* g,
+    device T* gx,
+    device T* gw,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  // Advance the input pointers
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+
+  // Allocate registers for the computation and accumulators
+  float thread_x[N_READS];
+  float thread_w[N_READS];
+  float thread_g[N_READS];
+  float sumx2 = 0;
+  float sumgwx = 0;
+
+  // Allocate shared memory to implement the reduction
+  constexpr int SIMD_SIZE = 32;
+  threadgroup float local_sumx2[SIMD_SIZE];
+  threadgroup float local_sumgwx[SIMD_SIZE];
+  threadgroup float local_normalizer[1];
+  threadgroup float local_meangwx[1];
+
+  // Read and accumulate locally
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      thread_x[i] = x[i];
+      thread_w[i] = w[w_stride * i];
+      thread_g[i] = g[i];
+
+      sumx2 += thread_x[i] * thread_x[i];
+      sumgwx += thread_x[i] * thread_w[i] * thread_g[i];
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        thread_x[i] = x[i];
+        thread_w[i] = w[w_stride * i];
+        thread_g[i] = g[i];
+
+        sumx2 += thread_x[i] * thread_x[i];
+        sumgwx += thread_x[i] * thread_w[i] * thread_g[i];
+      }
+    }
+  }
+
+  // Accumulate across threads
+  sumx2 = simd_sum(sumx2);
+  sumgwx = simd_sum(sumgwx);
+  if (simd_group_id == 0) {
+    local_sumx2[simd_lane_id] = 0;
+    local_sumgwx[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_lane_id == 0) {
+    local_sumx2[simd_group_id] = sumx2;
+    local_sumgwx[simd_group_id] = sumgwx;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    sumx2 = simd_sum(local_sumx2[simd_lane_id]);
+    sumgwx = simd_sum(local_sumgwx[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_meangwx[0] = sumgwx / axis_size;
+      local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  float meangwx = local_meangwx[0];
+  float normalizer = local_normalizer[0];
+  float normalizer3 = normalizer * normalizer * normalizer;
+
+  // Write the outputs
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      gx[i] = static_cast<T>(
+          thread_g[i] * thread_w[i] * normalizer -
+          thread_x[i] * meangwx * normalizer3);
+      if (has_w) {
+        gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+      }
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        gx[i] = static_cast<T>(
+            thread_g[i] * thread_w[i] * normalizer -
+            thread_x[i] * meangwx * normalizer3);
+        if (has_w) {
+          gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
+        }
+      }
+    }
+  }
+}
+
+template <typename T, int N_READS = RMS_N_READS>
+[[kernel]] void vjp_rms_looped(
+    const device T* x,
+    const device T* w,
+    const device T* g,
+    device T* gx,
+    device T* gw,
+    constant float& eps,
+    constant uint& axis_size,
+    constant uint& w_stride,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  // Advance the input pointers
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
+  w += w_stride * lid * N_READS;
+
+  // Allocate registers for the accumulators
+  float sumx2 = 0;
+  float sumgwx = 0;
+
+  // Allocate shared memory to implement the reduction
+  constexpr int SIMD_SIZE = 32;
+  threadgroup float local_sumx2[SIMD_SIZE];
+  threadgroup float local_sumgwx[SIMD_SIZE];
+  threadgroup float local_normalizer[1];
+  threadgroup float local_meangwx[1];
+
+  // Read and accumulate locally
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        float wi = w[w_stride * (i + r)];
+        float gi = g[i + r];
+
+        sumx2 += xi * xi;
+        sumgwx += xi * wi * gi;
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          float wi = w[w_stride * (i + r)];
+          float gi = g[i + r];
+
+          sumx2 += xi * xi;
+          sumgwx += xi * wi * gi;
+        }
+      }
+    }
+  }
+
+  // Accumulate across threads
+  sumx2 = simd_sum(sumx2);
+  sumgwx = simd_sum(sumgwx);
+  if (simd_group_id == 0) {
+    local_sumx2[simd_lane_id] = 0;
+    local_sumgwx[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_lane_id == 0) {
+    local_sumx2[simd_group_id] = sumx2;
+    local_sumgwx[simd_group_id] = sumgwx;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    sumx2 = simd_sum(local_sumx2[simd_lane_id]);
+    sumgwx = simd_sum(local_sumgwx[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_meangwx[0] = sumgwx / axis_size;
+      local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  float meangwx = local_meangwx[0];
+  float normalizer = local_normalizer[0];
+  float normalizer3 = normalizer * normalizer * normalizer;
+
+  // Write the outputs
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
+  for (uint r = 0; r < axis_size; r += lsize * N_READS) {
+    if (r + lid * N_READS + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        float xi = x[i + r];
+        float wi = w[w_stride * (i + r)];
+        float gi = g[i + r];
+
+        gx[i + r] =
+            static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
+        if (has_w) {
+          gw[i + r] = static_cast<T>(gi * xi * normalizer);
+        }
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if ((r + lid * N_READS + i) < axis_size) {
+          float xi = x[i + r];
+          float wi = w[w_stride * (i + r)];
+          float gi = g[i + r];
+
+          gx[i + r] =
+              static_cast<T>(gi * wi * normalizer - xi * meangwx * normalizer3);
+          if (has_w) {
+            gw[i + r] = static_cast<T>(gi * xi * normalizer);
+          }
+        }
+      }
+    }
+  }
+}
+
+// clang-format off
+#define instantiate_rms(name, itype)                                \
+  instantiate_kernel("rms" #name, rms_single_row, itype)            \
+  instantiate_kernel("vjp_rms" #name, vjp_rms_single_row, itype)    \
+  instantiate_kernel("rms_looped" #name, rms_looped, itype)         \
+  instantiate_kernel("vjp_rms_looped" #name, vjp_rms_looped, itype)
+
+instantiate_rms(float32, float)
+instantiate_rms(float16, half)
+instantiate_rms(bfloat16, bfloat16_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/rope.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/rope.metal
new file mode 100644
index 0000000..fce444a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/rope.metal
@@ -0,0 +1,229 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_math>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+constant bool forward [[function_constant(1)]];
+constant bool traditional [[function_constant(2)]];
+constant bool hs_transpose [[function_constant(3)]];
+
+template <typename T>
+void rope_single_impl(
+    const device T* in,
+    device T* out,
+    constant const int& offset,
+    const float inv_freq,
+    constant const float& scale,
+    constant const int64_t& stride,
+    uint2 pos,
+    uint2 grid) {
+  float L = scale * static_cast<float>(offset);
+
+  // Compute costheta, sintheta
+  float theta = L * inv_freq;
+  float costheta = metal::fast::cos(theta);
+  float sintheta = metal::fast::sin(theta);
+
+  // Compute the input and output indices
+  uint index_1, index_2;
+  if (traditional) {
+    index_1 = 2 * pos.x + pos.y * stride;
+    index_2 = index_1 + 1;
+  } else {
+    index_1 = pos.x + pos.y * stride;
+    index_2 = index_1 + grid.x;
+  }
+
+  // Read and write the output
+  float x1 = static_cast<float>(in[index_1]);
+  float x2 = static_cast<float>(in[index_2]);
+  float rx1;
+  float rx2;
+  if (forward) {
+    rx1 = x1 * costheta - x2 * sintheta;
+    rx2 = x1 * sintheta + x2 * costheta;
+  } else {
+    rx1 = x2 * sintheta + x1 * costheta;
+    rx2 = x2 * costheta - x1 * sintheta;
+  }
+  out[index_1] = static_cast<T>(rx1);
+  out[index_2] = static_cast<T>(rx2);
+}
+
+template <typename T>
+[[kernel]] void rope_single(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const int& offset,
+    constant const float& scale,
+    constant const int64_t& stride,
+    constant const float& base [[buffer(10)]],
+    uint2 pos [[thread_position_in_grid]],
+    uint2 grid [[threads_per_grid]]) {
+  float d = static_cast<float>(pos.x) / static_cast<float>(grid.x);
+  float inv_freq = metal::exp2(-d * base);
+  rope_single_impl<T>(in, out, offset, inv_freq, scale, stride, pos, grid);
+}
+
+template <typename T>
+[[kernel]] void rope_single_freqs(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const int& offset,
+    constant const float& scale,
+    constant const int64_t& stride,
+    const device float* freqs [[buffer(10)]],
+    constant const int64_t& freq_stride [[buffer(11)]],
+    uint2 pos [[thread_position_in_grid]],
+    uint2 grid [[threads_per_grid]]) {
+  float inv_freq = 1.0 / (freqs[freq_stride * pos.x]);
+  rope_single_impl<T>(in, out, offset, inv_freq, scale, stride, pos, grid);
+}
+
+template <typename T, typename IdxT, int N = 4>
+void rope_impl(
+    const device T* in,
+    device T* out,
+    const device int* offset,
+    const float inv_freq,
+    constant const float& scale,
+    constant const int64_t strides[3],
+    constant const int64_t out_strides[3],
+    constant const int64_t& offset_stride,
+    constant const int& n_head,
+    uint3 pos,
+    uint3 grid) {
+  auto n_head_up = N * ((n_head + N - 1) / N);
+  auto head_idx = static_cast<int>((pos.z * N) % n_head_up);
+  auto batch_idx = (pos.z * N) / n_head_up;
+  auto batch_offset = offset[batch_idx * offset_stride];
+  float L = scale * static_cast<float>(pos.y + batch_offset);
+  auto mat_idx = batch_idx * n_head + head_idx;
+
+  // Compute costheta, sintheta
+  float theta = L * inv_freq;
+  float costheta = metal::fast::cos(theta);
+  float sintheta = metal::fast::sin(theta);
+  // Compute the input and output indices
+  IdxT in_index_1;
+  if (hs_transpose) {
+    IdxT batch_stride = grid.y * IdxT(strides[1]);
+    in_index_1 =
+        batch_idx * batch_stride + pos.y * strides[1] + head_idx * strides[0];
+  } else {
+    in_index_1 = pos.y * IdxT(strides[1]) + mat_idx * IdxT(strides[0]);
+  }
+  IdxT in_index_2;
+  IdxT out_index_1 =
+      pos.y * IdxT(out_strides[1]) + mat_idx * IdxT(out_strides[0]);
+  IdxT out_index_2;
+  if (traditional) {
+    out_index_1 += 2 * pos.x * IdxT(out_strides[2]);
+    out_index_2 = out_index_1 + 1;
+    in_index_1 += 2 * pos.x * IdxT(strides[2]);
+    in_index_2 = in_index_1 + IdxT(strides[2]);
+  } else {
+    out_index_1 += pos.x * IdxT(out_strides[2]);
+    out_index_2 = out_index_1 + grid.x * IdxT(out_strides[2]);
+    in_index_1 += pos.x * IdxT(strides[2]);
+    in_index_2 = in_index_1 + grid.x * IdxT(strides[2]);
+  }
+  for (int i = 0; i < N && head_idx + i < n_head; ++i) {
+    // Read and write the output
+    float x1 = static_cast<float>(in[in_index_1]);
+    float x2 = static_cast<float>(in[in_index_2]);
+    float rx1;
+    float rx2;
+    if (forward) {
+      rx1 = x1 * costheta - x2 * sintheta;
+      rx2 = x1 * sintheta + x2 * costheta;
+    } else {
+      rx1 = x2 * sintheta + x1 * costheta;
+      rx2 = x2 * costheta - x1 * sintheta;
+    }
+    out[out_index_1] = static_cast<T>(rx1);
+    out[out_index_2] = static_cast<T>(rx2);
+    in_index_1 += IdxT(strides[0]);
+    in_index_2 += IdxT(strides[0]);
+    out_index_1 += IdxT(out_strides[0]);
+    out_index_2 += IdxT(out_strides[0]);
+  }
+}
+
+template <typename T, typename IdxT, int N = 4>
+[[kernel]] void rope(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const device int* offset,
+    constant const float& scale,
+    constant const int64_t strides[3],
+    constant const int64_t out_strides[3],
+    constant const int64_t& offset_stride,
+    constant const int& n_head,
+    constant const float& base [[buffer(10)]],
+    uint3 pos [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  float d = static_cast<float>(pos.x) / static_cast<float>(grid.x);
+  float inv_freq = metal::exp2(-d * base);
+  rope_impl<T, IdxT, N>(
+      in,
+      out,
+      offset,
+      inv_freq,
+      scale,
+      strides,
+      out_strides,
+      offset_stride,
+      n_head,
+      pos,
+      grid);
+}
+
+template <typename T, typename IdxT, int N = 4>
+[[kernel]] void rope_freqs(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const device int* offset,
+    constant const float& scale,
+    constant const int64_t strides[3],
+    constant const int64_t out_strides[3],
+    constant const int64_t& offset_stride,
+    constant const int& n_head,
+    const device float* freqs [[buffer(10)]],
+    constant const int64_t& freq_stride [[buffer(11)]],
+    uint3 pos [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  float inv_freq = 1.0 / (freqs[freq_stride * pos.x]);
+  rope_impl<T, IdxT, N>(
+      in,
+      out,
+      offset,
+      inv_freq,
+      scale,
+      strides,
+      out_strides,
+      offset_stride,
+      n_head,
+      pos,
+      grid);
+}
+
+// clang-format off
+#define instantiate_rope_g(name, type) \
+  instantiate_kernel("rope_" #name, rope, type, int32_t) \
+  instantiate_kernel("rope_freqs_" #name, rope_freqs, type, int32_t) \
+  instantiate_kernel("rope_large_" #name, rope, type, int64_t) \
+  instantiate_kernel("rope_freqs_large_" #name, rope_freqs, type, int64_t)
+
+#define instantiate_rope_s(name, type) \
+  instantiate_kernel("rope_single_" #name, rope_single, type) \
+  instantiate_kernel("rope_single_freqs_" #name, rope_single_freqs, type)
+
+#define instantiate_rope(name, type) \
+  instantiate_rope_s(name, type)     \
+  instantiate_rope_g(name, type)
+
+instantiate_rope(float16, half)
+instantiate_rope(bfloat16, bfloat16_t)
+instantiate_rope(float32, float) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal
new file mode 100644
index 0000000..c668d9d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal
@@ -0,0 +1,44 @@
+#include <metal_stdlib>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/sdpa_vector.h"
+
+using namespace metal;
+
+// SDPA vector instantiations
+#define instantiate_sdpa_vector_aggregation(type, value_dim) \
+  instantiate_kernel(                                        \
+      "sdpa_vector_2pass_2_" #type "_" #value_dim,           \
+      sdpa_vector_2pass_2,                                   \
+      type,                                                  \
+      value_dim)
+
+#define instantiate_sdpa_vector(type, qk_dim, value_dim)       \
+  instantiate_kernel(                                          \
+      "sdpa_vector_" #type "_" #qk_dim "_" #value_dim,         \
+      sdpa_vector,                                             \
+      type,                                                    \
+      qk_dim,                                                  \
+      value_dim)                                               \
+  instantiate_kernel(                                          \
+      "sdpa_vector_2pass_1_" #type "_" #qk_dim "_" #value_dim, \
+      sdpa_vector_2pass_1,                                     \
+      type,                                                    \
+      qk_dim,                                                  \
+      value_dim)
+
+#define instantiate_sdpa_vector_heads(type)      \
+  instantiate_sdpa_vector(type, 64, 64)          \
+  instantiate_sdpa_vector(type, 96, 96)          \
+  instantiate_sdpa_vector(type, 128, 128)        \
+  instantiate_sdpa_vector(type, 256, 256)        \
+  instantiate_sdpa_vector_aggregation(type, 64)  \
+  instantiate_sdpa_vector_aggregation(type, 96)  \
+  instantiate_sdpa_vector_aggregation(type, 128) \
+  instantiate_sdpa_vector_aggregation(type, 256)
+
+instantiate_sdpa_vector_heads(float)
+instantiate_sdpa_vector_heads(bfloat16_t)
+instantiate_sdpa_vector_heads(float16_t)
+    // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/scan.h b/macos/third_party/mlx/mlx/backend/metal/kernels/scan.h
new file mode 100644
index 0000000..1668261
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/scan.h
@@ -0,0 +1,514 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/binary_ops.h"
+
+#define DEFINE_SIMD_SCAN()                                               \
+  template <typename T, metal::enable_if_t<sizeof(T) < 8, bool> = true>  \
+  T simd_scan(T val) {                                                   \
+    return simd_scan_impl(val);                                          \
+  }                                                                      \
+                                                                         \
+  template <typename T, metal::enable_if_t<sizeof(T) == 8, bool> = true> \
+  T simd_scan(T val) {                                                   \
+    for (int i = 1; i <= 16; i *= 2) {                                   \
+      val = operator()(val, simd_shuffle_and_fill_up(val, init, i));     \
+    }                                                                    \
+    return val;                                                          \
+  }
+
+#define DEFINE_SIMD_EXCLUSIVE_SCAN()                                     \
+  template <typename T, metal::enable_if_t<sizeof(T) < 8, bool> = true>  \
+  T simd_exclusive_scan(T val) {                                         \
+    return simd_exclusive_scan_impl(val);                                \
+  }                                                                      \
+                                                                         \
+  template <typename T, metal::enable_if_t<sizeof(T) == 8, bool> = true> \
+  T simd_exclusive_scan(T val) {                                         \
+    val = simd_scan(val);                                                \
+    return simd_shuffle_and_fill_up(val, init, 1);                       \
+  }
+
+template <typename U>
+struct CumSum {
+  DEFINE_SIMD_SCAN()
+  DEFINE_SIMD_EXCLUSIVE_SCAN()
+
+  static constexpr constant U init = static_cast<U>(0);
+
+  template <typename T>
+  U operator()(U a, T b) {
+    return a + b;
+  }
+
+  U simd_scan_impl(U x) {
+    return simd_prefix_inclusive_sum(x);
+  }
+
+  U simd_exclusive_scan_impl(U x) {
+    return simd_prefix_exclusive_sum(x);
+  }
+};
+
+template <typename U>
+struct CumProd {
+  DEFINE_SIMD_SCAN()
+  DEFINE_SIMD_EXCLUSIVE_SCAN()
+
+  static constexpr constant U init = static_cast<U>(1.0f);
+
+  template <typename T>
+  U operator()(U a, T b) {
+    return a * b;
+  }
+
+  U simd_scan_impl(U x) {
+    return simd_prefix_inclusive_product(x);
+  }
+
+  U simd_exclusive_scan_impl(U x) {
+    return simd_prefix_exclusive_product(x);
+  }
+};
+
+template <>
+struct CumProd<bool> {
+  static constexpr constant bool init = true;
+
+  template <typename T>
+  bool operator()(bool a, T b) {
+    return a & static_cast<bool>(b);
+  }
+
+  bool simd_scan(bool x) {
+    for (int i = 1; i <= 16; i *= 2) {
+      bool other = simd_shuffle_and_fill_up(x, init, i);
+      x &= other;
+    }
+    return x;
+  }
+
+  bool simd_exclusive_scan(bool x) {
+    x = simd_scan(x);
+    return simd_shuffle_and_fill_up(x, init, 1);
+  }
+};
+
+template <typename U>
+struct CumMax {
+  static constexpr constant U init = Limits<U>::min;
+
+  template <typename T>
+  U operator()(U a, T b) {
+    return (a >= b) ? a : b;
+  }
+
+  U simd_scan(U x) {
+    for (int i = 1; i <= 16; i *= 2) {
+      U other = simd_shuffle_and_fill_up(x, init, i);
+      x = (x >= other) ? x : other;
+    }
+    return x;
+  }
+
+  U simd_exclusive_scan(U x) {
+    x = simd_scan(x);
+    return simd_shuffle_and_fill_up(x, init, 1);
+  }
+};
+
+template <typename U>
+struct CumMin {
+  static constexpr constant U init = Limits<U>::max;
+
+  template <typename T>
+  U operator()(U a, T b) {
+    return (a <= b) ? a : b;
+  }
+
+  U simd_scan(U x) {
+    for (int i = 1; i <= 16; i *= 2) {
+      U other = simd_shuffle_and_fill_up(x, init, i);
+      x = (x <= other) ? x : other;
+    }
+    return x;
+  }
+
+  U simd_exclusive_scan(U x) {
+    x = simd_scan(x);
+    return simd_shuffle_and_fill_up(x, init, 1);
+  }
+};
+
+template <typename U>
+struct CumLogaddexp {
+  static constexpr constant U init = Limits<U>::min;
+
+  template <typename T>
+  U operator()(U a, T b) {
+    return LogAddExp{}(a, static_cast<U>(b));
+  }
+
+  U simd_scan(U x) {
+    for (int i = 1; i <= 16; i *= 2) {
+      U other = simd_shuffle_and_fill_up(x, init, i);
+      x = LogAddExp{}(x, other);
+    }
+    return x;
+  }
+
+  U simd_exclusive_scan(U x) {
+    x = simd_scan(x);
+    return simd_shuffle_and_fill_up(x, init, 1);
+  }
+};
+
+template <typename T, typename U, int N_READS, bool reverse>
+inline void load_unsafe(U values[N_READS], const device T* input) {
+  if (reverse) {
+    for (int i = 0; i < N_READS; i++) {
+      values[N_READS - i - 1] = input[i];
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      values[i] = input[i];
+    }
+  }
+}
+
+template <typename T, typename U, int N_READS, bool reverse>
+inline void load_safe(
+    U values[N_READS],
+    const device T* input,
+    int start,
+    int total,
+    U init) {
+  if (reverse) {
+    for (int i = 0; i < N_READS; i++) {
+      values[N_READS - i - 1] =
+          (start + N_READS - i - 1 < total) ? input[i] : init;
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      values[i] = (start + i < total) ? input[i] : init;
+    }
+  }
+}
+
+template <typename U, int N_READS, bool reverse>
+inline void write_unsafe(U values[N_READS], device U* out) {
+  if (reverse) {
+    for (int i = 0; i < N_READS; i++) {
+      out[i] = values[N_READS - i - 1];
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      out[i] = values[i];
+    }
+  }
+}
+
+template <typename U, int N_READS, bool reverse>
+inline void write_safe(U values[N_READS], device U* out, int start, int total) {
+  if (reverse) {
+    for (int i = 0; i < N_READS; i++) {
+      if (start + N_READS - i - 1 < total) {
+        out[i] = values[N_READS - i - 1];
+      }
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if (start + i < total) {
+        out[i] = values[i];
+      }
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS,
+    bool inclusive,
+    bool reverse>
+[[kernel]] void contiguous_scan(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& axis_size [[buffer(2)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int simd_size = 32;
+  Op op;
+
+  // Position the pointers
+  size_t offset = (gid.y + gsize.y * size_t(gid.z)) * axis_size;
+  in += offset;
+  out += offset;
+
+  // Compute the number of simd_groups
+  uint simd_groups = lsize.x / simd_size;
+
+  // Allocate memory
+  U prefix = Op::init;
+  U values[N_READS];
+  threadgroup U simdgroup_sums[32];
+
+  // Loop over the reduced axis in blocks of size ceildiv(axis_size,
+  // N_READS*lsize)
+  //    Read block
+  //    Compute inclusive scan of the block
+  //      Compute inclusive scan per thread
+  //      Compute exclusive scan of thread sums in simdgroup
+  //      Write simdgroup sums in SM
+  //      Compute exclusive scan of simdgroup sums
+  //      Compute the output by scanning prefix, prev_simdgroup, prev_thread,
+  //      value
+  //    Write block
+
+  for (uint r = 0; r < ceildiv(axis_size, N_READS * lsize.x); r++) {
+    // Compute the block offset
+    uint offset = r * lsize.x * N_READS + lid.x * N_READS;
+
+    // Read the values
+    if (reverse) {
+      if ((offset + N_READS) < axis_size) {
+        load_unsafe<T, U, N_READS, reverse>(
+            values, in + axis_size - offset - N_READS);
+      } else {
+        load_safe<T, U, N_READS, reverse>(
+            values,
+            in + axis_size - offset - N_READS,
+            offset,
+            axis_size,
+            Op::init);
+      }
+    } else {
+      if ((offset + N_READS) < axis_size) {
+        load_unsafe<T, U, N_READS, reverse>(values, in + offset);
+      } else {
+        load_safe<T, U, N_READS, reverse>(
+            values, in + offset, offset, axis_size, Op::init);
+      }
+    }
+
+    // Compute an inclusive scan per thread
+    for (int i = 1; i < N_READS; i++) {
+      values[i] = op(values[i], values[i - 1]);
+    }
+
+    // Compute exclusive scan of thread sums
+    U prev_thread = op.simd_exclusive_scan(values[N_READS - 1]);
+
+    // Write simdgroup_sums to SM
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (simd_lane_id == simd_size - 1) {
+      simdgroup_sums[simd_group_id] = op(prev_thread, values[N_READS - 1]);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Compute exclusive scan of simdgroup_sums
+    if (simd_group_id == 0) {
+      U prev_simdgroup = op.simd_exclusive_scan(simdgroup_sums[simd_lane_id]);
+      simdgroup_sums[simd_lane_id] = prev_simdgroup;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Compute the output
+    for (int i = 0; i < N_READS; i++) {
+      values[i] = op(values[i], prefix);
+      values[i] = op(values[i], simdgroup_sums[simd_group_id]);
+      values[i] = op(values[i], prev_thread);
+    }
+
+    // Write the values
+    if (reverse) {
+      if (inclusive) {
+        if ((offset + N_READS) < axis_size) {
+          write_unsafe<U, N_READS, reverse>(
+              values, out + axis_size - offset - N_READS);
+        } else {
+          write_safe<U, N_READS, reverse>(
+              values, out + axis_size - offset - N_READS, offset, axis_size);
+        }
+      } else {
+        if (lid.x == 0 && offset == 0) {
+          out[axis_size - 1] = Op::init;
+        }
+        if ((offset + N_READS + 1) < axis_size) {
+          write_unsafe<U, N_READS, reverse>(
+              values, out + axis_size - offset - 1 - N_READS);
+        } else {
+          write_safe<U, N_READS, reverse>(
+              values,
+              out + axis_size - offset - 1 - N_READS,
+              offset + 1,
+              axis_size);
+        }
+      }
+    } else {
+      if (inclusive) {
+        if ((offset + N_READS) < axis_size) {
+          write_unsafe<U, N_READS, reverse>(values, out + offset);
+        } else {
+          write_safe<U, N_READS, reverse>(
+              values, out + offset, offset, axis_size);
+        }
+      } else {
+        if (lid.x == 0 && offset == 0) {
+          out[0] = Op::init;
+        }
+        if ((offset + N_READS + 1) < axis_size) {
+          write_unsafe<U, N_READS, reverse>(values, out + offset + 1);
+        } else {
+          write_safe<U, N_READS, reverse>(
+              values, out + offset + 1, offset + 1, axis_size);
+        }
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Share the prefix
+    if (simd_group_id == simd_groups - 1 && simd_lane_id == simd_size - 1) {
+      simdgroup_sums[0] = values[N_READS - 1];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    prefix = simdgroup_sums[0];
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N_READS,
+    bool inclusive,
+    bool reverse>
+[[kernel]] void strided_scan(
+    const device T* in [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant size_t& axis_size [[buffer(2)]],
+    const constant size_t& stride [[buffer(3)]],
+    const constant size_t& stride_blocks [[buffer(4)]],
+    uint3 gid [[threadgroup_position_in_grid]],
+    uint3 gsize [[threadgroups_per_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  constexpr int simd_size = 32;
+  constexpr int BM = 32;
+  constexpr int BN = 32;
+  constexpr int BN_pad = 32 + 16 / sizeof(U);
+  constexpr int n_simds = BN / N_READS;
+  constexpr int n_scans = BN / n_simds;
+  Op op;
+
+  threadgroup U read_buffer[BM * BN_pad];
+  U values[n_scans];
+  U prefix[n_scans];
+  for (int i = 0; i < n_scans; i++) {
+    prefix[i] = Op::init;
+  }
+
+  // Compute offsets
+  size_t full_gid = gid.y + gsize.y * size_t(gid.z);
+  size_t offset = full_gid / stride_blocks * axis_size * stride;
+  size_t global_index_x = full_gid % stride_blocks * BN;
+  uint read_offset_y = (lid.x * N_READS) / BN;
+  uint read_offset_x = (lid.x * N_READS) % BN;
+  uint scan_offset_y = simd_lane_id;
+  uint scan_offset_x = simd_group_id * n_scans;
+
+  uint stride_limit = stride - global_index_x;
+  in += offset + global_index_x + read_offset_x;
+  out += offset + global_index_x + read_offset_x;
+  threadgroup U* read_into =
+      read_buffer + read_offset_y * BN_pad + read_offset_x;
+  threadgroup U* read_from =
+      read_buffer + scan_offset_y * BN_pad + scan_offset_x;
+
+  for (uint j = 0; j < axis_size; j += BM) {
+    // Calculate the indices for the current thread
+    uint index_y = j + read_offset_y;
+    uint check_index_y = index_y;
+    if (reverse) {
+      index_y = axis_size - 1 - index_y;
+    }
+
+    // Read in SM
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (check_index_y < axis_size && (read_offset_x + N_READS) < stride_limit) {
+      for (int i = 0; i < N_READS; i++) {
+        read_into[i] = in[index_y * stride + i];
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if (check_index_y < axis_size && (read_offset_x + i) < stride_limit) {
+          read_into[i] = in[index_y * stride + i];
+        } else {
+          read_into[i] = Op::init;
+        }
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Read strided into registers
+    for (int i = 0; i < n_scans; i++) {
+      values[i] = read_from[i];
+    }
+    simdgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Perform the scan
+    for (int i = 0; i < n_scans; i++) {
+      values[i] = op.simd_scan(values[i]);
+      values[i] = op(values[i], prefix[i]);
+      prefix[i] = simd_shuffle(values[i], simd_size - 1);
+    }
+
+    // Write to SM
+    for (int i = 0; i < n_scans; i++) {
+      read_from[i] = values[i];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Write to device memory
+    if (!inclusive) {
+      if (check_index_y == 0) {
+        if ((read_offset_x + N_READS) < stride_limit) {
+          for (int i = 0; i < N_READS; i++) {
+            out[index_y * stride + i] = Op::init;
+          }
+        } else {
+          for (int i = 0; i < N_READS; i++) {
+            if ((read_offset_x + i) < stride_limit) {
+              out[index_y * stride + i] = Op::init;
+            }
+          }
+        }
+      }
+      if (reverse) {
+        index_y -= 1;
+        check_index_y += 1;
+      } else {
+        index_y += 1;
+        check_index_y += 1;
+      }
+    }
+    if (check_index_y < axis_size && (read_offset_x + N_READS) < stride_limit) {
+      for (int i = 0; i < N_READS; i++) {
+        out[index_y * stride + i] = read_into[i];
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if (check_index_y < axis_size && (read_offset_x + i) < stride_limit) {
+          out[index_y * stride + i] = read_into[i];
+        }
+      }
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/scan.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/scan.metal
new file mode 100644
index 0000000..854db9d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/scan.metal
@@ -0,0 +1,109 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_math>
+#include <metal_simdgroup>
+
+// clang-format off
+
+using namespace metal;
+
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/scan.h"
+
+#define instantiate_contiguous_scan(                                    \
+    name, itype, otype, op, inclusive, reverse, nreads)                 \
+  template [[host_name("contig_scan_" #name)]] [[kernel]] void          \
+  contiguous_scan<itype, otype, op<otype>, nreads, inclusive, reverse>( \
+      const device itype* in [[buffer(0)]],                             \
+      device otype* out [[buffer(1)]],                                  \
+      const constant size_t& axis_size [[buffer(2)]],                   \
+      uint3 gid [[threadgroup_position_in_grid]],                       \
+      uint3 gsize [[threadgroups_per_grid]],                            \
+      uint3 lid [[thread_position_in_threadgroup]],                     \
+      uint3 lsize [[threads_per_threadgroup]],                          \
+      uint simd_lane_id [[thread_index_in_simdgroup]],                  \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_strided_scan(                                    \
+    name, itype, otype, op, inclusive, reverse, nreads)              \
+  template [[host_name("strided_scan_" #name)]] [[kernel]] void      \
+  strided_scan<itype, otype, op<otype>, nreads, inclusive, reverse>( \
+      const device itype* in [[buffer(0)]],                          \
+      device otype* out [[buffer(1)]],                               \
+      const constant size_t& axis_size [[buffer(2)]],                \
+      const constant size_t& stride [[buffer(3)]],                   \
+      const constant size_t& stride_blocks [[buffer(4)]],            \
+      uint3 gid [[threadgroup_position_in_grid]],                    \
+      uint3 gsize [[threadgroups_per_grid]],                         \
+      uint3 lid [[thread_position_in_threadgroup]],                  \
+      uint simd_lane_id [[thread_index_in_simdgroup]],               \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_scan_helper(name, itype, otype, op, nreads)                                \
+  instantiate_contiguous_scan(inclusive_##name, itype, otype, op, true, false, nreads)         \
+  instantiate_contiguous_scan(exclusive_##name, itype, otype, op, false, false, nreads)        \
+  instantiate_contiguous_scan(reverse_inclusive_##name, itype, otype, op, true, true, nreads)  \
+  instantiate_contiguous_scan(reverse_exclusive_##name, itype, otype, op, false, true, nreads) \
+  instantiate_strided_scan(inclusive_##name, itype, otype, op, true, false, nreads)            \
+  instantiate_strided_scan(exclusive_##name, itype, otype, op, false, false, nreads)           \
+  instantiate_strided_scan(reverse_inclusive_##name, itype, otype, op, true, true, nreads)     \
+  instantiate_strided_scan(reverse_exclusive_##name, itype, otype, op, false, true, nreads)
+
+instantiate_scan_helper(sum_bool__int32,         bool,        int32_t,     CumSum, 4)
+instantiate_scan_helper(sum_bool__uint32,        bool,        uint32_t,    CumSum, 4)
+instantiate_scan_helper(sum_uint8_uint8,         uint8_t,     uint8_t,     CumSum, 4)
+instantiate_scan_helper(sum_uint16_uint16,       uint16_t,    uint16_t,    CumSum, 4)
+instantiate_scan_helper(sum_uint32_uint32,       uint32_t,    uint32_t,    CumSum, 4)
+instantiate_scan_helper(sum_uint64_uint64,       uint64_t,    uint64_t,    CumSum, 2)
+instantiate_scan_helper(sum_int8_int8,           int8_t,      int8_t,      CumSum, 4)
+instantiate_scan_helper(sum_int16_int16,         int16_t,     int16_t,     CumSum, 4)
+instantiate_scan_helper(sum_int32_int32,         int32_t,     int32_t,     CumSum, 4)
+instantiate_scan_helper(sum_int64_int64,         int64_t,     int64_t,     CumSum, 2)
+instantiate_scan_helper(sum_float16_float16,     half,        half,        CumSum, 4)
+instantiate_scan_helper(sum_float32_float32,     float,       float,       CumSum, 4)
+instantiate_scan_helper(sum_bfloat16_bfloat16,   bfloat16_t,  bfloat16_t,  CumSum, 4)
+instantiate_scan_helper(sum_complex64_complex64, complex64_t, complex64_t, CumSum, 2)
+instantiate_scan_helper(prod_bool__bool_,         bool,        bool,        CumProd, 4)
+instantiate_scan_helper(prod_uint8_uint8,         uint8_t,     uint8_t,     CumProd, 4)
+instantiate_scan_helper(prod_uint16_uint16,       uint16_t,    uint16_t,    CumProd, 4)
+instantiate_scan_helper(prod_uint32_uint32,       uint32_t,    uint32_t,    CumProd, 4)
+instantiate_scan_helper(prod_uint64_uint64,       uint64_t,    uint64_t,    CumProd, 2)
+instantiate_scan_helper(prod_int8_int8,           int8_t,      int8_t,      CumProd, 4)
+instantiate_scan_helper(prod_int16_int16,         int16_t,     int16_t,     CumProd, 4)
+instantiate_scan_helper(prod_int32_int32,         int32_t,     int32_t,     CumProd, 4)
+instantiate_scan_helper(prod_int64_int64,         int64_t,     int64_t,     CumProd, 2)
+instantiate_scan_helper(prod_float16_float16,     half,        half,        CumProd, 4)
+instantiate_scan_helper(prod_float32_float32,     float,       float,       CumProd, 4)
+instantiate_scan_helper(prod_bfloat16_bfloat16,   bfloat16_t,  bfloat16_t,  CumProd, 4)
+instantiate_scan_helper(prod_complex64_complex64, complex64_t, complex64_t, CumProd, 2)
+instantiate_scan_helper(max_bool__bool_,         bool,        bool,        CumMax, 4)
+instantiate_scan_helper(max_uint8_uint8,         uint8_t,     uint8_t,     CumMax, 4)
+instantiate_scan_helper(max_uint16_uint16,       uint16_t,    uint16_t,    CumMax, 4)
+instantiate_scan_helper(max_uint32_uint32,       uint32_t,    uint32_t,    CumMax, 4)
+instantiate_scan_helper(max_uint64_uint64,       uint64_t,    uint64_t,    CumMax, 2)
+instantiate_scan_helper(max_int8_int8,           int8_t,      int8_t,      CumMax, 4)
+instantiate_scan_helper(max_int16_int16,         int16_t,     int16_t,     CumMax, 4)
+instantiate_scan_helper(max_int32_int32,         int32_t,     int32_t,     CumMax, 4)
+instantiate_scan_helper(max_int64_int64,         int64_t,     int64_t,     CumMax, 2)
+instantiate_scan_helper(max_float16_float16,     half,        half,        CumMax, 4)
+instantiate_scan_helper(max_float32_float32,     float,       float,       CumMax, 4)
+instantiate_scan_helper(max_bfloat16_bfloat16,   bfloat16_t,  bfloat16_t,  CumMax, 4)
+instantiate_scan_helper(max_complex64_complex64, complex64_t, complex64_t, CumMax, 2)
+instantiate_scan_helper(min_bool__bool_,         bool,        bool,        CumMin, 4)
+instantiate_scan_helper(min_uint8_uint8,         uint8_t,     uint8_t,     CumMin, 4)
+instantiate_scan_helper(min_uint16_uint16,       uint16_t,    uint16_t,    CumMin, 4)
+instantiate_scan_helper(min_uint32_uint32,       uint32_t,    uint32_t,    CumMin, 4)
+instantiate_scan_helper(min_uint64_uint64,       uint64_t,    uint64_t,    CumMin, 2)
+instantiate_scan_helper(min_int8_int8,           int8_t,      int8_t,      CumMin, 4)
+instantiate_scan_helper(min_int16_int16,         int16_t,     int16_t,     CumMin, 4)
+instantiate_scan_helper(min_int32_int32,         int32_t,     int32_t,     CumMin, 4)
+instantiate_scan_helper(min_int64_int64,         int64_t,     int64_t,     CumMin, 2)
+instantiate_scan_helper(min_float16_float16,     half,        half,        CumMin, 4)
+instantiate_scan_helper(min_float32_float32,     float,       float,       CumMin, 4)
+instantiate_scan_helper(min_bfloat16_bfloat16,   bfloat16_t,  bfloat16_t,  CumMin, 4)
+instantiate_scan_helper(min_complex64_complex64, complex64_t, complex64_t, CumMin, 2)
+instantiate_scan_helper(logaddexp_float16_float16,     half,        half,        CumLogaddexp, 4)
+instantiate_scan_helper(logaddexp_float32_float32,     float,       float,       CumLogaddexp, 4)
+instantiate_scan_helper(logaddexp_bfloat16_bfloat16,   bfloat16_t,  bfloat16_t,  CumLogaddexp, 4)
+instantiate_scan_helper(logaddexp_complex64_complex64, complex64_t, complex64_t, CumLogaddexp, 2) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/sdpa_vector.h b/macos/third_party/mlx/mlx/backend/metal/kernels/sdpa_vector.h
new file mode 100644
index 0000000..1eec72b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/sdpa_vector.h
@@ -0,0 +1,394 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_simdgroup>
+
+using namespace metal;
+
+constant bool has_mask [[function_constant(20)]];
+constant bool query_transposed [[function_constant(21)]];
+constant bool do_causal [[function_constant(22)]];
+constant bool bool_mask [[function_constant(23)]];
+constant bool float_mask [[function_constant(24)]];
+constant bool has_sinks [[function_constant(25)]];
+constant int blocks [[function_constant(26)]];
+
+template <typename T, int D, int V = D>
+[[kernel]] void sdpa_vector(
+    const device T* queries [[buffer(0)]],
+    const device T* keys [[buffer(1)]],
+    const device T* values [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    const constant int& gqa_factor [[buffer(4)]],
+    const constant int& N [[buffer(5)]],
+    const constant size_t& k_head_stride [[buffer(6)]],
+    const constant size_t& k_seq_stride [[buffer(7)]],
+    const constant size_t& v_head_stride [[buffer(8)]],
+    const constant size_t& v_seq_stride [[buffer(9)]],
+    const constant float& scale [[buffer(10)]],
+    const device bool* bmask [[buffer(11), function_constant(bool_mask)]],
+    const device T* fmask [[buffer(12), function_constant(float_mask)]],
+    const constant int& mask_kv_seq_stride
+    [[buffer(13), function_constant(has_mask)]],
+    const constant int& mask_q_seq_stride
+    [[buffer(14), function_constant(has_mask)]],
+    const constant int& mask_head_stride
+    [[buffer(15), function_constant(has_mask)]],
+    const device T* sinks [[buffer(16), function_constant(has_sinks)]],
+    const constant int& num_q_heads
+    [[buffer(17), function_constant(has_sinks)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int BN = 32;
+  constexpr int BD = 32;
+  constexpr int qk_per_thread = D / BD;
+  constexpr int v_per_thread = V / BD;
+  int inner_k_stride = BN * int(k_seq_stride);
+  int inner_v_stride = BN * int(v_seq_stride);
+
+  typedef float U;
+
+  thread U q[qk_per_thread];
+  thread U k[qk_per_thread];
+  thread U o[v_per_thread];
+
+  threadgroup U outputs[BN * BD];
+  threadgroup U max_scores[BN];
+  threadgroup U sum_exp_scores[BN];
+
+  // Adjust positions
+  const int q_batch_head_idx = tid.x;
+  const int q_seq_idx = tid.y;
+  const int kv_head_idx = q_batch_head_idx / gqa_factor;
+  const int o_offset = q_batch_head_idx * tpg.y + q_seq_idx;
+  const int q_offset =
+      query_transposed ? tpg.x * q_seq_idx + q_batch_head_idx : o_offset;
+  queries += q_offset * D + simd_lid * qk_per_thread;
+  keys += kv_head_idx * k_head_stride + simd_gid * k_seq_stride +
+      simd_lid * qk_per_thread;
+  values += kv_head_idx * v_head_stride + simd_gid * v_seq_stride +
+      simd_lid * v_per_thread;
+  if (bool_mask) {
+    bmask += q_batch_head_idx * mask_head_stride +
+        simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride;
+  }
+  if (float_mask) {
+    fmask += q_batch_head_idx * mask_head_stride +
+        simd_gid * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride;
+  }
+
+  out += o_offset * V + simd_gid * v_per_thread;
+
+  // Read the query and 0 the output accumulator
+  for (int i = 0; i < qk_per_thread; i++) {
+    q[i] = static_cast<U>(scale) * queries[i];
+  }
+  for (int i = 0; i < v_per_thread; i++) {
+    o[i] = 0;
+  }
+
+  U max_score = Limits<U>::finite_min;
+  U sum_exp_score = 0;
+  if (has_sinks && simd_gid == 0) {
+    max_score = static_cast<U>(sinks[q_batch_head_idx % num_q_heads]);
+    sum_exp_score = 1;
+  }
+
+  // For each key
+  for (int i = simd_gid; i < N; i += BN) {
+    bool use_key = true;
+    if (do_causal) {
+      use_key = i <= (N - int(tpg.y) + int(q_seq_idx));
+    } else if (bool_mask) {
+      use_key = bmask[0];
+    } else if (float_mask) {
+      use_key = (fmask[0] >= Limits<T>::finite_min);
+    }
+    if (use_key) {
+      // Read the key
+      for (int j = 0; j < qk_per_thread; j++) {
+        k[j] = keys[j];
+      }
+
+      // Compute the i-th score
+      U score = 0;
+      for (int j = 0; j < qk_per_thread; j++) {
+        score += q[j] * k[j];
+      }
+      score = simd_sum(score);
+      if (float_mask) {
+        score += static_cast<U>(fmask[0]);
+      }
+
+      // Update the accumulators
+      U new_max = max(max_score, score);
+      U factor = fast::exp(max_score - new_max);
+      U exp_score = fast::exp(score - new_max);
+
+      max_score = new_max;
+      sum_exp_score = sum_exp_score * factor + exp_score;
+
+      // Update the output accumulator
+      for (int j = 0; j < v_per_thread; j++) {
+        o[j] = o[j] * factor + exp_score * values[j];
+      }
+    }
+
+    // Move the pointers to the next kv
+    keys += inner_k_stride;
+    values += inner_v_stride;
+    if (bool_mask) {
+      bmask += BN * mask_kv_seq_stride;
+    }
+    if (float_mask) {
+      fmask += BN * mask_kv_seq_stride;
+    }
+  }
+
+  // Each thread has a partial part of the output so we need to combine them.
+
+  // First let's communicate the max and sum_exp
+  if (simd_lid == 0) {
+    max_scores[simd_gid] = max_score;
+    sum_exp_scores[simd_gid] = sum_exp_score;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  max_score = max_scores[simd_lid];
+  U new_max = simd_max(max_score);
+  U factor = fast::exp(max_score - new_max);
+  sum_exp_score = simd_sum(sum_exp_scores[simd_lid] * factor);
+
+  // Now we need to aggregate all the outputs
+  for (int i = 0; i < v_per_thread; i++) {
+    outputs[simd_lid * BD + simd_gid] = o[i];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    o[i] = simd_sum(outputs[simd_gid * BD + simd_lid] * factor);
+    o[i] = sum_exp_score == 0 ? o[i] : (o[i] / sum_exp_score);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // And write the output
+  if (simd_lid == 0) {
+    for (int i = 0; i < v_per_thread; i++) {
+      out[i] = static_cast<T>(o[i]);
+    }
+  }
+}
+
+template <typename T, int D, int V = D>
+[[kernel]] void sdpa_vector_2pass_1(
+    const device T* queries [[buffer(0)]],
+    const device T* keys [[buffer(1)]],
+    const device T* values [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    device float* sums [[buffer(4)]],
+    device float* maxs [[buffer(5)]],
+    const constant int& N [[buffer(7)]],
+    const constant size_t& k_head_stride [[buffer(8)]],
+    const constant size_t& k_seq_stride [[buffer(9)]],
+    const constant size_t& v_head_stride [[buffer(10)]],
+    const constant size_t& v_seq_stride [[buffer(11)]],
+    const constant float& scale [[buffer(12)]],
+    const device bool* bmask [[buffer(13), function_constant(bool_mask)]],
+    const device T* fmask [[buffer(14), function_constant(float_mask)]],
+    const constant int& mask_kv_seq_stride
+    [[buffer(15), function_constant(has_mask)]],
+    const constant int& mask_q_seq_stride
+    [[buffer(16), function_constant(has_mask)]],
+    const constant int& mask_head_stride
+    [[buffer(17), function_constant(has_mask)]],
+    const device T* sinks [[buffer(18), function_constant(has_sinks)]],
+    uint3 tptg [[threads_per_threadgroup]],
+    uint3 tidtg [[thread_position_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int BD = 32;
+  constexpr int qk_per_thread = D / BD;
+  constexpr int v_per_thread = V / BD;
+
+  typedef float U;
+
+  thread U q[qk_per_thread];
+  thread U o[v_per_thread] = {0};
+
+  // Adjust positions
+  const int kv_head_idx = tid.x;
+  const int batch_idx = tid.y;
+  const int block_idx = tid.z;
+  const int gqa_factor = tptg.y;
+  const int q_seq_len = tptg.z;
+  const int q_seq_idx = tidtg.z;
+  const int q_head_idx = gqa_factor * kv_head_idx + tidtg.y;
+  const int num_kv_heads = tpg.x;
+  const int num_q_heads = num_kv_heads * gqa_factor;
+  const int q_batch_head_idx = (batch_idx * num_q_heads + q_head_idx);
+  const int o_offset = q_batch_head_idx * q_seq_len + q_seq_idx;
+  const int q_offset =
+      query_transposed ? num_q_heads * q_seq_idx + q_batch_head_idx : o_offset;
+
+  queries += q_offset * D + simd_lid * qk_per_thread;
+
+  const int kv_batch_head_idx = batch_idx * num_kv_heads + kv_head_idx;
+  keys += kv_batch_head_idx * k_head_stride + block_idx * k_seq_stride +
+      simd_lid * qk_per_thread;
+  values += kv_batch_head_idx * v_head_stride + block_idx * v_seq_stride +
+      simd_lid * v_per_thread;
+  out += o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread;
+  if (bool_mask) {
+    bmask += q_batch_head_idx * mask_head_stride +
+        block_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride;
+  }
+  if (float_mask) {
+    fmask += q_batch_head_idx * mask_head_stride +
+        block_idx * mask_kv_seq_stride + q_seq_idx * mask_q_seq_stride;
+  }
+  sums += o_offset * blocks + block_idx;
+  maxs += o_offset * blocks + block_idx;
+
+  // Read the query
+  for (int i = 0; i < qk_per_thread; i++) {
+    q[i] = static_cast<U>(scale) * queries[i];
+  }
+
+  U max_score = Limits<U>::finite_min;
+  U sum_exp_score = 0;
+  if (has_sinks && block_idx == 0) {
+    max_score = static_cast<U>(sinks[q_head_idx]);
+    sum_exp_score = 1;
+  }
+
+  // For each key
+  for (int i = block_idx; i < N; i += blocks) {
+    bool use_key = true;
+    if (do_causal) {
+      use_key = i <= (N - q_seq_len + int(q_seq_idx));
+    } else if (bool_mask) {
+      use_key = bmask[0];
+    } else if (float_mask) {
+      use_key = (fmask[0] >= Limits<T>::finite_min);
+    }
+    if (use_key) {
+      // Compute the i-th score
+      U score = 0;
+      for (int i = 0; i < qk_per_thread; i++) {
+        score += q[i] * keys[i];
+      }
+      score = simd_sum(score);
+
+      if (float_mask) {
+        score += fmask[0];
+      }
+
+      // Update the accumulators
+      U new_max = max(max_score, score);
+      U factor = fast::exp(max_score - new_max);
+      U exp_score = fast::exp(score - new_max);
+
+      max_score = new_max;
+      sum_exp_score = sum_exp_score * factor + exp_score;
+
+      // Update the output accumulator
+      for (int i = 0; i < v_per_thread; i++) {
+        o[i] = o[i] * factor + exp_score * values[i];
+      }
+    }
+
+    // Move the pointers to the next kv
+    keys += blocks * int(k_seq_stride);
+    values += blocks * int(v_seq_stride);
+    if (bool_mask) {
+      bmask += blocks * mask_kv_seq_stride;
+    }
+    if (float_mask) {
+      fmask += blocks * mask_kv_seq_stride;
+    }
+  }
+
+  // Write the sum and max and outputs
+  if (simd_lid == 0) {
+    sums[0] = sum_exp_score;
+    maxs[0] = max_score;
+  }
+
+  for (int i = 0; i < v_per_thread; i++) {
+    out[i] = static_cast<T>(o[i]);
+  }
+}
+
+template <typename T, int D>
+[[kernel]] void sdpa_vector_2pass_2(
+    const device T* partials [[buffer(0)]],
+    const device float* sums [[buffer(1)]],
+    const device float* maxs [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    const constant int& blocks [[buffer(4)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int BN = 32;
+  constexpr int BD = 32;
+  constexpr int elem_per_thread = D / BD;
+
+  typedef float U;
+
+  thread U o[elem_per_thread] = {0};
+  threadgroup U outputs[BN * BD];
+
+  // Adjust positions
+  const int head_idx = tid.x;
+  const int q_seq_idx = tid.y;
+  const int q_offset = head_idx * tpg.y + q_seq_idx;
+  partials += q_offset * blocks * D + simd_gid * D + simd_lid * elem_per_thread;
+  sums += q_offset * blocks;
+  maxs += q_offset * blocks;
+  out += q_offset * D + simd_gid * elem_per_thread;
+
+  // Set defaults
+  U sum_exp_score = 0.0;
+  U max_score = Limits<U>::finite_min;
+
+  // Reduce the max
+  for (int b = 0; b < blocks / BN; ++b) {
+    max_score = max(max_score, maxs[simd_lid + BN * b]);
+  }
+  max_score = simd_max(max_score);
+
+  // Reduce the d
+  for (int b = 0; b < blocks / BN; ++b) {
+    U factor = fast::exp(maxs[simd_lid + BN * b] - max_score);
+    sum_exp_score += factor * sums[simd_lid + BN * b];
+  }
+  sum_exp_score = simd_sum(sum_exp_score);
+
+  // Reduce the sum exp and partials
+  for (int b = 0; b < blocks / BN; ++b) {
+    U factor = fast::exp(maxs[simd_gid] - max_score);
+
+    // Update the output accumulator
+    for (int i = 0; i < elem_per_thread; i++) {
+      o[i] += factor * static_cast<U>(partials[i]);
+    }
+    maxs += BN;
+    sums += BN;
+    partials += BN * D;
+  }
+
+  // Use shared memory to transpose and reduce the final block
+  for (int i = 0; i < elem_per_thread; i++) {
+    outputs[simd_lid * BD + simd_gid] = o[i];
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    o[i] = simd_sum(outputs[simd_gid * BD + simd_lid]);
+    o[i] = sum_exp_score == 0 ? o[i] : (o[i] / sum_exp_score);
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // And write the output
+  if (simd_lid == 0) {
+    for (int i = 0; i < elem_per_thread; i++) {
+      out[i] = static_cast<T>(o[i]);
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/softmax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/softmax.h
new file mode 100644
index 0000000..6ea4ac7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/softmax.h
@@ -0,0 +1,190 @@
+// Copyright © 2023-2024 Apple Inc.
+
+template <typename T>
+inline T softmax_exp(T x) {
+  // Softmax doesn't need high precision exponential cause x is gonna be in
+  // (-oo, 0] anyway and subsequently it will be divided by sum(exp(x_i)).
+  return fast::exp(x);
+}
+
+template <typename T, typename AccT = T, int N_READS = SOFTMAX_N_READS>
+[[kernel]] void softmax_single_row(
+    const device T* in,
+    device T* out,
+    constant int& axis_size,
+    uint gid [[threadgroup_position_in_grid]],
+    uint _lid [[thread_position_in_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  int lid = _lid;
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup AccT local_max[SIMD_SIZE];
+  threadgroup AccT local_normalizer[SIMD_SIZE];
+
+  AccT ld[N_READS];
+
+  in += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      ld[i] = AccT(in[i]);
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      ld[i] =
+          ((lid * N_READS + i) < axis_size) ? AccT(in[i]) : Limits<AccT>::min;
+    }
+  }
+  if (simd_group_id == 0) {
+    local_max[simd_lane_id] = Limits<AccT>::min;
+    local_normalizer[simd_lane_id] = 0;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Get the max
+  AccT maxval = Limits<AccT>::finite_min;
+  for (int i = 0; i < N_READS; i++) {
+    maxval = (maxval < ld[i]) ? ld[i] : maxval;
+  }
+  maxval = simd_max(maxval);
+  if (simd_lane_id == 0) {
+    local_max[simd_group_id] = maxval;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    maxval = simd_max(local_max[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_max[0] = maxval;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  maxval = local_max[0];
+
+  // Compute exp(x_i - maxval) and store the partial sums in local_normalizer
+  AccT normalizer = 0;
+  for (int i = 0; i < N_READS; i++) {
+    AccT exp_x = softmax_exp(ld[i] - maxval);
+    ld[i] = exp_x;
+    normalizer += exp_x;
+  }
+  normalizer = simd_sum(normalizer);
+  if (simd_lane_id == 0) {
+    local_normalizer[simd_group_id] = normalizer;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_group_id == 0) {
+    normalizer = simd_sum(local_normalizer[simd_lane_id]);
+    if (simd_lane_id == 0) {
+      local_normalizer[0] = normalizer;
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  normalizer = 1 / local_normalizer[0];
+
+  // Normalize and write to the output
+  out += gid * size_t(axis_size) + lid * N_READS;
+  if (lid * N_READS + N_READS <= axis_size) {
+    for (int i = 0; i < N_READS; i++) {
+      out[i] = T(ld[i] * normalizer);
+    }
+  } else {
+    for (int i = 0; i < N_READS; i++) {
+      if ((lid * N_READS + i) < axis_size) {
+        out[i] = T(ld[i] * normalizer);
+      }
+    }
+  }
+}
+
+template <typename T, typename AccT = T, int N_READS = SOFTMAX_N_READS>
+[[kernel]] void softmax_looped(
+    const device T* in,
+    device T* out,
+    constant int& axis_size,
+    uint gid [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+  in += gid * size_t(axis_size);
+
+  constexpr int SIMD_SIZE = 32;
+
+  threadgroup AccT local_max[SIMD_SIZE];
+  threadgroup AccT local_normalizer[SIMD_SIZE];
+
+  // Get the max and the normalizer in one go
+  AccT prevmax;
+  AccT maxval = Limits<AccT>::finite_min;
+  AccT normalizer = 0;
+  for (int r = 0; r < static_cast<int>(ceildiv(axis_size, N_READS * lsize));
+       r++) {
+    int offset = r * lsize * N_READS + lid * N_READS;
+    AccT vals[N_READS];
+    if (offset + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        vals[i] = AccT(in[offset + i]);
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        vals[i] =
+            (offset + i < axis_size) ? AccT(in[offset + i]) : Limits<AccT>::min;
+      }
+    }
+    prevmax = maxval;
+    for (int i = 0; i < N_READS; i++) {
+      maxval = (maxval < vals[i]) ? vals[i] : maxval;
+    }
+    normalizer *= softmax_exp(prevmax - maxval);
+    for (int i = 0; i < N_READS; i++) {
+      normalizer += softmax_exp(vals[i] - maxval);
+    }
+  }
+  // Now we got partial normalizer of N_READS * ceildiv(axis_size, N_READS *
+  // lsize) parts. We need to combine them.
+  //    1. We start by finding the max across simd groups
+  //    2. We then change the partial normalizers to account for a possible
+  //       change in max
+  //    3. We sum all normalizers
+  prevmax = maxval;
+  maxval = simd_max(maxval);
+  normalizer *= softmax_exp(prevmax - maxval);
+  normalizer = simd_sum(normalizer);
+
+  // Now the normalizer and max value is correct for each simdgroup. We write
+  // them shared memory and combine them.
+  prevmax = maxval;
+  if (simd_lane_id == 0) {
+    local_max[simd_group_id] = maxval;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  maxval = simd_max(local_max[simd_lane_id]);
+  normalizer *= softmax_exp(prevmax - maxval);
+  if (simd_lane_id == 0) {
+    local_normalizer[simd_group_id] = normalizer;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  normalizer = simd_sum(local_normalizer[simd_lane_id]);
+  normalizer = 1 / normalizer;
+
+  // Finally given the normalizer and max value we can directly write the
+  // softmax output
+  out += gid * size_t(axis_size);
+  for (int r = 0; r < static_cast<int>(ceildiv(axis_size, N_READS * lsize));
+       r++) {
+    int offset = r * lsize * N_READS + lid * N_READS;
+    if (offset + N_READS <= axis_size) {
+      for (int i = 0; i < N_READS; i++) {
+        out[offset + i] = T(softmax_exp(in[offset + i] - maxval) * normalizer);
+      }
+    } else {
+      for (int i = 0; i < N_READS; i++) {
+        if (offset + i < axis_size) {
+          out[offset + i] =
+              T(softmax_exp(in[offset + i] - maxval) * normalizer);
+        }
+      }
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/softmax.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/softmax.metal
new file mode 100644
index 0000000..79d5d3f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/softmax.metal
@@ -0,0 +1,24 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_common>
+#include <metal_simdgroup>
+
+using namespace metal;
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/softmax.h"
+
+#define instantiate_softmax(name, itype)                                \
+  instantiate_kernel("block_softmax_" #name, softmax_single_row, itype) \
+  instantiate_kernel("looped_softmax_" #name, softmax_looped, itype)
+
+#define instantiate_softmax_precise(name, itype)                                       \
+  instantiate_kernel("block_softmax_precise_" #name, softmax_single_row, itype, float) \
+  instantiate_kernel("looped_softmax_precise_" #name, softmax_looped, itype, float)
+
+instantiate_softmax(float32, float)
+instantiate_softmax(float16, half)
+instantiate_softmax(bfloat16, bfloat16_t)
+instantiate_softmax_precise(float16, half)
+instantiate_softmax_precise(bfloat16, bfloat16_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/sort.h b/macos/third_party/mlx/mlx/backend/metal/kernels/sort.h
new file mode 100644
index 0000000..0d35733
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/sort.h
@@ -0,0 +1,719 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#define MLX_MTL_CONST static constant constexpr const
+#define MLX_MTL_LOOP_UNROLL _Pragma("clang loop unroll(full)")
+
+using namespace metal;
+
+// Based on GPU merge sort algorithm at
+// https://github.com/NVIDIA/cccl/tree/main/cub/cub
+
+///////////////////////////////////////////////////////////////////////////////
+// Thread-level sort
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+METAL_FUNC void thread_swap(thread T& a, thread T& b) {
+  T w = a;
+  a = b;
+  b = w;
+}
+
+template <typename T, typename = void>
+struct Init {
+  static constexpr constant T v = Limits<T>::max;
+};
+
+template <typename T>
+struct Init<T, metal::enable_if_t<metal::is_floating_point_v<T>>> {
+  static constexpr constant T v = metal::numeric_limits<T>::quiet_NaN();
+};
+
+template <typename T>
+struct LessThan {
+  static constexpr constant T init = Init<T>::v;
+  METAL_FUNC bool operator()(T a, T b) const {
+    if constexpr (
+        metal::is_floating_point_v<T> || metal::is_same_v<T, complex64_t>) {
+      bool an = isnan(a);
+      bool bn = isnan(b);
+      if (an | bn) {
+        return (!an) & bn;
+      }
+    }
+    return a < b;
+  }
+};
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    short N_PER_THREAD,
+    typename CompareOp>
+struct ThreadSort {
+  static METAL_FUNC void sort(
+      thread ValT (&vals)[N_PER_THREAD],
+      thread IdxT (&idxs)[N_PER_THREAD]) {
+    CompareOp op;
+    MLX_MTL_LOOP_UNROLL
+    for (short i = 0; i < N_PER_THREAD; ++i) {
+      MLX_MTL_LOOP_UNROLL
+      for (short j = i & 1; j < N_PER_THREAD - 1; j += 2) {
+        if (op(vals[j + 1], vals[j])) {
+          thread_swap(vals[j + 1], vals[j]);
+          if (ARG_SORT) {
+            thread_swap(idxs[j + 1], idxs[j]);
+          }
+        }
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Threadgroup-level sort
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD,
+    typename CompareOp>
+struct BlockMergeSort {
+  using thread_sort_t =
+      ThreadSort<ValT, IdxT, ARG_SORT, N_PER_THREAD, CompareOp>;
+  static METAL_FUNC int merge_partition(
+      const threadgroup ValT* As,
+      const threadgroup ValT* Bs,
+      short A_sz,
+      short B_sz,
+      short sort_md) {
+    CompareOp op;
+
+    short A_st = max(0, sort_md - B_sz);
+    short A_ed = min(sort_md, A_sz);
+
+    while (A_st < A_ed) {
+      short md = A_st + (A_ed - A_st) / 2;
+      auto a = As[md];
+      auto b = Bs[sort_md - 1 - md];
+
+      if (op(b, a)) {
+        A_ed = md;
+      } else {
+        A_st = md + 1;
+      }
+    }
+
+    return A_ed;
+  }
+
+  static METAL_FUNC void merge_step(
+      const threadgroup ValT* As,
+      const threadgroup ValT* Bs,
+      const threadgroup IdxT* As_idx,
+      const threadgroup IdxT* Bs_idx,
+      short A_sz,
+      short B_sz,
+      thread ValT (&vals)[N_PER_THREAD],
+      thread IdxT (&idxs)[N_PER_THREAD]) {
+    CompareOp op;
+    short a_idx = 0;
+    short b_idx = 0;
+
+    for (int i = 0; i < N_PER_THREAD; ++i) {
+      auto a = (a_idx < A_sz) ? As[a_idx] : ValT(CompareOp::init);
+      auto b = (b_idx < B_sz) ? Bs[b_idx] : ValT(CompareOp::init);
+      bool pred = (b_idx < B_sz) && (a_idx >= A_sz || op(b, a));
+
+      vals[i] = pred ? b : a;
+      if (ARG_SORT) {
+        if (pred) {
+          idxs[i] = Bs_idx[b_idx];
+        } else {
+          idxs[i] = (a_idx < A_sz) ? As_idx[a_idx] : IdxT(0);
+        }
+      }
+
+      b_idx += short(pred);
+      a_idx += short(!pred);
+    }
+  }
+
+  static METAL_FUNC void sort(
+      threadgroup ValT* tgp_vals [[threadgroup(0)]],
+      threadgroup IdxT* tgp_idxs [[threadgroup(1)]],
+      int size_sorted_axis,
+      uint3 lid [[thread_position_in_threadgroup]]) {
+    // Get thread location
+    int idx = lid.x * N_PER_THREAD;
+
+    // Load from shared memory
+    thread ValT thread_vals[N_PER_THREAD];
+    thread IdxT thread_idxs[N_PER_THREAD];
+    for (int i = 0; i < N_PER_THREAD; ++i) {
+      thread_vals[i] = tgp_vals[idx + i];
+      if (ARG_SORT) {
+        thread_idxs[i] = tgp_idxs[idx + i];
+      }
+    }
+
+    // Per thread sort
+    if (idx < size_sorted_axis) {
+      thread_sort_t::sort(thread_vals, thread_idxs);
+    }
+
+    // Do merges using threadgroup memory
+    for (int merge_threads = 2; merge_threads <= BLOCK_THREADS;
+         merge_threads *= 2) {
+      // Update threadgroup memory
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      for (int i = 0; i < N_PER_THREAD; ++i) {
+        tgp_vals[idx + i] = thread_vals[i];
+        if (ARG_SORT) {
+          tgp_idxs[idx + i] = thread_idxs[i];
+        }
+      }
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Find location in merge step
+      int merge_group = lid.x / merge_threads;
+      int merge_lane = lid.x % merge_threads;
+
+      int sort_sz = N_PER_THREAD * merge_threads;
+      int sort_st = N_PER_THREAD * merge_threads * merge_group;
+
+      // As = tgp_vals[A_st:A_ed] is sorted
+      // Bs = tgp_vals[B_st:B_ed] is sorted
+      int A_st = sort_st;
+      int A_ed = sort_st + sort_sz / 2;
+      int B_st = sort_st + sort_sz / 2;
+      int B_ed = sort_st + sort_sz;
+
+      const threadgroup ValT* As = tgp_vals + A_st;
+      const threadgroup ValT* Bs = tgp_vals + B_st;
+      int A_sz = A_ed - A_st;
+      int B_sz = B_ed - B_st;
+
+      // Find a partition of merge elements
+      //  Ci = merge(As[partition:], Bs[sort_md - partition:])
+      //       of size N_PER_THREAD for each merge lane i
+      //  C = [Ci] is sorted
+      int sort_md = N_PER_THREAD * merge_lane;
+      int partition = merge_partition(As, Bs, A_sz, B_sz, sort_md);
+
+      As += partition;
+      Bs += sort_md - partition;
+
+      A_sz -= partition;
+      B_sz -= sort_md - partition;
+
+      const threadgroup IdxT* As_idx =
+          ARG_SORT ? tgp_idxs + A_st + partition : nullptr;
+      const threadgroup IdxT* Bs_idx =
+          ARG_SORT ? tgp_idxs + B_st + sort_md - partition : nullptr;
+
+      // Merge starting at the partition and store results in thread registers
+      merge_step(As, Bs, As_idx, Bs_idx, A_sz, B_sz, thread_vals, thread_idxs);
+    }
+
+    // Write out to shared memory
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for (int i = 0; i < N_PER_THREAD; ++i) {
+      tgp_vals[idx + i] = thread_vals[i];
+      if (ARG_SORT) {
+        tgp_idxs[idx + i] = thread_idxs[i];
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Kernel sort
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename U,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD,
+    typename CompareOp = LessThan<T>>
+struct KernelMergeSort {
+  using ValT = T;
+  using IdxT = uint;
+  using block_merge_sort_t = BlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD,
+      CompareOp>;
+
+  MLX_MTL_CONST short N_PER_BLOCK = BLOCK_THREADS * N_PER_THREAD;
+
+  static METAL_FUNC void block_sort(
+      const device T* inp,
+      device U* out,
+      const constant int& size_sorted_axis,
+      const constant int& in_stride_sorted_axis,
+      const constant int& out_stride_sorted_axis,
+      const constant int& in_stride_segment_axis,
+      const constant int& out_stride_segment_axis,
+      threadgroup ValT* tgp_vals,
+      threadgroup IdxT* tgp_idxs,
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]]) {
+    // tid.y tells us the segment index
+    inp += tid.y * in_stride_segment_axis;
+    out += tid.y * out_stride_segment_axis;
+
+    // Copy into threadgroup memory
+    for (short i = lid.x; i < N_PER_BLOCK; i += BLOCK_THREADS) {
+      tgp_vals[i] = i < size_sorted_axis ? inp[i * in_stride_sorted_axis]
+                                         : ValT(CompareOp::init);
+      if (ARG_SORT) {
+        tgp_idxs[i] = i;
+      }
+    }
+
+    // Sort elements within the block
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    block_merge_sort_t::sort(tgp_vals, tgp_idxs, size_sorted_axis, lid);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Write output
+    for (int i = lid.x; i < size_sorted_axis; i += BLOCK_THREADS) {
+      if (ARG_SORT) {
+        out[i * out_stride_sorted_axis] = tgp_idxs[i];
+      } else {
+        out[i * out_stride_sorted_axis] = tgp_vals[i];
+      }
+    }
+  }
+};
+
+template <
+    typename T,
+    typename U,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD>
+[[kernel, max_total_threads_per_threadgroup(BLOCK_THREADS)]] void block_sort(
+    const device T* inp [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant int& size_sorted_axis [[buffer(2)]],
+    const constant int& in_stride_sorted_axis [[buffer(3)]],
+    const constant int& out_stride_sorted_axis [[buffer(4)]],
+    const constant int& in_stride_segment_axis [[buffer(5)]],
+    const constant int& out_stride_segment_axis [[buffer(6)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  using sort_kernel =
+      KernelMergeSort<T, U, ARG_SORT, BLOCK_THREADS, N_PER_THREAD>;
+  using ValT = typename sort_kernel::ValT;
+  using IdxT = typename sort_kernel::IdxT;
+
+  if (ARG_SORT) {
+    threadgroup ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    threadgroup IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        in_stride_segment_axis,
+        out_stride_segment_axis,
+        tgp_vals,
+        tgp_idxs,
+        tid,
+        lid);
+  } else {
+    threadgroup ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        in_stride_segment_axis,
+        out_stride_segment_axis,
+        tgp_vals,
+        nullptr,
+        tid,
+        lid);
+  }
+}
+
+constant constexpr const int zero_helper = 0;
+
+template <
+    typename T,
+    typename U,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD>
+[[kernel, max_total_threads_per_threadgroup(BLOCK_THREADS)]] void block_sort_nc(
+    const device T* inp [[buffer(0)]],
+    device U* out [[buffer(1)]],
+    const constant int& size_sorted_axis [[buffer(2)]],
+    const constant int& in_stride_sorted_axis [[buffer(3)]],
+    const constant int& out_stride_sorted_axis [[buffer(4)]],
+    const constant int& nc_dim [[buffer(5)]],
+    const constant int* nc_shape [[buffer(6)]],
+    const constant int64_t* in_nc_strides [[buffer(7)]],
+    const constant int64_t* out_nc_strides [[buffer(8)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  using sort_kernel =
+      KernelMergeSort<T, U, ARG_SORT, BLOCK_THREADS, N_PER_THREAD>;
+  using ValT = typename sort_kernel::ValT;
+  using IdxT = typename sort_kernel::IdxT;
+
+  auto in_block_idx = elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim);
+  auto out_block_idx = elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim);
+  inp += in_block_idx;
+  out += out_block_idx;
+
+  if (ARG_SORT) {
+    threadgroup ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    threadgroup IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        zero_helper,
+        zero_helper,
+        tgp_vals,
+        tgp_idxs,
+        tid,
+        lid);
+  } else {
+    threadgroup ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+    sort_kernel::block_sort(
+        inp,
+        out,
+        size_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        zero_helper,
+        zero_helper,
+        tgp_vals,
+        nullptr,
+        tid,
+        lid);
+  }
+}
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD,
+    typename CompareOp = LessThan<ValT>>
+struct KernelMultiBlockMergeSort {
+  using block_merge_sort_t = BlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD,
+      CompareOp>;
+
+  MLX_MTL_CONST short N_PER_BLOCK = BLOCK_THREADS * N_PER_THREAD;
+
+  static METAL_FUNC void block_sort(
+      const device ValT* inp,
+      device ValT* out_vals,
+      device IdxT* out_idxs,
+      const constant int& size_sorted_axis,
+      const constant int& stride_sorted_axis,
+      threadgroup ValT* tgp_vals,
+      threadgroup IdxT* tgp_idxs,
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]]) {
+    // tid.y tells us the segment index
+    int base_idx = tid.x * N_PER_BLOCK;
+
+    // Copy into threadgroup memory
+    for (short i = lid.x; i < N_PER_BLOCK; i += BLOCK_THREADS) {
+      int idx = base_idx + i;
+      tgp_vals[i] = idx < size_sorted_axis ? inp[idx * stride_sorted_axis]
+                                           : ValT(CompareOp::init);
+      tgp_idxs[i] = idx;
+    }
+
+    // Sort elements within the block
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    block_merge_sort_t::sort(tgp_vals, tgp_idxs, size_sorted_axis, lid);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Write output
+    for (int i = lid.x; i < N_PER_BLOCK; i += BLOCK_THREADS) {
+      int idx = base_idx + i;
+      if (idx < size_sorted_axis) {
+        out_vals[idx] = tgp_vals[i];
+        out_idxs[idx] = tgp_idxs[i];
+      }
+    }
+  }
+
+  static METAL_FUNC int merge_partition(
+      const device ValT* As,
+      const device ValT* Bs,
+      int A_sz,
+      int B_sz,
+      int sort_md) {
+    CompareOp op;
+
+    int A_st = max(0, sort_md - B_sz);
+    int A_ed = min(sort_md, A_sz);
+
+    while (A_st < A_ed) {
+      int md = A_st + (A_ed - A_st) / 2;
+      auto a = As[md];
+      auto b = Bs[sort_md - 1 - md];
+
+      if (op(b, a)) {
+        A_ed = md;
+      } else {
+        A_st = md + 1;
+      }
+    }
+
+    return A_ed;
+  }
+};
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD>
+[[kernel, max_total_threads_per_threadgroup(BLOCK_THREADS)]] void mb_block_sort(
+    const device ValT* inp [[buffer(0)]],
+    device ValT* out_vals [[buffer(1)]],
+    device IdxT* out_idxs [[buffer(2)]],
+    const constant int& size_sorted_axis [[buffer(3)]],
+    const constant int& stride_sorted_axis [[buffer(4)]],
+    const constant int& nc_dim [[buffer(5)]],
+    const constant int* nc_shape [[buffer(6)]],
+    const constant int64_t* nc_strides [[buffer(7)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  using sort_kernel = KernelMultiBlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD>;
+
+  auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim);
+  inp += block_idx;
+  out_vals += tid.y * size_sorted_axis;
+  out_idxs += tid.y * size_sorted_axis;
+
+  threadgroup ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+  threadgroup IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+
+  sort_kernel::block_sort(
+      inp,
+      out_vals,
+      out_idxs,
+      size_sorted_axis,
+      stride_sorted_axis,
+      tgp_vals,
+      tgp_idxs,
+      tid,
+      lid);
+}
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD>
+[[kernel]] void mb_block_partition(
+    device IdxT* block_partitions [[buffer(0)]],
+    const device ValT* dev_vals [[buffer(1)]],
+    const device IdxT* dev_idxs [[buffer(2)]],
+    const constant int& size_sorted_axis [[buffer(3)]],
+    const constant int& merge_tiles [[buffer(4)]],
+    const constant int& n_blocks [[buffer(5)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 tgp_dims [[threads_per_threadgroup]]) {
+  using sort_kernel = KernelMultiBlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD>;
+
+  block_partitions += tid.y * tgp_dims.x;
+  dev_vals += tid.y * size_sorted_axis;
+  dev_idxs += tid.y * size_sorted_axis;
+
+  for (int i = lid.x; i <= n_blocks; i += tgp_dims.x) {
+    // Find location in merge step
+    int merge_group = i / merge_tiles;
+    int merge_lane = i % merge_tiles;
+
+    int sort_sz = sort_kernel::N_PER_BLOCK * merge_tiles;
+    int sort_st = sort_kernel::N_PER_BLOCK * merge_tiles * merge_group;
+
+    int A_st = min(size_sorted_axis, sort_st);
+    int A_ed = min(size_sorted_axis, sort_st + sort_sz / 2);
+    int B_st = A_ed;
+    int B_ed = min(size_sorted_axis, B_st + sort_sz / 2);
+
+    int partition_at = min(B_ed - A_st, sort_kernel::N_PER_BLOCK * merge_lane);
+    int partition = sort_kernel::merge_partition(
+        dev_vals + A_st,
+        dev_vals + B_st,
+        A_ed - A_st,
+        B_ed - B_st,
+        partition_at);
+
+    block_partitions[i] = A_st + partition;
+  }
+}
+
+template <
+    typename ValT,
+    typename IdxT,
+    bool ARG_SORT,
+    short BLOCK_THREADS,
+    short N_PER_THREAD,
+    typename CompareOp = LessThan<ValT>>
+[[kernel, max_total_threads_per_threadgroup(BLOCK_THREADS)]] void
+mb_block_merge(
+    const device IdxT* block_partitions [[buffer(0)]],
+    const device ValT* dev_vals_in [[buffer(1)]],
+    const device IdxT* dev_idxs_in [[buffer(2)]],
+    device ValT* dev_vals_out [[buffer(3)]],
+    device IdxT* dev_idxs_out [[buffer(4)]],
+    const constant int& size_sorted_axis [[buffer(5)]],
+    const constant int& merge_tiles [[buffer(6)]],
+    const constant int& num_tiles [[buffer(7)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  using sort_kernel = KernelMultiBlockMergeSort<
+      ValT,
+      IdxT,
+      ARG_SORT,
+      BLOCK_THREADS,
+      N_PER_THREAD,
+      CompareOp>;
+
+  using block_sort_t = typename sort_kernel::block_merge_sort_t;
+
+  block_partitions += tid.y * (num_tiles + 1);
+  dev_vals_in += tid.y * size_sorted_axis;
+  dev_idxs_in += tid.y * size_sorted_axis;
+  dev_vals_out += tid.y * size_sorted_axis;
+  dev_idxs_out += tid.y * size_sorted_axis;
+
+  int block_idx = tid.x;
+  int merge_group = block_idx / merge_tiles;
+  int sort_st = sort_kernel::N_PER_BLOCK * merge_tiles * merge_group;
+  int sort_sz = sort_kernel::N_PER_BLOCK * merge_tiles;
+  int sort_md = sort_kernel::N_PER_BLOCK * block_idx - sort_st;
+
+  int A_st = block_partitions[block_idx + 0];
+  int A_ed = block_partitions[block_idx + 1];
+  int B_st = min(size_sorted_axis, 2 * sort_st + sort_sz / 2 + sort_md - A_st);
+  int B_ed = min(
+      size_sorted_axis,
+      2 * sort_st + sort_sz / 2 + sort_md + sort_kernel::N_PER_BLOCK - A_ed);
+
+  if ((block_idx % merge_tiles) == merge_tiles - 1) {
+    A_ed = min(size_sorted_axis, sort_st + sort_sz / 2);
+    B_ed = min(size_sorted_axis, sort_st + sort_sz);
+  }
+
+  int A_sz = A_ed - A_st;
+  int B_sz = B_ed - B_st;
+
+  // Load from global memory
+  thread ValT thread_vals[N_PER_THREAD];
+  thread IdxT thread_idxs[N_PER_THREAD];
+  for (int i = 0; i < N_PER_THREAD; i++) {
+    int idx = BLOCK_THREADS * i + lid.x;
+    if (idx < (A_sz + B_sz)) {
+      thread_vals[i] = (idx < A_sz) ? dev_vals_in[A_st + idx]
+                                    : dev_vals_in[B_st + idx - A_sz];
+      thread_idxs[i] = (idx < A_sz) ? dev_idxs_in[A_st + idx]
+                                    : dev_idxs_in[B_st + idx - A_sz];
+    } else {
+      thread_vals[i] = CompareOp::init;
+      thread_idxs[i] = 0;
+    }
+  }
+
+  // Write to shared memory
+  threadgroup ValT tgp_vals[sort_kernel::N_PER_BLOCK];
+  threadgroup IdxT tgp_idxs[sort_kernel::N_PER_BLOCK];
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  for (int i = 0; i < N_PER_THREAD; i++) {
+    int idx = BLOCK_THREADS * i + lid.x;
+    tgp_vals[idx] = thread_vals[i];
+    tgp_idxs[idx] = thread_idxs[i];
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Merge
+  int sort_md_local = min(A_sz + B_sz, N_PER_THREAD * int(lid.x));
+
+  int A_st_local = block_sort_t::merge_partition(
+      tgp_vals, tgp_vals + A_sz, A_sz, B_sz, sort_md_local);
+  int A_ed_local = A_sz;
+
+  int B_st_local = sort_md_local - A_st_local;
+  int B_ed_local = B_sz;
+
+  int A_sz_local = A_ed_local - A_st_local;
+  int B_sz_local = B_ed_local - B_st_local;
+
+  // Do merge
+  block_sort_t::merge_step(
+      tgp_vals + A_st_local,
+      tgp_vals + A_ed_local + B_st_local,
+      tgp_idxs + A_st_local,
+      tgp_idxs + A_ed_local + B_st_local,
+      A_sz_local,
+      B_sz_local,
+      thread_vals,
+      thread_idxs);
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  for (int i = 0; i < N_PER_THREAD; ++i) {
+    int idx = lid.x * N_PER_THREAD;
+    tgp_vals[idx + i] = thread_vals[i];
+    tgp_idxs[idx + i] = thread_idxs[i];
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  // Write output
+  int base_idx = tid.x * sort_kernel::N_PER_BLOCK;
+  for (int i = lid.x; i < sort_kernel::N_PER_BLOCK; i += BLOCK_THREADS) {
+    int idx = base_idx + i;
+    if (idx < size_sorted_axis) {
+      dev_vals_out[idx] = tgp_vals[i];
+      dev_idxs_out[idx] = tgp_idxs[i];
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/sort.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/sort.metal
new file mode 100644
index 0000000..7f198d5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/sort.metal
@@ -0,0 +1,80 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <metal_stdlib>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/sort.h"
+
+#define instantiate_block_sort(                                          \
+    name, itname, itype, otname, otype, arg_sort, bn, tn)                \
+  instantiate_kernel("c" #name "_" #itname "_" #otname "_bn" #bn "_tn" #tn, \
+                     block_sort, itype, otype, arg_sort, bn, tn) \
+  instantiate_kernel("nc" #name "_" #itname "_" #otname "_bn" #bn "_tn" #tn, \
+                     block_sort_nc, itype, otype, arg_sort, bn, tn)
+
+#define instantiate_arg_block_sort_base(itname, itype, bn, tn) \
+  instantiate_block_sort(                                      \
+      arg_block_sort, itname, itype, uint32, uint32_t, true, bn, tn)
+
+#define instantiate_block_sort_base(itname, itype, bn, tn) \
+  instantiate_block_sort(                                  \
+      _block_sort, itname, itype, itname, itype, false, bn, tn)
+
+#define instantiate_block_sort_tn(itname, itype, bn) \
+  instantiate_block_sort_base(itname, itype, bn, 4)  \
+  instantiate_arg_block_sort_base(itname, itype, bn, 4)
+
+#define instantiate_block_sort_bn(itname, itype) \
+  instantiate_block_sort_tn(itname, itype, 32)  \
+  instantiate_block_sort_tn(itname, itype, 64)  \
+  instantiate_block_sort_tn(itname, itype, 128)  \
+  instantiate_block_sort_tn(itname, itype, 256)  \
+  instantiate_block_sort_tn(itname, itype, 512)
+
+instantiate_block_sort_bn(uint8, uint8_t)
+instantiate_block_sort_bn(uint16, uint16_t)
+instantiate_block_sort_bn(uint32, uint32_t)
+instantiate_block_sort_bn(int8, int8_t)
+instantiate_block_sort_bn(int16, int16_t)
+instantiate_block_sort_bn(int32, int32_t)
+instantiate_block_sort_bn(float16, half)
+instantiate_block_sort_bn(float32, float)
+instantiate_block_sort_bn(bfloat16, bfloat16_t)
+
+#define instantiate_block_sort_long(itname, itype) \
+  instantiate_block_sort_tn(itname, itype, 32)     \
+  instantiate_block_sort_tn(itname, itype, 64)     \
+  instantiate_block_sort_tn(itname, itype, 128)    \
+  instantiate_block_sort_tn(itname, itype, 256)
+
+instantiate_block_sort_long(uint64, uint64_t)
+instantiate_block_sort_long(int64, int64_t)
+
+#define instantiate_multi_block_sort(                                      \
+    vtname, vtype, itname, itype, arg_sort, bn, tn)                        \
+  instantiate_kernel("sort_mbsort_" #vtname "_" #itname "_bn" #bn "_tn" #tn, \
+                     mb_block_sort, vtype, itype, arg_sort, bn, tn) \
+  instantiate_kernel("partition_mbsort_" #vtname "_" #itname "_bn" #bn "_tn" #tn, \
+                     mb_block_partition, vtype, itype, arg_sort, bn, tn) \
+  instantiate_kernel("merge_mbsort_" #vtname "_" #itname "_bn" #bn "_tn" #tn, \
+                     mb_block_merge, vtype, itype, arg_sort, bn, tn)
+
+#define instantiate_multi_block_sort_base(vtname, vtype) \
+  instantiate_multi_block_sort(vtname, vtype, uint32, uint32_t, true, 512, 4)
+
+instantiate_multi_block_sort_base(uint8, uint8_t)
+instantiate_multi_block_sort_base(uint16, uint16_t)
+instantiate_multi_block_sort_base(uint32, uint32_t)
+instantiate_multi_block_sort_base(int8, int8_t)
+instantiate_multi_block_sort_base(int16, int16_t)
+instantiate_multi_block_sort_base(int32, int32_t)
+instantiate_multi_block_sort_base(float16, half)
+instantiate_multi_block_sort_base(float32, float)
+instantiate_multi_block_sort_base(bfloat16, bfloat16_t)
+
+#define instantiate_multi_block_sort_long(vtname, vtype) \
+  instantiate_multi_block_sort(vtname, vtype, uint32, uint32_t, true, 256, 4)
+
+instantiate_multi_block_sort_long(uint64, uint64_t)
+instantiate_multi_block_sort_long(int64, int64_t) // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/attn.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/attn.h
new file mode 100644
index 0000000..991d4d6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/attn.h
@@ -0,0 +1,296 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/attn/loader.h"
+#include "mlx/backend/metal/kernels/steel/attn/mma.h"
+#include "mlx/backend/metal/kernels/steel/attn/params.h"
+#include "mlx/backend/metal/kernels/steel/attn/transforms.h"
+#include "mlx/backend/metal/kernels/steel/gemm/params.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernel class
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <bool M_aligned, bool N_aligned, bool K_aligned>
+struct LoopAlignment {};
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    bool MN_aligned,
+    bool K_aligned,
+    typename AccumType = typename AccumHelper<T>::accum_type,
+    typename Epilogue = TransformNone<U, AccumType>>
+struct GEMMKernel {
+  STEEL_CONST short tgp_padding_a = 16 / sizeof(T);
+  STEEL_CONST short tgp_padding_b = 16 / sizeof(T);
+  STEEL_CONST short tgp_mem_size_a =
+      transpose_a ? BK * (BM + tgp_padding_a) : BM * (BK + tgp_padding_a);
+  STEEL_CONST short tgp_mem_size_b =
+      transpose_b ? BN * (BK + tgp_padding_b) : BK * (BN + tgp_padding_b);
+  STEEL_CONST short tgp_mem_size = tgp_mem_size_a + tgp_mem_size_b;
+
+  STEEL_CONST short tgp_size = WM * WN * 32;
+
+  using loader_a_t = BlockLoader<
+      T,
+      transpose_a ? BK : BM,
+      transpose_a ? BM : BK,
+      transpose_a ? BM + tgp_padding_a : BK + tgp_padding_a,
+      !transpose_a,
+      tgp_size>;
+  using loader_b_t = BlockLoader<
+      T,
+      transpose_b ? BN : BK,
+      transpose_b ? BK : BN,
+      transpose_b ? BK + tgp_padding_b : BN + tgp_padding_b,
+      transpose_b,
+      tgp_size>;
+  using mma_t = BlockMMA<
+      T,
+      U,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      transpose_a ? BM + tgp_padding_a : BK + tgp_padding_a,
+      transpose_b ? BK + tgp_padding_b : BN + tgp_padding_b,
+      AccumType,
+      Epilogue>;
+
+  /* Main kernel function */
+  template <bool M_aligned, bool N_aligned, bool K_aligned_>
+  static METAL_FUNC void gemm_loop(
+      threadgroup T* As [[threadgroup(0)]],
+      threadgroup T* Bs [[threadgroup(1)]],
+      const int gemm_k_iterations,
+      thread loader_a_t& loader_a,
+      thread loader_b_t& loader_b,
+      thread mma_t& mma_op,
+      thread const short& tgp_bm,
+      thread const short& tgp_bn,
+      thread const short& lbk,
+      LoopAlignment<M_aligned, N_aligned, K_aligned_> l = {}) {
+    // Appease the compiler
+    (void)l;
+
+    short2 tile_dims_A = transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm);
+
+    short2 tile_dims_B = transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK);
+
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      // Load elements into threadgroup
+      if (M_aligned) {
+        loader_a.load_unsafe();
+      } else {
+        loader_a.load_safe(tile_dims_A);
+      }
+
+      if (N_aligned) {
+        loader_b.load_unsafe();
+      } else {
+        loader_b.load_safe(tile_dims_B);
+      }
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+
+    if (!K_aligned_) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      short2 tile_dims_A_last =
+          transpose_a ? short2(tgp_bm, lbk) : short2(lbk, tgp_bm);
+      short2 tile_dims_B_last =
+          transpose_b ? short2(lbk, tgp_bn) : short2(tgp_bn, lbk);
+
+      loader_a.load_safe(tile_dims_A_last);
+      loader_b.load_safe(tile_dims_B_last);
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      mma_op.mma(As, Bs);
+    }
+  }
+
+  /* Main kernel function */
+  static METAL_FUNC void run(
+      const device T* A [[buffer(0)]],
+      const device T* B [[buffer(1)]],
+      device U* D [[buffer(2)]],
+      const constant GEMMParams* params [[buffer(3)]],
+      threadgroup T* As [[threadgroup(0)]],
+      threadgroup T* Bs [[threadgroup(1)]],
+      uint simd_lane_id [[thread_index_in_simdgroup]],
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]]) {
+    // Pacifying compiler
+    (void)lid;
+
+    const int tid_y = ((tid.y) << params->swizzle_log) +
+        ((tid.x) & ((1 << params->swizzle_log) - 1));
+    const int tid_x = (tid.x) >> params->swizzle_log;
+
+    if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+      return;
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Find block in A, B, C
+    const int c_row = tid_y * BM;
+    const int c_col = tid_x * BN;
+    const size_t c_row_long = size_t(c_row);
+    const size_t c_col_long = size_t(c_col);
+
+    A += transpose_a ? c_row_long : c_row_long * params->lda;
+    B += transpose_b ? c_col_long * params->ldb : c_col_long;
+    D += c_row_long * params->ldd + c_col_long;
+
+    // Prepare threadgroup loading operations
+    thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+    thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+    // Prepare threadgroup mma operation
+    thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+    int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+    ///////////////////////////////////////////////////////////////////////////////
+    // MNK aligned loop
+    if (MN_aligned) {
+      for (int k = 0; k < gemm_k_iterations; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+
+      threadgroup_barrier(mem_flags::mem_none);
+
+      // Loop tail
+      if (!K_aligned) {
+        int lbk = params->K - params->gemm_k_iterations_aligned * BK;
+        short2 tile_dims_A = transpose_a ? short2(BM, lbk) : short2(lbk, BM);
+        short2 tile_dims_B = transpose_b ? short2(lbk, BN) : short2(BN, lbk);
+
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        mma_op.mma(As, Bs);
+      }
+
+      // Store results to device memory
+      mma_op.store_result(D, params->ldd);
+      return;
+
+    }
+    ///////////////////////////////////////////////////////////////////////////////
+    // MN unaligned loop
+    else { // Loop over K - unaligned case
+      short tgp_bm = min(BM, params->M - c_row);
+      short tgp_bn = min(BN, params->N - c_col);
+      short leftover_bk = params->K - params->gemm_k_iterations_aligned * BK;
+
+      if (tgp_bm == BM && tgp_bn == BN) {
+        gemm_loop<true, true, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result(D, params->ldd);
+        return;
+
+      } else if (tgp_bn == BN) {
+        gemm_loop<false, true, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+
+      } else if (tgp_bm == BM) {
+        gemm_loop<true, false, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+
+      } else {
+        gemm_loop<false, false, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+      }
+    }
+  }
+};
+
+} // namespace steel
+} // namespace mlx
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h
new file mode 100644
index 0000000..4de11b0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h
@@ -0,0 +1,476 @@
+// Copyright © 2024-25 Apple Inc.
+
+#include "mlx/backend/metal/kernels/steel/attn/attn.h"
+
+using namespace mlx::steel;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernels
+///////////////////////////////////////////////////////////////////////////////
+
+constant bool align_Q [[function_constant(200)]];
+constant bool align_K [[function_constant(201)]];
+
+constant bool has_mask [[function_constant(300)]];
+constant bool do_causal [[function_constant(301)]];
+constant bool has_sinks [[function_constant(302)]];
+
+template <typename T>
+struct TransformScale {
+  T scale;
+  METAL_FUNC TransformScale(T scale_) : scale(scale_) {}
+
+  METAL_FUNC T apply(T x) const {
+    return scale * x;
+  }
+};
+
+struct MaxOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return metal::max(x, y);
+  }
+};
+
+struct SumOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x + y;
+  }
+};
+
+struct MulOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x * y;
+  }
+};
+
+struct SubOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x - y;
+  }
+};
+
+struct ExpSubOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return fast::exp2(x - y);
+  }
+};
+
+struct DivOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x / y;
+  }
+};
+
+// clang-format off
+template <
+    typename T,
+    int BQ,
+    int BK,
+    int BD,
+    int WM,
+    int WN,
+    typename MaskType = float,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention(
+    const device T* Q [[buffer(0)]],
+    const device T* K [[buffer(1)]],
+    const device T* V [[buffer(2)]],
+    device T* O [[buffer(3)]],
+    const constant AttnParams* params [[buffer(4)]],
+    const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]],
+    const device MaskType* mask [[buffer(6), function_constant(has_mask)]],
+    const device T* sinks [[buffer(7), function_constant(has_sinks)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) { // clang-format on
+
+  // Pacifying compiler
+  (void)lid;
+
+  // Move to correct block
+  ulong3 tidl{tid.x, tid.y, tid.z};
+
+  Q += tidl.z * params->Q_strides[0] + // Batch
+      tidl.y * params->Q_strides[1] + // Head
+      tidl.x * BQ * params->Q_strides[2]; // Sequence
+
+  ulong kv_head_idx = int(tid.y) / params->gqa_factor;
+  K += tidl.z * params->K_strides[0] + // Batch
+      kv_head_idx * params->K_strides[1]; // Head
+
+  V += tidl.z * params->V_strides[0] + // Batch
+      kv_head_idx * params->V_strides[1]; // Head
+
+  O += tidl.z * params->O_strides[0] + // Batch
+      tidl.y * params->O_strides[1] + // Head
+      tidl.x * BQ * params->O_strides[2]; // Sequence
+
+  if (has_mask) {
+    mask += tidl.z * mask_params->M_strides[0] + // Batch
+        tidl.y * mask_params->M_strides[1]; // Head
+  }
+
+  // Prepare threadgroup memory
+  constexpr short padQ = 16 / sizeof(T);
+  constexpr short padK = 16 / sizeof(T);
+  constexpr short padV = 16 / sizeof(T);
+
+  constexpr short LDQ_tgp = BD + padQ;
+  constexpr short LDK_tgp = BK + padK;
+  constexpr short LDV_tgp = BD + padV;
+
+  constexpr short tgp_mem_0 = (BK + padK) * (BD);
+  constexpr short tgp_mem_1 = BK * (BD + padV);
+  constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1;
+
+  threadgroup T Q_smem[BQ * (BD + padQ)];
+  threadgroup T KV_smem[tgp_mem_s];
+
+  threadgroup T* Qs = Q_smem;
+  threadgroup T* Ks = KV_smem;
+  threadgroup T* Vs = KV_smem;
+
+  // Prepare block loaders
+  using QBlockLoader = BlockLoaderT<
+      /* typename T = */ T,
+      /* short BROWS = */ BQ,
+      /* short BCOLS = */ BD,
+      /* short kDstStrRow = */ LDQ_tgp,
+      /* short kDstStrCol = */ 1,
+      /* short reduction_dim = */ 1,
+      /* short tgp_size = */ WM * WN * 32>;
+
+  // K is loaded in transposed
+  using KBlockLoader = BlockLoaderT<
+      /* typename T = */ T,
+      /* short BROWS = */ BK,
+      /* short BCOLS = */ BD,
+      /* short kDstStrRow = */ 1,
+      /* short kDstStrCol = */ LDK_tgp,
+      /* short reduction_dim = */ 0,
+      /* short tgp_size = */ WM * WN * 32>;
+
+  using VBlockLoader = BlockLoaderT<
+      /* typename T = */ T,
+      /* short BROWS = */ BK,
+      /* short BCOLS = */ BD,
+      /* short kDstStrRow = */ LDV_tgp,
+      /* short kDstStrCol = */ 1,
+      /* short reduction_dim = */ 0,
+      /* short tgp_size = */ WM * WN * 32>;
+
+  QBlockLoader loader_q(
+      Q, params->Q_strides[2], Qs, simd_group_id, simd_lane_id);
+  KBlockLoader loader_k(
+      K, params->K_strides[2], Ks, simd_group_id, simd_lane_id);
+  VBlockLoader loader_v(
+      V, params->V_strides[2], Vs, simd_group_id, simd_lane_id);
+
+  TransformScale<T> ts(static_cast<T>(params->scale * M_LOG2E_F));
+
+  // Prepare MMA tiles
+  constexpr short kFragSize = 8; // MMAFrag size
+  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
+
+  constexpr int kNWarps = WM * WN;
+  static_assert(
+      BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0,
+      "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence.");
+
+  // Q seq frags per warp
+  constexpr int TQ = BQ / (kNWarps * kFragSize);
+  // KV sequence frags (all warps load the same frags)
+  constexpr int TK = BK / kFragSize;
+  // HeadDim frags (all warps load the same frags)
+  constexpr int TD = BD / kFragSize;
+
+  static_assert(TQ == 1, "Check TQ");
+
+  MMATile<AccumType, TQ, 1, MMAFrag_acc_t> Qtile;
+  MMATile<AccumType, 1, TK, MMAFrag_acc_t> Ktile;
+  MMATile<AccumType, TQ, TK, MMAFrag_acc_t> Stile;
+  MMATile<AccumType, 1, 1, MMAFrag_acc_t> Vtile;
+  MMATile<AccumType, TQ, TD, MMAFrag_acc_t> Otile;
+
+  Otile.clear();
+
+  // Prepare mma tile offsets
+  const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
+  const short sm = simd_coord.y;
+  const short sn = simd_coord.x;
+  const short tm = kFragSize * TQ * simd_group_id;
+
+  const short Qs_offset = (tm + sm) * LDQ_tgp + sn;
+  const short Ks_offset = sm * LDK_tgp + sn;
+  const short Vs_offset = sm * LDV_tgp + sn;
+
+  constexpr short Qs_tile_stride = kFragSize;
+  constexpr short Ks_tile_stride = kFragSize * LDK_tgp;
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Load Q blocks apply scale
+  if (!align_Q && int(tid.x) == (params->NQ_aligned)) {
+    loader_q.load_safe(short2(BD, params->qL_rem));
+  } else {
+    loader_q.load_unsafe();
+  }
+  loader_q.apply_inplace_op(ts);
+
+  // Init row reduction variables
+  constexpr short kRowsPT = decltype(Stile)::kRowsPerThread;
+
+  AccumType max_score[kRowsPT];
+  AccumType sum_score[kRowsPT] = {0};
+
+  // Init to -Inf
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < kRowsPT; ++i) {
+    max_score[i] = Limits<AccumType>::finite_min;
+  }
+
+  if (has_sinks) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      max_score[i] = M_LOG2E_F * static_cast<AccumType>(sinks[tidl.y]);
+      sum_score[i] = 1;
+    }
+  }
+
+  int kb_lim = params->NK;
+
+  if (do_causal) {
+    int q_max = (tid.x + 1) * BQ + params->qL_off;
+    kb_lim = (q_max + BK - 1) / BK;
+    kb_lim = min(params->NK, kb_lim);
+  }
+
+  // Loop over KV seq length
+  for (int kb = 0; kb < kb_lim; kb++) {
+    // Load K block and apply scale
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    if (!align_K && kb == (params->NK_aligned)) {
+      loader_k.load_safe(short2(BD, params->kL_rem));
+    } else {
+      loader_k.load_unsafe();
+    }
+
+    // Do S = Q @ K.T
+    Stile.clear();
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_UNROLL
+    for (short dd = 0; dd < TD; dd++) {
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Qtile.template load<T, 1, 1, LDQ_tgp, 1>(
+          &Qs[Qs_offset + dd * Qs_tile_stride]);
+      Ktile.template load<T, 1, 1, LDK_tgp, 1>(
+          &Ks[Ks_offset + dd * Ks_tile_stride]);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      tile_matmad(Stile, Qtile, Ktile, Stile);
+    }
+
+    // Mask out length sequence
+    if (!align_K && kb == (params->NK_aligned)) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      constexpr auto neg_inf = Limits<selem_t>::finite_min;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          short col_pos = sn + (j * stile_t::kFragCols);
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemCols; jj++) {
+            if ((col_pos + jj) >= params->kL_rem) {
+              Stile.frag_at(i, j)[jj] = neg_inf;
+            }
+          }
+        }
+      }
+    }
+
+    // Mask out if causal
+    if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      constexpr auto neg_inf = Limits<selem_t>::finite_min;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        const int row_pos =
+            tid.x * BQ + params->qL_off + tm + sm + (i * stile_t::kFragRows);
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          const int col_pos = kb * BK + sn + (j * stile_t::kFragCols);
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemCols; jj++) {
+            if (row_pos < (col_pos + jj)) {
+              Stile.frag_at(i, j)[jj] = neg_inf;
+            }
+          }
+        }
+      }
+    }
+
+    // Other masking as needed
+    if (has_mask) {
+      using stile_t = decltype(Stile);
+      using selem_t = typename stile_t::elem_type;
+      constexpr auto neg_inf = Limits<selem_t>::finite_min;
+
+      constexpr bool is_bool = is_same_v<MaskType, bool>;
+      using melem_t = typename metal::conditional_t<is_bool, bool, selem_t>;
+
+      using MMAFrag_mask_t = BaseMMAFrag<melem_t, kFragSize, kFragSize>;
+      using frag_t = typename MMAFrag_mask_t::frag_type;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < stile_t::kTileRows; i++) {
+        const int row_pos = tid.x * BQ + tm + sm + (i * stile_t::kFragRows);
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < stile_t::kTileCols; j++) {
+          const int col_pos = kb * BK + sn + (j * stile_t::kFragCols);
+
+          frag_t mfrag;
+
+          MMAFrag_mask_t::load_safe(
+              mfrag,
+              mask,
+              int64_t(mask_params->M_strides[2]),
+              Int<1>{},
+              params->qL,
+              params->kL,
+              row_pos,
+              col_pos);
+
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < stile_t::MMAFrag_t::kElemsPerFrag; jj++) {
+            if constexpr (is_bool) {
+              Stile.frag_at(i, j)[jj] =
+                  mfrag[jj] ? Stile.frag_at(i, j)[jj] : neg_inf;
+            } else {
+              Stile.frag_at(i, j)[jj] += M_LOG2E_F * selem_t(mfrag[jj]);
+            }
+          }
+        }
+      }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Load V blocks
+    if (!align_K && kb == (params->NK_aligned)) {
+      loader_v.load_safe(short2(BD, params->kL_rem));
+    } else {
+      loader_v.load_unsafe();
+    }
+
+    // Do softmax
+
+    // Temp variables
+    AccumType new_max[kRowsPT];
+    AccumType factor[kRowsPT];
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      new_max[i] = max_score[i];
+    }
+
+    // Row max
+    Stile.template row_reduce<MaxOp>(new_max);
+
+    // exp(Si - rowmax(Si))
+    Stile.template row_bin_op<ExpSubOp>(new_max);
+
+    // Factor exp(rowmax(Si) - rowmax(Si-1))
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      factor[i] = fast::exp2(max_score[i] - new_max[i]);
+    }
+
+    // Save max for next iteration
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      max_score[i] = new_max[i];
+    }
+
+    // Row Sum
+    AccumType sum_score_tmp[kRowsPT] = {0};
+    Stile.template row_reduce<SumOp>(sum_score_tmp);
+
+    // Update norm
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      sum_score[i] = sum_score[i] * factor[i] + sum_score_tmp[i];
+    }
+
+    // Update O
+    Otile.template row_bin_op<MulOp>(factor);
+
+    // Load V into registers
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_UNROLL
+    for (short iq = 0; iq < TQ; iq++) {
+      STEEL_PRAGMA_UNROLL
+      for (short id = 0; id < TD; id++) {
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
+
+          const short kk = ik * kFragSize;
+          const short dd = id * kFragSize;
+
+          Vtile.template load<T, 1, 1, LDV_tgp, 1>(
+              &Vs[Vs_offset + kk * LDV_tgp + dd]);
+
+          if constexpr (BD == 128) {
+            simdgroup_barrier(mem_flags::mem_none);
+          }
+
+          MMAFrag_acc_t::mma(
+              Otile.frag_at(iq, id),
+              Stile.frag_at(iq, ik),
+              Vtile.frag_at(0, 0),
+              Otile.frag_at(iq, id));
+        }
+      }
+    }
+
+    // Prepare for next iteration
+    loader_k.next();
+    loader_v.next();
+  }
+
+  // Normalize output
+  Otile.template row_bin_op<DivOp>(sum_score);
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Store results
+  O += (tm + sm) * params->O_strides[2] + sn;
+
+  if (!align_Q && int(tid.x) == (params->NQ_aligned)) {
+    auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm));
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    Otile.template store_safe<T, 1, 1>(O, params->O_strides[2], dst_tile_dims);
+  } else {
+    Otile.template store<T, 1, 1>(O, params->O_strides[2]);
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal
new file mode 100644
index 0000000..0ff9d91
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal
@@ -0,0 +1,27 @@
+// Copyright © 2024-25 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h"
+
+#define instantiate_attn(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \
+  instantiate_kernel(                                                    \
+      "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd            \
+      "_wm" #wm "_wn" #wn "_mask" #mname,                                \
+  attention, dtype, bq, bk, bd, wm, wn, mtype, float)
+
+#define instantiate_attn_shapes_helper(iname, itype, mname, mtype)  \
+    instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 32,  80, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 32,  64, 4, 1, mname, mtype)
+
+#define instantiate_attn_mask_helper(iname, itype) \
+    instantiate_attn_shapes_helper(iname, itype, iname, itype) \
+    instantiate_attn_shapes_helper(iname, itype, bool_, bool)
+
+instantiate_attn_mask_helper(float16, half);
+instantiate_attn_mask_helper(bfloat16, bfloat16_t);
+
+instantiate_attn_mask_helper(float32, float);
+// clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h
new file mode 100644
index 0000000..6abee21
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h
@@ -0,0 +1,481 @@
+// Copyright © 2024-25 Apple Inc.
+
+#include "mlx/backend/metal/kernels/steel/attn/nax.h"
+#include "mlx/backend/metal/kernels/steel/attn/params.h"
+#include "mlx/backend/metal/kernels/steel/attn/transforms.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace mlx::steel;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernels
+///////////////////////////////////////////////////////////////////////////////
+
+constant bool align_Q [[function_constant(200)]];
+constant bool align_K [[function_constant(201)]];
+
+constant bool has_mask [[function_constant(300)]];
+constant bool do_causal [[function_constant(301)]];
+constant bool has_sinks [[function_constant(302)]];
+
+template <typename T>
+struct TransformScale {
+  T scale;
+  METAL_FUNC TransformScale(T scale_) : scale(scale_) {}
+
+  METAL_FUNC T apply(T x) const {
+    return scale * x;
+  }
+};
+
+struct MaxOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return metal::max(x, y);
+  }
+};
+
+struct SumOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x + y;
+  }
+};
+
+struct MulOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x * y;
+  }
+};
+
+struct SubOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x - y;
+  }
+};
+
+struct ExpSubOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return fast::exp2(x - y);
+  }
+};
+
+struct DivOp {
+  template <typename T>
+  METAL_FUNC static constexpr T apply(T x, T y) {
+    return x / y;
+  }
+};
+
+// clang-format off
+template <
+    typename T,
+    int BQ,
+    int BK,
+    int BD,
+    int WM,
+    int WN,
+    typename MaskType = float,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention_nax(
+    const device T* Q [[buffer(0)]],
+    const device T* K [[buffer(1)]],
+    const device T* V [[buffer(2)]],
+    device T* O [[buffer(3)]],
+    const constant AttnParams* params [[buffer(4)]],
+    const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]],
+    const device MaskType* mask [[buffer(6), function_constant(has_mask)]],
+    const device T* sinks [[buffer(7), function_constant(has_sinks)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) { // clang-format on
+
+  // Pacifying compiler
+  (void)lid;
+  (void)simd_lane_id;
+
+  // Move to correct block
+  ulong3 tidl{tid.x, tid.y, tid.z};
+
+  Q += tidl.z * params->Q_strides[0] + // Batch
+      tidl.y * params->Q_strides[1] + // Head
+      tidl.x * BQ * params->Q_strides[2]; // Sequence
+
+  ulong kv_head_idx = int(tid.y) / params->gqa_factor;
+  K += tidl.z * params->K_strides[0] + // Batch
+      kv_head_idx * params->K_strides[1]; // Head
+
+  V += tidl.z * params->V_strides[0] + // Batch
+      kv_head_idx * params->V_strides[1]; // Head
+
+  O += tidl.z * params->O_strides[0] + // Batch
+      tidl.y * params->O_strides[1] + // Head
+      tidl.x * BQ * params->O_strides[2]; // Sequence
+
+  if (has_mask) {
+    mask += tidl.z * mask_params->M_strides[0] + // Batch
+        tidl.y * mask_params->M_strides[1]; // Head
+  }
+
+  const metal::uniform<float> scale2 =
+      make_uniform(params->scale) * make_uniform(1.44269504089f);
+
+  // Prepare MMA tiles
+  constexpr short UQ = 16;
+  constexpr short UD = 32;
+
+  constexpr int kNWarps = WM * WN;
+  static_assert(
+      BQ >= (kNWarps * UQ) && BQ % (kNWarps * UQ) == 0,
+      "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence.");
+
+  // Q seq frags per warp
+  constexpr int TQ = BQ / (kNWarps * UQ);
+  // HeadDim frags (all warps load the same frags)
+  constexpr int TD = BD / UD;
+
+  static_assert(TQ == 1, "Check TQ");
+
+  using OSubTile = NAXSubTile<AccumType, UQ, UD>;
+  NAXTile<AccumType, TQ, TD, OSubTile> Otile;
+
+  Otile.clear();
+
+  // Prepare mma tile offsets
+  const short2 simd_coord = OSubTile::NAXFrag_t::get_coord();
+  const short sm = simd_coord.y;
+  const short sn = simd_coord.x;
+  const short tm = UQ * TQ * simd_group_id;
+
+  Q += (tm + sm) * int(params->Q_strides[2]) + sn;
+  K += sm * int(params->K_strides[2]) + sn;
+  V += sm * int(params->V_strides[2]) + sn;
+
+  // Init row reduction variables
+  constexpr short kRowsPT = decltype(Otile)::kRowsPerThread;
+
+  metal::vec<AccumType, kRowsPT> max_score;
+  metal::vec<AccumType, kRowsPT> sum_score{0};
+
+  // Init to -Inf
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < kRowsPT; ++i) {
+    max_score[i] = Limits<AccumType>::finite_min;
+  }
+
+  if (has_sinks) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      max_score[i] = M_LOG2E_F * static_cast<AccumType>(sinks[tidl.y]);
+      sum_score[i] = 1;
+    }
+  }
+
+  int kb_lim = params->NK;
+
+  if (do_causal) {
+    int q_max = (tid.x + 1) * BQ + params->qL_off;
+    kb_lim = (q_max + BK - 1) / BK;
+    kb_lim = min(params->NK, kb_lim);
+  }
+
+  const bool is_last_bq = int(tid.x) == (params->NQ_aligned);
+  // const bool is_last_tq = int(simd_group_id) >= (params->qL_rem / UQ);
+  const bool is_last_q = is_last_bq;
+
+  const short lim_rows_q = params->qL_rem - (tm + sm);
+  const short lim_rows_k = params->kL_rem - sm;
+
+  // Loop over KV seq length
+  for (int kb = 0; kb < kb_lim; kb++) {
+    const int is_last_k = (kb == (params->NK_aligned));
+
+    // Do S = Q @ K.T
+    constexpr short UDs = 16;
+    constexpr short UKs = 32;
+
+    constexpr short TDs = BD / UDs;
+    constexpr short TKs = BK / UKs;
+
+    using SSubTile = NAXSubTile<AccumType, UQ, UKs>;
+    using QSubTile = NAXSubTile<T, UQ, UDs>;
+    using KSubTile = NAXSubTile<T, UKs, UDs>;
+
+    NAXTile<AccumType, TQ, TKs, SSubTile> Stile;
+
+    Stile.clear();
+
+    STEEL_PRAGMA_UNROLL
+    for (short iq = 0; iq < TQ; iq++) {
+      STEEL_PRAGMA_UNROLL
+      for (short ik = 0; ik < TKs; ik++) {
+        STEEL_PRAGMA_UNROLL
+        for (short id = 0; id < TDs; id++) {
+          NAXTile<T, 1, 1, QSubTile> Qtile;
+          NAXTile<T, 1, 1, KSubTile> Ktile;
+
+          const int Q_load_off = iq * UQ * int(params->Q_strides[2]) + id * UDs;
+          const int K_load_off =
+              ik * UKs * int(params->K_strides[2]) + id * UDs;
+
+          if (!align_Q && is_last_q) {
+            // Qtile.load_rows(
+            //     Q + Q_load_off,
+            //     int(params->Q_strides[2]),
+            //     lim_rows_q - iq * UQ);
+            Qtile.load_safe(
+                Q + Q_load_off,
+                int(params->Q_strides[2]),
+                short2(BD, lim_rows_q - iq * UQ));
+          } else {
+            Qtile.load(Q + Q_load_off, int(params->Q_strides[2]));
+          }
+
+          if (!align_K && is_last_k) {
+            // Ktile.load_rows(
+            //     K + K_load_off,
+            //     int(params->K_strides[2]),
+            //     lim_rows_k - ik * UKs);
+            Ktile.load_safe(
+                K + K_load_off,
+                int(params->K_strides[2]),
+                short2(BD, lim_rows_k - ik * UKs));
+          } else {
+            Ktile.load(K + K_load_off, int(params->K_strides[2]));
+          }
+
+          subtile_matmad_nax(
+              Stile.subtile_at(iq, ik),
+              Qtile.subtile_at(0, 0),
+              metal::false_type{},
+              Ktile.subtile_at(0, 0),
+              metal::true_type{});
+        }
+      }
+    }
+
+    // Scale S
+    STEEL_PRAGMA_UNROLL
+    for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) {
+      Stile.elems()[ii] *= float(scale2);
+    }
+
+    // Scale and Retile S
+    constexpr short UK = 16;
+    constexpr short TK = BK / UK;
+    using PSubTile = NAXSubTile<AccumType, UQ, UK>;
+
+    NAXTile<AccumType, TQ, TK, PSubTile> Ptile;
+
+    STEEL_PRAGMA_UNROLL
+    for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) {
+      Ptile.elems()[ii] = Stile.elems()[ii];
+    }
+
+    // Mask out length sequence
+    if (!align_K && is_last_k) {
+      constexpr auto neg_inf = Limits<AccumType>::finite_min;
+
+      STEEL_PRAGMA_UNROLL
+      for (short iq = 0; iq < TQ; iq++) {
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          const short col_pos = sn + ik * UK;
+
+          thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0);
+
+          STEEL_PRAGMA_UNROLL
+          for (short ii = 0; ii < PSubTile::kFragThrRows; ii++) {
+            STEEL_PRAGMA_UNROLL
+            for (short jj = 0; jj < PSubTile::kFragThrCols; jj++) {
+              const auto loc = ii * PSubTile::kFragThrCols + jj;
+              fg[loc] = ((col_pos + jj) >= params->kL_rem) ? neg_inf : fg[loc];
+            }
+          }
+        }
+      }
+    }
+
+    // Mask out if causal
+    if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) {
+      constexpr auto neg_inf = Limits<AccumType>::finite_min;
+
+      const int base_row = tid.x * BQ + params->qL_off + tm;
+      const int base_col = kb * BK;
+
+      STEEL_PRAGMA_UNROLL
+      for (short iq = 0; iq < TQ; iq++) {
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          const short row_pos = base_row + iq * UQ;
+          const short col_pos = base_col + ik * UK;
+
+          thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0);
+
+          STEEL_PRAGMA_UNROLL
+          for (short ii = 0; ii < PSubTile::kFragThrRows; ii++) {
+            STEEL_PRAGMA_UNROLL
+            for (short jj = 0; jj < PSubTile::kFragThrCols; jj++) {
+              const auto r = row_pos + ii * PSubTile::kFragRowsJump + sm;
+              const auto c = col_pos + jj + sn;
+              const auto loc = ii * PSubTile::kFragThrCols + jj;
+              fg[loc] = (r < c) ? neg_inf : fg[loc];
+            }
+          }
+        }
+      }
+    }
+
+    // Other masking as needed
+    if (has_mask) {
+      constexpr auto neg_inf = Limits<AccumType>::finite_min;
+
+      const int base_row = tid.x * BQ + tm;
+      const int base_col = kb * BK;
+
+      constexpr bool is_bool = is_same_v<MaskType, bool>;
+      using melem_t = typename metal::conditional_t<is_bool, bool, AccumType>;
+      using MSubTile = NAXSubTile<melem_t, UQ, UK>;
+
+      STEEL_PRAGMA_UNROLL
+      for (short iq = 0; iq < TQ; iq++) {
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          const short row_pos = base_row + iq * UQ + sm;
+          const short col_pos = base_col + ik * UK + sn;
+
+          MSubTile mfrag;
+          mfrag.load_safe(
+              mask,
+              int64_t(mask_params->M_strides[2]),
+              Int<1>{},
+              params->qL,
+              params->kL,
+              row_pos,
+              col_pos);
+
+          thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0);
+
+          STEEL_PRAGMA_UNROLL
+          for (short jj = 0; jj < MSubTile::kElemsPerFrag; jj++) {
+            if constexpr (is_bool) {
+              fg[jj] = mfrag.elems()[jj] ? fg[jj] : neg_inf;
+            } else {
+              fg[jj] += M_LOG2E_F * AccumType(mfrag.elems()[jj]);
+            }
+          }
+        }
+      }
+    }
+
+    // Do softmax
+
+    // Temp variables
+    metal::vec<AccumType, kRowsPT> new_max;
+    metal::vec<AccumType, kRowsPT> factor;
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      new_max[i] = max_score[i];
+    }
+
+    // Row max
+    Ptile.template row_reduce<MaxOp>(new_max);
+
+    // exp(Si - rowmax(Si))
+    Ptile.template row_bin_op<ExpSubOp>(new_max);
+
+    // Factor exp(rowmax(Si) - rowmax(Si-1))
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      factor[i] = fast::exp2(max_score[i] - new_max[i]);
+      max_score[i] = new_max[i];
+    }
+
+    // Row Sum
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      sum_score[i] = sum_score[i] * factor[i];
+    }
+
+    Ptile.template row_reduce<SumOp>(sum_score);
+
+    // Update O
+    Otile.template row_bin_op<MulOp>(factor);
+
+    simdgroup_barrier(mem_flags::mem_none);
+
+    // Do O = P @ V
+    STEEL_PRAGMA_UNROLL
+    for (short iq = 0; iq < TQ; iq++) {
+      STEEL_PRAGMA_UNROLL
+      for (short id = 0; id < TD; id++) {
+        if constexpr (BD == 128) {
+          if (id == 2) {
+            threadgroup_barrier(mem_flags::mem_none);
+          }
+        }
+
+        STEEL_PRAGMA_UNROLL
+        for (short ik = 0; ik < TK; ik++) {
+          using VSubTile = NAXSubTile<T, UK, UD>;
+          NAXTile<T, 1, 1, VSubTile> Vtile;
+
+          const int V_load_off = ik * UK * int(params->V_strides[2]) + id * UD;
+
+          if (!align_K && is_last_k) {
+            // Vtile.load_rows(
+            //     V + V_load_off,
+            //     int(params->V_strides[2]),
+            //     lim_rows_k - ik * UK);
+            Vtile.load_safe(
+                V + V_load_off,
+                int(params->V_strides[2]),
+                short2(BD, lim_rows_k - ik * UK));
+          } else {
+            Vtile.load(V + V_load_off, int(params->V_strides[2]));
+          }
+
+          subtile_matmad_nax(
+              Otile.subtile_at(iq, id),
+              Ptile.subtile_at(iq, ik),
+              metal::bool_constant<false>{},
+              Vtile.subtile_at(0, 0),
+              metal::bool_constant<false>{});
+        }
+      }
+    }
+
+    // Prepare for next iteration
+    K += BK * int(params->K_strides[2]);
+    V += BK * int(params->V_strides[2]);
+  }
+
+  // Normalize output
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  metal::vec<AccumType, kRowsPT> rcp;
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < kRowsPT; ++i) {
+    rcp[i] = 1.f / sum_score[i];
+  }
+
+  Otile.template row_bin_op<MulOp>(rcp);
+
+  // Store results
+  O += (tm + sm) * int(params->O_strides[2]) + sn;
+
+  if (!align_Q && is_last_q) {
+    if (lim_rows_q <= 0)
+      return;
+
+    // Otile.store_rows(O, params->O_strides[2], lim_rows_q);
+    Otile.store_safe(O, params->O_strides[2], short2(BD, lim_rows_q));
+  } else {
+    Otile.store(O, int(params->O_strides[2]));
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal
new file mode 100644
index 0000000..d40df50
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal
@@ -0,0 +1,28 @@
+// Copyright © 2024-25 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h"
+
+#define instantiate_attn(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \
+  instantiate_kernel(                                                    \
+      "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd            \
+      "_wm" #wm "_wn" #wn "_mask" #mname,                                \
+  attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float)
+
+#define instantiate_attn_shapes_helper(iname, itype, mname, mtype)  \
+    instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 64, 32,  64, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 64, 64,  64, 4, 1, mname, mtype)
+
+#define instantiate_attn_mask_helper(iname, itype) \
+    instantiate_attn_shapes_helper(iname, itype, iname, itype) \
+    instantiate_attn_shapes_helper(iname, itype, bool_, bool)
+
+instantiate_attn_mask_helper(float16, half);
+instantiate_attn_mask_helper(bfloat16, bfloat);
+
+instantiate_attn_mask_helper(float32, float);
+// clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/loader.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/loader.h
new file mode 100644
index 0000000..7ec7981
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/loader.h
@@ -0,0 +1,264 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short alignment = 1,
+    short n_reads = (BCOLS * BROWS) / (tgp_size),
+    short TCOLS = BCOLS / n_reads,
+    short TROWS = tgp_size / TCOLS>
+struct BlockLoader {
+  STEEL_CONST short n_rows = (BROWS + TROWS - 1) / TROWS;
+  STEEL_CONST short vec_size = n_reads;
+
+  // Leading dimension for src
+  const int src_ld;
+  const int tile_stride;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  struct alignas(alignment * sizeof(T)) ReadVector {
+    uint8_t v[sizeof(T) * vec_size];
+  };
+
+  /* Constructor */
+  METAL_FUNC BlockLoader(
+      const device T* src_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(reduction_dim ? BCOLS : BROWS * src_ld),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj) {}
+
+  /* Apply operation to threadgroup without bound checking */
+  template <typename UnaryOp>
+  METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = op.apply(dst[i * dst_ld + j]);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      *((threadgroup ReadVector*)(&dst[i * dst_ld])) =
+          *((const device ReadVector*)(&src[i * src_ld]));
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - with bound checking */
+  METAL_FUNC void load_safe(short2 src_tile_dim) const {
+    src_tile_dim = src_tile_dim - short2(bj, bi);
+
+    // Skip loading if thread has no valid reads
+    if (src_tile_dim.x <= 0 || src_tile_dim.y <= 0) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+      return;
+    }
+
+    // Use fast thread memory for bound checks
+    bool tmp_idx[vec_size];
+    T tmp_val[vec_size];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      // Make sure tmp_idx only contains valid indices
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_idx[j] = (i < src_tile_dim.y) && (j < src_tile_dim.x);
+      }
+
+      // Read valid indices into tmp_val
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = src[(tmp_idx[j] ? i * src_ld + j : 0)];
+      }
+
+      // Zero out unneeded values
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = tmp_idx[j] ? tmp_val[j] : T(0);
+      }
+
+      // Copy values to threadgroup memory
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = tmp_val[j];
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    src += tile_stride;
+  }
+};
+
+template <int R, int C>
+struct CShape {
+  STEEL_CONST int kRows = R;
+  STEEL_CONST int kCols = C;
+};
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short kDstStrRow,
+    short kDstStrCol,
+    short reduction_dim,
+    short tgp_size,
+    short n_reads = (BCOLS * BROWS) / (tgp_size),
+    short TCOLS = BCOLS / n_reads,
+    short TROWS = tgp_size / TCOLS>
+struct BlockLoaderT {
+  STEEL_CONST short n_rows = (BROWS + TROWS - 1) / TROWS;
+  STEEL_CONST short vec_size = n_reads;
+
+  // Leading dimension for src
+  const int src_ld;
+  const int tile_stride;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  /* Constructor */
+  METAL_FUNC BlockLoaderT(
+      const device T* src_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(reduction_dim ? BCOLS : BROWS * src_ld),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * kDstStrRow + bj * kDstStrCol),
+        src(src_ + bi * src_ld + bj) {}
+
+  /* Apply operation to threadgroup without bound checking */
+  template <typename UnaryOp>
+  METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * kDstStrRow + j * kDstStrCol] =
+            op.apply(dst[i * kDstStrRow + j * kDstStrCol]);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * kDstStrRow + j * kDstStrCol] = src[i * src_ld + j];
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - with bound checking */
+  METAL_FUNC void load_safe(short2 src_tile_dim) const {
+    src_tile_dim = src_tile_dim - short2(bj, bi);
+
+    // Skip loading if thread has no valid reads
+    if (src_tile_dim.x <= 0 || src_tile_dim.y <= 0) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * kDstStrRow + j * kDstStrCol] = T(0);
+        }
+      }
+      return;
+    }
+
+    // Use fast thread memory for bound checks
+    bool tmp_idx[vec_size];
+    T tmp_val[vec_size];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      // Make sure tmp_idx only contains valid indices
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_idx[j] = (i < src_tile_dim.y) && (j < src_tile_dim.x);
+      }
+
+      // Read valid indices into tmp_val
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = src[(tmp_idx[j] ? i * src_ld + j : 0)];
+      }
+
+      // Zero out unneeded values
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = tmp_idx[j] ? tmp_val[j] : T(0);
+      }
+
+      // Copy values to threadgroup memory
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * kDstStrRow + j * kDstStrCol] = tmp_val[j];
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    src += tile_stride;
+  }
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/mma.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/mma.h
new file mode 100644
index 0000000..737e930
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/mma.h
@@ -0,0 +1,750 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_simdgroup>
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/steel/attn/transforms.h"
+#include "mlx/backend/metal/kernels/steel/defines.h"
+#include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// MMA helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <typename RInt, typename CInt>
+struct Shape2D {
+  RInt r;
+  CInt c;
+
+  Shape2D(RInt r_, CInt c_) : r(r_), c(c_) {}
+};
+
+template <typename Shape, typename Layout>
+struct Layout2D {
+  Shape shape;
+  Layout layout;
+};
+
+template <typename T, int kFragRows_, int kFragCols_>
+struct BaseMMAFrag {
+  static_assert(
+      kFragRows_ == 8,
+      "Only 8 x 8 fragment matrices are currently supported");
+  static_assert(
+      kFragCols_ == 8,
+      "Only 8 x 8 fragment matrices are currently supported");
+};
+
+template <typename T>
+struct BaseMMAFrag<T, 8, 8> {
+  STEEL_CONST int kFragRows = 8;
+  STEEL_CONST int kFragCols = 8;
+
+  STEEL_CONST int kElemsPerFrag = (kFragRows * kFragCols) / 32;
+
+  STEEL_CONST int kElemRows = 1;
+  STEEL_CONST int kElemCols = 2;
+
+  static_assert(
+      kElemRows * kElemCols == kElemsPerFrag,
+      "MMAFrag shape is not consistent with MMAFrag size");
+
+  typedef metal::simdgroup_matrix<T, kFragRows, kFragCols> mat_type;
+  typedef metal::vec<T, kElemsPerFrag> frag_type;
+  typedef metal::vec<T, kElemRows> row_frag_type;
+  typedef metal::vec<T, kElemCols> col_frag_type;
+
+  template <typename U>
+  using dtype_mat_t = typename metal::simdgroup_matrix<U, kFragRows, kFragCols>;
+
+  template <typename U>
+  using dtype_frag_t = typename metal::vec<U, kElemsPerFrag>;
+
+  METAL_FUNC static constexpr short2 get_coord(
+      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
+    const short qid = simd_lane_id / 4;
+    const short fm = (qid & 4) + ((simd_lane_id / 2) % 4);
+    const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
+    return short2{fn, fm};
+  }
+
+  template <typename SrcPtrType, typename StrX, typename StrY>
+  METAL_FUNC static constexpr void
+  load(thread frag_type& dst, SrcPtrType src, StrX str_x, StrY str_y) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        dst[i * kElemCols + j] = static_cast<T>(src[i * str_x + j * str_y]);
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX,
+      typename OffY>
+  METAL_FUNC static constexpr void load_safe(
+      thread frag_type& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    src += off_x * str_x + off_y * str_y;
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
+          dst[i * kElemCols + j] = static_cast<T>(src[0]);
+        } else {
+          dst[i * kElemCols + j] = T(0);
+        }
+        src += str_y;
+      }
+      src -= kElemCols * str_y;
+      src += str_x;
+    }
+  }
+
+  template <typename DstPtrType, typename StrX, typename StrY>
+  METAL_FUNC static constexpr void
+  store(const thread frag_type& src, DstPtrType dst, StrX str_x, StrY str_y) {
+    using U = pointer_element_t<DstPtrType>;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        dst[i * str_x + j * str_y] = static_cast<U>(src[i * kElemCols + j]);
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX,
+      typename OffY>
+  METAL_FUNC static constexpr void store_safe(
+      const thread frag_type& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
+          dst[(off_x + i) * str_x + (off_y + j) * str_y] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  template <typename Atype, typename Btype, typename Ctype>
+  METAL_FUNC static constexpr void mma(
+      thread frag_type& D,
+      thread dtype_frag_t<Atype>& A,
+      thread dtype_frag_t<Btype>& B,
+      thread dtype_frag_t<Ctype>& C) {
+    mat_type D_mat;
+    dtype_mat_t<Atype> A_mat;
+    dtype_mat_t<Btype> B_mat;
+    dtype_mat_t<Ctype> C_mat;
+
+    reinterpret_cast<thread dtype_frag_t<Atype>&>(A_mat.thread_elements()) = A;
+    reinterpret_cast<thread dtype_frag_t<Btype>&>(B_mat.thread_elements()) = B;
+    reinterpret_cast<thread dtype_frag_t<Ctype>&>(C_mat.thread_elements()) = C;
+
+    mma(D_mat, A_mat, B_mat, C_mat);
+
+    D = reinterpret_cast<thread frag_type&>(D_mat.thread_elements());
+  }
+
+  template <typename Atype, typename Btype, typename Ctype>
+  METAL_FUNC static constexpr void mma(
+      thread mat_type& D,
+      thread dtype_mat_t<Atype>& A,
+      thread dtype_mat_t<Btype>& B,
+      thread dtype_mat_t<Ctype>& C) {
+    simdgroup_multiply_accumulate(D, A, B, C);
+  }
+
+  template <typename Op>
+  METAL_FUNC static constexpr void row_reduce(
+      thread const frag_type& inp_vals,
+      thread T* reduced_vals) {
+    T thr_reduce = Op::apply(inp_vals.x, inp_vals.y);
+
+    T qgr_reduce = simd_shuffle_xor(thr_reduce, ushort(1));
+    qgr_reduce = Op::apply(thr_reduce, qgr_reduce);
+
+    T sgr_reduce = simd_shuffle_xor(qgr_reduce, ushort(8));
+    sgr_reduce = Op::apply(qgr_reduce, sgr_reduce);
+
+    reduced_vals[0] = Op::apply(reduced_vals[0], sgr_reduce);
+  }
+
+  template <typename Op>
+  METAL_FUNC static constexpr void row_bin_op(
+      thread frag_type& inp_vals,
+      thread T* row_vals) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        inp_vals[i * kElemCols + j] =
+            Op::apply(inp_vals[i * kElemCols + j], row_vals[i]);
+      }
+    }
+  }
+};
+
+template <
+    typename T,
+    int kTileRows_,
+    int kTileCols_,
+    class MMAFrag_ = BaseMMAFrag<T, 8, 8>>
+struct MMATile {
+  using MMAFrag_t = MMAFrag_;
+  using elem_type = T;
+  STEEL_CONST int kFragRows = MMAFrag_t::kFragRows;
+  STEEL_CONST int kFragCols = MMAFrag_t::kFragCols;
+  STEEL_CONST int kElemsPerFrag = MMAFrag_t::kElemsPerFrag;
+
+  STEEL_CONST int kTileRows = kTileRows_;
+  STEEL_CONST int kTileCols = kTileCols_;
+
+  STEEL_CONST int kRows = kTileRows * kFragRows;
+  STEEL_CONST int kCols = kTileCols * kFragCols;
+
+  STEEL_CONST int kNumFrags = kTileRows * kTileCols;
+  STEEL_CONST int kElemsPerTile = kNumFrags * kElemsPerFrag;
+
+  STEEL_CONST int kRowsPerThread = kTileRows * MMAFrag_t::kElemRows;
+  STEEL_CONST int kColsPerThread = kTileCols * MMAFrag_t::kElemCols;
+
+  typedef typename MMAFrag_t::mat_type mat_type;
+  typedef typename MMAFrag_t::frag_type frag_type;
+
+  frag_type val_frags[kNumFrags]; // = {frag_type(0)};
+
+  METAL_FUNC MMATile() thread {}
+
+  METAL_FUNC constexpr void clear() {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kNumFrags; ++i) {
+      val_frags[i] = frag_type(0);
+    }
+  }
+
+  METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) {
+    return val_frags[i * kTileCols + j];
+  }
+
+  METAL_FUNC constexpr const thread frag_type& frag_at(
+      const short i,
+      const short j) const {
+    return val_frags[i * kTileCols + j];
+  }
+
+  METAL_FUNC mat_type mat_at(const short i, const short j) {
+    mat_type val_mat;
+    STEEL_PRAGMA_UNROLL
+    for (short ii = 0; ii < kElemsPerFrag; ++ii) {
+      val_mat.thread_elements()[ii] = frag_at(i, j)[ii];
+    }
+    return val_mat;
+  }
+
+  METAL_FUNC thread elem_type* elems() {
+    return reinterpret_cast<thread elem_type*>(val_frags);
+  }
+
+  METAL_FUNC const thread elem_type* elems() const {
+    return reinterpret_cast<const thread elem_type*>(val_frags);
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::template row_reduce<Op>(
+            frag_at(i, j), &vals[i * MMAFrag_t::kElemRows]);
+      }
+    }
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::template row_bin_op<Op>(
+            frag_at(i, j), &vals[i * MMAFrag_t::kElemRows]);
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y, int str_x, int str_y>
+  METAL_FUNC void load(const threadgroup U* src) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load(
+            frag_at(i, j),
+            &(
+                src[(i * kFragRows) * w_x * str_x +
+                    (j * kFragCols) * w_y * str_y]),
+            Int<str_x>{},
+            Int<str_y>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y, int str_x, int str_y>
+  METAL_FUNC void store(threadgroup U* dst) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store(
+            frag_at(i, j),
+            &(
+                dst[(i * kFragRows) * w_x * str_x +
+                    (j * kFragCols) * w_y * str_y]),
+            Int<str_x>{},
+            Int<str_y>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void load(const device U* src, const int ld) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load(
+            frag_at(i, j),
+            &(src[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
+            ld,
+            Int<1>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void store(device U* dst, const int ld) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store(
+            frag_at(i, j),
+            &(dst[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
+            ld,
+            Int<1>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void
+  load_safe(const device U* src, const int ld, const short2 src_tile_dims) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load_safe(
+            frag_at(i, j),
+            src,
+            ld,
+            Int<1>{},
+            src_tile_dims.y,
+            src_tile_dims.x,
+            (i * kFragRows) * w_x,
+            (j * kFragCols) * w_y);
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void
+  store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store_safe(
+            frag_at(i, j),
+            dst,
+            ld,
+            Int<1>{},
+            dst_tile_dims.y,
+            dst_tile_dims.x,
+            (i * kFragRows) * w_x,
+            (j * kFragCols) * w_y);
+      }
+    }
+  }
+};
+
+template <
+    typename Dtype,
+    typename Atype,
+    typename Btype,
+    typename Ctype,
+    int M,
+    int N,
+    int K,
+    class MMAFragD,
+    class MMAFragA,
+    class MMAFragB,
+    class MMAFragC>
+METAL_FUNC void tile_matmad(
+    thread MMATile<Dtype, M, N, MMAFragD>& D,
+    thread MMATile<Atype, M, K, MMAFragA>& A,
+    thread MMATile<Btype, K, N, MMAFragB>& B,
+    thread MMATile<Ctype, M, N, MMAFragC>& C) {
+  STEEL_PRAGMA_UNROLL
+  for (short m = 0; m < M; ++m) {
+    STEEL_PRAGMA_UNROLL
+    for (short n = 0; n < N; ++n) {
+      short m_serp = m; //(n % 2) ? (M - 1 - m) : m;
+      short n_serp = (m % 2) ? (N - 1 - n) : n;
+
+      STEEL_PRAGMA_UNROLL
+      for (short k = 0; k < K; ++k) {
+        MMAFragD::mma(
+            D.frag_at(m_serp, n_serp),
+            A.frag_at(m_serp, k),
+            B.frag_at(k, n_serp),
+            C.frag_at(m_serp, n_serp));
+      }
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    short lda_tgp,
+    short ldb_tgp,
+    typename AccumType = float,
+    typename Epilogue = TransformNone<U, AccumType>>
+struct BlockMMA {
+  // MMAFrag size
+  STEEL_CONST short kFragSize = 8;
+  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
+
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TM_stride = kFragSize * WM;
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TN_stride = kFragSize * WN;
+
+  // Warp tile size along M
+  STEEL_CONST short TM = BM / TM_stride;
+  // Warp tile size along N
+  STEEL_CONST short TN = BN / TN_stride;
+
+  // Threadgroup A strides
+  STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M
+  STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K
+
+  // Threadgroup B strides
+  STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K
+  STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N
+
+  // Threadgroup strides along K
+  STEEL_CONST short tile_stride_a = kFragSize * A_str_k;
+  STEEL_CONST short tile_stride_b = kFragSize * B_str_k;
+
+  // Simdgroup matrices
+  MMATile<AccumType, TM, 1, MMAFrag_acc_t> Atile;
+  MMATile<AccumType, 1, TN, MMAFrag_acc_t> Btile;
+  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile;
+
+  // Offsets within threadgroup
+  short sm;
+  short sn;
+
+  short As_offset;
+  short Bs_offset;
+
+  /* Constructor */
+  METAL_FUNC BlockMMA(
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
+    // Determine thread position in simdgroup matrix
+    short tm = kFragSize * (simd_group_id / WN);
+    short tn = kFragSize * (simd_group_id % WN);
+
+    short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
+    sm = simd_coord.y;
+    sn = simd_coord.x;
+
+    // Determine thread and simdgroup offset
+    As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // M, K
+    Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // K, N
+
+    sm += tm;
+    sn += tn;
+  }
+
+  /* (BM, BK) X (BK, BN) multiply accumulate function */
+  METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) {
+    // Adjust for simdgroup and thread location
+    As += As_offset;
+    Bs += Bs_offset;
+
+    // Iterate over BK in blocks of kFragSize
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < BK; kk += kFragSize) {
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Atile.template load<T, WM, 1, A_str_m, A_str_k>(As);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Btile.template load<T, 1, WN, B_str_k, B_str_n>(Bs);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      tile_matmad(Ctile, Atile, Btile, Ctile);
+
+      // Progress to next simdgroup tile
+      As += tile_stride_a;
+      Bs += tile_stride_b;
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(device U* D, const int ldd) {
+    // Apply epilogue
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
+    }
+
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+
+    Ctile.template store<U, WM, WN>(D, ldd);
+  }
+
+  METAL_FUNC void
+  store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) {
+    // Apply epilogue
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
+    }
+
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    Ctile.template store_safe<U, WM, WN>(D, ldd, dst_tile_dims);
+  }
+
+  /* Apply epilogue */
+  template <typename UnaryEpilogue>
+  METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) {
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = epilogue_op.apply(Ctile.elems()[i]);
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < decltype(Ctile)::kElemsPerFrag; k++) {
+          accum[k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue_safe(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+        // Read C
+        U c_elems[kelems] = {0};
+
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          if ((j * TN_stride + k) < dst_tile_dims.x) {
+            c_elems[k] = C[offset_c + k * fdc];
+          }
+        }
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          accum[k] = epilogue_op.apply(accum[k], c_elems[k]);
+        }
+      }
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+
+    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread const auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+        int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          D[offset_d + k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void store_result_safe(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < TM; i++) {
+      if (i * TM_stride < dst_tile_dims.y) {
+        STEEL_PRAGMA_UNROLL
+        for (int j = 0; j < TN; j++) {
+          // Get accumulated result and associated offset in C
+          thread const auto& accum = Ctile.frag_at(i, j);
+          int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+          int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+          // Apply epilogue
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < kelems; k++) {
+            if ((j * TN_stride + k) < dst_tile_dims.x) {
+              D[offset_d + k] =
+                  epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+            }
+          }
+        }
+      }
+    }
+  }
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/nax.h
new file mode 100644
index 0000000..c8f3ea5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/nax.h
@@ -0,0 +1,1076 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <metal_simdgroup>
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+#include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
+
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// MMA helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+///////////////////////////////////////////////////////////////////////////////
+// NAX Steel with new tiles
+///////////////////////////////////////////////////////////////////////////////
+
+struct BaseNAXFrag {
+  STEEL_CONST short kFragRows = 16;
+  STEEL_CONST short kFragCols = 16;
+
+  STEEL_CONST short kElemsPerFrag = (kFragRows * kFragCols) / 32;
+
+  STEEL_CONST short kElemRows = 2;
+  STEEL_CONST short kElemCols = 4;
+
+  STEEL_CONST short kElemRowsJump = 8;
+
+  static_assert(
+      kElemRows * kElemCols == kElemsPerFrag,
+      "MMAFrag shape is not consistent with MMAFrag size");
+
+  template <typename U>
+  using dtype_frag_t = typename metal::vec<U, kElemsPerFrag>;
+
+  METAL_FUNC static short2 get_coord() {
+    const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort());
+    const short qid = simd_lane_id >> 2;
+    const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3));
+    const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4;
+    return short2{fn, fm};
+  }
+
+  METAL_FUNC static short2 get_coord(short idx) {
+    const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort());
+    const short qid = simd_lane_id >> 2;
+    const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)) + (idx >> 2) * 8;
+    const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4 + idx % 4;
+    return short2{fn, fm};
+  }
+
+  template <
+      typename T,
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void load(
+      thread dtype_frag_t<T>& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    const short2 sc = short2{0, 0}; // get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if constexpr (metal::is_same_v<StrY, Int<1>>) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[i * kElemCols + j] = static_cast<T>(src[r * str_x + c + j]);
+        }
+      } else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[i * kElemCols + j] =
+              static_cast<T>(src[r * str_x + (c + j) * str_y]);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void load_rows(
+      thread dtype_frag_t<T>& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    const short2 sc = short2{0, 0}; // get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if (r < lim_x) {
+        if constexpr (metal::is_same_v<StrY, Int<1>>) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[i * kElemCols + j] = static_cast<T>(src[r * str_x + (c + j)]);
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[i * kElemCols + j] =
+                static_cast<T>(src[r * str_x + (c + j) * str_y]);
+          }
+        }
+
+      } else {
+        dst = dtype_frag_t<T>(0);
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void load_safe(
+      thread dtype_frag_t<T>& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    const short2 sc = short2{0, 0}; // get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if (r < lim_x && (c + j) < lim_y) {
+          dst[i * kElemCols + j] =
+              static_cast<T>(src[r * str_x + (c + j) * str_y]);
+        } else {
+          dst[i * kElemCols + j] = T(0);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = short2{0, 0}; // get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if constexpr (metal::is_same_v<StrY, Int<1>>) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[r * str_x + c + j] = static_cast<U>(src[i * kElemCols + j]);
+        }
+      } else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[r * str_x + (c + j) * str_y] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store_rows(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = short2{0, 0}; // get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if (r < lim_x) {
+        if constexpr (metal::is_same_v<StrY, Int<1>>) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[r * str_x + c + j] = static_cast<U>(src[i * kElemCols + j]);
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[r * str_x + (c + j) * str_y] =
+                static_cast<U>(src[i * kElemCols + j]);
+          }
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store_safe(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = short2{0, 0}; // get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if (r < lim_x && (c + j) < lim_y) {
+          dst[r * str_x + (c + j) * str_y] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename StartX,
+      typename StopX,
+      typename StartY,
+      typename StopY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store_slice(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      StartX start_x,
+      StopX stop_x,
+      StartY start_y,
+      StopY stop_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = short2{0, 0}; // get_coord();
+
+    const_for_loop<0, kElemRows, 1>([&](auto idx_row) {
+      const auto r = off_x + idx_row * Int<kElemRowsJump>{};
+      if (r >= stop_x - sc.y || r < start_x - sc.y) {
+        return;
+      }
+
+      const_for_loop<0, kElemCols, 1>([&](auto idx_col) {
+        const auto c = off_y + idx_col;
+        if (c >= stop_y - sc.x || c < start_y - sc.x) {
+          return;
+        }
+
+        const auto src_idx = idx_row * Int<kElemCols>{} + idx_col;
+        dst[(r + sc.y) * str_x + (c + sc.x) * str_y] =
+            static_cast<U>(src[src_idx]);
+      });
+    });
+  }
+
+  template <typename Op, typename T>
+  METAL_FUNC static constexpr void row_reduce(
+      thread const dtype_frag_t<T>& inp_vals,
+      thread T* reduced_vals) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      T thr_reduce = Op::apply(
+          Op::apply(inp_vals[i * kElemCols + 0], inp_vals[i * kElemCols + 1]),
+          Op::apply(inp_vals[i * kElemCols + 2], inp_vals[i * kElemCols + 3]));
+
+      T qgr_reduce = simd_shuffle_xor(thr_reduce, ushort(1));
+      qgr_reduce = Op::apply(thr_reduce, qgr_reduce);
+
+      T sgr_reduce = simd_shuffle_xor(qgr_reduce, ushort(8));
+      sgr_reduce = Op::apply(qgr_reduce, sgr_reduce);
+
+      reduced_vals[i] = Op::apply(reduced_vals[i], sgr_reduce);
+    }
+  }
+
+  template <typename Op, typename T>
+  METAL_FUNC static constexpr void row_bin_op(
+      thread dtype_frag_t<T>& inp_vals,
+      thread T* row_vals) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        inp_vals[i * kElemCols + j] =
+            Op::apply(inp_vals[i * kElemCols + j], row_vals[i]);
+      }
+    }
+  }
+};
+
+template <
+    typename T,
+    short kRows_,
+    short kCols_,
+    typename NAXFrag_ = BaseNAXFrag>
+struct NAXSubTile {
+  using NAXFrag_t = NAXFrag_;
+  STEEL_CONST short kRows = kRows_;
+  STEEL_CONST short kCols = kCols_;
+
+  STEEL_CONST short kFragRows = NAXFrag_t::kFragRows;
+  STEEL_CONST short kFragCols = NAXFrag_t::kFragCols;
+  STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag;
+
+  STEEL_CONST short kSubTileRows = kRows / kFragRows;
+  STEEL_CONST short kSubTileCols = kCols / kFragCols;
+
+  STEEL_CONST short kNumFrags = kSubTileRows * kSubTileCols;
+  STEEL_CONST short kElemsPerSubTile = kNumFrags * kElemsPerFrag;
+
+  STEEL_CONST int kRowsPerThread = kSubTileRows * NAXFrag_t::kElemRows;
+  STEEL_CONST int kColsPerThread = kSubTileCols * NAXFrag_t::kElemCols;
+
+  STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows;
+  STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols;
+  STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump;
+
+  using frag_type = typename NAXFrag_t::template dtype_frag_t<T>;
+
+  frag_type val_frags[kNumFrags];
+
+  METAL_FUNC constexpr void clear() {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kNumFrags; ++i) {
+      val_frags[i] = frag_type(0);
+    }
+  }
+
+  METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  METAL_FUNC constexpr const thread frag_type& frag_at(
+      const short i,
+      const short j) const {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  template <int i, int j>
+  METAL_FUNC constexpr thread frag_type& frag_at() {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  template <int i, int j>
+  METAL_FUNC constexpr const thread frag_type& frag_at() const {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  METAL_FUNC thread T* elems() {
+    return reinterpret_cast<thread T*>(val_frags);
+  }
+
+  METAL_FUNC const thread T* elems() const {
+    return reinterpret_cast<const thread T*>(val_frags);
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_reduce(thread metal::vec<T, kRowsPerThread>& vals) const {
+    thread T* vptr = (thread T*)(&vals);
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::template row_reduce<Op>(
+            frag_at(i, j), &vptr[i * kFragThrRows]);
+      }
+    }
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_bin_op(thread metal::vec<T, kRowsPerThread>& vals) {
+    thread T* vptr = (thread T*)(&vals);
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::template row_bin_op<Op>(
+            frag_at(i, j), &vptr[i * kFragThrRows]);
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void load(
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::load(
+            frag_at(i, j),
+            src,
+            str_x,
+            str_y,
+            off_x + i * kFragRows,
+            off_y + j * kFragCols);
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::store(
+            frag_at(i, j),
+            dst,
+            str_x,
+            str_y,
+            off_x + i * kFragRows,
+            off_y + j * kFragCols);
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void load_rows(
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::load_rows(
+            frag_at(i, j),
+            src,
+            str_x,
+            str_y,
+            lim_x,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void load_safe(
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::load_safe(
+            frag_at(i, j),
+            src,
+            str_x,
+            str_y,
+            lim_x,
+            lim_y,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store_rows(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::store_safe(
+            frag_at(i, j),
+            dst,
+            str_x,
+            str_y,
+            lim_x,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store_safe(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::store_safe(
+            frag_at(i, j),
+            dst,
+            str_x,
+            str_y,
+            lim_x,
+            lim_y,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename StartX,
+      typename StopX,
+      typename StartY,
+      typename StopY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store_slice(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      StartX start_x,
+      StopX stop_x,
+      StartY start_y,
+      StopY stop_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) const {
+    const_for_loop<0, kSubTileRows, 1>([&](auto idx_row) {
+      const_for_loop<0, kSubTileCols, 1>([&](auto idx_col) {
+        NAXFrag_t::store_slice(
+            frag_at<idx_row.value, idx_col.value>(),
+            dst,
+            str_x,
+            str_y,
+            start_x,
+            stop_x,
+            start_y,
+            stop_y,
+            off_x + idx_row * Int<kFragRows>{},
+            off_y + idx_col * Int<kFragCols>{});
+      });
+    });
+  }
+};
+
+template <
+    short RC,
+    short CC,
+    short RA,
+    short CA,
+    short RB,
+    short CB,
+    typename CType,
+    typename AType,
+    typename BType,
+    bool transpose_a,
+    bool transpose_b,
+    typename NAXFrag_t = BaseNAXFrag>
+METAL_FUNC void subtile_matmad_nax(
+    thread NAXSubTile<CType, RC, CC, NAXFrag_t>& C,
+    thread NAXSubTile<AType, RA, CA, NAXFrag_t>& A,
+    metal::bool_constant<transpose_a>,
+    thread NAXSubTile<BType, RB, CB, NAXFrag_t>& B,
+    metal::bool_constant<transpose_b>) {
+  // Static checks
+  constexpr short FMa = transpose_a ? CA : RA;
+  constexpr short FMc = RC;
+  static_assert(FMa == FMc, "NAX matmul: M dimensions do not match");
+
+  constexpr short FNb = transpose_b ? RB : CB;
+  constexpr short FNc = CC;
+  static_assert(FNb == FNc, "NAX matmul: N dimensions do not match");
+
+  constexpr short FKa = transpose_a ? RA : CA;
+  constexpr short FKb = transpose_b ? CB : RB;
+  static_assert(FKa == FKb, "NAX matmul: N dimensions do not match");
+
+  constexpr short FM = FMc;
+  constexpr short FN = FNc;
+  constexpr short FK = FKa;
+
+  constexpr int TM = FM / 16;
+  constexpr int TN = FN / 16;
+  constexpr int TK = FK / 16;
+
+  constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor(
+      FM,
+      FN,
+      FK,
+      transpose_a,
+      transpose_b,
+      true,
+      mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate);
+
+  mpp::tensor_ops::matmul2d<desc, metal::execution_simdgroup> gemm_op;
+
+  auto ct_a =
+      gemm_op.template get_left_input_cooperative_tensor<AType, BType, CType>();
+  auto ct_b =
+      gemm_op
+          .template get_right_input_cooperative_tensor<AType, BType, CType>();
+  auto ct_c = gemm_op.template get_destination_cooperative_tensor<
+      decltype(ct_a),
+      decltype(ct_b),
+      CType>();
+
+  STEEL_PRAGMA_UNROLL
+  for (short mm = 0; mm < TM; mm++) {
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < TK; kk++) {
+      const short fi = transpose_a ? kk : mm;
+      const short fj = transpose_a ? mm : kk;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < 8; i++) {
+        ct_a[(TK * mm + kk) * 8 + i] = A.frag_at(fi, fj)[i];
+      }
+    }
+  }
+
+  STEEL_PRAGMA_UNROLL
+  for (short nn = 0; nn < TN; nn++) {
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < TK; kk++) {
+      const short fi = transpose_b ? nn : kk;
+      const short fj = transpose_b ? kk : nn;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < 8; i++) {
+        ct_b[(TN * kk + nn) * 8 + i] = B.frag_at(fi, fj)[i];
+      }
+    }
+  }
+
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < ct_c.get_capacity(); i++) {
+    ct_c[i] = C.elems()[i];
+  }
+
+  gemm_op.run(ct_a, ct_b, ct_c);
+
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < ct_c.get_capacity(); i++) {
+    C.elems()[i] = ct_c[i];
+  }
+}
+
+template <typename T, short kTileRows_, short kTileCols_, class NAXSubTile_>
+struct NAXTile {
+  using NAXSubTile_t = NAXSubTile_;
+  using elem_type = T;
+  STEEL_CONST short kSubTileRows = NAXSubTile_t::kRows;
+  STEEL_CONST short kSubTileCols = NAXSubTile_t::kCols;
+  STEEL_CONST short kElemsPerSubTile = NAXSubTile_t::kElemsPerSubTile;
+
+  STEEL_CONST short kTileRows = kTileRows_;
+  STEEL_CONST short kTileCols = kTileCols_;
+
+  STEEL_CONST short kRows = kTileRows * kSubTileRows;
+  STEEL_CONST short kCols = kTileCols * kSubTileCols;
+
+  STEEL_CONST short kSubTiles = kTileRows * kTileCols;
+  STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile;
+
+  STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread;
+  STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread;
+
+  STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread;
+  STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread;
+
+  NAXSubTile_t val_subtiles[kSubTiles];
+
+  METAL_FUNC NAXTile() thread {}
+
+  METAL_FUNC constexpr void clear() {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTiles; ++i) {
+      val_subtiles[i].clear();
+    }
+  }
+
+  METAL_FUNC constexpr thread NAXSubTile_t& subtile_at(
+      const short i,
+      const short j) {
+    return val_subtiles[i * kTileCols + j];
+  }
+
+  METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at(
+      const short i,
+      const short j) const {
+    return val_subtiles[i * kTileCols + j];
+  }
+
+  template <int i, int j>
+  METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const {
+    return val_subtiles[i * kTileCols + j];
+  }
+
+  METAL_FUNC thread elem_type* elems() {
+    return reinterpret_cast<thread elem_type*>(val_subtiles[0].elems());
+  }
+
+  METAL_FUNC const thread elem_type* elems() const {
+    return reinterpret_cast<const thread elem_type*>(val_subtiles[0].elems());
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_reduce(thread metal::vec<T, kRowsPerThread>& vals) const {
+    auto sub_rows = (thread metal::vec<T, kSubTileThrRows>*)(&vals);
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).template row_reduce<Op>(sub_rows[i]);
+      }
+    }
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_bin_op(thread metal::vec<T, kRowsPerThread>& vals) {
+    auto sub_rows = (thread metal::vec<T, kSubTileThrRows>*)(&vals);
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).template row_bin_op<Op>(sub_rows[i]);
+      }
+    }
+  }
+
+  template <typename U, int str_x, int str_y>
+  METAL_FUNC void load(const threadgroup U* src) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load(
+            src,
+            Int<str_x>{},
+            Int<str_y>{},
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U, int str_x, int str_y>
+  METAL_FUNC void store(threadgroup U* dst) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store(
+            dst,
+            Int<str_x>{},
+            Int<str_y>{},
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void load(const device U* src, const int ld) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load(
+            &src[(i * kSubTileRows) * ld + (j * kSubTileCols)], ld, Int<1>{});
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void store(device U* dst, const int ld) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store(
+            &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)], ld, Int<1>{});
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void
+  load_safe(const device U* src, const int ld, const short2 src_tile_dims) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load_safe(
+            src,
+            ld,
+            Int<1>{},
+            src_tile_dims.y,
+            src_tile_dims.x,
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void
+  load_rows(const device U* src, const int ld, const short n_rows) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load_rows(
+            &src[(i * kSubTileRows) * ld + (j * kSubTileCols)],
+            ld,
+            Int<1>{},
+            n_rows - i * kSubTileRows);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void
+  store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store_safe(
+            dst,
+            ld,
+            Int<1>{},
+            dst_tile_dims.y,
+            dst_tile_dims.x,
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows)
+      const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store_rows(
+            &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)],
+            ld,
+            Int<1>{},
+            n_rows - i * kSubTileRows);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void store_slice(
+      device U* dst,
+      const int ld,
+      const short2 start,
+      const short2 stop) const {
+    const_for_loop<0, kTileRows, 1>([&](auto idx_row) {
+      const_for_loop<0, kTileCols, 1>([&](auto idx_col) {
+        subtile_at<idx_row.value, idx_col.value>().store_slice(
+            dst,
+            ld,
+            Int<1>{},
+            start.y,
+            stop.y,
+            start.x,
+            stop.x,
+            idx_row * Int<kSubTileRows>{},
+            idx_col * Int<kSubTileCols>{});
+      });
+    });
+  }
+};
+
+template <
+    class CTile,
+    class ATile,
+    class BTile,
+    bool transpose_a,
+    bool transpose_b>
+METAL_FUNC void tile_matmad_nax(
+    thread CTile& C,
+    thread ATile& A,
+    metal::bool_constant<transpose_a>,
+    thread BTile& B,
+    metal::bool_constant<transpose_b>) {
+  // Static checks
+  constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows;
+  constexpr short TMc = CTile::kTileRows;
+  static_assert(TMa == TMc, "NAX tile matmul: M dimensions do not match");
+
+  constexpr short FMa = transpose_a ? ATile::kSubTileCols : ATile::kSubTileRows;
+  constexpr short FMc = CTile::kSubTileRows;
+  static_assert(FMa == FMc, "NAX subtile matmul: M dimensions do not match");
+
+  constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols;
+  constexpr short TNc = CTile::kTileCols;
+  static_assert(TNb == TNc, "NAX tile matmul: N dimensions do not match");
+
+  constexpr short FNb = transpose_b ? BTile::kSubTileRows : BTile::kSubTileCols;
+  constexpr short FNc = CTile::kSubTileCols;
+  static_assert(FNb == FNc, "NAX subtile matmul: N dimensions do not match");
+
+  constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols;
+  constexpr short TKb = transpose_b ? BTile::kTileCols : BTile::kTileRows;
+  static_assert(TKa == TKb, "NAX tile matmul: K dimensions do not match");
+
+  constexpr short FKa = transpose_a ? ATile::kSubTileRows : ATile::kSubTileCols;
+  constexpr short FKb = transpose_b ? BTile::kSubTileCols : BTile::kSubTileRows;
+  static_assert(FKa == FKb, "NAX subtile matmul: K dimensions do not match");
+
+  constexpr short TM = TMc;
+  constexpr short TN = TNc;
+  constexpr short TK = TKa;
+
+  // Do matmul here
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < TM; ++i) {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < TN; ++j) {
+      STEEL_PRAGMA_UNROLL
+      for (short k = 0; k < TK; ++k) {
+        const short ra = transpose_a ? k : i;
+        const short ca = transpose_a ? i : k;
+        const short rb = transpose_b ? j : k;
+        const short cb = transpose_b ? k : j;
+
+        subtile_matmad_nax(
+            C.subtile_at(i, j),
+            A.subtile_at(ra, ca),
+            metal::bool_constant<transpose_a>{},
+            B.subtile_at(rb, cb),
+            metal::bool_constant<transpose_b>{});
+      }
+    }
+  }
+}
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/params.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/params.h
new file mode 100644
index 0000000..f1cf09f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/params.h
@@ -0,0 +1,44 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+///////////////////////////////////////////////////////////////////////////////
+// Attn param classes
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+struct AttnParams {
+  int B; ///< Batch Size
+  int H; ///< Heads
+  int D; ///< Head Dim
+
+  int qL; ///< Query Sequence Length
+  int kL; ///< Key Sequence Length
+
+  int gqa_factor; ///< Group Query factor
+  float scale; ///< Attention scale
+
+  int NQ; ///< Number of query blocks
+  int NK; ///< Number of key/value blocks
+
+  int NQ_aligned; ///< Number of full query blocks
+  int NK_aligned; ///< Number of full key/value blocks
+
+  int qL_rem; ///< Remainder in last query block
+  int kL_rem; ///< Remainder in last key/value block
+  int qL_off; ///< Offset in query sequence start
+
+  int64_t Q_strides[3]; ///< Query  strides (B, H, L, D = 1)
+  int64_t K_strides[3]; ///< Key    strides (B, H, L, D = 1)
+  int64_t V_strides[3]; ///< Value  strides (B, H, L, D = 1)
+  int64_t O_strides[3]; ///< Output strides (B, H, L, D = 1)
+};
+
+struct AttnMaskParams {
+  int64_t M_strides[3]; ///< Mask  strides (B, H, qL, kL = 1)
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/transforms.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/transforms.h
new file mode 100644
index 0000000..c0624d2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/attn/transforms.h
@@ -0,0 +1,71 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Transforms and Epilogues
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <typename OutT, typename InT>
+struct TransformNone {
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT) {
+    return static_cast<OutT>(x);
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAdd {
+  TransformAdd(const float, const float) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT c) {
+    return static_cast<OutT>(x) + c;
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAxpby {
+  const float alpha;
+  const float beta;
+
+  TransformAxpby(const float alpha_, const float beta_)
+      : alpha(alpha_), beta(beta_) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  METAL_FUNC OutT apply(InT x, OutT c) const {
+    return static_cast<OutT>(x * alpha + (beta * c));
+  }
+};
+
+template <typename T>
+struct AccumHelper {
+  typedef float accum_type;
+};
+
+struct BlockSwizzle {
+  static METAL_FUNC int2
+  swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) {
+    const int tid_x = (tid.x) >> swizzle_log;
+    const int tid_y =
+        ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1));
+    return int2(tid_x, tid_y);
+  }
+};
+
+} // namespace steel
+} // namespace mlx
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/conv.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/conv.h
new file mode 100644
index 0000000..d2e718f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/conv.h
@@ -0,0 +1,13 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/loader.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+
+using namespace metal;
+using namespace mlx::steel;
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h
new file mode 100644
index 0000000..850ec15
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h
@@ -0,0 +1,176 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_stdlib>
+
+using namespace metal;
+
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    int N_CHANNELS = 0,
+    bool SMALL_FILTER = false>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
+implicit_gemm_conv_2d(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device T* C [[buffer(2)]],
+    const constant MLXConvParams<2>* params [[buffer(3)]],
+    const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  using namespace mlx::steel;
+
+  (void)lid;
+
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = true;
+  constexpr short tgp_padding_a = 16 / sizeof(T);
+  constexpr short tgp_padding_b = 16 / sizeof(T);
+
+  constexpr short shape_a_cols = (transpose_a ? BM : BK) + tgp_padding_a;
+  constexpr short shape_b_cols = (transpose_b ? BK : BN) + tgp_padding_b;
+  constexpr short shape_a_rows = (transpose_a ? BK : BM);
+  constexpr short shape_b_rows = (transpose_b ? BN : BK);
+  constexpr short tgp_mem_size_a = shape_a_cols * shape_a_rows;
+  constexpr short tgp_mem_size_b = shape_b_cols * shape_b_rows;
+
+  constexpr short tgp_size = WM * WN * 32;
+
+  // Input loader
+
+  using loader_a_t = typename metal::conditional_t<
+      // Check for small channel specialization
+      N_CHANNELS != 0 && N_CHANNELS <= 4,
+
+      // Go to small channel specialization
+      Conv2DInputBlockLoaderSmallChannels<
+          T,
+          BM,
+          BN,
+          BK,
+          tgp_size,
+          N_CHANNELS,
+          tgp_padding_a>,
+
+      // Else go to general loader
+      typename metal::conditional_t<
+          // Check if filter size is small enough
+          SMALL_FILTER,
+
+          // Go to small filter specialization
+          Conv2DInputBlockLoaderSmallFilter<
+              T,
+              BM,
+              BN,
+              BK,
+              tgp_size,
+              tgp_padding_a>,
+
+          // Else go to large filter generalization
+          Conv2DInputBlockLoaderLargeFilter<
+              T,
+              BM,
+              BN,
+              BK,
+              tgp_size,
+              tgp_padding_a>>>;
+
+  // Weight loader
+  using loader_b_t = typename metal::conditional_t<
+      // Check for small channel specialization
+      N_CHANNELS != 0 && N_CHANNELS <= 4,
+
+      // Go to small channel specialization
+      Conv2DWeightBlockLoaderSmallChannels<
+          T,
+          BM,
+          BN,
+          BK,
+          tgp_size,
+          N_CHANNELS,
+          tgp_padding_b>,
+
+      // Else go to general loader
+      Conv2DWeightBlockLoader<T, BM, BN, BK, tgp_size, tgp_padding_b>>;
+
+  using mma_t = BlockMMA<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      shape_a_cols,
+      shape_b_cols>;
+
+  threadgroup T As[tgp_mem_size_a];
+  threadgroup T Bs[tgp_mem_size_b];
+
+  const int tid_y = ((tid.y) << gemm_params->swizzle_log) +
+      ((tid.x) & ((1 << gemm_params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> gemm_params->swizzle_log;
+
+  if (gemm_params->tiles_n <= tid_x || gemm_params->tiles_m <= tid_y) {
+    return;
+  }
+
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int K = gemm_params->K;
+  const int N = gemm_params->N;
+  const int C_per_group = params->C / params->groups;
+
+  // Groups
+  A += tid.z * C_per_group;
+  B += tid.z * N * K;
+  C += tid.z * N;
+
+  B += c_col * K;
+  C += c_row * (N * params->groups) + c_col;
+
+  const int2 offsets_a(0, c_row);
+  const int2 offsets_b(0, c_col);
+
+  // Prepare threadgroup loading operations
+  loader_a_t loader_a(
+      A, As, offsets_a, params, gemm_params, simd_gid, simd_lid);
+  loader_b_t loader_b(
+      B, Bs, offsets_b, params, gemm_params, simd_gid, simd_lid);
+
+  // Prepare threadgroup mma operation
+  mma_t mma_op(simd_gid, simd_lid);
+
+  int gemm_k_iterations = gemm_params->gemm_k_iterations;
+  for (int k = 0; k < gemm_k_iterations; k++) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Load elements into threadgroup
+    loader_a.load_unsafe();
+    loader_b.load_unsafe();
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(As, Bs);
+
+    // Prepare for next iteration
+    loader_a.next();
+    loader_b.next();
+  }
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Store results to device memory
+  short tgp_bm = min(BM, gemm_params->M - c_row);
+  short tgp_bn = min(BN, gemm_params->N - c_col);
+  const int ldc = N * params->groups;
+  mma_op.store_result_safe(C, ldc, short2(tgp_bn, tgp_bm));
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal
new file mode 100644
index 0000000..7851e5e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal
@@ -0,0 +1,56 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_stdlib>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+#include "mlx/backend/metal/kernels/steel/conv/conv.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+#include "mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h"
+
+#define instantiate_implicit_conv_2d(                                          \
+    name,                                                                      \
+    itype,                                                                     \
+    bm,                                                                        \
+    bn,                                                                        \
+    bk,                                                                        \
+    wm,                                                                        \
+    wn,                                                                        \
+    channel_name,                                                              \
+    n_channels,                                                                \
+    filter_name,                                                               \
+    small_filter)                                                              \
+  template [[host_name("implicit_gemm_conv_2d_" #name "_bm" #bm "_bn" #bn      \
+                       "_bk" #bk "_wm" #wm "_wn" #wn "_channel_" #channel_name \
+                       "_filter_" #filter_name)]] [[kernel]] void              \
+  implicit_gemm_conv_2d<itype, bm, bn, bk, wm, wn, n_channels, small_filter>(  \
+      const device itype* A [[buffer(0)]],                                     \
+      const device itype* B [[buffer(1)]],                                     \
+      device itype* C [[buffer(2)]],                                           \
+      const constant MLXConvParams<2>* params [[buffer(3)]],                   \
+      const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]],      \
+      uint3 tid [[threadgroup_position_in_grid]],                              \
+      uint3 lid [[thread_position_in_threadgroup]],                            \
+      uint simd_gid [[simdgroup_index_in_threadgroup]],                        \
+      uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_implicit_2d_filter(name, itype, bm, bn, bk, wm, wn)           \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, l, 0, s, true)  \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, l, 0, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 1, 1, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 2, 2, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 3, 3, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 4, 4, l, false)
+
+#define instantiate_implicit_2d_blocks(name, itype)               \
+    instantiate_implicit_2d_filter(name, itype, 32,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 64,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 32, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 32, 64, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 64, 16, 2, 2)
+
+instantiate_implicit_2d_blocks(float32, float);
+instantiate_implicit_2d_blocks(float16, half);
+instantiate_implicit_2d_blocks(bfloat16, bfloat16_t); // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h
new file mode 100644
index 0000000..1241f77
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h
@@ -0,0 +1,225 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h"
+
+constant bool align_C [[function_constant(200)]];
+
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    typename AccumType = float,
+    typename Epilogue = TransformNone<T, AccumType>>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
+implicit_gemm_conv_2d_general(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device T* C [[buffer(2)]],
+    const constant MLXConvParams<2>* params [[buffer(3)]],
+    const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]],
+    const constant Conv2DGeneralJumpParams* jump_params [[buffer(5)]],
+    const constant Conv2DGeneralBaseInfo* base_h [[buffer(6)]],
+    const constant Conv2DGeneralBaseInfo* base_w [[buffer(7)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  (void)lid;
+
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = true;
+  constexpr short tgp_padding_a = 16 / sizeof(T);
+  constexpr short tgp_padding_b = 16 / sizeof(T);
+
+  constexpr short shape_a_cols = (transpose_a ? BM : BK) + tgp_padding_a;
+  constexpr short shape_b_cols = (transpose_b ? BK : BN) + tgp_padding_b;
+  constexpr short shape_a_rows = (transpose_a ? BK : BM);
+  constexpr short shape_b_rows = (transpose_b ? BN : BK);
+  constexpr short tgp_mem_size_a = shape_a_cols * shape_a_rows;
+  constexpr short tgp_mem_size_b = shape_b_cols * shape_b_rows;
+
+  constexpr short tgp_size = WM * WN * 32;
+
+  // Input loader
+  using loader_a_t =
+      Conv2DInputBlockLoaderGeneral<T, BM, BN, BK, tgp_size, tgp_padding_a>;
+
+  // Weight loader
+  using loader_b_t =
+      Conv2DWeightBlockLoaderGeneral<T, BM, BN, BK, tgp_size, tgp_padding_b>;
+
+  using mma_t = BlockMMA<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      shape_a_cols,
+      shape_b_cols>;
+
+  threadgroup T As[tgp_mem_size_a];
+  threadgroup T Bs[tgp_mem_size_b];
+
+  const int tid_y = ((tid.y) << gemm_params->swizzle_log) +
+      ((tid.x) & ((1 << gemm_params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> gemm_params->swizzle_log;
+
+  if (gemm_params->tiles_n <= tid_x || gemm_params->tiles_m <= tid_y) {
+    return;
+  }
+
+  const int tid_z = tid.z;
+
+  const int base_oh = tid_z / jump_params->f_out_jump_w;
+  const int base_ow = tid_z % jump_params->f_out_jump_w;
+
+  const int base_wh = base_h[base_oh].weight_base;
+  const int base_ww = base_w[base_ow].weight_base;
+
+  const int base_wh_size = base_h[base_oh].weight_size;
+  const int base_ww_size = base_w[base_ow].weight_size;
+
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int K = gemm_params->K;
+
+  B += c_col * K;
+
+  const int4 offsets_a(0, c_row, base_oh, base_ow);
+  const int2 offsets_b(0, c_col);
+
+  // Prepare threadgroup loading operations
+  loader_a_t loader_a(
+      A,
+      As,
+      offsets_a,
+      params,
+      jump_params,
+      base_wh,
+      base_ww,
+      simd_gid,
+      simd_lid);
+  loader_b_t loader_b(
+      B,
+      Bs,
+      offsets_b,
+      params,
+      jump_params,
+      base_wh,
+      base_ww,
+      simd_gid,
+      simd_lid);
+
+  // Prepare threadgroup mma operation
+  mma_t mma_op(simd_gid, simd_lid);
+
+  if (align_C) {
+    int gemm_k_iterations =
+        base_wh_size * base_ww_size * gemm_params->gemm_k_iterations;
+
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      // Load elements into threadgroup
+      loader_a.load_unsafe();
+      loader_b.load_unsafe();
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+  }
+
+  else {
+    for (int k = 1; k < gemm_params->gemm_k_iterations; k++) {
+      for (int j = 0; j < base_wh_size * base_ww_size; j++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+    }
+    const short remaining_k = params->C % BK;
+    for (int j = 0; j < base_wh_size * base_ww_size; j++) {
+      // Load elements into threadgroup
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      loader_a.load_safe(remaining_k);
+      loader_b.load_safe(remaining_k);
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+  }
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Store results to device memory
+  {
+    // Adjust for simdgroup and thread location
+    int offset_m = c_row + mma_op.sm;
+    int offset_n = c_col + mma_op.sn;
+    C += offset_n;
+
+    if (offset_n >= gemm_params->N)
+      return;
+
+    short diff = gemm_params->N - offset_n;
+
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < mma_t::TM; i++) {
+      int cm = offset_m + i * mma_t::TM_stride;
+
+      int n = cm / jump_params->adj_out_hw;
+      int hw = cm % jump_params->adj_out_hw;
+      int oh =
+          (hw / jump_params->adj_out_w) * jump_params->f_out_jump_h + base_oh;
+      int ow =
+          (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + base_ow;
+
+      if (n < params->N && oh < params->oS[0] && ow < params->oS[1]) {
+        int offset_cm = n * params->out_strides[0] +
+            oh * params->out_strides[1] + ow * params->out_strides[2];
+
+        STEEL_PRAGMA_UNROLL
+        for (int j = 0; j < mma_t::TN; j++) {
+          // Get accumulated result and associated offset in C
+          thread const auto& accum = mma_op.Ctile.frag_at(i, j);
+          int offset = offset_cm + (j * mma_t::TN_stride);
+
+          constexpr short kelems = decltype(mma_op.Ctile)::kElemsPerFrag;
+
+          // Apply epilogue and output C
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < kelems; k++) {
+            if ((j * mma_t::TN_stride + k) < diff) {
+              C[offset + k] = Epilogue::apply(accum[k]);
+            }
+          }
+        }
+      }
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.metal
new file mode 100644
index 0000000..23e8575
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.metal
@@ -0,0 +1,47 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_stdlib>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+#include "mlx/backend/metal/kernels/steel/conv/conv.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+#include "mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h"
+
+using namespace metal;
+using namespace mlx::steel;
+
+#define instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn)         \
+  template                                                                    \
+      [[host_name("implicit_gemm_conv_2d_general_" #name "_bm" #bm "_bn" #bn  \
+                  "_bk" #bk "_wm" #wm "_wn" #wn)]] [[kernel]] void            \
+      implicit_gemm_conv_2d_general<itype, bm, bn, bk, wm, wn>(               \
+          const device itype* A [[buffer(0)]],                                \
+          const device itype* B [[buffer(1)]],                                \
+          device itype* C [[buffer(2)]],                                      \
+          const constant MLXConvParams<2>* params [[buffer(3)]],              \
+          const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]], \
+          const constant Conv2DGeneralJumpParams* jump_params [[buffer(5)]],  \
+          const constant Conv2DGeneralBaseInfo* base_h [[buffer(6)]],         \
+          const constant Conv2DGeneralBaseInfo* base_w [[buffer(7)]],         \
+          uint3 tid [[threadgroup_position_in_grid]],                         \
+          uint3 lid [[thread_position_in_threadgroup]],                       \
+          uint simd_gid [[simdgroup_index_in_threadgroup]],                   \
+          uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_implicit_2d_filter(name, itype, bm, bn, bk, wm, wn) \
+  instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn)
+
+#define instantiate_implicit_2d_blocks(name, itype)               \
+    instantiate_implicit_2d_filter(name, itype, 32,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 64,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 32, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 32, 64, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 64, 16, 2, 2)
+
+instantiate_implicit_2d_blocks(float32, float);
+instantiate_implicit_2d_blocks(float16, half);
+instantiate_implicit_2d_blocks(bfloat16, bfloat16_t); // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loader.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loader.h
new file mode 100644
index 0000000..f84a640
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loader.h
@@ -0,0 +1,6 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h"
+#include "mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h"
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h
new file mode 100644
index 0000000..d52642b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h
@@ -0,0 +1,451 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderLargeFilter {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant ImplicitGemmConv2DParams* gemm_params;
+
+  short weight_h;
+  short weight_w;
+
+  const device T* src[n_rows];
+
+  int read_n[n_rows];
+  int read_ih[n_rows];
+  int read_iw[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderLargeFilter(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        gemm_params(gemm_params_),
+        weight_h(0),
+        weight_w(0) {
+    int out_n_pixels = params->oS[0] * params->oS[1];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / out_n_pixels;
+      int hw = offset_nhw % out_n_pixels;
+      int oh = hw / params->oS[1];
+      int ow = hw % params->oS[1];
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+
+      // Adjust for flip
+      if (params->flip) {
+        ih += (params->wS[0] - 1) * params->kdil[0];
+        iw += (params->wS[1] - 1) * params->kdil[1];
+      }
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] +
+          iw * params->in_strides[2] + bj;
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Find bounds
+      int n = read_n[i];
+      int ih = read_ih[i] + weight_h * params->kdil[0];
+      int iw = read_iw[i] + weight_w * params->kdil[1];
+
+      // Read from input if in bounds
+      if ((n < params->N) && (ih >= 0 && ih < params->iS[0]) &&
+          (iw >= 0 && iw < params->iS[1])) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = src[i][j];
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    if (++weight_w < params->wS[1]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_w;
+      }
+
+      return;
+    }
+
+    weight_w = 0;
+
+    if (++weight_h < params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_h;
+      }
+
+      return;
+    }
+
+    weight_h = 0;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; i++) {
+      src[i] += gemm_params->inp_jump_c;
+    }
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderSmallFilter {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  using mask_t = short;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant ImplicitGemmConv2DParams* gemm_params;
+
+  short weight_h;
+  short weight_w;
+
+  const device T* src[n_rows];
+
+  mask_t mask_h[n_rows];
+  mask_t mask_w[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderSmallFilter(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        gemm_params(gemm_params_),
+        weight_h(0),
+        weight_w(0) {
+    int out_n_pixels = params->oS[0] * params->oS[1];
+
+    int read_n[n_rows];
+    int read_ih[n_rows];
+    int read_iw[n_rows];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / out_n_pixels;
+      int hw = offset_nhw % out_n_pixels;
+      int oh = hw / params->oS[1];
+      int ow = hw % params->oS[1];
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+
+      // Adjust for flip
+      if (params->flip) {
+        ih += (params->wS[0] - 1) * params->kdil[0];
+        iw += (params->wS[1] - 1) * params->kdil[1];
+      }
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] +
+          iw * params->in_strides[2] + bj;
+    }
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      mask_h[i] = 0;
+      mask_w[i] = 0;
+    }
+
+    for (short kh = 0; kh < params->wS[0]; kh++) {
+      short flip_h = params->flip ? params->wS[0] - kh - 1 : kh;
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; ++i) {
+        int n = read_n[i];
+        int ih = read_ih[i] + flip_h * params->kdil[0];
+
+        bool in_bounds = n < params->N && ih >= 0 && ih < params->iS[0];
+
+        mask_h[i] |= (in_bounds << kh);
+      }
+    }
+
+    for (short kw = 0; kw < params->wS[1]; kw++) {
+      short flip_w = params->flip ? params->wS[1] - kw - 1 : kw;
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; ++i) {
+        int iw = read_iw[i] + flip_w * params->kdil[1];
+
+        bool in_bounds = iw >= 0 && iw < params->iS[1];
+
+        mask_w[i] |= (in_bounds << kw);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    mask_t h_mask = mask_t(1) << weight_h;
+    mask_t w_mask = mask_t(1) << weight_w;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Read from input if in bounds
+      if ((mask_h[i] & h_mask) && (mask_w[i] & w_mask)) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = src[i][j];
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    if (++weight_w < params->wS[1]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_w;
+      }
+
+      return;
+    }
+
+    weight_w = 0;
+
+    if (++weight_h < params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_h;
+      }
+
+      return;
+    }
+
+    weight_h = 0;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; i++) {
+      src[i] += gemm_params->inp_jump_c;
+    }
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DWeightBlockLoader {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BN;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size =
+      (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4);
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Leading dimension for src
+  const int src_ld;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  const constant MLXConvParams<2>* params;
+
+  int weight_hw;
+  int weight_step;
+
+  const int read_n;
+  const bool do_read;
+
+  /* Constructor */
+  METAL_FUNC Conv2DWeightBlockLoader(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(params_->wt_strides[0]),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj),
+        params(params_),
+        weight_hw(0),
+        weight_step(params->C / params->groups),
+        read_n(offsets.y + bi),
+        do_read(read_n + n_rows * TROWS <= gemm_params_->N) {}
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    if (BN != 8 || do_read) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BN; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = src[i * src_ld + j];
+        }
+      }
+    } else {
+      for (short i = 0; i < BN; i += TROWS) {
+        if ((read_n + i) < params->O) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = src[i * src_ld + j];
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    if (++weight_hw < (params->wS[1] * params->wS[0])) {
+      src += weight_step;
+      return;
+    }
+
+    weight_hw = 0;
+
+    src += BK - (params->wS[1] * params->wS[0] - 1) * weight_step;
+  }
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h
new file mode 100644
index 0000000..2312e1c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h
@@ -0,0 +1,319 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <short n_channels_>
+struct ChannelHelper {
+  STEEL_CONST short n_channels = n_channels_;
+  STEEL_CONST short vec_size = n_channels_ <= 4 ? 4 : 8;
+  STEEL_CONST short excess = vec_size - n_channels_;
+};
+
+template <>
+struct ChannelHelper<1> {
+  STEEL_CONST short n_channels = 1;
+  STEEL_CONST short vec_size = 1;
+  STEEL_CONST short excess = 0;
+};
+
+template <>
+struct ChannelHelper<2> {
+  STEEL_CONST short n_channels = 2;
+  STEEL_CONST short vec_size = 2;
+  STEEL_CONST short excess = 0;
+};
+
+template <>
+struct ChannelHelper<3> {
+  STEEL_CONST short n_channels = 3;
+  STEEL_CONST short vec_size = 4;
+  STEEL_CONST short excess = 1;
+};
+
+template <>
+struct ChannelHelper<4> {
+  STEEL_CONST short n_channels = 4;
+  STEEL_CONST short vec_size = 4;
+  STEEL_CONST short excess = 0;
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short n_channels,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderSmallChannels {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = ChannelHelper<n_channels>::vec_size;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant ImplicitGemmConv2DParams* gemm_params;
+
+  int weight_hw;
+
+  const device T* src[n_rows];
+
+  int read_n[n_rows];
+  int read_ih[n_rows];
+  int read_iw[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderSmallChannels(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        gemm_params(gemm_params_),
+        weight_hw(thread_idx % TCOLS) {
+    int out_n_pixels = params->oS[0] * params->oS[1];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / out_n_pixels;
+      int hw = offset_nhw % out_n_pixels;
+      int oh = hw / params->oS[1];
+      int ow = hw % params->oS[1];
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] +
+          iw * params->in_strides[2];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    if (weight_hw >= params->wS[1] * params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+      return;
+    }
+
+    int wh = (weight_hw / params->wS[1]);
+    int ww = (weight_hw % params->wS[1]);
+
+    int flip_h = params->flip ? params->wS[0] - wh - 1 : wh;
+    int flip_w = params->flip ? params->wS[1] - ww - 1 : ww;
+
+    int weight_h = flip_h * params->kdil[0];
+    int weight_w = flip_w * params->kdil[1];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Find bounds
+      int n = read_n[i];
+      int ih = read_ih[i] + weight_h;
+      int iw = read_iw[i] + weight_w;
+
+      // Read from input if in bounds
+      if ((n < params->N) && (ih >= 0 && ih < params->iS[0]) &&
+          (iw >= 0 && iw < params->iS[1])) {
+        const device T* curr_src = src[i] + weight_h * params->in_strides[1] +
+            weight_w * params->in_strides[2];
+
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < n_channels; ++j) {
+          dst[is * dst_ld + j] = curr_src[j];
+        }
+
+        STEEL_PRAGMA_UNROLL
+        for (short j = n_channels; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_hw += TCOLS;
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short n_channels,
+    short tgp_padding = 0>
+struct Conv2DWeightBlockLoaderSmallChannels {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BN;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = ChannelHelper<n_channels>::vec_size;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Leading dimension for src
+  const int src_ld;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  const constant MLXConvParams<2>* params;
+
+  int weight_hw;
+
+  const int read_n;
+  const bool do_read;
+
+  /* Constructor */
+  METAL_FUNC Conv2DWeightBlockLoaderSmallChannels(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(params_->wt_strides[0]),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld),
+        params(params_),
+        weight_hw(thread_idx % TCOLS),
+        read_n(offsets.y + bi),
+        do_read(read_n + BN <= gemm_params_->N) {}
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    if (bi >= BROWS || bj >= BCOLS)
+      return;
+
+    if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+
+      return;
+    }
+
+    const device T* curr_src = src + weight_hw * (params->C / params->groups);
+
+    if (BN != 8 || do_read) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < n_channels; j++) {
+          dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+        }
+
+        STEEL_PRAGMA_UNROLL
+        for (short j = n_channels; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+    } else {
+      for (short i = 0; i < BROWS; i += TROWS) {
+        if (((read_n + i) < params->O)) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < n_channels; j++) {
+            dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+          }
+
+          STEEL_PRAGMA_UNROLL
+          for (short j = n_channels; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_hw += TCOLS;
+  }
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h
new file mode 100644
index 0000000..9b7ddc2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h
@@ -0,0 +1,381 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderGeneral {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant Conv2DGeneralJumpParams* jump_params;
+
+  const short base_wh;
+  const short base_ww;
+
+  short weight_h;
+  short weight_w;
+
+  const device T* src[n_rows];
+
+  int read_n[n_rows];
+  int read_ih[n_rows];
+  int read_iw[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderGeneral(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int4 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant Conv2DGeneralJumpParams* jump_params_,
+      const short base_wh_,
+      const short base_ww_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        jump_params(jump_params_),
+        base_wh(base_wh_),
+        base_ww(base_ww_),
+        weight_h(base_wh_),
+        weight_w(base_ww_) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / jump_params->adj_out_hw;
+      int hw = offset_nhw % jump_params->adj_out_hw;
+      int oh =
+          (hw / jump_params->adj_out_w) * jump_params->f_out_jump_h + offsets.z;
+      int ow =
+          (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + offsets.w;
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + bj;
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Find bounds
+      int n = read_n[i];
+
+      int h_flip = params->flip ? params->wS[0] - weight_h - 1 : weight_h;
+      int w_flip = params->flip ? params->wS[1] - weight_w - 1 : weight_w;
+
+      int ih_dil = read_ih[i] + h_flip * params->kdil[0];
+      int iw_dil = read_iw[i] + w_flip * params->kdil[1];
+
+      int ih = ih_dil / params->idil[0];
+      int iw = iw_dil / params->idil[1];
+
+      size_t offset = ih * params->in_strides[1] + iw * params->in_strides[2];
+
+      // Read from input if in bounds
+      if ((n < params->N) && (ih_dil >= 0 && ih < params->iS[0]) &&
+          (iw_dil >= 0 && iw < params->iS[1])) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = (src[i])[offset + j];
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void load_safe(const short remaining_k) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Find bounds
+      int n = read_n[i];
+
+      int h_flip = params->flip ? params->wS[0] - weight_h - 1 : weight_h;
+      int w_flip = params->flip ? params->wS[1] - weight_w - 1 : weight_w;
+
+      int ih_dil = read_ih[i] + h_flip * params->kdil[0];
+      int iw_dil = read_iw[i] + w_flip * params->kdil[1];
+
+      int ih = ih_dil / params->idil[0];
+      int iw = iw_dil / params->idil[1];
+
+      size_t offset = ih * params->in_strides[1] + iw * params->in_strides[2];
+
+      // Read from input if in bounds
+      if ((n < params->N) && (ih_dil >= 0 && ih < params->iS[0]) &&
+          (iw_dil >= 0 && iw < params->iS[1])) {
+        if (bj + vec_size <= remaining_k) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; ++j) {
+            dst[is * dst_ld + j] = (src[i])[offset + j];
+          }
+        } else {
+          for (short j = 0; j < vec_size; ++j) {
+            if (bj + j < remaining_k) {
+              dst[is * dst_ld + j] = (src[i])[offset + j];
+            } else {
+              dst[is * dst_ld + j] = T(0);
+            }
+          }
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_w += jump_params->f_wgt_jump_w;
+    if (weight_w < params->wS[1]) {
+      return;
+    }
+
+    weight_w = base_ww;
+
+    weight_h += jump_params->f_wgt_jump_h;
+    if (weight_h < params->wS[0]) {
+      return;
+    }
+
+    weight_h = base_wh;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; i++) {
+      src[i] += BK;
+    }
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DWeightBlockLoaderGeneral {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BN;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size =
+      (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4);
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Leading dimension for src
+  const int src_ld;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  const constant MLXConvParams<2>* params;
+  const constant Conv2DGeneralJumpParams* jump_params;
+
+  const short base_wh;
+  const short base_ww;
+
+  short weight_h;
+  short weight_w;
+
+  const int start_row;
+
+  /* Constructor */
+  METAL_FUNC Conv2DWeightBlockLoaderGeneral(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant Conv2DGeneralJumpParams* jump_params_,
+      const short base_wh_,
+      const short base_ww_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(params_->wt_strides[0]),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj),
+        params(params_),
+        jump_params(jump_params_),
+        base_wh(base_wh_),
+        base_ww(base_ww_),
+        weight_h(base_wh_),
+        weight_w(base_ww_),
+        start_row(offsets.y + bi) {}
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    const device T* curr_src = src + weight_h * params->wt_strides[1] +
+        weight_w * params->wt_strides[2];
+
+    if ((start_row + BN <= params->O)) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BN; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+        }
+      }
+    } else {
+      for (short i = 0; i < BN; i += TROWS) {
+        if ((start_row + i) < params->O) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void load_safe(const short remaining_k) const {
+    const device T* curr_src = src + weight_h * params->wt_strides[1] +
+        weight_w * params->wt_strides[2];
+
+    if ((start_row + BN <= params->O)) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BN; i += TROWS) {
+        if (bj + vec_size <= remaining_k) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+          }
+        } else {
+          for (short j = 0; j < vec_size; j++) {
+            if (bj + j < remaining_k) {
+              dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+            } else {
+              dst[i * dst_ld + j] = T(0);
+            }
+          }
+        }
+      }
+    } else {
+      for (short i = 0; i < BN; i += TROWS) {
+        if ((start_row + i) < params->O) {
+          if (bj + vec_size <= remaining_k) {
+            STEEL_PRAGMA_UNROLL
+            for (short j = 0; j < vec_size; j++) {
+              dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+            }
+          } else {
+            for (short j = 0; j < vec_size; j++) {
+              if (bj + j < remaining_k) {
+                dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+              } else {
+                dst[i * dst_ld + j] = T(0);
+              }
+            }
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_w += jump_params->f_wgt_jump_w;
+    if (weight_w < params->wS[1]) {
+      return;
+    }
+
+    weight_w = base_ww;
+
+    weight_h += jump_params->f_wgt_jump_h;
+    if (weight_h < params->wS[0]) {
+      return;
+    }
+
+    weight_h = base_wh;
+
+    src += BK;
+  }
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/params.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/params.h
new file mode 100644
index 0000000..61b8474
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/conv/params.h
@@ -0,0 +1,62 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+template <int NDIM>
+struct MLXConvParams {
+  const int N; // Batch size
+  const int C; // In channels
+  const int O; // Out channels
+  const int iS[NDIM]; // Input spatial dim
+  const int wS[NDIM]; // Weight spatial dim
+  const int oS[NDIM]; // Output spatial dim
+  const int str[NDIM]; // Kernel strides
+  const int pad[NDIM]; // Input padding
+  const int kdil[NDIM]; // Kernel dilation
+  const int idil[NDIM]; // Input dilation
+  const int64_t in_strides[NDIM + 2]; // In strides
+  const int64_t wt_strides[NDIM + 2]; // Wt strides
+  const int64_t out_strides[NDIM + 2]; // Out strides
+  const int groups; // Input channel groups
+  const bool flip;
+};
+
+namespace mlx {
+namespace steel {
+
+struct ImplicitGemmConv2DParams {
+  const int M;
+  const int N;
+  const int K;
+
+  const int gemm_k_iterations;
+
+  const int inp_jump_w;
+  const int inp_jump_h;
+  const int inp_jump_c;
+
+  const int tiles_n;
+  const int tiles_m;
+  const int swizzle_log;
+};
+
+struct Conv2DGeneralJumpParams {
+  const int f_wgt_jump_h;
+  const int f_wgt_jump_w;
+
+  const int f_out_jump_h;
+  const int f_out_jump_w;
+
+  const int adj_out_h;
+  const int adj_out_w;
+  const int adj_out_hw;
+  const int adj_implicit_m;
+};
+
+struct Conv2DGeneralBaseInfo {
+  int weight_base;
+  int weight_size;
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/defines.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/defines.h
new file mode 100644
index 0000000..f5657ee
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/defines.h
@@ -0,0 +1,7 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#define STEEL_CONST static constant constexpr const
+#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
+#define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)")
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/gemm.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/gemm.h
new file mode 100644
index 0000000..bbe1d96
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/gemm.h
@@ -0,0 +1,295 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/gemm/loader.h"
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+#include "mlx/backend/metal/kernels/steel/gemm/params.h"
+#include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernel class
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <bool M_aligned, bool N_aligned, bool K_aligned>
+struct LoopAlignment {};
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    bool MN_aligned,
+    bool K_aligned,
+    typename AccumType = typename AccumHelper<T>::accum_type,
+    typename Epilogue = TransformNone<U, AccumType>>
+struct GEMMKernel {
+  STEEL_CONST short tgp_padding_a = 16 / sizeof(T);
+  STEEL_CONST short tgp_padding_b = 16 / sizeof(T);
+  STEEL_CONST short tgp_mem_size_a =
+      transpose_a ? BK * (BM + tgp_padding_a) : BM * (BK + tgp_padding_a);
+  STEEL_CONST short tgp_mem_size_b =
+      transpose_b ? BN * (BK + tgp_padding_b) : BK * (BN + tgp_padding_b);
+  STEEL_CONST short tgp_mem_size = tgp_mem_size_a + tgp_mem_size_b;
+
+  STEEL_CONST short tgp_size = WM * WN * 32;
+
+  using loader_a_t = BlockLoader<
+      T,
+      transpose_a ? BK : BM,
+      transpose_a ? BM : BK,
+      transpose_a ? BM + tgp_padding_a : BK + tgp_padding_a,
+      !transpose_a,
+      tgp_size>;
+  using loader_b_t = BlockLoader<
+      T,
+      transpose_b ? BN : BK,
+      transpose_b ? BK : BN,
+      transpose_b ? BK + tgp_padding_b : BN + tgp_padding_b,
+      transpose_b,
+      tgp_size>;
+  using mma_t = BlockMMA<
+      T,
+      U,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      transpose_a ? BM + tgp_padding_a : BK + tgp_padding_a,
+      transpose_b ? BK + tgp_padding_b : BN + tgp_padding_b,
+      AccumType,
+      Epilogue>;
+
+  /* Main kernel function */
+  template <bool M_aligned, bool N_aligned, bool K_aligned_>
+  static METAL_FUNC void gemm_loop(
+      threadgroup T* As [[threadgroup(0)]],
+      threadgroup T* Bs [[threadgroup(1)]],
+      const int gemm_k_iterations,
+      thread loader_a_t& loader_a,
+      thread loader_b_t& loader_b,
+      thread mma_t& mma_op,
+      thread const short& tgp_bm,
+      thread const short& tgp_bn,
+      thread const short& lbk,
+      LoopAlignment<M_aligned, N_aligned, K_aligned_> l = {}) {
+    // Appease the compiler
+    (void)l;
+
+    short2 tile_dims_A = transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm);
+
+    short2 tile_dims_B = transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK);
+
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      // Load elements into threadgroup
+      if (M_aligned) {
+        loader_a.load_unsafe();
+      } else {
+        loader_a.load_safe(tile_dims_A);
+      }
+
+      if (N_aligned) {
+        loader_b.load_unsafe();
+      } else {
+        loader_b.load_safe(tile_dims_B);
+      }
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+
+    if (!K_aligned_) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      short2 tile_dims_A_last =
+          transpose_a ? short2(tgp_bm, lbk) : short2(lbk, tgp_bm);
+      short2 tile_dims_B_last =
+          transpose_b ? short2(lbk, tgp_bn) : short2(tgp_bn, lbk);
+
+      loader_a.load_safe(tile_dims_A_last);
+      loader_b.load_safe(tile_dims_B_last);
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      mma_op.mma(As, Bs);
+    }
+  }
+
+  /* Main kernel function */
+  static METAL_FUNC void run(
+      const device T* A [[buffer(0)]],
+      const device T* B [[buffer(1)]],
+      device U* D [[buffer(2)]],
+      const constant GEMMParams* params [[buffer(3)]],
+      threadgroup T* As [[threadgroup(0)]],
+      threadgroup T* Bs [[threadgroup(1)]],
+      uint simd_lane_id [[thread_index_in_simdgroup]],
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]]) {
+    // Pacifying compiler
+    (void)lid;
+
+    const int tid_y = ((tid.y) << params->swizzle_log) +
+        ((tid.x) & ((1 << params->swizzle_log) - 1));
+    const int tid_x = (tid.x) >> params->swizzle_log;
+
+    if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+      return;
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Find block in A, B, C
+    const int c_row = tid_y * BM;
+    const int c_col = tid_x * BN;
+    const size_t c_row_long = size_t(c_row);
+    const size_t c_col_long = size_t(c_col);
+
+    A += transpose_a ? c_row_long : c_row_long * params->lda;
+    B += transpose_b ? c_col_long * params->ldb : c_col_long;
+    D += c_row_long * params->ldd + c_col_long;
+
+    // Prepare threadgroup loading operations
+    thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+    thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+    // Prepare threadgroup mma operation
+    thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+    int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+    ///////////////////////////////////////////////////////////////////////////////
+    // MNK aligned loop
+    if (MN_aligned) {
+      for (int k = 0; k < gemm_k_iterations; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+
+      threadgroup_barrier(mem_flags::mem_none);
+
+      // Loop tail
+      if (!K_aligned) {
+        int lbk = params->K - params->gemm_k_iterations_aligned * BK;
+        short2 tile_dims_A = transpose_a ? short2(BM, lbk) : short2(lbk, BM);
+        short2 tile_dims_B = transpose_b ? short2(lbk, BN) : short2(BN, lbk);
+
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        mma_op.mma(As, Bs);
+      }
+
+      // Store results to device memory
+      mma_op.store_result(D, params->ldd);
+      return;
+
+    }
+    ///////////////////////////////////////////////////////////////////////////////
+    // MN unaligned loop
+    else { // Loop over K - unaligned case
+      short tgp_bm = min(BM, params->M - c_row);
+      short tgp_bn = min(BN, params->N - c_col);
+      short leftover_bk = params->K - params->gemm_k_iterations_aligned * BK;
+
+      if (tgp_bm == BM && tgp_bn == BN) {
+        gemm_loop<true, true, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result(D, params->ldd);
+        return;
+
+      } else if (tgp_bn == BN) {
+        gemm_loop<false, true, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+
+      } else if (tgp_bm == BM) {
+        gemm_loop<true, false, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+
+      } else {
+        gemm_loop<false, false, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+      }
+    }
+  }
+};
+
+} // namespace steel
+} // namespace mlx
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h
new file mode 100644
index 0000000..26645cf
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h
@@ -0,0 +1,157 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/gemm/nax.h"
+#include "mlx/backend/metal/kernels/steel/gemm/params.h"
+#include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace metal;
+
+namespace mlx::steel {
+
+template <
+    typename T,
+    short SM,
+    short SN,
+    short SK,
+    short BK,
+    bool transpose_a,
+    bool transpose_b,
+    bool kAlignedM,
+    bool kAlignedN,
+    bool kAlignedK,
+    short UM,
+    short UN,
+    short UK,
+    typename AccumType = float>
+auto gemm_loop(
+    const device T* A,
+    const device T* B,
+    int lda,
+    int ldb,
+    int K,
+    int gemm_k_iterations_aligned,
+    const short sgp_sm,
+    const short sgp_sn) {
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  constexpr int RA = transpose_a ? TK : TM;
+  constexpr int CA = transpose_a ? TM : TK;
+
+  constexpr int RB = transpose_b ? TN : TK;
+  constexpr int CB = transpose_b ? TK : TN;
+
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+  using ASubTile =
+      NAXSubTile<T, (transpose_a ? UK : UM), (transpose_a ? UM : UK)>;
+  using BSubTile =
+      NAXSubTile<T, (transpose_b ? UN : UK), (transpose_b ? UK : UN)>;
+
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+  Dtile.clear();
+
+  int gemm_k_iterations_ = gemm_k_iterations_aligned;
+
+  STEEL_PRAGMA_NO_UNROLL
+  for (int kk0 = 0; kk0 < gemm_k_iterations_; kk0++) {
+    threadgroup_barrier(mem_flags::mem_none);
+
+    STEEL_PRAGMA_NO_UNROLL
+    for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+      NAXTile<T, RA, CA, ASubTile> Atile;
+      NAXTile<T, RB, CB, BSubTile> Btile;
+      const int k = kk1;
+
+      volatile int compiler_barrier;
+
+      const int A_offset = transpose_a ? k * lda : k;
+      const int B_offset = transpose_b ? k : k * ldb;
+
+      if constexpr (kAlignedM) {
+        Atile.load(A + A_offset, lda);
+      } else {
+        const short rmax = transpose_a ? SK : sgp_sm;
+        const short cmax = transpose_a ? sgp_sm : SK;
+        Atile.load_safe(A + A_offset, lda, short2(cmax, rmax));
+      }
+
+      if constexpr (kAlignedN) {
+        Btile.load(B + B_offset, ldb);
+      } else {
+        const short rmax = transpose_b ? sgp_sn : SK;
+        const short cmax = transpose_b ? SK : sgp_sn;
+        Btile.load_safe(B + B_offset, ldb, short2(cmax, rmax));
+      }
+
+      tile_matmad_nax(
+          Dtile,
+          Atile,
+          metal::bool_constant<transpose_a>{},
+          Btile,
+          metal::bool_constant<transpose_b>{});
+
+      (void)compiler_barrier;
+    }
+
+    A += transpose_a ? (BK * lda) : BK;
+    B += transpose_b ? BK : (BK * ldb);
+  }
+
+  if constexpr (!kAlignedK) {
+    simdgroup_barrier(mem_flags::mem_none);
+
+    const short rem_bk = K - gemm_k_iterations_ * BK;
+
+    STEEL_PRAGMA_NO_UNROLL
+    for (int kk1 = 0; kk1 < rem_bk; kk1 += SK) {
+      NAXTile<T, 1, 1, ASubTile> Atile;
+      NAXTile<T, 1, 1, BSubTile> Btile;
+
+      STEEL_PRAGMA_UNROLL
+      for (int mm = 0; mm < TM; mm++) {
+        STEEL_PRAGMA_UNROLL
+        for (int nn = 0; nn < TN; nn++) {
+          STEEL_PRAGMA_UNROLL
+          for (int kk = 0; kk < TK; kk++) {
+            const int m = mm * UM;
+            const int n = nn * UN;
+            const int k = kk1 + kk * UK;
+            const short psk = max(0, rem_bk - k);
+
+            const int A_offset = transpose_a ? (m + k * lda) : (m * lda + k);
+            const int B_offset = transpose_b ? (k + n * ldb) : (k * ldb + n);
+
+            {
+              const short psm = kAlignedM ? SM : max(0, sgp_sm - m);
+              const short rmax = transpose_a ? psk : psm;
+              const short cmax = transpose_a ? psm : psk;
+              Atile.load_safe(A + A_offset, lda, short2(cmax, rmax));
+            }
+
+            {
+              const short psn = kAlignedN ? SN : max(0, sgp_sn - n);
+              const short rmax = transpose_b ? psn : psk;
+              const short cmax = transpose_b ? psk : psn;
+              Btile.load_safe(B + B_offset, ldb, short2(cmax, rmax));
+            }
+
+            subtile_matmad_nax(
+                Dtile.subtile_at(mm, nn),
+                Atile.subtile_at(0, 0),
+                metal::bool_constant<transpose_a>{},
+                Btile.subtile_at(0, 0),
+                metal::bool_constant<transpose_b>{});
+          }
+        }
+      }
+    }
+  }
+
+  return Dtile;
+}
+
+} // namespace mlx::steel
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.h
new file mode 100644
index 0000000..8583087
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.h
@@ -0,0 +1,346 @@
+// Copyright © 2024 Apple Inc.
+
+using namespace mlx::steel;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernels
+///////////////////////////////////////////////////////////////////////////////
+
+constant bool has_batch [[function_constant(10)]];
+
+constant bool use_out_source [[function_constant(100)]];
+constant bool do_axpby [[function_constant(110)]];
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+// clang-format off
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gemm(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    const device T* C [[buffer(2), function_constant(use_out_source)]],
+    device T* D [[buffer(3)]],
+    const constant GEMMParams* params [[buffer(4)]],
+    const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]],
+    const constant int* batch_shape [[buffer(6), function_constant(has_batch)]],
+    const constant int64_t* batch_strides [[buffer(7), function_constant(has_batch)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) { // clang-format on
+  // Pacifying compiler
+  (void)lid;
+
+  using gemm_kernel = GEMMKernel<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      true,
+      true,
+      AccumType>;
+
+  using loader_a_t = typename gemm_kernel::loader_a_t;
+  using loader_b_t = typename gemm_kernel::loader_b_t;
+  using mma_t = typename gemm_kernel::mma_t;
+
+  // Find block
+  const int tid_y = ((tid.y) << params->swizzle_log) +
+      ((tid.x) & ((1 << params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> params->swizzle_log;
+
+  // Exit early if out of bounds
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  // Adjust for batch
+  if (has_batch) {
+    const constant auto* A_bstrides = batch_strides;
+    const constant auto* B_bstrides = batch_strides + params->batch_ndim;
+
+    ulong2 batch_offsets = elem_to_loc_broadcast(
+        tid.z, batch_shape, A_bstrides, B_bstrides, params->batch_ndim);
+
+    A += batch_offsets.x;
+    B += batch_offsets.y;
+
+    if (use_out_source) {
+      const constant auto* C_bstrides = B_bstrides + params->batch_ndim;
+      C += elem_to_loc(tid.z, batch_shape, C_bstrides, params->batch_ndim);
+    }
+  } else {
+    A += params->batch_stride_a * tid.z;
+    B += params->batch_stride_b * tid.z;
+
+    if (use_out_source) {
+      C += addmm_params->batch_stride_c * tid.z;
+    }
+  }
+
+  D += params->batch_stride_d * tid.z;
+
+  // Prepare threadgroup memory
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Find block in A, B, C
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  D += c_row_long * params->ldd + c_col_long;
+
+  if (use_out_source) {
+    C += c_row_long * addmm_params->ldc + c_col_long * addmm_params->fdc;
+  }
+
+  // Prepare threadgroup mma operation
+  thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup loading operations
+  thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+  thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col));
+
+  // Prepare iterations
+  int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+  // Do unaligned K iterations first
+  if (!align_K) {
+    const int k_last = params->gemm_k_iterations_aligned * BK;
+    const int k_remain = params->K - k_last;
+    const size_t k_jump_a =
+        transpose_a ? params->lda * size_t(k_last) : size_t(k_last);
+    const size_t k_jump_b =
+        transpose_b ? size_t(k_last) : params->ldb * size_t(k_last);
+
+    // Move loader source ahead to end
+    loader_a.src += k_jump_a;
+    loader_b.src += k_jump_b;
+
+    // Load tile
+    const short2 tile_dims_A =
+        transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+    const short2 tile_dims_B =
+        transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+    loader_a.load_safe(tile_dims_A);
+    loader_b.load_safe(tile_dims_B);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Do matmul
+    mma_op.mma(As, Bs);
+
+    // Reset source back to start
+    loader_a.src -= k_jump_a;
+    loader_b.src -= k_jump_b;
+  }
+
+  const TransformAdd<AccumType, AccumType> epilogue_op_add(
+      addmm_params->alpha, addmm_params->beta);
+  const TransformAxpby<AccumType, AccumType> epilogue_op_axpby(
+      addmm_params->alpha, addmm_params->beta);
+
+  ///////////////////////////////////////////////////////////////////////////////
+  // MNK aligned loop
+  if (align_M && align_N) {
+    // Do gemm
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      // Load elements into threadgroup
+      loader_a.load_unsafe();
+      loader_b.load_unsafe();
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Do epilogue
+    if (use_out_source) {
+      if (do_axpby) {
+        mma_op.apply_epilogue(
+            C, addmm_params->ldc, addmm_params->fdc, epilogue_op_axpby);
+      } else {
+        mma_op.apply_epilogue(
+            C, addmm_params->ldc, addmm_params->fdc, epilogue_op_add);
+      }
+    }
+
+    // Store results to device memory
+    return mma_op.store_result(D, params->ldd);
+
+  }
+  ///////////////////////////////////////////////////////////////////////////////
+  // MN unaligned loop
+  else { // Loop over K - unaligned case
+    const int leftover_bk = 0;
+
+    if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) {
+      // Do gemm
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          leftover_bk,
+          LoopAlignment<true, true, true>{});
+
+      // Do epilogue
+      if (use_out_source) {
+        if (do_axpby) {
+          mma_op.apply_epilogue(
+              C, addmm_params->ldc, addmm_params->fdc, epilogue_op_axpby);
+        } else {
+          mma_op.apply_epilogue(
+              C, addmm_params->ldc, addmm_params->fdc, epilogue_op_add);
+        }
+      }
+
+      // Store results to device memory
+      return mma_op.store_result(D, params->ldd);
+
+    } else if (align_N || tgp_bn == BN) {
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          leftover_bk,
+          LoopAlignment<false, true, true>{});
+
+      // Do epilogue
+      if (use_out_source) {
+        if (do_axpby) {
+          mma_op.apply_epilogue_safe(
+              C,
+              addmm_params->ldc,
+              addmm_params->fdc,
+              short2(tgp_bn, tgp_bm),
+              epilogue_op_axpby);
+        } else {
+          mma_op.apply_epilogue_safe(
+              C,
+              addmm_params->ldc,
+              addmm_params->fdc,
+              short2(tgp_bn, tgp_bm),
+              epilogue_op_add);
+        }
+      }
+
+      // Store results to device memory
+      return mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+
+    } else if (align_M || tgp_bm == BM) {
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          leftover_bk,
+          LoopAlignment<true, false, true>{});
+
+      // Do epilogue
+      if (use_out_source) {
+        if (do_axpby) {
+          mma_op.apply_epilogue_safe(
+              C,
+              addmm_params->ldc,
+              addmm_params->fdc,
+              short2(tgp_bn, tgp_bm),
+              epilogue_op_axpby);
+        } else {
+          mma_op.apply_epilogue_safe(
+              C,
+              addmm_params->ldc,
+              addmm_params->fdc,
+              short2(tgp_bn, tgp_bm),
+              epilogue_op_add);
+        }
+      }
+
+      // Store results to device memory
+      return mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+
+    } else {
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          leftover_bk,
+          LoopAlignment<false, false, true>{});
+
+      // Do epilogue
+      if (use_out_source) {
+        if (do_axpby) {
+          mma_op.apply_epilogue_safe(
+              C,
+              addmm_params->ldc,
+              addmm_params->fdc,
+              short2(tgp_bn, tgp_bm),
+              epilogue_op_axpby);
+        } else {
+          mma_op.apply_epilogue_safe(
+              C,
+              addmm_params->ldc,
+              addmm_params->fdc,
+              short2(tgp_bn, tgp_bm),
+              epilogue_op_add);
+        }
+      }
+
+      // Store results to device memory
+      return mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.metal
new file mode 100644
index 0000000..5549580
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.metal
@@ -0,0 +1,34 @@
+// Copyright © 2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.h"
+
+#define instantiate_gemm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                                             \
+      "steel_gemm_fused_" #tname "_"  #iname "_" #oname                                           \
+      "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn,                                          \
+  gemm, itype, bm, bn, bk, wm, wn, trans_a, trans_b, float)
+
+#define instantiate_gemm_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gemm_shapes_helper(iname, itype, oname, otype) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 1, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 32, 32, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 64, 16, 1, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 32, 16, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 32,  8, 4, 1)
+
+instantiate_gemm_shapes_helper(float16, half, float16, half);
+instantiate_gemm_shapes_helper(bfloat16, bfloat16_t, bfloat16, bfloat16_t);
+
+instantiate_gemm_shapes_helper(float32, float, float32, float);
+instantiate_gemm_shapes_helper(complex64, complex64_t, complex64, complex64_t);
+// clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h
new file mode 100644
index 0000000..4ff9260
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h
@@ -0,0 +1,219 @@
+// Copyright © 2025 Apple Inc.
+
+using namespace mlx::steel;
+
+constant bool has_batch [[function_constant(10)]];
+
+constant bool use_out_source [[function_constant(100)]];
+constant bool do_axpby [[function_constant(110)]];
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+// clang-format off
+template <
+    bool kAlignedM,
+    bool kAlignedN,
+    typename NAXTile_t,
+    typename T>
+void gemm_epilogue(
+    thread NAXTile_t& Dtile,
+    const device T* C,
+    const constant GEMMParams* params,
+    const constant GEMMAddMMParams* addmm_params,
+    const short sgp_sm, 
+    const short sgp_sn) { // clang-format on
+
+  (void)params;
+
+  constexpr short UM = NAXTile_t::kSubTileRows;
+  constexpr short UN = NAXTile_t::kSubTileCols;
+  using CSubTile = NAXSubTile<T, UM, UN>;
+
+  using V = typename NAXTile_t::elem_type;
+
+  constexpr short TM = NAXTile_t::kTileRows;
+  constexpr short TN = NAXTile_t::kTileCols;
+  constexpr short kElemsPerSubTile = NAXTile_t::kElemsPerSubTile;
+
+  STEEL_PRAGMA_UNROLL
+  for (short mm = 0; mm < TM; mm++) {
+    STEEL_PRAGMA_UNROLL
+    for (short nn = 0; nn < TN; nn++) {
+      const short m = mm * UM;
+      const short n = nn * UN;
+
+      CSubTile CTile;
+
+      if constexpr (kAlignedM && kAlignedN) {
+        CTile.load(C, addmm_params->ldc, addmm_params->fdc, m, n);
+      } else {
+        CTile.load_safe(
+            C, addmm_params->ldc, addmm_params->fdc, sgp_sm, sgp_sn, m, n);
+      }
+
+      auto delems = Dtile.subtile_at(mm, nn).elems();
+      auto celems = CTile.elems();
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < kElemsPerSubTile; i++) {
+        if (do_axpby) {
+          delems[i] = addmm_params->alpha * delems[i] +
+              addmm_params->beta * static_cast<V>(celems[i]);
+        } else {
+          delems[i] += static_cast<V>(celems[i]);
+        }
+      }
+    }
+  }
+}
+
+// clang-format off
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gemm(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    const device T* C [[buffer(2), function_constant(use_out_source)]],
+    device T* D [[buffer(3)]],
+    const constant GEMMParams* params [[buffer(4)]],
+    const constant GEMMAddMMParams* addmm_params [[buffer(5), function_constant(use_out_source)]],
+    const constant int* batch_shape [[buffer(6), function_constant(has_batch)]],
+    const constant int64_t* batch_strides [[buffer(7), function_constant(has_batch)]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on
+  // Find block
+  const int tid_y = ((tid.y) << params->swizzle_log) +
+      ((tid.x) & ((1 << params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> params->swizzle_log;
+
+  // Exit early if out of bounds
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  // Adjust for batch
+  if (has_batch) {
+    const constant auto* A_bstrides = batch_strides;
+    const constant auto* B_bstrides = batch_strides + params->batch_ndim;
+
+    ulong2 batch_offsets = elem_to_loc_broadcast(
+        tid.z, batch_shape, A_bstrides, B_bstrides, params->batch_ndim);
+
+    A += batch_offsets.x;
+    B += batch_offsets.y;
+
+    if (use_out_source) {
+      const constant auto* C_bstrides = B_bstrides + params->batch_ndim;
+      C += elem_to_loc(tid.z, batch_shape, C_bstrides, params->batch_ndim);
+    }
+  } else {
+    A += params->batch_stride_a * tid.z;
+    B += params->batch_stride_b * tid.z;
+
+    if (use_out_source) {
+      C += addmm_params->batch_stride_c * tid.z;
+    }
+  }
+
+  D += params->batch_stride_d * tid.z;
+
+  // Prepare threadgroup memory
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Find block in A, B, C
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  D += c_row_long * params->ldd + c_col_long;
+
+  if (use_out_source) {
+    C += c_row_long * addmm_params->ldc + c_col_long * addmm_params->fdc;
+  }
+
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+
+  const short tm = SM * (simd_group_id / WN);
+  const short tn = SN * (simd_group_id % WN);
+
+  const int sgp_sm_int =
+      align_M ? int(SM) : min(int(SM), params->M - (c_row + tm));
+  const short sgp_sm = short(sgp_sm_int);
+  const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
+
+  const int sgp_sn_int =
+      align_N ? int(SN) : min(int(SN), params->N - (c_col + tn));
+  const short sgp_sn = short(sgp_sn_int);
+  const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN);
+
+  A += transpose_a ? tm : (tm * params->lda);
+  B += transpose_b ? (tn * params->ldb) : tn;
+  D += tm * params->ldd + tn;
+
+  if (use_out_source) {
+    C += tm * addmm_params->ldc + tn * addmm_params->fdc;
+  }
+
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  dispatch_bool(align_K, [&](auto kAlignedK) {
+    dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+      dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) {
+        Dtile = gemm_loop<
+            T,
+            SM,
+            SN,
+            SK,
+            BK,
+            transpose_a,
+            transpose_b,
+            kAlignedM.value,
+            kAlignedN.value,
+            kAlignedK.value,
+            UM,
+            UN,
+            UK,
+            AccumType>(
+            A,
+            B,
+            params->lda,
+            params->ldb,
+            params->K,
+            params->gemm_k_iterations_aligned,
+            sgp_sm,
+            sgp_sn);
+        if (use_out_source) {
+          gemm_epilogue<kAlignedM.value, kAlignedN.value>(
+              Dtile, C, params, addmm_params, sgp_sm, sgp_sn);
+        }
+        if constexpr (kAlignedM && kAlignedN) {
+          Dtile.store(D, int(params->ldd));
+        } else {
+          Dtile.store_safe(D, int(params->ldd), short2(sgp_sn, sgp_sm));
+        }
+      });
+    });
+  });
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal
new file mode 100644
index 0000000..366f125
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal
@@ -0,0 +1,30 @@
+// Copyright © 2025 Apple Inc.
+
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h"
+
+// clang-format off
+#define instantiate_gemm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                                             \
+      "steel_gemm_fused_nax_" #tname "_"  #iname "_" #oname                                       \
+      "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn,                                          \
+  gemm, itype, bm, bn, bk, wm, wn, trans_a, trans_b, float)
+
+#define instantiate_gemm_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gemm_shapes_helper(iname, itype, oname, otype) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype,  64,  64, 256, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 512, 4, 4)
+
+instantiate_gemm_shapes_helper(float16, half, float16, half);
+instantiate_gemm_shapes_helper(bfloat16, bfloat, bfloat16, bfloat);
+instantiate_gemm_shapes_helper(float32, float, float32, float);
+// clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h
new file mode 100644
index 0000000..4c055e6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h
@@ -0,0 +1,459 @@
+// Copyright © 2024 Apple Inc.
+
+using namespace mlx::steel;
+
+constant bool has_batch [[function_constant(10)]];
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    const device uint32_t* rhs_indices [[buffer(2)]],
+    device T* C [[buffer(3)]],
+    const constant GEMMParams* params [[buffer(4)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  using gemm_kernel = GEMMKernel<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      true,
+      true,
+      AccumType>;
+
+  using loader_a_t = typename gemm_kernel::loader_a_t;
+  using loader_b_t = typename gemm_kernel::loader_b_t;
+  using mma_t = typename gemm_kernel::mma_t;
+
+  if (params->tiles_n <= static_cast<int>(tid.x) ||
+      params->tiles_m <= static_cast<int>(tid.y)) {
+    return;
+  }
+
+  // Prepare threadgroup memory
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  // Find the block in A, B, C
+  const int c_row = tid.y * BM;
+  const int c_col = tid.x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col));
+
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  C += c_row_long * params->ldd + c_col_long;
+
+  // Do as many matmuls as necessary
+  uint32_t index;
+  short offset;
+  uint32_t index_next = rhs_indices[c_row];
+  short offset_next = 0;
+  int n = 0;
+  while (n < tgp_bm) {
+    n++;
+    offset = offset_next;
+    index = index_next;
+    offset_next = tgp_bm;
+    for (; n < tgp_bm; n++) {
+      if (rhs_indices[c_row + n] != index) {
+        offset_next = n;
+        index_next = rhs_indices[c_row + n];
+        break;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Prepare threadgroup mma operation
+    thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+    // Prepare threadgroup loading operations
+    thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+    thread loader_b_t loader_b(
+        B + index * params->batch_stride_b,
+        params->ldb,
+        Bs,
+        simd_group_id,
+        simd_lane_id);
+
+    // Prepare iterations
+    const int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+    // Do unaligned K iterations first
+    if (!align_K) {
+      const int k_last = params->gemm_k_iterations_aligned * BK;
+      const int k_remain = params->K - k_last;
+      const size_t k_jump_a =
+          transpose_a ? params->lda * size_t(k_last) : size_t(k_last);
+      const size_t k_jump_b =
+          transpose_b ? size_t(k_last) : params->ldb * size_t(k_last);
+
+      // Move loader source ahead to end
+      loader_a.src += k_jump_a;
+      loader_b.src += k_jump_b;
+
+      // Load tile
+      const short2 tile_dims_A =
+          transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+      const short2 tile_dims_B =
+          transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+      loader_a.load_safe(tile_dims_A);
+      loader_b.load_safe(tile_dims_B);
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Do matmul
+      mma_op.mma(As, Bs);
+
+      // Reset source back to start
+      loader_a.src -= k_jump_a;
+      loader_b.src -= k_jump_b;
+    }
+
+    // Matrix level aligned never check
+    if (align_M && align_N) {
+      for (int k = 0; k < gemm_k_iterations; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+
+      // Store results to device memory
+      if (offset_next - offset == BM) {
+        mma_op.store_result(C, params->ldd);
+      } else {
+        mma_op.store_result_slice(
+            C, params->ldd, short2(0, offset), short2(BN, offset_next));
+      }
+    } else {
+      const short lbk = 0;
+
+      // Tile aligned don't check
+      if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) {
+        gemm_kernel::gemm_loop(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            lbk,
+            LoopAlignment<true, true, true>{});
+        if (offset_next - offset == BM) {
+          mma_op.store_result(C, params->ldd);
+        } else {
+          mma_op.store_result_slice(
+              C, params->ldd, short2(0, offset), short2(BN, offset_next));
+        }
+      }
+
+      // Tile partially aligned check rows
+      else if (align_N || tgp_bn == BN) {
+        gemm_kernel::gemm_loop(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            lbk,
+            LoopAlignment<false, true, true>{});
+        mma_op.store_result_slice(
+            C, params->ldd, short2(0, offset), short2(BN, offset_next));
+      }
+
+      // Tile partially aligned check cols
+      else if (align_M || tgp_bm == BM) {
+        gemm_kernel::gemm_loop(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            lbk,
+            LoopAlignment<true, false, true>{});
+        mma_op.store_result_slice(
+            C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next));
+      }
+
+      // Nothing aligned so check both rows and cols
+      else {
+        gemm_kernel::gemm_loop(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            lbk,
+            LoopAlignment<false, false, true>{});
+        mma_op.store_result_slice(
+            C, params->ldd, short2(0, offset), short2(tgp_bn, offset_next));
+      }
+    }
+  }
+}
+
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    const device uint32_t* lhs_indices [[buffer(2)]],
+    const device uint32_t* rhs_indices [[buffer(3)]],
+    device T* C [[buffer(4)]],
+    const constant GEMMParams* params [[buffer(5)]],
+    const constant int* indices_shape [[buffer(6)]],
+    const constant int64_t* lhs_strides [[buffer(7)]],
+    const constant int64_t* rhs_strides [[buffer(8)]],
+    const constant int& batch_ndim_a [[buffer(9)]],
+    const constant int* batch_shape_a [[buffer(10)]],
+    const constant int64_t* batch_strides_a [[buffer(11)]],
+    const constant int& batch_ndim_b [[buffer(12)]],
+    const constant int* batch_shape_b [[buffer(13)]],
+    const constant int64_t* batch_strides_b [[buffer(14)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  using gemm_kernel = GEMMKernel<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      true,
+      true,
+      AccumType>;
+
+  using loader_a_t = typename gemm_kernel::loader_a_t;
+  using loader_b_t = typename gemm_kernel::loader_b_t;
+  using mma_t = typename gemm_kernel::mma_t;
+
+  if (params->tiles_n <= static_cast<int>(tid.x) ||
+      params->tiles_m <= static_cast<int>(tid.y)) {
+    return;
+  }
+
+  // Move A and B to the locations pointed by lhs_indices and rhs_indices.
+  uint32_t indx_A, indx_B;
+  if (has_batch) {
+    ulong2 indices_offsets = elem_to_loc_broadcast(
+        tid.z, indices_shape, lhs_strides, rhs_strides, params->batch_ndim);
+    indx_A = lhs_indices[indices_offsets.x];
+    indx_B = rhs_indices[indices_offsets.y];
+  } else {
+    indx_A = lhs_indices[params->batch_stride_a * tid.z];
+    indx_B = rhs_indices[params->batch_stride_b * tid.z];
+  }
+  A += elem_to_loc(indx_A, batch_shape_a, batch_strides_a, batch_ndim_a);
+  B += elem_to_loc(indx_B, batch_shape_b, batch_strides_b, batch_ndim_b);
+  C += params->batch_stride_d * tid.z;
+
+  // Prepare threadgroup memory
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  // Just make sure everybody's finished with the indexing math above.
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Find block in A, B, C
+  const int c_row = tid.y * BM;
+  const int c_col = tid.x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  C += c_row_long * params->ldd + c_col_long;
+
+  // Prepare threadgroup mma operation
+  thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup loading operations
+  thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+  thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col));
+
+  // Prepare iterations
+  int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+  // Do unaligned K iterations first
+  if (!align_K) {
+    const int k_last = params->gemm_k_iterations_aligned * BK;
+    const int k_remain = params->K - k_last;
+    const size_t k_jump_a =
+        transpose_a ? params->lda * size_t(k_last) : size_t(k_last);
+    const size_t k_jump_b =
+        transpose_b ? size_t(k_last) : params->ldb * size_t(k_last);
+
+    // Move loader source ahead to end
+    loader_a.src += k_jump_a;
+    loader_b.src += k_jump_b;
+
+    // Load tile
+    const short2 tile_dims_A =
+        transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+    const short2 tile_dims_B =
+        transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+    loader_a.load_safe(tile_dims_A);
+    loader_b.load_safe(tile_dims_B);
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Do matmul
+    mma_op.mma(As, Bs);
+
+    // Reset source back to start
+    loader_a.src -= k_jump_a;
+    loader_b.src -= k_jump_b;
+  }
+
+  // Matrix level aligned never check
+  if (align_M && align_N) {
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Load elements into threadgroup
+      loader_a.load_unsafe();
+      loader_b.load_unsafe();
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+
+    // Store results to device memory
+    mma_op.store_result(C, params->ldd);
+  } else {
+    const short lbk = 0;
+
+    // Tile aligned don't check
+    if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) {
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          lbk,
+          LoopAlignment<true, true, true>{});
+      mma_op.store_result(C, params->ldd);
+    }
+
+    // Tile partially aligned check rows
+    else if (align_N || tgp_bn == BN) {
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          lbk,
+          LoopAlignment<false, true, true>{});
+      mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+
+    // Tile partially aligned check cols
+    else if (align_M || tgp_bm == BM) {
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          lbk,
+          LoopAlignment<true, false, true>{});
+      mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+
+    // Nothing aligned so check both rows and cols
+    else {
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iterations,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          lbk,
+          LoopAlignment<false, false, true>{});
+      mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.metal
new file mode 100644
index 0000000..f8e5a2a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.metal
@@ -0,0 +1,59 @@
+// Copyright © 2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h"
+
+#define instantiate_gather_mm_rhs(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                         \
+      "steel_gather_mm_rhs_" #tname "_" #iname "_" #oname "_bm" #bm "_bn" #bn \
+      "_bk" #bk "_wm" #wm "_wn" #wn,                                          \
+      gather_mm_rhs,                                                          \
+      itype,                                                                  \
+      bm,                                                                     \
+      bn,                                                                     \
+      bk,                                                                     \
+      wm,                                                                     \
+      wn,                                                                     \
+      trans_a,                                                                \
+      trans_b,                                                                \
+      float)
+
+#define instantiate_gather_mm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                     \
+      "steel_gather_mm_" #tname "_" #iname "_" #oname "_bm" #bm "_bn" #bn \
+      "_bk" #bk "_wm" #wm "_wn" #wn,                                      \
+      gather_mm,                                                          \
+      itype,                                                              \
+      bm,                                                                 \
+      bn,                                                                 \
+      bk,                                                                 \
+      wm,                                                                 \
+      wn,                                                                 \
+      trans_a,                                                            \
+      trans_b,                                                            \
+      float)
+
+#define instantiate_gather_mm_rhs_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_gather_mm_rhs(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn)  \
+  instantiate_gather_mm_rhs(nt, false,  true, iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gather_mm_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_gather_mm(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn)      \
+  instantiate_gather_mm(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn)      \
+  instantiate_gather_mm(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn)      \
+  instantiate_gather_mm(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gather_mm_shapes_helper(iname, itype, oname, otype)                     \
+  instantiate_gather_mm_rhs_transpose_helper(iname, itype, oname, otype, 16, 64, 16, 1, 2)  \
+  instantiate_gather_mm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 2, 2)      \
+  instantiate_gather_mm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 1, 2)      \
+  instantiate_gather_mm_transpose_helper(iname, itype, oname, otype, 64, 32, 32, 2, 2)      \
+  instantiate_gather_mm_transpose_helper(iname, itype, oname, otype, 32, 64, 16, 1, 2)      \
+  instantiate_gather_mm_transpose_helper(iname, itype, oname, otype, 32, 32, 16, 2, 2)
+// clang-format on
+
+instantiate_gather_mm_shapes_helper(float16, half, float16, half);
+instantiate_gather_mm_shapes_helper(bfloat16, bfloat16_t, bfloat16, bfloat16_t);
+instantiate_gather_mm_shapes_helper(float32, float, float32, float);
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h
new file mode 100644
index 0000000..67cd737
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h
@@ -0,0 +1,143 @@
+// Copyright © 2024 Apple Inc.
+
+using namespace mlx::steel;
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+constant bool align_K [[function_constant(202)]];
+
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
+gather_mm_rhs_nax(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    const device uint32_t* rhs_indices [[buffer(2)]],
+    device T* C [[buffer(3)]],
+    const constant GEMMParams* params [[buffer(4)]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+
+  if (params->tiles_n <= static_cast<int>(tid.x) ||
+      params->tiles_m <= static_cast<int>(tid.y)) {
+    return;
+  }
+
+  // Find the block in A, B, C
+  const int c_row = tid.y * BM;
+  const int c_col = tid.x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  C += c_row_long * params->ldd + c_col_long;
+  rhs_indices += c_row;
+
+  const short tm = SM * (simd_group_id / WN);
+  const short tn = SN * (simd_group_id % WN);
+
+  const int sgp_sm_int =
+      align_M ? int(SM) : min(int(SM), params->M - (c_row + tm));
+  const short sgp_sm = short(sgp_sm_int);
+  const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
+
+  const int sgp_sn_int =
+      align_N ? int(SN) : min(int(SN), params->N - (c_col + tn));
+  const short sgp_sn = short(sgp_sn_int);
+  const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN);
+
+  A += transpose_a ? tm : (tm * params->lda);
+  B += transpose_b ? (tn * params->ldb) : tn;
+  C += tm * params->ldd + tn;
+  rhs_indices += tm;
+
+  // Do as many matmuls as necessary
+  uint32_t index;
+  short offset;
+  uint32_t index_next = rhs_indices[0];
+  short offset_next = 0;
+  int n = 0;
+  while (n < sgp_sm) {
+    n++;
+    offset = offset_next;
+    index = index_next;
+    offset_next = sgp_sm;
+    for (; n < sgp_sm; n++) {
+      if (rhs_indices[n] != index) {
+        offset_next = n;
+        index_next = rhs_indices[n];
+        break;
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_none);
+
+    using DSubTile = NAXSubTile<AccumType, UM, UN>;
+    NAXTile<AccumType, TM, TN, DSubTile> Ctile;
+
+    dispatch_bool(align_K, [&](auto kAlignedK) {
+      dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+        dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) {
+          auto do_gemm = gemm_loop<
+              T,
+              SM,
+              SN,
+              SK,
+              BK,
+              transpose_a,
+              transpose_b,
+              kAlignedM.value,
+              kAlignedN.value,
+              kAlignedK.value,
+              UM,
+              UN,
+              UK,
+              AccumType>;
+          Ctile = do_gemm(
+              A,
+              B + index * params->batch_stride_b,
+              params->lda,
+              params->ldb,
+              params->K,
+              params->gemm_k_iterations_aligned,
+              sgp_sm,
+              sgp_sn);
+
+          if constexpr (kAlignedN.value) {
+            if (offset_next - offset == SM) {
+              Ctile.store(C, int(params->ldd));
+            } else {
+              Ctile.store_slice(
+                  C,
+                  int(params->ldd),
+                  short2(0, offset),
+                  short2(SN, offset_next));
+            }
+          } else {
+            Ctile.store_slice(
+                C,
+                int(params->ldd),
+                short2(0, offset),
+                short2(sgp_sn, offset_next));
+          }
+        });
+      });
+    });
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.metal
new file mode 100644
index 0000000..690c7a3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.metal
@@ -0,0 +1,37 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h"
+
+// clang-format off
+#define instantiate_gather_mm_rhs(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                             \
+      "steel_gather_mm_rhs_nax_" #tname "_" #iname "_" #oname "_bm" #bm "_bn" #bn \
+      "_bk" #bk "_wm" #wm "_wn" #wn,                                              \
+      gather_mm_rhs_nax,                                                          \
+      itype,                                                                      \
+      bm,                                                                         \
+      bn,                                                                         \
+      bk,                                                                         \
+      wm,                                                                         \
+      wn,                                                                         \
+      trans_a,                                                                    \
+      trans_b,                                                                    \
+      float)
+
+#define instantiate_gather_mm_rhs_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_gather_mm_rhs(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn)  \
+  instantiate_gather_mm_rhs(nt, false,  true, iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gather_mm_shapes_helper(iname, itype, oname, otype)                      \
+  instantiate_gather_mm_rhs_transpose_helper(iname, itype, oname, otype, 16, 128, 128, 1, 4) \
+  instantiate_gather_mm_rhs_transpose_helper(iname, itype, oname, otype, 32, 128, 128, 1, 4) \
+  instantiate_gather_mm_rhs_transpose_helper(iname, itype, oname, otype, 64, 128, 128, 2, 4)
+// clang-format on
+
+instantiate_gather_mm_shapes_helper(float16, half, float16, half);
+instantiate_gather_mm_shapes_helper(bfloat16, bfloat, bfloat16, bfloat);
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h
new file mode 100644
index 0000000..cc3ddd9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h
@@ -0,0 +1,719 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+using namespace metal;
+using namespace mlx::steel;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernels
+///////////////////////////////////////////////////////////////////////////////
+
+struct _NoMask {
+  char x;
+
+  constexpr METAL_FUNC operator bool() {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const threadgroup {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const device {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const constant {
+    return true;
+  }
+};
+
+template <typename OutT, typename InT = OutT>
+struct ScaleOp {
+  OutT scale;
+
+  METAL_FUNC OutT apply(InT x) const {
+    return static_cast<OutT>(x) * scale;
+  }
+};
+
+typedef struct _NoMask nomask_t;
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    bool MN_aligned,
+    bool K_aligned>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
+block_masked_gemm(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device T* D [[buffer(3)]],
+    const constant GEMMParams* params [[buffer(4)]],
+    const constant int* batch_shape [[buffer(6)]],
+    const constant int64_t* batch_strides [[buffer(7)]],
+    const device out_mask_t* out_mask [[buffer(10)]],
+    const device op_mask_t* lhs_mask [[buffer(11)]],
+    const device op_mask_t* rhs_mask [[buffer(12)]],
+    const constant int* mask_strides [[buffer(13)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  // Appease the compiler
+  (void)lid;
+
+  static_assert(
+      BM == BN,
+      "block_masked_gemm must have the same block M and block N size");
+  static_assert(BM % BK == 0, "block_masked_gemm must have BM % BK == 0");
+
+  constexpr bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  constexpr bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  constexpr bool has_mul_operand_mask =
+      has_operand_mask && !metal::is_same_v<op_mask_t, bool>;
+  constexpr bool has_mul_output_mask =
+      has_output_mask && !metal::is_same_v<out_mask_t, bool>;
+
+  constexpr short k_mask_factor = short(BM / BK);
+
+  using gemm_kernel = GEMMKernel<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      MN_aligned,
+      K_aligned>;
+
+  const int tid_y = ((tid.y) << params->swizzle_log) +
+      ((tid.x) & ((1 << params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> params->swizzle_log;
+
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  const constant auto* mask_batch_strides =
+      batch_strides + 2 * params->batch_ndim;
+
+  if (params->batch_ndim > 1) {
+    if (has_output_mask) {
+      out_mask += elem_to_loc(
+          tid.z, batch_shape, mask_batch_strides, params->batch_ndim);
+
+      mask_batch_strides += params->batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      const constant auto* mask_strides_lhs = mask_batch_strides;
+      const constant auto* mask_strides_rhs =
+          mask_strides_lhs + params->batch_ndim;
+
+      ulong2 batch_offsets = elem_to_loc_broadcast(
+          tid.z,
+          batch_shape,
+          mask_strides_lhs,
+          mask_strides_rhs,
+          params->batch_ndim);
+
+      lhs_mask += batch_offsets.x;
+      rhs_mask += batch_offsets.y;
+    }
+  } else {
+    if (has_output_mask) {
+      out_mask += tid.z * mask_batch_strides[0];
+      mask_batch_strides += params->batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      lhs_mask += tid.z * mask_batch_strides[0];
+      rhs_mask += tid.z * mask_batch_strides[params->batch_ndim];
+    }
+  }
+
+  // Adjust for batch
+  if (params->batch_ndim > 1) {
+    const constant auto* A_bstrides = batch_strides;
+    const constant auto* B_bstrides = batch_strides + params->batch_ndim;
+
+    ulong2 batch_offsets = elem_to_loc_broadcast(
+        tid.z, batch_shape, A_bstrides, B_bstrides, params->batch_ndim);
+
+    A += batch_offsets.x;
+    B += batch_offsets.y;
+
+  } else {
+    A += params->batch_stride_a * tid.z;
+    B += params->batch_stride_b * tid.z;
+  }
+
+  D += params->batch_stride_d * tid.z;
+
+  // Find block in A, B, C
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  D += c_row_long * params->ldd + c_col_long;
+
+  const constant int* out_mask_strides = mask_strides;
+  const constant int* lhs_mask_strides =
+      mask_strides + (has_output_mask ? 2 : 0);
+  const constant int* rhs_mask_strides =
+      lhs_mask_strides + (has_operand_mask ? 2 : 0);
+
+  const int out_mask_offset = !has_output_mask
+      ? 0
+      : tid_y * out_mask_strides[1] + tid_x * out_mask_strides[0];
+  int lhs_mask_offset = !has_operand_mask ? 0 : tid_y * lhs_mask_strides[1];
+  int rhs_mask_offset = !has_operand_mask ? 0 : tid_x * rhs_mask_strides[0];
+  const int lhs_mask_step = !has_operand_mask ? 0 : lhs_mask_strides[0];
+  const int rhs_mask_step = !has_operand_mask ? 0 : rhs_mask_strides[1];
+  short k_factor_cnt = k_mask_factor;
+
+  ScaleOp<float> out_mask_op;
+  ScaleOp<T> lhs_mask_op;
+  ScaleOp<T> rhs_mask_op;
+
+  if (has_output_mask) {
+    auto mask_out = out_mask[out_mask_offset];
+
+    if (has_mul_output_mask) {
+      out_mask_op.scale = float(mask_out);
+    }
+
+    // Write zeros and return
+    if (!mask_out) {
+      constexpr short tgp_size = WM * WN * 32;
+      constexpr short vec_size = 4;
+
+      // Tile threads in threadgroup
+      constexpr short TN = BN / vec_size;
+      constexpr short TM = tgp_size / TN;
+
+      const short thread_idx = simd_group_id * 32 + simd_lane_id;
+      const short bi = thread_idx / TN;
+      const short bj = vec_size * (thread_idx % TN);
+
+      D += bi * params->ldd + bj;
+
+      short tgp_bm = min(BM, params->M - c_row);
+      short tgp_bn = min(BN, params->N - c_col);
+
+      if (MN_aligned || (tgp_bm == BM && tgp_bn == BN)) {
+        for (short ti = 0; ti < BM; ti += TM) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            D[ti * params->ldd + j] = T(0.);
+          }
+        }
+      } else {
+        short jmax = tgp_bn - bj;
+        jmax = jmax < vec_size ? jmax : vec_size;
+        for (short ti = 0; (bi + ti) < tgp_bm; ti += TM) {
+          for (short j = 0; j < jmax; j++) {
+            D[ti * params->ldd + j] = T(0.);
+          }
+        }
+      }
+
+      return;
+    }
+  }
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Prepare threadgroup mma operation
+  thread typename gemm_kernel::mma_t mma_op(simd_group_id, simd_lane_id);
+
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  // Prepare threadgroup loading operations
+  thread typename gemm_kernel::loader_a_t loader_a(
+      A, params->lda, As, simd_group_id, simd_lane_id);
+  thread typename gemm_kernel::loader_b_t loader_b(
+      B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm =
+      MN_aligned ? short(BM) : short(min(BM, params->M - c_row));
+  const short tgp_bn =
+      MN_aligned ? short(BN) : short(min(BN, params->N - c_col));
+
+  int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+  ///////////////////////////////////////////////////////////////////////////////
+  // Do unaligned K iterations first
+  if (!K_aligned) {
+    const int k_last = params->gemm_k_iterations_aligned * BK;
+    const int mask_idx_last = k_last / BM;
+
+    if (!has_operand_mask ||
+        (bool(lhs_mask[lhs_mask_offset + mask_idx_last * lhs_mask_step]) &&
+         bool(rhs_mask[rhs_mask_offset + mask_idx_last * rhs_mask_step]))) {
+      if (has_mul_operand_mask) {
+        lhs_mask_op.scale =
+            lhs_mask[lhs_mask_offset + mask_idx_last * lhs_mask_step];
+        rhs_mask_op.scale =
+            rhs_mask[rhs_mask_offset + mask_idx_last * rhs_mask_step];
+      }
+
+      // Move loader source ahead to end
+      const int k_remain = params->K - k_last;
+      const size_t k_jump_a =
+          transpose_a ? params->lda * size_t(k_last) : size_t(k_last);
+      const size_t k_jump_b =
+          transpose_b ? size_t(k_last) : params->ldb * size_t(k_last);
+
+      loader_a.src += k_jump_a;
+      loader_b.src += k_jump_b;
+
+      // Load tile
+      const short2 tile_dims_A =
+          transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+      const short2 tile_dims_B =
+          transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+
+      loader_a.load_safe(tile_dims_A);
+      loader_b.load_safe(tile_dims_B);
+
+      if (has_mul_operand_mask) {
+        loader_a.apply_inplace_op(lhs_mask_op);
+        loader_b.apply_inplace_op(rhs_mask_op);
+      }
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Do matmul
+      mma_op.mma(As, Bs);
+
+      // Reset source back to start
+      loader_a.src -= k_jump_a;
+      loader_b.src -= k_jump_b;
+    }
+  }
+
+  ///////////////////////////////////////////////////////////////////////////////
+  // MNK aligned loop
+  if (MN_aligned) {
+    for (; gemm_k_iterations > 0; gemm_k_iterations--) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      if (!has_operand_mask ||
+          (bool(lhs_mask[lhs_mask_offset]) &&
+           bool(rhs_mask[rhs_mask_offset]))) {
+        if (has_mul_operand_mask) {
+          lhs_mask_op.scale = lhs_mask[lhs_mask_offset];
+          rhs_mask_op.scale = rhs_mask[rhs_mask_offset];
+        }
+
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        if (has_mul_operand_mask) {
+          loader_a.apply_inplace_op(lhs_mask_op);
+          loader_b.apply_inplace_op(rhs_mask_op);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+      }
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+
+      k_factor_cnt--;
+      lhs_mask_offset += k_factor_cnt == 0 ? lhs_mask_step : 0;
+      rhs_mask_offset += k_factor_cnt == 0 ? rhs_mask_step : 0;
+      k_factor_cnt = k_factor_cnt == 0 ? k_mask_factor : k_factor_cnt;
+    }
+
+    if (has_mul_output_mask) {
+      mma_op.apply_epilogue(out_mask_op);
+    }
+
+    // Store results to device memory
+    mma_op.store_result(D, params->ldd);
+    return;
+
+  }
+  ///////////////////////////////////////////////////////////////////////////////
+  // MN unaligned loop
+  else {
+    const bool M_aligned = (tgp_bm == BM);
+    const bool N_aligned = (tgp_bn == BN);
+
+    const short2 tile_dims_A =
+        transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm);
+    const short2 tile_dims_B =
+        transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK);
+
+    for (; gemm_k_iterations > 0; gemm_k_iterations--) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      if (!has_operand_mask ||
+          (bool(lhs_mask[lhs_mask_offset]) &&
+           bool(rhs_mask[rhs_mask_offset]))) {
+        if (has_mul_operand_mask) {
+          lhs_mask_op.scale = lhs_mask[lhs_mask_offset];
+          rhs_mask_op.scale = rhs_mask[rhs_mask_offset];
+        }
+
+        // Load elements into threadgroup
+        if (M_aligned) {
+          loader_a.load_unsafe();
+        } else {
+          loader_a.load_safe(tile_dims_A);
+        }
+
+        if (N_aligned) {
+          loader_b.load_unsafe();
+        } else {
+          loader_b.load_safe(tile_dims_B);
+        }
+
+        if (has_mul_operand_mask) {
+          loader_a.apply_inplace_op(lhs_mask_op);
+          loader_b.apply_inplace_op(rhs_mask_op);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+      }
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+
+      k_factor_cnt--;
+      lhs_mask_offset += k_factor_cnt == 0 ? lhs_mask_step : 0;
+      rhs_mask_offset += k_factor_cnt == 0 ? rhs_mask_step : 0;
+      k_factor_cnt = k_factor_cnt == 0 ? k_mask_factor : k_factor_cnt;
+    }
+
+    if (has_mul_output_mask) {
+      mma_op.apply_epilogue(out_mask_op);
+    }
+
+    if (M_aligned && N_aligned) {
+      mma_op.store_result(D, params->ldd);
+    } else {
+      mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+  }
+}
+
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    bool MN_aligned,
+    bool K_aligned,
+    bool has_operand_mask = false>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void
+block_masked_gemm(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device T* D [[buffer(3)]],
+    const constant GEMMParams* params [[buffer(4)]],
+    const constant int* batch_shape [[buffer(6)]],
+    const constant int64_t* batch_strides [[buffer(7)]],
+    const device bool* out_mask [[buffer(10)]],
+    const device bool* lhs_mask [[buffer(11)]],
+    const device bool* rhs_mask [[buffer(12)]],
+    const constant int* mask_strides [[buffer(13)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  // Appease the compiler
+  (void)lid;
+
+  using gemm_kernel = GEMMKernel<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      MN_aligned,
+      K_aligned>;
+
+  const int tid_y = ((tid.y) << params->swizzle_log) +
+      ((tid.x) & ((1 << params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> params->swizzle_log;
+
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  if (params->batch_ndim > 1) {
+    const constant auto* mask_batch_strides =
+        batch_strides + 2 * params->batch_ndim;
+    out_mask +=
+        elem_to_loc(tid.z, batch_shape, mask_batch_strides, params->batch_ndim);
+
+    if (has_operand_mask) {
+      const constant auto* mask_strides_lhs =
+          mask_batch_strides + params->batch_ndim;
+      const constant auto* mask_strides_rhs =
+          mask_strides_lhs + params->batch_ndim;
+
+      ulong2 batch_offsets = elem_to_loc_broadcast(
+          tid.z,
+          batch_shape,
+          mask_strides_lhs,
+          mask_strides_rhs,
+          params->batch_ndim);
+
+      lhs_mask += batch_offsets.x;
+      rhs_mask += batch_offsets.y;
+    }
+  } else {
+    out_mask += tid.z * batch_strides[2 * params->batch_ndim];
+    if (has_operand_mask) {
+      lhs_mask += tid.z * batch_strides[3 * params->batch_ndim];
+      rhs_mask += tid.z * batch_strides[4 * params->batch_ndim];
+    }
+  }
+
+  // Adjust for batch
+  if (params->batch_ndim > 1) {
+    const constant auto* A_bstrides = batch_strides;
+    const constant auto* B_bstrides = batch_strides + params->batch_ndim;
+
+    ulong2 batch_offsets = elem_to_loc_broadcast(
+        tid.z, batch_shape, A_bstrides, B_bstrides, params->batch_ndim);
+
+    A += batch_offsets.x;
+    B += batch_offsets.y;
+
+  } else {
+    A += params->batch_stride_a * tid.z;
+    B += params->batch_stride_b * tid.z;
+  }
+
+  D += params->batch_stride_d * tid.z;
+
+  // Find block in A, B, C
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  D += c_row_long * params->ldd + c_col_long;
+
+  bool mask_out = out_mask[tid_y * mask_strides[1] + tid_x * mask_strides[0]];
+
+  // Write zeros and return
+  if (!mask_out) {
+    constexpr short tgp_size = WM * WN * 32;
+    constexpr short vec_size = 4;
+
+    // Tile threads in threadgroup
+    constexpr short TN = BN / vec_size;
+    constexpr short TM = tgp_size / TN;
+
+    const short thread_idx = simd_group_id * 32 + simd_lane_id;
+    const short bi = thread_idx / TN;
+    const short bj = vec_size * (thread_idx % TN);
+
+    D += bi * params->ldd + bj;
+
+    short tgp_bm = min(BM, params->M - c_row);
+    short tgp_bn = min(BN, params->N - c_col);
+
+    if (MN_aligned || (tgp_bm == BM && tgp_bn == BN)) {
+      for (short ti = 0; ti < BM; ti += TM) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          D[ti * params->ldd + j] = T(0.);
+        }
+      }
+    } else {
+      short jmax = tgp_bn - bj;
+      jmax = jmax < vec_size ? jmax : vec_size;
+      for (short ti = 0; (bi + ti) < tgp_bm; ti += TM) {
+        for (short j = 0; j < jmax; j++) {
+          D[ti * params->ldd + j] = T(0.);
+        }
+      }
+    }
+
+    return;
+  }
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Prepare threadgroup mma operation
+  thread typename gemm_kernel::mma_t mma_op(simd_group_id, simd_lane_id);
+
+  int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  // Prepare threadgroup loading operations
+  thread typename gemm_kernel::loader_a_t loader_a(
+      A, params->lda, As, simd_group_id, simd_lane_id);
+  thread typename gemm_kernel::loader_b_t loader_b(
+      B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+  ///////////////////////////////////////////////////////////////////////////////
+  // MNK aligned loop
+  if (MN_aligned) {
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      if (!has_operand_mask ||
+          (lhs_mask
+               [tid_y * mask_strides[3] + ((k * BK) / BM) * mask_strides[2]] &&
+           rhs_mask
+               [((k * BK) / BM) * mask_strides[5] + tid_x * mask_strides[4]])) {
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+      }
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Loop tail
+    if (!K_aligned) {
+      if (!has_operand_mask ||
+          (lhs_mask
+               [tid_y * mask_strides[3] + (params->K / BM) * mask_strides[2]] &&
+           rhs_mask
+               [(params->K / BM) * mask_strides[5] +
+                tid_x * mask_strides[4]])) {
+        int lbk = params->K - params->gemm_k_iterations_aligned * BK;
+        short2 tile_dims_A = transpose_a ? short2(BM, lbk) : short2(lbk, BM);
+        short2 tile_dims_B = transpose_b ? short2(lbk, BN) : short2(BN, lbk);
+
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        mma_op.mma(As, Bs);
+      }
+    }
+
+    // Store results to device memory
+    mma_op.store_result(D, params->ldd);
+    return;
+
+  }
+  ///////////////////////////////////////////////////////////////////////////////
+  // MN unaligned loop
+  else { // Loop over K - unaligned case
+    short tgp_bm = min(BM, params->M - c_row);
+    short tgp_bn = min(BN, params->N - c_col);
+    short lbk = params->K - params->gemm_k_iterations_aligned * BK;
+
+    bool M_aligned = (tgp_bm == BM);
+    bool N_aligned = (tgp_bn == BN);
+
+    short2 tile_dims_A = transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm);
+    short2 tile_dims_B = transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK);
+
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      if (!has_operand_mask ||
+          (lhs_mask
+               [tid_y * mask_strides[3] + ((k * BK) / BM) * mask_strides[2]] &&
+           rhs_mask
+               [((k * BK) / BM) * mask_strides[5] + tid_x * mask_strides[4]])) {
+        // Load elements into threadgroup
+        if (M_aligned) {
+          loader_a.load_unsafe();
+        } else {
+          loader_a.load_safe(tile_dims_A);
+        }
+
+        if (N_aligned) {
+          loader_b.load_unsafe();
+        } else {
+          loader_b.load_safe(tile_dims_B);
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+      }
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+
+    if (!K_aligned) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      if (!has_operand_mask ||
+          (lhs_mask
+               [tid_y * mask_strides[3] + (params->K / BM) * mask_strides[2]] &&
+           rhs_mask
+               [(params->K / BM) * mask_strides[5] +
+                tid_x * mask_strides[4]])) {
+        short2 tile_dims_A_last =
+            transpose_a ? short2(tgp_bm, lbk) : short2(lbk, tgp_bm);
+        short2 tile_dims_B_last =
+            transpose_b ? short2(lbk, tgp_bn) : short2(tgp_bn, lbk);
+
+        loader_a.load_safe(tile_dims_A_last);
+        loader_b.load_safe(tile_dims_B_last);
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        mma_op.mma(As, Bs);
+      }
+    }
+
+    if (M_aligned && N_aligned) {
+      mma_op.store_result(D, params->ldd);
+    } else {
+      mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.metal
new file mode 100644
index 0000000..aa1274c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.metal
@@ -0,0 +1,76 @@
+// Copyright © 2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h"
+
+#define instantiate_gemm(                                              \
+    outmaskname,                                                       \
+    outmasktype,                                                       \
+    opmaskname,                                                        \
+    opmasktype,                                                        \
+    tname,                                                             \
+    trans_a,                                                           \
+    trans_b,                                                           \
+    iname,                                                             \
+    itype,                                                             \
+    oname,                                                             \
+    otype,                                                             \
+    bm,                                                                \
+    bn,                                                                \
+    bk,                                                                \
+    wm,                                                                \
+    wn,                                                                \
+    aname,                                                             \
+    mn_aligned,                                                        \
+    kname,                                                             \
+    k_aligned)                                                         \
+  instantiate_kernel(                                                  \
+    "steel_gemm_block_outmask_" #outmaskname                           \
+      "_opmask_" #opmaskname "_" #tname "_" #iname "_" #oname          \
+      "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn                \
+      "_MN_" #aname "_K_" #kname,                                      \
+    block_masked_gemm,                                                 \
+      itype,                                                           \
+      outmasktype,                                                     \
+      opmasktype,                                                      \
+      bm,                                                              \
+      bn,                                                              \
+      bk,                                                              \
+      wm,                                                              \
+      wn,                                                              \
+      trans_a,                                                         \
+      trans_b,                                                         \
+      mn_aligned,                                                      \
+      k_aligned)
+
+#define instantiate_gemm_mask_helper(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)                \
+  instantiate_gemm(bool_, bool, bool_, bool, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)        \
+  instantiate_gemm(iname, itype, iname, itype, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)      \
+  instantiate_gemm(bool_, bool, iname, itype, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)       \
+  instantiate_gemm(iname, itype, bool_, bool, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)       \
+  instantiate_gemm(nomask, nomask_t, bool_, bool, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)   \
+  instantiate_gemm(nomask, nomask_t, iname, itype, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)  \
+  instantiate_gemm(bool_, bool, nomask, nomask_t, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)   \
+  instantiate_gemm(iname, itype, nomask, nomask_t, tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, aname, mn_aligned, kname, k_aligned)
+
+#define instantiate_gemm_aligned_helper(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn)                         \
+  instantiate_gemm_mask_helper(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, taligned, true, taligned, true)  \
+  instantiate_gemm_mask_helper(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, taligned, true, naligned, false) \
+  instantiate_gemm_mask_helper(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, naligned, false, taligned, true) \
+  instantiate_gemm_mask_helper(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, naligned, false, naligned, false)
+
+#define instantiate_gemm_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn)             \
+    instantiate_gemm_aligned_helper(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_aligned_helper(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_aligned_helper(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_aligned_helper(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gemm_shapes_helper(iname, itype, oname, otype)                  \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 32, 16, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 2, 2)
+
+instantiate_gemm_shapes_helper(float16, half, float16, half);
+instantiate_gemm_shapes_helper(bfloat16, bfloat16_t, bfloat16, bfloat16_t);
+instantiate_gemm_shapes_helper(float32, float, float32, float); // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h
new file mode 100644
index 0000000..5a43e22
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h
@@ -0,0 +1,266 @@
+// Copyright © 2025 Apple Inc.
+
+using namespace mlx::steel;
+
+constant bool segments_contiguous [[function_constant(199)]];
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void segmented_mm(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    const device uint32_t* segments [[buffer(2)]],
+    device T* C [[buffer(3)]],
+    const constant GEMMParams* params [[buffer(4)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]]) {
+  using gemm_kernel = GEMMKernel<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      true,
+      true,
+      AccumType>;
+
+  using loader_a_t = typename gemm_kernel::loader_a_t;
+  using loader_b_t = typename gemm_kernel::loader_b_t;
+  using mma_t = typename gemm_kernel::mma_t;
+
+  if (params->tiles_n <= static_cast<int>(tid.x) ||
+      params->tiles_m <= static_cast<int>(tid.y)) {
+    return;
+  }
+
+  // Prepare threadgroup memory
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  // Find the block in A, B, C
+  const int c_row = tid.y * BM;
+  const int c_col = tid.x * BN;
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+
+  // Prepare threadgroup bounds
+  const short tgp_bm = align_M ? BM : short(min(BM, params->M - c_row));
+  const short tgp_bn = align_N ? BN : short(min(BN, params->N - c_col));
+
+  // Move the pointers to the output tile
+  A += transpose_a ? c_row_long : c_row_long * params->lda;
+  B += transpose_b ? c_col_long * params->ldb : c_col_long;
+  C += c_row_long * params->ldd + c_col_long;
+
+  // Move the pointers to the start of the segment
+  uint32_t k_start, k_end;
+  if (segments_contiguous) {
+    k_start = segments[2 * tid.z];
+    k_end = segments[2 * tid.z + 1];
+  } else {
+    // We accept either contiguous (above) or weird strides where the beginning
+    // of the next one is the previous one. Basically the last two strides are
+    // both 1!
+    k_start = segments[tid.z];
+    k_end = segments[tid.z + 1];
+  }
+  A += transpose_a ? k_start * params->lda : k_start;
+  B += transpose_b ? k_start : k_start * params->ldb;
+  C += tid.z * params->batch_stride_d;
+
+  // Prepare threadgroup mma operation
+  thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup loading operations
+  thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+  thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+  // Matrix level alignment so only check K
+  if (align_M && align_N) {
+    uint32_t k = k_start + BK;
+    for (; k <= k_end; k += BK) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Load elements into threadgroup
+      loader_a.load_unsafe();
+      loader_b.load_unsafe();
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+    short k_remain = BK - short(k - k_end);
+    const short2 tile_dims_A =
+        transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+    const short2 tile_dims_B =
+        transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+    if (k_remain > 0) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      loader_a.load_safe(tile_dims_A);
+      loader_b.load_safe(tile_dims_B);
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      mma_op.mma(As, Bs);
+    }
+    mma_op.store_result(C, params->ldd);
+  } else {
+    // Tile aligned do the same as above
+    if ((align_M || tgp_bm == BM) && (align_N || tgp_bn == BN)) {
+      uint32_t k = k_start + BK;
+      for (; k <= k_end; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+      short k_remain = BK - short(k - k_end);
+      const short2 tile_dims_A =
+          transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+      const short2 tile_dims_B =
+          transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+      if (k_remain > 0) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(As, Bs);
+      }
+      mma_op.store_result(C, params->ldd);
+    }
+
+    // Tile partially aligned check rows
+    else if (align_N || tgp_bn == BN) {
+      uint32_t k = k_start + BK;
+      for (; k <= k_end; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Load elements into threadgroup
+        loader_a.load_safe(
+            transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm));
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+      short k_remain = BK - short(k - k_end);
+      const short2 tile_dims_A =
+          transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+      const short2 tile_dims_B =
+          transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+      if (k_remain > 0) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(As, Bs);
+      }
+      mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+
+    // Tile partially aligned check cols
+    else if (align_M || tgp_bm == BM) {
+      uint32_t k = k_start + BK;
+      for (; k <= k_end; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_safe(
+            transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK));
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+      short k_remain = BK - short(k - k_end);
+      const short2 tile_dims_A =
+          transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+      const short2 tile_dims_B =
+          transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+      if (k_remain > 0) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(As, Bs);
+      }
+      mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+
+    // Nothing aligned so check both rows and cols
+    else {
+      uint32_t k = k_start + BK;
+      for (; k <= k_end; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Load elements into threadgroup
+        loader_a.load_safe(
+            transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm));
+        loader_b.load_safe(
+            transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK));
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+      short k_remain = BK - short(k - k_end);
+      const short2 tile_dims_A =
+          transpose_a ? short2(tgp_bm, k_remain) : short2(k_remain, tgp_bm);
+      const short2 tile_dims_B =
+          transpose_b ? short2(k_remain, tgp_bn) : short2(tgp_bn, k_remain);
+      if (k_remain > 0) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(As, Bs);
+      }
+      mma_op.store_result_safe(C, params->ldd, short2(tgp_bn, tgp_bm));
+    }
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.metal
new file mode 100644
index 0000000..a7515c3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.metal
@@ -0,0 +1,43 @@
+// Copyright © 2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h"
+
+#define instantiate_segmented_mm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                         \
+      "steel_segmented_mm_" #tname "_" #iname "_" #oname "_bm" #bm "_bn" #bn  \
+      "_bk" #bk "_wm" #wm "_wn" #wn,                                          \
+      segmented_mm,                                                           \
+      itype,                                                                  \
+      bm,                                                                     \
+      bn,                                                                     \
+      bk,                                                                     \
+      wm,                                                                     \
+      wn,                                                                     \
+      trans_a,                                                                \
+      trans_b,                                                                \
+      float)
+
+#define instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_segmented_mm(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn)      \
+  instantiate_segmented_mm(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn)      \
+  instantiate_segmented_mm(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn)      \
+  instantiate_segmented_mm(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_segmented_mm_shapes_helper(iname, itype, oname, otype)                 \
+  instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 2, 2)  \
+  instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 64, 64, 16, 1, 2)  \
+  instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 64, 32, 32, 2, 2)  \
+  instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 32, 64, 16, 1, 2)  \
+  instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 32, 32, 16, 2, 2)
+// clang-format on
+
+instantiate_segmented_mm_shapes_helper(float16, half, float16, half);
+instantiate_segmented_mm_shapes_helper(
+    bfloat16,
+    bfloat16_t,
+    bfloat16,
+    bfloat16_t);
+instantiate_segmented_mm_shapes_helper(float32, float, float32, float);
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h
new file mode 100644
index 0000000..a372e93
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h
@@ -0,0 +1,227 @@
+// Copyright © 2024 Apple Inc.
+
+using namespace mlx::steel;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernels
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    bool MN_aligned,
+    bool K_aligned>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gemm_splitk(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device U* C [[buffer(2)]],
+    const constant GEMMSpiltKParams* params [[buffer(3)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  (void)lid;
+
+  using gemm_kernel = GEMMKernel<
+      T,
+      U,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      MN_aligned,
+      K_aligned>;
+  using loader_a_t = typename gemm_kernel::loader_a_t;
+  using loader_b_t = typename gemm_kernel::loader_b_t;
+  using mma_t = typename gemm_kernel::mma_t;
+
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  const int tid_x = tid.x;
+  const int tid_y = tid.y;
+  const int tid_z = tid.z;
+
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  // Find block in A, B, C
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int k_start = params->split_k_partition_size * tid_z;
+
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+  const size_t k_start_long = size_t(k_start);
+
+  A += transpose_a ? (c_row_long + k_start_long * params->lda)
+                   : (k_start_long + c_row_long * params->lda);
+  B += transpose_b ? (k_start_long + c_col_long * params->ldb)
+                   : (c_col_long + k_start_long * params->ldb);
+  C += (size_t(params->split_k_partition_stride) * tid_z) +
+      (c_row_long * params->ldc + c_col_long);
+
+  // Prepare threadgroup loading operations
+  thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+  thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup mma operation
+  thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+  int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+  short tgp_bm = min(BM, params->M - c_row);
+  short tgp_bn = min(BN, params->N - c_col);
+  short leftover_bk = params->K % BK;
+
+  if (MN_aligned || (tgp_bm == BM && tgp_bn == BN)) {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<true, true, true>{});
+  } else if (tgp_bn == BN) {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<false, true, true>{});
+  } else if (tgp_bm == BM) {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<true, false, true>{});
+  } else {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<false, false, true>{});
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  if ((tid_z + 1) == (params->split_k_partitions)) {
+    int gemm_k_iter_remaining =
+        (params->K - (k_start + params->split_k_partition_size)) / BK;
+    if (!K_aligned || gemm_k_iter_remaining > 0)
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iter_remaining,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          leftover_bk,
+          LoopAlignment<false, false, K_aligned>{});
+  }
+
+  if (MN_aligned || (tgp_bm == BM && tgp_bn == BN)) {
+    mma_op.store_result(C, params->ldc);
+  } else {
+    mma_op.store_result_safe(C, params->ldc, short2(tgp_bn, tgp_bm));
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Split k accumulation kernel
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename AccT,
+    typename OutT,
+    typename Epilogue = TransformNone<OutT, AccT>>
+[[kernel]] void gemm_splitk_accum(
+    const device AccT* C_split [[buffer(0)]],
+    device OutT* D [[buffer(1)]],
+    const constant int& k_partitions [[buffer(2)]],
+    const constant int& partition_stride [[buffer(3)]],
+    const constant int& ldd [[buffer(4)]],
+    uint2 gid [[thread_position_in_grid]]) {
+  // Ajust D and C
+  D += gid.x + gid.y * size_t(ldd);
+  C_split += gid.x + gid.y * size_t(ldd);
+
+  size_t offset = 0;
+  AccT out = 0;
+
+  for (int i = 0; i < k_partitions; i++) {
+    out += C_split[offset];
+    offset += partition_stride;
+  }
+
+  // Write output
+  D[0] = Epilogue::apply(out);
+}
+
+template <
+    typename AccT,
+    typename OutT,
+    typename Epilogue = TransformAxpby<OutT, AccT>>
+[[kernel]] void gemm_splitk_accum_axpby(
+    const device AccT* C_split [[buffer(0)]],
+    device OutT* D [[buffer(1)]],
+    const constant int& k_partitions [[buffer(2)]],
+    const constant int& partition_stride [[buffer(3)]],
+    const constant int& ldd [[buffer(4)]],
+    const device OutT* C [[buffer(5)]],
+    const constant int& ldc [[buffer(6)]],
+    const constant int& fdc [[buffer(7)]],
+    const constant float& alpha [[buffer(8)]],
+    const constant float& beta [[buffer(9)]],
+    uint2 gid [[thread_position_in_grid]]) {
+  // Ajust D and C
+  C += gid.x * size_t(fdc) + gid.y * size_t(ldc);
+  D += gid.x + gid.y * size_t(ldd);
+  C_split += gid.x + gid.y * size_t(ldd);
+
+  size_t offset = 0;
+  AccT out = 0;
+
+  for (int i = 0; i < k_partitions; i++) {
+    out += C_split[offset];
+    offset += partition_stride;
+  }
+
+  // Write output
+  Epilogue op(alpha, beta);
+  D[0] = op.apply(out, *C);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.metal
new file mode 100644
index 0000000..8c60b42
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.metal
@@ -0,0 +1,76 @@
+// Copyright © 2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/steel/gemm/gemm.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h"
+
+#define instantiate_gemm(                                     \
+    tname,                                                    \
+    trans_a,                                                  \
+    trans_b,                                                  \
+    iname,                                                    \
+    itype,                                                    \
+    oname,                                                    \
+    otype,                                                    \
+    bm,                                                       \
+    bn,                                                       \
+    bk,                                                       \
+    wm,                                                       \
+    wn,                                                       \
+    aname,                                                    \
+    mn_aligned,                                               \
+    kname,                                                    \
+    k_aligned)                                                \
+  instantiate_kernel(                                         \
+      "steel_gemm_splitk_" #tname "_" #iname "_" #oname       \
+         "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn    \
+         "_MN_" #aname "_K_" #kname,                          \
+  gemm_splitk,                                                \
+      itype,                                                  \
+      otype,                                                  \
+      bm,                                                     \
+      bn,                                                     \
+      bk,                                                     \
+      wm,                                                     \
+      wn,                                                     \
+      trans_a,                                                \
+      trans_b,                                                \
+      mn_aligned,                                             \
+      k_aligned)
+
+#define instantiate_gemm_aligned_helper(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn)             \
+  instantiate_gemm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, taligned, true, taligned, true)  \
+  instantiate_gemm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, taligned, true, naligned, false) \
+  instantiate_gemm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, naligned, false, taligned, true) \
+  instantiate_gemm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn, naligned, false, naligned, false)
+
+#define instantiate_gemm_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn)             \
+    instantiate_gemm_aligned_helper(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_aligned_helper(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_aligned_helper(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_aligned_helper(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gemm_shapes_helper(iname, itype, oname, otype)                  \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 16, 16, 16, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 16, 32, 16, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 16, 16, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 32, 32, 16, 2, 2)
+
+instantiate_gemm_shapes_helper(float16, half, float32, float);
+instantiate_gemm_shapes_helper(bfloat16, bfloat16_t, float32, float);
+instantiate_gemm_shapes_helper(float32, float, float32, float);
+instantiate_gemm_shapes_helper(complex64, complex64_t, complex64, complex64_t);
+
+#define instantiate_accum(oname, otype, aname, atype)      \
+  instantiate_kernel(                                      \
+    "steel_gemm_splitk_accum_" #oname "_" #aname,          \
+    gemm_splitk_accum, atype, otype)                       \
+  instantiate_kernel(                                      \
+    "steel_gemm_splitk_accum_" #oname "_" #aname "_axbpy", \
+  gemm_splitk_accum_axpby, atype, otype)                   \
+
+instantiate_accum(bfloat16, bfloat16_t, float32, float);
+instantiate_accum(float16, half, float32, float);
+instantiate_accum(float32, float, float32, float);
+instantiate_accum(complex64, complex64_t, complex64, complex64_t); // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h
new file mode 100644
index 0000000..1b6b828
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h
@@ -0,0 +1,152 @@
+// Copyright © 2026 Apple Inc.
+
+using namespace mlx::steel;
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+
+///////////////////////////////////////////////////////////////////////////////
+// NAX Split-K GEMM kernel
+///////////////////////////////////////////////////////////////////////////////
+
+// clang-format off
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gemm_splitk_nax(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device AccumType* C [[buffer(2)]],
+    const constant GEMMSpiltKParams* params [[buffer(3)]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on
+
+  const int linear_tid = tid.x;
+
+  // Compute swizzled tile dimensions
+  const int tn_swizzled = params->tiles_n << params->swizzle_log;
+  const int tm_swizzled =
+      (params->tiles_m + (1 << params->swizzle_log) - 1) >> params->swizzle_log;
+  const int tiles_per_partition = tn_swizzled * tm_swizzled;
+
+  const int tid_z = linear_tid / tiles_per_partition;
+  const int xy_flat = linear_tid % tiles_per_partition;
+
+  // Decode 2D grid coordinates in swizzled space
+  const int grid_x = xy_flat % tn_swizzled;
+  const int grid_y = xy_flat / tn_swizzled;
+
+  // Apply X-Y swizzle
+  const int tid_y = (grid_y << params->swizzle_log) +
+      (grid_x & ((1 << params->swizzle_log) - 1));
+  const int tid_x = grid_x >> params->swizzle_log;
+
+  // Exit early
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  // Calculate partition bounds
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int k_start = params->split_k_partition_size * tid_z;
+  const int k_end = min(k_start + params->split_k_partition_size, params->K);
+
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+  const size_t k_start_long = size_t(k_start);
+
+  // Adjust pointers for split-K partition
+  A += transpose_a ? (c_row_long + k_start_long * params->lda)
+                   : (k_start_long + c_row_long * params->lda);
+  B += transpose_b ? (k_start_long + c_col_long * params->ldb)
+                   : (c_col_long + k_start_long * params->ldb);
+  C += (size_t(params->split_k_partition_stride) * tid_z) +
+      (c_row_long * params->ldc + c_col_long);
+
+  // NAX tile configuration
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+
+  // Calculate simdgroup offsets and alignment
+  const short tm = SM * (simd_group_id / WN);
+  const short tn = SN * (simd_group_id % WN);
+
+  const int sgp_sm_int =
+      align_M ? int(SM) : min(int(SM), params->M - (c_row + tm));
+  const short sgp_sm = short(sgp_sm_int);
+  const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
+
+  const int sgp_sn_int =
+      align_N ? int(SN) : min(int(SN), params->N - (c_col + tn));
+  const short sgp_sn = short(sgp_sn_int);
+  const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN);
+
+  A += transpose_a ? tm : (tm * params->lda);
+  B += transpose_b ? (tn * params->ldb) : tn;
+  C += tm * params->ldc + tn;
+
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  // gemm_loop through the partition
+  // Check K-alignment at runtime (partition-specific)
+  const int partition_k_size = k_end - k_start;
+  const int partition_k_iters = partition_k_size / BK;
+  const bool partition_k_aligned = (partition_k_size % BK) == 0;
+
+  dispatch_bool(partition_k_aligned, [&](auto kAlignedK) {
+    dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+      dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) {
+        Dtile = gemm_loop<
+            T,
+            SM,
+            SN,
+            SK,
+            BK,
+            transpose_a,
+            transpose_b,
+            kAlignedM.value,
+            kAlignedN.value,
+            kAlignedK.value,
+            UM,
+            UN,
+            UK,
+            AccumType>(
+            A,
+            B,
+            params->lda,
+            params->ldb,
+            partition_k_size,
+            partition_k_iters,
+            sgp_sm,
+            sgp_sn);
+      });
+    });
+  });
+
+  // Store result
+  dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+    dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) {
+      if constexpr (kAlignedM && kAlignedN) {
+        Dtile.store(C, int(params->ldc));
+      } else {
+        Dtile.store_safe(C, int(params->ldc), short2(sgp_sn, sgp_sm));
+      }
+    });
+  });
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal
new file mode 100644
index 0000000..b1b6b89
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.metal
@@ -0,0 +1,30 @@
+// Copyright © 2026 Apple Inc.
+
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h"
+
+// clang-format off
+#define instantiate_gemm_splitk(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                                             \
+      "steel_gemm_splitk_nax_" #tname "_"  #iname "_" #oname                                      \
+      "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn,                                          \
+  gemm_splitk_nax, itype, bm, bn, bk, wm, wn, trans_a, trans_b, float)
+
+#define instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gemm_splitk_shapes_helper(iname, itype, oname, otype) \
+    instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype,  64,  64, 256, 2, 2) \
+    instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, 128, 128, 512, 4, 4)
+
+instantiate_gemm_splitk_shapes_helper(float16, half, float32, float);
+instantiate_gemm_splitk_shapes_helper(bfloat16, bfloat, float32, float);
+instantiate_gemm_splitk_shapes_helper(float32, float, float32, float);
+// clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/loader.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/loader.h
new file mode 100644
index 0000000..d421b2d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/loader.h
@@ -0,0 +1,137 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short alignment = 1,
+    short n_reads = (BCOLS * BROWS) / (tgp_size),
+    short TCOLS = BCOLS / n_reads,
+    short TROWS = tgp_size / TCOLS>
+struct BlockLoader {
+  STEEL_CONST short n_rows = (BROWS + TROWS - 1) / TROWS;
+  STEEL_CONST short vec_size = n_reads;
+
+  // Leading dimension for src
+  const int src_ld;
+  const int tile_stride;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  struct alignas(alignment * sizeof(T)) ReadVector {
+    uint8_t v[sizeof(T) * vec_size];
+  };
+
+  /* Constructor */
+  METAL_FUNC BlockLoader(
+      const device T* src_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(reduction_dim ? BCOLS : BROWS * src_ld),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj) {}
+
+  /* Apply operation to threadgroup without bound checking */
+  template <typename UnaryOp>
+  METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = op.apply(dst[i * dst_ld + j]);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      *((threadgroup ReadVector*)(&dst[i * dst_ld])) =
+          *((const device ReadVector*)(&src[i * src_ld]));
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - with bound checking */
+  METAL_FUNC void load_safe(short2 src_tile_dim) const {
+    src_tile_dim = src_tile_dim - short2(bj, bi);
+
+    // Skip loading if thread has no valid reads
+    if (src_tile_dim.x <= 0 || src_tile_dim.y <= 0) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+      return;
+    }
+
+    // Use fast thread memory for bound checks
+    bool tmp_idx[vec_size];
+    T tmp_val[vec_size];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      // Make sure tmp_idx only contains valid indices
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_idx[j] = (i < src_tile_dim.y) && (j < src_tile_dim.x);
+      }
+
+      // Read valid indices into tmp_val
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = src[(tmp_idx[j] ? i * src_ld + j : 0)];
+      }
+
+      // Zero out unneeded values
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = tmp_idx[j] ? tmp_val[j] : T(0);
+      }
+
+      // Copy values to threadgroup memory
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = tmp_val[j];
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    src += tile_stride;
+  }
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/mma.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/mma.h
new file mode 100644
index 0000000..1131994
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/mma.h
@@ -0,0 +1,1146 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_simdgroup>
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+#include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
+#include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// MMA helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <typename T, int kFragRows_, int kFragCols_>
+struct BaseMMAFrag {
+  static_assert(
+      kFragRows_ == 8,
+      "Only 8 x 8 fragment matrices are currently supported");
+  static_assert(
+      kFragCols_ == 8,
+      "Only 8 x 8 fragment matrices are currently supported");
+};
+
+template <typename T>
+struct BaseMMAFrag<T, 8, 8> {
+  STEEL_CONST int kFragRows = 8;
+  STEEL_CONST int kFragCols = 8;
+
+  STEEL_CONST int kElemsPerFrag = (kFragRows * kFragCols) / 32;
+
+  STEEL_CONST int kElemRows = 1;
+  STEEL_CONST int kElemCols = 2;
+
+  static_assert(
+      kElemRows * kElemCols == kElemsPerFrag,
+      "MMAFrag shape is not consistent with MMAFrag size");
+
+  typedef metal::simdgroup_matrix<T, kFragRows, kFragCols> mat_type;
+  typedef metal::vec<T, kElemsPerFrag> frag_type;
+
+  METAL_FUNC static constexpr short2 get_coord(
+      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
+    const short qid = simd_lane_id / 4;
+    const short fm = (qid & 4) + ((simd_lane_id / 2) % 4);
+    const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
+    return short2{fn, fm};
+  }
+
+  template <typename SrcPtrType, typename StrX, typename StrY>
+  METAL_FUNC static constexpr void
+  load(thread frag_type& dst, SrcPtrType src, StrX str_x, StrY str_y) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        dst[i * kElemCols + j] = static_cast<T>(src[i * str_x + j * str_y]);
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX,
+      typename OffY>
+  METAL_FUNC static constexpr void load_safe(
+      thread frag_type& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
+          dst[i * kElemCols + j] =
+              static_cast<T>(src[(off_x + i) * str_x + (off_x + j) * str_y]);
+        } else {
+          dst[i * kElemCols + j] = T(0);
+        }
+      }
+    }
+  }
+
+  template <typename DstPtrType, typename StrX, typename StrY>
+  METAL_FUNC static constexpr void
+  store(const thread frag_type& src, DstPtrType dst, StrX str_x, StrY str_y) {
+    using U = pointer_element_t<DstPtrType>;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        dst[i * str_x + j * str_y] = static_cast<U>(src[i * kElemCols + j]);
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX,
+      typename OffY>
+  METAL_FUNC static constexpr void store_safe(
+      const thread frag_type& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if ((off_x + i) < lim_x && (off_y + j) < lim_y) {
+          dst[(off_x + i) * str_x + (off_y + j) * str_y] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename StartX,
+      typename StopX,
+      typename StartY,
+      typename StopY,
+      typename OffX,
+      typename OffY>
+  METAL_FUNC static constexpr void store_slice(
+      const thread frag_type& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      StartX start_x,
+      StopX stop_x,
+      StartY start_y,
+      StopY stop_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if ((off_x + i) < stop_x && (off_x + i) >= start_x &&
+            (off_y + j) < stop_y && (off_y + j) >= start_y) {
+          dst[(off_x + i) * str_x + (off_y + j) * str_y] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  METAL_FUNC static constexpr void mma(
+      thread frag_type& D,
+      thread frag_type& A,
+      thread frag_type& B,
+      thread frag_type& C) {
+    mat_type D_mat;
+    mat_type A_mat;
+    mat_type B_mat;
+    mat_type C_mat;
+
+    reinterpret_cast<thread frag_type&>(A_mat.thread_elements()) = A;
+    reinterpret_cast<thread frag_type&>(B_mat.thread_elements()) = B;
+    reinterpret_cast<thread frag_type&>(C_mat.thread_elements()) = C;
+
+    mma(D_mat, A_mat, B_mat, C_mat);
+
+    D = reinterpret_cast<thread frag_type&>(D_mat.thread_elements());
+  }
+
+  METAL_FUNC static constexpr void mma(
+      thread mat_type& D,
+      thread mat_type& A,
+      thread mat_type& B,
+      thread mat_type& C) {
+    simdgroup_multiply_accumulate(D, A, B, C);
+  }
+};
+
+template <
+    typename T,
+    int kTileRows_,
+    int kTileCols_,
+    class MMAFrag_ = BaseMMAFrag<T, 8, 8>>
+struct MMATile {
+  using MMAFrag_t = MMAFrag_;
+  using elem_type = T;
+  STEEL_CONST int kFragRows = MMAFrag_t::kFragRows;
+  STEEL_CONST int kFragCols = MMAFrag_t::kFragCols;
+  STEEL_CONST int kElemsPerFrag = MMAFrag_t::kElemsPerFrag;
+
+  STEEL_CONST int kTileRows = kTileRows_;
+  STEEL_CONST int kTileCols = kTileCols_;
+
+  STEEL_CONST int kRows = kTileRows * kFragRows;
+  STEEL_CONST int kCols = kTileCols * kFragCols;
+
+  STEEL_CONST int kNumFrags = kTileRows * kTileCols;
+  STEEL_CONST int kElemsPerTile = kNumFrags * kElemsPerFrag;
+
+  typedef typename MMAFrag_t::mat_type mat_type;
+  typedef typename MMAFrag_t::frag_type frag_type;
+
+  frag_type val_frags[kNumFrags] = {frag_type(0)};
+
+  METAL_FUNC MMATile() thread {}
+
+  METAL_FUNC constexpr void clear() {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kNumFrags; ++i) {
+      val_frags[i] = frag_type(0);
+    }
+  }
+
+  METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) {
+    return val_frags[i * kTileCols + j];
+  }
+
+  METAL_FUNC constexpr const thread frag_type& frag_at(
+      const short i,
+      const short j) const {
+    return val_frags[i * kTileCols + j];
+  }
+
+  METAL_FUNC mat_type mat_at(const short i, const short j) {
+    mat_type val_mat;
+    STEEL_PRAGMA_UNROLL
+    for (short ii = 0; ii < kElemsPerFrag; ++ii) {
+      val_mat.thread_elements()[ii] = frag_at(i, j)[ii];
+    }
+    return val_mat;
+  }
+
+  METAL_FUNC thread elem_type* elems() {
+    return reinterpret_cast<thread elem_type*>(val_frags);
+  }
+
+  METAL_FUNC const thread elem_type* elems() const {
+    return reinterpret_cast<const thread elem_type*>(val_frags);
+  }
+
+  template <typename U, int w_x, int w_y, int str_x, int str_y>
+  METAL_FUNC void load(const threadgroup U* src) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load(
+            frag_at(i, j),
+            &(
+                src[(i * kFragRows) * w_x * str_x +
+                    (j * kFragCols) * w_y * str_y]),
+            Int<str_x>{},
+            Int<str_y>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y, int str_x, int str_y>
+  METAL_FUNC void store(threadgroup U* dst) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store(
+            frag_at(i, j),
+            &(
+                dst[(i * kFragRows) * w_x * str_x +
+                    (j * kFragCols) * w_y * str_y]),
+            Int<str_x>{},
+            Int<str_y>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void load(const device U* src, const int ld) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load(
+            frag_at(i, j),
+            &(src[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
+            ld,
+            Int<1>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void store(device U* dst, const int ld) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store(
+            frag_at(i, j),
+            &(dst[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]),
+            ld,
+            Int<1>{});
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void
+  load_safe(const device U* src, const int ld, const short2 src_tile_dims) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::load_safe(
+            frag_at(i, j),
+            src,
+            ld,
+            Int<1>{},
+            src_tile_dims.y,
+            src_tile_dims.x,
+            (i * kFragRows) * w_x,
+            (j * kFragCols) * w_y);
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void
+  store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store_safe(
+            frag_at(i, j),
+            dst,
+            ld,
+            Int<1>{},
+            dst_tile_dims.y,
+            dst_tile_dims.x,
+            (i * kFragRows) * w_x,
+            (j * kFragCols) * w_y);
+      }
+    }
+  }
+
+  template <typename U, int w_x, int w_y>
+  METAL_FUNC void store_slice(
+      device U* dst,
+      const int ld,
+      const short2 start,
+      const short2 stop) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        MMAFrag_t::store_slice(
+            frag_at(i, j),
+            dst,
+            ld,
+            Int<1>{},
+            start.y,
+            stop.y,
+            start.x,
+            stop.x,
+            (i * kFragRows) * w_x,
+            (j * kFragCols) * w_y);
+      }
+    }
+  }
+};
+
+template <typename T, typename U, int M, int N, int K>
+METAL_FUNC void tile_matmad(
+    thread MMATile<T, M, N>& D,
+    thread MMATile<U, M, K>& A,
+    thread MMATile<U, K, N>& B,
+    thread MMATile<T, M, N>& C) {
+  STEEL_PRAGMA_UNROLL
+  for (short m = 0; m < M; ++m) {
+    STEEL_PRAGMA_UNROLL
+    for (short n = 0; n < N; ++n) {
+      short n_serp = (m % 2) ? (N - 1 - n) : n;
+      STEEL_PRAGMA_UNROLL
+      for (short k = 0; k < K; ++k) {
+        MMATile<T, M, N>::MMAFrag_t::mma(
+            D.frag_at(m, n_serp),
+            A.frag_at(m, k),
+            B.frag_at(k, n_serp),
+            C.frag_at(m, n_serp));
+      }
+    }
+  }
+}
+
+template <typename InT>
+struct TransformNone<complex64_t, InT> {
+  static METAL_FUNC complex64_t apply(complex64_t x) {
+    return x;
+  }
+  static METAL_FUNC complex64_t apply(complex64_t x, complex64_t) {
+    return x;
+  }
+};
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    short lda_tgp,
+    short ldb_tgp,
+    typename AccumType = float,
+    typename Epilogue = TransformNone<U, AccumType>>
+struct BlockMMA {
+  // MMAFrag size
+  STEEL_CONST short kFragSize = 8;
+  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
+
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TM_stride = kFragSize * WM;
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TN_stride = kFragSize * WN;
+
+  // Warp tile size along M
+  STEEL_CONST short TM = BM / (kFragSize * WM);
+  // Warp tile size along N
+  STEEL_CONST short TN = BN / (kFragSize * WN);
+
+  // Threadgroup A strides
+  STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M
+  STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K
+
+  // Threadgroup B strides
+  STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K
+  STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N
+
+  // Threadgroup strides along K
+  STEEL_CONST short tile_stride_a = kFragSize * A_str_k;
+  STEEL_CONST short tile_stride_b = kFragSize * B_str_k;
+
+  // Simdgroup matrices
+  MMATile<AccumType, TM, 1, MMAFrag_acc_t> Atile;
+  MMATile<AccumType, 1, TN, MMAFrag_acc_t> Btile;
+  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile;
+
+  // Offsets within threadgroup
+  short sm;
+  short sn;
+
+  short As_offset;
+  short Bs_offset;
+
+  /* Constructor */
+  METAL_FUNC BlockMMA(
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
+    // Determine thread position in simdgroup matrix
+    short tm = kFragSize * (simd_group_id / WN);
+    short tn = kFragSize * (simd_group_id % WN);
+
+    short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
+    sm = simd_coord.y;
+    sn = simd_coord.x;
+
+    // Determine thread and simdgroup offset
+    As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // M, K
+    Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // K, N
+
+    sm += tm;
+    sn += tn;
+  }
+
+  /* (BM, BK) X (BK, BN) multiply accumulate function */
+  METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) {
+    // Adjust for simdgroup and thread location
+    As += As_offset;
+    Bs += Bs_offset;
+
+    // Iterate over BK in blocks of kFragSize
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < BK; kk += kFragSize) {
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Atile.template load<T, WM, 1, A_str_m, A_str_k>(As);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      Btile.template load<T, 1, WN, B_str_k, B_str_n>(Bs);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      tile_matmad(Ctile, Atile, Btile, Ctile);
+
+      // Progress to next simdgroup tile
+      As += tile_stride_a;
+      Bs += tile_stride_b;
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(device U* D, const int ldd) {
+    // Apply epilogue
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
+    }
+
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+
+    Ctile.template store<U, WM, WN>(D, ldd);
+  }
+
+  METAL_FUNC void
+  store_result_slice(device U* D, const int ldd, short2 start, short2 stop) {
+    // Apply epilogue
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
+    }
+
+    D += sm * ldd + sn;
+    start -= short2(sn, sm);
+    stop -= short2(sn, sm);
+
+    // TODO: Check the start as well
+    if (stop.y <= 0 || stop.x <= 0) {
+      return;
+    }
+
+    Ctile.template store_slice<U, WM, WN>(D, ldd, start, stop);
+  }
+
+  METAL_FUNC void
+  store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) {
+    // Apply epilogue
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]);
+    }
+
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    Ctile.template store_safe<U, WM, WN>(D, ldd, dst_tile_dims);
+  }
+
+  /* Apply epilogue */
+  template <typename UnaryEpilogue>
+  METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) {
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) {
+      Ctile.elems()[i] = epilogue_op.apply(Ctile.elems()[i]);
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < decltype(Ctile)::kElemsPerFrag; k++) {
+          accum[k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue_safe(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+        // Read C
+        U c_elems[kelems] = {0};
+
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          if ((j * TN_stride + k) < dst_tile_dims.x) {
+            c_elems[k] = C[offset_c + k * fdc];
+          }
+        }
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          accum[k] = epilogue_op.apply(accum[k], c_elems[k]);
+        }
+      }
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+
+    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in C
+        thread const auto& accum = Ctile.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+        int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          D[offset_d + k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void store_result_safe(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    constexpr short kelems = decltype(Ctile)::kElemsPerFrag;
+
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < TM; i++) {
+      if (i * TM_stride < dst_tile_dims.y) {
+        STEEL_PRAGMA_UNROLL
+        for (int j = 0; j < TN; j++) {
+          // Get accumulated result and associated offset in C
+          thread const auto& accum = Ctile.frag_at(i, j);
+          int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+          int offset_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+          // Apply epilogue
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < kelems; k++) {
+            if ((j * TN_stride + k) < dst_tile_dims.x) {
+              D[offset_d + k] =
+                  epilogue_op.apply(accum[k], C[offset_c + k * fdc]);
+            }
+          }
+        }
+      }
+    }
+  }
+};
+
+template <
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    short lda_tgp,
+    short ldb_tgp,
+    typename AccumType,
+    typename Epilogue>
+struct BlockMMA<
+    complex64_t,
+    U,
+    BM,
+    BN,
+    BK,
+    WM,
+    WN,
+    transpose_a,
+    transpose_b,
+    lda_tgp,
+    ldb_tgp,
+    AccumType,
+    Epilogue> {
+  static_assert(
+      metal::is_same_v<AccumType, float>,
+      "BlockMMA<complex64_t,...> expects float accumulators");
+  static_assert(
+      metal::is_same_v<U, complex64_t>,
+      "For complex BlockMMA, U must be complex64_t; use a different epilogue for projections");
+  // MMAFrag size
+  STEEL_CONST short kFragSize = 8;
+  using MMAFrag_acc_t = BaseMMAFrag<AccumType, kFragSize, kFragSize>;
+
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TM_stride = kFragSize * WM;
+  // Warp tile simdgroup matrix strides along M
+  STEEL_CONST short TN_stride = kFragSize * WN;
+
+  // Warp tile size along M
+  STEEL_CONST short TM = BM / (kFragSize * WM);
+  // Warp tile size along N
+  STEEL_CONST short TN = BN / (kFragSize * WN);
+
+  // Threadgroup A strides
+  STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M
+  STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K
+
+  // Threadgroup B strides
+  STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K
+  STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N
+
+  // Threadgroup strides along K
+  STEEL_CONST short tile_stride_a = kFragSize * A_str_k;
+  STEEL_CONST short tile_stride_b = kFragSize * B_str_k;
+
+  // When indexing complex as float[2]
+  STEEL_CONST short A_str_m_f = A_str_m * 2;
+  STEEL_CONST short A_str_k_f = A_str_k * 2;
+  STEEL_CONST short B_str_k_f = B_str_k * 2;
+  STEEL_CONST short B_str_n_f = B_str_n * 2;
+  STEEL_CONST short tile_stride_a_f = tile_stride_a * 2;
+  STEEL_CONST short tile_stride_b_f = tile_stride_b * 2;
+
+  // Accumulators (real/imag)
+  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile_r;
+  MMATile<AccumType, TM, TN, MMAFrag_acc_t> Ctile_i;
+
+  // Offsets within threadgroup
+  short sm, sn;
+  short As_offset, Bs_offset;
+
+  /* Constructor */
+  METAL_FUNC BlockMMA(
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]]) {
+    // Determine thread position in simdgroup matrix
+    short tm = kFragSize * (simd_group_id / WN);
+    short tn = kFragSize * (simd_group_id % WN);
+
+    short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id);
+    sm = simd_coord.y;
+    sn = simd_coord.x;
+
+    // Determine thread and simdgroup offset
+    As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // (M,K)
+    Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // (K,N)
+
+    sm += tm;
+    sn += tn;
+  }
+
+  /* Karatsuba MMA: 3 real MMAs per K-chunk */
+  METAL_FUNC void mma(
+      const threadgroup complex64_t* As,
+      const threadgroup complex64_t* Bs) {
+    // Adjust for simdgroup and thread location
+    As += As_offset;
+    Bs += Bs_offset;
+    threadgroup const float* As_f =
+        reinterpret_cast<threadgroup const float*>(As);
+    threadgroup const float* Bs_f =
+        reinterpret_cast<threadgroup const float*>(Bs);
+
+    // Iterate over BK in blocks of kFragSize
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < BK; kk += kFragSize) {
+      simdgroup_barrier(mem_flags::mem_none);
+
+      MMATile<AccumType, TM, 1, MMAFrag_acc_t> Ar, Ai;
+      Ar.template load<float, WM, 1, A_str_m_f, A_str_k_f>(As_f + 0);
+      Ai.template load<float, WM, 1, A_str_m_f, A_str_k_f>(As_f + 1);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      MMATile<AccumType, 1, TN, MMAFrag_acc_t> Br, Bi;
+      Br.template load<float, 1, WN, B_str_k_f, B_str_n_f>(Bs_f + 0);
+      Bi.template load<float, 1, WN, B_str_k_f, B_str_n_f>(Bs_f + 1);
+
+      simdgroup_barrier(mem_flags::mem_none);
+
+      // P = Ar*Br ; Q = Ai*Bi ; R = (Ar+Ai)*(Br+Bi)
+      MMATile<AccumType, TM, TN, MMAFrag_acc_t> P, Q, R;
+
+      tile_matmad(P, Ar, Br, P);
+      tile_matmad(Q, Ai, Bi, Q);
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < decltype(Ar)::kElemsPerTile; ++i)
+        Ar.elems()[i] += Ai.elems()[i];
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < decltype(Br)::kElemsPerTile; ++i)
+        Br.elems()[i] += Bi.elems()[i];
+
+      tile_matmad(R, Ar, Br, R);
+
+      // C_r += P - Q ; C_i -= Q
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; ++i) {
+        const auto p = P.elems()[i];
+        const auto q = Q.elems()[i];
+        const auto r = R.elems()[i];
+        Ctile_r.elems()[i] += (p - q);
+        Ctile_i.elems()[i] += (r - p - q);
+      }
+
+      // Progress to next simdgroup tile
+      As_f += tile_stride_a_f;
+      Bs_f += tile_stride_b_f;
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(device U* D, const int ldd) {
+    // Adjust for simdgroup and thread location
+    D += sm * ldd + sn;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        thread const auto& r = Ctile_r.frag_at(i, j);
+        thread const auto& im = Ctile_i.frag_at(i, j);
+        int off = (i * TM_stride) * ldd + (j * TN_stride);
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) {
+          D[off + k] = Epilogue::apply(complex64_t(r[k], im[k]));
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void
+  store_result_slice(device U* D, const int ldd, short2 start, short2 stop) {
+    D += sm * ldd + sn;
+    start -= short2(sn, sm);
+    stop -= short2(sn, sm);
+
+    if (stop.y <= 0 || stop.x <= 0)
+      return;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; ++i) {
+      const int row = i * TM_stride;
+      if (row >= start.y && row < stop.y) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < TN; ++j) {
+          const int off = row * ldd + (j * TN_stride);
+          thread const auto& r = Ctile_r.frag_at(i, j);
+          thread const auto& im = Ctile_i.frag_at(i, j);
+
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; ++k) {
+            const int col = j * TN_stride + k;
+            if (col >= start.x && col < stop.x) {
+              D[off + k] = Epilogue::apply(complex64_t(r[k], im[k]));
+            }
+          }
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void
+  store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) {
+    D += sm * ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      if (i * TM_stride < dst_tile_dims.y) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < TN; j++) {
+          int off = (i * TM_stride) * ldd + (j * TN_stride);
+          thread const auto& r = Ctile_r.frag_at(i, j);
+          thread const auto& im = Ctile_i.frag_at(i, j);
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) {
+            if ((j * TN_stride + k) < dst_tile_dims.x) {
+              D[off + k] = Epilogue::apply(complex64_t(r[k], im[k]));
+            }
+          }
+        }
+      }
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename UnaryEpilogue>
+  METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; i++) {
+      complex64_t out = epilogue_op.apply(
+          complex64_t(Ctile_r.elems()[i], Ctile_i.elems()[i]));
+      Ctile_r.elems()[i] = out.real;
+      Ctile_i.elems()[i] = out.imag;
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in Cr, Ci
+        thread auto& r = Ctile_r.frag_at(i, j);
+        thread auto& im = Ctile_i.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) {
+          complex64_t out = epilogue_op.apply(
+              complex64_t(r[k], im[k]), C[offset_c + k * fdc]);
+          r[k] = out.real;
+          im[k] = out.imag;
+        }
+      }
+    }
+  }
+
+  /* Apply epilogue */
+  template <typename BinaryEpilogue>
+  METAL_FUNC void apply_epilogue_safe(
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const BinaryEpilogue& epilogue_op) {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in Cr, Ci
+        thread auto& r = Ctile_r.frag_at(i, j);
+        thread auto& im = Ctile_i.frag_at(i, j);
+        int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+
+        constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag;
+        complex64_t tmp[kelems];
+
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          if ((j * TN_stride + k) < dst_tile_dims.x &&
+              (i * TM_stride) < dst_tile_dims.y) {
+            tmp[k] = C[offset_c + k * fdc];
+          } else {
+            tmp[k] = complex64_t(0.0f, 0.0f);
+          }
+        }
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          complex64_t out = epilogue_op.apply(complex64_t(r[k], im[k]), tmp[k]);
+          r[k] = out.real;
+          im[k] = out.imag;
+        }
+      }
+    }
+  }
+
+  /* Store results from simdgroup_matrix results into device memory */
+  METAL_FUNC void store_result(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+
+    constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag;
+
+    // Loop over all simdgroup tiles
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < TM; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < TN; j++) {
+        // Get accumulated result and associated offset in Cr, Ci
+        thread const auto& r = Ctile_r.frag_at(i, j);
+        thread const auto& im = Ctile_i.frag_at(i, j);
+        int off_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+        int off_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+        // Apply epilogue
+        STEEL_PRAGMA_UNROLL
+        for (short k = 0; k < kelems; k++) {
+          D[off_d + k] =
+              epilogue_op.apply(complex64_t(r[k], im[k]), C[off_c + k * fdc]);
+        }
+      }
+    }
+  }
+
+  METAL_FUNC void store_result_safe(
+      device U* D,
+      const int ldd,
+      const device U* C,
+      const int ldc,
+      const int fdc,
+      short2 dst_tile_dims,
+      thread const Epilogue& epilogue_op) const {
+    // Adjust for simdgroup and thread location
+    C += (sm)*ldc + (sn)*fdc;
+    D += (sm)*ldd + sn;
+    dst_tile_dims -= short2(sn, sm);
+
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
+    constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag;
+
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < TM; i++) {
+      if (i * TM_stride < dst_tile_dims.y) {
+        STEEL_PRAGMA_UNROLL
+        for (int j = 0; j < TN; j++) {
+          // Get accumulated result and associated offset in Cr, Ci
+          thread const auto& r = Ctile_r.frag_at(i, j);
+          thread const auto& im = Ctile_i.frag_at(i, j);
+          int off_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc;
+          int off_d = (i * TM_stride) * ldd + (j * TN_stride);
+
+          // Apply epilogue
+          STEEL_PRAGMA_UNROLL
+          for (short k = 0; k < kelems; k++) {
+            if ((j * TN_stride + k) < dst_tile_dims.x) {
+              D[off_d + k] = epilogue_op.apply(
+                  complex64_t(r[k], im[k]), C[off_c + k * fdc]);
+            }
+          }
+        }
+      }
+    }
+  }
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/nax.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/nax.h
new file mode 100644
index 0000000..5839176
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/nax.h
@@ -0,0 +1,1084 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <metal_simdgroup>
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+#include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
+#include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"
+
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// MMA helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+///////////////////////////////////////////////////////////////////////////////
+// NAX Steel with new tiles
+///////////////////////////////////////////////////////////////////////////////
+
+struct BaseNAXFrag {
+  STEEL_CONST short kFragRows = 16;
+  STEEL_CONST short kFragCols = 16;
+
+  STEEL_CONST short kElemsPerFrag = (kFragRows * kFragCols) / 32;
+
+  STEEL_CONST short kElemRows = 2;
+  STEEL_CONST short kElemCols = 4;
+
+  STEEL_CONST short kElemRowsJump = 8;
+
+  static_assert(
+      kElemRows * kElemCols == kElemsPerFrag,
+      "MMAFrag shape is not consistent with MMAFrag size");
+
+  template <typename U>
+  using dtype_frag_t = typename metal::vec<U, kElemsPerFrag>;
+
+  METAL_FUNC static short2 get_coord() {
+    const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort());
+    const short qid = simd_lane_id >> 2;
+    const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3));
+    const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4;
+    return short2{fn, fm};
+  }
+
+  METAL_FUNC static short2 get_coord(short idx) {
+    const ushort simd_lane_id = __metal_get_thread_index_in_simdgroup(ushort());
+    const short qid = simd_lane_id >> 2;
+    const short fm = ((qid & 4) | ((simd_lane_id >> 1) & 3)) + (idx >> 2) * 8;
+    const short fn = ((qid & 2) | (simd_lane_id & 1)) * 4 + idx % 4;
+    return short2{fn, fm};
+  }
+
+  template <
+      typename T,
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void load(
+      thread dtype_frag_t<T>& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    const short2 sc = get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if constexpr (metal::is_same_v<StrY, Int<1>>) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[i * kElemCols + j] = static_cast<T>(src[r * str_x + c + j]);
+        }
+      } else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[i * kElemCols + j] =
+              static_cast<T>(src[r * str_x + (c + j) * str_y]);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void load_rows(
+      thread dtype_frag_t<T>& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    const short2 sc = get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if (r < lim_x) {
+        if constexpr (metal::is_same_v<StrY, Int<1>>) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[i * kElemCols + j] = static_cast<T>(src[r * str_x + (c + j)]);
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[i * kElemCols + j] =
+                static_cast<T>(src[r * str_x + (c + j) * str_y]);
+          }
+        }
+
+      } else {
+        dst = dtype_frag_t<T>(0);
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void load_safe(
+      thread dtype_frag_t<T>& dst,
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    const short2 sc = get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if (r < lim_x && (c + j) < lim_y) {
+          dst[i * kElemCols + j] =
+              static_cast<T>(src[r * str_x + (c + j) * str_y]);
+        } else {
+          dst[i * kElemCols + j] = T(0);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if constexpr (metal::is_same_v<StrY, Int<1>>) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[r * str_x + c + j] = static_cast<U>(src[i * kElemCols + j]);
+        }
+      } else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < kElemCols; j++) {
+          dst[r * str_x + (c + j) * str_y] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store_rows(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      if (r < lim_x) {
+        if constexpr (metal::is_same_v<StrY, Int<1>>) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[r * str_x + c + j] = static_cast<U>(src[i * kElemCols + j]);
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < kElemCols; j++) {
+            dst[r * str_x + (c + j) * str_y] =
+                static_cast<U>(src[i * kElemCols + j]);
+          }
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store_safe(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = get_coord();
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      const auto r = off_x + i * kElemRowsJump + sc.y;
+      const auto c = off_y + sc.x;
+
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        if (r < lim_x && (c + j) < lim_y) {
+          dst[r * str_x + (c + j) * str_y] =
+              static_cast<U>(src[i * kElemCols + j]);
+        }
+      }
+    }
+  }
+
+  template <
+      typename T,
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename StartX,
+      typename StopX,
+      typename StartY,
+      typename StopY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC static constexpr void store_slice(
+      const thread dtype_frag_t<T>& src,
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      StartX start_x,
+      StopX stop_x,
+      StartY start_y,
+      StopY stop_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) {
+    using U = pointer_element_t<DstPtrType>;
+
+    const short2 sc = get_coord();
+
+    const_for_loop<0, kElemRows, 1>([&](auto idx_row) {
+      const auto r = off_x + idx_row * Int<kElemRowsJump>{};
+      if (r >= stop_x - sc.y || r < start_x - sc.y) {
+        return;
+      }
+
+      const_for_loop<0, kElemCols, 1>([&](auto idx_col) {
+        const auto c = off_y + idx_col;
+        if (c >= stop_y - sc.x || c < start_y - sc.x) {
+          return;
+        }
+
+        const auto src_idx = idx_row * Int<kElemCols>{} + idx_col;
+        dst[(r + sc.y) * str_x + (c + sc.x) * str_y] =
+            static_cast<U>(src[src_idx]);
+      });
+    });
+  }
+
+  template <typename Op, typename T>
+  METAL_FUNC static constexpr void row_reduce(
+      thread const dtype_frag_t<T>& inp_vals,
+      thread T* reduced_vals) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      T thr_reduce = Op::apply(
+          Op::apply(inp_vals[i * kElemCols + 0], inp_vals[i * kElemCols + 1]),
+          Op::apply(inp_vals[i * kElemCols + 2], inp_vals[i * kElemCols + 3]));
+
+      T qgr_reduce = simd_shuffle_xor(thr_reduce, ushort(1));
+      qgr_reduce = Op::apply(thr_reduce, qgr_reduce);
+
+      T sgr_reduce = simd_shuffle_xor(qgr_reduce, ushort(8));
+      sgr_reduce = Op::apply(qgr_reduce, sgr_reduce);
+
+      reduced_vals[i] = Op::apply(reduced_vals[i], sgr_reduce);
+    }
+  }
+
+  template <typename Op, typename T>
+  METAL_FUNC static constexpr void row_bin_op(
+      thread dtype_frag_t<T>& inp_vals,
+      thread T* row_vals) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kElemRows; i++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kElemCols; j++) {
+        inp_vals[i * kElemCols + j] =
+            Op::apply(inp_vals[i * kElemCols + j], row_vals[i]);
+      }
+    }
+  }
+};
+
+template <
+    typename T,
+    short kRows_,
+    short kCols_,
+    typename NAXFrag_t = BaseNAXFrag>
+struct NAXSubTile {
+  STEEL_CONST short kRows = kRows_;
+  STEEL_CONST short kCols = kCols_;
+
+  STEEL_CONST short kFragRows = NAXFrag_t::kFragRows;
+  STEEL_CONST short kFragCols = NAXFrag_t::kFragCols;
+  STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag;
+
+  STEEL_CONST short kSubTileRows = kRows / kFragRows;
+  STEEL_CONST short kSubTileCols = kCols / kFragCols;
+
+  STEEL_CONST short kNumFrags = kSubTileRows * kSubTileCols;
+  STEEL_CONST short kElemsPerSubTile = kNumFrags * kElemsPerFrag;
+
+  STEEL_CONST int kRowsPerThread = kSubTileRows * NAXFrag_t::kElemRows;
+  STEEL_CONST int kColsPerThread = kSubTileCols * NAXFrag_t::kElemCols;
+
+  STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows;
+  STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols;
+  STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump;
+
+  using frag_type = typename NAXFrag_t::template dtype_frag_t<T>;
+
+  frag_type val_frags[kNumFrags];
+
+  METAL_FUNC constexpr void clear() {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kNumFrags; ++i) {
+      val_frags[i] = frag_type(0);
+    }
+  }
+
+  METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  METAL_FUNC constexpr const thread frag_type& frag_at(
+      const short i,
+      const short j) const {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  template <int i, int j>
+  METAL_FUNC constexpr thread frag_type& frag_at() {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  template <int i, int j>
+  METAL_FUNC constexpr const thread frag_type& frag_at() const {
+    return val_frags[i * kSubTileCols + j];
+  }
+
+  METAL_FUNC thread T* elems() {
+    return reinterpret_cast<thread T*>(val_frags);
+  }
+
+  METAL_FUNC const thread T* elems() const {
+    return reinterpret_cast<const thread T*>(val_frags);
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_reduce(thread metal::vec<T, kRowsPerThread>& vals) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::template row_reduce<Op>(
+            frag_at(i, j), &vals[i * kFragThrRows]);
+      }
+    }
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_bin_op(thread metal::vec<T, kRowsPerThread>& vals) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::template row_bin_op<Op>(
+            frag_at(i, j), &vals[i * kFragThrRows]);
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void load(
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::load(
+            frag_at(i, j),
+            src,
+            str_x,
+            str_y,
+            off_x + i * kFragRows,
+            off_y + j * kFragCols);
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      OffX off_x = {},
+      OffY off_y = {}) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::store(
+            frag_at(i, j),
+            dst,
+            str_x,
+            str_y,
+            off_x + i * kFragRows,
+            off_y + j * kFragCols);
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void load_rows(
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::load_rows(
+            frag_at(i, j),
+            src,
+            str_x,
+            str_y,
+            lim_x,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename SrcPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void load_safe(
+      SrcPtrType src,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::load_safe(
+            frag_at(i, j),
+            src,
+            str_x,
+            str_y,
+            lim_x,
+            lim_y,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename LimY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store_safe(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      LimY lim_y,
+      OffX off_x = {},
+      OffY off_y = {}) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::store_safe(
+            frag_at(i, j),
+            dst,
+            str_x,
+            str_y,
+            lim_x,
+            lim_y,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename LimX,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store_rows(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      LimX lim_x,
+      OffX off_x = {},
+      OffY off_y = {}) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kSubTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kSubTileCols; ++j) {
+        NAXFrag_t::store_safe(
+            frag_at(i, j),
+            dst,
+            str_x,
+            str_y,
+            lim_x,
+            off_x + (i * kFragRows),
+            off_y + (j * kFragCols));
+      }
+    }
+  }
+
+  template <
+      typename DstPtrType,
+      typename StrX,
+      typename StrY,
+      typename StartX,
+      typename StopX,
+      typename StartY,
+      typename StopY,
+      typename OffX = Int<0>,
+      typename OffY = Int<0>>
+  METAL_FUNC constexpr void store_slice(
+      DstPtrType dst,
+      StrX str_x,
+      StrY str_y,
+      StartX start_x,
+      StopX stop_x,
+      StartY start_y,
+      StopY stop_y,
+      OffX off_x = Int<0>{},
+      OffY off_y = Int<0>{}) const {
+    const_for_loop<0, kSubTileRows, 1>([&](auto idx_row) {
+      const_for_loop<0, kSubTileCols, 1>([&](auto idx_col) {
+        NAXFrag_t::store_slice(
+            frag_at<idx_row.value, idx_col.value>(),
+            dst,
+            str_x,
+            str_y,
+            start_x,
+            stop_x,
+            start_y,
+            stop_y,
+            off_x + idx_row * Int<kFragRows>{},
+            off_y + idx_col * Int<kFragCols>{});
+      });
+    });
+  }
+};
+
+template <
+    short RC,
+    short CC,
+    short RA,
+    short CA,
+    short RB,
+    short CB,
+    typename CType,
+    typename AType,
+    typename BType,
+    bool transpose_a,
+    bool transpose_b,
+    typename NAXFrag_t = BaseNAXFrag>
+METAL_FUNC void subtile_matmad_nax(
+    thread NAXSubTile<CType, RC, CC, NAXFrag_t>& C,
+    thread NAXSubTile<AType, RA, CA, NAXFrag_t>& A,
+    metal::bool_constant<transpose_a>,
+    thread NAXSubTile<BType, RB, CB, NAXFrag_t>& B,
+    metal::bool_constant<transpose_b>) {
+  // Static checks
+  constexpr short FMa = transpose_a ? CA : RA;
+  constexpr short FMc = RC;
+  static_assert(FMa == FMc, "NAX matmul: M dimensions do not match");
+
+  constexpr short FNb = transpose_b ? RB : CB;
+  constexpr short FNc = CC;
+  static_assert(FNb == FNc, "NAX matmul: N dimensions do not match");
+
+  constexpr short FKa = transpose_a ? RA : CA;
+  constexpr short FKb = transpose_b ? CB : RB;
+  static_assert(FKa == FKb, "NAX matmul: N dimensions do not match");
+
+  constexpr short FM = FMc;
+  constexpr short FN = FNc;
+  constexpr short FK = FKa;
+
+  constexpr int TM = FM / 16;
+  constexpr int TN = FN / 16;
+  constexpr int TK = FK / 16;
+
+  // Create Matmul descriptor
+  constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor(
+      FM,
+      FN,
+      FK,
+      transpose_a,
+      transpose_b,
+      true,
+      mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate);
+
+  // Create matmul op
+  mpp::tensor_ops::matmul2d<desc, metal::execution_simdgroup> gemm_op;
+
+  // Create matmul operands in registers
+  auto ct_a =
+      gemm_op.template get_left_input_cooperative_tensor<AType, BType, CType>();
+  auto ct_b =
+      gemm_op
+          .template get_right_input_cooperative_tensor<AType, BType, CType>();
+
+  // Create matmul output in register
+  auto ct_c = gemm_op.template get_destination_cooperative_tensor<
+      decltype(ct_a),
+      decltype(ct_b),
+      CType>();
+
+  // Load A in to left operand registers
+  STEEL_PRAGMA_UNROLL
+  for (short mm = 0; mm < TM; mm++) {
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < TK; kk++) {
+      const short fi = transpose_a ? kk : mm;
+      const short fj = transpose_a ? mm : kk;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < 8; i++) {
+        ct_a[(TK * mm + kk) * 8 + i] = A.frag_at(fi, fj)[i];
+      }
+    }
+  }
+
+  // Load B into right operand registers
+  STEEL_PRAGMA_UNROLL
+  for (short nn = 0; nn < TN; nn++) {
+    STEEL_PRAGMA_UNROLL
+    for (short kk = 0; kk < TK; kk++) {
+      const short fi = transpose_b ? nn : kk;
+      const short fj = transpose_b ? kk : nn;
+
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < 8; i++) {
+        ct_b[(TN * kk + nn) * 8 + i] = B.frag_at(fi, fj)[i];
+      }
+    }
+  }
+
+  // Load C into output registers (op handles accumulation)
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < ct_c.get_capacity(); i++) {
+    ct_c[i] = C.elems()[i];
+  }
+
+  // Do matmul
+  gemm_op.run(ct_a, ct_b, ct_c);
+
+  // Copy out results
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < ct_c.get_capacity(); i++) {
+    C.elems()[i] = ct_c[i];
+  }
+}
+
+template <typename T, short kTileRows_, short kTileCols_, class NAXSubTile_>
+struct NAXTile {
+  using NAXSubTile_t = NAXSubTile_;
+  using elem_type = T;
+  STEEL_CONST short kSubTileRows = NAXSubTile_t::kRows;
+  STEEL_CONST short kSubTileCols = NAXSubTile_t::kCols;
+  STEEL_CONST short kElemsPerSubTile = NAXSubTile_t::kElemsPerSubTile;
+
+  STEEL_CONST short kTileRows = kTileRows_;
+  STEEL_CONST short kTileCols = kTileCols_;
+
+  STEEL_CONST short kRows = kTileRows * kSubTileRows;
+  STEEL_CONST short kCols = kTileCols * kSubTileCols;
+
+  STEEL_CONST short kSubTiles = kTileRows * kTileCols;
+  STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile;
+
+  STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread;
+  STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread;
+
+  STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread;
+  STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread;
+
+  NAXSubTile_t val_subtiles[kSubTiles];
+
+  METAL_FUNC NAXTile() thread {}
+
+  METAL_FUNC constexpr void clear() {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kSubTiles; ++i) {
+      val_subtiles[i].clear();
+    }
+  }
+
+  METAL_FUNC constexpr thread NAXSubTile_t& subtile_at(
+      const short i,
+      const short j) {
+    return val_subtiles[i * kTileCols + j];
+  }
+
+  METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at(
+      const short i,
+      const short j) const {
+    return val_subtiles[i * kTileCols + j];
+  }
+
+  template <int i, int j>
+  METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const {
+    return val_subtiles[i * kTileCols + j];
+  }
+
+  METAL_FUNC thread elem_type* elems() {
+    return reinterpret_cast<thread elem_type*>(val_subtiles[0].elems());
+  }
+
+  METAL_FUNC const thread elem_type* elems() const {
+    return reinterpret_cast<const thread elem_type*>(val_subtiles[0].elems());
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_reduce(thread metal::vec<T, kRowsPerThread>& vals) const {
+    auto sub_rows = (thread metal::vec<T, kSubTileThrRows>*)(&vals);
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).template row_reduce<Op>(sub_rows[i]);
+      }
+    }
+  }
+
+  template <typename Op>
+  METAL_FUNC void row_bin_op(thread metal::vec<T, kRowsPerThread>& vals) {
+    auto sub_rows = (thread metal::vec<T, kSubTileThrRows>*)(&vals);
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).template row_bin_op<Op>(sub_rows[i]);
+      }
+    }
+  }
+
+  template <typename U, int str_x, int str_y>
+  METAL_FUNC void load(const threadgroup U* src) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load(
+            src,
+            Int<str_x>{},
+            Int<str_y>{},
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U, int str_x, int str_y>
+  METAL_FUNC void store(threadgroup U* dst) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store(
+            dst,
+            Int<str_x>{},
+            Int<str_y>{},
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void load(const device U* src, const int ld) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load(
+            &src[(i * kSubTileRows * ld + j * kSubTileCols)], ld, Int<1>{});
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void store(device U* dst, const int ld) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store(
+            &dst[(i * kSubTileRows * ld + j * kSubTileCols)], ld, Int<1>{});
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void
+  load_rows(const device U* src, const int ld, const short n_rows) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load_rows(
+            &src[(i * kSubTileRows) * ld + (j * kSubTileCols)],
+            ld,
+            Int<1>{},
+            n_rows - i * kSubTileRows);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void
+  load_safe(const device U* src, const int ld, const short2 src_tile_dims) {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).load_safe(
+            src,
+            ld,
+            Int<1>{},
+            src_tile_dims.y,
+            src_tile_dims.x,
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows)
+      const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store_rows(
+            &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)],
+            ld,
+            Int<1>{},
+            n_rows - i * kSubTileRows);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void
+  store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const {
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < kTileRows; ++i) {
+      STEEL_PRAGMA_UNROLL
+      for (int j = 0; j < kTileCols; ++j) {
+        subtile_at(i, j).store_safe(
+            dst,
+            ld,
+            Int<1>{},
+            dst_tile_dims.y,
+            dst_tile_dims.x,
+            i * kSubTileRows,
+            j * kSubTileCols);
+      }
+    }
+  }
+
+  template <typename U>
+  METAL_FUNC void store_slice(
+      device U* dst,
+      const int ld,
+      const short2 start,
+      const short2 stop) const {
+    const_for_loop<0, kTileRows, 1>([&](auto idx_row) {
+      const_for_loop<0, kTileCols, 1>([&](auto idx_col) {
+        subtile_at<idx_row.value, idx_col.value>().store_slice(
+            dst,
+            ld,
+            Int<1>{},
+            start.y,
+            stop.y,
+            start.x,
+            stop.x,
+            idx_row * Int<kSubTileRows>{},
+            idx_col * Int<kSubTileCols>{});
+      });
+    });
+  }
+};
+
+template <
+    class CTile,
+    class ATile,
+    class BTile,
+    bool transpose_a,
+    bool transpose_b>
+METAL_FUNC void tile_matmad_nax(
+    thread CTile& C,
+    thread ATile& A,
+    metal::bool_constant<transpose_a>,
+    thread BTile& B,
+    metal::bool_constant<transpose_b>) {
+  // Static checks
+  constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows;
+  constexpr short TMc = CTile::kTileRows;
+  static_assert(TMa == TMc, "NAX tile matmul: M dimensions do not match");
+
+  constexpr short FMa = transpose_a ? ATile::kSubTileCols : ATile::kSubTileRows;
+  constexpr short FMc = CTile::kSubTileRows;
+  static_assert(FMa == FMc, "NAX subtile matmul: M dimensions do not match");
+
+  constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols;
+  constexpr short TNc = CTile::kTileCols;
+  static_assert(TNb == TNc, "NAX tile matmul: N dimensions do not match");
+
+  constexpr short FNb = transpose_b ? BTile::kSubTileRows : BTile::kSubTileCols;
+  constexpr short FNc = CTile::kSubTileCols;
+  static_assert(FNb == FNc, "NAX subtile matmul: N dimensions do not match");
+
+  constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols;
+  constexpr short TKb = transpose_b ? BTile::kTileCols : BTile::kTileRows;
+  static_assert(TKa == TKb, "NAX tile matmul: K dimensions do not match");
+
+  constexpr short FKa = transpose_a ? ATile::kSubTileRows : ATile::kSubTileCols;
+  constexpr short FKb = transpose_b ? BTile::kSubTileCols : BTile::kSubTileRows;
+  static_assert(FKa == FKb, "NAX subtile matmul: K dimensions do not match");
+
+  constexpr short TM = TMc;
+  constexpr short TN = TNc;
+  constexpr short TK = TKa;
+
+  // Do matmul here
+  STEEL_PRAGMA_UNROLL
+  for (short i = 0; i < TM; ++i) {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < TN; ++j) {
+      STEEL_PRAGMA_UNROLL
+      for (short k = 0; k < TK; ++k) {
+        const short ra = transpose_a ? k : i;
+        const short ca = transpose_a ? i : k;
+        const short rb = transpose_b ? j : k;
+        const short cb = transpose_b ? k : j;
+
+        subtile_matmad_nax(
+            C.subtile_at(i, j),
+            A.subtile_at(ra, ca),
+            metal::bool_constant<transpose_a>{},
+            B.subtile_at(rb, cb),
+            metal::bool_constant<transpose_b>{});
+      }
+    }
+  }
+}
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/params.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/params.h
new file mode 100644
index 0000000..b0ba07d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/params.h
@@ -0,0 +1,65 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM param classes
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+struct GEMMParams {
+  const int M;
+  const int N;
+  const int K;
+
+  const int lda;
+  const int ldb;
+  const int ldd;
+
+  const int tiles_n;
+  const int tiles_m;
+
+  const int64_t batch_stride_a;
+  const int64_t batch_stride_b;
+  const int64_t batch_stride_d;
+
+  const int swizzle_log;
+  const int gemm_k_iterations_aligned;
+
+  const int batch_ndim;
+};
+
+struct GEMMSpiltKParams {
+  const int M;
+  const int N;
+  const int K;
+
+  const int lda;
+  const int ldb;
+  const int ldc;
+
+  const int tiles_n;
+  const int tiles_m;
+
+  const int split_k_partitions;
+  const int split_k_partition_stride;
+  const int split_k_partition_size;
+
+  const int swizzle_log;
+  const int gemm_k_iterations_aligned;
+};
+
+struct GEMMAddMMParams {
+  const int ldc;
+  const int fdc;
+
+  const int64_t batch_stride_c;
+
+  const float alpha;
+  const float beta;
+};
+
+} // namespace steel
+} // namespace mlx
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/transforms.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/transforms.h
new file mode 100644
index 0000000..0282a12
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/gemm/transforms.h
@@ -0,0 +1,72 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Transforms and Epilogues
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <typename OutT, typename InT>
+struct TransformNone {
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT) {
+    return static_cast<OutT>(x);
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAdd {
+  TransformAdd(const float, const float) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT c) {
+    return static_cast<OutT>(x) + c;
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAxpby {
+  const float alpha;
+  const float beta;
+
+  TransformAxpby(const float alpha_, const float beta_)
+      : alpha(alpha_), beta(beta_) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  METAL_FUNC OutT apply(InT x, OutT c) const {
+    return static_cast<OutT>(
+        x * static_cast<InT>(alpha) + (static_cast<OutT>(beta) * c));
+  }
+};
+
+template <typename T>
+struct AccumHelper {
+  typedef float accum_type;
+};
+
+struct BlockSwizzle {
+  static METAL_FUNC int2
+  swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) {
+    const int tid_x = (tid.x) >> swizzle_log;
+    const int tid_y =
+        ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1));
+    return int2(tid_x, tid_y);
+  }
+};
+
+} // namespace steel
+} // namespace mlx
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils.h
new file mode 100644
index 0000000..55720a2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils.h
@@ -0,0 +1,42 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+METAL_FUNC ulong2 elem_to_loc_broadcast(
+    uint elem,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    int ndim) {
+  ulong loc_a{0};
+  ulong loc_b{0};
+  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
+    int pos_in_dim = (elem % shape[i]);
+    elem /= shape[i];
+    loc_a += pos_in_dim * a_strides[i];
+    loc_b += pos_in_dim * b_strides[i];
+  }
+  return ulong2(loc_a, loc_b);
+}
+
+METAL_FUNC ulong3 elem_to_loc_broadcast(
+    uint elem,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    constant const int64_t* c_strides,
+    int ndim) {
+  ulong loc_a{0};
+  ulong loc_b{0};
+  ulong loc_c{0};
+  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
+    int pos_in_dim = (elem % shape[i]);
+    elem /= shape[i];
+    loc_a += pos_in_dim * a_strides[i];
+    loc_b += pos_in_dim * b_strides[i];
+    loc_c += pos_in_dim * c_strides[i];
+  }
+  return ulong3(loc_a, loc_b, loc_c);
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils/integral_constant.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils/integral_constant.h
new file mode 100644
index 0000000..526f561
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils/integral_constant.h
@@ -0,0 +1,134 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+#include "mlx/backend/metal/kernels/steel/utils/type_traits.h"
+
+#pragma METAL internals : enable
+
+namespace mlx {
+namespace steel {
+
+///////////////////////////////////////////////////////////////////////////////
+// Integral constant with casting
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, T v>
+struct integral_constant {
+  static constexpr constant T value = v;
+  using value_type = T;
+  using type = integral_constant;
+
+  METAL_FUNC constexpr operator value_type() const noexcept {
+    return value;
+  }
+
+  // METAL_FUNC constexpr value_type operator()() const noexcept {
+  //   return value;
+  // }
+};
+
+template <bool B>
+using bool_constant = integral_constant<bool, B>;
+using true_type = bool_constant<true>;
+using false_type = bool_constant<false>;
+
+template <class T>
+struct is_integral : bool_constant<metal::is_integral<T>::value> {};
+
+template <class T, T v>
+struct is_integral<integral_constant<T, v>>
+    : bool_constant<metal::is_integral<T>::value> {};
+
+template <typename T>
+constexpr constant bool is_integral_v = is_integral<T>::value;
+
+template <int val>
+using Int = integral_constant<int, val>;
+
+///////////////////////////////////////////////////////////////////////////////
+// Binary Operators on Integral constants
+///////////////////////////////////////////////////////////////////////////////
+
+#define integral_const_binop(__op__, __operator__)          \
+  template <typename T, T tv, typename U, U uv>             \
+  METAL_FUNC constexpr auto __operator__(                   \
+      integral_constant<T, tv>, integral_constant<U, uv>) { \
+    constexpr auto res = tv __op__ uv;                      \
+    return integral_constant<decltype(res), res>{};         \
+  }
+
+integral_const_binop(+, operator+);
+integral_const_binop(-, operator-);
+integral_const_binop(*, operator*);
+integral_const_binop(/, operator/);
+
+integral_const_binop(==, operator==);
+integral_const_binop(!=, operator!=);
+integral_const_binop(<, operator<);
+integral_const_binop(>, operator>);
+integral_const_binop(<=, operator<=);
+integral_const_binop(>=, operator>=);
+
+integral_const_binop(&&, operator&&);
+integral_const_binop(||, operator||);
+
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator||(true_type, T) {
+  return true_type{};
+}
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator||(T, true_type) {
+  return true_type{};
+}
+
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator&&(false_type, T) {
+  return false_type{};
+}
+
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator&&(T, false_type) {
+  return false_type{};
+}
+
+// Dispatch utilities
+template <typename F>
+void dispatch_bool(bool v, F f) {
+  if (v) {
+    f(true_type{});
+  } else {
+    f(false_type{});
+  }
+}
+
+template <int start, int stop, int step, typename F>
+constexpr void const_for_loop(F f) {
+  if constexpr (start < stop) {
+    constexpr auto idx = Int<start>{};
+    f(idx);
+    const_for_loop<start + step, stop, step, F>(f);
+  }
+}
+
+#undef integral_const_binop
+
+///////////////////////////////////////////////////////////////////////////////
+// Reduction operators
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+METAL_FUNC constexpr T sum(T x) {
+  return x;
+}
+
+template <typename T, typename... Us>
+METAL_FUNC constexpr auto sum(T x, Us... us) {
+  return x + sum(us...);
+}
+
+} // namespace steel
+} // namespace mlx
+
+#pragma METAL internals : disable
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils/type_traits.h b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils/type_traits.h
new file mode 100644
index 0000000..f004dc8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/steel/utils/type_traits.h
@@ -0,0 +1,55 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+#pragma METAL internals : enable
+
+namespace metal {
+
+template <typename T>
+struct is_empty : metal::bool_constant<__is_empty(T)> {};
+
+#ifdef __cpp_variable_templates
+template <typename T>
+constexpr constant bool is_empty_v = is_empty<T>::value;
+#endif
+
+template <typename... Ts>
+struct make_void {
+  typedef void type;
+};
+
+template <typename... Ts>
+using void_t = typename make_void<Ts...>::type;
+
+template <class T>
+struct is_static : metal::bool_constant<is_empty<remove_cv_t<T>>::value> {};
+
+template <typename T>
+struct pointer_element {};
+
+template <typename T>
+struct pointer_element<thread T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<device T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<constant T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<threadgroup T*> {
+  using type = remove_cv_t<T>;
+};
+
+template <typename T>
+using pointer_element_t = typename pointer_element<remove_cv_t<T>>::type;
+
+} // namespace metal
+
+#pragma METAL internals : disable
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/ternary.h b/macos/third_party/mlx/mlx/backend/metal/kernels/ternary.h
new file mode 100644
index 0000000..705b73e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/ternary.h
@@ -0,0 +1,145 @@
+// Copyright © 2024 Apple Inc.
+
+template <
+    typename T,
+    typename Op,
+    bool BSCALAR,
+    bool CSCALAR,
+    int N = WorkPerThread<T>::n>
+[[kernel]] void ternary_v(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      auto bidx = BSCALAR ? 0 : index + i;
+      auto cidx = CSCALAR ? 0 : index + i;
+      d[index + i] = Op()(a[index + i], b[bidx], c[cidx]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto bidx = BSCALAR ? 0 : index + i;
+      auto cidx = CSCALAR ? 0 : index + i;
+      d[index + i] = Op()(a[index + i], b[bidx], c[cidx]);
+    }
+  }
+}
+
+template <
+    typename T,
+    typename Op,
+    bool BSCALAR,
+    bool CSCALAR,
+    int N = WorkPerThread<T>::n>
+[[kernel]] void ternary_v2(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      auto bidx = BSCALAR ? 0 : offset + i;
+      auto cidx = CSCALAR ? 0 : offset + i;
+      d[offset + i] = Op()(a[offset + i], b[bidx], c[cidx]);
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      auto bidx = BSCALAR ? 0 : offset + i;
+      auto cidx = CSCALAR ? 0 : offset + i;
+      d[offset + i] = Op()(a[offset + i], b[bidx], c[cidx]);
+    }
+  }
+}
+
+template <typename T, typename Op, typename IdxT = int64_t>
+[[kernel]] void ternary_g_nd1(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const int64_t& a_strides,
+    constant const int64_t& b_strides,
+    constant const int64_t& c_strides,
+    uint index [[thread_position_in_grid]]) {
+  auto a_idx = elem_to_loc_1<IdxT>(index, a_strides);
+  auto b_idx = elem_to_loc_1<IdxT>(index, b_strides);
+  auto c_idx = elem_to_loc_1<IdxT>(index, c_strides);
+  d[index] = Op()(a[a_idx], b[b_idx], c[c_idx]);
+}
+
+template <typename T, typename Op, typename IdxT = int64_t>
+[[kernel]] void ternary_g_nd2(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const int64_t a_strides[2],
+    constant const int64_t b_strides[2],
+    constant const int64_t c_strides[2],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_2<IdxT>(index, a_strides);
+  auto b_idx = elem_to_loc_2<IdxT>(index, b_strides);
+  auto c_idx = elem_to_loc_2<IdxT>(index, c_strides);
+  IdxT out_idx = index.x + IdxT(grid_dim.x) * index.y;
+  d[out_idx] = Op()(a[a_idx], b[b_idx], c[c_idx]);
+}
+
+template <typename T, typename Op, typename IdxT = int64_t>
+[[kernel]] void ternary_g_nd3(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const int64_t a_strides[3],
+    constant const int64_t b_strides[3],
+    constant const int64_t c_strides[3],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_3<IdxT>(index, a_strides);
+  auto b_idx = elem_to_loc_3<IdxT>(index, b_strides);
+  auto c_idx = elem_to_loc_3<IdxT>(index, c_strides);
+  IdxT out_idx = index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z);
+  d[out_idx] = Op()(a[a_idx], b[b_idx], c[c_idx]);
+}
+
+template <typename T, typename Op, int N = 1, typename IdxT = int64_t>
+[[kernel]] void ternary_g(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    constant const int64_t* c_strides,
+    constant const int& ndim,
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto idx = elem_to_loc_3_nd<IdxT>(
+      {N * index.x, index.y, index.z},
+      shape,
+      a_strides,
+      b_strides,
+      c_strides,
+      ndim);
+  auto xshape = shape[ndim - 1];
+  IdxT out_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z);
+  IdxT a_xstride = a_strides[ndim - 1];
+  IdxT b_xstride = b_strides[ndim - 1];
+  IdxT c_xstride = c_strides[ndim - 1];
+  for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) {
+    d[out_idx++] = Op()(a[idx.x], b[idx.y], c[idx.z]);
+    idx.x += a_xstride;
+    idx.y += b_xstride;
+    idx.z += c_xstride;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/ternary.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/ternary.metal
new file mode 100644
index 0000000..ac1cd41
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/ternary.metal
@@ -0,0 +1,48 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_integer>
+#include <metal_math>
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/ternary_ops.h"
+#include "mlx/backend/metal/kernels/ternary.h"
+
+#define instantiate_ternary_base(op, tname, type)                    \
+  instantiate_kernel("v_" #op #tname, ternary_v, type, op, false, false, 1) \
+  instantiate_kernel("v2_" #op #tname, ternary_v2, type, op, false, false)  \
+  instantiate_kernel("vs_" #op #tname, ternary_v, type, op, false, true, 1) \
+  instantiate_kernel("vs2_" #op #tname, ternary_v2, type, op, false, true)  \
+  instantiate_kernel("sv_" #op #tname, ternary_v, type, op, true, false, 1) \
+  instantiate_kernel("sv2_" #op #tname, ternary_v2, type, op, true, false)  \
+  instantiate_kernel("gn2_" #op #tname, ternary_g, type, op, 2, int) \
+  instantiate_kernel("g1_" #op #tname, ternary_g_nd1, type, op, int) \
+  instantiate_kernel("g2_" #op #tname, ternary_g_nd2, type, op, int) \
+  instantiate_kernel("g3_" #op #tname, ternary_g_nd3, type, op, int) \
+  instantiate_kernel("g1large_" #op #tname, ternary_g_nd1, type, op) \
+  instantiate_kernel("g2large_" #op #tname, ternary_g_nd2, type, op) \
+  instantiate_kernel("g3large_" #op #tname, ternary_g_nd3, type, op) \
+  instantiate_kernel("gn4large_" #op #tname, ternary_g, type, op, 4) \
+
+#define instantiate_ternary_all(op, tname, type)            \
+  instantiate_kernel("vn_" #op #tname, ternary_v, type, op, false, false) \
+  instantiate_kernel("vsn_" #op #tname, ternary_v, type, op, false, true) \
+  instantiate_kernel("svn_" #op #tname, ternary_v, type, op, true, false) \
+  instantiate_ternary_base(op, tname, type)
+
+#define instantiate_ternary_types(op)               \
+  instantiate_ternary_all(op, bool_, bool)          \
+  instantiate_ternary_all(op, uint8, uint8_t)       \
+  instantiate_ternary_all(op, uint16, uint16_t)     \
+  instantiate_ternary_all(op, uint32, uint32_t)     \
+  instantiate_ternary_base(op, uint64, uint64_t)    \
+  instantiate_ternary_all(op, int8, int8_t)         \
+  instantiate_ternary_all(op, int16, int16_t)       \
+  instantiate_ternary_all(op, int32, int32_t)       \
+  instantiate_ternary_base(op, int64, int64_t)      \
+  instantiate_ternary_all(op, float16, half)        \
+  instantiate_ternary_all(op, float32, float)       \
+  instantiate_ternary_all(op, bfloat16, bfloat16_t) \
+  instantiate_ternary_base(op, complex64, complex64_t) // clang-format on
+
+instantiate_ternary_types(Select)
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/ternary_ops.h b/macos/third_party/mlx/mlx/backend/metal/kernels/ternary_ops.h
new file mode 100644
index 0000000..e0235d9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/ternary_ops.h
@@ -0,0 +1,10 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+struct Select {
+  template <typename T>
+  T operator()(bool condition, T x, T y) {
+    return condition ? x : y;
+  }
+};
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/unary.h b/macos/third_party/mlx/mlx/backend/metal/kernels/unary.h
new file mode 100644
index 0000000..db7be3d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/unary.h
@@ -0,0 +1,63 @@
+// Copyright © 2024 Apple Inc.
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void unary_v(
+    device const T* in,
+    device U* out,
+    constant uint& size,
+    uint index [[thread_position_in_grid]]) {
+  index *= N;
+  if (N > 1 && index + N > size) {
+    for (int i = 0; index + i < size; ++i) {
+      out[index + i] = static_cast<U>(Op()(in[index + i]));
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      out[index + i] = static_cast<U>(Op()(in[index + i]));
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N = WorkPerThread<T>::n>
+[[kernel]] void unary_v2(
+    device const T* in,
+    device U* out,
+    constant int64_t& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y));
+  if (N > 1 && offset + N > size) {
+    for (int i = 0; offset + i < size; ++i) {
+      out[offset + i] = static_cast<U>(Op()(in[offset + i]));
+    }
+  } else {
+    for (int i = 0; i < N; ++i) {
+      out[offset + i] = static_cast<U>(Op()(in[offset + i]));
+    }
+  }
+}
+
+template <
+    typename T,
+    typename U,
+    typename Op,
+    int N = 1,
+    typename IdxT = int64_t>
+[[kernel]] void unary_g(
+    device const T* in,
+    device U* out,
+    constant const int* in_shape,
+    constant const int64_t* in_strides,
+    device const int& ndim,
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto idx = elem_to_loc<IdxT>(
+      {N * index.x, index.y, index.z}, in_shape, in_strides, ndim);
+  auto xshape = in_shape[ndim - 1];
+  IdxT xstride = in_strides[ndim - 1];
+  IdxT out_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z);
+  for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) {
+    out[out_idx++] = static_cast<U>(Op()(in[idx]));
+    idx += xstride;
+  }
+}
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/unary.metal b/macos/third_party/mlx/mlx/backend/metal/kernels/unary.metal
new file mode 100644
index 0000000..54a0f56
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/unary.metal
@@ -0,0 +1,115 @@
+// Copyright © 2024 Apple Inc.
+
+// clang-format off
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/unary_ops.h"
+#include "mlx/backend/metal/kernels/unary.h"
+
+#define instantiate_unary_work_per_thread(op, in_tname, out_tname, in_type, out_type) \
+  instantiate_kernel("vn_" #op #in_tname #out_tname, unary_v, in_type, out_type, op)
+
+#define instantiate_unary_base(op, in_tname, out_tname, in_type, out_type)             \
+  instantiate_kernel("v_" #op #in_tname #out_tname, unary_v, in_type, out_type, op, 1) \
+  instantiate_kernel("v2_" #op #in_tname #out_tname, unary_v2, in_type, out_type, op)  \
+  instantiate_kernel(                                                                  \
+      "gn1_" #op #in_tname #out_tname, unary_g, in_type, out_type, op, 1, int)         \
+  instantiate_kernel(                                                                  \
+      "gn4large_" #op #in_tname #out_tname, unary_g, in_type, out_type, op, 4)
+
+#define instantiate_unary_all(op, in_tname, out_tname, in_type, out_type)       \
+  instantiate_unary_base(op, in_tname, out_tname, in_type, out_type)            \
+  instantiate_unary_work_per_thread(op, in_tname, out_tname, in_type, out_type)
+
+#define instantiate_unary_all_same(op, tname, type)   \
+  instantiate_unary_all(op, tname, tname, type, type)
+
+#define instantiate_unary_base_same(op, tname, type)   \
+  instantiate_unary_base(op, tname, tname, type, type)
+
+#define instantiate_unary_float(op)                    \
+  instantiate_unary_all_same(op, float16, half)        \
+  instantiate_unary_all_same(op, float32, float)       \
+  instantiate_unary_all_same(op, bfloat16, bfloat16_t)
+
+#define instantiate_unary_int(op)                   \
+  instantiate_unary_all_same(op, uint8, uint8_t)    \
+  instantiate_unary_all_same(op, uint16, uint16_t)  \
+  instantiate_unary_all_same(op, uint32, uint32_t)  \
+  instantiate_unary_base_same(op, uint64, uint64_t) \
+  instantiate_unary_all_same(op, int8, int8_t)      \
+  instantiate_unary_all_same(op, int16, int16_t)    \
+  instantiate_unary_all_same(op, int32, int32_t)    \
+  instantiate_unary_base_same(op, int64, int64_t)
+
+#define instantiate_unary_types(op)                \
+  instantiate_unary_all_same(op, bool_, bool)      \
+  instantiate_unary_int(op)                        \
+  instantiate_unary_float(op)
+
+instantiate_unary_types(Abs)
+instantiate_unary_float(ArcCos)
+instantiate_unary_float(ArcCosh)
+instantiate_unary_float(ArcSin)
+instantiate_unary_float(ArcSinh)
+instantiate_unary_float(ArcTan)
+instantiate_unary_float(ArcTanh)
+instantiate_unary_types(Ceil)
+instantiate_unary_float(Cos)
+instantiate_unary_float(Cosh)
+instantiate_unary_float(Exp)
+instantiate_unary_float(Expm1)
+instantiate_unary_types(Floor)
+instantiate_unary_float(Log)
+instantiate_unary_float(Log2)
+instantiate_unary_float(Log10)
+instantiate_unary_float(Log1p)
+instantiate_unary_types(Negative)
+instantiate_unary_float(Sigmoid)
+instantiate_unary_float(Erf)
+instantiate_unary_float(ErfInv)
+instantiate_unary_types(Sign)
+instantiate_unary_float(Sin)
+instantiate_unary_float(Sinh)
+instantiate_unary_types(Square)
+instantiate_unary_float(Sqrt)
+instantiate_unary_float(Rsqrt)
+instantiate_unary_float(Tan)
+instantiate_unary_float(Tanh)
+instantiate_unary_float(Round)
+instantiate_unary_int(BitwiseInvert)
+
+instantiate_unary_base_same(Abs, complex64, complex64_t)
+instantiate_unary_base_same(ArcCos, complex64, complex64_t)
+instantiate_unary_base_same(ArcSin, complex64, complex64_t)
+instantiate_unary_base_same(ArcTan, complex64, complex64_t)
+instantiate_unary_base_same(Conjugate, complex64, complex64_t)
+instantiate_unary_base_same(Cos, complex64, complex64_t)
+instantiate_unary_base_same(Cosh, complex64, complex64_t)
+instantiate_unary_base_same(Exp, complex64, complex64_t)
+instantiate_unary_base_same(Log, complex64, complex64_t)
+instantiate_unary_base_same(Log1p, complex64, complex64_t)
+instantiate_unary_base_same(Log2, complex64, complex64_t)
+instantiate_unary_base_same(Log10, complex64, complex64_t)
+instantiate_unary_base_same(Negative, complex64, complex64_t)
+instantiate_unary_base_same(Sign, complex64, complex64_t)
+instantiate_unary_base_same(Sin, complex64, complex64_t)
+instantiate_unary_base_same(Sinh, complex64, complex64_t)
+instantiate_unary_base_same(Square, complex64, complex64_t)
+instantiate_unary_base_same(Sqrt, complex64, complex64_t)
+instantiate_unary_base_same(Rsqrt, complex64, complex64_t)
+instantiate_unary_base_same(Tan, complex64, complex64_t)
+instantiate_unary_base_same(Tanh, complex64, complex64_t)
+instantiate_unary_base_same(Round, complex64, complex64_t)
+instantiate_unary_base(Real, complex64, float32, complex64_t, float)
+instantiate_unary_base(Imag, complex64, float32, complex64_t, float)
+
+instantiate_unary_all_same(LogicalNot, bool_, bool)
+
+instantiate_unary_all(ToFP8, float16, uint8, float16_t, uint8_t)
+instantiate_unary_all(ToFP8, bfloat16, uint8, bfloat16_t, uint8_t)
+instantiate_unary_all(ToFP8, float32, uint8, float, uint8_t)
+instantiate_unary_all(FromFP8, uint8, float16, uint8_t, float16_t)
+instantiate_unary_all(FromFP8, uint8, bfloat16, uint8_t, bfloat16_t)
+instantiate_unary_all(FromFP8, uint8, float32, uint8_t, float)
+
+    // clang-format on
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/unary_ops.h b/macos/third_party/mlx/mlx/backend/metal/kernels/unary_ops.h
new file mode 100644
index 0000000..327bb5a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/unary_ops.h
@@ -0,0 +1,454 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <metal_integer>
+#include <metal_math>
+
+#include "mlx/backend/metal/kernels/cexpf.h"
+#include "mlx/backend/metal/kernels/erf.h"
+#include "mlx/backend/metal/kernels/expm1f.h"
+#include "mlx/backend/metal/kernels/fp8.h"
+
+namespace {
+constant float inf = metal::numeric_limits<float>::infinity();
+}
+
+struct Abs {
+  template <typename T>
+  T operator()(T x) {
+    return metal::abs(x);
+  };
+  uint8_t operator()(uint8_t x) {
+    return x;
+  };
+  uint16_t operator()(uint16_t x) {
+    return x;
+  };
+  uint32_t operator()(uint32_t x) {
+    return x;
+  };
+  uint64_t operator()(uint64_t x) {
+    return x;
+  };
+  bool operator()(bool x) {
+    return x;
+  };
+  complex64_t operator()(complex64_t x) {
+    return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0};
+  };
+};
+
+struct ArcCos {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::acos(x);
+  };
+
+  complex64_t operator()(complex64_t x);
+};
+
+struct ArcCosh {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::acosh(x);
+  };
+};
+
+struct ArcSin {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::asin(x);
+  };
+
+  complex64_t operator()(complex64_t x);
+};
+
+struct ArcSinh {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::asinh(x);
+  };
+};
+
+struct ArcTan {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::atan(x);
+  };
+
+  complex64_t operator()(complex64_t x);
+};
+
+struct ArcTanh {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::atanh(x);
+  };
+};
+
+struct BitwiseInvert {
+  template <typename T>
+  T operator()(T x) {
+    return ~x;
+  };
+};
+
+struct Ceil {
+  template <typename T>
+  T operator()(T x) {
+    return metal::ceil(x);
+  };
+  int8_t operator()(int8_t x) {
+    return x;
+  };
+  int16_t operator()(int16_t x) {
+    return x;
+  };
+  int32_t operator()(int32_t x) {
+    return x;
+  };
+  int64_t operator()(int64_t x) {
+    return x;
+  };
+  uint8_t operator()(uint8_t x) {
+    return x;
+  };
+  uint16_t operator()(uint16_t x) {
+    return x;
+  };
+  uint32_t operator()(uint32_t x) {
+    return x;
+  };
+  uint64_t operator()(uint64_t x) {
+    return x;
+  };
+  bool operator()(bool x) {
+    return x;
+  };
+};
+
+struct Cos {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::cos(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    return {
+        metal::precise::cos(x.real) * metal::precise::cosh(x.imag),
+        -metal::precise::sin(x.real) * metal::precise::sinh(x.imag)};
+  };
+};
+
+struct Cosh {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::cosh(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    return {
+        metal::precise::cosh(x.real) * metal::precise::cos(x.imag),
+        metal::precise::sinh(x.real) * metal::precise::sin(x.imag)};
+  };
+};
+
+struct Conjugate {
+  complex64_t operator()(complex64_t x) {
+    return complex64_t{x.real, -x.imag};
+  }
+};
+
+struct Erf {
+  template <typename T>
+  T operator()(T x) {
+    return static_cast<T>(erf(static_cast<float>(x)));
+  };
+};
+
+struct ErfInv {
+  template <typename T>
+  T operator()(T x) {
+    return static_cast<T>(erfinv(static_cast<float>(x)));
+  };
+};
+
+struct Exp {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::exp(x);
+  };
+  complex64_t operator()(complex64_t x) {
+    return cexpf(x);
+  }
+};
+
+struct Expm1 {
+  template <typename T>
+  T operator()(T x) {
+    return static_cast<T>(expm1f(static_cast<float>(x)));
+  };
+};
+
+struct Floor {
+  template <typename T>
+  T operator()(T x) {
+    return metal::floor(x);
+  };
+  int8_t operator()(int8_t x) {
+    return x;
+  };
+  int16_t operator()(int16_t x) {
+    return x;
+  };
+  int32_t operator()(int32_t x) {
+    return x;
+  };
+  int64_t operator()(int64_t x) {
+    return x;
+  };
+  uint8_t operator()(uint8_t x) {
+    return x;
+  };
+  uint16_t operator()(uint16_t x) {
+    return x;
+  };
+  uint32_t operator()(uint32_t x) {
+    return x;
+  };
+  uint64_t operator()(uint64_t x) {
+    return x;
+  };
+  bool operator()(bool x) {
+    return x;
+  };
+};
+
+struct Imag {
+  float operator()(complex64_t x) {
+    return x.imag;
+  };
+};
+
+struct Log {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::log(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    auto r = metal::precise::log(Abs{}(x).real);
+    auto i = metal::precise::atan2(x.imag, x.real);
+    return {r, i};
+  };
+};
+
+struct Log2 {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::log2(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    auto y = Log{}(x);
+    return {y.real / M_LN2_F, y.imag / M_LN2_F};
+  };
+};
+
+struct Log10 {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::log10(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    auto y = Log{}(x);
+    return {y.real / M_LN10_F, y.imag / M_LN10_F};
+  };
+};
+
+struct Log1p {
+  template <typename T>
+  T operator()(T x) {
+    return log1p(x);
+  };
+};
+
+struct LogicalNot {
+  template <typename T>
+  T operator()(T x) {
+    return !x;
+  };
+};
+
+struct Negative {
+  template <typename T>
+  T operator()(T x) {
+    return -x;
+  };
+};
+
+struct Real {
+  float operator()(complex64_t x) {
+    return x.real;
+  };
+};
+
+struct Round {
+  template <typename T>
+  T operator()(T x) {
+    return metal::rint(x);
+  };
+  complex64_t operator()(complex64_t x) {
+    return {metal::rint(x.real), metal::rint(x.imag)};
+  };
+};
+
+struct Sigmoid {
+  template <typename T>
+  T operator()(T x) {
+    auto y = 1 / (1 + metal::exp(metal::abs(x)));
+    return (x < 0) ? y : 1 - y;
+  }
+};
+
+struct Sign {
+  template <typename T>
+  T operator()(T x) {
+    return (x > T(0)) - (x < T(0));
+  };
+  uint32_t operator()(uint32_t x) {
+    return x != 0;
+  };
+  complex64_t operator()(complex64_t x) {
+    if (x == complex64_t(0)) {
+      return x;
+    }
+    return x /
+        (complex64_t)metal::precise::sqrt(x.real * x.real + x.imag * x.imag);
+  };
+};
+
+struct Sin {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::sin(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    return {
+        metal::precise::sin(x.real) * metal::precise::cosh(x.imag),
+        metal::precise::cos(x.real) * metal::precise::sinh(x.imag)};
+  };
+};
+
+struct Sinh {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::sinh(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    return {
+        metal::precise::sinh(x.real) * metal::precise::cos(x.imag),
+        metal::precise::cosh(x.real) * metal::precise::sin(x.imag)};
+  };
+};
+
+struct Square {
+  template <typename T>
+  T operator()(T x) {
+    return x * x;
+  };
+};
+
+struct Sqrt {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::sqrt(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    if (x.real == 0.0 && x.imag == 0.0) {
+      return {0.0, 0.0};
+    }
+    auto r = Abs{}(x).real;
+    auto a = metal::precise::sqrt((r + x.real) / 2.0);
+    auto b_abs = metal::precise::sqrt((r - x.real) / 2.0);
+    auto b = metal::copysign(b_abs, x.imag);
+    return {a, b};
+  }
+};
+
+struct Rsqrt {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::rsqrt(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    return 1.0 / Sqrt{}(x);
+  }
+};
+
+struct Tan {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::tan(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    float tan_a = metal::precise::tan(x.real);
+    float tanh_b = metal::precise::tanh(x.imag);
+    float t1 = tan_a * tanh_b;
+    float denom = 1. + t1 * t1;
+    return {(tan_a - tanh_b * t1) / denom, (tanh_b + tan_a * t1) / denom};
+  };
+};
+
+struct Tanh {
+  template <typename T>
+  T operator()(T x) {
+    return metal::precise::tanh(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    float tanh_a = metal::precise::tanh(x.real);
+    float tan_b = metal::precise::tan(x.imag);
+    float t1 = tanh_a * tan_b;
+    float denom = 1. + t1 * t1;
+    return {(tanh_a + tan_b * t1) / denom, (tan_b - tanh_a * t1) / denom};
+  };
+};
+
+complex64_t ArcCos::operator()(complex64_t x) {
+  auto i = complex64_t{0.0, 1.0};
+  auto y = Log{}(x + i * Sqrt{}(1.0 - x * x));
+  return {y.imag, -y.real};
+};
+
+complex64_t ArcSin::operator()(complex64_t x) {
+  auto i = complex64_t{0.0, 1.0};
+  auto y = Log{}(i * x + Sqrt{}(1.0 - x * x));
+  return {y.imag, -y.real};
+};
+
+complex64_t ArcTan::operator()(complex64_t x) {
+  auto i = complex64_t{0.0, 1.0};
+  auto ix = i * x;
+  return (1.0 / complex64_t{0.0, 2.0}) * Log{}((1.0 + ix) / (1.0 - ix));
+};
+
+struct ToFP8 {
+  template <typename T>
+  uint8_t operator()(T f) {
+    return fp8_e4m3(f).bits;
+  }
+};
+
+struct FromFP8 {
+  float operator()(uint8_t x) {
+    return float(*(thread fp8_e4m3*)(&x));
+  }
+};
diff --git a/macos/third_party/mlx/mlx/backend/metal/kernels/utils.h b/macos/third_party/mlx/mlx/backend/metal/kernels/utils.h
new file mode 100644
index 0000000..d356450
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/kernels/utils.h
@@ -0,0 +1,445 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <metal_math>
+
+#include "mlx/backend/metal/kernels/bf16.h"
+#include "mlx/backend/metal/kernels/bf16_math.h"
+#include "mlx/backend/metal/kernels/complex.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/logging.h"
+
+typedef half float16_t;
+
+// Work per thread values for different types. The values here are expected to
+// match get_work_per_thread in mlx/backend/metal/utils.h
+template <typename U>
+struct WorkPerThread {
+  static_assert(sizeof(U) <= 8, "Type too large");
+  static constexpr int constant n = 8 / sizeof(U);
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Type limits utils
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename U>
+struct Limits {
+  static const constant U max = metal::numeric_limits<U>::max();
+  static const constant U min = metal::numeric_limits<U>::min();
+  static const constant U finite_max = metal::numeric_limits<U>::max();
+  static const constant U finite_min = metal::numeric_limits<U>::min();
+};
+
+#define instantiate_default_limit(type)                                      \
+  template <>                                                                \
+  struct Limits<type> {                                                      \
+    static constexpr constant type max = metal::numeric_limits<type>::max(); \
+    static constexpr constant type min = metal::numeric_limits<type>::min(); \
+    static constexpr constant type finite_max =                              \
+        metal::numeric_limits<type>::max();                                  \
+    static constexpr constant type finite_min =                              \
+        metal::numeric_limits<type>::min();                                  \
+  };
+
+instantiate_default_limit(uint8_t);
+instantiate_default_limit(uint16_t);
+instantiate_default_limit(uint32_t);
+instantiate_default_limit(uint64_t);
+instantiate_default_limit(int8_t);
+instantiate_default_limit(int16_t);
+instantiate_default_limit(int32_t);
+instantiate_default_limit(int64_t);
+
+#define instantiate_float_limit(type)             \
+  template <>                                     \
+  struct Limits<type> {                           \
+    static constexpr constant type max =          \
+        metal::numeric_limits<type>::infinity();  \
+    static constexpr constant type min =          \
+        -metal::numeric_limits<type>::infinity(); \
+    static constexpr constant type finite_max =   \
+        metal::numeric_limits<type>::max();       \
+    static constexpr constant type finite_min =   \
+        -metal::numeric_limits<type>::max();      \
+  };
+
+instantiate_float_limit(half);
+instantiate_float_limit(float);
+instantiate_float_limit(bfloat16_t);
+
+template <>
+struct Limits<bool> {
+  static constexpr constant bool max = true;
+  static constexpr constant bool min = false;
+};
+
+template <>
+struct Limits<complex64_t> {
+  static constexpr constant complex64_t max = complex64_t(
+      metal::numeric_limits<float>::infinity(),
+      metal::numeric_limits<float>::infinity());
+  static constexpr constant complex64_t min = complex64_t(
+      -metal::numeric_limits<float>::infinity(),
+      -metal::numeric_limits<float>::infinity());
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Indexing utils
+///////////////////////////////////////////////////////////////////////////////
+
+#define MLX_MTL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
+
+///////////////////////////////////////////////////////////////////////////////
+// Single Array with generic dims
+
+template <typename IdxT = int64_t>
+METAL_FUNC IdxT elem_to_loc(
+    IdxT elem,
+    constant const int* shape,
+    constant const int64_t* strides,
+    int ndim) {
+  IdxT loc = 0;
+  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
+    loc += (elem % shape[i]) * IdxT(strides[i]);
+    elem /= shape[i];
+  }
+  return loc;
+}
+
+// Non templated version to handle arbitrary dims
+template <typename IdxT = int64_t>
+METAL_FUNC IdxT elem_to_loc(
+    uint3 elem,
+    constant const int* shape,
+    constant const int64_t* strides,
+    int ndim) {
+  IdxT loc =
+      elem.x * IdxT(strides[ndim - 1]) + elem.y * IdxT(strides[ndim - 2]);
+  for (int d = ndim - 3; d >= 0; --d) {
+    loc += (elem.z % shape[d]) * IdxT(strides[d]);
+    elem.z /= shape[d];
+  }
+  return loc;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Single Array with fixed N dims
+
+template <typename IdxT = int64_t>
+METAL_FUNC IdxT elem_to_loc_1(uint elem, constant const int64_t& stride) {
+  return elem * IdxT(stride);
+}
+
+template <typename IdxT = int64_t>
+METAL_FUNC IdxT elem_to_loc_2(uint2 elem, constant const int64_t strides[2]) {
+  return elem.x * IdxT(strides[1]) + elem.y * IdxT(strides[0]);
+}
+
+template <typename IdxT = int64_t>
+METAL_FUNC IdxT elem_to_loc_3(uint3 elem, constant const int64_t strides[3]) {
+  return elem.x * IdxT(strides[2]) + elem.y * IdxT(strides[1]) +
+      elem.z * IdxT(strides[0]);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Multiple Arrays with generic dims
+
+template <typename IdxT = int64_t>
+METAL_FUNC vec<IdxT, 2> elem_to_loc_2_nd(
+    uint3 elem,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    int ndim) {
+  vec<IdxT, 2> loc = {
+      IdxT(
+          elem.x * IdxT(a_strides[ndim - 1]) +
+          IdxT(elem.y) * IdxT(a_strides[ndim - 2])),
+      IdxT(
+          elem.x * IdxT(b_strides[ndim - 1]) +
+          elem.y * IdxT(b_strides[ndim - 2]))};
+  for (int d = ndim - 3; d >= 0; --d) {
+    uint l = elem.z % shape[d];
+    loc.x += l * IdxT(a_strides[d]);
+    loc.y += l * IdxT(b_strides[d]);
+    elem.z /= shape[d];
+  }
+  return loc;
+}
+
+template <typename IdxT = int64_t>
+METAL_FUNC vec<IdxT, 3> elem_to_loc_3_nd(
+    uint3 elem,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    constant const int64_t* c_strides,
+    int ndim) {
+  vec<IdxT, 3> loc = {
+      IdxT(elem.x * IdxT(a_strides[ndim - 1])) +
+          IdxT(elem.y * IdxT(a_strides[ndim - 2])),
+      IdxT(elem.x * IdxT(b_strides[ndim - 1])) +
+          IdxT(elem.y * IdxT(b_strides[ndim - 2])),
+      IdxT(elem.x * IdxT(c_strides[ndim - 1])) +
+          IdxT(elem.y * IdxT(c_strides[ndim - 2]))};
+  for (int d = ndim - 3; d >= 0; --d) {
+    uint l = elem.z % shape[d];
+    loc.x += l * IdxT(a_strides[d]);
+    loc.y += l * IdxT(b_strides[d]);
+    loc.z += l * IdxT(c_strides[d]);
+    elem.z /= shape[d];
+  }
+  return loc;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Elem to loc in a loop utils
+///////////////////////////////////////////////////////////////////////////////
+
+template <int DIM, typename OffsetT = size_t, bool General = true>
+struct LoopedElemToLoc {
+  int dim;
+  LoopedElemToLoc<DIM - 1, OffsetT, General> inner_looper;
+  OffsetT offset{0};
+  int index{0};
+
+  LoopedElemToLoc(int dim) : dim(dim), inner_looper(dim - 1) {}
+
+  void next(const constant int* shape, const constant int64_t* strides) {
+    if (dim == 0) {
+      return;
+    }
+    index++;
+    offset += OffsetT(strides[dim - 1]);
+    if (index >= shape[dim - 1]) {
+      index = 0;
+      inner_looper.next(shape, strides);
+      offset = inner_looper.offset;
+    }
+  }
+
+  void next(int n, const constant int* shape, const constant int64_t* strides) {
+    if (dim == 0) {
+      return;
+    }
+    index += n;
+    offset += n * OffsetT(strides[dim - 1]);
+
+    if (index >= shape[dim - 1]) {
+      int extra = index - shape[dim - 1];
+      if (extra >= shape[dim - 1]) {
+        inner_looper.next(1 + extra / shape[dim - 1], shape, strides);
+        extra = extra % shape[dim - 1];
+      } else {
+        inner_looper.next(shape, strides);
+      }
+      index = 0;
+      offset = inner_looper.offset;
+      if (extra > 0) {
+        next(extra, shape, strides);
+      }
+    }
+  }
+
+  OffsetT location() {
+    return offset;
+  }
+};
+
+template <typename OffsetT>
+struct LoopedElemToLoc<1, OffsetT, true> {
+  int dim;
+  OffsetT offset{0};
+  uint index{0};
+
+  LoopedElemToLoc(int dim) : dim(dim) {}
+
+  void next(const constant int* shape, const constant int64_t* strides) {
+    index++;
+    if (dim > 1) {
+      offset = elem_to_loc<OffsetT>(index, shape, strides, dim);
+    } else {
+      offset += OffsetT(strides[0]);
+    }
+  }
+
+  void next(int n, const constant int* shape, const constant int64_t* strides) {
+    index += n;
+    if (dim > 1) {
+      offset = elem_to_loc<OffsetT>(index, shape, strides, dim);
+    } else {
+      offset = index * OffsetT(strides[0]);
+    }
+  }
+
+  OffsetT location() {
+    return offset;
+  }
+};
+
+template <typename OffsetT>
+struct LoopedElemToLoc<1, OffsetT, false> {
+  OffsetT offset{0};
+
+  LoopedElemToLoc(int) {}
+
+  void next(const constant int*, const constant int64_t* strides) {
+    offset += OffsetT(strides[0]);
+  }
+
+  void next(int n, const constant int*, const constant int64_t* strides) {
+    offset += n * OffsetT(strides[0]);
+  }
+
+  OffsetT location() {
+    return offset;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Calculation utils
+///////////////////////////////////////////////////////////////////////////////
+
+/** Compute ceil((float)N/(float)M) */
+template <typename T, typename U>
+inline T ceildiv(T N, U M) {
+  return (N + M - 1) / M;
+}
+
+// https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html#1202
+inline float log1p(float x) {
+  float xp1 = 1.0f + x;
+  if (xp1 == Limits<float>::max) {
+    return Limits<float>::max;
+  }
+  if (xp1 == 1.0f) {
+    return x;
+  }
+
+  return x * (metal::log(xp1) / (xp1 - 1.0f));
+}
+
+inline bfloat16_t log1p(bfloat16_t x) {
+  float xp1 = 1.0f + static_cast<float>(x);
+  if (xp1 == Limits<float>::max) {
+    return Limits<bfloat16_t>::max;
+  }
+  if (xp1 == 1.0f) {
+    return x;
+  }
+
+  return bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f)));
+}
+
+inline complex64_t log1p(complex64_t in) {
+  float x = in.real;
+  float y = in.imag;
+  float zabs = metal::precise::sqrt(x * x + y * y);
+  float theta = metal::atan2(y, x + 1);
+  if (zabs < 0.5f) {
+    float r = x * (2 + x) + y * y;
+    if (r == 0) { // handle underflow
+      return {x, theta};
+    }
+    return {0.5f * log1p(r), theta};
+  } else {
+    auto z0 = metal::sqrt((x + 1) * (x + 1) + y * y);
+    return {metal::log(z0), theta};
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// SIMD shuffle ops
+///////////////////////////////////////////////////////////////////////////////
+
+inline uint64_t simd_shuffle_down(uint64_t data, uint16_t delta) {
+  return as_type<uint64_t>(
+      metal::simd_shuffle_down(as_type<uint2>(data), delta));
+}
+
+inline int64_t simd_shuffle_down(int64_t data, uint16_t delta) {
+  return as_type<int64_t>(
+      metal::simd_shuffle_down(as_type<uint2>(data), delta));
+}
+
+inline bool simd_shuffle_down(bool data, uint16_t delta) {
+  return simd_shuffle_down(static_cast<uint32_t>(data), delta);
+}
+
+inline complex64_t simd_shuffle_down(complex64_t data, uint16_t delta) {
+  return complex64_t(
+      simd_shuffle_down(data.real, delta), simd_shuffle_down(data.imag, delta));
+}
+
+inline uint64_t simd_shuffle_up(uint64_t data, uint16_t delta) {
+  return as_type<uint64_t>(metal::simd_shuffle_up(as_type<uint2>(data), delta));
+}
+
+inline int64_t simd_shuffle_up(int64_t data, uint16_t delta) {
+  return as_type<int64_t>(metal::simd_shuffle_up(as_type<uint2>(data), delta));
+}
+
+inline bool simd_shuffle_up(bool data, uint16_t delta) {
+  return simd_shuffle_up(static_cast<uint32_t>(data), delta);
+}
+
+inline complex64_t simd_shuffle_up(complex64_t data, uint16_t delta) {
+  return complex64_t(
+      simd_shuffle_up(data.real, delta), simd_shuffle_up(data.imag, delta));
+}
+
+inline uint64_t
+simd_shuffle_and_fill_up(uint64_t data, uint64_t filling, uint16_t delta) {
+  return as_type<uint64_t>(metal::simd_shuffle_and_fill_up(
+      as_type<uint2>(data), as_type<uint2>(filling), delta));
+}
+
+inline int64_t
+simd_shuffle_and_fill_up(int64_t data, int64_t filling, uint16_t delta) {
+  return as_type<int64_t>(metal::simd_shuffle_and_fill_up(
+      as_type<uint2>(data), as_type<uint2>(filling), delta));
+}
+
+inline bool simd_shuffle_and_fill_up(bool data, bool filling, uint16_t delta) {
+  return simd_shuffle_and_fill_up(
+      static_cast<uint32_t>(data), static_cast<uint32_t>(filling), delta);
+}
+
+inline complex64_t simd_shuffle_and_fill_up(
+    complex64_t data,
+    complex64_t filling,
+    uint16_t delta) {
+  return complex64_t(
+      simd_shuffle_and_fill_up(data.real, filling.real, delta),
+      simd_shuffle_and_fill_up(data.imag, filling.imag, delta));
+}
+
+inline uint64_t simd_shuffle(uint64_t data, uint16_t lane) {
+  return as_type<uint64_t>(metal::simd_shuffle(as_type<uint2>(data), lane));
+}
+
+inline int64_t simd_shuffle(int64_t data, uint16_t lane) {
+  return as_type<int64_t>(metal::simd_shuffle(as_type<uint2>(data), lane));
+}
+
+inline bool simd_shuffle(bool data, uint16_t lane) {
+  return simd_shuffle(static_cast<uint32_t>(data), lane);
+}
+
+inline complex64_t simd_shuffle(complex64_t data, uint16_t lane) {
+  return complex64_t(
+      simd_shuffle(data.real, lane), simd_shuffle(data.imag, lane));
+}
+
+// std::conditional is not included with Metal
+template <bool condition, typename T, typename U>
+struct ConditionalType {
+  using type = U;
+};
+
+template <typename T, typename U>
+struct ConditionalType<true, T, U> {
+  using type = T;
+};
diff --git a/macos/third_party/mlx/mlx/backend/metal/logsumexp.cpp b/macos/third_party/mlx/mlx/backend/metal/logsumexp.cpp
new file mode 100644
index 0000000..2cfdcdc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/logsumexp.cpp
@@ -0,0 +1,95 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <algorithm>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+constexpr int LOGSUMEXP_LOOPED_LIMIT = 4096;
+
+void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  if (!issubdtype(out.dtype(), floating)) {
+    throw std::runtime_error(
+        "[logsumexp] Does not support non-floating point types.");
+  }
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  // Make sure that the last dimension is contiguous
+  auto ensure_contiguous = [&s, &d](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      d.add_temporary(x_copy, s.index);
+      return x_copy;
+    }
+  };
+
+  auto in = ensure_contiguous(inputs[0]);
+  if (in.flags().row_contiguous) {
+    out.set_data(allocator::malloc(out.nbytes()));
+  } else {
+    auto n = in.shape(-1);
+    auto flags = in.flags();
+    auto strides = in.strides();
+    for (auto& s : strides) {
+      s /= n;
+    }
+    bool col_contig = strides[0] == 1;
+    for (int i = 1; col_contig && i < strides.size(); ++i) {
+      col_contig &=
+          (out.shape(i) == 1 || strides[i - 1] == out.shape(i) * strides[i]);
+    }
+    flags.col_contiguous = col_contig;
+    out.set_data(
+        allocator::malloc(in.nbytes() / n),
+        in.data_size() / n,
+        std::move(strides),
+        flags);
+  }
+
+  int axis_size = in.shape().back();
+  int n_rows = in.data_size() / axis_size;
+
+  const int simd_size = 32;
+  const int n_reads = 4;
+  const int looped_limit = LOGSUMEXP_LOOPED_LIMIT;
+
+  std::string kernel_name = (axis_size > looped_limit) ? "looped_" : "block_";
+  kernel_name += "logsumexp_";
+  kernel_name += type_to_name(out);
+
+  auto kernel = get_logsumexp_kernel(d, kernel_name, out);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  {
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(axis_size, 2);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/make_compiled_preamble.sh b/macos/third_party/mlx/mlx/backend/metal/make_compiled_preamble.sh
new file mode 100644
index 0000000..bb55ed3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/make_compiled_preamble.sh
@@ -0,0 +1,120 @@
+#!/bin/bash
+#
+# This script generates a C++ function that provides the Metal source code
+# at runtime for use with kernel generation.
+#
+# The steps executed are as follows 
+# - Take as input a metal header file in the mlx metal backend 
+# - Use the metal compiler to expand the dependency headers 
+# - Sort the headers in order of inclusion 
+# - Expand the headers in order of inclusion 
+# - Export the generated source code content as a C++ function
+#
+# Doing the expansion this way allows us to retain macros, comments, and 
+# formatting in the expanded source. This adds user readibility, and also 
+# enables use of the metal macros in the source code which can then be 
+# handled by the metal runtime compiler
+#
+# Copyright © 2023-25 Apple Inc.
+
+OUTPUT_DIR=$1
+CC=$2
+SRC_DIR=$3
+SRC_FILE=$4
+CFLAGS=$5
+SRC_NAME=$(basename -- "${SRC_FILE}")
+JIT_INCLUDES=${SRC_DIR}/mlx/backend/metal/kernels/jit
+INPUT_FILE=${SRC_DIR}/mlx/backend/metal/kernels/${SRC_FILE}.h
+OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp
+
+# Prepare output
+mkdir -p "$OUTPUT_DIR"
+
+# Use the metal compiler to get a list of headers (with depth)
+CCC="xcrun -sdk macosx metal -x metal"
+HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null )
+
+# Remove any included system frameworks (for MetalPerformancePrimitive headers)
+HDRS=$(echo "$HDRS" | grep -v "Xcode")
+
+# Use the header depth to sort the files in order of inclusion
+declare -a HDRS_LIST=($HDRS)
+declare -a HDRS_STACK=()
+declare -a HDRS_SORTED=()
+
+length=${#HDRS_LIST[@]}
+
+HDRS_LIST+=(".")
+
+for ((i=0; i<${length}; i+=2));
+do 
+
+  header="${HDRS_LIST[$i+1]#$SRC_DIR/}"
+
+  str_this="${HDRS_LIST[$i]}"
+  str_next="${HDRS_LIST[$i + 2]}"
+
+  depth_this=${#str_this}
+  depth_next=${#str_next}
+
+  # If we have a dependency then we stack it
+  if [ $depth_next -gt $depth_this ]; then 
+    HDRS_STACK=($header ${HDRS_STACK[@]})
+
+  # If we are done with this level 
+  else 
+    # We add the header to out list
+    HDRS_SORTED+=($header) 
+
+    # Pop the stacked up dependencies
+    pop_len=$((depth_this - depth_next))
+    for popped_header in "${HDRS_STACK[@]:0:$pop_len}"
+    do 
+      HDRS_SORTED+=($popped_header)
+    done 
+
+    HDRS_STACK=(${HDRS_STACK[@]:$pop_len})
+  fi  
+
+done
+
+# Make sure the given metal header is also expanded in the source content
+HDRS_SORTED+=("${INPUT_FILE#$SRC_DIR/}")
+
+# Expand the headers in order of inclusion 
+CONTENT=$(
+echo "// Copyright © 2025 Apple Inc."
+echo "" 
+echo "// Auto generated source for ${INPUT_FILE#$SRC_DIR/}"
+echo ""
+
+for header in "${HDRS_SORTED[@]}"
+do 
+  echo "///////////////////////////////////////////////////////////////////////////////"
+  echo "// Contents from \"${header}\""
+  echo "///////////////////////////////////////////////////////////////////////////////"
+  echo ""
+
+  echo "#line 1 \"${header}\""
+
+  grep -h -v -G -e "#include \".*.h\"" -e "#pragma once" "${SRC_DIR}/${header}" 
+  
+  echo ""
+  
+done
+
+echo "///////////////////////////////////////////////////////////////////////////////"
+)
+
+# Export the generated source code content as a C++ function
+cat << EOF > "$OUTPUT_FILE"
+namespace mlx::core::metal {
+
+const char* $SRC_NAME() {
+  return R"preamble(
+$CONTENT
+)preamble";
+}
+
+} // namespace mlx::core::metal
+EOF
diff --git a/macos/third_party/mlx/mlx/backend/metal/matmul.cpp b/macos/third_party/mlx/mlx/backend/metal/matmul.cpp
new file mode 100644
index 0000000..98d7b66
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/matmul.cpp
@@ -0,0 +1,2572 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <algorithm>
+#include <cassert>
+#include <numeric>
+#include <sstream>
+
+#include "mlx/backend/common/broadcasting.h"
+#include "mlx/backend/common/matmul.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/binary.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/steel/gemm/params.h"
+#include "mlx/backend/metal/matmul.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+std::tuple<bool, int64_t, array> check_transpose(
+    std::vector<array>& copies,
+    const Stream& s,
+    const array& arr,
+    bool is_vector) {
+  auto stx = arr.strides()[arr.ndim() - 2];
+  auto sty = arr.strides()[arr.ndim() - 1];
+  if (sty == 1 && (!is_vector || stx == arr.shape(-1))) {
+    return std::make_tuple(false, stx, arr);
+  } else if (stx == 1 && (!is_vector || sty == arr.shape(-2))) {
+    return std::make_tuple(true, sty, arr);
+  } else {
+    array arr_copy = contiguous_copy_gpu(arr, s);
+    copies.push_back(arr_copy);
+    return std::make_tuple(false, arr.shape(-1), arr_copy);
+  }
+};
+
+inline array
+ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    d.add_temporary(x_copy, s.index);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+inline std::tuple<bool, int64_t, array>
+ensure_batch_contiguous(const array& x, metal::Device& d, const Stream& s) {
+  if (x.flags().row_contiguous) {
+    return std::make_tuple(false, x.strides()[x.ndim() - 2], x);
+  }
+
+  bool rc = true;
+  for (int i = 0; i < x.ndim() - 3; i++) {
+    rc &= x.strides()[i + 1] * x.shape(i) == x.strides()[i];
+  }
+  if (rc) {
+    auto stx = x.strides()[x.ndim() - 2];
+    auto sty = x.strides()[x.ndim() - 1];
+    auto K = x.shape(-2);
+    auto N = x.shape(-1);
+    if (sty == 1 && (N != 1 || stx == N)) {
+      return std::make_tuple(false, stx, x);
+    }
+    if (stx == 1 && (N != 1 || sty == K)) {
+      return std::make_tuple(true, sty, x);
+    }
+  }
+
+  array x_copy = contiguous_copy_gpu(x, s);
+  d.add_temporary(x_copy, s.index);
+  return std::make_tuple(false, x_copy.strides()[x_copy.ndim() - 2], x_copy);
+}
+
+} // namespace
+
+///////////////////////////////////////////////////////////////////////////////
+// Steel matmul fallback
+///////////////////////////////////////////////////////////////////////////////
+
+#define GEMM_TPARAM_MACRO(devc)                                           \
+  if (devc == 'g' || devc == 'p') { /* Small device */                    \
+    if (out.dtype() == complex64) {                                       \
+      bm = 64;                                                            \
+      bn = 32;                                                            \
+      bk = 8;                                                             \
+      wm = 4;                                                             \
+      wn = 1;                                                             \
+    } else if (!transpose_a && transpose_b) { /* nt */                    \
+      bm = 64;                                                            \
+      bn = 32;                                                            \
+      bk = 32;                                                            \
+      wm = 2;                                                             \
+      wn = 2;                                                             \
+    } else if (out.dtype() != float32) { /* half and bfloat */            \
+      bm = 64;                                                            \
+      bn = 64;                                                            \
+      bk = 16;                                                            \
+      wm = 1;                                                             \
+      wn = 2;                                                             \
+    }                                                                     \
+  } else if (devc == 'd') { /* Large device */                            \
+    if ((size_t)batch_size_out * M * N >= 1ul << 20) { /* large matmul */ \
+      if (out.dtype() != float32) { /* half and bfloat */                 \
+        if (2 * std::max(M, N) > K) { /* Reasonable K */                  \
+          bm = 64;                                                        \
+          bn = 64;                                                        \
+          bk = 16;                                                        \
+          wm = 1;                                                         \
+          wn = 2;                                                         \
+        } else if (!transpose_a && transpose_b) { /* nt with large k */   \
+          bm = 64;                                                        \
+          bn = 32;                                                        \
+          bk = 32;                                                        \
+          wm = 2;                                                         \
+          wn = 2;                                                         \
+        } else { /* nn with large K */                                    \
+          bm = 32;                                                        \
+          bn = 64;                                                        \
+          bk = 16;                                                        \
+          wm = 1;                                                         \
+          wn = 2;                                                         \
+        }                                                                 \
+      } /* float takes default */                                         \
+    } else { /* smaller matmul */                                         \
+      if (out.dtype() != float32) { /* half and bfloat */                 \
+        if (!transpose_a && transpose_b) { /* nt */                       \
+          bm = 64;                                                        \
+          bn = 32;                                                        \
+          bk = 32;                                                        \
+          wm = 2;                                                         \
+          wn = 2;                                                         \
+        } else { /* nn */                                                 \
+          bm = 64;                                                        \
+          bn = 64;                                                        \
+          bk = 16;                                                        \
+          wm = 1;                                                         \
+          wn = 2;                                                         \
+        }                                                                 \
+      } else { /* floats */                                               \
+        if (!transpose_a && transpose_b) { /* nt */                       \
+          bm = 32;                                                        \
+          bn = 64;                                                        \
+          bk = 16;                                                        \
+          wm = 1;                                                         \
+          wn = 2;                                                         \
+        } else { /* nn */                                                 \
+          bm = 64;                                                        \
+          bn = 32;                                                        \
+          bk = 32;                                                        \
+          wm = 2;                                                         \
+          wn = 2;                                                         \
+        }                                                                 \
+      }                                                                   \
+    }                                                                     \
+  } else { /* Medium device */                                            \
+    bm = 64;                                                              \
+    bn = 64;                                                              \
+    bk = 16;                                                              \
+    wm = 2;                                                               \
+    wn = 2;                                                               \
+  }
+
+///////////////////////////////////////////////////////////////////////////////
+// Regular steel matmul dispatch
+///////////////////////////////////////////////////////////////////////////////
+
+template <bool CHECK_AB>
+void steel_matmul_regular_axpby_nax(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    int ldd,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape,
+    Strides batch_strides,
+    int64_t A_batch_stride,
+    int64_t B_batch_stride,
+    int64_t matrix_stride_out,
+    int64_t C_batch_stride /* = 0*/,
+    float alpha /* = 1.0f */,
+    float beta /* = 0.0f */) {
+  using namespace mlx::steel;
+
+  // Determine dispatch kernel
+  int bm = 128, bn = 128, bk = 512;
+  int wm = 4, wn = 4;
+
+  // Prepare kernel name
+  std::ostringstream kname;
+
+  // clang-format off
+  kname << "steel_gemm_fused_nax_"
+        << (transpose_a ? 't' : 'n')
+        << (transpose_b ? 't' : 'n')
+        << "_" << type_to_name(a)
+        << "_" << type_to_name(out)
+        << "_bm" << bm << "_bn" << bn << "_bk" << bk
+        << "_wm" << wm << "_wn" << wn; // clang-format on
+
+  std::string base_name = kname.str();
+
+  const bool has_batch = (batch_shape.size() > 1);
+  const bool use_out_source = CHECK_AB && (alpha != 0.0f || beta != 1.0f);
+  const bool do_axpby = use_out_source && (alpha != 1.0f || beta != 1.0f);
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  metal::MTLFCList func_consts = {
+      {&has_batch, MTL::DataType::DataTypeBool, 10},
+      {&use_out_source, MTL::DataType::DataTypeBool, 100},
+      {&do_axpby, MTL::DataType::DataTypeBool, 110},
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+      {&align_K, MTL::DataType::DataTypeBool, 202},
+  };
+
+  // clang-format off
+  kname << "_has_batch_" << (has_batch ? 't' : 'n')
+        << "_use_out_source_" << (use_out_source ? 't' : 'n')
+        << "_do_axpby_" << (do_axpby ? 't' : 'n')
+        << "_align_M_" << (align_M ? 't' : 'n')
+        << "_align_N_" << (align_N ? 't' : 'n')
+        << "_align_K_" << (align_K ? 't' : 'n'); // clang-format on
+
+  std::string hash_name = kname.str();
+
+  // Encode and dispatch kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_fused_nax_kernel(
+      /* metal::Device& d = */ d,
+      /* const std::string& kernel_name = */ base_name,
+      /* const std::string& hash_name = */ hash_name,
+      /* const metal::MTLFCList& func_consts = */ func_consts,
+      /* const array& out = */ out,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* int bm = */ bm,
+      /* int bn = */ bn,
+      /* int bk = */ bk,
+      /* int wm = */ wm,
+      /* int wn = */ wn);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Use problem size to determine threadblock swizzle
+  int tn = (N + bn - 1) / bn;
+  int tm = (M + bm - 1) / bm;
+
+  // TODO: Explore device-based tuning for swizzle
+  int swizzle_log = tm <= 3 ? 0 : 1;
+
+  // Prepare steel matmul params
+  GEMMParams params{/* const int M = */ M,
+                    /* const int N = */ N,
+                    /* const int K = */ K,
+                    /* const int lda = */ lda,
+                    /* const int ldb = */ ldb,
+                    /* const int ldd = */ ldd,
+                    /* const int tiles_n = */ tn,
+                    /* const int tiles_m = */ tm,
+                    /* const int64_t batch_stride_a = */ A_batch_stride,
+                    /* const int64_t batch_stride_b = */ B_batch_stride,
+                    /* const int64_t batch_stride_d = */ matrix_stride_out,
+                    /* const int swizzle_log = */ swizzle_log,
+                    /* const int gemm_k_iterations_aligned = */ (K / bk),
+                    /* const int batch_ndim = */ int(batch_shape.size())};
+
+  // Prepare launch grid params
+  int tile = 1 << swizzle_log;
+  tm = (tm + tile - 1) / tile;
+  tn = tn * tile;
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(tn, tm, batch_size_out);
+
+  // Launch kernel
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_output_array(out, 3);
+
+  compute_encoder.set_bytes(params, 4);
+
+  if (has_batch) {
+    compute_encoder.set_vector_bytes(batch_shape, 6);
+    compute_encoder.set_vector_bytes(batch_strides, 7);
+  }
+
+  if (use_out_source) {
+    int ldc = c.strides()[c.ndim() - 2];
+    int fdc = c.strides()[c.ndim() - 1];
+
+    GEMMAddMMParams params{/* const int ldc = */ ldc,
+                           /* const int fdc = */ fdc,
+                           /* const int64_t batch_stride_c = */ C_batch_stride,
+                           /* const float alpha = */ alpha,
+                           /* const float beta = */ beta};
+
+    compute_encoder.set_input_array(c, 2);
+    compute_encoder.set_bytes(params, 5);
+  }
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  // Record copies
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+template <bool CHECK_AB>
+void steel_matmul_regular_axpby(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    int ldd,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape,
+    Strides batch_strides,
+    int64_t A_batch_stride,
+    int64_t B_batch_stride,
+    int64_t matrix_stride_out,
+    int64_t C_batch_stride /* = 0*/,
+    float alpha /* = 1.0f */,
+    float beta /* = 0.0f */) {
+  if (metal::is_nax_available() && !issubdtype(a.dtype(), complexfloating) &&
+      (env::enable_tf32() || a.dtype() != float32)) {
+    return steel_matmul_regular_axpby_nax<CHECK_AB>(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& a = */ a,
+        /* const array& b = */ b,
+        /* const array& c = */ c,
+        /* array& out = */ out,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* int batch_size_out = */ batch_size_out,
+        /* int lda = */ lda,
+        /* int ldb = */ ldb,
+        /* int ldd = */ ldd,
+        /* bool transpose_a = */ transpose_a,
+        /* bool transpose_b = */ transpose_b,
+        /* std::vector<array>& copies = */ copies,
+        /* Shape batch_shape = */ batch_shape,
+        /* Strides batch_strides = */ batch_strides,
+        /* int64_t A_batch_stride = */ A_batch_stride,
+        /* int64_t B_batch_stride = */ B_batch_stride,
+        /* int64_t matrix_stride_out = */ matrix_stride_out,
+        /* int64_t C_batch_stride = */ C_batch_stride,
+        /* float alpha = */ alpha,
+        /* float beta = */ beta);
+  }
+
+  using namespace mlx::steel;
+
+  // Determine dispatch kernel
+  int bm = 64, bn = 64, bk = 16;
+  int wm = 2, wn = 2;
+
+  char devc = d.get_architecture().back();
+  GEMM_TPARAM_MACRO(devc)
+
+  // Prepare kernel name
+  std::ostringstream kname;
+
+  // clang-format off
+  kname << "steel_gemm_fused_"
+        << (transpose_a ? 't' : 'n')
+        << (transpose_b ? 't' : 'n')
+        << "_" << type_to_name(a)
+        << "_" << type_to_name(out)
+        << "_bm" << bm << "_bn" << bn << "_bk" << bk
+        << "_wm" << wm << "_wn" << wn; // clang-format on
+
+  std::string base_name = kname.str();
+
+  const bool has_batch = (batch_shape.size() > 1);
+  const bool use_out_source = CHECK_AB && (alpha != 0.0f || beta != 1.0f);
+  const bool do_axpby = use_out_source && (alpha != 1.0f || beta != 1.0f);
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  metal::MTLFCList func_consts = {
+      {&has_batch, MTL::DataType::DataTypeBool, 10},
+      {&use_out_source, MTL::DataType::DataTypeBool, 100},
+      {&do_axpby, MTL::DataType::DataTypeBool, 110},
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+      {&align_K, MTL::DataType::DataTypeBool, 202},
+  };
+
+  // clang-format off
+  kname << "_has_batch_" << (has_batch ? 't' : 'n')
+        << "_use_out_source_" << (use_out_source ? 't' : 'n')
+        << "_do_axpby_" << (do_axpby ? 't' : 'n')
+        << "_align_M_" << (align_M ? 't' : 'n')
+        << "_align_N_" << (align_N ? 't' : 'n')
+        << "_align_K_" << (align_K ? 't' : 'n'); // clang-format on
+
+  std::string hash_name = kname.str();
+
+  // Encode and dispatch kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_fused_kernel(
+      /* metal::Device& d = */ d,
+      /* const std::string& kernel_name = */ base_name,
+      /* const std::string& hash_name = */ hash_name,
+      /* const metal::MTLFCList& func_consts = */ func_consts,
+      /* const array& out = */ out,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* int bm = */ bm,
+      /* int bn = */ bn,
+      /* int bk = */ bk,
+      /* int wm = */ wm,
+      /* int wn = */ wn);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Use problem size to determine threadblock swizzle
+  int tn = (N + bn - 1) / bn;
+  int tm = (M + bm - 1) / bm;
+
+  // TODO: Explore device-based tuning for swizzle
+  int swizzle_log = 0; // tm >= 6 ? 3 : (tm <= 3 ? 0 : 2);
+
+  // Prepare steel matmul params
+  GEMMParams params{/* const int M = */ M,
+                    /* const int N = */ N,
+                    /* const int K = */ K,
+                    /* const int lda = */ lda,
+                    /* const int ldb = */ ldb,
+                    /* const int ldd = */ ldd,
+                    /* const int tiles_n = */ tn,
+                    /* const int tiles_m = */ tm,
+                    /* const int64_t batch_stride_a = */ A_batch_stride,
+                    /* const int64_t batch_stride_b = */ B_batch_stride,
+                    /* const int64_t batch_stride_d = */ matrix_stride_out,
+                    /* const int swizzle_log = */ swizzle_log,
+                    /* const int gemm_k_iterations_aligned = */ (K / bk),
+                    /* const int batch_ndim = */ int(batch_shape.size())};
+
+  // Prepare launch grid params
+  int tile = 1 << swizzle_log;
+  tm = (tm + tile - 1) / tile;
+  tn = tn * tile;
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(tn, tm, batch_size_out);
+
+  // Launch kernel
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_output_array(out, 3);
+
+  compute_encoder.set_bytes(params, 4);
+
+  if (has_batch) {
+    compute_encoder.set_vector_bytes(batch_shape, 6);
+    compute_encoder.set_vector_bytes(batch_strides, 7);
+  }
+
+  if (use_out_source) {
+    int ldc = c.strides()[c.ndim() - 2];
+    int fdc = c.strides()[c.ndim() - 1];
+
+    GEMMAddMMParams params{/* const int ldc = */ ldc,
+                           /* const int fdc = */ fdc,
+                           /* const int64_t batch_stride_c = */ C_batch_stride,
+                           /* const float alpha = */ alpha,
+                           /* const float beta = */ beta};
+
+    compute_encoder.set_input_array(c, 2);
+    compute_encoder.set_bytes(params, 5);
+  }
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  // Record copies
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Split k steel matmul
+///////////////////////////////////////////////////////////////////////////////
+
+template <bool CHECK_AB = true>
+void steel_gemm_splitk_axpby(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    float alpha = 1.0f,
+    float beta = 0.0f) {
+  using namespace mlx::steel;
+
+  int _tm = (M + 32 - 1) / 32;
+  int _tn = (N + 32 - 1) / 32;
+  int _tk = K / 16;
+
+  int bm = M < 40 ? 16 : 32;
+  int bn = N < 40 ? 16 : 32;
+  int bk = 16;
+  int wm = 2, wn = 2;
+
+  // As _tk grows use more partitions, as _tm * _tn grow use fewer partitions
+  int split_k_partitions =
+      std::min(std::max(2, next_power_of_2(_tk / (_tm * _tn))), 32);
+  int split_k_partition_stride = M * N;
+  int gemm_k_iterations = (K / bk) / split_k_partitions;
+  int split_k_partition_size = gemm_k_iterations * bk;
+
+  array C_split(
+      {split_k_partitions, M, N},
+      issubdtype(out.dtype(), complexfloating) ? complex64 : float32,
+      nullptr,
+      {});
+  C_split.set_data(allocator::malloc(C_split.nbytes()));
+  copies.push_back(C_split);
+
+  bool mn_aligned = M % bm == 0 && N % bn == 0;
+  bool k_aligned = K % bk == 0;
+  std::ostringstream kname;
+
+  // clang-format off
+  kname << "steel_gemm_splitk_"
+        << (transpose_a ? 't' : 'n')
+        << (transpose_b ? 't' : 'n')
+        << "_" << type_to_name(a)
+        << "_" << type_to_name(C_split)
+        << "_bm" << bm << "_bn" << bn << "_bk" << bk
+        << "_wm" << wm << "_wn" << wn
+        << "_MN_" << (mn_aligned ? "t" : "n") << "aligned"
+        << "_K_" << (k_aligned ? "t" : "n") << "aligned"; // clang-format on
+
+  // Encode and dispatch gemm kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_splitk_kernel(
+      /* metal::Device& d = */ d,
+      /* const std::string& kernel_name = */ kname.str(),
+      /* const array& in = */ a,
+      /* const array& out = */ C_split,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* int bm = */ bm,
+      /* int bn = */ bn,
+      /* int bk = */ bk,
+      /* int wm = */ wm,
+      /* int wn = */ wn,
+      /* bool mn_aligned = */ mn_aligned,
+      /* bool k_aligned = */ k_aligned);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int tn = (N + bn - 1) / bn;
+  int tm = (M + bm - 1) / bm;
+
+  GEMMSpiltKParams params{
+      /* const int M = */ M,
+      /* const int N = */ N,
+      /* const int K = */ K,
+      /* const int lda = */ lda,
+      /* const int ldb = */ ldb,
+      /* const int ldc = */ N,
+      /* const int tiles_n = */ tn,
+      /* const int tiles_m = */ tm,
+      /* const int split_k_partitions = */ split_k_partitions,
+      /* const int split_k_partition_stride = */ split_k_partition_stride,
+      /* const int split_k_partition_size = */ split_k_partition_size,
+      /* const int swizzle_log = */ 0, // no swizzle
+      /* const int gemm_k_iterations_aligned = */ gemm_k_iterations};
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(tn, tm, split_k_partitions);
+
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_output_array(C_split, 2);
+
+  compute_encoder.set_bytes(params, 3);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  // Do accum kernel
+  {
+    const bool do_axpby = CHECK_AB && (alpha != 1.0f || beta != 0.0f);
+
+    auto kernel_name = "steel_gemm_splitk_accum_" + type_to_name(out) + "_" +
+        type_to_name(C_split);
+
+    if (do_axpby) {
+      kernel_name = kernel_name + "_axbpy";
+    }
+
+    auto kernel = get_steel_gemm_splitk_accum_kernel(
+        /* metal::Device& d = */ d,
+        /* const std::string& kernel_name = */ kernel_name,
+        /* const array& in = */ C_split,
+        /* const array& out = */ out,
+        /* bool axbpy = */ do_axpby);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    // Set the arguments for the kernel
+    compute_encoder.set_input_array(C_split, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(split_k_partitions, 2);
+    compute_encoder.set_bytes(split_k_partition_stride, 3);
+    compute_encoder.set_bytes(N, 4);
+
+    if (do_axpby) {
+      int ldc = c.strides()[c.ndim() - 2];
+      int fdc = c.strides()[c.ndim() - 1];
+
+      compute_encoder.set_input_array(c, 5);
+      compute_encoder.set_bytes(ldc, 6);
+      compute_encoder.set_bytes(fdc, 7);
+      compute_encoder.set_bytes(alpha, 8);
+      compute_encoder.set_bytes(beta, 9);
+    }
+
+    // Launch enough thread groups for each output
+    MTL::Size grid_dims = MTL::Size(N, M, 1);
+    auto group_dims = get_block_dims(N, M, 1);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// NAX Split k steel matmul
+///////////////////////////////////////////////////////////////////////////////
+
+template <bool CHECK_AB = true>
+void steel_gemm_splitk_axpby_nax(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    float alpha = 1.0f,
+    float beta = 0.0f) {
+  using namespace mlx::steel;
+
+  constexpr int bm = 128, bn = 128, bk = 512;
+  constexpr int wm = 4, wn = 4;
+
+  // Determine how many partitions to split K into
+  constexpr int split_k_partition_size = 3072;
+  int split_k_partitions =
+      (K + split_k_partition_size - 1) / split_k_partition_size;
+
+  const int bk_iters_per_partition = split_k_partition_size / bk;
+  const int split_k_partition_stride = M * N;
+
+  array C_split({split_k_partitions, M, N}, float32, nullptr, {});
+  C_split.set_data(allocator::malloc(C_split.nbytes()));
+  copies.push_back(C_split);
+
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  // Per-tile align_K is checked at runtime; only the last tile can be unaligned
+  metal::MTLFCList func_consts = {
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201}};
+
+  std::ostringstream kname;
+
+  // clang-format off
+  kname << "steel_gemm_splitk_nax_"
+        << (transpose_a ? 't' : 'n')
+        << (transpose_b ? 't' : 'n')
+        << "_" << type_to_name(a)
+        << "_" << type_to_name(C_split)
+        << "_bm" << bm << "_bn" << bn << "_bk" << bk
+        << "_wm" << wm << "_wn" << wn; // clang-format on
+
+  std::string base_name = kname.str();
+
+  // clang-format off
+  kname << "_align_M_" << (align_M ? 't' : 'n')
+        << "_align_N_" << (align_N ? 't' : 'n')
+        << "_align_K_" << (align_K ? 't' : 'n'); // clang-format on
+
+  std::string hash_name = kname.str();
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_splitk_nax_kernel(
+      /* metal::Device& d = */ d,
+      /* const std::string& kernel_name = */ base_name,
+      /* const std::string& hash_name = */ hash_name,
+      /* const metal::MTLFCList& func_consts = */ func_consts,
+      /* const array& out = */ C_split,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* int bm = */ bm,
+      /* int bn = */ bn,
+      /* int bk = */ bk,
+      /* int wm = */ wm,
+      /* int wn = */ wn);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int tn = (N + bn - 1) / bn;
+  int tm = (M + bm - 1) / bm;
+
+  int swizzle_log = tm <= 3 ? 0 : 1;
+
+  // Compute swizzled tile counts
+  int tile = 1 << swizzle_log;
+  int tm_swizzled = (tm + tile - 1) / tile;
+  int tn_swizzled = tn * tile;
+
+  GEMMSpiltKParams params{
+      /* const int M = */ M,
+      /* const int N = */ N,
+      /* const int K = */ K,
+      /* const int lda = */ lda,
+      /* const int ldb = */ ldb,
+      /* const int ldc = */ N,
+      /* const int tiles_n = */ tn,
+      /* const int tiles_m = */ tm,
+      /* const int split_k_partitions = */ split_k_partitions,
+      /* const int split_k_partition_stride = */ split_k_partition_stride,
+      /* const int split_k_partition_size = */ split_k_partition_size,
+      /* const int swizzle_log = */ swizzle_log,
+      /* const int gemm_k_iterations_aligned = */ bk_iters_per_partition};
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  // Use 1D grid with K-partition-major layout: [Partition0: M×N
+  // tiles][Partition1: M×N tiles]... Grid size is 1D to prevent driver/HW from
+  // using its own heuristic to exploit 2D locality by launching threadgroups in
+  // a non-linear order
+  MTL::Size grid_dims =
+      MTL::Size(tn_swizzled * tm_swizzled * split_k_partitions, 1, 1);
+
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_output_array(C_split, 2);
+
+  compute_encoder.set_bytes(params, 3);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  // Do accum kernel
+  {
+    const bool do_axpby = CHECK_AB && (alpha != 1.0f || beta != 0.0f);
+
+    auto kernel_name = "steel_gemm_splitk_accum_" + type_to_name(out) + "_" +
+        type_to_name(C_split);
+
+    if (do_axpby) {
+      kernel_name = kernel_name + "_axbpy";
+    }
+
+    auto kernel = get_steel_gemm_splitk_accum_kernel(
+        /* metal::Device& d = */ d,
+        /* const std::string& kernel_name = */ kernel_name,
+        /* const array& in = */ C_split,
+        /* const array& out = */ out,
+        /* bool axbpy = */ do_axpby);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    // Set the arguments for the kernel
+    compute_encoder.set_input_array(C_split, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(split_k_partitions, 2);
+    compute_encoder.set_bytes(split_k_partition_stride, 3);
+    compute_encoder.set_bytes(N, 4);
+
+    if (do_axpby) {
+      int ldc = c.strides()[c.ndim() - 2];
+      int fdc = c.strides()[c.ndim() - 1];
+
+      compute_encoder.set_input_array(c, 5);
+      compute_encoder.set_bytes(ldc, 6);
+      compute_encoder.set_bytes(fdc, 7);
+      compute_encoder.set_bytes(alpha, 8);
+      compute_encoder.set_bytes(beta, 9);
+    }
+
+    // Launch enough thread groups for each output
+    MTL::Size grid_dims = MTL::Size(N, M, 1);
+    auto group_dims = get_block_dims(N, M, 1);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Split matmul routing
+///////////////////////////////////////////////////////////////////////////////
+
+template <bool CHECK_AB>
+void steel_matmul_axpby(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape /* = {} */,
+    Strides A_batch_stride /* = {} */,
+    Strides B_batch_stride /* = {} */,
+    Strides C_batch_stride /* = {} */,
+    float alpha /* = 1.0f */,
+    float beta /* = 0.0f */) {
+  if (batch_shape.empty()) {
+    /////////////////////////////////////////////////////////////////////////////
+    // Check and collapse batch dimensions
+    if constexpr (CHECK_AB) {
+      auto [batch_shape_, A_bstride_, B_bstride_, C_bstride_] =
+          collapse_batches(a, b, c);
+
+      batch_shape = batch_shape_;
+      A_batch_stride = A_bstride_;
+      B_batch_stride = B_bstride_;
+      C_batch_stride = C_bstride_;
+      // Collapse batches into M if needed
+      if (batch_size_out > 1 && !transpose_a && batch_shape.size() == 1 &&
+          a.strides()[a.ndim() - 2] == K && A_batch_stride.back() == M * K &&
+          C_batch_stride.back() == M * c.strides()[c.ndim() - 2] &&
+          B_batch_stride.back() == 0) {
+        M *= batch_shape.back();
+        batch_size_out = 1;
+
+        A_batch_stride = {0};
+        B_batch_stride = {0};
+        C_batch_stride = {0};
+        batch_shape = {1};
+      }
+    } else {
+      auto [batch_shape_, A_bstride_, B_bstride_] = collapse_batches(a, b);
+
+      batch_shape = batch_shape_;
+      A_batch_stride = A_bstride_;
+      B_batch_stride = B_bstride_;
+      // Collapse batches into M if needed
+      if (batch_size_out > 1 && !transpose_a && batch_shape.size() == 1 &&
+          a.strides()[a.ndim() - 2] == K && A_batch_stride.back() == M * K &&
+          B_batch_stride.back() == 0) {
+        M *= batch_shape.back();
+        batch_size_out = 1;
+
+        A_batch_stride = {0};
+        B_batch_stride = {0};
+        batch_shape = {1};
+      }
+    }
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Split K specialization
+
+  int _tm = (M + 16 - 1) / 16;
+  int _tn = (N + 16 - 1) / 16;
+  int _tk = K / 16;
+
+  // Case 1: Small M×N with large K, use SIMD split-K
+  char devc = d.get_architecture().back();
+  // Max and Ultra dispatch larger sizes to splitk
+  int min_tmn_threshold = (devc == 's' || devc == 'd') ? 2048 : 1024;
+  if (batch_size_out == 1 && (_tm * _tn) <= min_tmn_threshold && _tk >= 8 &&
+      K >= std::max(M, N)) {
+    return steel_gemm_splitk_axpby<CHECK_AB>(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& a = */ a,
+        /* const array& b = */ b,
+        /* const array& c = */ c,
+        /* array& out = */ out,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* int batch_size_out = */ batch_size_out,
+        /* int lda = */ lda,
+        /* int ldb = */ ldb,
+        /* bool transpose_a = */ transpose_a,
+        /* bool transpose_b = */ transpose_b,
+        /* std::vector<array>& copies = */ copies,
+        /* float alpha = */ alpha,
+        /* float beta = */ beta);
+  }
+
+  // Case 2: Large K with sufficient M, N, and NAX is available, use NAX split-K
+  // TODO: Add device-specific tuning for more NAX GPUs in the future
+  constexpr int min_mn_threshold = 2048 * 2048;
+  constexpr int min_k_threshold = 10240;
+  if (batch_size_out == 1 && metal::is_nax_available() &&
+      !issubdtype(a.dtype(), complexfloating) &&
+      (env::enable_tf32() || a.dtype() != float32) &&
+      int64_t(M) * N >= min_mn_threshold && K >= min_k_threshold &&
+      K >= (3 * std::max(M, N))) {
+    return steel_gemm_splitk_axpby_nax<CHECK_AB>(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& a = */ a,
+        /* const array& b = */ b,
+        /* const array& c = */ c,
+        /* array& out = */ out,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* int batch_size_out = */ batch_size_out,
+        /* int lda = */ lda,
+        /* int ldb = */ ldb,
+        /* bool transpose_a = */ transpose_a,
+        /* bool transpose_b = */ transpose_b,
+        /* std::vector<array>& copies = */ copies,
+        /* float alpha = */ alpha,
+        /* float beta = */ beta);
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Regular kernel dispatch
+  auto batch_strides = A_batch_stride;
+  batch_strides.insert(
+      batch_strides.end(), B_batch_stride.begin(), B_batch_stride.end());
+  if (CHECK_AB && !C_batch_stride.empty()) {
+    batch_strides.insert(
+        batch_strides.end(), C_batch_stride.begin(), C_batch_stride.end());
+  }
+
+  int64_t A_batch_stride_ = A_batch_stride.empty() ? 0 : A_batch_stride.back();
+  int64_t B_batch_stride_ = B_batch_stride.empty() ? 0 : B_batch_stride.back();
+  int64_t C_batch_stride_ = C_batch_stride.empty() ? 0 : C_batch_stride.back();
+
+  return steel_matmul_regular_axpby<CHECK_AB>(
+      /* const Stream& s = */ s,
+      /* metal::Device& d = */ d,
+      /* const array& a = */ a,
+      /* const array& b = */ b,
+      /* const array& c = */ c,
+      /* array& out = */ out,
+      /* int M = */ M,
+      /* int N = */ N,
+      /* int K = */ K,
+      /* int batch_size_out = */ batch_size_out,
+      /* int lda = */ lda,
+      /* int ldb = */ ldb,
+      /* int ldd = */ N,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* std::vector<array>& copies = */ copies,
+      /* Shape batch_shape = */ std::move(batch_shape),
+      /* Strides batch_strides = */ std::move(batch_strides),
+      /* int64_t A_batch_stride = */ A_batch_stride_,
+      /* int64_t B_batch_stride = */ B_batch_stride_,
+      /* int64_t matrix_stride_out = */ int64_t(M) * N,
+      /* int64_t C_batch_stride = */ C_batch_stride_,
+      /* float alpha = */ alpha,
+      /* float beta = */ beta);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMV dispatch
+///////////////////////////////////////////////////////////////////////////////
+
+template <bool CHECK_AB = true>
+void gemv_axbpy(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape = {},
+    Strides A_batch_stride = {},
+    Strides B_batch_stride = {},
+    Strides C_batch_stride = {},
+    float alpha = 1.0f,
+    float beta = 0.0f) {
+  // Collect problem info
+  bool is_b_matrix = N != 1;
+
+  auto& mat = is_b_matrix ? b : a;
+  auto& vec = is_b_matrix ? a : b;
+  bool transpose_mat = is_b_matrix ? !transpose_b : transpose_a;
+  int in_vector_len = K;
+  int out_vector_len = is_b_matrix ? N : M;
+
+  int mat_ld = is_b_matrix ? ldb : lda;
+
+  auto batch_strides_mat = is_b_matrix ? B_batch_stride : A_batch_stride;
+  auto batch_strides_vec = is_b_matrix ? A_batch_stride : B_batch_stride;
+
+  // Determine if inputs have simple batching / broadcasting
+  bool contiguous_kernel = (batch_shape.size() == 1);
+
+  int batch_ndim = batch_shape.size();
+
+  // Determine dispatch kernel
+  int tm = 4, tn = 4;
+  int sm = 1, sn = 32;
+  int bm = 1, bn = 1;
+  int n_out_per_tgp;
+  std::ostringstream kname;
+
+  if (transpose_mat) {
+    if (in_vector_len >= 8192 && out_vector_len >= 2048) {
+      sm = 4;
+      sn = 8;
+    } else {
+      sm = 8;
+      sn = 4;
+    }
+
+    if (out_vector_len >= 2048) {
+      bn = 16;
+    } else if (out_vector_len >= 512) {
+      bn = 4;
+    } else {
+      bn = 2;
+    }
+
+    // Specialized kernel for very small outputs
+    tn = out_vector_len < tn ? 1 : tn;
+
+    n_out_per_tgp = bn * sn * tn;
+    kname << "gemv_t_" << type_to_name(out);
+
+  } else {
+    bm = out_vector_len >= 4096 ? 8 : 4;
+    sn = 32;
+
+    if (K <= 64) {
+      bm = 1;
+      sm = 8;
+      sn = 4;
+    } else if (K >= 16 * out_vector_len) {
+      bm = 1;
+      bn = 8;
+    }
+
+    // Specialized kernel for very small outputs
+    tm = out_vector_len < tm ? 1 : tm;
+
+    n_out_per_tgp = bm * sm * tm;
+    kname << "gemv_" << type_to_name(out);
+  }
+
+  const bool do_axpby = CHECK_AB && (alpha != 1.0f || beta != 0.0f);
+
+  // clang-format off
+  kname << "_bm" << bm << "_bn" << bn
+        << "_sm" << sm << "_sn" << sn
+        << "_tm" << tm << "_tn" << tn
+        << "_nc" << !contiguous_kernel
+        << "_axpby" << do_axpby; // clang-format on
+
+  // Encode and dispatch kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname.str());
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int n_tgp = (out_vector_len + n_out_per_tgp - 1) / n_out_per_tgp;
+  MTL::Size group_dims = MTL::Size(32, bn, bm);
+  MTL::Size grid_dims = MTL::Size(n_tgp, 1, batch_size_out);
+
+  compute_encoder.set_input_array(mat, 0);
+  compute_encoder.set_input_array(vec, 1);
+  compute_encoder.set_output_array(out, 3);
+
+  compute_encoder.set_bytes(in_vector_len, 4);
+  compute_encoder.set_bytes(out_vector_len, 5);
+  compute_encoder.set_bytes(mat_ld, 6);
+
+  compute_encoder.set_bytes(batch_ndim, 9);
+  compute_encoder.set_vector_bytes(batch_shape, 10);
+  compute_encoder.set_vector_bytes(batch_strides_vec, 11);
+  compute_encoder.set_vector_bytes(batch_strides_mat, 12);
+
+  if (do_axpby) {
+    compute_encoder.set_input_array(c, 2);
+
+    compute_encoder.set_bytes(alpha, 7);
+    compute_encoder.set_bytes(beta, 8);
+
+    compute_encoder.set_vector_bytes(C_batch_stride, 13);
+
+    int bias_stride = c.strides()[c.ndim() - 1];
+    compute_encoder.set_bytes(bias_stride, 14);
+  }
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+inline void gemv(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape = {},
+    Strides A_batch_stride = {},
+    Strides B_batch_stride = {}) {
+  return gemv_axbpy<false>(
+      /* const Stream& s = */ s,
+      /* metal::Device& d = */ d,
+      /* const array& a = */ a,
+      /* const array& b = */ b,
+      /* const array& c = */ b,
+      /* array& out = */ out,
+      /* int M = */ M,
+      /* int N = */ N,
+      /* int K = */ K,
+      /* int batch_size_out = */ batch_size_out,
+      /* int lda = */ lda,
+      /* int ldb = */ ldb,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* std::vector<array>& copies = */ copies,
+      /* Shape batch_shape = */ batch_shape,
+      /* Strides A_batch_stride = */ A_batch_stride,
+      /* Strides B_batch_stride = */ B_batch_stride);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Matmul implementation
+///////////////////////////////////////////////////////////////////////////////
+
+void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (!issubdtype(out.dtype(), inexact)) {
+    throw std::runtime_error("[matmul] dtype must be inexact.");
+  }
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+  // Return 0s if either input is empty
+  if (a_pre.size() == 0 || b_pre.size() == 0) {
+    array zero = array(0, a_pre.dtype());
+    fill_gpu(zero, out, s);
+    d.add_temporary(std::move(zero), s.index);
+    return;
+  }
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Init checks and prep
+
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int K = a_pre.shape(-1);
+
+  // Keep a vector with copies to be cleared in the completed buffer to release
+  // the arrays
+  std::vector<array> copies;
+  auto [a_transposed, a_cols, a] = check_transpose(copies, s, a_pre, M == 1);
+  auto [b_transposed, b_cols, b] = check_transpose(copies, s, b_pre, N == 1);
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Check and collapse batch dimensions
+
+  auto [batch_shape, A_batch_stride, B_batch_stride] = collapse_batches(a, b);
+
+  auto batch_size_out = out.size() / (size_t(M) * size_t(N));
+
+  // Collapse batches into M if needed
+  if (batch_size_out > 1 && !a_transposed && batch_shape.size() == 1 &&
+      a.strides()[a.ndim() - 2] == K && A_batch_stride.back() == M * K &&
+      B_batch_stride.back() == 0) {
+    M *= batch_shape.back();
+    batch_size_out = 1;
+
+    A_batch_stride = {0};
+    B_batch_stride = {0};
+    batch_shape = {1};
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Gemv specialization
+
+  // Route to gemv if needed
+  if (std::min(M, N) == 1) {
+    return gemv(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& a = */ a,
+        /* const array& b = */ b,
+        /* array& out = */ out,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* int batch_size_out = */ batch_size_out,
+        /* int lda = */ a_cols,
+        /* int ldb = */ b_cols,
+        /* bool transpose_a = */ a_transposed,
+        /* bool transpose_b = */ b_transposed,
+        /* std::vector<array>& copies = */ copies,
+        /* Shape batch_shape = */ std::move(batch_shape),
+        /* Strides A_batch_stride = */ std::move(A_batch_stride),
+        /* Strides B_batch_stride = */ std::move(B_batch_stride));
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Gemm specialization
+
+  return steel_matmul(
+      /* const Stream& s = */ s,
+      /* metal::Device& d = */ d,
+      /* const array& a = */ a,
+      /* const array& b = */ b,
+      /* array& out = */ out,
+      /* int M = */ M,
+      /* int N = */ N,
+      /* int K = */ K,
+      /* int batch_size_out = */ batch_size_out,
+      /* int lda = */ a_cols,
+      /* int ldb = */ b_cols,
+      /* bool transpose_a = */ a_transposed,
+      /* bool transpose_b = */ b_transposed,
+      /* std::vector<array>& copies = */ copies,
+      /* Shape batch_shape = */ std::move(batch_shape),
+      /* Strides A_batch_stride = */ std::move(A_batch_stride),
+      /* Strides B_batch_stride = */ std::move(B_batch_stride));
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// AddMM implementation
+///////////////////////////////////////////////////////////////////////////////
+
+void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 3);
+  if (!issubdtype(out.dtype(), floating)) {
+    throw std::runtime_error(
+        "[matmul] Does not yet support non-floating point types.");
+  }
+
+  // Return 0s if either input is empty
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(out.nbytes()));
+    return;
+  }
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  // Handle empty matrix case (K=0)
+  if (inputs[0].shape(-1) == 0) {
+    auto& c = inputs[2];
+    if (beta_ == 1.0f) {
+      copy_gpu(
+          c,
+          out,
+          c.flags().row_contiguous ? CopyType::Vector : CopyType::General,
+          s);
+    } else {
+      array beta_scalar = array(beta_, c.dtype());
+      binary_op_gpu({c, beta_scalar}, out, "Multiply", s);
+      d.add_temporary(std::move(beta_scalar), s.index);
+    }
+    return;
+  }
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+  auto& c_pre = inputs[2];
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Init checks and prep
+
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int K = a_pre.shape(-1);
+
+  // Keep a vector with copies to be cleared in the completed buffer to release
+  // the arrays
+  std::vector<array> copies;
+  auto [transpose_a, a_cols, a] = check_transpose(copies, s, a_pre, M == 1);
+  auto [transpose_b, b_cols, b] = check_transpose(copies, s, b_pre, N == 1);
+
+  array c = c_pre;
+
+  int lda = a_cols;
+  int ldb = b_cols;
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Check and collapse batch dimensions
+  auto [batch_shape, A_batch_stride, B_batch_stride, C_batch_stride] =
+      collapse_batches(a, b, c);
+
+  int64_t matrix_stride_out = M * static_cast<int64_t>(N);
+  auto batch_size_out = out.size() / (matrix_stride_out);
+
+  // Collapse batches into M if needed
+  if (batch_size_out > 1 && !transpose_a && batch_shape.size() == 1 &&
+      a.strides()[a.ndim() - 2] == K && A_batch_stride.back() == M * K &&
+      C_batch_stride.back() == M * c.strides()[c.ndim() - 2] &&
+      B_batch_stride.back() == 0) {
+    M *= batch_shape.back();
+    batch_size_out = 1;
+
+    A_batch_stride = {0};
+    B_batch_stride = {0};
+    C_batch_stride = {0};
+    batch_shape = {1};
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Gemv specialization
+
+  // Route to gemv if needed
+  if (std::min(M, N) == 1) {
+    return gemv_axbpy(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& a = */ a,
+        /* const array& b = */ b,
+        /* const array& c = */ c,
+        /* array& out = */ out,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* int batch_size_out = */ batch_size_out,
+        /* int lda = */ lda,
+        /* int ldb = */ ldb,
+        /* bool transpose_a = */ transpose_a,
+        /* bool transpose_b = */ transpose_b,
+        /* std::vector<array>& copies = */ copies,
+        /* Shape batch_shape = */ batch_shape,
+        /* Strides A_batch_stride = */ A_batch_stride,
+        /* Strides B_batch_stride = */ B_batch_stride,
+        /* Strides C_batch_stride = */ C_batch_stride,
+        /* float alpha = */ alpha_,
+        /* float beta = */ beta_);
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Regular addmm dispatch
+
+  return steel_matmul_axpby(
+      /* const Stream& s = */ s,
+      /* metal::Device& d = */ d,
+      /* const array& a = */ a,
+      /* const array& b = */ b,
+      /* const array& c = */ c,
+      /* array& out = */ out,
+      /* int M = */ M,
+      /* int N = */ N,
+      /* int K = */ K,
+      /* int batch_size_out = */ batch_size_out,
+      /* int lda = */ lda,
+      /* int ldb = */ ldb,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* std::vector<array>& copies = */ copies,
+      /* Shape batch_shape = */ batch_shape,
+      /* Strides A_batch_stride = */ A_batch_stride,
+      /* Strides B_batch_stride = */ B_batch_stride,
+      /* Strides B_batch_stride = */ C_batch_stride,
+      /* float alpha = */ alpha_,
+      /* float beta = */ beta_);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// BlockMaskedMM implementation
+///////////////////////////////////////////////////////////////////////////////
+
+void BlockMaskedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  using namespace mlx::steel;
+  // assert(inputs.size() == 2);
+  if (!issubdtype(out.dtype(), floating)) {
+    throw std::runtime_error(
+        "[matmul] Does not yet support non-floating point types.");
+  }
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+  // Return 0s if either input is empty
+  if (a_pre.size() == 0 || b_pre.size() == 0) {
+    array zero = array(0, a_pre.dtype());
+    fill_gpu(zero, out, s);
+    d.add_temporary(std::move(zero), s.index);
+    return;
+  }
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Init checks and prep
+
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int K = a_pre.shape(-1);
+
+  // Keep a vector with copies to be cleared in the completed buffer to release
+  // the arrays
+  std::vector<array> copies;
+  auto [transpose_a, a_cols, a] = check_transpose(copies, s, a_pre, M == 1);
+  auto [transpose_b, b_cols, b] = check_transpose(copies, s, b_pre, N == 1);
+
+  int lda = a_cols;
+  int ldb = b_cols;
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Check and collapse batch dimensions
+
+  bool has_op_mask = inputs.size() > 3;
+  bool has_out_mask = inputs.size() == 3 || inputs.size() == 5;
+
+  // Prepare kernel name
+  std::string out_mask_nm = has_out_mask ? type_to_name(inputs[2]) : "nomask";
+  std::string op_mask_nm = has_op_mask ? type_to_name(inputs.back()) : "nomask";
+
+  Shape batch_shape{1};
+  Strides A_batch_stride{0};
+  Strides B_batch_stride{0};
+  Strides outmask_bstride{0};
+  Strides Amask_bstride{0};
+  Strides Bmask_bstride{0};
+  int64_t A_batch_str = 0;
+  int64_t B_batch_str = 0;
+
+  Strides batch_strides;
+
+  if (out.ndim() > 2) {
+    Shape bshape{out.shape().begin(), out.shape().end() - 2};
+    std::vector<Strides> bstrides;
+
+    for (auto& arr : inputs) {
+      bstrides.emplace_back(arr.strides().begin(), arr.strides().end() - 2);
+    }
+
+    // auto [bshape_c, bstrides_c] = collapse_contiguous_dims(bshape, bstrides);
+    batch_shape = bshape;
+    A_batch_str = bstrides[0].back();
+    B_batch_str = bstrides[1].back();
+
+    for (auto& bstr : bstrides) {
+      batch_strides.insert(batch_strides.end(), bstr.begin(), bstr.end());
+    }
+
+    A_batch_stride = bstrides[0];
+    B_batch_stride = bstrides[1];
+
+    if (has_out_mask) {
+      outmask_bstride = bstrides[2];
+    }
+    if (has_op_mask) {
+      Amask_bstride = bstrides[has_out_mask + 2];
+      Bmask_bstride = bstrides[has_out_mask + 3];
+    }
+
+  } else {
+    batch_strides = Strides(inputs.size(), 0);
+  }
+
+  int64_t matrix_stride_out = static_cast<int64_t>(M) * N;
+  size_t batch_size_out = out.size() / (matrix_stride_out);
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Gemv specialization
+
+  // Route to gemv if needed
+  if (std::min(M, N) == 1) {
+    // Collect problem info
+    bool is_b_matrix = N != 1;
+
+    auto& mat = is_b_matrix ? b : a;
+    auto& vec = is_b_matrix ? a : b;
+    bool transpose_mat = is_b_matrix ? !transpose_b : transpose_a;
+    int in_vector_len = K;
+    int out_vector_len = is_b_matrix ? N : M;
+
+    int mat_ld = is_b_matrix ? b_cols : a_cols;
+
+    auto batch_strides_mat = is_b_matrix ? B_batch_stride : A_batch_stride;
+    auto batch_strides_vec = is_b_matrix ? A_batch_stride : B_batch_stride;
+
+    auto mask_bstrides_mat = is_b_matrix ? Bmask_bstride : Amask_bstride;
+    auto mask_bstrides_vec = is_b_matrix ? Amask_bstride : Bmask_bstride;
+
+    auto mat_mask_idx = int(has_out_mask) + (is_b_matrix ? 3 : 2);
+    auto vec_mask_idx = int(has_out_mask) + (is_b_matrix ? 2 : 3);
+
+    // Determine if inputs have simple batching / broadcasting
+    bool contiguous_kernel = (batch_shape.size() == 1);
+
+    int batch_ndim = batch_shape.size();
+
+    // Determine dispatch kernel
+    int tm = 4, tn = 4;
+    int sm = 1, sn = 32;
+    int bm = 1, bn = 1;
+    int n_out_per_tgp;
+    std::ostringstream kname;
+
+    if (transpose_mat) {
+      sm = 8;
+      sn = 4;
+      bm = 1;
+      bn = (block_size_ == 64 && out_vector_len >= 2048) ? 4 : 2;
+      tm = block_size_ == 32 ? 4 : 8;
+      tn = 4;
+
+      // Specialized kernel for very small outputs
+      tn = out_vector_len < tn ? 1 : tn;
+
+      n_out_per_tgp = bn * sn * tn;
+      kname << "gemv_t";
+
+    } else {
+      if (block_size_ == 32) {
+        sm = 4;
+        sn = 8;
+        bm = 2;
+      } else {
+        sm = 2;
+        sn = 16;
+        bm = out_vector_len >= 512 ? 4 : 2;
+      }
+
+      // Specialized kernel for very small outputs
+      tm = out_vector_len < tm ? 1 : tm;
+
+      n_out_per_tgp = bm * sm * tm;
+      kname << "gemv";
+    }
+
+    kname << "_outmask_" << out_mask_nm;
+    kname << "_opmask_" << op_mask_nm;
+    kname << "_" << type_to_name(out);
+    kname << "_bm" << bm << "_bn" << bn;
+    kname << "_sm" << sm << "_sn" << sn;
+    kname << "_tm" << tm << "_tn" << tn;
+    kname << "_nc" << !contiguous_kernel;
+
+    // Encode and dispatch kernel
+    auto kernel = get_gemv_masked_kernel(
+        d,
+        kname.str(),
+        out,
+        has_out_mask ? std::optional<array>{inputs[2]} : std::nullopt,
+        has_op_mask ? std::optional<array>{inputs.back()} : std::nullopt,
+        transpose_mat,
+        bm,
+        bn,
+        sm,
+        sn,
+        tm,
+        tn,
+        contiguous_kernel);
+
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    int n_tgp = (out_vector_len + n_out_per_tgp - 1) / n_out_per_tgp;
+    MTL::Size group_dims = MTL::Size(32, bn, bm);
+    MTL::Size grid_dims = MTL::Size(n_tgp, 1, batch_size_out);
+
+    // Get mask params
+    std::vector<int> mask_strides;
+    Strides mask_batch_strides;
+    if (has_out_mask) {
+      auto& out_mask = inputs[2];
+
+      if (transpose_mat) {
+        mask_strides.push_back(out_mask.strides(out.shape(-2) == 1 ? -1 : -2));
+        mask_strides.push_back(out_mask.strides(out.shape(-2) == 1 ? -2 : -1));
+      } else {
+        mask_strides.push_back(out_mask.strides(out.shape(-1) == 1 ? -1 : -2));
+        mask_strides.push_back(out_mask.strides(out.shape(-1) == 1 ? -2 : -1));
+      }
+
+      mask_batch_strides.insert(
+          mask_batch_strides.end(),
+          outmask_bstride.begin(),
+          outmask_bstride.end());
+
+      compute_encoder.set_input_array(out_mask, 20);
+    }
+
+    if (has_op_mask) {
+      auto& mat_mask = inputs[mat_mask_idx];
+
+      if (transpose_mat) {
+        mask_strides.push_back(mat_mask.strides(!is_b_matrix ? -2 : -1));
+        mask_strides.push_back(mat_mask.strides(!is_b_matrix ? -1 : -2));
+      } else {
+        mask_strides.push_back(mat_mask.strides(is_b_matrix ? -2 : -1));
+        mask_strides.push_back(mat_mask.strides(is_b_matrix ? -1 : -2));
+      }
+
+      mask_batch_strides.insert(
+          mask_batch_strides.end(),
+          mask_bstrides_mat.begin(),
+          mask_bstrides_mat.end());
+
+      compute_encoder.set_input_array(mat_mask, 21);
+
+      auto& vec_mask = inputs[vec_mask_idx];
+      if (transpose_mat) {
+        mask_strides.push_back(vec_mask.strides(vec.shape(-2) == 1 ? -1 : -2));
+        mask_strides.push_back(vec_mask.strides(vec.shape(-2) == 1 ? -2 : -1));
+      } else {
+        mask_strides.push_back(vec_mask.strides(vec.shape(-1) == 1 ? -1 : -2));
+        mask_strides.push_back(vec_mask.strides(vec.shape(-1) == 1 ? -2 : -1));
+      }
+
+      mask_batch_strides.insert(
+          mask_batch_strides.end(),
+          mask_bstrides_vec.begin(),
+          mask_bstrides_vec.end());
+
+      compute_encoder.set_input_array(vec_mask, 22);
+    }
+
+    // Get gemv params
+    compute_encoder.set_input_array(mat, 0);
+    compute_encoder.set_input_array(vec, 1);
+    compute_encoder.set_output_array(out, 3);
+
+    compute_encoder.set_bytes(in_vector_len, 4);
+    compute_encoder.set_bytes(out_vector_len, 5);
+    compute_encoder.set_bytes(mat_ld, 6);
+    compute_encoder.set_bytes(batch_ndim, 9);
+    compute_encoder.set_vector_bytes(batch_shape, 10);
+    compute_encoder.set_vector_bytes(batch_strides_vec, 11);
+    compute_encoder.set_vector_bytes(batch_strides_mat, 12);
+
+    compute_encoder.set_vector_bytes(mask_strides, 23);
+    compute_encoder.set_vector_bytes(mask_batch_strides, 24);
+
+    compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+    d.add_temporaries(std::move(copies), s.index);
+    return;
+  }
+
+  /////////////////////////////////////////////////////////////////////////////
+  // Regular kernel dispatch
+
+  // Determine dispatch kernel
+  int bm = block_size_, bn = block_size_, bk = 16;
+  int wm = 2, wn = 2;
+  bool mn_aligned = M % bm == 0 && N % bn == 0;
+  bool k_aligned = K % bk == 0;
+
+  std::ostringstream kname;
+  kname << "steel_gemm_block_outmask_" << out_mask_nm << "_opmask_"
+        << op_mask_nm << "_" << (transpose_a ? 't' : 'n')
+        << (transpose_b ? 't' : 'n') << "_" << type_to_name(a) << "_"
+        << type_to_name(out) << "_bm" << bm << "_bn" << bn << "_bk" << bk
+        << "_wm" << wm << "_wn" << wn << "_MN_" << (mn_aligned ? "t" : "n")
+        << "aligned"
+        << "_K_" << (k_aligned ? "t" : "n") << "aligned";
+
+  // Encode and dispatch kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_masked_kernel(
+      d,
+      kname.str(),
+      out,
+      has_out_mask ? std::optional<array>{inputs[2]} : std::nullopt,
+      has_op_mask ? std::optional<array>{inputs.back()} : std::nullopt,
+      transpose_a,
+      transpose_b,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn,
+      mn_aligned,
+      k_aligned);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Use problem size to determine threadblock swizzle
+  int tn = (N + bn - 1) / bn;
+  int tm = (M + bm - 1) / bm;
+
+  // TODO: Explore device-based tuning for swizzle
+  int swizzle_log = 0; // tm >= 6 ? 3 : (tm <= 3 ? 0 : 2);
+
+  // Prepare steel matmul params
+  GEMMParams params{/* const int M = */ M,
+                    /* const int N = */ N,
+                    /* const int K = */ K,
+                    /* const int lda = */ lda,
+                    /* const int ldb = */ ldb,
+                    /* const int ldd = */ N,
+                    /* const int tiles_n = */ tn,
+                    /* const int tiles_m = */ tm,
+                    /* const int64_t batch_stride_a = */ A_batch_str,
+                    /* const int64_t batch_stride_b = */ B_batch_str,
+                    /* const int64_t batch_stride_d = */ matrix_stride_out,
+                    /* const int swizzle_log = */ swizzle_log,
+                    /* const int gemm_k_iterations_aligned = */ (K / bk),
+                    /* const int batch_ndim = */ int(batch_shape.size())};
+
+  // Prepare launch grid params
+  int tile = 1 << swizzle_log;
+  tm = (tm + tile - 1) / tile;
+  tn = tn * tile;
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(tn, tm, batch_size_out);
+
+  std::vector<int> mask_strides;
+
+  if (has_out_mask) {
+    auto& out_mask = inputs[2];
+    mask_strides.push_back(*(out_mask.strides().end() - 1));
+    mask_strides.push_back(*(out_mask.strides().end() - 2));
+
+    compute_encoder.set_input_array(out_mask, 10);
+  }
+
+  if (has_op_mask) {
+    auto& lhs_mask = inputs[2 + has_out_mask];
+    mask_strides.push_back(*(lhs_mask.strides().end() - 1));
+    mask_strides.push_back(*(lhs_mask.strides().end() - 2));
+
+    compute_encoder.set_input_array(lhs_mask, 11);
+
+    auto& rhs_mask = inputs[3 + has_out_mask];
+    mask_strides.push_back(*(rhs_mask.strides().end() - 1));
+    mask_strides.push_back(*(rhs_mask.strides().end() - 2));
+
+    compute_encoder.set_input_array(rhs_mask, 12);
+  }
+
+  // Launch kernel
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_output_array(out, 3);
+
+  compute_encoder.set_bytes(params, 4);
+
+  compute_encoder.set_vector_bytes(batch_shape, 6);
+  compute_encoder.set_vector_bytes(batch_strides, 7);
+
+  compute_encoder.set_vector_bytes(mask_strides, 13);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// GatherMM implementation
+///////////////////////////////////////////////////////////////////////////////
+
+void gather_mm_rhs(
+    const array& a_,
+    const array& b_,
+    const array& indices_,
+    array& out,
+    metal::Device& d,
+    const Stream& s) {
+  array indices = ensure_row_contiguous(indices_, d, s);
+  auto [transpose_b, ldb, b] = ensure_batch_contiguous(b_, d, s);
+
+  // Broadcast a with indices. If we are here that means lhs_indices were not
+  // provided so the lhs_indices are implied to be the shape of a broadcasted
+  // with rhs_indices. We need only broadcast a and copy it as if applying the
+  // lhs_indices.
+  auto broadcast_with_indices = [&d, &s, &indices](const array& x) {
+    if (x.size() / x.shape(-2) / x.shape(-1) == indices.size()) {
+      return ensure_row_contiguous(x, d, s);
+    }
+
+    auto x_shape = indices.shape();
+    x_shape.push_back(x.shape(-2));
+    x_shape.push_back(x.shape(-1));
+    array new_x(std::move(x_shape), x.dtype(), nullptr, {});
+    broadcast(x, new_x);
+    return ensure_row_contiguous(new_x, d, s);
+  };
+  array a = broadcast_with_indices(a_);
+
+  // Extract the matmul shapes
+  int K = a.shape(-1);
+  int M = a.size() / K;
+  int N = b.shape(-1);
+  int lda = a.strides()[a.ndim() - 2]; // should be K
+
+  // Define the dispatch blocks
+  int bm = 16, bn = 64, bk = 16;
+  int wm = 1, wn = 2;
+
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  // Define the kernel name
+  std::string base_name;
+  base_name.reserve(64);
+  concatenate(
+      base_name,
+      "steel_gather_mm_rhs_n",
+      transpose_b ? 't' : 'n',
+      '_',
+      type_to_name(a),
+      '_',
+      type_to_name(out),
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn);
+
+  metal::MTLFCList func_consts = {
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+      {&align_K, MTL::DataType::DataTypeBool, 202},
+  };
+
+  // And the kernel hash that includes the function constants
+  std::string hash_name;
+  hash_name.reserve(128);
+  concatenate(
+      hash_name,
+      base_name,
+      "_align_M_",
+      align_M ? 't' : 'n',
+      "_align_N_",
+      align_N ? 't' : 'n',
+      "_align_K_",
+      align_K ? 't' : 'n');
+
+  // Get and set the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_gather_kernel(
+      d,
+      base_name,
+      hash_name,
+      func_consts,
+      out,
+      false,
+      transpose_b,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn,
+      true);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Prepare the matmul params
+  auto batch_stride_b = b.ndim() > 2 ? b.strides()[b.ndim() - 3] : b.size();
+  steel::GEMMParams params{
+      /* const int M = */ M,
+      /* const int N = */ N,
+      /* const int K = */ K,
+      /* const int lda = */ lda,
+      /* const int ldb = */ static_cast<int>(ldb),
+      /* const int ldd = */ N,
+      /* const int tiles_n = */ (N + bn - 1) / bn,
+      /* const int tiles_m = */ (M + bm - 1) / bm,
+      /* const int64_t batch_stride_a = */ 0,
+      /* const int64_t batch_stride_b = */ static_cast<int64_t>(batch_stride_b),
+      /* const int64_t batch_stride_d = */ 0,
+      /* const int swizzle_log = */ 0,
+      /* const int gemm_k_iterations_aligned = */ (K / bk),
+      /* const int batch_ndim = */ 0};
+
+  // Prepare the grid
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(params.tiles_n, params.tiles_m, 1);
+
+  // Launch kernel
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_input_array(indices, 2);
+  compute_encoder.set_output_array(out, 3);
+  compute_encoder.set_bytes(params, 4);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_mm_rhs_nax(
+    const array& a_,
+    const array& b_,
+    const array& indices_,
+    array& out,
+    metal::Device& d,
+    const Stream& s) {
+  array indices = ensure_row_contiguous(indices_, d, s);
+  auto [transpose_b, ldb, b] = ensure_batch_contiguous(b_, d, s);
+
+  // Broadcast a with indices. If we are here that means lhs_indices were not
+  // provided so the lhs_indices are implied to be the shape of a broadcasted
+  // with rhs_indices. We need only broadcast a and copy it as if applying the
+  // lhs_indices.
+  auto broadcast_with_indices = [&d, &s, &indices](const array& x) {
+    if (x.size() / x.shape(-2) / x.shape(-1) == indices.size()) {
+      return ensure_row_contiguous(x, d, s);
+    }
+
+    auto x_shape = indices.shape();
+    x_shape.push_back(x.shape(-2));
+    x_shape.push_back(x.shape(-1));
+    array new_x(std::move(x_shape), x.dtype(), nullptr, {});
+    broadcast(x, new_x);
+    return ensure_row_contiguous(new_x, d, s);
+  };
+  array a = broadcast_with_indices(a_);
+
+  // Extract the matmul shapes
+  int K = a.shape(-1);
+  int M = a.size() / K;
+  int N = b.shape(-1);
+  int lda = a.strides()[a.ndim() - 2]; // should be K
+  int E = b.shape(0);
+
+  // Define the dispatch blocks
+  int bm, bn = 128, bk = 128, wm, wn = 4;
+  if (M / E > 48) {
+    bm = 64;
+    wm = 2;
+  } else if (M / E > 24) {
+    bm = 32l;
+    wm = 1;
+  } else {
+    bm = 16;
+    wm = 1;
+  }
+
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  // Define the kernel name
+  std::string base_name;
+  base_name.reserve(64);
+  concatenate(
+      base_name,
+      "steel_gather_mm_rhs_nax_n",
+      transpose_b ? 't' : 'n',
+      '_',
+      type_to_name(a),
+      '_',
+      type_to_name(out),
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn);
+
+  metal::MTLFCList func_consts = {
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+      {&align_K, MTL::DataType::DataTypeBool, 202},
+  };
+
+  // And the kernel hash that includes the function constants
+  std::string hash_name;
+  hash_name.reserve(128);
+  concatenate(
+      hash_name,
+      base_name,
+      "_align_M_",
+      align_M ? 't' : 'n',
+      "_align_N_",
+      align_N ? 't' : 'n',
+      "_align_K_",
+      align_K ? 't' : 'n');
+
+  // Get and set the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_gather_nax_kernel(
+      d,
+      base_name,
+      hash_name,
+      func_consts,
+      out,
+      false,
+      transpose_b,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn,
+      true);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Prepare the matmul params
+  auto batch_stride_b = b.ndim() > 2 ? b.strides()[b.ndim() - 3] : b.size();
+  steel::GEMMParams params{
+      /* const int M = */ M,
+      /* const int N = */ N,
+      /* const int K = */ K,
+      /* const int lda = */ lda,
+      /* const int ldb = */ static_cast<int>(ldb),
+      /* const int ldd = */ N,
+      /* const int tiles_n = */ (N + bn - 1) / bn,
+      /* const int tiles_m = */ (M + bm - 1) / bm,
+      /* const int64_t batch_stride_a = */ 0,
+      /* const int64_t batch_stride_b = */ static_cast<int64_t>(batch_stride_b),
+      /* const int64_t batch_stride_d = */ 0,
+      /* const int swizzle_log = */ 0,
+      /* const int gemm_k_iterations_aligned = */ (K / bk),
+      /* const int batch_ndim = */ 0};
+
+  // Prepare the grid
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(params.tiles_n, params.tiles_m, 1);
+
+  // Launch kernel
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_input_array(indices, 2);
+  compute_encoder.set_output_array(out, 3);
+  compute_encoder.set_bytes(params, 4);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_mv(
+    const array& mat_,
+    const array& vec_,
+    const array& mat_indices_,
+    const array& vec_indices_,
+    array& out,
+    int N,
+    int K,
+    bool is_mv,
+    metal::Device& d,
+    const Stream& s) {
+  // Copy if needed
+  std::vector<array> copies;
+  auto [transpose_mat, mat_cols, mat] =
+      check_transpose(copies, s, mat_, N == 1);
+  auto [transpose_vec, vec_cols, vec] = check_transpose(copies, s, vec_, true);
+  d.add_temporaries(std::move(copies), s.index);
+
+  // If we are doing vector matrix instead of matrix vector we need to flip the
+  // matrix transposition. Basically m @ v = v @ m.T assuming that v is treated
+  // as a one dimensional array.
+  transpose_mat = (!is_mv) ^ transpose_mat;
+
+  // Define some shapes
+  int in_vector_len = K;
+  int out_vector_len = N;
+  int mat_ld = mat_cols;
+
+  int batch_size_out = out.size() / N;
+  int batch_ndim = out.ndim() - 2;
+  int batch_ndim_mat = mat.ndim() - 2;
+  int batch_ndim_vec = vec.ndim() - 2;
+  Strides index_strides = vec_indices_.strides();
+  index_strides.insert(
+      index_strides.end(),
+      mat_indices_.strides().begin(),
+      mat_indices_.strides().end());
+
+  // Determine dispatch kernel
+  int tm = 4, tn = 4;
+  int sm = 1, sn = 32;
+  int bm = 1, bn = 1;
+  int n_out_per_tgp;
+  std::ostringstream kname;
+
+  if (transpose_mat) {
+    if (in_vector_len >= 8192 && out_vector_len >= 2048) {
+      sm = 4;
+      sn = 8;
+    } else {
+      sm = 8;
+      sn = 4;
+    }
+
+    if (out_vector_len >= 2048) {
+      bn = 16;
+    } else if (out_vector_len >= 512) {
+      bn = 4;
+    } else {
+      bn = 2;
+    }
+
+    // Specialized kernel for very small outputs
+    tn = out_vector_len < tn ? 1 : tn;
+
+    n_out_per_tgp = bn * sn * tn;
+    kname << "gemv_t_gather_" << type_to_name(out);
+
+  } else {
+    bm = out_vector_len >= 4096 ? 8 : 4;
+    sn = 32;
+
+    // Specialized kernel for very small outputs
+    tm = out_vector_len < tm ? 1 : tm;
+
+    n_out_per_tgp = bm * sm * tm;
+    kname << "gemv_gather_" << type_to_name(out);
+  }
+
+  kname << "_bm" << bm << "_bn" << bn << "_sm" << sm << "_sn" << sn << "_tm"
+        << tm << "_tn" << tn;
+
+  // Encode and dispatch kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname.str());
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int n_tgp = (out_vector_len + n_out_per_tgp - 1) / n_out_per_tgp;
+  MTL::Size group_dims = MTL::Size(32, bn, bm);
+  MTL::Size grid_dims = MTL::Size(n_tgp, 1, batch_size_out);
+
+  compute_encoder.set_input_array(mat, 0);
+  compute_encoder.set_input_array(vec, 1);
+  compute_encoder.set_output_array(out, 3);
+
+  compute_encoder.set_bytes(in_vector_len, 4);
+  compute_encoder.set_bytes(out_vector_len, 5);
+  compute_encoder.set_bytes(mat_ld, 6);
+
+  compute_encoder.set_bytes(batch_ndim, 9);
+  compute_encoder.set_vector_bytes(out.shape(), 10);
+  compute_encoder.set_vector_bytes(index_strides, 11);
+
+  compute_encoder.set_bytes(batch_ndim_vec, 12);
+  compute_encoder.set_vector_bytes(vec.shape(), 13);
+  compute_encoder.set_vector_bytes(vec.strides(), 14);
+
+  compute_encoder.set_bytes(batch_ndim_mat, 15);
+  compute_encoder.set_vector_bytes(mat.shape(), 16);
+  compute_encoder.set_vector_bytes(mat.strides(), 17);
+
+  compute_encoder.set_input_array(vec_indices_, 18);
+  compute_encoder.set_input_array(mat_indices_, 19);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_mm(
+    const array& a_,
+    const array& b_,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s) {
+  // Copy if needed
+  std::vector<array> copies;
+  auto [transpose_a, lda, a] = check_transpose(copies, s, a_, false);
+  auto [transpose_b, ldb, b] = check_transpose(copies, s, b_, false);
+  d.add_temporaries(std::move(copies), s.index);
+
+  // Determine dispatch kernel
+  int bm = 64, bn = 64, bk = 16;
+  int wm = 2, wn = 2;
+  size_t batch_size_out = out.size() / M / N;
+  int batch_ndim = out.ndim() - 2;
+  int batch_ndim_a = a.ndim() - 2;
+  int batch_ndim_b = b.ndim() - 2;
+
+  char devc = d.get_architecture().back();
+  GEMM_TPARAM_MACRO(devc)
+
+  const bool has_batch = batch_ndim > 1;
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  // Define the kernel name
+  std::string base_name;
+  base_name.reserve(128);
+  concatenate(
+      base_name,
+      "steel_gather_mm_",
+      transpose_a ? 't' : 'n',
+      transpose_b ? 't' : 'n',
+      "_",
+      type_to_name(a),
+      "_",
+      type_to_name(out),
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn);
+
+  metal::MTLFCList func_consts = {
+      {&has_batch, MTL::DataType::DataTypeBool, 10},
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+      {&align_K, MTL::DataType::DataTypeBool, 202},
+  };
+
+  // And the kernel hash that includes the function constants
+  std::string hash_name;
+  hash_name.reserve(128);
+  concatenate(
+      hash_name,
+      base_name,
+      "_has_batch_",
+      has_batch ? 't' : 'n',
+      "_align_M_",
+      align_M ? 't' : 'n',
+      "_align_N_",
+      align_N ? 't' : 'n',
+      "_align_K_",
+      align_K ? 't' : 'n');
+
+  // Get and set the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_gather_kernel(
+      d,
+      base_name,
+      hash_name,
+      func_consts,
+      out,
+      transpose_a,
+      transpose_b,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn,
+      false);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Prepare the matmul params
+  steel::GEMMParams params{/* const int M = */ M,
+                           /* const int N = */ N,
+                           /* const int K = */ K,
+                           /* const int lda = */ static_cast<int>(lda),
+                           /* const int ldb = */ static_cast<int>(ldb),
+                           /* const int ldd = */ N,
+                           /* const int tiles_n = */ (N + bn - 1) / bn,
+                           /* const int tiles_m = */ (M + bm - 1) / bm,
+                           /* const int64_t batch_stride_a = */
+                           (batch_ndim > 0) ? lhs_indices.strides()[0] : 0,
+                           /* const int64_t batch_stride_b = */
+                           (batch_ndim > 0) ? rhs_indices.strides()[0] : 0,
+                           /* const int64_t batch_stride_d = */ M * N,
+                           /* const int swizzle_log = */ 0,
+                           /* const int gemm_k_iterations_aligned = */ (K / bk),
+                           /* const int batch_ndim = */ batch_ndim};
+
+  // Prepare the grid
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims =
+      MTL::Size(params.tiles_n, params.tiles_m, batch_size_out);
+
+  // Launch kernel
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_input_array(lhs_indices, 2);
+  compute_encoder.set_input_array(rhs_indices, 3);
+  compute_encoder.set_output_array(out, 4);
+  compute_encoder.set_bytes(params, 5);
+  compute_encoder.set_vector_bytes(lhs_indices.shape(), 6);
+  compute_encoder.set_vector_bytes(lhs_indices.strides(), 7);
+  compute_encoder.set_vector_bytes(rhs_indices.strides(), 8);
+  compute_encoder.set_bytes(batch_ndim_a, 9);
+  compute_encoder.set_vector_bytes(a.shape(), 10);
+  compute_encoder.set_vector_bytes(a.strides(), 11);
+  compute_encoder.set_bytes(batch_ndim_b, 12);
+  compute_encoder.set_vector_bytes(b.shape(), 13);
+  compute_encoder.set_vector_bytes(b.strides(), 14);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void GatherMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& lhs_indices = inputs[2];
+  auto& rhs_indices = inputs[3];
+
+  // Return 0s if either input is empty
+  if (a.size() == 0 || b.size() == 0) {
+    array zero = array(0, a.dtype());
+    fill_gpu(zero, out, s);
+    d.add_temporary(std::move(zero), s.index);
+    return;
+  }
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  // Extract shapes from inputs.
+  int M = a.shape(-2);
+  int N = b.shape(-1);
+  int K = a.shape(-1);
+
+  // We are walking a in order and b is also in order so we can batch up the
+  // matmuls and reuse reading a and b.
+  if (M == 1 && right_sorted_ == true) {
+    if (metal::is_nax_available() &&
+        (env::enable_tf32() || a.dtype() != float32)) {
+      return gather_mm_rhs_nax(a, b, rhs_indices, out, d, s);
+    }
+    gather_mm_rhs(a, b, rhs_indices, out, d, s);
+    return;
+  }
+
+  // Route to gather gemv if any of a or b are vectors
+  if (M == 1) {
+    gather_mv(b, a, rhs_indices, lhs_indices, out, N, K, false, d, s);
+    return;
+  }
+  if (N == 1) {
+    gather_mv(a, b, lhs_indices, rhs_indices, out, M, K, true, d, s);
+    return;
+  }
+
+  // Route to non specialized gather mm
+  gather_mm(a, b, lhs_indices, rhs_indices, out, M, N, K, d, s);
+}
+
+void segmented_mm(
+    const array& a_,
+    const array& b_,
+    const array& segments_,
+    array& out,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s) {
+  auto check_segments_layout = [&d, &s](const array& x) {
+    // Contiguous so return early
+    if (x.flags().row_contiguous) {
+      return std::make_tuple(true, x);
+    }
+
+    bool rc = true;
+    for (int i = 0; i < x.ndim() - 2; i++) {
+      rc &=
+          (x.strides(i + 1) * x.shape(i) == x.strides(i)) || (x.shape(i) == 1);
+    }
+    rc &= x.strides(x.ndim() - 1) == 1;
+    if (x.ndim() > 1) {
+      rc &= x.strides(x.ndim() - 2) == 1;
+    }
+
+    if (rc) {
+      return std::make_tuple(false, x);
+    }
+
+    array x_copy = contiguous_copy_gpu(x, s);
+    d.add_temporary(x_copy, s.index);
+    return std::make_tuple(true, x_copy);
+  };
+
+  // Copy if needed
+  std::vector<array> copies;
+  auto [transpose_a, lda, a] = check_transpose(copies, s, a_, false);
+  auto [transpose_b, ldb, b] = check_transpose(copies, s, b_, false);
+  auto [segments_contiguous, segments] = check_segments_layout(segments_);
+  d.add_temporaries(std::move(copies), s.index);
+
+  // Determine dispatch kernel
+  int bm = 64, bn = 64, bk = 16;
+  int wm = 2, wn = 2;
+  size_t batch_size_out = out.size() / M / N;
+
+  char devc = d.get_architecture().back();
+  GEMM_TPARAM_MACRO(devc)
+
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+
+  // Define the kernel name
+  std::string base_name;
+  base_name.reserve(128);
+  concatenate(
+      base_name,
+      "steel_segmented_mm_",
+      transpose_a ? 't' : 'n',
+      transpose_b ? 't' : 'n',
+      "_",
+      type_to_name(a),
+      "_",
+      type_to_name(out),
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn);
+
+  metal::MTLFCList func_consts = {
+      {&segments_contiguous, MTL::DataType::DataTypeBool, 199},
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+  };
+
+  // And the kernel hash that includes the function constants
+  std::string hash_name;
+  hash_name.reserve(128);
+  concatenate(
+      hash_name,
+      base_name,
+      "_segments_contiguous_",
+      segments_contiguous ? 't' : 'n',
+      "_align_M_",
+      align_M ? 't' : 'n',
+      "_align_N_",
+      align_N ? 't' : 'n');
+
+  // Get and set the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_segmented_kernel(
+      d,
+      base_name,
+      hash_name,
+      func_consts,
+      out,
+      transpose_a,
+      transpose_b,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Prepare the matmul params
+  steel::GEMMParams params{/* const int M = */ M,
+                           /* const int N = */ N,
+                           /* const int K = */ K,
+                           /* const int lda = */ static_cast<int>(lda),
+                           /* const int ldb = */ static_cast<int>(ldb),
+                           /* const int ldd = */ N,
+                           /* const int tiles_n = */ (N + bn - 1) / bn,
+                           /* const int tiles_m = */ (M + bm - 1) / bm,
+                           /* const int64_t batch_stride_a = */ 0,
+                           /* const int64_t batch_stride_b = */ 0,
+                           /* const int64_t batch_stride_d = */ M * N,
+                           /* const int swizzle_log = */ 0,
+                           /* const int gemm_k_iterations_aligned = */ 0,
+                           /* const int batch_ndim = */ 0};
+
+  // Prepare the grid
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims =
+      MTL::Size(params.tiles_n, params.tiles_m, batch_size_out);
+
+  // Launch kernel
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_input_array(segments, 2);
+  compute_encoder.set_output_array(out, 3);
+  compute_encoder.set_bytes(params, 4);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void SegmentedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& segments = inputs[2];
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  // Extract shapes from inputs.
+  int M = a.shape(-2);
+  int N = b.shape(-1);
+  int K = a.shape(-1);
+
+  segmented_mm(a, b, segments, out, M, N, K, d, s);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/matmul.h b/macos/third_party/mlx/mlx/backend/metal/matmul.h
new file mode 100644
index 0000000..218664b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/matmul.h
@@ -0,0 +1,144 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/device.h"
+
+namespace mlx::core {
+
+template <bool CHECK_AB = true>
+void steel_matmul_regular_axpby(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    int ldd,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape,
+    Strides batch_strides,
+    int64_t A_batch_stride,
+    int64_t B_batch_stride,
+    int64_t matrix_stride_out,
+    int64_t C_batch_stride = 0,
+    float alpha = 1.0f,
+    float beta = 0.0f);
+
+inline void steel_matmul_regular(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    int ldd,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape,
+    Strides batch_strides,
+    int64_t A_batch_stride,
+    int64_t B_batch_stride,
+    int64_t matrix_stride_out) {
+  return steel_matmul_regular_axpby<false>(
+      /* const Stream& s = */ s,
+      /* metal::Device& d = */ d,
+      /* const array& a = */ a,
+      /* const array& b = */ b,
+      /* const array& c = */ b,
+      /* array& out = */ out,
+      /* int M = */ M,
+      /* int N = */ N,
+      /* int K = */ K,
+      /* int batch_size_out = */ batch_size_out,
+      /* int lda = */ lda,
+      /* int ldb = */ ldb,
+      /* int ldd = */ ldd,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* std::vector<array>& copies = */ copies,
+      /* Shape batch_shape = */ batch_shape,
+      /* Strides batch_strides = */ batch_strides,
+      /* int64_t A_batch_stride = */ A_batch_stride,
+      /* int64_t B_batch_stride = */ B_batch_stride,
+      /* int64_t matrix_stride_out = */ matrix_stride_out);
+}
+
+template <bool CHECK_AB = true>
+void steel_matmul_axpby(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape = {},
+    Strides A_batch_stride = {},
+    Strides B_batch_stride = {},
+    Strides C_batch_stride = {},
+    float alpha = 1.0f,
+    float beta = 0.0f);
+
+inline void steel_matmul(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    Shape batch_shape = {},
+    Strides A_batch_stride = {},
+    Strides B_batch_stride = {}) {
+  return steel_matmul_axpby<false>(
+      /* const Stream& s = */ s,
+      /* metal::Device& d = */ d,
+      /* const array& a = */ a,
+      /* const array& b = */ b,
+      /* const array& c = */ b,
+      /* array& out = */ out,
+      /* int M = */ M,
+      /* int N = */ N,
+      /* int K = */ K,
+      /* int batch_size_out = */ batch_size_out,
+      /* int lda = */ lda,
+      /* int ldb = */ ldb,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* std::vector<array>& copies = */ copies,
+      /* Shape batch_shape = */ batch_shape,
+      /* Strides A_batch_stride = */ A_batch_stride,
+      /* Strides B_batch_stride = */ B_batch_stride);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/metal.cpp b/macos/third_party/mlx/mlx/backend/metal/metal.cpp
new file mode 100644
index 0000000..51bd2e6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/metal.cpp
@@ -0,0 +1,50 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <memory>
+
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/backend/metal/utils.h"
+
+namespace mlx::core::metal {
+
+bool is_available() {
+  return true;
+}
+
+void start_capture(std::string path, NS::Object* object) {
+  auto pool = new_scoped_memory_pool();
+
+  auto descriptor = MTL::CaptureDescriptor::alloc()->init();
+  descriptor->setCaptureObject(object);
+
+  if (!path.empty()) {
+    auto string = NS::String::string(path.c_str(), NS::UTF8StringEncoding);
+    auto url = NS::URL::fileURLWithPath(string);
+    descriptor->setDestination(MTL::CaptureDestinationGPUTraceDocument);
+    descriptor->setOutputURL(url);
+  }
+
+  auto manager = MTL::CaptureManager::sharedCaptureManager();
+  NS::Error* error;
+  bool started = manager->startCapture(descriptor, &error);
+  descriptor->release();
+  if (!started) {
+    std::ostringstream msg;
+    msg << "[metal::start_capture] Failed to start: "
+        << error->localizedDescription()->utf8String();
+    throw std::runtime_error(msg.str());
+  }
+}
+
+void start_capture(std::string path) {
+  auto& device = metal::device(mlx::core::Device::gpu);
+  return start_capture(path, device.mtl_device());
+}
+
+void stop_capture() {
+  auto pool = new_scoped_memory_pool();
+  auto manager = MTL::CaptureManager::sharedCaptureManager();
+  manager->stopCapture();
+}
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/metal.h b/macos/third_party/mlx/mlx/backend/metal/metal.h
new file mode 100644
index 0000000..6662e21
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/metal.h
@@ -0,0 +1,25 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+
+namespace mlx::core::metal {
+
+/* Check if the Metal backend is available. */
+MLX_API bool is_available();
+
+/** Capture a GPU trace, saving it to an absolute file `path` */
+MLX_API void start_capture(std::string path = "");
+MLX_API void stop_capture();
+
+/** Get information about the GPU and system settings. */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info();
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/no_metal.cpp b/macos/third_party/mlx/mlx/backend/metal/no_metal.cpp
new file mode 100644
index 0000000..9d6a734
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/no_metal.cpp
@@ -0,0 +1,42 @@
+// Copyright © 2025 Apple Inc.
+
+#include <stdexcept>
+
+#include "mlx/backend/metal/metal.h"
+#include "mlx/fast.h"
+
+namespace mlx::core {
+
+namespace metal {
+
+bool is_available() {
+  return false;
+}
+
+void start_capture(std::string) {}
+void stop_capture() {}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info() {
+  throw std::runtime_error(
+      "[metal::device_info] Cannot get device info without metal backend");
+};
+
+} // namespace metal
+
+namespace fast {
+
+CustomKernelFunction metal_kernel(
+    const std::string&,
+    const std::vector<std::string>&,
+    const std::vector<std::string>&,
+    const std::string&,
+    const std::string&,
+    bool,
+    bool) {
+  throw std::runtime_error("[metal_kernel] No Metal back-end.");
+}
+
+} // namespace fast
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/nojit_kernels.cpp b/macos/third_party/mlx/mlx/backend/metal/nojit_kernels.cpp
new file mode 100644
index 0000000..533b192
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/nojit_kernels.cpp
@@ -0,0 +1,414 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+MTL::ComputePipelineState* get_arange_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_unary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype,
+    Dtype,
+    const char*) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_binary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype,
+    Dtype,
+    const char*) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_binary_two_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype,
+    Dtype,
+    const char*) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_ternary_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    Dtype,
+    const char*) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const array&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_dynamic_copy_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const array&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_softmax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool,
+    const array&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_logsumexp_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_scan_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    bool,
+    bool,
+    const std::string&,
+    const array&,
+    const array&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const array&,
+    int,
+    int) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_mb_sort_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const array&,
+    int,
+    int) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_reduce_init_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string&,
+    const std::string&,
+    const Dtype&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_reduce_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string&,
+    const std::string&,
+    const Dtype&,
+    const Dtype&,
+    const std::string&,
+    int,
+    int,
+    int) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_fused_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool,
+    bool) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_accum_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const array&,
+    bool) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool,
+    bool) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_gather_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_segmented_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_gemv_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    const std::optional<array>&,
+    const std::optional<array>&,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_steel_conv_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array&,
+    int,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_steel_conv_general_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    int,
+    int,
+    int,
+    int,
+    int) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_fft_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const std::string&) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_quantized_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string&,
+    const std::string&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_gather_qmm_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    int,
+    int,
+    const std::string&,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_fused_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string&,
+    const std::string&) {
+  return d.get_kernel(kernel_name);
+}
+
+MTL::ComputePipelineState* get_gather_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    int,
+    int,
+    const std::string&,
+    int,
+    int,
+    int,
+    int,
+    int,
+    bool) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_attention_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    int,
+    int,
+    int,
+    int,
+    int,
+    const array&) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_attention_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    int,
+    int,
+    int,
+    int,
+    int,
+    const array&) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/normalization.cpp b/macos/third_party/mlx/mlx/backend/metal/normalization.cpp
new file mode 100644
index 0000000..da0160b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/normalization.cpp
@@ -0,0 +1,435 @@
+// Copyright © 2024 Apple Inc.
+#include <algorithm>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/reduce.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+
+namespace mlx::core::fast {
+
+bool RMSNorm::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+void RMSNorm::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& out = outputs[0];
+
+  // Make sure that the last dimension is contiguous
+  auto set_output = [&s, &out](const array& x) {
+    bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
+    if (no_copy && x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back() || x.shape(-2) == 1);
+    }
+    if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  const array x = set_output(inputs[0]);
+  const array& w = inputs[1];
+
+  auto axis_size = static_cast<uint32_t>(x.shape().back());
+  int n_rows = x.data_size() / axis_size;
+
+  const int simd_size = 32;
+  const int n_reads = RMS_N_READS;
+  const int looped_limit = RMS_LOOPED_LIMIT;
+  std::string op_name = "rms";
+  if (axis_size > looped_limit) {
+    op_name += "_looped";
+  }
+  op_name += type_to_name(out);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  {
+    auto kernel = d.get_kernel(op_name);
+
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(
+        x.data_shared_ptr() == nullptr ? out : x, 0);
+    compute_encoder.set_input_array(w, 1);
+    compute_encoder.set_output_array(out, 2);
+    compute_encoder.set_bytes(eps_, 3);
+    compute_encoder.set_bytes(axis_size, 4);
+    compute_encoder.set_bytes(w_stride, 5);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void RMSNormVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  // Ensure row contiguity. We could relax this step by checking that the array
+  // is contiguous (no broadcasts or holes) and that the input strides are the
+  // same as the cotangent strides but for now this is simpler.
+  auto check_input = [&s](const array& x) -> std::pair<array, bool> {
+    if (x.flags().row_contiguous) {
+      return {x, false};
+    }
+    array x_copy = contiguous_copy_gpu(x, s);
+    return {x_copy, true};
+  };
+  bool donate_g = inputs[2].is_donatable();
+  auto [x, copied] = check_input(inputs[0]);
+  const array& w = inputs[1];
+  auto [g, g_copied] = check_input(inputs[2]);
+  donate_g |= g_copied;
+  array& gx = outputs[0];
+  array& gw = outputs[1];
+
+  // Check whether we had a weight
+  bool has_w = w.ndim() != 0;
+
+  // Allocate space for the outputs
+  bool g_in_gx = false;
+  if (x.is_donatable()) {
+    gx.copy_shared_buffer(x);
+  } else if (g.is_donatable()) {
+    gx.copy_shared_buffer(g);
+    g_in_gx = true;
+  } else {
+    gx.set_data(allocator::malloc(gx.nbytes()));
+  }
+  if (g_copied && !g_in_gx) {
+    d.add_temporary(g, s.index);
+  }
+
+  auto axis_size = static_cast<uint32_t>(x.shape().back());
+  int n_rows = x.data_size() / axis_size;
+
+  // Allocate the gradient accumulator gw and a temporary to store the
+  // gradients before they are accumulated.
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
+  if (has_w) {
+    if (!g_in_gx && donate_g) {
+      gw_temp.copy_shared_buffer(g);
+    } else {
+      gw_temp.set_data(allocator::malloc(gw_temp.nbytes()));
+      d.add_temporary(gw_temp, s.index);
+    }
+  }
+  gw.set_data(allocator::malloc(gw.nbytes()));
+
+  const int simd_size = 32;
+  const int n_reads = RMS_N_READS;
+  const int looped_limit = RMS_LOOPED_LIMIT;
+  std::string op_name = "vjp_rms";
+  if (axis_size > looped_limit) {
+    op_name += "_looped";
+  }
+  op_name += type_to_name(gx);
+
+  std::string hash_name = op_name + ((has_w) ? "_w" : "_now");
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+  };
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  {
+    auto kernel = d.get_kernel(op_name, hash_name, func_consts);
+
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(x, 0);
+    compute_encoder.set_input_array(w, 1);
+    compute_encoder.set_input_array(g, 2);
+    compute_encoder.set_output_array(gx, 3);
+    compute_encoder.set_output_array(gw_temp, 4);
+    compute_encoder.set_bytes(eps_, 5);
+    compute_encoder.set_bytes(axis_size, 6);
+    compute_encoder.set_bytes(w_stride, 7);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+
+  if (has_w) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    strided_reduce_general_dispatch(
+        gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
+  }
+}
+
+bool LayerNorm::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+void LayerNorm::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& out = outputs[0];
+
+  // Make sure that the last dimension is contiguous
+  auto set_output = [&s, &out](const array& x) {
+    bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1;
+    if (no_copy && x.ndim() > 1) {
+      auto s = x.strides()[x.ndim() - 2];
+      no_copy &= (s == 0 || s == x.shape().back() || x.shape(-2) == 1);
+    }
+    if (no_copy) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  const array x = set_output(inputs[0]);
+  const array& w = inputs[1];
+  const array& b = inputs[2];
+
+  auto axis_size = static_cast<uint32_t>(x.shape().back());
+  int n_rows = x.data_size() / axis_size;
+
+  int simd_size = 32;
+  int n_reads = 8;
+  int looped_limit = 6656;
+  std::string op_name = "layer_norm";
+  if (axis_size > looped_limit) {
+    op_name += "_looped";
+    n_reads = 4;
+  }
+  op_name += type_to_name(out);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  {
+    auto kernel = d.get_kernel(op_name);
+
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      if (threadgroup_size > kernel->maxTotalThreadsPerThreadgroup()) {
+        std::ostringstream msg;
+        msg << "[layer_norm] Threadgroup size " << threadgroup_size
+            << " is larger than the maximum allowed threadgroup size "
+            << kernel->maxTotalThreadsPerThreadgroup();
+        throw std::runtime_error(msg.str());
+      }
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+    uint32_t b_stride = (b.ndim() == 1) ? b.strides()[0] : 0;
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(
+        x.data_shared_ptr() == nullptr ? out : x, 0);
+    compute_encoder.set_input_array(w, 1);
+    compute_encoder.set_input_array(b, 2);
+    compute_encoder.set_output_array(out, 3);
+    compute_encoder.set_bytes(eps_, 4);
+    compute_encoder.set_bytes(axis_size, 5);
+    compute_encoder.set_bytes(w_stride, 6);
+    compute_encoder.set_bytes(b_stride, 7);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void LayerNormVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  // Ensure row contiguity. We could relax this step by checking that the array
+  // is contiguous (no broadcasts or holes) and that the input strides are the
+  // same as the cotangent strides but for now this is simpler.
+  auto check_input = [&s](const array& x) -> std::pair<array, bool> {
+    if (x.flags().row_contiguous) {
+      return {x, false};
+    }
+    array x_copy = contiguous_copy_gpu(x, s);
+    return {x_copy, true};
+  };
+  bool donate_x = inputs[0].is_donatable();
+  bool donate_g = inputs[3].is_donatable();
+  auto [x, copied] = check_input(inputs[0]);
+  donate_x |= copied;
+  const array& w = inputs[1];
+  auto [g, g_copied] = check_input(inputs[3]);
+  donate_g |= g_copied;
+  array& gx = outputs[0];
+  array& gw = outputs[1];
+  array& gb = outputs[2];
+
+  // Check whether we had a weight
+  bool has_w = w.ndim() != 0;
+
+  // Allocate space for the outputs
+  bool g_in_gx = false;
+  if (donate_x) {
+    gx.copy_shared_buffer(x);
+  } else if (donate_g) {
+    gx.copy_shared_buffer(g);
+    g_in_gx = true;
+  } else {
+    gx.set_data(allocator::malloc(gx.nbytes()));
+  }
+  if (g_copied && !g_in_gx) {
+    d.add_temporary(g, s.index);
+  }
+
+  auto axis_size = static_cast<uint32_t>(x.shape().back());
+  int n_rows = x.data_size() / axis_size;
+
+  // Allocate a temporary to store the gradients for w and allocate the output
+  // gradient accumulators.
+  array gw_temp =
+      (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
+  if (has_w) {
+    if (!g_in_gx && donate_g) {
+      gw_temp.copy_shared_buffer(g);
+    } else {
+      gw_temp.set_data(allocator::malloc(gw_temp.nbytes()));
+      d.add_temporary(gw_temp, s.index);
+    }
+  }
+  gw.set_data(allocator::malloc(gw.nbytes()));
+  gb.set_data(allocator::malloc(gb.nbytes()));
+
+  // Finish with the gradient for b in case we had a b
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  if (gb.ndim() == 1 && gb.size() == axis_size) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    strided_reduce_general_dispatch(
+        g, gb, "sum", plan, {0}, compute_encoder, d, s);
+  }
+
+  int simd_size = 32;
+  int n_reads = 8;
+  int looped_limit = 8192;
+  std::string op_name = "vjp_layer_norm";
+  if (axis_size > looped_limit) {
+    op_name += "_looped";
+    n_reads = 4;
+  }
+  op_name += type_to_name(gx);
+
+  std::string hash_name = op_name + ((has_w) ? "_w" : "_now");
+  metal::MTLFCList func_consts = {
+      {&has_w, MTL::DataType::DataTypeBool, 20},
+  };
+
+  {
+    auto kernel = d.get_kernel(op_name, hash_name, func_consts);
+
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      if (threadgroup_size > kernel->maxTotalThreadsPerThreadgroup()) {
+        std::ostringstream msg;
+        msg << "[vjp_layer_norm] Threadgroup size " << threadgroup_size
+            << " is larger than the maximum allowed threadgroup size "
+            << kernel->maxTotalThreadsPerThreadgroup();
+        throw std::runtime_error(msg.str());
+      }
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0;
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(x, 0);
+    compute_encoder.set_input_array(w, 1);
+    compute_encoder.set_input_array(g, 2);
+    compute_encoder.set_output_array(gx, 3);
+    compute_encoder.set_output_array(gw_temp, 4);
+    compute_encoder.set_bytes(eps_, 5);
+    compute_encoder.set_bytes(axis_size, 6);
+    compute_encoder.set_bytes(w_stride, 7);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+
+  if (has_w) {
+    ReductionPlan plan(
+        ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
+    strided_reduce_general_dispatch(
+        gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
+  }
+}
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/backend/metal/primitives.cpp b/macos/third_party/mlx/mlx/backend/metal/primitives.cpp
new file mode 100644
index 0000000..5f6376c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/primitives.cpp
@@ -0,0 +1,242 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <algorithm>
+#include <cassert>
+#include <numeric>
+#include <sstream>
+
+#include "mlx/backend/common/slicing.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/gpu/slicing.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+template <typename T>
+void arange_set_scalars(T start, T next, metal::CommandEncoder& enc) {
+  enc.set_bytes(start, 0);
+  T step = next - start;
+  enc.set_bytes(step, 1);
+}
+
+void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 0);
+  out.set_data(allocator::malloc(out.nbytes()));
+  if (out.size() == 0) {
+    return;
+  }
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto kernel = get_arange_kernel(d, "arange" + type_to_name(out), out);
+  size_t nthreads = out.size();
+  MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
+  MTL::Size group_dims = MTL::Size(
+      std::min(nthreads, kernel->maxTotalThreadsPerThreadgroup()), 1, 1);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  switch (out.dtype()) {
+    case bool_: // unsupported
+      throw std::runtime_error("[Arange::eval_gpu] Does not support bool");
+    case uint8:
+      arange_set_scalars<uint8_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case uint16:
+      arange_set_scalars<uint16_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case uint32:
+      arange_set_scalars<uint32_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case uint64:
+      arange_set_scalars<uint64_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case int8:
+      arange_set_scalars<int8_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case int16:
+      arange_set_scalars<int16_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case int32:
+      arange_set_scalars<int32_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case int64:
+      arange_set_scalars<int64_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case float16:
+      arange_set_scalars<float16_t>(start_, start_ + step_, compute_encoder);
+      break;
+    case float32:
+      arange_set_scalars<float>(start_, start_ + step_, compute_encoder);
+      break;
+    case bfloat16:
+      arange_set_scalars<bfloat16_t>(start_, start_ + step_, compute_encoder);
+      break;
+    default:
+      throw std::runtime_error("[Arange::eval_gpu] Does not support type.");
+  }
+
+  compute_encoder.set_output_array(out, 2);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  std::string op_name;
+  switch (reduce_type_) {
+    case ArgReduce::ArgMin:
+      op_name = "argmin_";
+      break;
+    case ArgReduce::ArgMax:
+      op_name = "argmax_";
+      break;
+  }
+
+  // Prepare the shapes, strides and axis arguments.
+  auto in_strides = in.strides();
+  auto shape = in.shape();
+  auto out_strides = out.strides();
+  auto axis_stride = in_strides[axis_];
+  size_t axis_size = shape[axis_];
+  if (out_strides.size() == in_strides.size()) {
+    out_strides.erase(out_strides.begin() + axis_);
+  }
+  in_strides.erase(in_strides.begin() + axis_);
+  shape.erase(shape.begin() + axis_);
+  size_t ndim = shape.size();
+
+  // ArgReduce
+  int simd_size = 32;
+  int n_reads = 4;
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  {
+    auto kernel = d.get_kernel(op_name + type_to_name(in));
+    NS::UInteger thread_group_size = std::min(
+        (axis_size + n_reads - 1) / n_reads,
+        kernel->maxTotalThreadsPerThreadgroup());
+    // round up to the closest number divisible by simd_size
+    thread_group_size =
+        (thread_group_size + simd_size - 1) / simd_size * simd_size;
+    assert(thread_group_size <= kernel->maxTotalThreadsPerThreadgroup());
+
+    auto gd = get_2d_grid_dims(out.shape(), out.strides());
+    MTL::Size grid_dims = MTL::Size(thread_group_size, gd.width, gd.height);
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(out, 1);
+    if (ndim == 0) {
+      // Pass place holders so metal doesn't complain
+      int shape_ = 0;
+      int64_t stride_ = 0;
+      compute_encoder.set_bytes(shape_, 2);
+      compute_encoder.set_bytes(stride_, 3);
+      compute_encoder.set_bytes(stride_, 4);
+    } else {
+      compute_encoder.set_vector_bytes(shape, 2);
+      compute_encoder.set_vector_bytes(in_strides, 3);
+      compute_encoder.set_vector_bytes(out_strides, 4);
+    }
+    compute_encoder.set_bytes(ndim, 5);
+    compute_encoder.set_bytes(axis_stride, 6);
+    compute_encoder.set_bytes(axis_size, 7);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void Load::eval_gpu(const std::vector<array>& inputs, array& out) {
+  throw std::runtime_error("[Load::eval_gpu] Not implemented.");
+}
+
+void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  // keys has shape (N1, ..., NK, 2)
+  // out has shape (N1, ..., NK, M1, M2, ...)
+  auto& keys = inputs[0];
+  size_t num_keys = keys.size() / 2;
+
+  size_t elems_per_key = out.size() / num_keys;
+  size_t bytes_per_key = out.itemsize() * elems_per_key;
+  out.set_data(allocator::malloc(out.nbytes()));
+  if (out.size() == 0) {
+    return;
+  }
+
+  size_t out_per_key = (bytes_per_key + 4 - 1) / 4;
+  size_t half_size = out_per_key / 2;
+  bool odd = out_per_key % 2;
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  std::string kname = keys.flags().row_contiguous ? "rbitsc" : "rbits";
+  auto kernel = d.get_kernel(kname);
+
+  // organize into grid nkeys x elem_per_key
+  MTL::Size grid_dims = MTL::Size(num_keys, half_size + odd, 1);
+  auto group_dims = get_block_dims(num_keys, half_size + odd, 1);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(keys, 0);
+  compute_encoder.set_output_array(out, 1);
+  compute_encoder.set_bytes(odd, 2);
+  compute_encoder.set_bytes(bytes_per_key, 3);
+
+  if (!keys.flags().row_contiguous) {
+    int ndim = keys.ndim();
+    compute_encoder.set_bytes(ndim, 4);
+    compute_encoder.set_vector_bytes(keys.shape(), 5);
+    compute_encoder.set_vector_bytes(keys.strides(), 6);
+  }
+
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void QRF::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("[QRF::eval_gpu] Metal QR factorization NYI.");
+}
+
+void SVD::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("[SVD::eval_gpu] Metal SVD NYI.");
+}
+
+void Inverse::eval_gpu(const std::vector<array>& inputs, array& output) {
+  throw std::runtime_error("[Inverse::eval_gpu] Metal inversion NYI.");
+}
+
+void Cholesky::eval_gpu(const std::vector<array>& inputs, array& out) {
+  throw std::runtime_error(
+      "[Cholesky::eval_gpu] Metal Cholesky decomposition NYI.");
+}
+
+void Eig::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("[Eig::eval_gpu] Metal Eig NYI.");
+}
+
+void Eigh::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("[Eigh::eval_gpu] Metal Eigh NYI.");
+}
+
+void LUF::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI.");
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/quantized.cpp b/macos/third_party/mlx/mlx/backend/metal/quantized.cpp
new file mode 100644
index 0000000..cb67c74
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/quantized.cpp
@@ -0,0 +1,1651 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/common/broadcasting.h"
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/reduce.h"
+#include "mlx/backend/metal/unary.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename... Args>
+auto get_quantized_kernel_wrapped(
+    metal::Device& d,
+    const std::string& name,
+    const std::string& func,
+    const std::string& mode,
+    const std::string& type,
+    int group_size,
+    int bits,
+    Args... args) {
+  std::string template_def;
+  std::string fname = ((mode == "affine") ? "affine_" : "fp_") + func;
+  template_def = get_template_definition(
+      name, fname, type, group_size, bits, std::forward<Args>(args)...);
+  return get_quantized_kernel(d, name, template_def, mode);
+}
+
+template <typename... Args>
+auto get_qmm_nax_kernel_wrapped(
+    metal::Device& d,
+    const std::string& name,
+    const std::string& func,
+    const std::string& mode,
+    const std::string& type,
+    int group_size,
+    int bits,
+    Args... args) {
+  std::string template_def;
+  std::string fname = ((mode == "affine") ? "affine_" : "fp_") + func;
+  template_def = get_template_definition(
+      name, fname, type, group_size, bits, std::forward<Args>(args)...);
+  return get_qmm_nax_kernel(d, name, template_def, mode);
+}
+
+inline array
+ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    d.add_temporary(x_copy, s.index);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+inline array ensure_row_contiguous_matrix(
+    const array& x,
+    metal::Device& d,
+    const Stream& s) {
+  if (x.ndim() < 2) {
+    if (x.strides()[0] == 1) {
+      return x;
+    }
+  } else {
+    auto stride_0 = x.strides()[x.ndim() - 2];
+    auto stride_1 = x.strides()[x.ndim() - 1];
+    if (stride_0 == x.shape(-1) && stride_1 == 1) {
+      return x;
+    }
+  }
+  array x_copy = contiguous_copy_gpu(x, s);
+  d.add_temporary(x_copy, s.index);
+  return x_copy;
+}
+
+inline int get_qmv_batch_limit(int D, int O, metal::Device& d) {
+  auto arch = d.get_architecture();
+  auto arch_size = arch.back();
+  auto arch_gen = arch.substr(arch.size() - 3, 2);
+  if (arch_gen == "13" || arch_gen == "14") {
+    switch (arch_size) {
+      case 'd':
+        if (D <= 2048 && O <= 2048) {
+          return 32;
+        } else if (D <= 4096 && O <= 4096) {
+          return 18;
+        } else {
+          return 12;
+        }
+      default:
+        if (D <= 2048 && O <= 2048) {
+          return 14;
+        } else if (D <= 4096 && O <= 4096) {
+          return 10;
+        } else {
+          return 6;
+        }
+    }
+  } else {
+    switch (arch_size) {
+      case 'd':
+        if (D <= 2048 && O <= 2048) {
+          return 32;
+        } else if (D <= 4096 && O <= 4096) {
+          return 18;
+        } else {
+          return 12;
+        }
+      default:
+        if (D <= 2048 && O <= 2048) {
+          return 18;
+        } else if (D <= 4096 && O <= 4096) {
+          return 12;
+        } else {
+          return 10;
+        }
+    }
+  }
+}
+
+inline int add_strides_and_shapes(
+    CommandEncoder& compute_encoder,
+    bool skip,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    int offset) {
+  if (skip) {
+    return 0;
+  }
+
+  // TODO: Collapse batch dimensions
+
+  int x_batch_ndims = x.ndim() - 2;
+  int w_batch_ndims = w.ndim() - 2;
+  compute_encoder.set_bytes(x_batch_ndims, offset++);
+  compute_encoder.set_vector_bytes(x.shape(), offset++);
+  compute_encoder.set_vector_bytes(x.strides(), offset++);
+  compute_encoder.set_bytes(w_batch_ndims, offset++);
+  compute_encoder.set_vector_bytes(w.shape(), offset++);
+  compute_encoder.set_vector_bytes(w.strides(), offset++);
+  compute_encoder.set_vector_bytes(scales.strides(), offset++);
+  if (biases) {
+    compute_encoder.set_vector_bytes(biases->strides(), offset++);
+  }
+
+  return offset;
+}
+
+inline int add_gather_strides_and_shapes(
+    CommandEncoder& compute_encoder,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int offset) {
+  auto [shape, strides] = collapse_contiguous_dims(
+      lhs_indices.shape(), {lhs_indices.strides(), rhs_indices.strides()});
+  int ndims = shape.size();
+
+  compute_encoder.set_bytes(ndims, offset++);
+  compute_encoder.set_vector_bytes(shape, offset++);
+  compute_encoder.set_vector_bytes(strides[0], offset++);
+  compute_encoder.set_vector_bytes(strides[1], offset++);
+
+  return offset;
+}
+
+} // namespace
+
+void qmv_quad(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int B = out.size() / M / N;
+
+  constexpr int quads_per_simd = 8;
+  constexpr int results_per_quadgroup = 8;
+  int bn = quads_per_simd * results_per_quadgroup;
+  int simdgroup_size = 32;
+  MTL::Size group_dims(simdgroup_size, 1, 1);
+  MTL::Size grid_dims(M, (N + bn - 1) / bn, B);
+
+  std::string kname;
+  kname.reserve(64);
+  std::string type_string = get_type_string(x.dtype());
+
+  concatenate(
+      kname,
+      mode + "_qmv_quad_",
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      "_d_",
+      K,
+      B > 1 ? "_batch_1" : "_batch_0");
+  auto kernel = get_quantized_kernel_wrapped(
+      d, kname, "qmv_quad", mode, type_string, group_size, bits, K, B > 1);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c++);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int B = out.size() / M / N;
+
+  int bn = 8;
+  int bk = 32;
+  MTL::Size group_dims(bk, 2, 1);
+  MTL::Size grid_dims(M, (N + bn - 1) / bn, B);
+
+  std::string kname;
+  kname.reserve(64);
+  std::string type_string = get_type_string(x.dtype());
+  bool fast = N % bn == 0 && K % 512 == 0;
+
+  concatenate(
+      kname,
+      mode + (fast ? "_qmv_fast_" : "_qmv_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      B > 1 ? "_batch_1" : "_batch_0");
+  auto kernel = get_quantized_kernel_wrapped(
+      d,
+      kname,
+      (fast ? "qmv_fast" : "qmv"),
+      mode,
+      type_string,
+      group_size,
+      bits,
+      B > 1);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void qvm_split_k(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int split_k = K > 8192 ? 32 : 8;
+  int split_D = (K + split_k - 1) / split_k;
+  int B = out.size() / M / N;
+  B *= split_k;
+
+  constexpr int num_simdgroups = 2;
+  constexpr int bk = 32;
+  int bn = std::min(group_size, 32) * num_simdgroups;
+  MTL::Size group_dims = MTL::Size(bk, num_simdgroups, 1);
+  MTL::Size grid_dims = MTL::Size(M, N / bn, B);
+
+  auto x_shape = x.shape();
+  auto x_strides = x.strides();
+  if (x_shape.size() == 1) {
+    x_shape.insert(x_shape.begin(), 1);
+    x_strides.insert(x_strides.begin(), 0);
+  }
+
+  int x_ndim = x_shape.size();
+  int x_batch_ndims = x_ndim - 2;
+  int w_batch_ndims = w.ndim() - 2;
+  auto w_shape = w.shape();
+  auto w_strides = w.strides();
+  auto s_strides = scales.strides();
+
+  // Add split_k dim with reshapes
+  x_shape.insert(x_shape.end() - 2, split_k);
+  x_shape.back() /= split_k;
+  x_strides.insert(x_strides.end() - 2, split_D);
+  x_strides[x_ndim - 1] = split_D;
+  x_batch_ndims += 1;
+
+  w_shape.insert(w_shape.end() - 2, split_k);
+  w_shape[w.ndim() - 1] /= split_k;
+  w_strides.insert(w_strides.end() - 2, split_D * w.shape(-1));
+  w_batch_ndims += 1;
+  s_strides.insert(s_strides.end() - 2, split_D * scales.shape(-1));
+
+  int final_block_size = K - (split_k - 1) * split_D;
+
+  auto temp_shape = out.shape();
+  if (temp_shape.size() == 1) {
+    temp_shape.insert(temp_shape.begin(), 1);
+  }
+  temp_shape.insert(temp_shape.end() - 2, split_k);
+  array intermediate(temp_shape, x.dtype(), nullptr, {});
+  intermediate.set_data(allocator::malloc(intermediate.nbytes()));
+  d.add_temporary(intermediate, s.index);
+
+  std::string type_string = get_type_string(x.dtype());
+  std::string kname;
+  kname.reserve(64);
+  concatenate(
+      kname,
+      mode + "_qvm_split_k_",
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      "_spk_",
+      split_k);
+
+  // Encode and dispatch kernel
+  auto kernel = get_quantized_kernel_wrapped(
+      d, kname, "qvm_split_k", mode, type_string, group_size, bits, split_k);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_output_array(intermediate, c++);
+  compute_encoder.set_bytes(split_D, c++);
+  compute_encoder.set_bytes(N, c++);
+
+  compute_encoder.set_bytes(x_batch_ndims, c++);
+  compute_encoder.set_vector_bytes(x_shape, c++);
+  compute_encoder.set_vector_bytes(x_strides, c++);
+  compute_encoder.set_bytes(w_batch_ndims, c++);
+  compute_encoder.set_vector_bytes(w_shape, c++);
+  compute_encoder.set_vector_bytes(w_strides, c++);
+  compute_encoder.set_vector_bytes(s_strides, c++);
+  if (biases) {
+    auto b_strides = biases->strides();
+    b_strides.insert(b_strides.end() - 2, split_D * biases->shape(-1));
+    compute_encoder.set_vector_bytes(b_strides, c++);
+  }
+  compute_encoder.set_bytes(final_block_size, c++);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  int axis = intermediate.ndim() - 3;
+  ReductionPlan plan(
+      ReductionOpType::ContiguousStridedReduce,
+      {intermediate.shape(axis)},
+      {intermediate.strides(axis)});
+  strided_reduce_general_dispatch(
+      intermediate, out, "sum", plan, {axis}, compute_encoder, d, s);
+}
+
+void qvm(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int B = out.size() / M / N;
+
+  constexpr int num_simdgroups = 2;
+  constexpr int bk = 32;
+  int bn = std::min(group_size, 32) * num_simdgroups;
+  MTL::Size group_dims(bk, num_simdgroups, 1);
+  MTL::Size grid_dims(M, (N + bn - 1) / bn, B);
+
+  std::string kname;
+  kname.reserve(64);
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode + "_qvm_",
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      B > 1 ? "_batch_1" : "_batch_0");
+  auto kernel = get_quantized_kernel_wrapped(
+      d, kname, "qvm", mode, type_string, group_size, bits, B > 1);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c++);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void qmm_nax(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    bool transpose,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int B = out.size() / M / N;
+
+  int wm = 2;
+  int wn = 2;
+  int bm = 64;
+  int bn = 64;
+  int bk = 64;
+  MTL::Size group_dims(32, wn, wm);
+  MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B);
+
+  std::string kname;
+  kname.reserve(64);
+  bool aligned = N % 64 == 0;
+  bool batched = B > 1;
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode + (transpose ? "_qmm_t_nax_" : "_qmm_n_nax_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn,
+      transpose ? (aligned ? "_alN_true" : "_alN_false") : "",
+      batched ? "_batch_1" : "_batch_0");
+  std::string template_def;
+  MTL::ComputePipelineState* kernel;
+  if (transpose) {
+    kernel = get_qmm_nax_kernel_wrapped(
+        d,
+        kname,
+        "qmm_t_nax",
+        mode,
+        type_string,
+        group_size,
+        bits,
+        aligned,
+        batched,
+        bm,
+        bk,
+        bn,
+        wm,
+        wn);
+  } else {
+    kernel = get_qmm_nax_kernel_wrapped(
+        d,
+        kname,
+        "qmm_n_nax",
+        mode,
+        type_string,
+        group_size,
+        bits,
+        batched,
+        bm,
+        bk,
+        bn,
+        wm,
+        wn);
+  }
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  compute_encoder.set_bytes(M, c++);
+  add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_qmm_nax(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    bool transpose,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int B = out.size() / M / N;
+
+  int wm = 2;
+  int wn = 2;
+  int bm = 64;
+  int bn = 64;
+  int bk = 32;
+  MTL::Size group_dims(32, wn, wm);
+  MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B);
+
+  std::string kname;
+  kname.reserve(64);
+  bool aligned = N % 64 == 0;
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode + (transpose ? "_gather_qmm_t_nax_" : "_gather_qmm_n_nax_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      "_bm",
+      bm,
+      "_bn",
+      bn,
+      "_bk",
+      bk,
+      "_wm",
+      wm,
+      "_wn",
+      wn,
+      transpose ? (aligned ? "_alN_true" : "_alN_false") : "");
+  MTL::ComputePipelineState* kernel;
+  if (transpose) {
+    kernel = get_qmm_nax_kernel_wrapped(
+        d,
+        kname,
+        "gather_qmm_t_nax_",
+        mode,
+        type_string,
+        group_size,
+        bits,
+        aligned,
+        bm,
+        bk,
+        bn,
+        wm,
+        wn);
+  } else {
+    kernel = get_qmm_nax_kernel_wrapped(
+        d,
+        kname,
+        "gather_qmm_n_nax_",
+        mode,
+        type_string,
+        group_size,
+        bits,
+        bm,
+        bk,
+        bn,
+        wm,
+        wn);
+  }
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_input_array(lhs_indices, c++);
+  compute_encoder.set_input_array(rhs_indices, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  compute_encoder.set_bytes(M, c++);
+  c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c);
+  add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void qmm(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    bool transpose,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  if (metal::is_nax_available() && transpose && (K % 64 == 0) &&
+      (env::enable_tf32() || x.dtype() != float32)) {
+    return qmm_nax(
+        /* const array& x = */ x,
+        /* const array& w = */ w,
+        /* const array& scales = */ scales,
+        /* const std::optional<array>& biases = */ biases,
+        /* array& out = */ out,
+        /* bool transpose = */ transpose,
+        /* int group_size = */ group_size,
+        /* int bits = */ bits,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* metal::Device& d = */ d,
+        /* const Stream& s = */ s,
+        /* const std::string& mode = */ mode);
+  }
+
+  int B = out.size() / M / N;
+
+  int wm = 2;
+  int wn = 2;
+  int bm = 32;
+  int bn = 32;
+  MTL::Size group_dims(32, wn, wm);
+  MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B);
+
+  std::string kname;
+  kname.reserve(64);
+  bool aligned = N % 32 == 0;
+  bool batched = B > 1;
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode + (transpose ? "_qmm_t_" : "_qmm_n_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      transpose ? (aligned ? "_alN_true" : "_alN_false") : "",
+      batched ? "_batch_1" : "_batch_0");
+  std::string template_def;
+  MTL::ComputePipelineState* kernel;
+  if (transpose) {
+    kernel = get_quantized_kernel_wrapped(
+        d,
+        kname,
+        "qmm_t",
+        mode,
+        type_string,
+        group_size,
+        bits,
+        aligned,
+        batched);
+  } else {
+    kernel = get_quantized_kernel_wrapped(
+        d, kname, "qmm_n", mode, type_string, group_size, bits, batched);
+  }
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  compute_encoder.set_bytes(M, c++);
+  add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_qmm(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    bool transpose,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  if (metal::is_nax_available() && transpose && (K % 64 == 0) &&
+      (env::enable_tf32() || x.dtype() != float32)) {
+    return gather_qmm_nax(
+        /* const array& x = */ x,
+        /* const array& w = */ w,
+        /* const array& scales = */ scales,
+        /* const std::optional<array>& biases = */ biases,
+        /* const array& lhs_indices = */ lhs_indices,
+        /* const array& rhs_indices = */ rhs_indices,
+        /* array& out = */ out,
+        /* bool transpose = */ transpose,
+        /* int group_size = */ group_size,
+        /* int bits = */ bits,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* metal::Device& d = */ d,
+        /* const Stream& s = */ s,
+        /* const std::string& mode = */ mode);
+  }
+
+  int B = out.size() / M / N;
+
+  int wm = 2;
+  int wn = 2;
+  int bm = 32;
+  int bn = 32;
+  MTL::Size group_dims(32, wn, wm);
+  MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B);
+
+  std::string kname;
+  kname.reserve(64);
+  bool aligned = N % 32 == 0;
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode + (transpose ? "_gather_qmm_t_" : "_gather_qmm_n_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      transpose ? (aligned ? "_alN_true" : "_alN_false") : "");
+  MTL::ComputePipelineState* kernel;
+  if (transpose) {
+    kernel = get_quantized_kernel_wrapped(
+        d, kname, "gather_qmm_t", mode, type_string, group_size, bits, aligned);
+  } else {
+    kernel = get_quantized_kernel_wrapped(
+        d, kname, "gather_qmm_n", mode, type_string, group_size, bits);
+  }
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_input_array(lhs_indices, c++);
+  compute_encoder.set_input_array(rhs_indices, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  compute_encoder.set_bytes(M, c++);
+  c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c);
+  add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int B = out.size() / M / N;
+
+  int bn = 8;
+  int bk = 32;
+  MTL::Size group_dims(bk, 2, 1);
+  MTL::Size grid_dims(M, (N + bn - 1) / bn, B);
+
+  std::string kname;
+  kname.reserve(64);
+  std::string type_string = get_type_string(x.dtype());
+  bool fast = N % bn == 0 && K % 512 == 0;
+  concatenate(
+      kname,
+      mode + (fast ? "_gather_qmv_fast_" : "_gather_qmv_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits);
+
+  auto kernel = get_quantized_kernel_wrapped(
+      d,
+      kname,
+      (fast ? "gather_qmv_fast" : "gather_qmv"),
+      mode,
+      type_string,
+      group_size,
+      bits);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_input_array(lhs_indices, c++);
+  compute_encoder.set_input_array(rhs_indices, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c);
+  add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_qvm(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  int B = out.size() / M / N;
+
+  constexpr int num_simdgroups = 2;
+  constexpr int bk = 32;
+  int bn = std::min(group_size, 32) * num_simdgroups;
+  MTL::Size group_dims(bk, num_simdgroups, 1);
+  MTL::Size grid_dims(M, (N + bn - 1) / bn, B);
+
+  std::string kname;
+  kname.reserve(64);
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode + "_gather_qvm_",
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits);
+  auto kernel = get_quantized_kernel_wrapped(
+      d, kname, "gather_qvm", mode, type_string, group_size, bits);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int c = 0;
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases) {
+    compute_encoder.set_input_array(*biases, c++);
+  }
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_input_array(lhs_indices, c++);
+  compute_encoder.set_input_array(rhs_indices, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(K, c++);
+  compute_encoder.set_bytes(N, c++);
+  c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c++);
+  add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_qmm_rhs_nax(
+    const array& x_,
+    const array& w_,
+    const array& scales_,
+    const std::optional<array>& biases_,
+    const array& indices_,
+    array& out,
+    bool transpose,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string mode) {
+  // Start by normalizing the indices
+  array indices = ensure_row_contiguous(indices_, d, s);
+
+  // Broadcast x with indices. If we are here that means lhs_indices were not
+  // provided so the lhs_indices are implied to be the shape of x broadcasted
+  // with rhs_indices. We need only broadcast x and copy it as if applying the
+  // lhs_indices.
+  auto broadcast_with_indices = [&d, &s, &indices](const array& x) {
+    if (x.size() / x.shape(-2) / x.shape(-1) == indices.size()) {
+      return ensure_row_contiguous(x, d, s);
+    }
+
+    auto x_shape = indices.shape();
+    x_shape.push_back(x.shape(-2));
+    x_shape.push_back(x.shape(-1));
+    array new_x(std::move(x_shape), x.dtype(), nullptr, {});
+    broadcast(x, new_x);
+    return ensure_row_contiguous(new_x, d, s);
+  };
+
+  // Normalize the input arrays
+  array x = broadcast_with_indices(x_);
+  array w = ensure_row_contiguous(w_, d, s);
+  array scales = ensure_row_contiguous(scales_, d, s);
+
+  // TODO: Tune the block sizes
+  int bm = 64, bn = 64, bk = 64;
+  int wm = 2, wn = 2;
+
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  // Make the kernel name
+  std::string kname;
+  kname.reserve(64);
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode +
+          (transpose ? "_gather_qmm_rhs_nax_nt_" : "_gather_qmm_rhs_nax_nn_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      "_bm_",
+      bm,
+      "_bn_",
+      bn,
+      "_bk_",
+      bk,
+      "_wm_",
+      wm,
+      "_wn_",
+      wn);
+
+  metal::MTLFCList func_consts = {
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+      {&align_K, MTL::DataType::DataTypeBool, 202},
+  };
+
+  // And the kernel hash that includes the function constants
+  std::string hash_name;
+  hash_name.reserve(128);
+  concatenate(
+      hash_name,
+      kname,
+      "_align_M_",
+      align_M ? 't' : 'n',
+      "_align_N_",
+      align_N ? 't' : 'n',
+      "_align_K_",
+      align_K ? 't' : 'n');
+
+  // Get and set the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_gather_qmm_nax_kernel(
+      d,
+      kname,
+      hash_name,
+      func_consts,
+      x,
+      group_size,
+      bits,
+      mode,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn,
+      transpose);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  MTL::Size group_dims(32, wn, wm);
+  MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1);
+
+  int c = 0;
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases_) {
+    array biases = ensure_row_contiguous(*biases_, d, s);
+    compute_encoder.set_input_array(biases, c++);
+  }
+  compute_encoder.set_input_array(indices, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(M, c++);
+  compute_encoder.set_bytes(N, c++);
+  compute_encoder.set_bytes(K, c++);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void gather_qmm_rhs(
+    const array& x_,
+    const array& w_,
+    const array& scales_,
+    const std::optional<array>& biases_,
+    const array& indices_,
+    array& out,
+    bool transpose,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string mode) {
+  if (metal::is_nax_available() && transpose &&
+      (env::enable_tf32() || x_.dtype() != float32)) {
+    return gather_qmm_rhs_nax(
+        /* const array& x_ = */ x_,
+        /* const array& w_ = */ w_,
+        /* const array& scales_ = */ scales_,
+        /* const std::optional<array>& biases_ = */ biases_,
+        /* const array& indices_ = */ indices_,
+        /* array& out = */ out,
+        /* bool transpose = */ transpose,
+        /* int group_size = */ group_size,
+        /* int bits = */ bits,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* metal::Device& d = */ d,
+        /* const Stream& s = */ s,
+        /* const std::string mode = */ mode);
+  }
+
+  // Start by normalizing the indices
+  array indices = ensure_row_contiguous(indices_, d, s);
+
+  // Broadcast x with indices. If we are here that means lhs_indices were not
+  // provided so the lhs_indices are implied to be the shape of x broadcasted
+  // with rhs_indices. We need only broadcast x and copy it as if applying the
+  // lhs_indices.
+  auto broadcast_with_indices = [&d, &s, &indices](const array& x) {
+    if (x.size() / x.shape(-2) / x.shape(-1) == indices.size()) {
+      return ensure_row_contiguous(x, d, s);
+    }
+
+    auto x_shape = indices.shape();
+    x_shape.push_back(x.shape(-2));
+    x_shape.push_back(x.shape(-1));
+    array new_x(std::move(x_shape), x.dtype(), nullptr, {});
+    broadcast(x, new_x);
+    return ensure_row_contiguous(new_x, d, s);
+  };
+
+  // Normalize the input arrays
+  array x = broadcast_with_indices(x_);
+  array w = ensure_row_contiguous(w_, d, s);
+  array scales = ensure_row_contiguous(scales_, d, s);
+
+  // TODO: Tune the block sizes
+  int bm = 16, bn = 32, bk = 32;
+  int wm = 1, wn = 2;
+
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  // Make the kernel name
+  std::string kname;
+  kname.reserve(64);
+  std::string type_string = get_type_string(x.dtype());
+  concatenate(
+      kname,
+      mode + (transpose ? "_gather_qmm_rhs_nt_" : "_gather_qmm_rhs_nn_"),
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits,
+      "_bm_",
+      bm,
+      "_bn_",
+      bn,
+      "_bk_",
+      bk,
+      "_wm_",
+      wm,
+      "_wn_",
+      wn);
+
+  metal::MTLFCList func_consts = {
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201},
+      {&align_K, MTL::DataType::DataTypeBool, 202},
+  };
+
+  // And the kernel hash that includes the function constants
+  std::string hash_name;
+  hash_name.reserve(128);
+  concatenate(
+      hash_name,
+      kname,
+      "_align_M_",
+      align_M ? 't' : 'n',
+      "_align_N_",
+      align_N ? 't' : 'n',
+      "_align_K_",
+      align_K ? 't' : 'n');
+
+  // Get and set the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_gather_qmm_kernel(
+      d,
+      kname,
+      hash_name,
+      func_consts,
+      x,
+      group_size,
+      bits,
+      mode,
+      bm,
+      bn,
+      bk,
+      wm,
+      wn,
+      transpose);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  MTL::Size group_dims(32, wn, wm);
+  MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1);
+
+  int c = 0;
+  compute_encoder.set_input_array(x, c++);
+  compute_encoder.set_input_array(w, c++);
+  compute_encoder.set_input_array(scales, c++);
+  if (biases_) {
+    array biases = ensure_row_contiguous(*biases_, d, s);
+    compute_encoder.set_input_array(biases, c++);
+  }
+  compute_encoder.set_input_array(indices, c++);
+  compute_encoder.set_output_array(out, c++);
+  compute_encoder.set_bytes(M, c++);
+  compute_encoder.set_bytes(N, c++);
+  compute_encoder.set_bytes(K, c++);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void dispatch_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int group_size,
+    int bits,
+    int M,
+    int N,
+    int K,
+    metal::Device& d,
+    const Stream& s,
+    const std::string& mode) {
+  // It is a qmv with a small inner dimension so route to qmv_quad kernel
+  if ((K == 128 || K == 64) && is_power_of_2(bits)) {
+    qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode);
+    return;
+  }
+
+  // Run of the mill qmv
+  qmv(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode);
+}
+
+void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  // Make sure the last two dims of x and w, s, b are contiguous. This should
+  // be relaxed for x.
+  array x = ensure_row_contiguous_matrix(inputs[0], d, s);
+  array w = ensure_row_contiguous_matrix(inputs[1], d, s);
+  array scales = ensure_row_contiguous_matrix(inputs[2], d, s);
+  std::optional<array> biases = std::nullopt;
+  if (inputs.size() == 4) {
+    biases = ensure_row_contiguous_matrix(inputs[3], d, s);
+  }
+
+  // Extract the matmul shapes
+  bool non_batched = w.ndim() == 2 && x.flags().row_contiguous;
+  int K = x.shape(-1);
+  int M = non_batched ? x.size() / K : x.shape(-2);
+  int N = out.shape(-1);
+
+  int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4;
+  auto mode = quantization_mode_to_string(mode_);
+  // It is a matrix matrix product.
+  if (M >= vector_limit) {
+    qmm(x,
+        w,
+        scales,
+        biases,
+        out,
+        transpose_,
+        group_size_,
+        bits_,
+        M,
+        N,
+        K,
+        d,
+        s,
+        mode);
+    return;
+  }
+
+  // Run of the mill qmv
+  if (transpose_) {
+    dispatch_qmv(
+        x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
+    return;
+  }
+
+  // Run of the mill qvm
+  if (K < 1024) {
+    qvm(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
+    return;
+  }
+
+  // Qvm with large dimension so route to a split K kernel for more parallelism
+  qvm_split_k(
+      x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
+  return;
+}
+
+void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  array x = ensure_row_contiguous_matrix(inputs[0], d, s);
+  array w = ensure_row_contiguous_matrix(inputs[1], d, s);
+  array scales = ensure_row_contiguous_matrix(inputs[2], d, s);
+  std::optional<array> biases = std::nullopt;
+  if (inputs.size() == 6) {
+    biases = ensure_row_contiguous_matrix(inputs[3], d, s);
+  }
+  const array& lhs_indices = inputs[inputs.size() - 2];
+  const array& rhs_indices = inputs[inputs.size() - 1];
+
+  int K = x.shape(-1);
+  int M = x.shape(-2);
+  int N = out.shape(-1);
+  int B = out.size() / M / N;
+  int E = w.size() / w.shape(-1) / w.shape(-2);
+  int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4;
+  auto mode = quantization_mode_to_string(mode_);
+
+  // We are walking x in order and w is also in order so we can batch up the
+  // matmuls and reuse reading x and w.
+  //
+  // TODO: Tune 16 and 4 here a bit better.
+  if (M == 1 && B >= 16 && right_sorted_ == true && B / E >= 4) {
+    gather_qmm_rhs(
+        x,
+        w,
+        scales,
+        biases,
+        rhs_indices,
+        out,
+        transpose_,
+        group_size_,
+        bits_,
+        x.size() / K,
+        N,
+        K,
+        d,
+        s,
+        mode);
+    return;
+  }
+
+  // It is a matrix matrix product
+  if (M >= vector_limit) {
+    gather_qmm(
+        x,
+        w,
+        scales,
+        biases,
+        lhs_indices,
+        rhs_indices,
+        out,
+        transpose_,
+        group_size_,
+        bits_,
+        M,
+        N,
+        K,
+        d,
+        s,
+        mode);
+    return;
+  }
+
+  if (transpose_) {
+    gather_qmv(
+        x,
+        w,
+        scales,
+        biases,
+        lhs_indices,
+        rhs_indices,
+        out,
+        group_size_,
+        bits_,
+        M,
+        N,
+        K,
+        d,
+        s,
+        mode);
+    return;
+  }
+
+  gather_qvm(
+      x,
+      w,
+      scales,
+      biases,
+      lhs_indices,
+      rhs_indices,
+      out,
+      group_size_,
+      bits_,
+      M,
+      N,
+      K,
+      d,
+      s,
+      mode);
+}
+
+void quantize_dequantize(
+    const array& in,
+    array& out,
+    std::string mode,
+    int group_size,
+    int bits,
+    metal::Device& d,
+    const Stream& s) {
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  auto w = ensure_row_contiguous(in, d, s);
+  compute_encoder.set_input_array(w, 0);
+  compute_encoder.set_output_array(out, 1);
+  auto type_string = get_type_string(in.dtype());
+  std::string kname;
+  concatenate(
+      kname,
+      mode + "_quantize_dequantize_",
+      type_string,
+      "_gs_",
+      group_size,
+      "_b_",
+      bits);
+  auto kernel = get_quantized_kernel_wrapped(
+      d, kname, "quantize_dequantize", mode, type_string, group_size, bits);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  constexpr int uint8_per_uint32 = 4;
+  constexpr int simd_size = 32;
+  int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits;
+  int per_thread = std::max(group_size / simd_size, 1);
+  size_t nthreads = w.size() / per_thread;
+
+  NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (thread_group_size > nthreads) {
+    thread_group_size = nthreads;
+  }
+  auto group_dims = MTL::Size(thread_group_size, 1, 1);
+  bool use_2d = nthreads > UINT_MAX;
+  auto grid_shape = w.shape();
+  grid_shape.back() /= per_thread;
+  MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides())
+                               : MTL::Size(nthreads, 1, 1);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto mode = quantization_mode_to_string(mode_);
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    out.set_data(allocator::malloc(out.nbytes()));
+
+    bool donate_x = inputs[0].is_donatable();
+    array x = ensure_row_contiguous(inputs[0], d, s);
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype());
+    quantize_dequantize(x, xhat, mode, group_size_, bits_, d, s);
+
+    // Make sure the last two dims of w and s are contiguous
+    array w = ensure_row_contiguous_matrix(inputs[1], d, s);
+    array scales = ensure_row_contiguous_matrix(inputs[2], d, s);
+
+    bool non_batched = w.ndim() == 2;
+    int K = x.shape(-1);
+    int M = non_batched ? x.size() / K : x.shape(-2);
+    int N = out.shape(-1);
+    dispatch_qmv(
+        xhat,
+        w,
+        scales,
+        std::nullopt,
+        out,
+        group_size_,
+        bits_,
+        M,
+        N,
+        K,
+        d,
+        s,
+        mode);
+    return;
+  } else {
+    throw std::runtime_error("[QQMatmul] NYI for the general case");
+  }
+}
+
+void fast::Quantize::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& w_pre = inputs[0];
+  auto& out = outputs[0];
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  auto w = ensure_row_contiguous(w_pre, d, s);
+  if (dequantize_) {
+    auto scales = ensure_row_contiguous(inputs[1], d, s);
+    if (mode_ == QuantizationMode::Affine) {
+      auto biases = ensure_row_contiguous(inputs[2], d, s);
+      compute_encoder.set_input_array(biases, 2);
+    }
+    compute_encoder.set_input_array(w, 0);
+    compute_encoder.set_input_array(scales, 1);
+    compute_encoder.set_output_array(out, 3);
+  } else {
+    auto& scales = outputs[1];
+    scales.set_data(allocator::malloc(scales.nbytes()));
+    if (mode_ == QuantizationMode::Affine) {
+      auto& biases = outputs[2];
+      biases.set_data(allocator::malloc(biases.nbytes()));
+      compute_encoder.set_output_array(biases, 3);
+    }
+    compute_encoder.set_input_array(w, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_output_array(scales, 2);
+  }
+
+  auto type_string = dequantize_ ? get_type_string(out.dtype())
+                                 : get_type_string(w_pre.dtype());
+  auto mode = quantization_mode_to_string(mode_);
+  std::string kname;
+  concatenate(
+      kname,
+      mode + (dequantize_ ? "_dequantize" : "_quantize"),
+      "_",
+      type_string,
+      "_gs_",
+      group_size_,
+      "_b_",
+      bits_);
+  auto kernel = get_quantized_kernel_wrapped(
+      d,
+      kname,
+      dequantize_ ? "dequantize" : "quantize",
+      mode,
+      type_string,
+      group_size_,
+      bits_);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Treat uint32 as uint8 in kernel
+  constexpr int uint8_per_uint32 = 4;
+  constexpr int simd_size = 32;
+  int packs_per_int = (bits_ == 3 || bits_ == 5) ? 8
+      : bits_ == 6                               ? 4
+                                                 : 8 / bits_;
+  int per_thread =
+      dequantize_ ? packs_per_int : std::max(group_size_ / simd_size, 1);
+  size_t nthreads =
+      dequantize_ ? out.size() / packs_per_int : w.size() / per_thread;
+
+  NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (thread_group_size > nthreads) {
+    thread_group_size = nthreads;
+  }
+  auto group_dims = MTL::Size(thread_group_size, 1, 1);
+  bool use_2d = nthreads > UINT_MAX;
+  auto grid_shape = w.shape();
+  if (dequantize_) {
+    grid_shape.back() *= uint8_per_uint32;
+  } else {
+    grid_shape.back() /= per_thread;
+  }
+  MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides())
+                               : MTL::Size(nthreads, 1, 1);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void fast::ConvertFP8::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& in = inputs[0];
+  auto& out = outputs[0];
+  unary_op_gpu(inputs, out, name(), stream());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/reduce.cpp b/macos/third_party/mlx/mlx/backend/metal/reduce.cpp
new file mode 100644
index 0000000..504943d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/reduce.cpp
@@ -0,0 +1,1038 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <algorithm>
+#include <cassert>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/reduce.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+struct RowReduceArgs {
+  // Input shape and strides not including the reduction axes
+  Shape shape;
+  Strides strides;
+  int ndim;
+
+  // Input shape and strides for the reduction axes
+  Shape reduce_shape;
+  Strides reduce_strides;
+  int reduce_ndim;
+
+  // The number of rows we are reducing. Namely prod(reduce_shape).
+  size_t non_row_reductions;
+
+  // The size of the row.
+  size_t row_size;
+
+  RowReduceArgs(
+      const array& in,
+      const ReductionPlan& plan,
+      const std::vector<int>& axes) {
+    row_size = plan.shape.back();
+
+    reduce_shape = plan.shape;
+    reduce_strides = plan.strides;
+    reduce_shape.pop_back();
+    reduce_strides.pop_back();
+    reduce_ndim = reduce_shape.size();
+
+    non_row_reductions = 1;
+    for (auto s : reduce_shape) {
+      non_row_reductions *= s;
+    }
+
+    std::tie(shape, strides) = shapes_without_reduction_axes(in, axes);
+    std::tie(shape, strides) = collapse_contiguous_dims(shape, strides);
+    ndim = shape.size();
+  }
+
+  void encode(CommandEncoder& compute_encoder) {
+    // Push 0s to avoid encoding empty vectors.
+    if (reduce_ndim == 0) {
+      reduce_shape.push_back(0);
+      reduce_strides.push_back(0);
+    }
+    if (ndim == 0) {
+      shape.push_back(0);
+      strides.push_back(0);
+    }
+
+    compute_encoder.set_bytes(row_size, 2);
+    compute_encoder.set_bytes(non_row_reductions, 3);
+    compute_encoder.set_vector_bytes(shape, 4);
+    compute_encoder.set_vector_bytes(strides, 5);
+    compute_encoder.set_bytes(ndim, 6);
+    compute_encoder.set_vector_bytes(reduce_shape, 7);
+    compute_encoder.set_vector_bytes(reduce_strides, 8);
+    compute_encoder.set_bytes(reduce_ndim, 9);
+
+    if (reduce_ndim == 0) {
+      reduce_shape.pop_back();
+      reduce_strides.pop_back();
+    }
+    if (ndim == 0) {
+      shape.pop_back();
+      strides.pop_back();
+    }
+  }
+};
+
+struct ColReduceArgs {
+  // Input shape and strides not including the reduction axes
+  Shape shape;
+  Strides strides;
+  int ndim;
+
+  // Input shape and strides for the reduction axes
+  Shape reduce_shape;
+  Strides reduce_strides;
+  int reduce_ndim;
+
+  // The number of column reductions we are doing. Namely prod(reduce_shape).
+  size_t non_col_reductions;
+
+  // The size of the contiguous column reduction.
+  size_t reduction_size;
+  int64_t reduction_stride;
+
+  ColReduceArgs(
+      const array& in,
+      const ReductionPlan& plan,
+      const std::vector<int>& axes) {
+    reduction_size = plan.shape.back();
+    reduction_stride = plan.strides.back();
+
+    reduce_shape = plan.shape;
+    reduce_strides = plan.strides;
+    reduce_shape.pop_back();
+    reduce_strides.pop_back();
+    reduce_ndim = reduce_shape.size();
+
+    non_col_reductions = 1;
+    for (auto s : reduce_shape) {
+      non_col_reductions *= s;
+    }
+
+    // We 'll use a stride_back variable because strides.back() could be 0 but
+    // yet we may have removed the appropriate amount of elements. It is safe
+    // to compute the stride by multiplying shapes (while < reduction_stride)
+    // because it is a contiguous section.
+    int64_t stride_back = 1;
+    std::tie(shape, strides) = shapes_without_reduction_axes(in, axes);
+    while (!shape.empty() && stride_back < reduction_stride) {
+      stride_back *= shape.back();
+      shape.pop_back();
+      strides.pop_back();
+    }
+    std::tie(shape, strides) = collapse_contiguous_dims(shape, strides);
+    ndim = shape.size();
+  }
+
+  /**
+   * Create the col reduce arguments for reducing the 1st axis of the row
+   * contiguous intermediate array.
+   */
+  ColReduceArgs(const array& intermediate) {
+    assert(intermediate.flags().row_contiguous);
+
+    reduction_size = intermediate.shape(0);
+    reduction_stride = intermediate.size() / reduction_size;
+    non_col_reductions = 1;
+    reduce_ndim = 0;
+    ndim = 0;
+  }
+
+  void encode(CommandEncoder& compute_encoder) {
+    // Push 0s to avoid encoding empty vectors.
+    if (reduce_ndim == 0) {
+      reduce_shape.push_back(0);
+      reduce_strides.push_back(0);
+    }
+    if (ndim == 0) {
+      shape.push_back(0);
+      strides.push_back(0);
+    }
+
+    compute_encoder.set_bytes(reduction_size, 2);
+    compute_encoder.set_bytes(reduction_stride, 3);
+    compute_encoder.set_vector_bytes(shape, 4);
+    compute_encoder.set_vector_bytes(strides, 5);
+    compute_encoder.set_bytes(ndim, 6);
+    compute_encoder.set_vector_bytes(reduce_shape, 7);
+    compute_encoder.set_vector_bytes(reduce_strides, 8);
+    compute_encoder.set_bytes(reduce_ndim, 9);
+    compute_encoder.set_bytes(non_col_reductions, 10);
+
+    if (reduce_ndim == 0) {
+      reduce_shape.pop_back();
+      reduce_strides.pop_back();
+    }
+    if (ndim == 0) {
+      shape.pop_back();
+      strides.pop_back();
+    }
+  }
+};
+
+} // namespace
+
+inline auto safe_div(size_t n, size_t m) {
+  return m == 0 ? 0 : (n + m - 1) / m;
+}
+
+inline auto safe_divup(size_t n, size_t m) {
+  return safe_div(n, m) * m;
+}
+
+inline bool is_64b_int(Dtype dtype) {
+  return dtype == int64 || dtype == uint64;
+}
+
+inline bool is_64b_dtype(Dtype dtype) {
+  return dtype == int64 || dtype == uint64 || dtype == complex64;
+}
+
+inline int get_kernel_reduce_ndim(int reduce_ndim) {
+  if (reduce_ndim <= 1) {
+    return 1;
+  } else if (reduce_ndim == 2) {
+    return 2;
+  } else {
+    return 5;
+  }
+}
+
+inline int threadgroup_size_from_row_size(int row_size) {
+  // 1 simdgroup per row smallish rows
+  if (row_size <= 512) {
+    return 32;
+  }
+
+  // 2 simdgroups per row for medium rows
+  if (row_size <= 1024) {
+    return 128;
+  }
+
+  // up to 32 simdgroups after that
+  int thread_group_size;
+  thread_group_size = (row_size + REDUCE_N_READS - 1) / REDUCE_N_READS;
+  thread_group_size = ((thread_group_size + 31) / 32) * 32;
+  thread_group_size = std::min(1024, thread_group_size);
+  return thread_group_size;
+}
+
+inline auto output_grid_for_col_reduce(
+    const array& out,
+    const ColReduceArgs& args) {
+  auto out_shape = out.shape();
+  auto out_strides = out.strides();
+  while (!out_shape.empty() && out_strides.back() < args.reduction_stride) {
+    out_shape.pop_back();
+    out_strides.pop_back();
+  }
+  return get_2d_grid_dims(out_shape, out_strides);
+}
+
+std::pair<Dtype, Dtype> remap_reduce_types(
+    const array& in,
+    const std::string& op_name) {
+  if (op_name == "sum" || op_name == "prod") {
+    if (issubdtype(in.dtype(), integer)) {
+      switch (in.dtype()) {
+        case uint8:
+          return {uint8, uint32};
+        case uint16:
+          return {uint16, uint32};
+        case uint32:
+          return {uint32, uint32};
+        case uint64:
+          return {uint64, uint64};
+        case int8:
+          return {int8, int32};
+        case int16:
+          return {int16, int32};
+        case int32:
+          return {int32, int32};
+        case int64:
+          return {int64, int64};
+        default:
+          throw std::runtime_error("Unsupported integer type");
+      }
+    }
+    if (in.dtype() == bool_) {
+      return {int8, int32};
+    }
+    return {in.dtype(), in.dtype()};
+  } else if (op_name == "and" || op_name == "or") {
+    if (in.dtype().size() == 1) {
+      return {bool_, bool_};
+    } else if (in.dtype().size() == 2) {
+      return {int16, bool_};
+    } else if (in.dtype().size() == 4) {
+      return {int32, bool_};
+    } else {
+      return {int64, bool_};
+    }
+  }
+  return {in.dtype(), in.dtype()};
+}
+
+void init_reduce(
+    array& out,
+    const std::string& op_name,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  auto [_, out_type] = remap_reduce_types(out, op_name);
+  const std::string func_name = "init_reduce";
+  std::string kname = func_name;
+  concatenate(kname, "_", op_name, type_to_name(out_type));
+  auto kernel = get_reduce_init_kernel(d, kname, func_name, op_name, out_type);
+  size_t nthreads = out.size();
+  MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
+  NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (thread_group_size > nthreads) {
+    thread_group_size = nthreads;
+  }
+  MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_output_array(out, 0);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void all_reduce_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  // Set the kernel
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+  const std::string func_name = "all_reduce";
+  std::string kname = func_name;
+  concatenate(kname, "_", op_name, type_to_name(in_type));
+  auto kernel = get_reduce_kernel(
+      d, kname, func_name, op_name, in_type, out_type, "int64_t");
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  size_t in_size = in.size();
+
+  // Small array so dispatch a single threadgroup
+  if (in_size <= REDUCE_N_READS * 1024) {
+    int threadgroup_size = (in_size + REDUCE_N_READS - 1) / REDUCE_N_READS;
+    threadgroup_size = ((threadgroup_size + 31) / 32) * 32;
+    MTL::Size grid_dims(threadgroup_size, 1, 1);
+
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(in_size, 2);
+    compute_encoder.set_bytes(in_size, 3);
+    compute_encoder.dispatch_threads(grid_dims, grid_dims);
+  }
+
+  // We need multiple threadgroups so we 'll do it in 2 passes.
+  else {
+    int n_rows, threadgroup_2nd_pass;
+    // Less than 2**26 bytes
+    if (in.nbytes() <= (1 << 26)) {
+      n_rows = 32 * REDUCE_N_READS;
+      threadgroup_2nd_pass = 32;
+    }
+
+    // Really large matrix so parallelize as much as possible
+    else {
+      n_rows = 1024 * REDUCE_N_READS;
+      threadgroup_2nd_pass = 1024;
+    }
+
+    // Allocate an intermediate tensor to hold results if needed
+    array intermediate({n_rows}, out_type, nullptr, {});
+    intermediate.set_data(allocator::malloc(intermediate.nbytes()));
+    d.add_temporary(intermediate, s.index);
+
+    // 1st pass
+    size_t row_size = (in_size + n_rows - 1) / n_rows;
+    int threadgroup_size =
+        std::min((row_size + REDUCE_N_READS - 1) / REDUCE_N_READS, 1024ul);
+    threadgroup_size = ((threadgroup_size + 31) / 32) * 32;
+    MTL::Size grid_dims(threadgroup_size, n_rows, 1);
+    MTL::Size group_dims(threadgroup_size, 1, 1);
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(intermediate, 1);
+    compute_encoder.set_bytes(in_size, 2);
+    compute_encoder.set_bytes(row_size, 3);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+    // 2nd pass
+    std::string kname_2nd_pass = func_name;
+    concatenate(kname_2nd_pass, "_", op_name, type_to_name(intermediate));
+    auto kernel_2nd_pass = get_reduce_kernel(
+        d, kname_2nd_pass, func_name, op_name, out_type, out_type, "int64_t");
+    compute_encoder.set_compute_pipeline_state(kernel_2nd_pass);
+    size_t intermediate_size = n_rows;
+    grid_dims = MTL::Size(threadgroup_2nd_pass, 1, 1);
+    group_dims = MTL::Size(threadgroup_2nd_pass, 1, 1);
+    compute_encoder.set_input_array(intermediate, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(intermediate_size, 2);
+    compute_encoder.set_bytes(intermediate_size, 3);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void row_reduce_small(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    RowReduceArgs& args,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  // Set the kernel
+  int n = get_kernel_reduce_ndim(args.reduce_ndim);
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+  const std::string func_name = "row_reduce_small";
+  std::string kname = func_name;
+  bool large = in.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(
+      kname,
+      "_",
+      std::to_string(n),
+      "_reduce_",
+      op_name,
+      type_to_name(in_type));
+  auto kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      in_type,
+      out_type,
+      large ? "size_t" : "int",
+      n);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Figure out the grid dims
+  MTL::Size grid_dims;
+  MTL::Size group_dims;
+  if ((args.non_row_reductions < 32 && args.row_size <= 8) ||
+      args.non_row_reductions <= 8) {
+    grid_dims = get_2d_grid_dims(out.shape(), out.strides());
+    group_dims =
+        MTL::Size((grid_dims.width < 1024) ? grid_dims.width : 1024, 1, 1);
+  } else {
+    auto out_grid_size = get_2d_grid_dims(out.shape(), out.strides());
+    grid_dims = MTL::Size(32, out_grid_size.width, out_grid_size.height);
+    group_dims = MTL::Size(32, 1, 1);
+  }
+
+  // Launch
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  args.encode(compute_encoder);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void row_reduce_simple(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    RowReduceArgs& args,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  // Set the kernel
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+  const std::string func_name = "row_reduce_simple";
+  std::string kname = func_name;
+  concatenate(kname, "_", op_name, type_to_name(in_type));
+
+  auto kernel = get_reduce_kernel(
+      d, kname, func_name, op_name, in_type, out_type, "size_t");
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Figure out the grid dims
+  size_t row_size = args.row_size;
+  size_t out_size = out.size();
+  auto out_grid_size = get_2d_grid_dims(out.shape(), out.strides());
+  out_grid_size.width =
+      (out_grid_size.width + REDUCE_N_WRITES - 1) / REDUCE_N_WRITES;
+  int threadgroup_size = threadgroup_size_from_row_size(row_size);
+  if (in.itemsize() == 8) {
+    threadgroup_size = std::min(threadgroup_size, 512);
+  }
+  MTL::Size grid_dims(
+      threadgroup_size, out_grid_size.width, out_grid_size.height);
+  MTL::Size group_dims(threadgroup_size, 1, 1);
+
+  // Launch
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  compute_encoder.set_bytes(row_size, 2);
+  compute_encoder.set_bytes(out_size, 3);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void row_reduce_looped(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    RowReduceArgs& args,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+
+  // Set the kernel
+  int n = get_kernel_reduce_ndim(args.reduce_ndim);
+  const std::string func_name = "row_reduce_looped";
+  std::string kname = func_name;
+  bool large = in.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(
+      kname,
+      "_",
+      std::to_string(n),
+      "_reduce_",
+      op_name,
+      type_to_name(in_type));
+  auto kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      in_type,
+      out_type,
+      large ? "size_t" : "int",
+      n);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Figure out the grid
+  auto out_grid_size = get_2d_grid_dims(out.shape(), out.strides());
+  int threadgroup_size = threadgroup_size_from_row_size(args.row_size);
+  MTL::Size grid_dims(
+      threadgroup_size, out_grid_size.width, out_grid_size.height);
+  MTL::Size group_dims(threadgroup_size, 1, 1);
+
+  // Launch
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  args.encode(compute_encoder);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void row_reduce_general_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    const ReductionPlan& plan,
+    const std::vector<int>& axes,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  // Prepare the arguments for the kernel
+  RowReduceArgs args(in, plan, axes);
+
+  // Case 1: The row is small
+  if (args.row_size <= 64) {
+    return row_reduce_small(in, out, op_name, args, compute_encoder, d, s);
+  }
+
+  // Case 2: Contiguous reduce without non-row reductions
+  if (plan.type == ContiguousReduce && args.reduce_ndim == 0 &&
+      in.size() / args.row_size >= 32) {
+    return row_reduce_simple(in, out, op_name, args, compute_encoder, d, s);
+  }
+
+  // Case 3: General row reduce including non-row reductions
+  return row_reduce_looped(in, out, op_name, args, compute_encoder, d, s);
+}
+
+void strided_reduce_small(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    ColReduceArgs& args,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+
+  // Figure out the grid dims
+  MTL::Size grid_dims, group_dims;
+
+  // Prepare the arguments for the kernel
+  args.reduce_shape.push_back(args.reduction_size);
+  args.reduce_strides.push_back(args.reduction_stride);
+  args.reduce_ndim++;
+
+  int n = get_kernel_reduce_ndim(args.reduce_ndim);
+  const std::string func_name = "col_reduce_small";
+  std::string kname = func_name;
+  bool large = in.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(
+      kname,
+      "_",
+      std::to_string(n),
+      "_reduce_",
+      op_name,
+      type_to_name(in_type));
+  auto kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      in_type,
+      out_type,
+      large ? "size_t" : "int",
+      n);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  const int n_reads = 4;
+  size_t reduction_stride_blocks =
+      (args.reduction_stride + n_reads - 1) / n_reads;
+  size_t total = args.reduction_size * args.non_col_reductions;
+  size_t threadgroup_x = std::min(reduction_stride_blocks, 32ul);
+  size_t threadgroup_y = std::min(
+      8ul,
+      std::min(kernel->maxTotalThreadsPerThreadgroup() / threadgroup_x, total));
+
+  group_dims = MTL::Size(threadgroup_x, threadgroup_y, 1);
+  grid_dims = output_grid_for_col_reduce(out, args);
+  grid_dims = MTL::Size(
+      (reduction_stride_blocks + threadgroup_x - 1) / threadgroup_x,
+      grid_dims.width,
+      grid_dims.height);
+
+  // Launch
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  args.encode(compute_encoder);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void strided_reduce_longcolumn(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    ColReduceArgs& args,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+  size_t total_reduction_size = args.reduction_size * args.non_col_reductions;
+  size_t outer_blocks = 32;
+  if (total_reduction_size >= 32768) {
+    outer_blocks = 128;
+  }
+
+  // Prepare the temporary accumulator
+  Shape intermediate_shape;
+  intermediate_shape.reserve(out.ndim() + 1);
+  intermediate_shape.push_back(outer_blocks);
+  intermediate_shape.insert(
+      intermediate_shape.end(), out.shape().begin(), out.shape().end());
+  array intermediate(std::move(intermediate_shape), out_type, nullptr, {});
+  intermediate.set_data(allocator::malloc(intermediate.nbytes()));
+  d.add_temporary(intermediate, s.index);
+
+  // Prepare the arguments for the kernel
+  args.reduce_shape.push_back(args.reduction_size);
+  args.reduce_strides.push_back(args.reduction_stride);
+  args.reduce_ndim++;
+
+  // Figure out the grid dims
+  size_t out_size = out.size();
+  size_t threadgroup_x = args.reduction_stride;
+  size_t threadgroup_y =
+      (args.non_col_reductions * args.reduction_size + outer_blocks - 1) /
+      outer_blocks;
+  threadgroup_y = std::min(32ul, threadgroup_y);
+
+  auto out_grid_size = output_grid_for_col_reduce(out, args);
+  MTL::Size grid_dims(out_grid_size.width, out_grid_size.height, outer_blocks);
+  MTL::Size group_dims(threadgroup_x, threadgroup_y, 1);
+
+  // Set the kernel
+  int n = get_kernel_reduce_ndim(args.reduce_ndim);
+  std::string func_name = "col_reduce_longcolumn";
+  std::string kname = func_name;
+  bool large = in.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(
+      kname,
+      "_",
+      std::to_string(n),
+      "_reduce_",
+      op_name,
+      type_to_name(in_type));
+  auto kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      in_type,
+      out_type,
+      large ? "int64_t" : "int",
+      n);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Launch
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(intermediate, 1);
+  args.encode(compute_encoder);
+  compute_encoder.set_bytes(out_size, 11);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  // Make the 2nd pass arguments and grid_dims
+  ColReduceArgs second_args(intermediate);
+  second_args.reduce_shape.push_back(outer_blocks);
+  second_args.reduce_strides.push_back(out.size());
+  second_args.reduce_ndim++;
+  int BN = 32;
+  grid_dims = MTL::Size(256 * ((out.size() + BN - 1) / BN), 1, 1);
+  group_dims = MTL::Size(256, 1, 1);
+
+  // Set the 2nd kernel
+  func_name = "col_reduce_looped";
+  kname = func_name;
+  large = intermediate.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(kname, "_1_32_32_reduce_", op_name, type_to_name(intermediate));
+  kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      intermediate.dtype(),
+      out_type,
+      large ? "int64_t" : "int",
+      1,
+      32,
+      32);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  compute_encoder.set_input_array(intermediate, 0);
+  compute_encoder.set_output_array(out, 1);
+  second_args.encode(compute_encoder);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void strided_reduce_looped(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    ColReduceArgs& args,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+
+  // Prepare the arguments for the kernel
+  args.reduce_shape.push_back(args.reduction_size);
+  args.reduce_strides.push_back(args.reduction_stride);
+  args.reduce_ndim++;
+
+  // Figure out the grid dims
+  auto out_grid_size = output_grid_for_col_reduce(out, args);
+  int BN = 32;
+  int BM = 1024 / BN;
+  int threadgroup_size = 8 * 32;
+  MTL::Size grid_dims(
+      threadgroup_size * ((args.reduction_stride + BN - 1) / BN),
+      out_grid_size.width,
+      out_grid_size.height);
+  MTL::Size group_dims(threadgroup_size, 1, 1);
+
+  // Set the kernel
+  int n = get_kernel_reduce_ndim(args.reduce_ndim);
+  std::string func_name = "col_reduce_looped";
+  std::string kname = func_name;
+  bool large = in.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(
+      kname,
+      "_",
+      std::to_string(n),
+      "_",
+      std::to_string(BM),
+      "_",
+      std::to_string(BN),
+      "_reduce_",
+      op_name,
+      type_to_name(in_type));
+  auto kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      in_type,
+      out_type,
+      large ? "int64_t" : "int",
+      n,
+      BM,
+      BN);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Launch
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  args.encode(compute_encoder);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void strided_reduce_2pass(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    ColReduceArgs& args,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  auto [in_type, out_type] = remap_reduce_types(in, op_name);
+
+  // Prepare the temporary accumulator
+  Shape intermediate_shape;
+  intermediate_shape.reserve(out.ndim() + 1);
+  intermediate_shape.push_back(32);
+  intermediate_shape.insert(
+      intermediate_shape.end(), out.shape().begin(), out.shape().end());
+  array intermediate(std::move(intermediate_shape), out_type, nullptr, {});
+  intermediate.set_data(allocator::malloc(intermediate.nbytes()));
+  d.add_temporary(intermediate, s.index);
+
+  // Prepare the arguments for the kernel
+  args.reduce_shape.push_back(args.reduction_size);
+  args.reduce_strides.push_back(args.reduction_stride);
+  args.reduce_ndim++;
+
+  // Figure out the grid dims
+  size_t out_size = out.size() / args.reduction_stride;
+  auto out_grid_size = output_grid_for_col_reduce(out, args);
+  int outer_blocks = 32;
+  int BN = 32;
+  int BM = 1024 / BN;
+  int threadgroup_size = 8 * 32;
+  MTL::Size grid_dims(
+      threadgroup_size * ((args.reduction_stride + BN - 1) / BN),
+      out_grid_size.width * outer_blocks,
+      out_grid_size.height);
+  MTL::Size group_dims(threadgroup_size, 1, 1);
+
+  // Set the kernel
+  int n = get_kernel_reduce_ndim(args.reduce_ndim);
+  std::string func_name = "col_reduce_2pass";
+  std::string kname = func_name;
+  bool large = in.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(
+      kname,
+      "_",
+      std::to_string(n),
+      "_",
+      std::to_string(BM),
+      "_",
+      std::to_string(BN),
+      "_reduce_",
+      op_name,
+      type_to_name(in_type));
+  auto kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      in_type,
+      out_type,
+      large ? "int64_t" : "int",
+      n,
+      BM,
+      BN);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Launch
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(intermediate, 1);
+  args.encode(compute_encoder);
+  compute_encoder.set_bytes(out_size, 11);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+  // Make the 2nd pass arguments and grid_dims
+  ColReduceArgs second_args(intermediate);
+  second_args.reduce_shape.push_back(outer_blocks);
+  second_args.reduce_strides.push_back(out.size());
+  second_args.reduce_ndim++;
+  grid_dims = MTL::Size(threadgroup_size * ((out.size() + BN - 1) / BN), 1, 1);
+
+  // Set the 2nd kernel
+  func_name = "col_reduce_looped";
+  kname = func_name;
+  large = intermediate.size() > INT32_MAX;
+  if (large) {
+    kname += "_large";
+  }
+  concatenate(kname, "_1_32_32_reduce_", op_name, type_to_name(intermediate));
+  kernel = get_reduce_kernel(
+      d,
+      kname,
+      func_name,
+      op_name,
+      intermediate.dtype(),
+      out_type,
+      large ? "int64_t" : "int",
+      1,
+      32,
+      32);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  compute_encoder.set_input_array(intermediate, 0);
+  compute_encoder.set_output_array(out, 1);
+  second_args.encode(compute_encoder);
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+void strided_reduce_general_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    const ReductionPlan& plan,
+    const std::vector<int>& axes,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s) {
+  // Prepare the arguments for the kernel
+  ColReduceArgs args(in, plan, axes);
+
+  // Small column
+  if (args.reduction_size * args.non_col_reductions < 32) {
+    return strided_reduce_small(in, out, op_name, args, compute_encoder, d, s);
+  }
+
+  // Long column but small row
+  if (args.reduction_stride < 32 &&
+      args.reduction_size * args.non_col_reductions >= 1024) {
+    return strided_reduce_longcolumn(
+        in, out, op_name, args, compute_encoder, d, s);
+  }
+
+  if (args.reduction_size * args.non_col_reductions > 256 &&
+      out.size() / 32 < 1024) {
+    return strided_reduce_2pass(in, out, op_name, args, compute_encoder, d, s);
+  }
+
+  return strided_reduce_looped(in, out, op_name, args, compute_encoder, d, s);
+}
+
+void Reduce::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  array in = inputs[0];
+
+  // Make sure no identity reductions trickle down here
+  assert(!axes_.empty());
+  assert(out.size() != in.size());
+
+  // Continue with reduction operation
+  // Minimum of 4 bytes since we use size 4 structs for all reduce
+  // and metal will complain o/w
+  size_t min_bytes = std::max(out.nbytes(), 4ul);
+  out.set_data(allocator::malloc(min_bytes));
+  std::string op_name;
+  switch (reduce_type_) {
+    case Reduce::And:
+      op_name = "and";
+      break;
+    case Reduce::Or:
+      op_name = "or";
+      break;
+    case Reduce::Sum:
+      op_name = "sum";
+      break;
+    case Reduce::Prod:
+      op_name = "prod";
+      break;
+    case Reduce::Min:
+      op_name = out.dtype() == bool_ ? "and" : "min";
+      break;
+    case Reduce::Max:
+      op_name = out.dtype() == bool_ ? "or" : "max";
+      break;
+  }
+
+  // Initialize output
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  // Reduce
+  if (in.size() > 0) {
+    ReductionPlan plan = get_reduction_plan(in, axes_);
+
+    // If it is a general reduce then copy the input to a contiguous array and
+    // recompute the plan.
+    //
+    // TODO: This can be avoided by making the output have the same strides as
+    //       input for the axes with stride smaller than the minimum reduction
+    //       stride.
+    if (plan.type == GeneralReduce) {
+      array in_copy = contiguous_copy_gpu(in, s);
+      d.add_temporary(in_copy, s.index);
+      in = in_copy;
+      plan = get_reduction_plan(in, axes_);
+    }
+
+    // Reducing over everything and the data is all there no broadcasting or
+    // slicing etc.
+    if (plan.type == ContiguousAllReduce) {
+      all_reduce_dispatch(in, out, op_name, compute_encoder, d, s);
+    }
+
+    // At least the last dimension is row contiguous and we are reducing over
+    // the last dim.
+    else if (
+        plan.type == ContiguousReduce || plan.type == GeneralContiguousReduce) {
+      row_reduce_general_dispatch(
+          in, out, op_name, plan, axes_, compute_encoder, d, s);
+    }
+
+    // At least the last two dimensions are contiguous and we are doing a
+    // strided reduce over these.
+    else if (
+        plan.type == ContiguousStridedReduce ||
+        plan.type == GeneralStridedReduce) {
+      strided_reduce_general_dispatch(
+          in, out, op_name, plan, axes_, compute_encoder, d, s);
+    }
+  }
+
+  // Nothing to reduce just initialize the output
+  else {
+    init_reduce(out, op_name, compute_encoder, d, s);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/reduce.h b/macos/third_party/mlx/mlx/backend/metal/reduce.h
new file mode 100644
index 0000000..a997d7e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/reduce.h
@@ -0,0 +1,41 @@
+// Copyright @ 2023 - 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/common/reduce.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/stream.h"
+
+namespace mlx::core {
+
+using metal::CommandEncoder;
+
+void all_reduce_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s);
+
+void row_reduce_general_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    const ReductionPlan& plan,
+    const std::vector<int>& axes,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s);
+
+void strided_reduce_general_dispatch(
+    const array& in,
+    array& out,
+    const std::string& op_name,
+    const ReductionPlan& plan,
+    const std::vector<int>& axes,
+    CommandEncoder& compute_encoder,
+    metal::Device& d,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/resident.cpp b/macos/third_party/mlx/mlx/backend/metal/resident.cpp
new file mode 100644
index 0000000..798824c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/resident.cpp
@@ -0,0 +1,100 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/metal/resident.h"
+
+namespace mlx::core::metal {
+
+ResidencySet::ResidencySet(MTL::Device* d) {
+  if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
+    return;
+  } else if (__builtin_available(macOS 15, iOS 18, *)) {
+    auto pool = new_scoped_memory_pool();
+    auto desc = MTL::ResidencySetDescriptor::alloc()->init();
+    NS::Error* error;
+    wired_set_ = d->newResidencySet(desc, &error);
+    desc->release();
+    if (!wired_set_) {
+      std::ostringstream msg;
+      msg << "[metal::Device] Unable to construct residency set.\n";
+      if (error) {
+        msg << error->localizedDescription()->utf8String() << "\n";
+      }
+      throw std::runtime_error(msg.str());
+    }
+    wired_set_->requestResidency();
+  }
+}
+
+void ResidencySet::insert(MTL::Allocation* buf) {
+  if (!wired_set_) {
+    return;
+  }
+  if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) {
+    wired_set_->addAllocation(buf);
+    wired_set_->commit();
+  } else {
+    unwired_set_.insert(buf);
+  }
+}
+
+void ResidencySet::erase(MTL::Allocation* buf) {
+  if (!wired_set_) {
+    return;
+  }
+  if (auto it = unwired_set_.find(buf); it != unwired_set_.end()) {
+    unwired_set_.erase(it);
+  } else {
+    wired_set_->removeAllocation(buf);
+    wired_set_->commit();
+  }
+}
+
+void ResidencySet::resize(size_t size) {
+  if (!wired_set_) {
+    return;
+  }
+
+  if (capacity_ == size) {
+    return;
+  }
+  capacity_ = size;
+
+  size_t current_size = wired_set_->allocatedSize();
+
+  if (current_size < size) {
+    auto pool = new_scoped_memory_pool();
+    // Add unwired allocations to the set
+    for (auto it = unwired_set_.begin(); it != unwired_set_.end();) {
+      auto buf_size = (*it)->allocatedSize();
+      if (current_size + buf_size > size) {
+        it++;
+      } else {
+        current_size += buf_size;
+        wired_set_->addAllocation(*it);
+        unwired_set_.erase(it++);
+      }
+    }
+    wired_set_->commit();
+  } else if (current_size > size) {
+    auto pool = new_scoped_memory_pool();
+    // Remove wired allocations until under capacity
+    auto allocations = wired_set_->allAllocations();
+    auto num_allocations = wired_set_->allocationCount();
+    for (int i = 0; i < num_allocations && current_size > size; ++i) {
+      auto buf = static_cast<const MTL::Allocation*>(allocations->object(i));
+      wired_set_->removeAllocation(buf);
+      current_size -= buf->allocatedSize();
+      unwired_set_.insert(buf);
+    }
+    wired_set_->commit();
+  }
+}
+
+ResidencySet::~ResidencySet() {
+  if (wired_set_) {
+    auto pool = new_scoped_memory_pool();
+    wired_set_->release();
+  }
+}
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/resident.h b/macos/third_party/mlx/mlx/backend/metal/resident.h
new file mode 100644
index 0000000..5db5582
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/resident.h
@@ -0,0 +1,32 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/device.h"
+
+namespace mlx::core::metal {
+
+class ResidencySet {
+ public:
+  ResidencySet(MTL::Device* d);
+  ~ResidencySet();
+
+  ResidencySet(const ResidencySet&) = delete;
+  ResidencySet& operator=(const ResidencySet&) = delete;
+
+  const MTL::ResidencySet* mtl_residency_set() {
+    return wired_set_;
+  }
+
+  void insert(MTL::Allocation* buf);
+  void erase(MTL::Allocation* buf);
+
+  void resize(size_t size);
+
+ private:
+  MTL::ResidencySet* wired_set_{nullptr};
+  std::unordered_set<const MTL::Allocation*> unwired_set_;
+  size_t capacity_{0};
+};
+
+} // namespace mlx::core::metal
diff --git a/macos/third_party/mlx/mlx/backend/metal/rope.cpp b/macos/third_party/mlx/mlx/backend/metal/rope.cpp
new file mode 100644
index 0000000..ca0a662
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/rope.cpp
@@ -0,0 +1,165 @@
+// Copyright © 2023-2024 Apple Inc.
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+
+namespace mlx::core::fast {
+
+constexpr int n_per_thread = 4;
+
+bool RoPE::use_fallback(Stream s) {
+  return s.device == Device::cpu;
+}
+
+void RoPE::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  assert(outputs.size() == 1);
+  auto& in = inputs[0];
+  auto& out = outputs[0];
+
+  auto& s = out.primitive().stream();
+  auto& d = metal::device(s.device);
+
+  int64_t strides[3];
+  int64_t out_strides[3];
+  bool donated = false;
+  int ndim = in.ndim();
+  int B = in.shape(0);
+  int T = in.shape(-2);
+  int D = in.shape(-1);
+  size_t mat_size = T * D;
+  bool large = in.data_size() > INT32_MAX || in.size() > INT32_MAX;
+
+  int dispatch_ndim = ndim;
+  while (in.shape(-dispatch_ndim) == 1 && dispatch_ndim > 3) {
+    dispatch_ndim--;
+  }
+
+  int N = 1;
+  for (int i = 1; i < (ndim - 2); ++i) {
+    N *= in.shape(i);
+  }
+
+  bool head_seq_transpose = false;
+
+  if (dims_ < D) {
+    donated = true;
+    auto ctype =
+        (in.flags().row_contiguous) ? CopyType::Vector : CopyType::General;
+    copy_gpu(in, out, ctype, s);
+    strides[0] = mat_size;
+    strides[1] = out.strides()[ndim - 2];
+    strides[2] = out.strides()[ndim - 1];
+  } else if (in.flags().row_contiguous) {
+    if (in.is_donatable()) {
+      donated = true;
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(allocator::malloc(out.nbytes()));
+    }
+    strides[0] = mat_size;
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
+  } else if (dispatch_ndim == 3) {
+    // Handle non-contiguous 3D inputs
+    out.set_data(allocator::malloc(out.nbytes()));
+    strides[0] = in.strides()[ndim - 3];
+    strides[1] = in.strides()[ndim - 2];
+    strides[2] = in.strides()[ndim - 1];
+  } else if (
+      ndim == 4 &&
+      // batch dim is regularly strided
+      in.strides()[0] == T * N * D &&
+      // sequence and head dimensions are transposed
+      in.strides()[1] == D && in.strides()[2] == N * D) {
+    head_seq_transpose = true;
+    out.set_data(allocator::malloc(out.nbytes()));
+    strides[0] = in.strides()[1];
+    strides[1] = in.strides()[2];
+    strides[2] = in.strides()[3];
+  } else {
+    // Copy non-contiguous > 3D inputs into the output and treat
+    // input as donated
+    donated = true;
+    copy_gpu(in, out, CopyType::General, s);
+    strides[0] = mat_size;
+    strides[1] = out.strides()[ndim - 2];
+    strides[2] = out.strides()[ndim - 1];
+  }
+  out_strides[0] = mat_size;
+  out_strides[1] = out.strides()[ndim - 2];
+  out_strides[2] = out.strides()[ndim - 1];
+
+  // Special case for inference (single time step, contiguous, one offset)
+  auto& offset = inputs[1];
+  bool single = in.flags().row_contiguous && T == 1 && offset.size() == 1;
+
+  bool with_freqs = inputs.size() == 3;
+  std::string kname;
+  concatenate(
+      kname,
+      "rope_",
+      single ? "single_" : "",
+      (with_freqs) ? "freqs_" : "",
+      large ? "large_" : "",
+      type_to_name(in));
+  std::string hash_name;
+  concatenate(
+      hash_name,
+      kname,
+      "_",
+      forward_ ? "" : "vjp_",
+      traditional_ ? "traditional_" : "",
+      head_seq_transpose ? "transpose" : "");
+  metal::MTLFCList func_consts = {
+      {&forward_, MTL::DataType::DataTypeBool, 1},
+      {&traditional_, MTL::DataType::DataTypeBool, 2},
+      {&head_seq_transpose, MTL::DataType::DataTypeBool, 3}};
+
+  auto kernel = d.get_kernel(kname, hash_name, func_consts);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  float base = std::log2(base_);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(donated ? out : in, 0);
+  compute_encoder.set_output_array(out, 1);
+
+  compute_encoder.set_input_array(offset, 2);
+  compute_encoder.set_bytes(scale_, 3);
+
+  MTL::Size group_dims;
+  MTL::Size grid_dims;
+  if (single) {
+    compute_encoder.set_bytes(out_strides, 1, 4);
+    uint32_t dim0 = dims_ / 2;
+    group_dims = get_block_dims(dim0, N, 1);
+    grid_dims = MTL::Size(dim0, N, 1);
+  } else {
+    compute_encoder.set_bytes(strides, 3, 4);
+    compute_encoder.set_bytes(out_strides, 3, 5);
+    int64_t offset_stride = 0;
+    if (offset.ndim() > 0) {
+      offset_stride = offset.strides()[0];
+    }
+    compute_encoder.set_bytes(offset_stride, 6);
+    compute_encoder.set_bytes(N, 7);
+    uint32_t dim0 = dims_ / 2;
+    uint32_t dim1 = T;
+    uint32_t dim2 = B * ((N + n_per_thread - 1) / n_per_thread);
+    group_dims = get_block_dims(dim0, dim1, dim2);
+    grid_dims = MTL::Size(dim0, dim1, dim2);
+  }
+
+  if (with_freqs) {
+    auto& freqs = inputs[2];
+    compute_encoder.set_input_array(freqs, 10);
+    auto freq_stride = freqs.strides()[0];
+    compute_encoder.set_bytes(freq_stride, 11);
+  } else {
+    compute_encoder.set_bytes(base, 10);
+  }
+  compute_encoder.dispatch_threads(grid_dims, group_dims);
+}
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp b/macos/third_party/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp
new file mode 100644
index 0000000..c990cda
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp
@@ -0,0 +1,794 @@
+// Copyright © 2024 Apple Inc.
+#include <sstream>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/steel/attn/params.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fast {
+
+namespace {
+
+void sdpa_full_self_attention_nax(
+    const Stream& s,
+    metal::Device& d,
+    const array& q,
+    const array& k,
+    const array& v,
+    const float scale,
+    array& o,
+    bool do_causal_,
+    const std::optional<array>& mask,
+    const std::optional<array>& sinks) {
+  using namespace mlx::steel;
+
+  int wm = 4;
+  int wn = 1;
+
+  int bd = q.shape(-1);
+  int bq = 64;
+  int bk = 32;
+
+  int B = q.shape(0);
+  int H = q.shape(1);
+  int D = q.shape(3);
+  int gqa_factor = q.shape(1) / k.shape(1);
+
+  int qL = q.shape(2);
+  int kL = k.shape(2);
+
+  const bool align_Q = (qL % bq) == 0;
+  const bool align_K = (kL % bk) == 0;
+  const bool has_mask = mask.has_value();
+  const bool do_causal = do_causal_;
+  const bool has_sinks = sinks.has_value();
+
+  metal::MTLFCList func_consts = {
+      {&align_Q, MTL::DataType::DataTypeBool, 200},
+      {&align_K, MTL::DataType::DataTypeBool, 201},
+      {&has_mask, MTL::DataType::DataTypeBool, 300},
+      {&do_causal, MTL::DataType::DataTypeBool, 301},
+      {&has_sinks, MTL::DataType::DataTypeBool, 302}};
+
+  std::string base_name;
+  concatenate(
+      base_name,
+      "steel_attention_",
+      type_to_name(q),
+      "_bq",
+      bq,
+      "_bk",
+      bk,
+      "_bd",
+      bd,
+      "_wm",
+      wm,
+      "_wn",
+      wn,
+      "_mask",
+      type_to_name(has_mask ? *mask : q));
+
+  std::string hash_name;
+  concatenate(
+      hash_name,
+      base_name,
+      "_align_Q_",
+      (align_Q ? 't' : 'n'),
+      "_align_K_",
+      (align_K ? 't' : 'n'),
+      "_has_mask_",
+      (has_mask ? 't' : 'n'),
+      "_do_causal_",
+      (do_causal ? 't' : 'n'),
+      "_has_sinks_",
+      (has_sinks ? 't' : 'n'));
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  auto kernel = get_steel_attention_nax_kernel(
+      d,
+      base_name,
+      hash_name,
+      func_consts,
+      q,
+      bq,
+      bk,
+      bd,
+      wm,
+      wn,
+      (has_mask ? *mask : q));
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  const int NQ = (qL + bq - 1) / bq;
+  const int NK = (kL + bk - 1) / bk;
+
+  const int NQ_aligned = qL / bq;
+  const int NK_aligned = kL / bk;
+
+  AttnParams params{
+      /* int B = */ B,
+      /* int H = */ H,
+      /* int D = */ D,
+
+      /* int qL = */ qL,
+      /* int kL = */ kL,
+
+      /* int gqa_factor = */ gqa_factor,
+      /* float scale = */ scale,
+
+      /* int NQ = */ NQ,
+      /* int NK = */ NK,
+
+      /* int NQ_aligned = */ NQ_aligned,
+      /* int NK_aligned = */ NK_aligned,
+
+      /* int qL_rem = */ (qL - NQ_aligned * bq),
+      /* int kL_rem = */ (kL - NK_aligned * bk),
+      /* int qL_off = */ (kL - qL),
+
+      /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)},
+      /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)},
+      /* int64_t V_strides[3] = */ {v.strides(0), v.strides(1), v.strides(2)},
+      /* int64_t O_strides[3] = */ {o.strides(0), o.strides(1), o.strides(2)}};
+
+  compute_encoder.set_input_array(q, 0);
+  compute_encoder.set_input_array(k, 1);
+  compute_encoder.set_input_array(v, 2);
+  compute_encoder.set_output_array(o, 3);
+  compute_encoder.set_bytes(params, 4);
+
+  if (has_mask) {
+    auto& m = *mask;
+
+    AttnMaskParams mask_params{/* int64_t M_strides[3] = */ {
+        m.strides(0), m.strides(1), m.strides(2)}};
+
+    compute_encoder.set_bytes(mask_params, 5);
+    compute_encoder.set_input_array(m, 6);
+  }
+  if (has_sinks) {
+    compute_encoder.set_input_array(*sinks, 7);
+  }
+
+  MTL::Size grid_dims = MTL::Size(NQ, H, B);
+  MTL::Size group_dims = MTL::Size(32, wm, wn);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void sdpa_full_self_attention_metal(
+    const Stream& s,
+    metal::Device& d,
+    const array& q,
+    const array& k,
+    const array& v,
+    const float scale,
+    array& o,
+    bool do_causal_,
+    const std::optional<array>& mask,
+    const std::optional<array>& sinks) {
+  if (metal::is_nax_available() && q.shape(3) != 80 &&
+      (env::enable_tf32() || q.dtype() != float32)) {
+    return sdpa_full_self_attention_nax(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& q = */ q,
+        /* const array& k = */ k,
+        /* const array& v = */ v,
+        /* const float scale = */ scale,
+        /* array& o = */ o,
+        /* bool do_causal_ = */ do_causal_,
+        /* const std::optional<array>& mask = */ mask,
+        /* const std::optional<array>& sinks = */ sinks);
+  }
+
+  using namespace mlx::steel;
+
+  int wm = 4;
+  int wn = 1;
+
+  int bd = q.shape(-1);
+  int bq = 32;
+  int bk = bd < 128 ? 32 : 16;
+
+  int B = q.shape(0);
+  int H = q.shape(1);
+  int D = q.shape(3);
+  int gqa_factor = q.shape(1) / k.shape(1);
+
+  int qL = q.shape(2);
+  int kL = k.shape(2);
+
+  const bool align_Q = (qL % bq) == 0;
+  const bool align_K = (kL % bk) == 0;
+  const bool has_mask = mask.has_value();
+  const bool do_causal = do_causal_;
+  const bool has_sinks = sinks.has_value();
+
+  metal::MTLFCList func_consts = {
+      {&align_Q, MTL::DataType::DataTypeBool, 200},
+      {&align_K, MTL::DataType::DataTypeBool, 201},
+      {&has_mask, MTL::DataType::DataTypeBool, 300},
+      {&do_causal, MTL::DataType::DataTypeBool, 301},
+      {&has_sinks, MTL::DataType::DataTypeBool, 302}};
+
+  std::string base_name;
+  concatenate(
+      base_name,
+      "steel_attention_",
+      type_to_name(q),
+      "_bq",
+      bq,
+      "_bk",
+      bk,
+      "_bd",
+      bd,
+      "_wm",
+      wm,
+      "_wn",
+      wn,
+      "_mask",
+      type_to_name(has_mask ? *mask : q));
+
+  std::string hash_name;
+  concatenate(
+      hash_name,
+      base_name,
+      "_align_Q_",
+      (align_Q ? 't' : 'n'),
+      "_align_K_",
+      (align_K ? 't' : 'n'),
+      "_has_mask_",
+      (has_mask ? 't' : 'n'),
+      "_do_causal_",
+      (do_causal ? 't' : 'n'),
+      "_has_sinks_",
+      (has_sinks ? 't' : 'n'));
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  auto kernel = get_steel_attention_kernel(
+      d,
+      base_name,
+      hash_name,
+      func_consts,
+      q,
+      bq,
+      bk,
+      bd,
+      wm,
+      wn,
+      (has_mask ? *mask : q));
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  const int NQ = (qL + bq - 1) / bq;
+  const int NK = (kL + bk - 1) / bk;
+
+  const int NQ_aligned = qL / bq;
+  const int NK_aligned = kL / bk;
+
+  AttnParams params{
+      /* int B = */ B,
+      /* int H = */ H,
+      /* int D = */ D,
+
+      /* int qL = */ qL,
+      /* int kL = */ kL,
+
+      /* int gqa_factor = */ gqa_factor,
+      /* float scale = */ scale,
+
+      /* int NQ = */ NQ,
+      /* int NK = */ NK,
+
+      /* int NQ_aligned = */ NQ_aligned,
+      /* int NK_aligned = */ NK_aligned,
+
+      /* int qL_rem = */ (qL - NQ_aligned * bq),
+      /* int kL_rem = */ (kL - NK_aligned * bk),
+      /* int qL_off = */ (kL - qL),
+
+      /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)},
+      /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)},
+      /* int64_t V_strides[3] = */ {v.strides(0), v.strides(1), v.strides(2)},
+      /* int64_t O_strides[3] = */ {o.strides(0), o.strides(1), o.strides(2)}};
+
+  compute_encoder.set_input_array(q, 0);
+  compute_encoder.set_input_array(k, 1);
+  compute_encoder.set_input_array(v, 2);
+  compute_encoder.set_output_array(o, 3);
+  compute_encoder.set_bytes(params, 4);
+
+  if (has_mask) {
+    auto& m = *mask;
+
+    AttnMaskParams mask_params{/* int64_t M_strides[3] = */ {
+        m.strides(0), m.strides(1), m.strides(2)}};
+
+    compute_encoder.set_bytes(mask_params, 5);
+    compute_encoder.set_input_array(m, 6);
+  }
+  if (has_sinks) {
+    compute_encoder.set_input_array(*sinks, 7);
+  }
+
+  MTL::Size grid_dims = MTL::Size(NQ, H, B);
+  MTL::Size group_dims = MTL::Size(32, wm, wn);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void sdpa_vector(
+    const Stream& s,
+    metal::Device& d,
+    const array& q,
+    const array& k,
+    const array& v,
+    array& out,
+    float scale,
+    bool do_causal,
+    const std::optional<array>& mask,
+    const std::optional<array>& sinks) {
+  // Set the kernel name
+  std::string kname;
+  kname.reserve(64);
+  kname += "sdpa_vector_";
+  kname += get_type_string(q.dtype());
+  kname += "_";
+  kname += std::to_string(q.shape(-1));
+  kname += "_";
+  kname += std::to_string(v.shape(-1));
+
+  // Compute the necessary sizes
+  int gqa_factor = q.shape(1) / k.shape(1);
+  int N = k.shape(2);
+  size_t k_head_stride = k.shape(1) == 1 ? k.strides(0) : k.strides(1);
+  size_t k_seq_stride = k.strides()[2];
+  size_t v_head_stride = v.shape(1) == 1 ? v.strides(0) : v.strides(1);
+  size_t v_seq_stride = v.strides()[2];
+
+  MTL::Size group_dims(1024, 1, 1);
+  MTL::Size grid_dims(q.shape(0) * q.shape(1), q.shape(2), 1);
+
+  bool has_mask = mask.has_value();
+  bool bool_mask = has_mask && (*mask).dtype() == bool_;
+  bool float_mask = has_mask && !bool_mask;
+  bool query_transposed = !q.flags().row_contiguous;
+  bool has_sinks = sinks.has_value();
+  metal::MTLFCList func_consts = {
+      {&has_mask, MTL::DataType::DataTypeBool, 20},
+      {&query_transposed, MTL::DataType::DataTypeBool, 21},
+      {&do_causal, MTL::DataType::DataTypeBool, 22},
+      {&bool_mask, MTL::DataType::DataTypeBool, 23},
+      {&float_mask, MTL::DataType::DataTypeBool, 24},
+      {&has_sinks, MTL::DataType::DataTypeBool, 25},
+  };
+  std::string hash_name = kname;
+  hash_name += has_mask ? (bool_mask ? "_boolmask" : "_floatmask") : "_nomask";
+  hash_name += query_transposed ? "_qt" : "_qnt";
+  hash_name += do_causal ? "_c" : "_nc";
+  hash_name += has_sinks ? "_sinks" : "_nosinks";
+
+  // Get the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname, hash_name, func_consts);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Set its arguments
+  compute_encoder.set_input_array(q, 0);
+  compute_encoder.set_input_array(k, 1);
+  compute_encoder.set_input_array(v, 2);
+  compute_encoder.set_output_array(out, 3);
+  compute_encoder.set_bytes(gqa_factor, 4);
+  compute_encoder.set_bytes(N, 5);
+  compute_encoder.set_bytes(k_head_stride, 6);
+  compute_encoder.set_bytes(k_seq_stride, 7);
+  compute_encoder.set_bytes(v_head_stride, 8);
+  compute_encoder.set_bytes(v_seq_stride, 9);
+
+  compute_encoder.set_bytes(scale, 10);
+  if (has_mask) {
+    auto& m = *mask;
+    compute_encoder.set_input_array(m, 11 + float_mask);
+    int32_t kv_seq_stride = m.shape(3) > 1 ? m.strides(3) : 0;
+    int32_t q_seq_stride = m.shape(2) > 1 ? m.strides(2) : 0;
+    int32_t head_stride =
+        m.shape(1) > 1 ? m.strides(1) : (m.shape(0) > 1 ? m.strides(0) : 0);
+    compute_encoder.set_bytes(kv_seq_stride, 13);
+    compute_encoder.set_bytes(q_seq_stride, 14);
+    compute_encoder.set_bytes(head_stride, 15);
+  }
+  if (has_sinks) {
+    compute_encoder.set_input_array(*sinks, 16);
+    compute_encoder.set_bytes(q.shape(1), 17);
+  }
+
+  // Launch
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void sdpa_vector_2pass(
+    const Stream& s,
+    metal::Device& d,
+    const array& q,
+    const array& k,
+    const array& v,
+    array& out,
+    float scale,
+    bool do_causal,
+    const std::optional<array>& mask,
+    const std::optional<array>& sinks) {
+  // Set the kernel name
+  std::string kname;
+  kname.reserve(64);
+  kname += "sdpa_vector_2pass_1_";
+  kname += get_type_string(q.dtype());
+  kname += "_";
+  kname += std::to_string(q.shape(-1));
+  kname += "_";
+  kname += std::to_string(v.shape(-1));
+
+  // Compute the necessary sizes
+  int gqa_factor = q.shape(1) / k.shape(1);
+  int n_simds = gqa_factor * q.shape(2);
+
+  char devc = d.get_architecture().back();
+  int N = k.shape(2);
+  int blocks;
+  if (devc == 's') {
+    blocks = 64;
+    if (N > 1024 && n_simds > 4) {
+      if (N <= 8192) {
+        blocks = 128;
+      } else if (N <= 32768) {
+        blocks = 256;
+      } else if (N <= 65536) {
+        blocks = 512;
+      } else {
+        blocks = 1024;
+      }
+    }
+  } else if (devc == 'd') {
+    blocks = 128;
+    if (n_simds <= 2 && N > 8192) {
+      blocks = 256;
+    } else if (n_simds >= 6) {
+      if (N >= 16384 && N < 65536) {
+        blocks = 512;
+      } else if (N >= 65536) {
+        blocks = 1024;
+      }
+    }
+  } else {
+    if (n_simds >= 4) {
+      blocks = 64;
+    } else {
+      blocks = 32;
+    }
+  }
+  size_t k_head_stride = k.shape(1) == 1 ? k.strides(0) : k.strides(1);
+  size_t k_seq_stride = k.strides()[2];
+  size_t v_head_stride = v.shape(1) == 1 ? v.strides(0) : v.strides(1);
+  size_t v_seq_stride = v.strides()[2];
+  MTL::Size group_dims(32, gqa_factor, q.shape(2));
+  MTL::Size grid_dims(k.shape(1), q.shape(0), blocks);
+
+  // Allocate the intermediates
+  Shape intermediate_shape;
+  intermediate_shape.reserve(out.ndim() + 1);
+  intermediate_shape.insert(
+      intermediate_shape.end(), out.shape().begin(), out.shape().end() - 1);
+  intermediate_shape.push_back(blocks);
+  intermediate_shape.push_back(out.shape().back());
+  array intermediate(intermediate_shape, q.dtype(), nullptr, {});
+  intermediate_shape.pop_back();
+  array sums(intermediate_shape, float32, nullptr, {});
+  array maxs(std::move(intermediate_shape), float32, nullptr, {});
+  intermediate.set_data(allocator::malloc(intermediate.nbytes()));
+  sums.set_data(allocator::malloc(sums.nbytes()));
+  maxs.set_data(allocator::malloc(maxs.nbytes()));
+  d.add_temporary(intermediate, s.index);
+  d.add_temporary(sums, s.index);
+  d.add_temporary(maxs, s.index);
+
+  bool has_mask = mask.has_value();
+  bool bool_mask = has_mask && (*mask).dtype() == bool_;
+  bool float_mask = has_mask && !bool_mask;
+  bool query_transposed = !q.flags().row_contiguous;
+  bool has_sinks = sinks.has_value();
+  metal::MTLFCList func_consts = {
+      {&has_mask, MTL::DataType::DataTypeBool, 20},
+      {&query_transposed, MTL::DataType::DataTypeBool, 21},
+      {&do_causal, MTL::DataType::DataTypeBool, 22},
+      {&bool_mask, MTL::DataType::DataTypeBool, 23},
+      {&float_mask, MTL::DataType::DataTypeBool, 24},
+      {&has_sinks, MTL::DataType::DataTypeBool, 25},
+      {&blocks, MTL::DataType::DataTypeInt, 26},
+  };
+  std::string hash_name = kname;
+  hash_name += has_mask ? (bool_mask ? "_boolmask" : "_floatmask") : "_nomask";
+  hash_name += query_transposed ? "_qt" : "_qnt";
+  hash_name += do_causal ? "_c" : "_nc";
+  hash_name += has_sinks ? "_sinks_" : "_nosinks_";
+  hash_name += std::to_string(blocks);
+
+  // Get the kernel
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname, hash_name, func_consts);
+  check_kernel_threadgroup_size(kernel, group_dims, hash_name);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Set its arguments
+  compute_encoder.set_input_array(q, 0);
+  compute_encoder.set_input_array(k, 1);
+  compute_encoder.set_input_array(v, 2);
+  compute_encoder.set_output_array(intermediate, 3);
+  compute_encoder.set_output_array(sums, 4);
+  compute_encoder.set_output_array(maxs, 5);
+  compute_encoder.set_bytes(N, 7);
+  compute_encoder.set_bytes(k_head_stride, 8);
+  compute_encoder.set_bytes(k_seq_stride, 9);
+  compute_encoder.set_bytes(v_head_stride, 10);
+  compute_encoder.set_bytes(v_seq_stride, 11);
+  compute_encoder.set_bytes(scale, 12);
+  if (has_mask) {
+    auto& m = *mask;
+    compute_encoder.set_input_array(m, 13 + float_mask);
+    int32_t kv_seq_stride = m.shape(3) > 1 ? m.strides(3) : 0;
+    int32_t q_seq_stride = m.shape(2) > 1 ? m.strides(2) : 0;
+    int32_t head_stride =
+        m.shape(1) > 1 ? m.strides(1) : (m.shape(0) > 1 ? m.strides(0) : 0);
+    compute_encoder.set_bytes(kv_seq_stride, 15);
+    compute_encoder.set_bytes(q_seq_stride, 16);
+    compute_encoder.set_bytes(head_stride, 17);
+  }
+  if (has_sinks) {
+    compute_encoder.set_input_array(*sinks, 18);
+  }
+
+  // Launch
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  // Final pass
+  kname.clear();
+  kname = "sdpa_vector_2pass_2_";
+  kname += get_type_string(q.dtype());
+  kname += "_";
+  kname += std::to_string(v.shape(-1));
+
+  // Get the kernel
+  kernel = d.get_kernel(kname);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Set its arguments
+  compute_encoder.set_input_array(intermediate, 0);
+  compute_encoder.set_input_array(sums, 1);
+  compute_encoder.set_input_array(maxs, 2);
+  compute_encoder.set_output_array(out, 3);
+  compute_encoder.set_bytes(blocks, 4);
+
+  // Launch
+  group_dims = MTL::Size(1024, 1, 1);
+  grid_dims = MTL::Size(q.shape(0) * q.shape(1), q.shape(2), 1);
+  check_kernel_threadgroup_size(kernel, group_dims, kname);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+} // namespace
+
+bool ScaledDotProductAttention::use_fallback(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool is_training,
+    bool output_logsumexp,
+    Stream s) {
+  if (is_training) {
+    // It's faster for training on Metal to use the unfused SDPA for both
+    // forward and backward.
+    return true;
+  }
+  if (output_logsumexp) {
+    return true;
+  }
+  if (s.device == Device::cpu) {
+    return true;
+  }
+
+  const int value_head_dim = v.shape(-1);
+  const int query_head_dim = q.shape(-1);
+  const int query_sequence_length = q.shape(2);
+  const int key_sequence_length = k.shape(2);
+  const int num_query_heads = q.shape(1);
+  const int num_kv_heads = k.shape(1);
+  const int gqa_factor = num_query_heads / num_kv_heads;
+
+  const bool sdpa_vector_supported_head_dim =
+      query_head_dim == value_head_dim &&
+      (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 ||
+       query_head_dim == 256);
+  const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim &&
+      (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128);
+
+  const bool sdpa_full_supported_mask = !has_mask || has_arr_mask ||
+      (query_sequence_length <= key_sequence_length && do_causal);
+
+  const bool supports_sdpa_full = query_sequence_length > 8 &&
+      sdpa_full_supported_mask && sdpa_full_supported_head_dim;
+
+  const bool supports_sdpa_vector = (query_sequence_length <= 8) &&
+      (query_sequence_length <= key_sequence_length) &&
+      sdpa_vector_supported_head_dim &&
+      (query_sequence_length * gqa_factor) <= 32;
+
+  return !(supports_sdpa_full || supports_sdpa_vector);
+}
+
+bool ScaledDotProductAttention::supports_bool_mask() {
+  return true;
+}
+
+void ScaledDotProductAttention::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto& q_pre = inputs[0];
+  auto& k_pre = inputs[1];
+  auto& v_pre = inputs[2];
+  auto& o = outputs[0];
+
+  std::vector<array> copies;
+
+  // Define some copy functions to ensure the layout of the inputs is as
+  // expected.
+  copies.reserve(inputs.size());
+  auto copy_unless = [&copies, &s](
+                         auto predicate, const array& arr) -> const array& {
+    if (!predicate(arr)) {
+      array arr_copy = contiguous_copy_gpu(arr, s);
+      copies.push_back(std::move(arr_copy));
+      return copies.back();
+    } else {
+      return arr;
+    }
+  };
+
+  // Checks that the headdim dimension has stride 1.
+  auto is_matrix_contiguous = [](const array& arr) {
+    return arr.strides(-1) == 1;
+  };
+
+  std::optional<array> sinks = std::nullopt;
+  if (has_sinks_) {
+    sinks = copy_unless(is_matrix_contiguous, inputs.back());
+  }
+  bool has_arr_mask = inputs.size() > (3 + has_sinks_);
+
+  // We are in vector mode ie single query
+  if (q_pre.shape(2) <= 8) {
+    auto q_copy_unless = [](const array& arr) {
+      if (arr.flags().row_contiguous) {
+        return true;
+      }
+      auto& strides = arr.strides();
+      auto& shape = arr.shape();
+      if (shape[0] == 1 || shape[1] == 1) {
+        // If either the batch or head dimension is a singleton, the other can
+        // be transposed with the sequence dimension
+        auto bidx = shape[0] == 1 ? 1 : 0;
+        return (strides[3] == 1) && (strides[2] == shape[3] * shape[bidx]) &&
+            (strides[bidx] == shape[3]);
+      }
+      return false;
+    };
+
+    auto kv_copy_unless = [](const array& arr) {
+      // keys and values should be copied if:
+      // - the last dimension is not contiguous
+      // - the batch and head dim are not contiguous
+      auto& strides = arr.strides();
+      auto& shape = arr.shape();
+      if (strides.back() != 1) {
+        return false;
+      }
+      if (shape[0] == 1 || shape[1] == 1) {
+        return true;
+      }
+      return (strides[0] == strides[1] * shape[1]);
+    };
+
+    const auto& q = copy_unless(q_copy_unless, q_pre);
+    const auto& k = copy_unless(kv_copy_unless, k_pre);
+    const auto& v = copy_unless(kv_copy_unless, v_pre);
+
+    // Donate the query if possible
+    if (q.is_donatable() && q.flags().row_contiguous && q.size() == o.size()) {
+      o.copy_shared_buffer(q);
+    } else {
+      o.set_data(allocator::malloc(o.nbytes()));
+    }
+
+    auto mask_copy_unless = [&q](const array& arr) {
+      auto& strides = arr.strides();
+      auto& shape = arr.shape();
+      return arr.flags().row_contiguous || q.shape(0) == 1 || q.shape(1) == 1 ||
+          (strides[0] == strides[1] * shape[1]);
+    };
+
+    auto mask = has_arr_mask
+        ? std::optional<array>{copy_unless(mask_copy_unless, inputs[3])}
+        : std::nullopt;
+
+    // We route to the 2 pass fused attention if
+    // - The device is large and the sequence length long
+    // - The sequence length is even longer and we have gqa
+    bool do_causal = do_causal_ && q.shape(2) > 1;
+    char devc = d.get_architecture().back();
+    if (((devc == 'd' || devc == 's') && k.shape(2) >= 1024) ||
+        (k.shape(1) < q.shape(1) && k.shape(2) >= 4096)) {
+      sdpa_vector_2pass(s, d, q, k, v, o, scale_, do_causal, mask, sinks);
+    } else {
+      sdpa_vector(s, d, q, k, v, o, scale_, do_causal, mask, sinks);
+    }
+  }
+
+  // Full attention mode
+  else {
+    const auto& q = copy_unless(is_matrix_contiguous, q_pre);
+    const auto& k = copy_unless(is_matrix_contiguous, k_pre);
+    const auto& v = copy_unless(is_matrix_contiguous, v_pre);
+
+    int64_t str_oD = 1;
+    int64_t str_oH = o.shape(3);
+    int64_t str_oL = o.shape(1) * str_oH;
+    int64_t str_oB = o.shape(2) * str_oL;
+    size_t data_size = o.shape(0) * str_oB;
+
+    array::Flags flags{
+        /* bool contiguous = */ 1,
+        /* bool row_contiguous = */ 0,
+        /* bool col_contiguous = */ 0,
+    };
+
+    o.set_data(
+        allocator::malloc(o.nbytes()),
+        data_size,
+        {str_oB, str_oH, str_oL, str_oD},
+        flags);
+
+    auto mask = has_arr_mask
+        ? std::optional<array>{copy_unless(is_matrix_contiguous, inputs[3])}
+        : std::nullopt;
+
+    sdpa_full_self_attention_metal(
+        s, d, q, k, v, scale_, o, do_causal_, mask, sinks);
+  }
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) {
+  return true;
+}
+
+void ScaledDotProductAttentionVJP::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error("NYI");
+}
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/backend/metal/scan.cpp b/macos/third_party/mlx/mlx/backend/metal/scan.cpp
new file mode 100644
index 0000000..b35de64
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/scan.cpp
@@ -0,0 +1,146 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+#include <sstream>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/scan.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void scan_gpu_inplace(
+    array in,
+    array& out,
+    Scan::ReduceType reduce_type,
+    int axis,
+    bool reverse,
+    bool inclusive,
+    const Stream& s) {
+  auto& d = metal::device(s.device);
+
+  bool contiguous = in.strides()[axis] == 1;
+
+  std::string reduce_type_str;
+  switch (reduce_type) {
+    case Scan::Sum:
+      reduce_type_str = "sum";
+      break;
+    case Scan::Prod:
+      reduce_type_str = "prod";
+      break;
+    case Scan::Max:
+      reduce_type_str = "max";
+      break;
+    case Scan::Min:
+      reduce_type_str = "min";
+      break;
+    case Scan::LogAddExp:
+      reduce_type_str = "logaddexp";
+      break;
+  }
+
+  std::string kname;
+  concatenate(
+      kname,
+      contiguous ? "contig_" : "strided_",
+      "scan_",
+      reverse ? "reverse_" : "",
+      inclusive ? "inclusive_" : "exclusive_",
+      reduce_type_str,
+      "_",
+      type_to_name(in),
+      "_",
+      type_to_name(out));
+
+  auto kernel =
+      get_scan_kernel(d, kname, reverse, inclusive, reduce_type_str, in, out);
+
+  if (contiguous) {
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(out, 1);
+    size_t size = in.shape(axis);
+    compute_encoder.set_bytes(size, 2);
+
+    // Compute the thread grid
+    int n_reads = (in.itemsize() <= 4) ? 4 : 2;
+    constexpr int simd_size = 32;
+    int elements_per_simd = n_reads * simd_size;
+    int thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    if (size <= n_reads * 1024) {
+      thread_group_size =
+          ((size + elements_per_simd - 1) / elements_per_simd) * simd_size;
+    } else if (size <= n_reads * 2048) {
+      thread_group_size =
+          ((size / 2 + elements_per_simd - 1) / elements_per_simd) * simd_size;
+    }
+    thread_group_size = std::min(
+        thread_group_size,
+        static_cast<int>(kernel->maxTotalThreadsPerThreadgroup()));
+    auto tmp_grid_dims =
+        get_2d_grid_dims(in.shape(), in.strides(), /*divisor=*/size);
+    MTL::Size grid_dims(
+        thread_group_size, tmp_grid_dims.width, tmp_grid_dims.height);
+    MTL::Size group_dims(thread_group_size, 1, 1);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(
+        in.data_shared_ptr() == nullptr ? out : in, 0);
+    compute_encoder.set_output_array(out, 1);
+    size_t size = in.shape(axis);
+    size_t stride = in.strides()[axis];
+    int bn = 32;
+    size_t stride_blocks = (stride + bn - 1) / bn;
+    compute_encoder.set_bytes(size, 2);
+    compute_encoder.set_bytes(stride, 3);
+    compute_encoder.set_bytes(stride_blocks, 4);
+
+    // Compute the thread grid
+    int n_reads = (in.itemsize() <= 4) ? 4 : 2;
+    int n_simdgroups = bn / n_reads;
+    int thread_group_size = n_simdgroups * 32;
+    auto tmp_grid_dims =
+        get_2d_grid_dims(in.shape(), in.strides(), /*divisor=*/size * stride);
+    if (tmp_grid_dims.width * stride_blocks <= UINT_MAX) {
+      tmp_grid_dims.width *= stride_blocks;
+    } else {
+      tmp_grid_dims.height *= stride_blocks;
+    }
+    MTL::Size grid_dims(
+        thread_group_size, tmp_grid_dims.width, tmp_grid_dims.height);
+    MTL::Size group_dims(thread_group_size, 1, 1);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  auto in = inputs[0];
+  if (in.flags().contiguous && in.strides()[axis_] != 0) {
+    if (in.is_donatable() && in.itemsize() == out.itemsize()) {
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(
+          allocator::malloc(in.data_size() * out.itemsize()),
+          in.data_size(),
+          in.strides(),
+          in.flags());
+    }
+  } else {
+    in = contiguous_copy_gpu(in, stream());
+    out.copy_shared_buffer(in);
+  }
+
+  scan_gpu_inplace(
+      in, out, reduce_type_, axis_, reverse_, inclusive_, stream());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/scan.h b/macos/third_party/mlx/mlx/backend/metal/scan.h
new file mode 100644
index 0000000..dab79c5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/scan.h
@@ -0,0 +1,17 @@
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void scan_gpu_inplace(
+    array in,
+    array& out,
+    Scan::ReduceType reduce_type,
+    int axis,
+    bool reverse,
+    bool inclusive,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/slicing.cpp b/macos/third_party/mlx/mlx/backend/metal/slicing.cpp
new file mode 100644
index 0000000..1e14c35
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/slicing.cpp
@@ -0,0 +1,99 @@
+// Copyright © 2024 Apple Inc.
+
+#include <numeric>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/gpu/slicing.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+
+namespace mlx::core {
+
+void concatenate_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    int axis,
+    const Stream& s) {
+  std::vector<int> sizes;
+  sizes.push_back(0);
+  for (auto& p : inputs) {
+    sizes.push_back(p.shape(axis));
+  }
+  std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto strides = out.strides();
+  auto flags = out.flags();
+  flags.row_contiguous = false;
+  flags.col_contiguous = false;
+  flags.contiguous = false;
+  auto& d = metal::device(s.device);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto concurrent_ctx = compute_encoder.start_concurrent();
+  for (int i = 0; i < inputs.size(); i++) {
+    array out_slice(inputs[i].shape(), out.dtype(), nullptr, {});
+    size_t data_offset = strides[axis] * sizes[i];
+    out_slice.copy_shared_buffer(
+        out, strides, flags, out_slice.size(), data_offset);
+    copy_gpu_inplace(inputs[i], out_slice, CopyType::GeneralGeneral, s);
+  }
+}
+
+array compute_dynamic_offset(
+    const array& indices,
+    const Strides& strides,
+    const std::vector<int>& axes,
+    const Stream& s) {
+  auto& d = metal::device(s.device);
+
+  // Kernel to compute offset here.
+  array offset({1}, int64, nullptr, {});
+  bool donate = indices.is_donatable() &&
+      (indices.data_size() * indices.itemsize()) >= offset.itemsize();
+  if (donate) {
+    offset.copy_shared_buffer(indices);
+  } else {
+    offset.set_data(allocator::malloc(offset.itemsize()));
+  }
+  d.add_temporary(offset, s.index);
+
+  auto dtype = indices.dtype();
+  std::string lib_name = "compute_dynamic_offset_" + type_to_name(dtype);
+  auto lib = d.get_library(lib_name, [dtype]() {
+    return fmt::format(
+        R"(
+        [[kernel]] void compute_dynamic_offset_{0}(
+            constant const {1}* indices [[buffer(0)]],
+            device int64_t& offset [[buffer(1)]],
+            constant const int64_t* strides [[buffer(2)]],
+            constant const int* axes [[buffer(3)]],
+            constant const int& n_axes [[buffer(4)]],
+            uint index [[thread_position_in_grid]]) {{
+          int64_t acc = 0;
+          for (int i = 0; i < n_axes; ++i) {{
+            acc += indices[i] * strides[axes[i]];
+          }}
+          offset = acc;
+        }})",
+        type_to_name(dtype),
+        get_type_string(dtype));
+  });
+  auto kernel = d.get_kernel(lib_name, lib);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(indices, 0);
+  compute_encoder.set_output_array(offset, 1);
+  compute_encoder.set_vector_bytes(strides, 2);
+  compute_encoder.set_vector_bytes(axes, 3);
+  int n_axes = axes.size();
+  compute_encoder.set_bytes(n_axes, 4);
+  MTL::Size dims = MTL::Size(1, 1, 1);
+  compute_encoder.dispatch_threads(dims, dims);
+  return offset;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/softmax.cpp b/macos/third_party/mlx/mlx/backend/metal/softmax.cpp
new file mode 100644
index 0000000..0b1a184
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/softmax.cpp
@@ -0,0 +1,87 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <algorithm>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+constexpr int SOFTMAX_LOOPED_LIMIT = 4096;
+
+void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  if (!issubdtype(out.dtype(), floating)) {
+    throw std::runtime_error(
+        "[softmax] Does not support non-floating point types.");
+  }
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  // Make sure that the last dimension is contiguous
+  auto set_output = [&s, &out](const array& x) {
+    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
+      if (x.is_donatable()) {
+        out.copy_shared_buffer(x);
+      } else {
+        out.set_data(
+            allocator::malloc(x.data_size() * x.itemsize()),
+            x.data_size(),
+            x.strides(),
+            x.flags());
+      }
+      return x;
+    } else {
+      array x_copy = contiguous_copy_gpu(x, s);
+      out.copy_shared_buffer(x_copy);
+      return x_copy;
+    }
+  };
+
+  const array in = set_output(inputs[0]);
+
+  int axis_size = in.shape().back();
+  int n_rows = in.data_size() / axis_size;
+
+  const int simd_size = 32;
+  const int n_reads = SOFTMAX_N_READS;
+  const int looped_limit = SOFTMAX_LOOPED_LIMIT;
+
+  std::string kernel_name = (axis_size > looped_limit) ? "looped_" : "block_";
+  kernel_name += "softmax_";
+  if (in.dtype() != float32 && precise_) {
+    kernel_name += "precise_";
+  }
+  kernel_name += type_to_name(out);
+
+  auto kernel = get_softmax_kernel(d, kernel_name, precise_, out);
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  {
+    MTL::Size grid_dims, group_dims;
+    if (axis_size <= looped_limit) {
+      size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
+      size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
+      size_t threadgroup_size = simd_size * simds_needed;
+      assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    } else {
+      size_t threadgroup_size = kernel->maxTotalThreadsPerThreadgroup();
+      size_t n_threads = n_rows * threadgroup_size;
+      grid_dims = MTL::Size(n_threads, 1, 1);
+      group_dims = MTL::Size(threadgroup_size, 1, 1);
+    }
+
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(axis_size, 2);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/sort.cpp b/macos/third_party/mlx/mlx/backend/metal/sort.cpp
new file mode 100644
index 0000000..3c84022
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/sort.cpp
@@ -0,0 +1,368 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <algorithm>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+void single_block_sort(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    array& out,
+    int axis,
+    int bn,
+    int tn,
+    bool argsort) {
+  // Prepare shapes
+  int n_rows = in.size() / in.shape(axis);
+
+  auto in_nc_str = in.strides();
+  in_nc_str.erase(in_nc_str.begin() + axis);
+
+  auto out_nc_str = out.strides();
+  out_nc_str.erase(out_nc_str.begin() + axis);
+
+  auto nc_shape = in.shape();
+  nc_shape.erase(nc_shape.begin() + axis);
+
+  int nc_dim = nc_shape.size();
+
+  int size_sorted_axis = in.shape(axis);
+  int in_stride_sorted_axis = in.strides()[axis];
+  int out_stride_sorted_axis = out.strides()[axis];
+
+  // We can only use the contiguous kernel if the sorted axis
+  // has the largest or smallest stride.
+  // We also need the input to be contiguous
+  bool contiguous = in.flags().contiguous;
+  auto check_strides = [](array x, int sort_stride) {
+    int min_stride = *std::min_element(x.strides().begin(), x.strides().end());
+    int max_stride = *std::max_element(x.strides().begin(), x.strides().end());
+    return sort_stride == min_stride || sort_stride == max_stride;
+  };
+  contiguous &= check_strides(in, in_stride_sorted_axis);
+  contiguous &= check_strides(out, out_stride_sorted_axis);
+
+  // Prepare kernel name
+  std::ostringstream kname;
+  kname << (contiguous ? "c" : "nc");
+  if (argsort) {
+    kname << "arg";
+  }
+
+  kname << "_block_sort_" << type_to_name(in) << "_" << type_to_name(out)
+        << "_bn" << bn << "_tn" << tn;
+  auto kernel = get_sort_kernel(d, kname.str(), in, out, bn, tn);
+
+  // Prepare command encoder
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Set inputs
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  compute_encoder.set_bytes(size_sorted_axis, 2);
+  compute_encoder.set_bytes(in_stride_sorted_axis, 3);
+  compute_encoder.set_bytes(out_stride_sorted_axis, 4);
+
+  if (contiguous) {
+    int in_stride_segment_axis = INT32_MAX;
+    int out_stride_segment_axis = INT32_MAX;
+    for (int i = 0; i < in_nc_str.size(); i++) {
+      if (nc_shape[i] == 1) {
+        continue;
+      }
+      if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) {
+        throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
+      }
+      in_stride_segment_axis =
+          std::min(in_stride_segment_axis, static_cast<int>(in_nc_str[i]));
+      out_stride_segment_axis =
+          std::min(out_stride_segment_axis, static_cast<int>(out_nc_str[i]));
+    }
+    compute_encoder.set_bytes(in_stride_segment_axis, 5);
+    compute_encoder.set_bytes(out_stride_segment_axis, 6);
+  } else {
+    compute_encoder.set_bytes(nc_dim, 5);
+    if (nc_shape.empty()) {
+      int shape = 0;
+      int64_t stride = 0;
+      compute_encoder.set_bytes(shape, 6);
+      compute_encoder.set_bytes(stride, 7);
+      compute_encoder.set_bytes(stride, 8);
+    } else {
+      compute_encoder.set_vector_bytes(nc_shape, 6);
+      compute_encoder.set_vector_bytes(in_nc_str, 7);
+      compute_encoder.set_vector_bytes(out_nc_str, 8);
+    }
+  }
+
+  MTL::Size group_dims = MTL::Size(bn, 1, 1);
+  MTL::Size grid_dims = MTL::Size(1, n_rows, 1);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void multi_block_sort(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    array& out,
+    int axis,
+    int bn,
+    int tn,
+    int n_blocks,
+    bool argsort) {
+  // Prepare shapes
+  int n_rows = in.size() / in.shape(axis);
+
+  auto nc_str = in.strides();
+  nc_str.erase(nc_str.begin() + axis);
+
+  auto nc_shape = in.shape();
+  nc_shape.erase(nc_shape.begin() + axis);
+
+  int nc_dim = nc_shape.size();
+
+  if (nc_dim == 0) {
+    nc_shape = {0};
+    nc_str = {1};
+  }
+
+  int size_sorted_axis = in.shape(axis);
+  int stride_sorted_axis = in.strides()[axis];
+
+  // Make temporary copies
+  array dev_vals_0({n_rows, size_sorted_axis}, in.dtype(), nullptr, {});
+  array dev_vals_1({n_rows, size_sorted_axis}, in.dtype(), nullptr, {});
+
+  array dev_idxs_0({n_rows, size_sorted_axis}, uint32, nullptr, {});
+  array dev_idxs_1({n_rows, size_sorted_axis}, uint32, nullptr, {});
+
+  array block_partitions({n_rows, n_blocks + 1}, uint32, nullptr, {});
+
+  // Do allocations
+  dev_vals_0.set_data(allocator::malloc(dev_vals_0.nbytes()));
+  dev_vals_1.set_data(allocator::malloc(dev_vals_1.nbytes()));
+  dev_idxs_0.set_data(allocator::malloc(dev_idxs_0.nbytes()));
+  dev_idxs_1.set_data(allocator::malloc(dev_idxs_1.nbytes()));
+  block_partitions.set_data(allocator::malloc(block_partitions.nbytes()));
+
+  std::vector<array> copies = {
+      dev_vals_0, dev_vals_1, dev_idxs_0, dev_idxs_1, block_partitions};
+
+  // Prepare command encoder
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  // Do blockwise sort
+  {
+    std::ostringstream kname;
+    kname << "sort_mbsort_" << type_to_name(dev_vals_0) << "_"
+          << type_to_name(dev_idxs_0) << "_bn" << bn << "_tn" << tn;
+    auto kernel =
+        get_mb_sort_kernel(d, kname.str(), dev_vals_0, dev_idxs_0, bn, tn);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(dev_vals_0, 1);
+    compute_encoder.set_output_array(dev_idxs_0, 2);
+    compute_encoder.set_bytes(size_sorted_axis, 3);
+    compute_encoder.set_bytes(stride_sorted_axis, 4);
+    compute_encoder.set_bytes(nc_dim, 5);
+    compute_encoder.set_vector_bytes(nc_shape, 6);
+    compute_encoder.set_vector_bytes(nc_str, 7);
+
+    MTL::Size group_dims = MTL::Size(bn, 1, 1);
+    MTL::Size grid_dims = MTL::Size(n_blocks, n_rows, 1);
+
+    compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+  }
+
+  // Do merges
+  bool ping = false;
+  array dev_vals_in = dev_vals_0;
+  array dev_idxs_in = dev_idxs_0;
+  array dev_vals_out = dev_vals_1;
+  array dev_idxs_out = dev_idxs_1;
+
+  int n_thr_per_group = (n_blocks + 1) < 1024 ? (n_blocks + 1) : 1024;
+
+  for (int merge_tiles = 2; (merge_tiles / 2) < n_blocks; merge_tiles *= 2) {
+    dev_vals_in = ping ? dev_vals_1 : dev_vals_0;
+    dev_idxs_in = ping ? dev_idxs_1 : dev_idxs_0;
+    dev_vals_out = ping ? dev_vals_0 : dev_vals_1;
+    dev_idxs_out = ping ? dev_idxs_0 : dev_idxs_1;
+    ping = !ping;
+
+    // Do partition
+    {
+      std::ostringstream kname;
+      kname << "partition_mbsort_" << type_to_name(dev_vals_in) << "_"
+            << type_to_name(dev_idxs_in) << "_bn" << bn << "_tn" << tn;
+
+      auto kernel =
+          get_mb_sort_kernel(d, kname.str(), dev_vals_0, dev_idxs_0, bn, tn);
+      compute_encoder.set_compute_pipeline_state(kernel);
+
+      compute_encoder.set_output_array(block_partitions, 0);
+      compute_encoder.set_input_array(dev_vals_in, 1);
+      compute_encoder.set_input_array(dev_idxs_in, 2);
+      compute_encoder.set_bytes(size_sorted_axis, 3);
+      compute_encoder.set_bytes(merge_tiles, 4);
+      compute_encoder.set_bytes(n_blocks, 5);
+
+      MTL::Size group_dims = MTL::Size(n_thr_per_group, 1, 1);
+      MTL::Size grid_dims = MTL::Size(1, n_rows, 1);
+
+      compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+    }
+
+    // Do merge
+    {
+      std::ostringstream kname;
+      kname << "merge_mbsort_" << type_to_name(dev_vals_in) << "_"
+            << type_to_name(dev_idxs_in) << "_bn" << bn << "_tn" << tn;
+
+      auto kernel =
+          get_mb_sort_kernel(d, kname.str(), dev_vals_0, dev_idxs_0, bn, tn);
+      compute_encoder.set_compute_pipeline_state(kernel);
+
+      compute_encoder.set_input_array(block_partitions, 0);
+      compute_encoder.set_input_array(dev_vals_in, 1);
+      compute_encoder.set_input_array(dev_idxs_in, 2);
+      compute_encoder.set_output_array(dev_vals_out, 3);
+      compute_encoder.set_output_array(dev_idxs_out, 4);
+      compute_encoder.set_bytes(size_sorted_axis, 5);
+      compute_encoder.set_bytes(merge_tiles, 6);
+      compute_encoder.set_bytes(n_blocks, 7);
+
+      MTL::Size group_dims = MTL::Size(bn, 1, 1);
+      MTL::Size grid_dims = MTL::Size(n_blocks, n_rows, 1);
+
+      compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+    }
+  }
+
+  // Copy outputs with appropriate strides
+  auto strides = out.strides();
+  for (int ax = axis + 1; ax < strides.size(); ax++) {
+    strides[ax] *= out.shape(axis);
+  }
+  strides[axis] = 1;
+  copy_gpu_inplace(
+      (argsort) ? dev_idxs_out : dev_vals_out,
+      out,
+      out.shape(),
+      strides,
+      out.strides(),
+      0,
+      0,
+      (axis == in.ndim() - 1) ? CopyType::Vector : CopyType::General,
+      s);
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+void gpu_merge_sort(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    array& out,
+    int axis_,
+    bool argsort) {
+  // Get size info
+  int axis = axis_ < 0 ? axis_ + in.ndim() : axis_;
+  int size_sorted_axis = in.shape(axis);
+
+  // Get kernel size
+  int tn = 4;
+  int potential_bn = (size_sorted_axis + tn - 1) / tn;
+
+  int bn;
+  if (potential_bn > 256) {
+    bn = 512;
+  } else if (potential_bn > 128) {
+    bn = 256;
+  } else if (potential_bn > 64) {
+    bn = 128;
+  } else if (potential_bn > 32) {
+    bn = 64;
+  } else {
+    bn = 32;
+  }
+
+  if (bn == 512 && size_of(in.dtype()) > 4) {
+    bn = 256;
+  }
+
+  int n_per_block = bn * tn;
+  int n_blocks = (size_sorted_axis + n_per_block - 1) / n_per_block;
+
+  if (n_blocks > 1) {
+    return multi_block_sort(s, d, in, out, axis, bn, tn, n_blocks, argsort);
+  } else {
+    return single_block_sort(s, d, in, out, axis, bn, tn, argsort);
+  }
+}
+
+} // namespace
+
+void ArgSort::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, true);
+}
+
+void Sort::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, false);
+}
+
+void ArgPartition::eval_gpu(const std::vector<array>& inputs, array& out) {
+  // We direct arg partition to sort for now
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, true);
+}
+
+void Partition::eval_gpu(const std::vector<array>& inputs, array& out) {
+  // We direct partition to sort for now
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, false);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/ternary.cpp b/macos/third_party/mlx/mlx/backend/metal/ternary.cpp
new file mode 100644
index 0000000..252815a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/ternary.cpp
@@ -0,0 +1,160 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/ternary.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void ternary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  assert(inputs.size() == 3);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& c = inputs[2];
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+
+  if (out.size() == 0) {
+    return;
+  }
+
+  // Try to collapse contiguous dims
+  auto maybe_collapse = [topt, &a, &b, &c, &out]() {
+    if (topt == TernaryOpType::General) {
+      auto [shape, strides] = collapse_contiguous_dims(a, b, c, out);
+      return std::make_tuple(
+          shape, strides[0], strides[1], strides[2], strides[3]);
+    } else {
+      Strides e;
+      return std::make_tuple(Shape{}, e, e, e, e);
+    }
+  };
+  auto [shape, strides_a, strides_b, strides_c, strides_out] = maybe_collapse();
+
+  bool large;
+  auto ndim = shape.size();
+  int work_per_thread;
+  if (topt == TernaryOpType::General) {
+    large = a.data_size() > INT32_MAX || b.data_size() > INT32_MAX ||
+        c.data_size() > INT32_MAX || out.size() > INT32_MAX;
+    work_per_thread = large ? 4 : 2;
+  } else {
+    large = out.data_size() > INT32_MAX;
+    work_per_thread = get_work_per_thread(b.dtype(), out.data_size());
+  }
+  std::string kernel_name;
+  if (topt == TernaryOpType::General) {
+    kernel_name = "g";
+    if (shape.size() <= 3) {
+      kernel_name += std::to_string(shape.size());
+    } else if (work_per_thread > 1) {
+      concatenate(kernel_name, "n", std::to_string(work_per_thread));
+    }
+    if (large) {
+      kernel_name += "large";
+    }
+  } else {
+    if (topt == TernaryOpType::VectorScalarVector) {
+      kernel_name = "sv";
+    } else if (topt == TernaryOpType::VectorVectorScalar) {
+      kernel_name = "vs";
+    } else {
+      kernel_name = "v";
+    }
+    if (large) {
+      kernel_name += "2";
+    } else if (work_per_thread > 1) {
+      kernel_name += "n";
+    }
+  }
+  concatenate(kernel_name, "_", op, type_to_name(b));
+
+  auto& d = metal::device(s.device);
+
+  auto kernel = get_ternary_kernel(d, kernel_name, out.dtype(), op);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_input_array(c, 2);
+  compute_encoder.set_output_array(out, 3);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (topt == TernaryOpType::General) {
+    // Launch up to 3D grid of threads
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = out.size() / (dim0 * dim1);
+
+    if (ndim > 3) {
+      compute_encoder.set_vector_bytes(shape, 4);
+      compute_encoder.set_vector_bytes(strides_a, 5);
+      compute_encoder.set_vector_bytes(strides_b, 6);
+      compute_encoder.set_vector_bytes(strides_c, 7);
+
+      compute_encoder.set_bytes(ndim, 8);
+      dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    } else {
+      // The shape is implicit in the grid for <= 3D
+      compute_encoder.set_vector_bytes(strides_a, 4);
+      compute_encoder.set_vector_bytes(strides_b, 5);
+      compute_encoder.set_vector_bytes(strides_c, 6);
+    }
+
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::ternary] Must use 1024 sized block");
+    }
+    MTL::Size group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    // Launch a 1D or 2D grid of threads
+    size_t nthreads = ceildiv(out.data_size(), work_per_thread);
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    MTL::Size grid_dims;
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(out.data_size(), 4);
+      grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+    } else {
+      compute_encoder.set_bytes<int>(out.data_size(), 4);
+      grid_dims = MTL::Size(nthreads, 1, 1);
+    }
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void ternary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& c = inputs[2];
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+  set_ternary_op_output_data(a, b, c, out, topt);
+  ternary_op_gpu_inplace(inputs, out, op, s);
+}
+
+void ternary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op) {
+  auto& s = out.primitive().stream();
+  ternary_op_gpu(inputs, out, op, s);
+}
+
+void Select::eval_gpu(const std::vector<array>& inputs, array& out) {
+  ternary_op_gpu(inputs, out, name());
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/ternary.h b/macos/third_party/mlx/mlx/backend/metal/ternary.h
new file mode 100644
index 0000000..91c6fbb
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/ternary.h
@@ -0,0 +1,21 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+void ternary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+void ternary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/unary.cpp b/macos/third_party/mlx/mlx/backend/metal/unary.cpp
new file mode 100644
index 0000000..833b23f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/unary.cpp
@@ -0,0 +1,161 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/unary.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+#define UNARY_GPU(func)                                               \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    unary_op_gpu(inputs, out, name());                                \
+  }
+
+namespace mlx::core {
+
+void unary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  auto& in = inputs[0];
+  bool contig = in.flags().contiguous;
+  if (in.size() == 0) {
+    return;
+  }
+
+  auto& d = metal::device(s.device);
+
+  auto maybe_collapse = [contig, &in]() {
+    if (!contig) {
+      return collapse_contiguous_dims(in);
+    } else {
+      return std::make_pair(Shape{}, Strides{});
+    }
+  };
+  auto [shape, strides] = maybe_collapse();
+  int ndim = shape.size();
+  bool large;
+  if (!contig) {
+    large = in.data_size() > INT32_MAX || out.size() > INT32_MAX;
+  } else {
+    large = in.data_size() > UINT32_MAX;
+  }
+  int work_per_thread;
+  std::string kernel_name;
+  if (contig) {
+    work_per_thread = get_work_per_thread(in.dtype(), in.data_size());
+    kernel_name = (large ? "v2" : (work_per_thread > 1 ? "vn" : "v"));
+  } else {
+    work_per_thread = large ? 4 : 1;
+    kernel_name = "gn" + std::to_string(work_per_thread);
+    if (large) {
+      kernel_name += "large";
+    }
+  }
+  concatenate(kernel_name, "_", op, type_to_name(in), type_to_name(out));
+  auto kernel = get_unary_kernel(d, kernel_name, in.dtype(), out.dtype(), op);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  if (!contig) {
+    // Launch up to 3D grid of threads
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = out.size() / (dim0 * dim1);
+    compute_encoder.set_vector_bytes(shape, 2);
+    compute_encoder.set_vector_bytes(strides, 3);
+    compute_encoder.set_bytes(ndim, 4);
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::unary] Must use 1024 sized block");
+    }
+    dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    auto group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    size_t nthreads = ceildiv(in.data_size(), work_per_thread);
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    MTL::Size grid_dims;
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(in.data_size(), 2);
+      grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+    } else {
+      compute_encoder.set_bytes<int>(in.data_size(), 2);
+      grid_dims = MTL::Size(nthreads, 1, 1);
+    }
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void unary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  set_unary_output_data(inputs[0], out);
+  unary_op_gpu_inplace(inputs, out, op, s);
+}
+
+void unary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op) {
+  auto& s = out.primitive().stream();
+  unary_op_gpu(inputs, out, op, s);
+}
+
+UNARY_GPU(Abs)
+UNARY_GPU(ArcCos)
+UNARY_GPU(ArcCosh)
+UNARY_GPU(ArcSin)
+UNARY_GPU(ArcSinh)
+UNARY_GPU(ArcTan)
+UNARY_GPU(ArcTanh)
+UNARY_GPU(BitwiseInvert)
+UNARY_GPU(Conjugate)
+UNARY_GPU(Cos)
+UNARY_GPU(Cosh)
+UNARY_GPU(Erf)
+UNARY_GPU(ErfInv)
+UNARY_GPU(Exp)
+UNARY_GPU(Expm1)
+UNARY_GPU(Imag)
+UNARY_GPU(Log1p)
+UNARY_GPU(LogicalNot)
+UNARY_GPU(Floor)
+UNARY_GPU(Ceil)
+UNARY_GPU(Negative)
+UNARY_GPU(Real)
+UNARY_GPU(Sigmoid)
+UNARY_GPU(Sign)
+UNARY_GPU(Sin)
+UNARY_GPU(Sinh)
+UNARY_GPU(Square)
+UNARY_GPU(Sqrt)
+UNARY_GPU(Tan)
+UNARY_GPU(Tanh)
+
+void Log::eval_gpu(const std::vector<array>& inputs, array& out) {
+  unary_op_gpu(inputs, out, name());
+}
+
+void Round::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_op_gpu(inputs, out, name());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/unary.h b/macos/third_party/mlx/mlx/backend/metal/unary.h
new file mode 100644
index 0000000..1d6ecf0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/unary.h
@@ -0,0 +1,21 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+void unary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+void unary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/utils.cpp b/macos/third_party/mlx/mlx/backend/metal/utils.cpp
new file mode 100644
index 0000000..9785018
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/utils.cpp
@@ -0,0 +1,77 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/metal/utils.h"
+#include "mlx/backend/common/utils.h"
+
+namespace mlx::core {
+
+std::string type_to_name(const Dtype& t) {
+  std::string tname;
+  switch (t) {
+    case bool_:
+      tname = "bool_";
+      break;
+    case uint8:
+      tname = "uint8";
+      break;
+    case uint16:
+      tname = "uint16";
+      break;
+    case uint32:
+      tname = "uint32";
+      break;
+    case uint64:
+      tname = "uint64";
+      break;
+    case int8:
+      tname = "int8";
+      break;
+    case int16:
+      tname = "int16";
+      break;
+    case int32:
+      tname = "int32";
+      break;
+    case int64:
+      tname = "int64";
+      break;
+    case float16:
+      tname = "float16";
+      break;
+    case float32:
+      tname = "float32";
+      break;
+    case float64:
+      tname = "double";
+      break;
+    case bfloat16:
+      tname = "bfloat16";
+      break;
+    case complex64:
+      tname = "complex64";
+      break;
+  }
+  return tname;
+}
+
+std::string type_to_name(const array& a) {
+  return type_to_name(a.dtype());
+}
+
+MTL::Size get_block_dims(int dim0, int dim1, int dim2, int pow2) {
+  Dims dims = get_block_dims_common(dim0, dim1, dim2, pow2);
+  return MTL::Size(std::get<0>(dims), std::get<1>(dims), std::get<2>(dims));
+}
+
+MTL::Size get_2d_grid_dims(const Shape& shape, const Strides& strides) {
+  Dims dims = get_2d_grid_dims_common(shape, strides);
+  return MTL::Size(std::get<0>(dims), std::get<1>(dims), std::get<2>(dims));
+}
+
+MTL::Size
+get_2d_grid_dims(const Shape& shape, const Strides& strides, size_t divisor) {
+  Dims dims = get_2d_grid_dims_common(shape, strides, divisor);
+  return MTL::Size(std::get<0>(dims), std::get<1>(dims), std::get<2>(dims));
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/metal/utils.h b/macos/third_party/mlx/mlx/backend/metal/utils.h
new file mode 100644
index 0000000..c4cef8c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/metal/utils.h
@@ -0,0 +1,99 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <type_traits>
+
+#include "mlx/array.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+MLX_API std::string type_to_name(const Dtype& t);
+MLX_API std::string type_to_name(const array& a);
+
+// Compute the grid and block dimensions, check backend/common/utils.h for docs.
+MTL::Size get_block_dims(int dim0, int dim1, int dim2, int pow2 = 10);
+MTL::Size get_2d_grid_dims(const Shape& shape, const Strides& strides);
+MTL::Size
+get_2d_grid_dims(const Shape& shape, const Strides& strides, size_t divisor);
+
+inline NS::String* make_string(std::ostringstream& os) {
+  std::string string = os.str();
+  return NS::String::string(string.c_str(), NS::UTF8StringEncoding);
+}
+
+inline void debug_set_stream_queue_label(MTL::CommandQueue* queue, int index) {
+#ifdef MLX_METAL_DEBUG
+  std::ostringstream label;
+  label << "Stream " << index;
+  queue->setLabel(make_string(label));
+#endif
+}
+
+inline void debug_set_primitive_buffer_label(
+    MTL::CommandBuffer* command_buffer,
+    Primitive& primitive) {
+#ifdef MLX_METAL_DEBUG
+  std::ostringstream label;
+  if (auto cbuf_label = command_buffer->label(); cbuf_label) {
+    label << cbuf_label->utf8String();
+  }
+  label << primitive.name();
+  command_buffer->setLabel(make_string(label));
+#endif
+}
+
+template <typename T>
+constexpr bool is_numeric_except_char = std::is_arithmetic_v<T> &&
+    !std::is_same_v<T, char> && !std::is_same_v<T, signed char> &&
+    !std::is_same_v<T, unsigned char> && !std::is_same_v<T, wchar_t>;
+
+template <typename T>
+void concatenate(std::string& acc, T first) {
+  if constexpr (is_numeric_except_char<T>) {
+    acc += std::to_string(first);
+  } else {
+    acc += first;
+  }
+}
+
+template <typename T, typename... Args>
+void concatenate(std::string& acc, T first, Args... args) {
+  if constexpr (is_numeric_except_char<T>) {
+    acc += std::to_string(first);
+  } else {
+    acc += first;
+  }
+  concatenate(acc, args...);
+}
+
+inline int get_work_per_thread(Dtype dtype) {
+  return std::max(1, 8 / dtype.size());
+}
+inline int get_work_per_thread(Dtype dtype, size_t size) {
+  constexpr size_t wpt_threshold = 1 << 16;
+  return size < wpt_threshold ? 1 : std::max(1, 8 / dtype.size());
+}
+
+inline size_t ceildiv(size_t n, size_t m) {
+  return (n + m - 1) / m;
+}
+
+inline void check_kernel_threadgroup_size(
+    const MTL::ComputePipelineState* kernel,
+    MTL::Size group_dims,
+    const std::string& name) {
+  auto max_size = kernel->maxTotalThreadsPerThreadgroup();
+  auto requested_size = group_dims.width * group_dims.height * group_dims.depth;
+
+  if (max_size < requested_size) {
+    std::ostringstream msg;
+    msg << "Maximum threads per threadgroup is " << max_size
+        << " but requested " << requested_size << " for kernel " << name << ".";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/no_cpu/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/no_cpu/CMakeLists.txt
new file mode 100644
index 0000000..c81a258
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_cpu/CMakeLists.txt
@@ -0,0 +1,7 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/../cpu/eval.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/../cpu/encoder.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp)
diff --git a/macos/third_party/mlx/mlx/backend/no_cpu/compiled.cpp b/macos/third_party/mlx/mlx/backend/no_cpu/compiled.cpp
new file mode 100644
index 0000000..2eeddab
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_cpu/compiled.cpp
@@ -0,0 +1,24 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/compile_impl.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+// GPU compile is always available if the GPU is available and since we are in
+// this file CPU compile is not available so check if the device is a GPU
+// device.
+namespace detail {
+bool compile_available_for_device(const Device& device) {
+  return device == Device::gpu;
+}
+} // namespace detail
+
+void Compiled::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error(
+      "[Compiled::eval_cpu] CPU compilation not supported on the platform.");
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/no_cpu/device_info.cpp b/macos/third_party/mlx/mlx/backend/no_cpu/device_info.cpp
new file mode 100644
index 0000000..0fc07fc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_cpu/device_info.cpp
@@ -0,0 +1,22 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cpu/device_info.h"
+
+namespace mlx::core::cpu {
+
+bool is_available() {
+  return false;
+}
+
+int device_count() {
+  return 0;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int /* device_index */) {
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+  return empty;
+}
+
+} // namespace mlx::core::cpu
diff --git a/macos/third_party/mlx/mlx/backend/no_cpu/primitives.cpp b/macos/third_party/mlx/mlx/backend/no_cpu/primitives.cpp
new file mode 100644
index 0000000..ae51dd9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_cpu/primitives.cpp
@@ -0,0 +1,146 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/primitives.h"
+#include "mlx/distributed/primitives.h"
+#include "mlx/fast_primitives.h"
+
+#define NO_CPU_MULTI(func)                                             \
+  void func::eval_cpu(                                                 \
+      const std::vector<array>& inputs, std::vector<array>& outputs) { \
+    throw std::runtime_error(#func " has no CPU implementation.");     \
+  }
+
+#define NO_CPU(func)                                                  \
+  void func::eval_cpu(const std::vector<array>& inputs, array& out) { \
+    throw std::runtime_error(#func " has no CPU implementation.");    \
+  }
+
+namespace mlx::core {
+
+NO_CPU(Abs)
+NO_CPU(Add)
+NO_CPU(AddMM)
+NO_CPU(Arange)
+NO_CPU(ArcCos)
+NO_CPU(ArcCosh)
+NO_CPU(ArcSin)
+NO_CPU(ArcSinh)
+NO_CPU(ArcTan)
+NO_CPU(ArcTan2)
+NO_CPU(ArcTanh)
+NO_CPU(ArgPartition)
+NO_CPU(ArgReduce)
+NO_CPU(ArgSort)
+NO_CPU(AsType)
+NO_CPU(AsStrided)
+NO_CPU(BitwiseBinary)
+NO_CPU(BitwiseInvert)
+NO_CPU(BlockMaskedMM)
+NO_CPU(Broadcast)
+NO_CPU(BroadcastAxes)
+NO_CPU(Ceil)
+NO_CPU(Cholesky)
+NO_CPU(Concatenate)
+NO_CPU(Conjugate)
+NO_CPU(Contiguous)
+NO_CPU(Convolution)
+NO_CPU(Copy)
+NO_CPU(Cos)
+NO_CPU(Cosh)
+NO_CPU_MULTI(CustomTransforms)
+NO_CPU_MULTI(Depends)
+NO_CPU(Divide)
+NO_CPU_MULTI(DivMod)
+NO_CPU(DynamicSlice)
+NO_CPU(DynamicSliceUpdate)
+NO_CPU(NumberOfElements)
+NO_CPU(Remainder)
+NO_CPU_MULTI(Eig)
+NO_CPU_MULTI(Eigh)
+NO_CPU(Equal)
+NO_CPU(Erf)
+NO_CPU(ErfInv)
+NO_CPU(Exp)
+NO_CPU(ExpandDims)
+NO_CPU(Expm1)
+NO_CPU(FFT)
+NO_CPU(Flatten)
+NO_CPU(Floor)
+NO_CPU(Full)
+NO_CPU(Gather)
+NO_CPU(GatherAxis)
+NO_CPU(GatherMM)
+NO_CPU(GatherQMM)
+NO_CPU(Greater)
+NO_CPU(GreaterEqual)
+NO_CPU(Hadamard)
+NO_CPU(Imag)
+NO_CPU(Less)
+NO_CPU(LessEqual)
+NO_CPU(Log)
+NO_CPU(Log1p)
+NO_CPU(LogicalNot)
+NO_CPU(LogicalAnd)
+NO_CPU(LogicalOr)
+NO_CPU(LogAddExp)
+NO_CPU(LogSumExp)
+NO_CPU_MULTI(LUF)
+NO_CPU(Matmul)
+NO_CPU(Maximum)
+NO_CPU(MaskedScatter)
+NO_CPU(Minimum)
+NO_CPU(Multiply)
+NO_CPU(Negative)
+NO_CPU(NotEqual)
+NO_CPU(Pad)
+NO_CPU(Partition)
+NO_CPU(Power)
+NO_CPU_MULTI(QRF)
+NO_CPU(QuantizedMatmul)
+NO_CPU(QQMatmul)
+NO_CPU(RandomBits)
+NO_CPU(Real)
+NO_CPU(Reduce)
+NO_CPU(Reshape)
+NO_CPU(Round)
+NO_CPU(Scan)
+NO_CPU(Scatter)
+NO_CPU(ScatterAxis)
+NO_CPU(Select)
+NO_CPU(SegmentedMM)
+NO_CPU(Sigmoid)
+NO_CPU(Sign)
+NO_CPU(Sin)
+NO_CPU(Sinh)
+NO_CPU(Slice)
+NO_CPU(SliceUpdate)
+NO_CPU(Softmax)
+NO_CPU(Sort)
+NO_CPU_MULTI(Split)
+NO_CPU(Square)
+NO_CPU(Squeeze)
+NO_CPU(Sqrt)
+NO_CPU(StopGradient)
+NO_CPU(Subtract)
+NO_CPU_MULTI(SVD)
+NO_CPU(Tan)
+NO_CPU(Tanh)
+NO_CPU(Transpose)
+NO_CPU(Unflatten)
+NO_CPU(Inverse)
+NO_CPU(View)
+
+namespace fast {
+NO_CPU_MULTI(Quantize)
+NO_CPU_MULTI(ConvertFP8)
+} // namespace fast
+
+namespace distributed {
+NO_CPU_MULTI(AllReduce)
+NO_CPU_MULTI(AllGather)
+NO_CPU_MULTI(Send)
+NO_CPU_MULTI(Recv)
+NO_CPU_MULTI(ReduceScatter)
+} // namespace distributed
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/CMakeLists.txt b/macos/third_party/mlx/mlx/backend/no_gpu/CMakeLists.txt
new file mode 100644
index 0000000..c777204
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/CMakeLists.txt
@@ -0,0 +1,8 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp)
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/allocator.cpp b/macos/third_party/mlx/mlx/backend/no_gpu/allocator.cpp
new file mode 100644
index 0000000..abb83e5
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/allocator.cpp
@@ -0,0 +1,134 @@
+// Copyright © 2023 Apple Inc.
+
+#include <algorithm>
+#include <mutex>
+
+#include "mlx/allocator.h"
+#include "mlx/memory.h"
+
+#ifdef __APPLE__
+#include "mlx/backend/no_gpu/apple_memory.h"
+#elif defined(__linux__)
+#include "mlx/backend/no_gpu/linux_memory.h"
+#else
+size_t get_memory_size() {
+  return 0;
+}
+#endif
+
+namespace mlx::core {
+
+namespace allocator {
+
+class CommonAllocator : public Allocator {
+  /** A general CPU allocator. */
+ public:
+  virtual Buffer malloc(size_t size) override;
+  virtual void free(Buffer buffer) override;
+  virtual size_t size(Buffer buffer) const override;
+
+  size_t get_active_memory() const {
+    return active_memory_;
+  };
+  size_t get_peak_memory() const {
+    return peak_memory_;
+  };
+  void reset_peak_memory() {
+    std::unique_lock lk(mutex_);
+    peak_memory_ = 0;
+  };
+  size_t get_memory_limit() {
+    return memory_limit_;
+  }
+  size_t set_memory_limit(size_t limit) {
+    std::unique_lock lk(mutex_);
+    std::swap(memory_limit_, limit);
+    return limit;
+  }
+
+ private:
+  size_t memory_limit_;
+  size_t active_memory_{0};
+  size_t peak_memory_{0};
+  std::mutex mutex_;
+  CommonAllocator() : memory_limit_(0.8 * get_memory_size()) {
+    if (memory_limit_ == 0) {
+      memory_limit_ = 1UL << 33;
+    }
+  };
+
+  friend CommonAllocator& common_allocator();
+};
+
+CommonAllocator& common_allocator() {
+  static CommonAllocator allocator_;
+  return allocator_;
+}
+
+Allocator& allocator() {
+  return common_allocator();
+}
+
+void* Buffer::raw_ptr() {
+  if (!ptr_) {
+    return nullptr;
+  }
+  return static_cast<size_t*>(ptr_) + 1;
+}
+
+Buffer CommonAllocator::malloc(size_t size) {
+  void* ptr = std::malloc(size + sizeof(size_t));
+  if (ptr != nullptr) {
+    *static_cast<size_t*>(ptr) = size;
+  }
+  std::unique_lock lk(mutex_);
+  active_memory_ += size;
+  peak_memory_ = std::max(active_memory_, peak_memory_);
+  return Buffer{ptr};
+}
+
+void CommonAllocator::free(Buffer buffer) {
+  auto sz = size(buffer);
+  std::free(buffer.ptr());
+  std::unique_lock lk(mutex_);
+  active_memory_ -= sz;
+}
+
+size_t CommonAllocator::size(Buffer buffer) const {
+  if (buffer.ptr() == nullptr) {
+    return 0;
+  }
+  return *static_cast<size_t*>(buffer.ptr());
+}
+
+} // namespace allocator
+
+size_t get_active_memory() {
+  return allocator::common_allocator().get_active_memory();
+}
+size_t get_peak_memory() {
+  return allocator::common_allocator().get_peak_memory();
+}
+void reset_peak_memory() {
+  return allocator::common_allocator().reset_peak_memory();
+}
+size_t set_memory_limit(size_t limit) {
+  return allocator::common_allocator().set_memory_limit(limit);
+}
+size_t get_memory_limit() {
+  return allocator::common_allocator().get_memory_limit();
+}
+
+// No-ops for common allocator
+size_t get_cache_memory() {
+  return 0;
+}
+size_t set_cache_limit(size_t) {
+  return 0;
+}
+size_t set_wired_limit(size_t) {
+  return 0;
+}
+void clear_cache() {}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/apple_memory.h b/macos/third_party/mlx/mlx/backend/no_gpu/apple_memory.h
new file mode 100644
index 0000000..7fdc530
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/apple_memory.h
@@ -0,0 +1,16 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <sys/sysctl.h>
+
+namespace {
+
+size_t get_memory_size() {
+  size_t memsize = 0;
+  size_t length = sizeof(memsize);
+  sysctlbyname("hw.memsize", &memsize, &length, NULL, 0);
+  return memsize;
+}
+
+} // namespace
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/device_info.cpp b/macos/third_party/mlx/mlx/backend/no_gpu/device_info.cpp
new file mode 100644
index 0000000..65d1cfc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/device_info.cpp
@@ -0,0 +1,22 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/gpu/device_info.h"
+
+namespace mlx::core::gpu {
+
+bool is_available() {
+  return false;
+}
+
+int device_count() {
+  return 0;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int /* device_index */) {
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+  return empty;
+}
+
+} // namespace mlx::core::gpu
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/eval.cpp b/macos/third_party/mlx/mlx/backend/no_gpu/eval.cpp
new file mode 100644
index 0000000..5d88970
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/eval.cpp
@@ -0,0 +1,24 @@
+// Copyright © 2025 Apple Inc.
+
+#include <stdexcept>
+
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/gpu/eval.h"
+
+namespace mlx::core::gpu {
+
+void new_stream(Stream) {}
+
+void eval(array&) {
+  throw std::runtime_error("[gpu::eval] GPU backend is not available");
+}
+
+void finalize(Stream) {
+  throw std::runtime_error("[gpu::finalize] GPU backend is not available");
+}
+
+void synchronize(Stream) {
+  throw std::runtime_error("[gpu::synchronize]  GPU backend is not available");
+}
+
+} // namespace mlx::core::gpu
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/event.cpp b/macos/third_party/mlx/mlx/backend/no_gpu/event.cpp
new file mode 100644
index 0000000..6dde047
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/event.cpp
@@ -0,0 +1,53 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/event.h"
+#include "mlx/scheduler.h"
+
+#include <condition_variable>
+#include <mutex>
+
+namespace mlx::core {
+
+struct EventCounter {
+  uint64_t value{0};
+  std::mutex mtx;
+  std::condition_variable cv;
+};
+
+Event::Event(Stream stream) : stream_(stream) {
+  auto dtor = [](void* ptr) { delete static_cast<EventCounter*>(ptr); };
+  event_ = std::shared_ptr<void>(new EventCounter{}, dtor);
+}
+
+void Event::wait() {
+  auto ec = static_cast<EventCounter*>(event_.get());
+  std::unique_lock<std::mutex> lk(ec->mtx);
+  if (ec->value >= value()) {
+    return;
+  }
+  ec->cv.wait(lk, [value = value(), ec] { return ec->value >= value; });
+}
+
+void Event::wait(Stream stream) {
+  scheduler::enqueue(stream, [*this]() mutable { wait(); });
+}
+
+void Event::signal(Stream stream) {
+  scheduler::enqueue(stream, [*this]() mutable {
+    auto ec = static_cast<EventCounter*>(event_.get());
+    {
+      std::lock_guard<std::mutex> lk(ec->mtx);
+      ec->value = value();
+    }
+    ec->cv.notify_all();
+  });
+}
+
+bool Event::is_signaled() const {
+  auto ec = static_cast<EventCounter*>(event_.get());
+  {
+    std::lock_guard<std::mutex> lk(ec->mtx);
+    return (ec->value >= value());
+  }
+}
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/fence.cpp b/macos/third_party/mlx/mlx/backend/no_gpu/fence.cpp
new file mode 100644
index 0000000..cd66d23
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/fence.cpp
@@ -0,0 +1,54 @@
+// Copyright © 2024 Apple Inc.
+
+#include <condition_variable>
+#include <mutex>
+
+#include "mlx/fence.h"
+#include "mlx/scheduler.h"
+
+namespace mlx::core {
+
+struct FenceImpl {
+  uint32_t count{0};
+  uint32_t value{0};
+  std::mutex mtx;
+  std::condition_variable cv;
+};
+
+Fence::Fence(Stream) {
+  auto dtor = [](void* ptr) { delete static_cast<FenceImpl*>(ptr); };
+  fence_ = std::shared_ptr<void>(new FenceImpl{}, dtor);
+}
+
+void Fence::wait(Stream stream, const array&) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      std::unique_lock<std::mutex> lk(f.mtx);
+      if (f.value >= count) {
+        return;
+      }
+      f.cv.wait(lk, [&f, count] { return f.value >= count; });
+    });
+  } else {
+    throw std::runtime_error("[Fence::wait] Invalid stream.");
+  }
+}
+
+void Fence::update(Stream stream, const array&, bool) {
+  auto& f = *static_cast<FenceImpl*>(fence_.get());
+  f.count++;
+  if (stream.device == Device::cpu) {
+    scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable {
+      auto& f = *static_cast<FenceImpl*>(fence_.get());
+      std::unique_lock<std::mutex> lk(f.mtx);
+      f.value = count;
+      f.cv.notify_all();
+    });
+  } else {
+    throw std::runtime_error("[Fence::update] Invalid stream.");
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/linux_memory.h b/macos/third_party/mlx/mlx/backend/no_gpu/linux_memory.h
new file mode 100644
index 0000000..f909edc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/linux_memory.h
@@ -0,0 +1,22 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <sys/sysinfo.h>
+
+namespace {
+
+size_t get_memory_size() {
+  struct sysinfo info;
+
+  if (sysinfo(&info) != 0) {
+    return 0;
+  }
+
+  size_t total_ram = info.totalram;
+  total_ram *= info.mem_unit;
+
+  return total_ram;
+}
+
+} // namespace
diff --git a/macos/third_party/mlx/mlx/backend/no_gpu/primitives.cpp b/macos/third_party/mlx/mlx/backend/no_gpu/primitives.cpp
new file mode 100644
index 0000000..4819ed2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/backend/no_gpu/primitives.cpp
@@ -0,0 +1,185 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/primitives.h"
+#include "mlx/distributed/primitives.h"
+#include "mlx/fast_primitives.h"
+
+#define NO_GPU_MULTI(func)                                             \
+  void func::eval_gpu(                                                 \
+      const std::vector<array>& inputs, std::vector<array>& outputs) { \
+    throw std::runtime_error(#func " has no GPU implementation.");     \
+  }
+
+#define NO_GPU_USE_FALLBACK(func)     \
+  bool func::use_fallback(Stream s) { \
+    return true;                      \
+  }                                   \
+  NO_GPU_MULTI(func)
+
+#define NO_GPU(func)                                                  \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    throw std::runtime_error(#func " has no GPU implementation.");    \
+  }
+
+namespace mlx::core {
+
+bool fast::ScaledDotProductAttention::use_fallback(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool is_training,
+    bool output_logsumexp,
+    Stream s) {
+  return true;
+}
+
+bool fast::ScaledDotProductAttention::supports_bool_mask() {
+  return false;
+}
+
+bool fast::ScaledDotProductAttentionVJP::use_fallback(
+    const array& q,
+    Stream s) {
+  return true;
+}
+
+NO_GPU(Abs)
+NO_GPU(Add)
+NO_GPU(AddMM)
+NO_GPU(Arange)
+NO_GPU(ArcCos)
+NO_GPU(ArcCosh)
+NO_GPU(ArcSin)
+NO_GPU(ArcSinh)
+NO_GPU(ArcTan)
+NO_GPU(ArcTan2)
+NO_GPU(ArcTanh)
+NO_GPU(ArgPartition)
+NO_GPU(ArgReduce)
+NO_GPU(ArgSort)
+NO_GPU(AsType)
+NO_GPU(AsStrided)
+NO_GPU(BitwiseBinary)
+NO_GPU(BitwiseInvert)
+NO_GPU(BlockMaskedMM)
+NO_GPU(Broadcast)
+NO_GPU(BroadcastAxes)
+NO_GPU(Ceil)
+NO_GPU_MULTI(Compiled)
+NO_GPU(Concatenate)
+NO_GPU(Conjugate)
+NO_GPU(Contiguous)
+NO_GPU(Convolution)
+NO_GPU(Copy)
+NO_GPU(Cos)
+NO_GPU(Cosh)
+NO_GPU_MULTI(CustomTransforms)
+NO_GPU_MULTI(Depends)
+NO_GPU(Divide)
+NO_GPU_MULTI(DivMod)
+NO_GPU(DynamicSlice)
+NO_GPU(DynamicSliceUpdate)
+NO_GPU(NumberOfElements)
+NO_GPU(Remainder)
+NO_GPU(Equal)
+NO_GPU(Erf)
+NO_GPU(ErfInv)
+NO_GPU(Exp)
+NO_GPU(ExpandDims)
+NO_GPU(Expm1)
+NO_GPU(FFT)
+NO_GPU(Flatten)
+NO_GPU(Floor)
+NO_GPU(Full)
+NO_GPU(Gather)
+NO_GPU(GatherAxis)
+NO_GPU(GatherMM)
+NO_GPU(GatherQMM)
+NO_GPU(Greater)
+NO_GPU(GreaterEqual)
+NO_GPU(Hadamard)
+NO_GPU(Imag)
+NO_GPU(Less)
+NO_GPU(LessEqual)
+NO_GPU(Load)
+NO_GPU(Log)
+NO_GPU(Log1p)
+NO_GPU(LogicalNot)
+NO_GPU(LogicalAnd)
+NO_GPU(LogicalOr)
+NO_GPU(LogAddExp)
+NO_GPU(LogSumExp)
+NO_GPU_MULTI(LUF)
+NO_GPU(Matmul)
+NO_GPU(Maximum)
+NO_GPU(Minimum)
+NO_GPU(Multiply)
+NO_GPU(Negative)
+NO_GPU(NotEqual)
+NO_GPU(Pad)
+NO_GPU(Partition)
+NO_GPU(Power)
+NO_GPU_MULTI(QRF)
+NO_GPU(QuantizedMatmul)
+NO_GPU(QQMatmul)
+NO_GPU(RandomBits)
+NO_GPU(Real)
+NO_GPU(Reduce)
+NO_GPU(Reshape)
+NO_GPU(Round)
+NO_GPU(Scan)
+NO_GPU(Scatter)
+NO_GPU(ScatterAxis)
+NO_GPU(Select)
+NO_GPU(SegmentedMM)
+NO_GPU(Sigmoid)
+NO_GPU(Sign)
+NO_GPU(Sin)
+NO_GPU(Sinh)
+NO_GPU(Slice)
+NO_GPU(SliceUpdate)
+NO_GPU(Softmax)
+NO_GPU(Sort)
+NO_GPU_MULTI(Split)
+NO_GPU(Square)
+NO_GPU(Squeeze)
+NO_GPU(Sqrt)
+NO_GPU(StopGradient)
+NO_GPU(Subtract)
+NO_GPU_MULTI(SVD)
+NO_GPU(Tan)
+NO_GPU(Tanh)
+NO_GPU(Transpose)
+NO_GPU(Unflatten)
+NO_GPU(Inverse)
+NO_GPU(Cholesky)
+NO_GPU_MULTI(Eigh)
+NO_GPU_MULTI(Eig)
+NO_GPU(View)
+NO_GPU(MaskedScatter)
+
+namespace fast {
+NO_GPU_USE_FALLBACK(LayerNorm)
+NO_GPU_MULTI(LayerNormVJP)
+NO_GPU_USE_FALLBACK(RMSNorm)
+NO_GPU_MULTI(RMSNormVJP)
+NO_GPU_USE_FALLBACK(RoPE)
+NO_GPU_MULTI(ScaledDotProductAttention)
+NO_GPU_MULTI(ScaledDotProductAttentionVJP)
+NO_GPU_MULTI(ConvertFP8)
+NO_GPU_MULTI(Quantize)
+NO_GPU_MULTI(CustomKernel)
+} // namespace fast
+
+namespace distributed {
+NO_GPU_MULTI(AllReduce)
+NO_GPU_MULTI(AllGather)
+NO_GPU_MULTI(Send)
+NO_GPU_MULTI(Recv)
+NO_GPU_MULTI(ReduceScatter)
+} // namespace distributed
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/compile.cpp b/macos/third_party/mlx/mlx/compile.cpp
new file mode 100644
index 0000000..ca5f069
--- /dev/null
+++ b/macos/third_party/mlx/mlx/compile.cpp
@@ -0,0 +1,1243 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <cstdlib>
+#include <map>
+#include <sstream>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/allocator.h"
+#include "mlx/backend/common/compiled.h"
+#include "mlx/compile.h"
+#include "mlx/compile_impl.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/graph_utils.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/transforms_impl.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+constexpr int max_compile_depth = 11;
+constexpr int max_compile_arrays = 24;
+
+bool is_unary(const Primitive& p) {
+  return (
+      typeid(p) == typeid(Abs) || typeid(p) == typeid(ArcCos) ||
+      typeid(p) == typeid(ArcCosh) || typeid(p) == typeid(ArcSin) ||
+      typeid(p) == typeid(ArcSinh) || typeid(p) == typeid(ArcTan) ||
+      typeid(p) == typeid(ArcTanh) || typeid(p) == typeid(AsType) ||
+      typeid(p) == typeid(Ceil) || typeid(p) == typeid(Cos) ||
+      typeid(p) == typeid(Conjugate) || typeid(p) == typeid(Cosh) ||
+      typeid(p) == typeid(Remainder) || typeid(p) == typeid(Erf) ||
+      typeid(p) == typeid(ErfInv) || typeid(p) == typeid(Exp) ||
+      typeid(p) == typeid(Floor) || typeid(p) == typeid(Log) ||
+      typeid(p) == typeid(Log1p) || typeid(p) == typeid(LogicalNot) ||
+      typeid(p) == typeid(Negative) || typeid(p) == typeid(Round) ||
+      typeid(p) == typeid(Sigmoid) || typeid(p) == typeid(Sign) ||
+      typeid(p) == typeid(Sin) || typeid(p) == typeid(Sinh) ||
+      typeid(p) == typeid(Square) || typeid(p) == typeid(Sqrt) ||
+      typeid(p) == typeid(Tan) || typeid(p) == typeid(Tanh) ||
+      typeid(p) == typeid(Expm1) || typeid(p) == typeid(Real) ||
+      typeid(p) == typeid(Imag) || typeid(p) == typeid(BitwiseInvert));
+}
+
+bool is_binary(const Primitive& p) {
+  return (
+      typeid(p) == typeid(Add) || typeid(p) == typeid(Divide) ||
+      typeid(p) == typeid(Equal) || typeid(p) == typeid(Greater) ||
+      typeid(p) == typeid(GreaterEqual) || typeid(p) == typeid(Less) ||
+      typeid(p) == typeid(LessEqual) || typeid(p) == typeid(LogicalNot) ||
+      typeid(p) == typeid(LogicalAnd) || typeid(p) == typeid(LogicalOr) ||
+      typeid(p) == typeid(LogAddExp) || typeid(p) == typeid(Maximum) ||
+      typeid(p) == typeid(Minimum) || typeid(p) == typeid(Multiply) ||
+      typeid(p) == typeid(NotEqual) || typeid(p) == typeid(Power) ||
+      typeid(p) == typeid(Subtract) || typeid(p) == typeid(BitwiseBinary) ||
+      typeid(p) == typeid(ArcTan2));
+}
+
+bool is_ternary(const Primitive& p) {
+  return typeid(p) == typeid(Select);
+}
+
+bool is_broadcast(const Primitive& p) {
+  return typeid(p) == typeid(Broadcast);
+}
+
+bool is_noop(const Primitive& p) {
+  return typeid(p) == typeid(Copy) || typeid(p) == typeid(StopGradient);
+}
+
+bool is_reduction(const Primitive& p) {
+  return typeid(p) == typeid(Reduce) || typeid(p) == typeid(ArgReduce);
+}
+
+bool is_fusable(const Primitive& p) {
+  return is_unary(p) || is_binary(p) || is_ternary(p) || is_broadcast(p);
+}
+
+Compiled::Compiled(
+    Stream stream,
+    std::vector<array> inputs,
+    std::vector<array> outputs,
+    std::vector<array> tape,
+    std::unordered_set<uintptr_t> constant_ids)
+    : Primitive(stream),
+      inputs_(std::move(inputs)),
+      outputs_(std::move(outputs)),
+      tape_(std::move(tape)),
+      constant_ids_(std::move(constant_ids)),
+      is_constant_([this](size_t i) {
+        return constant_ids_.find(inputs_[i].id()) != constant_ids_.end();
+      }) {
+  // Build the kernel name.
+  NodeNamer namer;
+  std::ostringstream os;
+  std::ostringstream constant_hasher;
+
+  std::unordered_set<uintptr_t> output_ids;
+  for (auto& o : outputs_) {
+    output_ids.insert(o.id());
+  }
+
+  // Fill the input names. This is not really necessary, I just like having A,
+  // B, C, ... as the inputs.
+  for (const auto& x : inputs_) {
+    namer.get_name(x);
+  }
+
+  // The primitives describing the tape. For unary and binary primitives this
+  // must be enough to describe the full computation.
+  for (const auto& a : tape_) {
+    // name and type of output
+    os << namer.get_name(a) << kindof(a.dtype()) << a.itemsize();
+    // whether or not it's an output
+    if (output_ids.find(a.id()) != output_ids.end()) {
+      os << "O";
+    } else {
+      os << "I";
+    }
+    // computation performed
+    os << a.primitive().name();
+    // name of inputs to the function
+    for (auto& inp : a.inputs()) {
+      os << namer.get_name(inp);
+    }
+  }
+  os << "_";
+
+  for (const auto& x : inputs_) {
+    if (constant_ids_.find(x.id()) != constant_ids_.end()) {
+      os << "C";
+      print_constant(constant_hasher, x);
+    } else {
+      os << (is_scalar(x) ? "S" : "V");
+    }
+  }
+  os << "_";
+  for (const auto& x : inputs) {
+    if (constant_ids.find(x.id()) != constant_ids.end()) {
+      continue;
+    }
+    os << kindof(x.dtype()) << x.itemsize();
+  }
+  os << "_" << std::hash<std::string>{}(constant_hasher.str());
+
+  kernel_lib_ = os.str();
+}
+
+std::vector<array> Compiled::vjp(
+    const std::vector<array>&,
+    const std::vector<array>&,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  throw std::runtime_error("[Compiled] Cannot vjp primitive.");
+}
+
+std::vector<array> Compiled::jvp(
+    const std::vector<array>&,
+    const std::vector<array>&,
+    const std::vector<int>&) {
+  throw std::runtime_error("[Compiled] Cannot jvp primitive.");
+}
+
+std::pair<std::vector<array>, std::vector<int>> Compiled::vmap(
+    const std::vector<array>&,
+    const std::vector<int>&) {
+  throw std::runtime_error("[Compiled] Cannot vmap primitive.");
+}
+
+bool Compiled::is_equivalent(const Primitive& other) const {
+  const Compiled& a_other = static_cast<const Compiled&>(other);
+  return std::equal(
+      tape_.begin(),
+      tape_.end(),
+      a_other.tape_.begin(),
+      a_other.tape_.end(),
+      [](const array& a1, const array& a2) {
+        auto& p1 = a1.primitive();
+        auto& p2 = a2.primitive();
+        return typeid(p1) == typeid(p2) && p1.is_equivalent(p2);
+      });
+}
+
+const char* Compiled::name() const {
+  if (name_.empty()) {
+    std::ostringstream os;
+    os << "Compiled";
+    for (auto& a : tape_) {
+      os << a.primitive().name();
+    }
+    name_ = os.str();
+  }
+  return name_.c_str();
+}
+
+std::vector<Shape> Compiled::output_shapes(const std::vector<array>& inputs) {
+  size_t nd = 0;
+  for (auto& in : inputs) {
+    nd = std::max(nd, in.ndim());
+  }
+  Shape out_shape(nd, 0);
+  for (auto& in : inputs) {
+    auto dd = nd - in.ndim();
+    for (auto i = dd; i < nd; ++i) {
+      out_shape[i] = std::max(out_shape[i], in.shape()[i - dd]);
+    }
+  }
+  // All outputs have the same shape
+  return std::vector<Shape>(outputs_.size(), out_shape);
+}
+
+namespace detail {
+
+CompileMode& compile_mode() {
+  auto get_val = []() {
+    if (std::getenv("MLX_DISABLE_COMPILE")) {
+      return CompileMode::disabled;
+    } else {
+      return CompileMode::enabled;
+    }
+  };
+  static CompileMode compile_mode_ = get_val();
+  return compile_mode_;
+}
+
+// Helper like below but only merges the two provided arrays. If the src has
+// siblings then these won't be merged to the dst.
+void merge_one(array& dst, array& src, ParentsMap& parents_map) {
+  auto src_parents = parents_map.find(src.id());
+  if (src_parents == parents_map.end()) {
+    return;
+  }
+  auto& pairs = parents_map[dst.id()];
+  for (auto& parent : src_parents->second) {
+    parent.first.inputs()[parent.second] = dst;
+    pairs.push_back(parent);
+  }
+
+  // If src is a parent of dst, remove it from dst's parents
+  for (auto it = pairs.begin(); it != pairs.end();) {
+    if (it->first.id() == src.id()) {
+      it = pairs.erase(it);
+    } else {
+      it++;
+    }
+  }
+  // Remove the source from the map to avoid fusing with it again
+  parents_map.erase(src_parents);
+}
+
+// Helper that merges two arrays in the graph by setting the parents of the
+// source to point to the destination. The arrays are assumed to be coming from
+// equivalent primitives so their siblings are merged as well.
+void merge(array& dst, array& src, ParentsMap& parents_map) {
+  // Canonicalize the order of the primitives outputs
+  auto sources = src.outputs();
+  auto dests = dst.outputs();
+  // For each src parent, point it to the corresponding dst
+  for (int i = 0; i < sources.size(); ++i) {
+    merge_one(dests[i], sources[i], parents_map);
+  }
+}
+
+// Any parent in the divider will continue to refer to `x` but any parent not
+// in the divider will refer to a copy of the operation.
+array split_one(
+    const array& x,
+    ParentsMap& parents_map,
+    const std::unordered_set<uintptr_t>& divider) {
+  array y(x.shape(), x.dtype(), x.primitive_ptr(), x.inputs());
+
+  auto& x_parents = parents_map[x.id()];
+  auto& y_parents = parents_map[y.id()];
+
+  for (auto it = x_parents.begin(); it != x_parents.end();) {
+    if (divider.find(it->first.id()) != divider.end()) {
+      it->first.inputs()[it->second] = y;
+      y_parents.emplace_back(std::move(*it));
+      it = x_parents.erase(it);
+    } else {
+      it++;
+    }
+  }
+
+  return y;
+}
+
+template <typename T, typename... U>
+std::uintptr_t get_function_address(const std::function<T(U...)>& fun) {
+  using FunType = T (*)(U...);
+  const FunType* fun_ptr = fun.template target<FunType>();
+  if (fun_ptr == nullptr) {
+    return 0;
+  }
+  return reinterpret_cast<std::uintptr_t>(*fun_ptr);
+}
+
+class CompilerCache {
+ public:
+  struct CacheEntry {
+    CacheEntry(Stream stream, bool shapeless)
+        : stream(stream), shapeless(shapeless) {};
+    Stream stream;
+    bool shapeless;
+    std::vector<array> inputs;
+    std::vector<array> outputs;
+    std::vector<array> tape;
+    bool empty{true};
+    std::vector<uint64_t> constants;
+    std::shared_ptr<void> extra;
+  };
+
+  // Returns a reference to a CacheEntry which can be updated
+  // by the caller to avoid copying large tapes / inputs / outputs
+  CacheEntry& find(
+      std::uintptr_t fun_id,
+      const std::vector<array>& inputs,
+      bool shapeless,
+      const std::vector<uint64_t>& constants) {
+    // Find the cache entries for |fun_id|.
+    std::vector<CacheEntry>& entries = cache_[fun_id];
+
+    // Compare if 2 arrays have same shape and dtype.
+    auto has_same_shape_and_dtype = [shapeless](
+                                        const std::vector<array>& in1,
+                                        const std::vector<array>& in2) {
+      if (in1.size() != in2.size()) {
+        return false;
+      }
+      for (size_t i = 0; i < in1.size(); ++i) {
+        if (in1[i].ndim() != in2[i].ndim()) {
+          return false;
+        }
+        if (!shapeless && in1[i].shape() != in2[i].shape()) {
+          return false;
+        }
+        if (in1[i].dtype() != in2[i].dtype()) {
+          return false;
+        }
+      }
+      return true;
+    };
+    // Loop over entries and check:
+    // - Default stream and device match the entry's default stream
+    // - Inputs match i.e. shapes and types must be equal.
+    auto stream = default_stream(default_device());
+    for (CacheEntry& entry : entries) {
+      // Check that the default stream and device match
+      if (entry.stream != stream) {
+        continue;
+      }
+      if (entry.shapeless != shapeless) {
+        continue;
+      }
+
+      // Check the inputs match and return if so
+      if (has_same_shape_and_dtype(inputs, entry.inputs) &&
+          constants == entry.constants) {
+        return entry;
+      }
+    }
+    // Otherwise append a new cache entry
+    entries.push_back(CacheEntry{stream, shapeless});
+    return entries.back();
+  }
+
+  void erase(std::uintptr_t fun_id) {
+    cache_.erase(fun_id);
+  }
+
+  void clear() {
+    cache_.clear();
+  }
+
+ private:
+  CompilerCache() {
+    // Make sure the allocator is fully
+    // initialized before the compiler cache
+    allocator::allocator();
+  }
+
+  friend CompilerCache& compiler_cache();
+  std::unordered_map<std::uintptr_t, std::vector<CacheEntry>> cache_;
+};
+
+CompilerCache& compiler_cache() {
+  static CompilerCache compiler_cache_;
+  return compiler_cache_;
+}
+
+std::tuple<std::vector<array>, std::vector<array>, std::shared_ptr<void>>
+compile_trace(
+    const ArrayFnWithExtra& fun,
+    const std::vector<array>& inputs,
+    bool shapeless) {
+  // Set the global tracing flag.
+  detail::InTracing in_tracing{shapeless};
+
+  // Run the function on placeholder inputs
+  // to get compute graph
+  std::vector<array> tracer_inputs;
+  for (int i = 0; i < inputs.size(); ++i) {
+    array in(inputs[i].shape(), inputs[i].dtype(), nullptr, {});
+    in.set_tracer(true);
+    tracer_inputs.push_back(std::move(in));
+  }
+
+  auto output = fun(tracer_inputs);
+  return {tracer_inputs, output.first, output.second};
+}
+
+// Traverses the graph to build a tape and a map of array ids to their parents
+std::pair<std::vector<array>, ParentsMap> compile_dfs(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const std::vector<array>& original_inputs) {
+  std::vector<array> tape;
+  std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>
+      parents_map;
+  {
+    std::function<void(const array&)> recurse;
+    std::unordered_set<std::uintptr_t> input_set;
+    std::unordered_set<std::uintptr_t> original_input_set;
+    for (int i = 0; i < inputs.size(); ++i) {
+      input_set.insert(inputs[i].id());
+      original_input_set.insert(original_inputs[i].id());
+    }
+
+    // DFS the graph to build the tape, and log parents and scalars
+    std::unordered_set<std::uintptr_t> cache;
+    recurse = [&](const array& a) {
+      auto id = a.id();
+      if (original_input_set.find(id) != original_input_set.end()) {
+        throw std::invalid_argument(
+            "[compile] Attempting to compile a function with uncaptured inputs is not allowed.");
+      }
+      if (cache.find(id) != cache.end()) {
+        return;
+      }
+      for (int i = 0; i < a.inputs().size(); i++) {
+        auto& in = a.inputs()[i];
+        parents_map[in.id()].push_back({a, i});
+        for (auto& s : a.siblings()) {
+          parents_map[in.id()].push_back({s, i});
+        }
+        // Don't recurse on inputs (but add them to the tape for the purpose
+        // of future optimizations)
+        if (input_set.find(a.id()) == input_set.end()) {
+          recurse(in);
+        }
+      }
+      cache.insert(id);
+      for (auto& s : a.siblings()) {
+        cache.insert(s.id());
+      }
+      tape.push_back(a);
+    };
+    for (auto& a : outputs) {
+      recurse(a);
+    }
+  }
+
+  // Deep copy the tape and parents map while preserving inputs and outputs
+  std::vector<array> new_tape;
+  std::unordered_set<uintptr_t> io_set;
+  std::unordered_map<uintptr_t, array> old_to_new;
+  for (auto& o : outputs) {
+    old_to_new.insert({o.id(), o});
+    io_set.insert(o.id());
+    for (auto& s : o.siblings()) {
+      old_to_new.insert({s.id(), s});
+      io_set.insert(s.id());
+    }
+  }
+  for (auto& i : inputs) {
+    io_set.insert(i.id());
+    old_to_new.insert({i.id(), i});
+  }
+
+  new_tape.reserve(tape.size());
+  for (auto& arr : tape) {
+    if (!arr.has_primitive() || (io_set.find(arr.id()) != io_set.end())) {
+      old_to_new.insert({arr.id(), arr});
+      new_tape.push_back(arr);
+      continue;
+    }
+    std::vector<array> inputs;
+    inputs.reserve(arr.inputs().size());
+    for (auto& i : arr.inputs()) {
+      inputs.push_back(old_to_new.find(i.id())->second);
+    }
+    if (arr.siblings().size() > 0) {
+      std::vector<Dtype> types;
+      std::vector<Shape> shapes;
+      auto out = arr.outputs();
+      for (auto& o : out) {
+        types.push_back(o.dtype());
+        shapes.push_back(o.shape());
+      }
+      auto as = array::make_arrays(
+          std::move(shapes), types, arr.primitive_ptr(), std::move(inputs));
+      for (int i = 0; i < out.size(); ++i) {
+        old_to_new.insert({out[i].id(), as[i]});
+      }
+      new_tape.push_back(as[arr.sibling_position()]);
+    } else {
+      auto a = array(
+          arr.shape(), arr.dtype(), arr.primitive_ptr(), std::move(inputs));
+      old_to_new.insert({arr.id(), a});
+      new_tape.push_back(a);
+    }
+  }
+  io_set.clear();
+  for (auto& o : outputs) {
+    if (!(io_set.insert(o.id()).second)) {
+      continue;
+    }
+    for (auto& i : o.inputs()) {
+      i = old_to_new.find(i.id())->second;
+    }
+    for (auto& s : o.siblings()) {
+      io_set.insert(s.id());
+      for (auto& i : s.inputs()) {
+        i = old_to_new.find(i.id())->second;
+      }
+    }
+  }
+  tape = std::move(new_tape);
+
+  std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>
+      new_parents_map;
+  for (auto& [id, vec] : parents_map) {
+    for (auto& [a, _] : vec) {
+      a = old_to_new.find(a.id())->second;
+    }
+    new_parents_map[old_to_new.find(id)->second.id()] = std::move(vec);
+  }
+  parents_map = std::move(new_parents_map);
+  return {tape, parents_map};
+}
+
+static inline uint64_t splitmix64(uint64_t x) noexcept {
+  x += 0x9e3779b97f4a7c15ull;
+  x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ull;
+  x = (x ^ (x >> 27)) * 0x94d049bb133111ebull;
+  return x ^ (x >> 31);
+}
+
+struct VecU64Hash {
+  size_t operator()(const std::vector<uint64_t>& s) const noexcept {
+    uint64_t h =
+        0x243f6a8885a308d3ull ^ (uint64_t)s.size() * 0x9e3779b97f4a7c15ull;
+    for (uint64_t x : s) {
+      h = splitmix64(x ^ splitmix64(h + 0x9e3779b97f4a7c15ull));
+    }
+    return (size_t)h;
+  }
+};
+
+// Simplify the tape. Note, this function modifies in-place both the tape,
+// the parents map to remove orphaned arrays, and potentially the outputs
+void compile_simplify(
+    std::vector<array>& tape,
+    ParentsMap& parents_map,
+    std::vector<array>& outputs,
+    int passes) {
+  // Helpers to identify identical scalars
+  std::map<std::pair<uint64_t, Dtype::Val>, array> scalars;
+  auto is_scalar = [](const array& a) {
+    // Condition for when it's safe to read an array
+    return a.is_available() && a.ndim() == 0;
+  };
+  auto get_scalar_rep = [](const array& a) {
+    uint64_t v = 0;
+    switch (a.dtype().size()) {
+      case 1:
+        v = *a.data<uint8_t>();
+        break;
+      case 2:
+        v = *a.data<uint16_t>();
+        break;
+      case 4:
+        v = *a.data<uint32_t>();
+        break;
+      case 8:
+        v = *a.data<uint64_t>();
+        break;
+    }
+    return std::make_pair(v, a.dtype().val());
+  };
+
+  for (auto& a : tape) {
+    if (is_scalar(a)) {
+      scalars.insert({get_scalar_rep(a), a});
+    }
+  }
+
+  // Depth-1 array equivalence check.
+  auto array_equivalent = [](const array& a, const array& b) {
+    if (!a.has_primitive() || !b.has_primitive()) {
+      return false;
+    }
+    if (a.primitive_id() == b.primitive_id()) {
+      return false;
+    }
+    const auto& pa = a.primitive();
+    const auto& pb = b.primitive();
+    if (typeid(pa) != typeid(pb)) {
+      return false;
+    }
+
+    if (a.inputs().size() != b.inputs().size()) {
+      return false;
+    }
+
+    for (int i = 0; i < a.inputs().size(); i++) {
+      if (a.inputs()[i].id() != b.inputs()[i].id()) {
+        return false;
+      }
+    }
+
+    return pa.is_equivalent(pb);
+  };
+
+  // Merge scalars
+  std::vector<array> new_tape;
+  for (auto& arr : tape) {
+    // Check if we can merge scalars
+    if (is_scalar(arr)) {
+      auto scalar = scalars.find(get_scalar_rep(arr));
+      if (scalar->second.id() != arr.id()) {
+        merge(scalar->second, arr, parents_map);
+        // Don't keep orphaned scalars in the tape
+        continue;
+      }
+    }
+    new_tape.push_back(std::move(arr));
+  }
+  tape = std::move(new_tape);
+
+  // Remove no-ops
+  {
+    std::unordered_map<uintptr_t, array> output_map;
+    for (auto& o : outputs) {
+      output_map.insert({o.id(), o});
+    }
+    for (auto& arr : tape) {
+      if (!arr.has_primitive() || !is_noop(arr.primitive())) {
+        new_tape.push_back(std::move(arr));
+        continue;
+      }
+      merge_one(arr.inputs()[0], arr, parents_map);
+      if (auto it = output_map.find(arr.id()); it != output_map.end()) {
+        it->second = arr.inputs()[0];
+      }
+    }
+    tape = std::move(new_tape);
+    for (auto& o : outputs) {
+      o = output_map.at(o.id());
+    }
+  }
+
+  std::unordered_map<std::uintptr_t, uint32_t> tape_order;
+  for (uint32_t i = 0; i < tape.size(); ++i) {
+    tape_order.insert({tape[i].id(), i});
+  }
+
+  std::unordered_set<uintptr_t> output_set;
+  for (auto& o : outputs) {
+    output_set.insert(o.id());
+  }
+
+  // Multi-pass merge only keeping non-orphaned arrays in the tape
+  for (int pass = 0; pass < passes; ++pass) {
+    for (auto& arr : tape) {
+      // Helper to check if we can merge the parents of the
+      // given array
+      auto maybe_merge_parents = [&](auto& a) {
+        auto parents = parents_map.find(a.id());
+        if (parents != parents_map.end()) {
+          auto N = parents->second.size();
+          std::vector<bool> mask(N, false);
+
+          auto try_merge = [&](int dst_idx, int src_idx) {
+            if (tape_order[parents->second[src_idx].first.id()] <
+                tape_order[parents->second[dst_idx].first.id()]) {
+              std::swap(src_idx, dst_idx);
+            }
+            auto& src = parents->second[src_idx].first;
+            auto& dst = parents->second[dst_idx].first;
+            if (src.id() != dst.id() && array_equivalent(src, dst) &&
+                output_set.find(src.id()) == output_set.end()) {
+              merge(dst, src, parents_map);
+              mask[src_idx] = true;
+            }
+          };
+
+          if (N > 100) {
+            std::unordered_map<
+                std::vector<uint64_t>,
+                std::vector<int>,
+                VecU64Hash>
+                dst_map;
+            // Find possibly mergeable groups
+            for (int i = 0; i < N; i++) {
+              // Make the hash key
+              std::vector<uint64_t> key;
+              auto& curr = parents->second[i].first;
+              key.reserve(curr.inputs().size() + 2);
+              for (auto& in : curr.inputs()) {
+                key.push_back(in.id());
+              }
+              auto& p = curr.primitive();
+              key.push_back(curr.inputs().size());
+              key.push_back(typeid(p).hash_code());
+              auto it = dst_map.find(key);
+              if (it == dst_map.end()) {
+                bool _;
+                std::tie(it, _) = dst_map.insert({key, std::vector<int>{}});
+              }
+              it->second.push_back(i);
+            }
+            for (auto& [_, group] : dst_map) {
+              for (int i = 0; i < group.size(); ++i) {
+                if (mask[group[i]]) {
+                  continue;
+                }
+                for (int j = i + 1; j < group.size(); ++j) {
+                  if (mask[group[j]]) {
+                    continue;
+                  }
+                  try_merge(group[i], group[j]);
+                }
+              }
+            }
+          } else {
+            for (int i = 0; i < N; ++i) {
+              if (mask[i]) {
+                continue;
+              }
+              for (int j = i + 1; j < N; ++j) {
+                if (mask[j]) {
+                  continue;
+                }
+                try_merge(i, j);
+              }
+            }
+          }
+
+          // Erase orphaned parents so we don't keep fusing with them
+          for (int i = N - 1; i >= 0; --i) {
+            if (mask[i]) {
+              parents->second.erase(parents->second.begin() + i);
+            }
+          }
+          return false;
+        } else {
+          return output_set.find(a.id()) == output_set.end();
+        }
+      };
+      bool discard = maybe_merge_parents(arr);
+      for (auto& s : arr.siblings()) {
+        discard &= maybe_merge_parents(s);
+      }
+      // If an array and its siblings have no parents, and none of them are
+      // outputs, it is safe to remove it from the tape
+      if (!discard) {
+        new_tape.push_back(std::move(arr));
+      }
+    }
+    tape = std::move(new_tape);
+  }
+}
+
+// Extract sub-graphs of the graph that can be compiled
+// and replace them with a Compiled Primitive.
+void compile_fuse(
+    std::vector<array>& tape,
+    ParentsMap& parents_map,
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  // Track outputs to replace with new compiled outputs
+  std::unordered_map<uintptr_t, array> output_map;
+  for (auto& o : outputs) {
+    output_map.insert({o.id(), o});
+  }
+
+  // Set of inputs to distinguish constants
+  std::unordered_set<uintptr_t> input_ids;
+  for (auto& in : inputs) {
+    input_ids.insert(in.id());
+  }
+
+  // Go through the tape in reverse order and check for fusable sub-graphs
+  std::vector<array> new_tape;
+  std::unordered_set<uintptr_t> global_cache;
+  for (int i = tape.size() - 1; i >= 0; --i) {
+    auto& arr = tape[i];
+
+    // Already compiled
+    if (global_cache.find(arr.id()) != global_cache.end()) {
+      continue;
+    }
+
+    // Two pass recursion:
+    // First pass:
+    //  - Collect all the primitives which we can fuse with
+    //  - Keeps a cache of fusable primitives which may be added out of
+    //    DAG order. We have to determine if all of a fused primitive's
+    //    outputs are also in the fused section, and this may not be the
+    //    case the first time we visit it.
+    // Second pass:
+    //  - Collect inputs to the new compiled primitive
+    //  - Add fusable primitives to a tape in the correct order
+
+    std::function<void(const array&, int, const Stream&, const Shape&)> recurse;
+    std::unordered_set<uintptr_t> cache;
+    std::unordered_set<uintptr_t> input_set;
+    recurse = [&](const array& a,
+                  int depth,
+                  const Stream& s,
+                  const Shape& shape) {
+      if (cache.find(a.id()) != cache.end()) {
+        return;
+      }
+
+      // Stop fusing if:
+      // - Depth limit exceeded
+      // - Constant input
+      // - Stream mismatch
+      // - Non fusable primitive
+      // - Is global output but has a different shape
+      if (depth >= max_compile_depth || !a.has_primitive() ||
+          a.primitive().stream() != s || !is_fusable(a.primitive()) ||
+          (output_map.find(a.id()) != output_map.end() && a.shape() != shape)) {
+        // Possible input
+        input_set.insert(a.id());
+        return;
+      }
+
+      bool all_parents_in = true;
+      if (depth > 0) {
+        // Guaranteed to have a parent since nested in the
+        // recursion.
+        auto& parents = parents_map.at(a.id());
+        for (auto& [p, idx] : parents) {
+          auto in_cache = cache.find(p.id()) != cache.end();
+          if (!in_cache) {
+            all_parents_in = false;
+            break;
+          }
+        }
+      }
+
+      // Arrays with a mix of parents outside the compilable section
+      // are not fusable except for broadcast which we can split to avoid
+      // stopping fusion
+      if (!all_parents_in) {
+        if (a.has_primitive() && is_broadcast(a.primitive())) {
+          array b = split_one(a, parents_map, cache);
+          recurse(b, depth, s, shape);
+        } else {
+          // Possible input
+          input_set.insert(a.id());
+        }
+        return;
+      }
+
+      if (output_map.find(a.id()) != output_map.end()) {
+        input_set.insert(a.id());
+      } else {
+        // Not an input anymore since fusing it
+        input_set.erase(a.id());
+      }
+      if (input_set.size() >= max_compile_arrays) {
+        return;
+      }
+      cache.insert({a.id()});
+
+      for (auto& in : a.inputs()) {
+        recurse(in, depth + 1, s, shape);
+      }
+    };
+
+    // This will be the result of the fused operation so it needs
+    //   a) to not be already computed ie have a primitive
+    //   b) that primitive to not be a broadcast since it will unnecessarily
+    //      cast to a contiguous array potentially blowing up memory
+    if (arr.has_primitive() && !is_broadcast(arr.primitive())) {
+      Stream s = arr.primitive().stream();
+      recurse(arr, 0, s, arr.shape());
+    }
+
+    // Not worth fusing a single primitive
+    if (cache.size() <= 1) {
+      new_tape.push_back(arr);
+      continue;
+    }
+
+    // Recurse a second time to build the tape in the right
+    // order and collect the inputs
+    input_set.clear();
+    std::vector<array> inputs;
+    std::vector<array> fused_tape;
+    std::unordered_set<uintptr_t> tape_set;
+    std::function<void(const array&)> recurse_tape;
+    recurse_tape = [&](const array& a) {
+      if (cache.find(a.id()) == cache.end()) {
+        if (input_set.find(a.id()) == input_set.end()) {
+          input_set.insert(a.id());
+          inputs.push_back(a);
+        }
+        return;
+      }
+      if (tape_set.find(a.id()) != tape_set.end()) {
+        return;
+      }
+      tape_set.insert(a.id());
+      for (auto& in : a.inputs()) {
+        recurse_tape(in);
+      }
+      fused_tape.push_back(a);
+    };
+    recurse_tape(arr);
+
+    std::vector<array> old_outputs;
+    // Add to global cache and add any global outputs to outputs
+    // of new primitive
+    for (int j = 0; j < fused_tape.size() - 1; ++j) {
+      auto& f = fused_tape[j];
+      if (output_map.find(f.id()) != output_map.end()) {
+        old_outputs.push_back(f);
+        // Parents are now siblings, update the parent map
+        auto& pairs = parents_map[f.id()];
+        pairs.erase(
+            std::remove_if(
+                pairs.begin(),
+                pairs.end(),
+                [&](auto& p) {
+                  return cache.find(p.first.id()) != cache.end();
+                }),
+            pairs.end());
+      } else {
+        // Remove inner fused arrays parents from the parents map
+        // to keep the parents map in a valid state
+        parents_map.erase(f.id());
+      }
+      global_cache.insert({f.id()});
+    }
+    old_outputs.push_back(arr);
+
+    std::vector<Shape> shapes;
+    std::vector<Dtype> types;
+    for (auto& o : old_outputs) {
+      if (o.shape() != old_outputs.back().shape()) {
+        throw std::runtime_error(
+            "[compile] Compilation failed. Tried to fuse operations with different output shapes");
+      }
+      shapes.push_back(o.shape());
+      types.push_back(o.dtype());
+    }
+    std::unordered_set<uintptr_t> constant_ids;
+    for (auto& in : inputs) {
+      // Scalar constant
+      if (in.size() == 1 && !in.has_primitive() &&
+          input_ids.find(in.id()) == input_ids.end()) {
+        constant_ids.insert(in.id());
+      }
+    }
+    auto compiled_outputs = array::make_arrays(
+        std::move(shapes),
+        types,
+        std::make_shared<Compiled>(
+            old_outputs.back().primitive().stream(),
+            inputs,
+            old_outputs,
+            std::move(fused_tape),
+            std::move(constant_ids)),
+        inputs);
+
+    // One output per primitive
+    new_tape.push_back(compiled_outputs.back());
+
+    // Replace inputs old parents with compiled_outputs
+    for (int i = 0; i < inputs.size(); ++i) {
+      auto& pairs = parents_map[inputs[i].id()];
+      pairs.erase(
+          std::remove_if(
+              pairs.begin(),
+              pairs.end(),
+              [&](auto& p) { return cache.find(p.first.id()) != cache.end(); }),
+          pairs.end());
+      for (auto& o : compiled_outputs) {
+        pairs.push_back({o, i});
+      }
+    }
+
+    // - Update outputs parents to point to compiled outputs
+    // - Update any overall graph outputs to be compiled outputs
+    for (int o = 0; o < old_outputs.size(); ++o) {
+      merge_one(compiled_outputs[o], old_outputs[o], parents_map);
+      if (auto it = output_map.find(old_outputs[o].id());
+          it != output_map.end()) {
+        it->second = compiled_outputs[o];
+      }
+    }
+  }
+
+  std::reverse(new_tape.begin(), new_tape.end());
+  tape = std::move(new_tape);
+
+  // Replace output with potentially compiled output
+  for (auto& o : outputs) {
+    o = output_map.at(o.id());
+  }
+}
+
+std::vector<array> compile_replace(
+    const std::vector<array>& tape,
+    const std::vector<array>& trace_inputs,
+    const std::vector<array>& trace_outputs,
+    const std::vector<array>& inputs,
+    bool shapeless) {
+  std::unordered_map<uintptr_t, array> trace_to_real;
+  for (int i = 0; i < inputs.size(); ++i) {
+    trace_to_real.insert({trace_inputs[i].id(), inputs[i]});
+  }
+
+  auto is_load = [](const Primitive& p) { return typeid(p) == typeid(Load); };
+
+  for (auto& a : tape) {
+    // Arrays in the tape without primitives are either:
+    // - inputs, which are already in the map
+    // - constants, which can be used directly
+    // - a load primitive which has no inputs and will become a constant
+    //   after the first eval
+    if (!a.has_primitive() || is_load(a.primitive())) {
+      trace_to_real.insert({a.id(), a});
+    } else {
+      // Find real inputs
+      std::vector<array> real_inputs;
+      for (auto& in : a.inputs()) {
+        real_inputs.push_back(trace_to_real.at(in.id()));
+      }
+      if (a.siblings().empty()) {
+        auto shape =
+            shapeless ? a.primitive().output_shapes(real_inputs)[0] : a.shape();
+        auto real_a = array(
+            std::move(shape),
+            a.dtype(),
+            a.primitive_ptr(),
+            std::move(real_inputs));
+        trace_to_real.insert({a.id(), std::move(real_a)});
+      } else {
+        // Ensure the order is correct for multi-output primitives
+        std::vector<Dtype> types;
+        auto trace_out = a.outputs();
+        for (auto& o : trace_out) {
+          types.push_back(o.dtype());
+        }
+        std::vector<Shape> shapes;
+        if (shapeless) {
+          shapes = a.primitive().output_shapes(real_inputs);
+        } else {
+          for (auto& o : trace_out) {
+            shapes.push_back(o.shape());
+          }
+        }
+        auto real_out = array::make_arrays(
+            std::move(shapes), types, a.primitive_ptr(), real_inputs);
+        for (int i = 0; i < trace_out.size(); ++i) {
+          trace_to_real.insert({trace_out[i].id(), std::move(real_out[i])});
+        }
+      }
+    }
+  }
+
+  std::vector<array> outputs;
+  for (auto& o : trace_outputs) {
+    outputs.push_back(trace_to_real.at(o.id()));
+  }
+  return outputs;
+}
+
+bool skip_compile() {
+  return compile_mode() == CompileMode::disabled ||
+      !(compile_available_for_device(default_device()));
+}
+
+ArrayFnWithExtra compile(
+    ArrayFnWithExtra fun,
+    std::uintptr_t fun_id,
+    bool shapeless /* = false */,
+    std::vector<uint64_t> constants /* = {} */) {
+  if (skip_compile()) {
+    return fun;
+  }
+  if (!fun) {
+    throw std::invalid_argument(
+        "[compile] Cannot compile a function without a target.");
+  }
+
+  return [fun = std::move(fun),
+          fun_id,
+          shapeless,
+          constants = std::move(constants)](const std::vector<array>& inputs) {
+    // If the inputs are tracers, trace the original graph
+    if (std::any_of(inputs.begin(), inputs.end(), [](auto& in) {
+          return in.is_tracer();
+        })) {
+      return fun(inputs);
+    }
+
+    // Find a cache entry with the correct inputs
+    auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants);
+
+    // No matching cache entry existed, so compile
+    if (entry.empty) {
+      // Mark the entry as not empty since we are about to fill it
+      entry.empty = false;
+      // Set the constants
+      entry.constants = std::move(constants);
+      // Trace to build the graph
+      std::tie(entry.inputs, entry.outputs, entry.extra) =
+          compile_trace(fun, inputs, shapeless);
+
+      // DFS the graph and get a tape, and a map of array id to (parent,
+      // position in parent inputs)
+      std::unordered_map<uintptr_t, std::vector<std::pair<array, int>>>
+          parents_map;
+      std::tie(entry.tape, parents_map) =
+          compile_dfs(entry.inputs, entry.outputs, inputs);
+
+      // Simplify the tape
+      if (compile_mode() != CompileMode::no_simplify) {
+        compile_simplify(
+            entry.tape, parents_map, entry.outputs, /* passes */ 3);
+      }
+
+      // Kernel fusion to generate Compiled primitives. The tape and
+      // new outputs must be updated accordingly
+      if (compile_mode() != CompileMode::no_fuse) {
+        compile_fuse(entry.tape, parents_map, entry.inputs, entry.outputs);
+      }
+    }
+
+    // At this point we must have a tape, now replace the placeholders
+    // with real arrays that can be evaluated
+    return ArraysAndExtra{
+        compile_replace(
+            entry.tape, entry.inputs, entry.outputs, inputs, shapeless),
+        entry.extra};
+  };
+}
+
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::uintptr_t fun_id,
+    bool shapeless /* = false */,
+    std::vector<uint64_t> constants /* = {} */) {
+  if (skip_compile()) {
+    return fun;
+  }
+  if (!fun) {
+    throw std::invalid_argument(
+        "[compile] Cannot compile a function without a target.");
+  }
+
+  ArrayFnWithExtra fun_with_extra =
+      [fun = std::move(fun)](const std::vector<array>& inputs) {
+        return ArraysAndExtra{fun(inputs), nullptr};
+      };
+
+  auto compiled_fun = compile(
+      std::move(fun_with_extra), fun_id, shapeless, std::move(constants));
+
+  return [compiled_fun =
+              std::move(compiled_fun)](const std::vector<array>& inputs) {
+    return compiled_fun(inputs).first;
+  };
+}
+
+void compile_erase(std::uintptr_t fun_id) {
+  detail::compiler_cache().erase(fun_id);
+}
+
+void compile_clear_cache() {
+  detail::compiler_cache().clear();
+}
+
+} // namespace detail
+
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    bool shapeless /* false */) {
+  if (detail::skip_compile()) {
+    return fun;
+  }
+  auto fun_id = detail::get_function_address(fun);
+  if (fun_id) {
+    // If the function has an addressable target then no need to manage it's
+    // lifetime
+    return detail::compile(std::move(fun), fun_id, shapeless);
+  } else {
+    auto pfun = std::shared_ptr<
+        std::function<std::vector<array>(const std::vector<array>&)>>(
+        new std::function<std::vector<array>(const std::vector<array>&)>{fun},
+        [](auto* p) {
+          detail::compile_erase(reinterpret_cast<std::uintptr_t>(p));
+          delete p;
+        });
+    fun_id = reinterpret_cast<std::uintptr_t>(pfun.get());
+    return detail::compile(
+        [pfun = std::move(pfun)](const auto& inputs) {
+          return (*pfun)(inputs);
+        },
+        fun_id,
+        shapeless);
+  }
+}
+
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::vector<array> (*fun)(const std::vector<array>&),
+    bool shapeless /* = false */) {
+  if (detail::skip_compile()) {
+    return fun;
+  }
+  return detail::compile(fun, reinterpret_cast<std::uintptr_t>(fun), shapeless);
+}
+
+void disable_compile() {
+  detail::compile_mode() = CompileMode::disabled;
+}
+
+void enable_compile() {
+  detail::compile_mode() = CompileMode::enabled;
+}
+
+void set_compile_mode(CompileMode mode) {
+  detail::compile_mode() = mode;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/compile.h b/macos/third_party/mlx/mlx/compile.h
new file mode 100644
index 0000000..eba0983
--- /dev/null
+++ b/macos/third_party/mlx/mlx/compile.h
@@ -0,0 +1,45 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+enum class CompileMode { disabled, no_simplify, no_fuse, enabled };
+
+/** Compile takes a function and returns a compiled function. */
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    bool shapeless = false);
+
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::vector<array> (*fun)(const std::vector<array>&),
+    bool shapeless = false);
+
+// Convert capture-less lambdas to function pointers.
+template <
+    typename F,
+    typename = std::enable_if_t<
+        std::is_convertible_v<F, decltype(+std::declval<F>())>>>
+std::function<std::vector<array>(const std::vector<array>&)> compile(
+    F&& f,
+    bool shapeless = false) {
+  return compile(+f, shapeless);
+}
+
+/** Globally disable compilation.
+ * Setting the environment variable ``MLX_DISABLE_COMPILE`` can also
+ * be used to disable compilation.
+ */
+MLX_API void disable_compile();
+
+/** Globally enable compilation.
+ * This will override the environment variable ``MLX_DISABLE_COMPILE``.
+ */
+MLX_API void enable_compile();
+
+/** Set the compiler mode to the given value. */
+MLX_API void set_compile_mode(CompileMode mode);
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/compile_impl.h b/macos/third_party/mlx/mlx/compile_impl.h
new file mode 100644
index 0000000..238a8b9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/compile_impl.h
@@ -0,0 +1,70 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core::detail {
+
+using ArraysAndExtra = std::pair<std::vector<array>, std::shared_ptr<void>>;
+using ArrayFnWithExtra =
+    std::function<ArraysAndExtra(const std::vector<array>&)>;
+
+// This is not part of the general C++ API as calling with a bad id is a bad
+// idea.
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> compile(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::uintptr_t fun_id,
+    bool shapeless = false,
+    std::vector<uint64_t> constants = {});
+
+MLX_API ArrayFnWithExtra compile(
+    ArrayFnWithExtra fun,
+    std::uintptr_t fun_id,
+    bool shapeless,
+    std::vector<uint64_t> constants);
+
+// Erase cached compile functions
+MLX_API void compile_erase(std::uintptr_t fun_id);
+
+// Clear the compiler cache causing a recompilation of all compiled functions
+// when called again.
+MLX_API void compile_clear_cache();
+
+bool compile_available_for_device(const Device& device);
+
+std::tuple<std::vector<array>, std::vector<array>, std::shared_ptr<void>>
+compile_trace(
+    const ArrayFnWithExtra& fun,
+    const std::vector<array>& inputs,
+    bool shapeless);
+
+using ParentsMap =
+    std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>;
+
+// Traverses the graph to build a tape and a map of array ids to their parents
+std::pair<std::vector<array>, ParentsMap> compile_dfs(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs,
+    const std::vector<array>& original_inputs);
+
+// Simplify the tape.
+void compile_simplify(
+    std::vector<array>& tape,
+    ParentsMap& parents_map,
+    std::vector<array>& outputs,
+    int passes);
+
+std::vector<array> compile_replace(
+    const std::vector<array>& tape,
+    const std::vector<array>& trace_inputs,
+    const std::vector<array>& trace_outputs,
+    const std::vector<array>& inputs,
+    bool shapeless);
+
+void compile_validate_shapeless(const std::vector<array>& tape);
+
+} // namespace mlx::core::detail
diff --git a/macos/third_party/mlx/mlx/device.cpp b/macos/third_party/mlx/mlx/device.cpp
new file mode 100644
index 0000000..4a62036
--- /dev/null
+++ b/macos/third_party/mlx/mlx/device.cpp
@@ -0,0 +1,72 @@
+// Copyright © 2023-2026 Apple Inc.
+
+#include <stdexcept>
+
+#include "mlx/backend/cpu/device_info.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/device.h"
+
+namespace mlx::core {
+
+Device& mutable_default_device() {
+  static Device default_device{gpu::is_available() ? Device::gpu : Device::cpu};
+  return default_device;
+}
+
+const Device& default_device() {
+  return mutable_default_device();
+}
+
+void set_default_device(const Device& d) {
+  if (!gpu::is_available() && d == Device::gpu) {
+    throw std::invalid_argument(
+        "[set_default_device] Cannot set gpu device without gpu backend.");
+  }
+  mutable_default_device() = d;
+}
+
+bool operator==(const Device& lhs, const Device& rhs) {
+  return lhs.type == rhs.type && lhs.index == rhs.index;
+}
+
+bool operator!=(const Device& lhs, const Device& rhs) {
+  return !(lhs == rhs);
+}
+
+bool is_available(const Device& d) {
+  switch (d.type) {
+    case Device::cpu:
+      return cpu::is_available();
+    case Device::gpu:
+      return gpu::is_available();
+  }
+  // appease compiler
+  return false;
+}
+
+int device_count(Device::DeviceType type) {
+  switch (type) {
+    case Device::cpu:
+      return cpu::device_count();
+    case Device::gpu:
+      return gpu::device_count();
+  }
+  // appease compiler
+  return 0;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(const Device& d) {
+  switch (d.type) {
+    case Device::cpu:
+      return cpu::device_info(d.index);
+    case Device::gpu:
+      return gpu::device_info(d.index);
+  }
+  // appease compiler
+  static std::unordered_map<std::string, std::variant<std::string, size_t>>
+      empty;
+  return empty;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/device.h b/macos/third_party/mlx/mlx/device.h
new file mode 100644
index 0000000..f89ad18
--- /dev/null
+++ b/macos/third_party/mlx/mlx/device.h
@@ -0,0 +1,56 @@
+// Copyright © 2023-2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/api.h"
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+namespace mlx::core {
+
+struct MLX_API Device {
+  enum class DeviceType {
+    cpu,
+    gpu,
+  };
+
+  static constexpr DeviceType cpu = DeviceType::cpu;
+  static constexpr DeviceType gpu = DeviceType::gpu;
+
+  Device(DeviceType type, int index = 0) : type(type), index(index) {}
+
+  DeviceType type;
+  int index;
+};
+
+MLX_API const Device& default_device();
+
+MLX_API void set_default_device(const Device& d);
+
+MLX_API bool operator==(const Device& lhs, const Device& rhs);
+MLX_API bool operator!=(const Device& lhs, const Device& rhs);
+
+MLX_API bool is_available(const Device& d);
+
+/** Get the number of available devices for the given device type. */
+MLX_API int device_count(Device::DeviceType type);
+
+/**
+ * Get information about a device.
+ *
+ * Returns a map of device properties. Keys vary by backend:
+ *   - device_name (string): Device name
+ *   - architecture (string): Architecture identifier
+ *   - total_memory/memory_size (size_t): Total device memory
+ *   - free_memory (size_t): Available memory (CUDA only)
+ *   - uuid (string): Device UUID (CUDA only)
+ *   - pci_bus_id (string): PCI bus ID (CUDA only)
+ *   - compute_capability_major/minor (size_t): Compute capability (CUDA only)
+ */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(const Device& d = default_device());
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/distributed/CMakeLists.txt b/macos/third_party/mlx/mlx/distributed/CMakeLists.txt
new file mode 100644
index 0000000..807fd0f
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/CMakeLists.txt
@@ -0,0 +1,14 @@
+target_sources(
+  mlx
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cpp)
+
+if(MLX_BUILD_CPU AND NOT WIN32)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
+endif()
+
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/mpi)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/ring)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/nccl)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/jaccl)
diff --git a/macos/third_party/mlx/mlx/distributed/distributed.cpp b/macos/third_party/mlx/mlx/distributed/distributed.cpp
new file mode 100644
index 0000000..3cde6a2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/distributed.cpp
@@ -0,0 +1,197 @@
+// Copyright © 2024 Apple Inc.
+
+#include <unordered_map>
+
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/jaccl/jaccl.h"
+#include "mlx/distributed/mpi/mpi.h"
+#include "mlx/distributed/nccl/nccl.h"
+#include "mlx/distributed/ring/ring.h"
+
+namespace mlx::core::distributed {
+
+namespace detail {
+
+Stream communication_stream(Group group, StreamOrDevice s /* = {} */) {
+  return group.raw_group()->communication_stream(s);
+}
+
+void all_sum(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_sum(input, output, stream);
+}
+
+void all_max(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_max(input, output, stream);
+}
+
+void all_min(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_min(input, output, stream);
+}
+
+void all_gather(Group group, const array& input, array& output, Stream stream) {
+  group.raw_group()->all_gather(input, output, stream);
+}
+
+void send(Group group, const array& input, int dst, Stream stream) {
+  group.raw_group()->send(input, dst, stream);
+}
+
+void recv(Group group, array& out, int src, Stream stream) {
+  group.raw_group()->recv(out, src, stream);
+}
+
+void sum_scatter(
+    Group group,
+    const array& input,
+    array& output,
+    Stream stream) {
+  group.raw_group()->sum_scatter(input, output, stream);
+}
+
+class EmptyGroup : public GroupImpl {
+ public:
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s);
+  }
+
+  int rank() override {
+    return 0;
+  }
+
+  int size() override {
+    return 1;
+  }
+
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    throw std::runtime_error("Cannot split the distributed group further.");
+  }
+
+  void all_sum(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void all_gather(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void send(const array&, int, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void recv(array&, int, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+
+  void all_max(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+
+  void all_min(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+  void sum_scatter(const array&, array&, Stream) override {
+    throw std::runtime_error(
+        "Communication not implemented in an empty distributed group.");
+  }
+};
+
+} // namespace detail
+
+bool is_available() {
+  return mpi::is_available() || ring::is_available() || nccl::is_available() ||
+      jaccl::is_available();
+}
+
+bool is_available(const std::string& bk) {
+  if (bk == "any") {
+    return is_available();
+  }
+  if (bk == "mpi") {
+    return mpi::is_available();
+  }
+  if (bk == "ring") {
+    return ring::is_available();
+  }
+  if (bk == "nccl") {
+    return nccl::is_available();
+  }
+  if (bk == "jaccl") {
+    return jaccl::is_available();
+  }
+  return false;
+}
+
+int Group::rank() const {
+  return group_->rank();
+}
+
+int Group::size() const {
+  return group_->size();
+}
+
+Group Group::split(int color, int key /* = -1 */) const {
+  return Group(group_->split(color, key));
+}
+
+Group init(bool strict /* = false */, const std::string& bk /* = "any" */) {
+  static std::unordered_map<std::string, std::shared_ptr<detail::GroupImpl>>
+      backends;
+
+  // Already initialized so return the group.
+  if (auto g = backends.find(bk); g != backends.end()) {
+    return Group(g->second);
+  }
+
+  // Create the requested communication group
+  std::shared_ptr<detail::GroupImpl> group{nullptr};
+  std::string bk_ = bk;
+  if (bk == "mpi") {
+    group = mpi::init(strict);
+  } else if (bk == "ring") {
+    group = ring::init(strict);
+  } else if (bk == "nccl") {
+    group = nccl::init(strict);
+  } else if (bk == "jaccl") {
+    group = jaccl::init(strict);
+  } else if (bk == "any") {
+    if (mlx::core::cu::is_available()) {
+      group = nccl::init(false);
+      bk_ = "nccl";
+    }
+    if (group == nullptr) {
+      group = ring::init(false);
+      bk_ = "ring";
+    }
+    if (group == nullptr) {
+      group = mpi::init(false);
+      bk_ = "mpi";
+    }
+    if (group == nullptr) {
+      group = jaccl::init(false);
+      bk_ = "jaccl";
+    }
+    if (group == nullptr && strict) {
+      throw std::runtime_error("[distributed] Couldn't initialize any backend");
+    }
+  } else {
+    std::ostringstream msg;
+    msg << "[distributed] The only valid values for backend are 'any', 'mpi', 'nccl', "
+        << "'jaccl' and 'ring' but '" << bk << "' was provided.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (group == nullptr) {
+    group = std::make_shared<detail::EmptyGroup>();
+  } else {
+    backends.insert({"any", group});
+  }
+  backends.insert({std::move(bk_), group});
+  return Group(group);
+}
+
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/distributed/distributed.h b/macos/third_party/mlx/mlx/distributed/distributed.h
new file mode 100644
index 0000000..00c7a80
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/distributed.h
@@ -0,0 +1,61 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <memory>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::distributed {
+
+// Forward declaration of the base group implementation.
+namespace detail {
+class GroupImpl;
+};
+
+/* Check if a communication backend is available */
+MLX_API bool is_available();
+MLX_API bool is_available(const std::string& bk);
+
+/**
+ * A distributed::Group represents a group of independent mlx processes that
+ * can communicate. We must also be able to create sub-groups from a group in
+ * order to define more granular communication.
+ */
+struct MLX_API Group {
+  Group(std::shared_ptr<detail::GroupImpl> group) : group_(std::move(group)) {}
+
+  int rank() const;
+  int size() const;
+
+  /**
+   * Split the group according to the provided color. Namely processes that use
+   * the same color will go to the same group.
+   *
+   * The key defines the rank of the processes in the new group. The smaller
+   * the key the smaller the rank. If the provided key is negative, then the
+   * rank in the current group is used.
+   */
+  Group split(int color, int key = -1) const;
+
+  const std::shared_ptr<detail::GroupImpl>& raw_group() const {
+    return group_;
+  }
+
+ private:
+  std::shared_ptr<detail::GroupImpl> group_{nullptr};
+};
+
+/**
+ * Initialize the distributed backend and return the group containing all
+ * discoverable processes.
+ *
+ * If strict is true then throw an error if we couldn't initialize the
+ * distributed subsystem. Otherwise simply return a singleton group which will
+ * render communication operations as no-op.
+ */
+MLX_API Group init(bool strict = false, const std::string& bk = "any");
+
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/distributed/distributed_impl.h b/macos/third_party/mlx/mlx/distributed/distributed_impl.h
new file mode 100644
index 0000000..d889587
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/distributed_impl.h
@@ -0,0 +1,59 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::detail {
+
+/**
+ * Abstract base class of a distributed group implementation.
+ */
+class GroupImpl {
+ public:
+  virtual ~GroupImpl() {}
+
+  // Choose the stream this communication group can operate on
+  virtual Stream communication_stream(StreamOrDevice s = {}) = 0;
+
+  // Group operations
+  virtual int rank() = 0;
+  virtual int size() = 0;
+  virtual std::shared_ptr<GroupImpl> split(int color, int key = -1) = 0;
+
+  // Actual communication operations
+  virtual void all_sum(const array& input, array& output, Stream stream) = 0;
+  virtual void all_gather(const array& input, array& output, Stream stream) = 0;
+  virtual void send(const array& input, int dst, Stream stream) = 0;
+  virtual void recv(array& out, int src, Stream stream) = 0;
+  virtual void all_max(const array& input, array& output, Stream stream) = 0;
+  virtual void all_min(const array& input, array& output, Stream stream) = 0;
+  virtual void
+  sum_scatter(const array& input, array& output, Stream stream) = 0;
+};
+
+/* Define the MLX stream that the communication should happen in. */
+Stream communication_stream(Group group, StreamOrDevice s = {});
+
+/* Perform an all reduce sum operation */
+void all_sum(Group group, const array& input, array& output, Stream stream);
+
+/* Perform an all gather operation */
+void all_gather(Group group, const array& input, array& output, Stream stream);
+
+/** Send an array to the dst rank */
+void send(Group group, const array& input, int dst, Stream stream);
+
+/** Recv an array from the src rank */
+void recv(Group group, array& out, int src, Stream stream);
+
+/** Max reduction */
+void all_max(Group group, const array& input, array& output, Stream stream);
+
+/** Min reduction */
+void all_min(Group group, const array& input, array& output, Stream stream);
+
+/** Reduce scatter with average operation */
+void sum_scatter(Group group, const array& input, array& output, Stream stream);
+
+} // namespace mlx::core::distributed::detail
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/CMakeLists.txt b/macos/third_party/mlx/mlx/distributed/jaccl/CMakeLists.txt
new file mode 100644
index 0000000..567ed67
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/CMakeLists.txt
@@ -0,0 +1,12 @@
+if(MLX_BUILD_CPU
+   AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
+   AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
+  target_sources(
+    mlx
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/mesh.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/ring.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_jaccl.cpp)
+endif()
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/jaccl.cpp b/macos/third_party/mlx/mlx/distributed/jaccl/jaccl.cpp
new file mode 100644
index 0000000..5935a98
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/jaccl.cpp
@@ -0,0 +1,178 @@
+// Copyright © 2025 Apple Inc.
+
+#include <fstream>
+#include <sstream>
+
+#include <json.hpp>
+
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/jaccl/mesh.h"
+#include "mlx/distributed/jaccl/ring.h"
+#include "mlx/distributed/jaccl/utils.h"
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+using json = nlohmann::json;
+
+namespace {
+
+struct DeviceFile {
+  DeviceFile(const char* dev_file) {
+    std::ifstream f(dev_file);
+    json devices = json::parse(f);
+    if (!devices.is_array()) {
+      throw std::runtime_error(
+          "[jaccl] The device file should start with an array");
+    }
+
+    devices_.resize(devices.size());
+    for (int rank = 0; rank < devices.size(); rank++) {
+      auto conn = devices[rank];
+      if (!conn.is_array()) {
+        throw std::runtime_error(
+            "[jaccl] The device file should have an array of arrays");
+      }
+      if (conn.size() != devices_.size()) {
+        std::ostringstream msg;
+        msg << "[jaccl] The device file should contain the connectivity of each rank to "
+            << "all other ranks but rank " << rank << " contains only "
+            << conn.size() << " entries.";
+        throw std::runtime_error(msg.str());
+      }
+
+      devices_[rank].resize(conn.size());
+      for (int dst = 0; dst < conn.size(); dst++) {
+        auto names = conn[dst];
+        if (names.is_string()) {
+          devices_[rank][dst].push_back(names);
+        } else if (names.is_array()) {
+          for (auto name_it = names.begin(); name_it != names.end();
+               name_it++) {
+            devices_[rank][dst].push_back(*name_it);
+          }
+        } else if (!names.is_null()) {
+          throw std::runtime_error(
+              "[jaccl] Device names should be null, a string or array of strings.");
+        }
+      }
+    }
+  }
+
+  int size() {
+    return devices_.size();
+  }
+
+  bool is_valid_mesh() {
+    for (int src = 0; src < size(); src++) {
+      for (int dst = 0; dst < size(); dst++) {
+        if (devices_[src][dst].size() != static_cast<size_t>(src != dst)) {
+          return false;
+        }
+      }
+    }
+
+    return true;
+  }
+
+  bool is_valid_ring() {
+    int num_connections = devices_[0][1].size();
+    if (num_connections == 0) {
+      return false;
+    }
+
+    for (int src = 0; src < size(); src++) {
+      int left = (src + size() - 1) % size();
+      int right = (src + 1) % size();
+      for (int dst = 0; dst < size(); dst++) {
+        if (dst != left && dst != right) {
+          if (devices_[src][dst].size() != 0) {
+            return false;
+          }
+        } else {
+          if (devices_[src][dst].size() != num_connections) {
+            return false;
+          }
+        }
+      }
+    }
+
+    return true;
+  }
+
+  std::vector<std::string> extract_mesh_connectivity(int rank) {
+    std::vector<std::string> devices(size());
+    for (int dst = 0; dst < size(); dst++) {
+      if (dst != rank) {
+        devices[dst] = devices_[rank][dst][0];
+      }
+    }
+    return devices;
+  }
+
+  std::pair<std::vector<std::string>, std::vector<std::string>>
+  extract_ring_connectivity(int rank) {
+    int left = (rank + size() - 1) % size();
+    int right = (rank + 1) % size();
+
+    return std::make_pair(devices_[rank][left], devices_[rank][right]);
+  }
+
+  std::vector<std::vector<std::vector<std::string>>> devices_;
+};
+
+} // namespace
+
+namespace mlx::core::distributed::jaccl {
+
+bool is_available() {
+  return ibv().is_available();
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  const char* dev_file = std::getenv("MLX_IBV_DEVICES");
+  const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR");
+  const char* rank_str = std::getenv("MLX_RANK");
+  const char* ring = std::getenv("MLX_JACCL_RING");
+
+  if (!is_available() || !dev_file || !coordinator || !rank_str) {
+    if (strict) {
+      std::ostringstream msg;
+      msg << "[jaccl] You need to provide via environment variables a rank (MLX_RANK), "
+          << "a device file (MLX_IBV_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) "
+          << "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "")
+          << "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "")
+          << "\" and MLX_JACCL_COORDINATOR=\""
+          << ((coordinator) ? coordinator : "");
+      throw std::runtime_error(msg.str());
+    }
+    return nullptr;
+  }
+
+  auto rank = std::atoi(rank_str);
+  bool prefer_ring = ring != nullptr;
+  DeviceFile devices(dev_file);
+
+  if (rank >= devices.size() || rank < 0) {
+    std::ostringstream msg;
+    msg << "[jaccl] Invalid rank " << rank << ". It should be between 0 and "
+        << devices.size();
+    throw std::runtime_error(msg.str());
+  }
+
+  if (prefer_ring && devices.is_valid_ring()) {
+    auto [left, right] = devices.extract_ring_connectivity(rank);
+    return std::make_shared<RingGroup>(
+        rank, devices.size(), left, right, coordinator);
+  } else if (devices.is_valid_mesh()) {
+    auto device_names = devices.extract_mesh_connectivity(rank);
+    return std::make_shared<MeshGroup>(rank, device_names, coordinator);
+  } else if (devices.is_valid_ring()) {
+    auto [left, right] = devices.extract_ring_connectivity(rank);
+    return std::make_shared<RingGroup>(
+        rank, devices.size(), left, right, coordinator);
+  } else {
+    throw std::runtime_error(
+        "[jaccl] The device file should define a valid mesh or a valid ring.");
+  }
+}
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/jaccl.h b/macos/third_party/mlx/mlx/distributed/jaccl/jaccl.h
new file mode 100644
index 0000000..d07f9cc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/jaccl.h
@@ -0,0 +1,12 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::jaccl {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available();
+std::shared_ptr<GroupImpl> init(bool strict = false);
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/mesh.cpp b/macos/third_party/mlx/mlx/distributed/jaccl/mesh.cpp
new file mode 100644
index 0000000..c8df4e6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/mesh.cpp
@@ -0,0 +1,451 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/distributed/jaccl/mesh.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/distributed/reduction_ops.h"
+#include "mlx/dtype_utils.h"
+
+constexpr int MAX_PEERS = 8;
+
+namespace mlx::core::distributed::jaccl {
+
+MeshGroup::MeshGroup(
+    int rank,
+    const std::vector<std::string>& device_names,
+    const char* coordinator_addr)
+    : rank_(rank),
+      size_(device_names.size()),
+      side_channel_(rank_, size_, coordinator_addr),
+      connections_(create_connections(device_names)) {
+  if (size_ > MAX_PEERS) {
+    std::ostringstream msg;
+    msg << "[jaccl] The JACCL mesh supports up to " << MAX_PEERS
+        << " peers but " << size_ << " were provided.";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Initialize all the connections and allocate buffers
+  initialize();
+
+  // Make sure every node has reached here before continuing
+  side_channel_.all_gather<int>(0);
+}
+
+void MeshGroup::initialize() {
+  // Create the queue pairs
+  for (auto& conn : connections_) {
+    if (conn.ctx == nullptr) {
+      continue;
+    }
+    conn.allocate_protection_domain();
+    conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
+    conn.create_queue_pair();
+  }
+
+  allocate_buffers();
+
+  // First init all connections
+  for (int peer = 0; peer < size_; peer++) {
+    if (peer == rank_) {
+      continue;
+    }
+    connections_[peer].queue_pair_init();
+  }
+
+  // Gather the information to be exchanged, this also serves as a barrier so
+  // that all peers have initialized their connections before attempting to
+  // transition to RTS.
+  std::vector<Destination> info;
+  for (auto& conn : connections_) {
+    info.emplace_back(conn.info());
+  }
+  auto all_infos = side_channel_.all_gather(info);
+
+  // Transition queue pairs to RTS
+  for (int peer = 0; peer < size_; peer++) {
+    if (peer == rank_) {
+      continue;
+    }
+    auto peer_info = all_infos[peer][rank_];
+    connections_[peer].queue_pair_rtr(peer_info);
+    connections_[peer].queue_pair_rts();
+  }
+}
+
+void MeshGroup::allocate_buffers() {
+  // Deregister any buffers and free the memory
+  buffers_.clear();
+
+  // Allocate the memory
+  for (int k = 0; k < BUFFER_SIZES; k++) {
+    for (int i = 0; i < NUM_BUFFERS; i++) {
+      for (int j = 0; j < size_; j++) {
+        buffers_.emplace_back(FRAME_SIZE * (1 << k));
+      }
+    }
+  }
+
+  for (int k = 0; k < BUFFER_SIZES; k++) {
+    for (int i = 0; i < NUM_BUFFERS; i++) {
+      for (int j = 0; j < size_; j++) {
+        // This is our send buffer so register it with all pds so we can send
+        // it to all connected devices.
+        if (j == rank_) {
+          for (auto& conn : connections_) {
+            if (conn.ctx != nullptr) {
+              buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
+                  .register_to_protection_domain(conn.protection_domain);
+            }
+          }
+        }
+
+        // This is the recv buffer from rank j so register it to rank j's
+        // protection domain.
+        else {
+          buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
+              .register_to_protection_domain(connections_[j].protection_domain);
+        }
+      }
+    }
+  }
+}
+
+void MeshGroup::all_sum(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::SumOp<T>{});
+  });
+}
+
+void MeshGroup::all_max(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::MaxOp<T>{});
+  });
+}
+
+void MeshGroup::all_min(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::MinOp<T>{});
+  });
+}
+
+void MeshGroup::all_gather(const array& input, array& output, Stream stream) {
+  auto in_ptr = input.data<char>();
+  auto out_ptr = output.data<char>();
+  size_t n_bytes = input.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.set_output_array(output);
+  encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() {
+    // Copy our data to the appropriate place
+    std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes);
+
+    // Fully connected all gather
+    char* data = out_ptr;
+    char* our_data = out_ptr + rank_ * n_bytes;
+    auto [sz, N] = buffer_size_from_message(n_bytes);
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2;
+    int64_t total = static_cast<int64_t>(n_bytes);
+    int num_peers = size_ - 1;
+
+    // Counters to maintain the state of transfers
+    int in_flight = 0;
+    int read_offset = 0;
+    int completed_send_count[PIPELINE] = {0};
+    int write_offset[MAX_PEERS] = {0};
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (read_offset < total && buff < PIPELINE) {
+      post_recv_all(sz, buff);
+      std::copy(
+          our_data + read_offset,
+          our_data + std::min(read_offset + N, total),
+          send_buffer(sz, buff).begin<char>());
+      post_send_all(sz, buff);
+
+      buff++;
+      in_flight += 2 * num_peers;
+      read_offset += N;
+    }
+
+    // Main loop
+    //
+    // Keep going until we have no longer data in flight.
+    while (in_flight > 0) {
+      ibv_wc wc[WC_NUM];
+      int n = poll(connections_, WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int work_type = wc[i].wr_id >> 16;
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int rank = wc[i].wr_id & 0xff;
+
+        in_flight--;
+
+        // Send completed. If all sends completed then send the next chunk.
+        if (work_type == SEND_WR && read_offset < total) {
+          completed_send_count[buff]++;
+          if (completed_send_count[buff] == num_peers) {
+            std::copy(
+                our_data + read_offset,
+                our_data + std::min(read_offset + N, total),
+                send_buffer(sz, buff).begin<char>());
+            post_send_all(sz, buff);
+
+            completed_send_count[buff] = 0;
+            in_flight += num_peers;
+            read_offset += N;
+          }
+        }
+
+        // Recv completed. If we have more chunks then post another recv.
+        else if (work_type == RECV_WR) {
+          std::copy(
+              recv_buffer(sz, buff, rank).begin<char>(),
+              recv_buffer(sz, buff, rank).begin<char>() +
+                  std::min(N, total - write_offset[rank]),
+              data + rank * n_bytes + write_offset[rank]);
+          write_offset[rank] += N;
+          if (write_offset[rank] + N * (PIPELINE - 1) < total) {
+            recv_from(sz, rank, buff);
+            in_flight++;
+          }
+        }
+      }
+    }
+  });
+}
+
+void MeshGroup::send(const array& input, int dst, Stream stream) {
+  auto data = input.data<char>();
+  int64_t n_bytes = input.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.dispatch([data, n_bytes, dst, this]() {
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE;
+    auto [sz, N] = buffer_size_from_message(n_bytes);
+
+    int in_flight = 0;
+    int64_t read_offset = 0;
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (read_offset < n_bytes && buff < PIPELINE) {
+      std::copy(
+          data + read_offset,
+          data + std::min(read_offset + N, n_bytes),
+          send_buffer(sz, buff).begin<char>());
+      send_to(sz, dst, buff);
+
+      buff++;
+      read_offset += N;
+      in_flight++;
+    }
+
+    // Main loop
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a send was completed and we have more data to send then go ahead
+      // and send them.
+      ibv_wc wc[WC_NUM];
+      int n = connections_[dst].poll(WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int rank = wc[i].wr_id & 0xff;
+
+        in_flight--;
+
+        if (read_offset < n_bytes) {
+          std::copy(
+              data + read_offset,
+              data + std::min(read_offset + N, n_bytes),
+              send_buffer(sz, buff).begin<char>());
+          send_to(sz, dst, buff);
+
+          read_offset += N;
+          in_flight++;
+        }
+      }
+    }
+  });
+}
+
+void MeshGroup::recv(array& out, int src, Stream stream) {
+  auto data = out.data<char>();
+  int64_t n_bytes = out.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([data, n_bytes, src, this]() {
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE;
+    auto [sz, N] = buffer_size_from_message(n_bytes);
+
+    int in_flight = 0;
+    int64_t write_offset = 0;
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (N * buff < n_bytes && buff < PIPELINE) {
+      recv_from(sz, src, buff);
+
+      in_flight++;
+      buff++;
+    }
+
+    // Main loop
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a recv was completed copy it to the output and if we have more
+      // data to fetch post another recv.
+      ibv_wc wc[WC_NUM];
+      int n = connections_[src].poll(WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int rank = wc[i].wr_id & 0xff;
+
+        in_flight--;
+
+        std::copy(
+            recv_buffer(sz, buff, src).begin<char>(),
+            recv_buffer(sz, buff, src).begin<char>() +
+                std::min(n_bytes - write_offset, static_cast<int64_t>(N)),
+            data + write_offset);
+        write_offset += N;
+
+        if (write_offset + (PIPELINE - 1) * N < n_bytes) {
+          recv_from(sz, src, buff);
+
+          in_flight++;
+        }
+      }
+    }
+  });
+}
+
+template <typename T, typename ReduceOp>
+void MeshGroup::all_reduce(
+    const array& input,
+    array& output,
+    Stream stream,
+    ReduceOp reduce_op) {
+  auto in_ptr = input.data<T>();
+  auto out_ptr = output.data<T>();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.set_output_array(output);
+  encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() {
+    // If not inplace all reduce then copy the input to the output first
+    if (in_ptr != out_ptr) {
+      std::memcpy(out_ptr, in_ptr, size * sizeof(T));
+    }
+
+    // Fully connected all reduce
+    T* data = out_ptr;
+    auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T));
+    int64_t N = buffer_size / sizeof(T);
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2;
+    int64_t total = static_cast<int64_t>(size);
+    int num_peers = size_ - 1;
+
+    // Counters to maintain the state of transfers
+    int in_flight = 0;
+    int64_t read_offset = 0;
+    int completed_send_count[PIPELINE] = {0};
+    int completed_recv_begin[MAX_PEERS] = {0};
+    int completed_recv_end[MAX_PEERS] = {0};
+
+    // Prefill the pipeline
+    int buff = 0;
+    while (read_offset < total && buff < PIPELINE) {
+      post_recv_all(sz, buff);
+      std::copy(
+          data + read_offset,
+          data + std::min(read_offset + N, total),
+          send_buffer(sz, buff).begin<T>());
+      post_send_all(sz, buff);
+
+      buff++;
+      in_flight += 2 * num_peers;
+      read_offset += N;
+    }
+
+    // Main loop
+    //
+    // Keep going until we have no longer data in flight.
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a send was completed mark how many completions we have received
+      // for that buffer. If we have sent the buffer to all peers we can
+      // reuse the buffer so copy the next chunk of data and send it to all.
+      //
+      // If a receive is completed then advance the pointer of completed
+      // receives.
+      ibv_wc wc[WC_NUM];
+      int n = poll(connections_, WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int work_type = wc[i].wr_id >> 16;
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int rank = wc[i].wr_id & 0xff;
+
+        in_flight--;
+
+        if (work_type == SEND_WR && read_offset < total) {
+          completed_send_count[buff]++;
+          if (completed_send_count[buff] == num_peers) {
+            std::copy(
+                data + read_offset,
+                data + std::min(read_offset + N, total),
+                send_buffer(sz, buff).begin<T>());
+            post_send_all(sz, buff);
+
+            completed_send_count[buff] = 0;
+            in_flight += num_peers;
+            read_offset += N;
+          }
+        }
+
+        else if (work_type == RECV_WR) {
+          completed_recv_end[rank]++;
+        }
+      }
+
+      // Process the completed recv
+      //
+      // For each rank we have a range of completed recv defined by a begin
+      // and end inclusive and exlusive in standard C++ fashion.
+      //
+      // When there is an unprocessed receive we first check if we have
+      // finished sending the write location. If so then we reduce in-place
+      // and then check if there is more to be received and post a recv.
+      for (int r = 0; r < size_; r++) {
+        int s = completed_recv_begin[r];
+        int e = completed_recv_end[r];
+        int w = s * N;
+        while (w < read_offset && e - s > 0) {
+          int buff = s % PIPELINE;
+          reduce_op(
+              recv_buffer(sz, buff, r).begin<T>(),
+              data + w,
+              std::min(N, total - w));
+          w += N;
+          s++;
+          if (w + (PIPELINE - 1) * N < total) {
+            recv_from(sz, r, buff);
+            in_flight++;
+          }
+        }
+        completed_recv_begin[r] = s;
+      }
+    }
+  });
+}
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/mesh.h b/macos/third_party/mlx/mlx/distributed/jaccl/mesh.h
new file mode 100644
index 0000000..6f779e9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/mesh.h
@@ -0,0 +1,122 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/jaccl/utils.h"
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+namespace mlx::core::distributed::jaccl {
+
+/**
+ * The JACCL communication group for a fully connected mesh. We expect one
+ * connection per peer and it should be the lowest latency communication group
+ * for small to medium size messages.
+ *
+ * Like all JACCL groups it uses a side channel to exchange the necessary
+ * information and then configure the connections to be ready for RDMA
+ * operations.
+ */
+class MeshGroup : public GroupImpl {
+ public:
+  MeshGroup(
+      int rank,
+      const std::vector<std::string>& device_names,
+      const char* coordinator_addr);
+
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s, Device::cpu);
+  }
+
+  int rank() override {
+    return rank_;
+  }
+
+  int size() override {
+    return size_;
+  }
+
+  void all_sum(const array& input, array& output, Stream stream) override;
+  void all_max(const array& input, array& output, Stream stream) override;
+  void all_min(const array& input, array& output, Stream stream) override;
+  void all_gather(const array& input, array& output, Stream stream) override;
+  void send(const array& input, int dst, Stream stream) override;
+  void recv(array& out, int src, Stream stream) override;
+
+  void sum_scatter(const array& input, array& output, Stream stream) override {
+    throw std::runtime_error("[jaccl] sum_scatter not supported.");
+  }
+
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    throw std::runtime_error("[jaccl] Group split not supported.");
+  }
+
+ private:
+  template <typename T, typename ReduceOp>
+  void all_reduce(
+      const array& input,
+      array& output,
+      Stream stream,
+      ReduceOp reduce_op);
+
+  /**
+   * Performs the connection initialization. Namely, after this call all
+   * Connection objects should have a queue pair in RTS state and all buffers
+   * should have been allocated.
+   */
+  void initialize();
+
+  /**
+   * Allocate all the buffers that we will use in the communication group.
+   */
+  void allocate_buffers();
+
+  void send_to(int sz, int rank, int buff) {
+    connections_[rank].post_send(
+        send_buffer(sz, buff), SEND_WR << 16 | buff << 8 | rank);
+  }
+
+  void recv_from(int sz, int rank, int buff) {
+    connections_[rank].post_recv(
+        recv_buffer(sz, buff, rank), RECV_WR << 16 | buff << 8 | rank);
+  }
+
+  SharedBuffer& send_buffer(int sz, int buff) {
+    return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank_];
+  }
+
+  SharedBuffer& recv_buffer(int sz, int buff, int rank) {
+    return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank];
+  }
+
+  void post_send_all(int sz, int buff) {
+    auto& b = send_buffer(sz, buff);
+    int wr_id = SEND_WR << 16 | buff << 8;
+    for (int i = 0; i < size_; i++) {
+      if (i == rank_) {
+        continue;
+      }
+      connections_[i].post_send(b, wr_id | i);
+    }
+  }
+
+  void post_recv_all(int sz, int buff) {
+    int b = sz * NUM_BUFFERS * size_ + buff * size_;
+    int wr_id = RECV_WR << 16 | buff << 8;
+    for (int i = 0; i < size_; i++) {
+      if (i == rank_) {
+        continue;
+      }
+      connections_[i].post_recv(buffers_[b + i], wr_id | i);
+    }
+  }
+
+  int rank_;
+  int size_;
+  SideChannel side_channel_;
+  std::vector<Connection> connections_;
+  std::vector<SharedBuffer> buffers_;
+};
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/no_jaccl.cpp b/macos/third_party/mlx/mlx/distributed/jaccl/no_jaccl.cpp
new file mode 100644
index 0000000..12fd6ab
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/no_jaccl.cpp
@@ -0,0 +1,20 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/distributed/jaccl/jaccl.h"
+
+namespace mlx::core::distributed::jaccl {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available() {
+  return false;
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  if (strict) {
+    throw std::runtime_error("Cannot initialize jaccl distributed backend.");
+  }
+  return nullptr;
+}
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/ring.cpp b/macos/third_party/mlx/mlx/distributed/jaccl/ring.cpp
new file mode 100644
index 0000000..fe27a2e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/ring.cpp
@@ -0,0 +1,692 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/distributed/jaccl/ring.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/distributed/reduction_ops.h"
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core::distributed::jaccl {
+
+RingGroup::RingGroup(
+    int rank,
+    int size,
+    const std::vector<std::string>& left_devices,
+    const std::vector<std::string>& right_devices,
+    const char* coordinator_addr)
+    : rank_(rank),
+      size_(size),
+      side_channel_(rank_, size_, coordinator_addr),
+      left_(create_connections(left_devices)),
+      right_(create_connections(right_devices)) {
+  if (left_.size() > MAX_CONNS || right_.size() > MAX_CONNS) {
+    std::ostringstream msg;
+    msg << "[jaccl] Up to " << MAX_CONNS << " per direction supported but "
+        << left_.size() << " were provided.";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Initialize all the connections and allocate buffers
+  initialize();
+
+  // Make sure every node has reached here before continuing
+  side_channel_.all_gather<int>(0);
+}
+
+void RingGroup::initialize() {
+  // Create the queue pairs
+  for (auto& conn : left_) {
+    conn.allocate_protection_domain();
+    conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
+    conn.create_queue_pair();
+  }
+  for (auto& conn : right_) {
+    conn.allocate_protection_domain();
+    conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
+    conn.create_queue_pair();
+  }
+
+  // Allocate the buffers
+  allocate_buffers();
+
+  // Initialize the conections
+  for (auto& conn : left_) {
+    conn.queue_pair_init();
+  }
+  for (auto& conn : right_) {
+    conn.queue_pair_init();
+  }
+
+  // Gather the information to be exchanged, this also serves as a barrier so
+  // that all peers have initialized their connections before attempting to
+  // transition to RTS.
+  std::vector<Destination> left_info;
+  for (auto& conn : left_) {
+    left_info.emplace_back(conn.info());
+  }
+  std::vector<Destination> right_info;
+  for (auto& conn : right_) {
+    right_info.emplace_back(conn.info());
+  }
+  auto all_left_infos = side_channel_.all_gather(left_info);
+  auto all_right_infos = side_channel_.all_gather(right_info);
+
+  // Transition queue pairs to RTS
+  int left_peer = (rank_ + size_ - 1) % size_;
+  for (int i = 0; i < left_.size(); i++) {
+    auto peer_info = all_right_infos[left_peer][i];
+    left_[i].queue_pair_rtr(peer_info);
+    left_[i].queue_pair_rts();
+  }
+  int right_peer = (rank_ + 1) % size_;
+  for (int i = 0; i < right_.size(); i++) {
+    auto peer_info = all_left_infos[right_peer][i];
+    right_[i].queue_pair_rtr(peer_info);
+    right_[i].queue_pair_rts();
+  }
+}
+
+void RingGroup::allocate_buffers() {
+  // Deregister any buffers and free the memory
+  send_buffers_.clear();
+  recv_buffers_.clear();
+
+  // Allocate the memory
+  for (int k = 0; k < BUFFER_SIZES; k++) {
+    for (int i = 0; i < NUM_BUFFERS; i++) {
+      for (int j = 0; j < MAX_CONNS * 2; j++) {
+        send_buffers_.emplace_back(FRAME_SIZE * (1 << k));
+        recv_buffers_.emplace_back(FRAME_SIZE * (1 << k));
+      }
+    }
+  }
+
+  // Register the buffers with the corresponding connections
+  for (int k = 0; k < BUFFER_SIZES; k++) {
+    for (int i = 0; i < NUM_BUFFERS; i++) {
+      for (int j = 0; j < MAX_CONNS * 2; j++) {
+        int wire = j % MAX_CONNS;
+        int lr = j / MAX_CONNS;
+        if (wire >= left_.size()) {
+          continue;
+        }
+        if (lr) {
+          send_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
+              .register_to_protection_domain(left_[wire].protection_domain);
+          recv_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
+              .register_to_protection_domain(right_[wire].protection_domain);
+        } else {
+          send_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
+              .register_to_protection_domain(right_[wire].protection_domain);
+          recv_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
+              .register_to_protection_domain(left_[wire].protection_domain);
+        }
+      }
+    }
+  }
+}
+
+void RingGroup::all_sum(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::SumOp<T>{});
+  });
+}
+
+void RingGroup::all_max(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::MaxOp<T>{});
+  });
+}
+
+void RingGroup::all_min(const array& input, array& output, Stream stream) {
+  dispatch_all_types(output.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    all_reduce<T>(input, output, stream, detail::MinOp<T>{});
+  });
+}
+
+void RingGroup::all_gather(const array& input, array& output, Stream stream) {
+  auto in_ptr = input.data<char>();
+  auto out_ptr = output.data<char>();
+  size_t n_bytes = input.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.set_output_array(output);
+  encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() {
+    // Copy our data to the appropriate place
+    std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes);
+
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE * MAX_CONNS * 2 * 2;
+    int n_wires = left_.size();
+    size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires);
+    size_t out_bytes = n_bytes * size_;
+    auto [sz, N] = buffer_size_from_message(n_bytes_per_wire);
+    int n_steps = (n_bytes_per_wire + N - 1) / N;
+
+    // Counters to maintain the state of transfers
+    int in_flight = 0;
+    int64_t send_offset[2];
+    int64_t recv_offset[2];
+    int64_t limits[2];
+    int send_count[2 * MAX_CONNS] = {0};
+    int recv_count[2 * MAX_CONNS] = {0};
+    send_offset[0] = send_offset[1] = rank_ * n_bytes;
+    recv_offset[0] = ((rank_ + size_ - 1) % size_) * n_bytes;
+    recv_offset[1] = ((rank_ + 1) % size_) * n_bytes;
+    limits[0] = n_wires * n_bytes_per_wire;
+    limits[1] = n_bytes;
+
+    // Possible perf improvement by not syncing at every step but running ahead
+    // as needed.
+    for (int k = 0; k < size_ - 1; k++) {
+      // Prefill the pipeline
+      int buff = 0;
+      while (buff < n_steps && buff < PIPELINE) {
+        post_recv_all(sz, buff);
+        for (int lr = 0; lr < 2; lr++) {
+          for (int lw = 0; lw < n_wires; lw++) {
+            int64_t offset = lw * N +
+                send_count[lr * MAX_CONNS + lw] * n_wires * N +
+                lr * n_wires * n_bytes_per_wire;
+            std::copy(
+                out_ptr + send_offset[lr] + offset,
+                out_ptr + send_offset[lr] +
+                    std::max(offset, std::min(offset + N, limits[lr])),
+                send_buffer(sz, buff, lr, lw).begin<char>());
+            send_count[lr * MAX_CONNS + lw]++;
+          }
+        }
+        post_send_all(sz, buff);
+
+        buff++;
+        in_flight += 2 * 2 * n_wires;
+      }
+
+      // Main loop
+      //
+      // Keep going until we have no longer data in flight.
+      while (in_flight > 0) {
+        ibv_wc wc[WC_NUM];
+        int n = poll(left_, right_, WC_NUM, wc);
+        for (int i = 0; i < n; i++) {
+          int work_type = wc[i].wr_id >> 16;
+          int buff = (wc[i].wr_id >> 8) & 0xff;
+          int wire = wc[i].wr_id & 0xff;
+          int lr = wire / MAX_CONNS;
+          int lw = wire % MAX_CONNS;
+
+          in_flight--;
+
+          if (work_type == SEND_WR && send_count[wire] < n_steps) {
+            int64_t offset = lw * N + send_count[wire] * n_wires * N +
+                lr * n_wires * n_bytes_per_wire;
+            std::copy(
+                out_ptr + send_offset[lr] + offset,
+                out_ptr + send_offset[lr] +
+                    std::max(offset, std::min(offset + N, limits[lr])),
+                send_buffer(sz, buff, lr, lw).begin<char>());
+            send_to(sz, buff, lr, lw);
+            in_flight++;
+            send_count[wire]++;
+          }
+
+          else if (work_type == RECV_WR) {
+            int64_t offset = lw * N + recv_count[wire] * n_wires * N +
+                lr * n_wires * n_bytes_per_wire;
+            std::copy(
+                recv_buffer(sz, buff, lr, lw).begin<char>(),
+                recv_buffer(sz, buff, lr, lw).begin<char>() +
+                    std::max<int64_t>(0, std::min(N, limits[lr] - offset)),
+                out_ptr + recv_offset[lr] + offset);
+            recv_count[wire]++;
+            if (recv_count[wire] + (PIPELINE - 1) < n_steps) {
+              recv_from(sz, buff, lr, lw);
+              in_flight++;
+            }
+          }
+        }
+      }
+
+      send_offset[0] = (send_offset[0] + out_bytes - n_bytes) % out_bytes;
+      recv_offset[0] = (recv_offset[0] + out_bytes - n_bytes) % out_bytes;
+      send_offset[1] = (send_offset[1] + n_bytes) % out_bytes;
+      recv_offset[1] = (recv_offset[1] + n_bytes) % out_bytes;
+      for (int i = 0; i < 2 * MAX_CONNS; i++) {
+        send_count[i] = recv_count[i] = 0;
+      }
+    }
+  });
+}
+
+void RingGroup::send(const array& input, int dst, Stream stream) {
+  int right = (rank_ + 1) % size_;
+  int left = (rank_ + size_ - 1) % size_;
+  if (dst != right && dst != left) {
+    std::ostringstream msg;
+    msg << "[jaccl] In ring mode send is only supported to direct neighbors "
+        << "but tried to send to " << dst << " from " << rank_ << std::endl;
+    throw std::runtime_error(msg.str());
+  }
+  auto data = input.data<char>();
+  int64_t n_bytes = input.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.dispatch([data, n_bytes, dst, left, this]() {
+    // In the case that size_ == 2 then left == right so we bias send towards
+    // left and recv towards right so that the selections will be correct for
+    // the 2 node case.
+    auto& conns = (dst == left) ? left_ : right_;
+    int dir = dst == left;
+
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE * MAX_CONNS;
+
+    int n_wires = conns.size();
+    int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires;
+    auto [sz, N] = buffer_size_from_message(bytes_per_wire);
+
+    int in_flight = 0;
+    int64_t read_offset[MAX_CONNS];
+    int64_t limits[MAX_CONNS];
+    for (int lw = 0; lw < n_wires; lw++) {
+      read_offset[lw] = std::min(lw * bytes_per_wire, n_bytes);
+      limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes);
+    }
+
+    // Prefill the pipeline
+    for (int lw = 0; lw < n_wires; lw++) {
+      int buff = 0;
+      while (read_offset[lw] < limits[lw] && buff < PIPELINE) {
+        std::copy(
+            data + read_offset[lw],
+            data + std::min(read_offset[lw] + N, limits[lw]),
+            send_buffer(sz, buff, dir, lw).begin<char>());
+        send_to(sz, buff, dir, lw);
+
+        buff++;
+        read_offset[lw] += N;
+        in_flight++;
+      }
+    }
+
+    // Main loop
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a send was completed and we have more data to send then go ahead
+      // and send them.
+      ibv_wc wc[WC_NUM];
+      int n = poll(conns, WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int wire = wc[i].wr_id & 0xff;
+        int lw = wire % MAX_CONNS;
+
+        in_flight--;
+
+        if (read_offset[lw] < limits[lw]) {
+          std::copy(
+              data + read_offset[lw],
+              data + std::min(read_offset[lw] + N, limits[lw]),
+              send_buffer(sz, buff, dir, lw).begin<char>());
+          send_to(sz, buff, dir, lw);
+
+          read_offset[lw] += N;
+          in_flight++;
+        }
+      }
+    }
+  });
+}
+
+void RingGroup::recv(array& out, int src, Stream stream) {
+  int right = (rank_ + 1) % size_;
+  int left = (rank_ + size_ - 1) % size_;
+  if (src != right && src != left) {
+    std::ostringstream msg;
+    msg << "[jaccl] In ring mode recv is only supported to direct neighbors "
+        << "but tried to recv from " << src << " to " << rank_ << std::endl;
+    throw std::runtime_error(msg.str());
+  }
+  auto data = out.data<char>();
+  int64_t n_bytes = out.nbytes();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_output_array(out);
+  encoder.dispatch([data, n_bytes, src, right, this]() {
+    // In the case that size_ == 2 then left == right so we bias send towards
+    // left and recv towards right so that the selections will be correct for
+    // the 2 node case.
+    auto& conns = (src == right) ? right_ : left_;
+    int dir = src == right;
+
+    constexpr int PIPELINE = 2;
+    constexpr int WC_NUM = PIPELINE * MAX_CONNS;
+
+    int n_wires = conns.size();
+    int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires;
+    auto [sz, N] = buffer_size_from_message(bytes_per_wire);
+
+    int in_flight = 0;
+    int64_t write_offset[MAX_CONNS];
+    int64_t limits[MAX_CONNS];
+    for (int lw = 0; lw < n_wires; lw++) {
+      write_offset[lw] = std::min(lw * bytes_per_wire, n_bytes);
+      limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes);
+    }
+
+    // Prefill the pipeline
+    for (int lw = 0; lw < n_wires; lw++) {
+      int buff = 0;
+      while (N * buff < limits[lw] && buff < PIPELINE) {
+        recv_from(sz, buff, dir, lw);
+
+        buff++;
+        in_flight++;
+      }
+    }
+
+    // Main loop
+    while (in_flight > 0) {
+      // Poll the hardware for completions.
+      //
+      // If a recv was completed copy it to the output and if we have more
+      // data to fetch post another recv.
+      ibv_wc wc[WC_NUM];
+      int n = poll(conns, WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int wire = wc[i].wr_id & 0xff;
+        int lw = wire % MAX_CONNS;
+
+        in_flight--;
+
+        std::copy(
+            recv_buffer(sz, buff, dir, lw).begin<char>(),
+            recv_buffer(sz, buff, dir, lw).begin<char>() +
+                std::max<int64_t>(
+                    0, std::min<int64_t>(limits[lw] - write_offset[lw], N)),
+            data + write_offset[lw]);
+        write_offset[lw] += N;
+
+        if (write_offset[lw] + (PIPELINE - 1) * N < limits[lw]) {
+          recv_from(sz, buff, dir, lw);
+
+          in_flight++;
+        }
+      }
+    }
+  });
+}
+
+template <typename T, typename ReduceOp>
+void RingGroup::all_reduce(
+    const array& input,
+    array& output,
+    Stream stream,
+    ReduceOp reduce_op) {
+  auto in_ptr = input.data<T>();
+  auto out_ptr = output.data<T>();
+  auto& encoder = cpu::get_command_encoder(stream);
+  encoder.set_input_array(input);
+  encoder.set_output_array(output);
+  encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() {
+    if (size < size_ * 2 * left_.size()) {
+      all_reduce_impl<1, T, ReduceOp>(in_ptr, out_ptr, size, 1, reduce_op);
+      return;
+    }
+
+    all_reduce_impl<2, T, ReduceOp>(
+        in_ptr, out_ptr, size, left_.size(), reduce_op);
+  });
+}
+
+template <int MAX_DIR, typename T, typename ReduceOp>
+void RingGroup::all_reduce_impl(
+    const T* in_ptr,
+    T* out_ptr,
+    int64_t size,
+    int n_wires,
+    ReduceOp reduce_op) {
+  // If not inplace all reduce then copy the input to the output first
+  if (in_ptr != out_ptr) {
+    std::memcpy(out_ptr, in_ptr, size * sizeof(T));
+  }
+
+  constexpr int PIPELINE = 2;
+  constexpr int WC_NUM = PIPELINE * MAX_CONNS * 2 * MAX_DIR;
+  int64_t chunk_size = (size + size_ - 1) / size_;
+  int64_t size_per_wire =
+      (chunk_size + (MAX_DIR * n_wires) - 1) / (MAX_DIR * n_wires);
+  auto [sz, N] = buffer_size_from_message(size_per_wire * sizeof(T));
+  N /= sizeof(T);
+  int64_t n_steps = (size_per_wire + N - 1) / N;
+
+  // Counters to maintain the state of transfers
+  int in_flight = 0;
+  int64_t chunk_multiple_size = size_ * chunk_size;
+  int64_t send_offset[MAX_DIR];
+  int64_t recv_offset[MAX_DIR];
+  int64_t send_limits[MAX_DIR];
+  int64_t recv_limits[MAX_DIR];
+  int send_count[MAX_DIR * MAX_CONNS] = {0};
+  int recv_count[MAX_DIR * MAX_CONNS] = {0};
+  send_offset[0] = rank_ * chunk_size;
+  recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size;
+  if constexpr (MAX_DIR == 2) {
+    send_offset[1] = rank_ * chunk_size;
+    recv_offset[1] = ((rank_ + 1) % size_) * chunk_size;
+    send_limits[0] = std::min(
+        n_wires * size_per_wire, std::max<int64_t>(0, size - send_offset[0]));
+    send_limits[1] =
+        std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
+    recv_limits[0] = std::min(
+        n_wires * size_per_wire, std::max<int64_t>(0, size - recv_offset[0]));
+    recv_limits[1] =
+        std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
+  } else {
+    send_limits[0] =
+        std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
+    recv_limits[0] =
+        std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
+  }
+
+  // First reduce scatter
+  //
+  // Possible perf improvement by not syncing at every step but running ahead
+  // as needed.
+  for (int k = 0; k < size_ - 1; k++) {
+    // Prefill the pipeline
+    int buff = 0;
+    while (buff < n_steps && buff < PIPELINE) {
+      post_recv_all<MAX_DIR>(sz, buff, n_wires);
+      for (int lr = 0; lr < MAX_DIR; lr++) {
+        for (int lw = 0; lw < n_wires; lw++) {
+          int64_t offset = lw * N +
+              send_count[lr * MAX_CONNS + lw] * n_wires * N +
+              lr * n_wires * size_per_wire;
+          std::copy(
+              out_ptr + send_offset[lr] + offset,
+              out_ptr + send_offset[lr] +
+                  std::max(offset, std::min(offset + N, send_limits[lr])),
+              send_buffer(sz, buff, lr, lw).begin<T>());
+          send_count[lr * MAX_CONNS + lw]++;
+        }
+      }
+      post_send_all<MAX_DIR>(sz, buff, n_wires);
+
+      buff++;
+      in_flight += 2 * MAX_DIR * n_wires;
+    }
+
+    // Main loop
+    //
+    // Keep going until we have no longer data in flight.
+    while (in_flight > 0) {
+      ibv_wc wc[WC_NUM];
+      int n = poll(left_, right_, WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int work_type = wc[i].wr_id >> 16;
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int wire = wc[i].wr_id & 0xff;
+        int lr = wire / MAX_CONNS;
+        int lw = wire % MAX_CONNS;
+
+        in_flight--;
+
+        if (work_type == SEND_WR && send_count[wire] < n_steps) {
+          int64_t offset = lw * N + send_count[wire] * n_wires * N +
+              lr * n_wires * size_per_wire;
+          std::copy(
+              out_ptr + send_offset[lr] + offset,
+              out_ptr + send_offset[lr] +
+                  std::max(offset, std::min(offset + N, send_limits[lr])),
+              send_buffer(sz, buff, lr, lw).begin<T>());
+          send_to(sz, buff, lr, lw);
+          in_flight++;
+          send_count[wire]++;
+        }
+
+        else if (work_type == RECV_WR) {
+          int64_t offset = lw * N + recv_count[wire] * n_wires * N +
+              lr * n_wires * size_per_wire;
+          reduce_op(
+              recv_buffer(sz, buff, lr, lw).begin<T>(),
+              out_ptr + recv_offset[lr] + offset,
+              std::max<int64_t>(0, std::min(N, recv_limits[lr] - offset)));
+          recv_count[wire]++;
+          if (recv_count[wire] + (PIPELINE - 1) < n_steps) {
+            recv_from(sz, buff, lr, lw);
+            in_flight++;
+          }
+        }
+      }
+    }
+
+    send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) %
+        chunk_multiple_size;
+    recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) %
+        chunk_multiple_size;
+    if constexpr (MAX_DIR == 2) {
+      send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size;
+      recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size;
+      send_limits[0] = std::min(
+          n_wires * size_per_wire, std::max<int64_t>(0, size - send_offset[0]));
+      send_limits[1] =
+          std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
+      recv_limits[0] = std::min(
+          n_wires * size_per_wire, std::max<int64_t>(0, size - recv_offset[0]));
+      recv_limits[1] =
+          std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
+    } else {
+      send_limits[0] =
+          std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
+      recv_limits[0] =
+          std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
+    }
+    for (int i = 0; i < MAX_DIR * MAX_CONNS; i++) {
+      send_count[i] = recv_count[i] = 0;
+    }
+  }
+
+  // Secondly all gather
+  //
+  // The offsets are correct from the scatter reduce
+  for (int k = 0; k < size_ - 1; k++) {
+    // Prefill the pipeline
+    int buff = 0;
+    while (buff < n_steps && buff < PIPELINE) {
+      post_recv_all<MAX_DIR>(sz, buff, n_wires);
+      for (int lr = 0; lr < MAX_DIR; lr++) {
+        for (int lw = 0; lw < n_wires; lw++) {
+          int64_t offset = lw * N +
+              send_count[lr * MAX_CONNS + lw] * n_wires * N +
+              lr * n_wires * size_per_wire;
+          std::copy(
+              out_ptr + send_offset[lr] + offset,
+              out_ptr + send_offset[lr] +
+                  std::max(offset, std::min(offset + N, send_limits[lr])),
+              send_buffer(sz, buff, lr, lw).begin<T>());
+          send_count[lr * MAX_CONNS + lw]++;
+        }
+      }
+      post_send_all<MAX_DIR>(sz, buff, n_wires);
+
+      buff++;
+      in_flight += 2 * MAX_DIR * n_wires;
+    }
+
+    // Main loop
+    //
+    // Keep going until we have no longer data in flight.
+    while (in_flight > 0) {
+      ibv_wc wc[WC_NUM];
+      int n = poll(left_, right_, WC_NUM, wc);
+      for (int i = 0; i < n; i++) {
+        int work_type = wc[i].wr_id >> 16;
+        int buff = (wc[i].wr_id >> 8) & 0xff;
+        int wire = wc[i].wr_id & 0xff;
+        int lr = wire / MAX_CONNS;
+        int lw = wire % MAX_CONNS;
+
+        in_flight--;
+
+        if (work_type == SEND_WR && send_count[wire] < n_steps) {
+          int64_t offset = lw * N + send_count[wire] * n_wires * N +
+              lr * n_wires * size_per_wire;
+          std::copy(
+              out_ptr + send_offset[lr] + offset,
+              out_ptr + send_offset[lr] +
+                  std::max(offset, std::min(offset + N, send_limits[lr])),
+              send_buffer(sz, buff, lr, lw).begin<T>());
+          send_to(sz, buff, lr, lw);
+          in_flight++;
+          send_count[wire]++;
+        }
+
+        else if (work_type == RECV_WR) {
+          int64_t offset = lw * N + recv_count[wire] * n_wires * N +
+              lr * n_wires * size_per_wire;
+          std::copy(
+              recv_buffer(sz, buff, lr, lw).begin<T>(),
+              recv_buffer(sz, buff, lr, lw).begin<T>() +
+                  std::max<int64_t>(0, std::min(N, recv_limits[lr] - offset)),
+              out_ptr + recv_offset[lr] + offset);
+          recv_count[wire]++;
+          if (recv_count[wire] + (PIPELINE - 1) < n_steps) {
+            recv_from(sz, buff, lr, lw);
+            in_flight++;
+          }
+        }
+      }
+    }
+
+    send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) %
+        chunk_multiple_size;
+    recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) %
+        chunk_multiple_size;
+    if constexpr (MAX_DIR == 2) {
+      send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size;
+      recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size;
+      send_limits[0] = std::min(
+          n_wires * size_per_wire, std::max<int64_t>(0, size - send_offset[0]));
+      send_limits[1] =
+          std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
+      recv_limits[0] = std::min(
+          n_wires * size_per_wire, std::max<int64_t>(0, size - recv_offset[0]));
+      recv_limits[1] =
+          std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
+    } else {
+      send_limits[0] =
+          std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
+      recv_limits[0] =
+          std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
+    }
+    for (int i = 0; i < MAX_DIR * MAX_CONNS; i++) {
+      send_count[i] = recv_count[i] = 0;
+    }
+  }
+}
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/ring.h b/macos/third_party/mlx/mlx/distributed/jaccl/ring.h
new file mode 100644
index 0000000..a59ceb3
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/ring.h
@@ -0,0 +1,178 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/jaccl/utils.h"
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+constexpr int MAX_CONNS = 4;
+
+namespace mlx::core::distributed::jaccl {
+
+/**
+ * The JACCL communication group for a ring where each node is connected to its
+ * two neighboring nodes. It should be the highest bandwidth communication
+ * group for large messages when many connections per peer are used.
+ *
+ * Like all JACCL groups it uses a side channel to exchange the necessary
+ * information and then configure the connections to be ready for RDMA
+ * operations.
+ */
+class RingGroup : public GroupImpl {
+ public:
+  RingGroup(
+      int rank,
+      int size,
+      const std::vector<std::string>& left_devices,
+      const std::vector<std::string>& right_devices,
+      const char* coordinator_addr);
+
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s, Device::cpu);
+  }
+
+  int rank() override {
+    return rank_;
+  }
+
+  int size() override {
+    return size_;
+  }
+
+  void all_sum(const array& input, array& output, Stream stream) override;
+  void all_max(const array& input, array& output, Stream stream) override;
+  void all_min(const array& input, array& output, Stream stream) override;
+  void all_gather(const array& input, array& output, Stream stream) override;
+  void send(const array& input, int dst, Stream stream) override;
+  void recv(array& out, int src, Stream stream) override;
+
+  void sum_scatter(const array& input, array& output, Stream stream) override {
+    throw std::runtime_error("[jaccl] sum_scatter not supported.");
+  }
+
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    throw std::runtime_error("[jaccl] Group split not supported.");
+  }
+
+ private:
+  template <typename T, typename ReduceOp>
+  void all_reduce(
+      const array& input,
+      array& output,
+      Stream stream,
+      ReduceOp reduce_op);
+
+  template <int MAX_DIR, typename T, typename ReduceOp>
+  void all_reduce_impl(
+      const T* in_ptr,
+      T* out_ptr,
+      int64_t size,
+      int n_wires,
+      ReduceOp reduce_op);
+
+  /**
+   * Performs the connection initialization. Namely, after this call all
+   * Connection objects should have a queue pair in RTS state and all buffers
+   * should have been allocated.
+   */
+  void initialize();
+
+  /**
+   * Allocate all the buffers that we will use in the communication group.
+   */
+  void allocate_buffers();
+
+  void send_to(int sz, int buff, int left_right, int wire) {
+    if (left_right) {
+      left_[wire].post_send(
+          send_buffer_left(sz, buff, wire),
+          SEND_WR << 16 | buff << 8 | (MAX_CONNS + wire));
+    } else {
+      right_[wire].post_send(
+          send_buffer_right(sz, buff, wire), SEND_WR << 16 | buff << 8 | wire);
+    }
+  }
+
+  void recv_from(int sz, int buff, int left_right, int wire) {
+    if (left_right) {
+      right_[wire].post_recv(
+          recv_buffer_right(sz, buff, wire),
+          RECV_WR << 16 | buff << 8 | (MAX_CONNS + wire));
+    } else {
+      left_[wire].post_recv(
+          recv_buffer_left(sz, buff, wire), RECV_WR << 16 | buff << 8 | wire);
+    }
+  }
+
+  SharedBuffer& send_buffer_right(int sz, int buff, int wire) {
+    return send_buffers_
+        [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + wire];
+  }
+
+  SharedBuffer& send_buffer_left(int sz, int buff, int wire) {
+    return send_buffers_
+        [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + MAX_CONNS +
+         wire];
+  }
+
+  SharedBuffer& send_buffer(int sz, int buff, int left_right, int wire) {
+    return send_buffers_
+        [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 +
+         left_right * MAX_CONNS + wire];
+  }
+
+  SharedBuffer& recv_buffer_left(int sz, int buff, int wire) {
+    return recv_buffers_
+        [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + wire];
+  }
+
+  SharedBuffer& recv_buffer_right(int sz, int buff, int wire) {
+    return recv_buffers_
+        [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + MAX_CONNS +
+         wire];
+  }
+
+  SharedBuffer& recv_buffer(int sz, int buff, int left_right, int wire) {
+    return recv_buffers_
+        [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 +
+         left_right * MAX_CONNS + wire];
+  }
+
+  template <int MAX_DIR>
+  void post_recv_all(int sz, int buff, int n_wires) {
+    for (int lr = 0; lr < MAX_DIR; lr++) {
+      for (int lw = 0; lw < n_wires; lw++) {
+        recv_from(sz, buff, lr, lw);
+      }
+    }
+  }
+
+  void post_recv_all(int sz, int buff) {
+    post_recv_all<2>(sz, buff, left_.size());
+  }
+
+  template <int MAX_DIR>
+  void post_send_all(int sz, int buff, int n_wires) {
+    for (int lr = 0; lr < MAX_DIR; lr++) {
+      for (int lw = 0; lw < n_wires; lw++) {
+        send_to(sz, buff, lr, lw);
+      }
+    }
+  }
+
+  void post_send_all(int sz, int buff) {
+    post_send_all<2>(sz, buff, left_.size());
+  }
+
+  int rank_;
+  int size_;
+  SideChannel side_channel_;
+  std::vector<Connection> left_;
+  std::vector<Connection> right_;
+  std::vector<SharedBuffer> send_buffers_;
+  std::vector<SharedBuffer> recv_buffers_;
+};
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/utils.cpp b/macos/third_party/mlx/mlx/distributed/jaccl/utils.cpp
new file mode 100644
index 0000000..dc87d42
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/utils.cpp
@@ -0,0 +1,329 @@
+// Copyright © 2025 Apple Inc.
+
+#include <dlfcn.h>
+#include <unistd.h>
+#include <iostream>
+#include <sstream>
+
+#include "mlx/distributed/jaccl/utils.h"
+
+#define LOAD_SYMBOL(symbol, variable)                               \
+  {                                                                 \
+    variable = (decltype(variable))dlsym(librdma_handle_, #symbol); \
+    char* error = dlerror();                                        \
+    if (error != nullptr) {                                         \
+      std::cerr << IBV_TAG << " " << error << std::endl;            \
+      librdma_handle_ = nullptr;                                    \
+      return;                                                       \
+    }                                                               \
+  }
+
+namespace {
+
+void* page_aligned_alloc(size_t num_bytes) {
+  static size_t page_size = sysconf(_SC_PAGESIZE);
+  void* buf;
+  if (posix_memalign(&buf, page_size, num_bytes)) {
+    return nullptr;
+  }
+  return buf;
+}
+
+} // namespace
+
+namespace mlx::core::distributed::jaccl {
+
+IBVWrapper::IBVWrapper() {
+  librdma_handle_ = dlopen("librdma.dylib", RTLD_NOW | RTLD_GLOBAL);
+  if (librdma_handle_ == nullptr) {
+    return;
+  }
+
+  LOAD_SYMBOL(ibv_get_device_list, get_device_list);
+  LOAD_SYMBOL(ibv_get_device_name, get_device_name);
+  LOAD_SYMBOL(ibv_open_device, open_device);
+  LOAD_SYMBOL(ibv_free_device_list, free_device_list);
+  LOAD_SYMBOL(ibv_close_device, close_device);
+
+  LOAD_SYMBOL(ibv_alloc_pd, alloc_pd);
+  LOAD_SYMBOL(ibv_create_qp, create_qp);
+  LOAD_SYMBOL(ibv_create_cq, create_cq);
+  LOAD_SYMBOL(ibv_destroy_cq, destroy_cq);
+  LOAD_SYMBOL(ibv_destroy_qp, destroy_qp);
+  LOAD_SYMBOL(ibv_dealloc_pd, dealloc_pd);
+
+  LOAD_SYMBOL(ibv_query_port, query_port);
+  LOAD_SYMBOL(ibv_query_gid, query_gid);
+  LOAD_SYMBOL(ibv_modify_qp, modify_qp);
+  LOAD_SYMBOL(ibv_reg_mr, reg_mr);
+  LOAD_SYMBOL(ibv_dereg_mr, dereg_mr);
+
+  // Not really symbols but leaving them here in case they become symbols in
+  // the future.
+  //
+  // LOAD_SYMBOL(ibv_post_send, post_send);
+  // LOAD_SYMBOL(ibv_post_recv, post_recv);
+  // LOAD_SYMBOL(ibv_poll_cq, poll_cq);
+}
+
+IBVWrapper& ibv() {
+  static IBVWrapper wrapper;
+  return wrapper;
+}
+
+SharedBuffer::SharedBuffer(size_t num_bytes)
+    : data_(page_aligned_alloc(num_bytes)), num_bytes_(num_bytes) {}
+
+SharedBuffer::SharedBuffer(SharedBuffer&& b) : data_(nullptr), num_bytes_(0) {
+  std::swap(data_, b.data_);
+  std::swap(num_bytes_, b.num_bytes_);
+  std::swap(memory_regions_, b.memory_regions_);
+}
+
+SharedBuffer::~SharedBuffer() {
+  for (auto& [pd, mr] : memory_regions_) {
+    ibv().dereg_mr(mr);
+  }
+  if (data_ != nullptr) {
+    std::free(data_);
+  }
+}
+
+void SharedBuffer::register_to_protection_domain(ibv_pd* protection_domain) {
+  auto [it, inserted] = memory_regions_.insert({protection_domain, nullptr});
+  if (!inserted) {
+    throw std::runtime_error(
+        "[jaccl] Buffer can be registered once per protection domain");
+  }
+
+  it->second = ibv().reg_mr(
+      protection_domain,
+      data_,
+      num_bytes_,
+      IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
+          IBV_ACCESS_REMOTE_WRITE);
+  if (!it->second) {
+    throw std::runtime_error("[jaccl] Register memory region failed");
+  }
+}
+
+Connection::Connection(ibv_context* ctx_)
+    : ctx(ctx_),
+      protection_domain(nullptr),
+      completion_queue(nullptr),
+      queue_pair(nullptr) {
+  src.local_id = -1;
+}
+
+Connection::Connection(Connection&& c) : Connection(nullptr) {
+  std::swap(ctx, c.ctx);
+  std::swap(protection_domain, c.protection_domain);
+  std::swap(completion_queue, c.completion_queue);
+  std::swap(queue_pair, c.queue_pair);
+  std::swap(src, c.src);
+}
+
+Connection::~Connection() {
+  if (queue_pair != nullptr) {
+    ibv().destroy_qp(queue_pair);
+  }
+  if (completion_queue != nullptr) {
+    ibv().destroy_cq(completion_queue);
+  }
+  if (protection_domain != nullptr) {
+    ibv().dealloc_pd(protection_domain);
+  }
+  if (ctx != nullptr) {
+    ibv().close_device(ctx);
+  }
+}
+
+void Connection::allocate_protection_domain() {
+  protection_domain = ibv().alloc_pd(ctx);
+  if (protection_domain == nullptr) {
+    throw std::runtime_error("[jaccl] Couldn't allocate protection domain");
+  }
+}
+
+void Connection::create_completion_queue(int num_entries) {
+  completion_queue = ibv().create_cq(ctx, num_entries, nullptr, nullptr, 0);
+  if (completion_queue == nullptr) {
+    throw std::runtime_error("[jaccl] Couldn't create completion queue");
+  }
+}
+
+void Connection::create_queue_pair() {
+  ibv_qp_init_attr init_attr;
+  init_attr.qp_context = ctx;
+  init_attr.qp_context = ctx;
+  init_attr.send_cq = completion_queue;
+  init_attr.recv_cq = completion_queue;
+  init_attr.srq = nullptr;
+  init_attr.cap.max_send_wr = MAX_SEND_WR;
+  init_attr.cap.max_recv_wr = MAX_RECV_WR;
+  init_attr.cap.max_send_sge = 1;
+  init_attr.cap.max_recv_sge = 1;
+  init_attr.cap.max_inline_data = 0;
+  init_attr.qp_type = IBV_QPT_UC;
+  init_attr.sq_sig_all = 0;
+
+  queue_pair = ibv().create_qp(protection_domain, &init_attr);
+
+  if (queue_pair == nullptr) {
+    throw std::runtime_error("[jaccl] Couldn't create queue pair");
+  }
+}
+
+const Destination& Connection::info() {
+  if (queue_pair == nullptr || src.local_id >= 0) {
+    return src;
+  }
+
+  ibv_port_attr port_attr;
+  ibv().query_port(ctx, 1, &port_attr);
+  ibv_gid gid;
+  ibv().query_gid(ctx, 1, 1, &gid);
+
+  src.local_id = port_attr.lid;
+  src.queue_pair_number = queue_pair->qp_num;
+  src.packet_sequence_number = 7; // TODO: Change to sth random
+  src.global_identifier = gid;
+
+  return src;
+}
+
+void Connection::queue_pair_init() {
+  ibv_qp_attr attr = {};
+  attr.qp_state = IBV_QPS_INIT;
+  attr.port_num = 1;
+  attr.pkey_index = 0;
+  attr.qp_access_flags =
+      IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
+
+  int mask =
+      IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS;
+
+  if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) {
+    std::ostringstream msg;
+    msg << "[jaccl] Changing queue pair to INIT failed with errno " << status;
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+void Connection::queue_pair_rtr(const Destination& dst) {
+  ibv_qp_attr attr = {};
+  memset(&attr, 0, sizeof(attr));
+  attr.qp_state = IBV_QPS_RTR;
+  attr.path_mtu = IBV_MTU_1024;
+  attr.rq_psn = dst.packet_sequence_number;
+  attr.dest_qp_num = dst.queue_pair_number;
+  attr.ah_attr.dlid = dst.local_id;
+  attr.ah_attr.sl = 0;
+  attr.ah_attr.src_path_bits = 0;
+  attr.ah_attr.port_num = 1;
+  attr.ah_attr.is_global = 0;
+
+  if (dst.global_identifier.global.interface_id) {
+    attr.ah_attr.is_global = 1;
+    attr.ah_attr.grh.hop_limit = 1;
+    attr.ah_attr.grh.dgid = dst.global_identifier;
+    attr.ah_attr.grh.sgid_index = 1;
+  }
+
+  int mask = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN |
+      IBV_QP_RQ_PSN;
+
+  if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) {
+    std::ostringstream msg;
+    msg << "[jaccl] Changing queue pair to RTR failed with errno " << status;
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+void Connection::queue_pair_rts() {
+  ibv_qp_attr attr = {};
+  attr.qp_state = IBV_QPS_RTS;
+  attr.sq_psn = src.packet_sequence_number;
+
+  int mask = IBV_QP_STATE | IBV_QP_SQ_PSN;
+
+  if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) {
+    std::ostringstream msg;
+    msg << "[jaccl] Changing queue pair to RTS failed with errno " << status;
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+std::vector<Connection> create_connections(
+    const std::vector<std::string>& device_names) {
+  std::vector<Connection> connections;
+  int num_devices = 0;
+  ibv_device** devices = ibv().get_device_list(&num_devices);
+  for (auto& name : device_names) {
+    // Empty so add a nullptr context
+    if (name.empty()) {
+      connections.emplace_back(nullptr);
+      continue;
+    }
+
+    // Search for the name and try to open the device
+    for (int i = 0; i < num_devices; i++) {
+      if (name == ibv().get_device_name(devices[i])) {
+        auto ctx = ibv().open_device(devices[i]);
+        if (ctx == nullptr) {
+          std::ostringstream msg;
+          msg << "[jaccl] Could not open device " << name;
+          throw std::runtime_error(msg.str());
+        }
+        connections.emplace_back(ctx);
+        break;
+      }
+    }
+  }
+  ibv().free_device_list(devices);
+
+  return connections;
+}
+
+SideChannel::SideChannel(int rank, int size, const char* addr)
+    : rank_(rank), size_(size) {
+  auto address = detail::parse_address(addr);
+
+  if (rank_ == 0) {
+    detail::TCPSocket server(IBV_TAG);
+    server.listen(IBV_TAG, address);
+
+    for (int i = 0; i < size - 1; i++) {
+      sockets_.push_back(server.accept(IBV_TAG));
+    }
+
+    std::vector<int> ranks(size - 1);
+    for (int i = 0; i < size - 1; i++) {
+      sockets_[i].recv(
+          IBV_TAG, reinterpret_cast<char*>(&ranks[i]), sizeof(int));
+      ranks[i]--;
+    }
+    for (int i = 0; i < size - 1; i++) {
+      while (i != ranks[i]) {
+        std::swap(sockets_[i], sockets_[ranks[i]]);
+        std::swap(ranks[i], ranks[ranks[i]]);
+      }
+    }
+  } else {
+    sockets_.push_back(
+        detail::TCPSocket::connect(
+            IBV_TAG, address, 4, 1000, [](int attempt, int wait) {
+              std::cerr << IBV_TAG << " Connection attempt " << attempt
+                        << " waiting " << wait << " ms" << std::endl;
+            }));
+    sockets_[0].send(IBV_TAG, reinterpret_cast<char*>(&rank_), sizeof(int));
+  }
+}
+
+SideChannel::SideChannel(SideChannel&& sc)
+    : rank_(sc.rank_), size_(sc.size_), sockets_(std::move(sc.sockets_)) {
+  sc.rank_ = -1;
+  sc.size_ = -1;
+}
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/jaccl/utils.h b/macos/third_party/mlx/mlx/distributed/jaccl/utils.h
new file mode 100644
index 0000000..5b700f4
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/jaccl/utils.h
@@ -0,0 +1,342 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <infiniband/verbs.h>
+
+#include <unordered_map>
+#include <vector>
+
+#include "mlx/distributed/utils.h"
+
+constexpr const char* IBV_TAG = "[jaccl]";
+constexpr int SEND_WR = 1;
+constexpr int RECV_WR = 2;
+constexpr int MAX_SEND_WR = 32;
+constexpr int MAX_RECV_WR = 32;
+constexpr int BUFFER_SIZES = 8;
+constexpr int NUM_BUFFERS = 2;
+constexpr int FRAME_SIZE = 4096;
+
+namespace detail = mlx::core::distributed::detail;
+
+namespace {
+
+template <typename T, typename = void>
+struct is_container : std::false_type {};
+
+template <typename T>
+struct is_container<
+    T,
+    std::void_t<typename T::value_type, typename T::iterator>>
+    : std::true_type {};
+
+inline std::pair<int, int64_t> buffer_size_from_message(int64_t msg) {
+  if (__builtin_available(macOS 26.3, iOS 26.3, tvOS 26.3, visionOS 26.3, *)) {
+    for (int k = BUFFER_SIZES - 1; k > 0; k--) {
+      if (msg >= FRAME_SIZE * (1 << k)) {
+        return {k, FRAME_SIZE * (1 << k)};
+      }
+    }
+  }
+  return {0, FRAME_SIZE};
+}
+
+} // namespace
+
+namespace mlx::core::distributed::jaccl {
+
+/**
+ * Wrapper for the ibverbs API.
+ */
+struct IBVWrapper {
+  IBVWrapper();
+  bool is_available() {
+    return librdma_handle_ != nullptr;
+  }
+
+  // API
+  ibv_device** (*get_device_list)(int*);
+  const char* (*get_device_name)(ibv_device*);
+  ibv_context* (*open_device)(ibv_device*);
+  void (*free_device_list)(ibv_device**);
+  int (*close_device)(ibv_context*);
+
+  ibv_pd* (*alloc_pd)(ibv_context*);
+  ibv_qp* (*create_qp)(ibv_pd*, ibv_qp_init_attr*);
+  ibv_cq* (*create_cq)(ibv_context*, int, void*, ibv_comp_channel*, int);
+  int (*destroy_cq)(ibv_cq*);
+  int (*destroy_qp)(ibv_qp*);
+  int (*dealloc_pd)(ibv_pd*);
+
+  int (*query_port)(ibv_context*, uint8_t, ibv_port_attr*);
+  int (*query_gid)(ibv_context*, uint8_t, int, ibv_gid*);
+  int (*modify_qp)(ibv_qp*, ibv_qp_attr*, int);
+  ibv_mr* (*reg_mr)(ibv_pd*, void*, size_t, int);
+  int (*dereg_mr)(ibv_mr*);
+
+ private:
+  void* librdma_handle_;
+};
+
+IBVWrapper& ibv();
+
+/**
+ * Contains the information that defines a destination to a remote device.
+ * Basically we can compute our own destination and share it with remote hosts
+ * over the side channel.
+ */
+struct Destination {
+  int local_id;
+  int queue_pair_number;
+  int packet_sequence_number;
+  ibv_gid global_identifier;
+};
+
+/**
+ * A buffer that can be registered to a number of protection domains.
+ */
+class SharedBuffer {
+ public:
+  SharedBuffer(size_t num_bytes);
+  SharedBuffer(SharedBuffer&& b);
+  ~SharedBuffer();
+
+  SharedBuffer(const SharedBuffer&) = delete;
+  SharedBuffer& operator=(const SharedBuffer&) = delete;
+
+  void register_to_protection_domain(ibv_pd* protection_domain);
+
+  size_t size() const {
+    return num_bytes_;
+  }
+
+  uint32_t local_key(ibv_pd* protection_domain) const {
+    return memory_regions_.at(protection_domain)->lkey;
+  }
+
+  ibv_sge to_scatter_gather_entry(ibv_pd* protection_domain) const {
+    ibv_sge entry;
+    entry.addr = reinterpret_cast<uintptr_t>(data_);
+    entry.length = size();
+    entry.lkey = local_key(protection_domain);
+    return entry;
+  }
+
+  template <typename T>
+  T* data() {
+    return static_cast<T*>(data_);
+  }
+
+  template <typename T>
+  T* begin() {
+    return static_cast<T*>(data_);
+  }
+
+  template <typename T>
+  T* end() {
+    return static_cast<T*>(data_) + size() / sizeof(T);
+  }
+
+ private:
+  void* data_;
+  size_t num_bytes_;
+  std::unordered_map<ibv_pd*, ibv_mr*> memory_regions_;
+};
+
+/**
+ * Manipulates an RDMA connection. Enables (among other things)
+ *
+ *   - Creating a queue pair
+ *   - Sending and receiving
+ *   - Checking completion
+ */
+struct Connection {
+  ibv_context* ctx;
+  ibv_pd* protection_domain;
+  ibv_cq* completion_queue;
+  ibv_qp* queue_pair;
+  Destination src; // holds the local information
+
+  Connection(ibv_context* ctx_);
+  Connection(Connection&& c);
+
+  Connection(const Connection&) = delete;
+  Connection& operator=(Connection&) = delete;
+
+  ~Connection();
+  void allocate_protection_domain();
+  void create_completion_queue(int num_entries);
+  void create_queue_pair();
+
+  const Destination& info();
+  void queue_pair_init();
+  void queue_pair_rtr(const Destination& dst);
+  void queue_pair_rts();
+
+  void post_send(const SharedBuffer& buff, uint64_t work_request_id) {
+    ibv_send_wr work_request, *bad_work_request;
+
+    auto entry = buff.to_scatter_gather_entry(protection_domain);
+    work_request.wr_id = work_request_id;
+    work_request.sg_list = &entry;
+    work_request.num_sge = 1;
+    work_request.opcode = IBV_WR_SEND;
+    work_request.send_flags = IBV_SEND_SIGNALED;
+    work_request.next = nullptr;
+
+    if (int status =
+            ibv_post_send(queue_pair, &work_request, &bad_work_request);
+        status != 0) {
+      std::ostringstream msg;
+      msg << "[jaccl] Send failed with error code " << status;
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  void post_recv(const SharedBuffer& buff, uint64_t work_request_id) {
+    ibv_recv_wr work_request, *bad_work_request;
+
+    auto entry = buff.to_scatter_gather_entry(protection_domain);
+    work_request.wr_id = work_request_id;
+    work_request.sg_list = &entry;
+    work_request.num_sge = 1;
+    work_request.next = nullptr;
+
+    if (int status =
+            ibv_post_recv(queue_pair, &work_request, &bad_work_request);
+        status != 0) {
+      std::ostringstream msg;
+      msg << "[jaccl] Recv failed with error code " << status;
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  int poll(int num_completions, ibv_wc* work_completions) {
+    return ibv_poll_cq(completion_queue, num_completions, work_completions);
+  }
+};
+
+std::vector<Connection> create_connections(
+    const std::vector<std::string>& device_names);
+
+inline int poll(
+    const std::vector<Connection>& connections,
+    int num_completions,
+    ibv_wc* work_completions) {
+  int completions = 0;
+  for (auto& c : connections) {
+    if (c.ctx == nullptr) {
+      continue;
+    }
+    if (completions >= num_completions) {
+      return completions;
+    }
+
+    int n = ibv_poll_cq(
+        c.completion_queue,
+        num_completions - completions,
+        work_completions + completions);
+
+    completions += n;
+  }
+  return completions;
+}
+
+inline int poll(
+    const std::vector<Connection>& connections_1,
+    const std::vector<Connection>& connections_2,
+    int num_completions,
+    ibv_wc* work_completions) {
+  int completions = 0;
+  completions += poll(connections_1, num_completions, work_completions);
+  completions += poll(
+      connections_2,
+      num_completions - completions,
+      work_completions + completions);
+  return completions;
+}
+
+/**
+ * Implement a TCP side channel to exchange information about the RDMA
+ * connections.
+ *
+ * Implements a simple all gather where every node sends to rank 0 and rank 0
+ * broadcasts to every node.
+ */
+class SideChannel {
+ public:
+  SideChannel(int rank, int size, const char* addr);
+  SideChannel(SideChannel&& sc);
+
+  SideChannel(const SideChannel&) = delete;
+  SideChannel& operator=(const SideChannel&) = delete;
+
+  template <typename T>
+  std::vector<T> all_gather(const T& v) {
+    std::vector<T> result(size_);
+
+    // T is a container of stuff like std::vector or std::string
+    if constexpr (is_container<T>::value) {
+      using U = typename T::value_type;
+
+      // Share the lengths first and set the communication size to be the
+      // maximum length of the containers.
+      auto lengths = all_gather<int>(v.size());
+      auto max_len = *std::max_element(lengths.begin(), lengths.end());
+      for (auto& s : result) {
+        s.resize(max_len);
+      }
+
+      // All gather of length max_len
+      if (rank_ == 0) {
+        std::copy(v.begin(), v.end(), result[rank_].begin());
+        for (int i = 1; i < size_; i++) {
+          sockets_[i - 1].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len);
+        }
+        for (int i = 1; i < size_; i++) {
+          for (int j = 0; j < size_; j++) {
+            sockets_[i - 1].send(
+                IBV_TAG, result[j].data(), sizeof(U) * max_len);
+          }
+        }
+      } else {
+        std::copy(v.begin(), v.end(), result[rank_].begin());
+        sockets_[0].send(IBV_TAG, result[rank_].data(), sizeof(U) * max_len);
+        for (int i = 0; i < size_; i++) {
+          sockets_[0].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len);
+        }
+      }
+
+      // Resize the outputs back to the original length
+      for (int i = 0; i < size_; i++) {
+        result[i].resize(lengths[i]);
+      }
+    }
+
+    // T is a scalar
+    else {
+      if (rank_ == 0) {
+        result[rank_] = v;
+        for (int i = 1; i < size_; i++) {
+          sockets_[i - 1].recv(IBV_TAG, &result[i], sizeof(T));
+        }
+        for (int i = 1; i < size_; i++) {
+          sockets_[i - 1].send(IBV_TAG, result.data(), size_ * sizeof(T));
+        }
+      } else {
+        sockets_[0].send(IBV_TAG, &v, sizeof(T));
+        sockets_[0].recv(IBV_TAG, result.data(), size_ * sizeof(T));
+      }
+    }
+
+    return result;
+  }
+
+ private:
+  int rank_;
+  int size_;
+  std::vector<detail::TCPSocket> sockets_;
+};
+
+} // namespace mlx::core::distributed::jaccl
diff --git a/macos/third_party/mlx/mlx/distributed/mpi/CMakeLists.txt b/macos/third_party/mlx/mlx/distributed/mpi/CMakeLists.txt
new file mode 100644
index 0000000..842f70b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/mpi/CMakeLists.txt
@@ -0,0 +1,5 @@
+if(MLX_BUILD_CPU)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/mpi.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_mpi.cpp)
+endif()
diff --git a/macos/third_party/mlx/mlx/distributed/mpi/mpi.cpp b/macos/third_party/mlx/mlx/distributed/mpi/mpi.cpp
new file mode 100644
index 0000000..3b176e6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/mpi/mpi.cpp
@@ -0,0 +1,501 @@
+// Copyright © 2024 Apple Inc.
+
+#include <dlfcn.h>
+#include <cstdlib>
+#include <iostream>
+
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/mpi/mpi.h"
+#include "mlx/distributed/mpi/mpi_declarations.h"
+
+#define LOAD_SYMBOL(symbol, variable)                              \
+  {                                                                \
+    variable = (decltype(variable))dlsym(libmpi_handle_, #symbol); \
+    char* error = dlerror();                                       \
+    if (error != nullptr) {                                        \
+      libmpi_handle_ = nullptr;                                    \
+      return;                                                      \
+    }                                                              \
+  }
+
+static const char* get_libmpi_name() {
+  const char* libname = std::getenv("MLX_MPI_LIBNAME");
+  if (libname != nullptr) {
+    return libname;
+  }
+#ifdef __APPLE__
+  return "libmpi.dylib";
+#else
+  return "libmpi.so";
+#endif
+}
+
+namespace mlx::core::distributed::mpi {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+namespace {
+
+template <typename T>
+void simple_sum(
+    void* input,
+    void* accumulator,
+    int* len,
+    MPI_Datatype* datatype) {
+  T* in = (T*)input;
+  T* acc = (T*)accumulator;
+  int N = *len;
+
+  while (N-- > 0) {
+    *acc += *in;
+    acc++;
+    in++;
+  }
+}
+template void simple_sum<float16_t>(void*, void*, int*, MPI_Datatype*);
+template void simple_sum<bfloat16_t>(void*, void*, int*, MPI_Datatype*);
+
+template <typename T>
+void simple_max(
+    void* input,
+    void* accumulator,
+    int* len,
+    MPI_Datatype* datatype) {
+  T* in = (T*)input;
+  T* acc = (T*)accumulator;
+  int N = *len;
+
+  while (N-- > 0) {
+    *acc = std::max(*acc, *in);
+    acc++;
+    in++;
+  }
+}
+template void simple_max<float16_t>(void*, void*, int*, MPI_Datatype*);
+template void simple_max<bfloat16_t>(void*, void*, int*, MPI_Datatype*);
+template void simple_max<complex64_t>(void*, void*, int*, MPI_Datatype*);
+
+template <typename T>
+void simple_min(
+    void* input,
+    void* accumulator,
+    int* len,
+    MPI_Datatype* datatype) {
+  T* in = (T*)input;
+  T* acc = (T*)accumulator;
+  int N = *len;
+
+  while (N-- > 0) {
+    *acc = std::min(*acc, *in);
+    acc++;
+    in++;
+  }
+}
+template void simple_min<float16_t>(void*, void*, int*, MPI_Datatype*);
+template void simple_min<bfloat16_t>(void*, void*, int*, MPI_Datatype*);
+template void simple_min<complex64_t>(void*, void*, int*, MPI_Datatype*);
+
+struct MPIWrapper {
+  MPIWrapper() {
+    initialized_ = false;
+
+    libmpi_handle_ = dlopen(get_libmpi_name(), RTLD_NOW | RTLD_GLOBAL);
+    if (libmpi_handle_ == nullptr) {
+      return;
+    }
+
+    // Check library version and warn if it isn't Open MPI
+    int (*get_version)(char*, int*);
+    LOAD_SYMBOL(MPI_Get_library_version, get_version);
+    char version_ptr[MPI_MAX_LIBRARY_VERSION_STRING];
+    int version_length = 0;
+    get_version(version_ptr, &version_length);
+    std::string_view version(version_ptr, version_length);
+    if (version.find("Open MPI") == std::string::npos) {
+      std::cerr << "[mpi] MPI found but it does not appear to be Open MPI."
+                << "MLX requires Open MPI but this is " << version << std::endl;
+      libmpi_handle_ = nullptr;
+      return;
+    }
+
+    // API
+    LOAD_SYMBOL(MPI_Init, init);
+    LOAD_SYMBOL(MPI_Finalize, finalize);
+    LOAD_SYMBOL(MPI_Comm_rank, rank);
+    LOAD_SYMBOL(MPI_Comm_size, size);
+    LOAD_SYMBOL(MPI_Comm_split, comm_split);
+    LOAD_SYMBOL(MPI_Comm_free, comm_free);
+    LOAD_SYMBOL(MPI_Allreduce, all_reduce);
+    LOAD_SYMBOL(MPI_Allgather, all_gather);
+    LOAD_SYMBOL(MPI_Send, send);
+    LOAD_SYMBOL(MPI_Recv, recv);
+    LOAD_SYMBOL(MPI_Type_contiguous, mpi_type_contiguous);
+    LOAD_SYMBOL(MPI_Type_commit, mpi_type_commit);
+    LOAD_SYMBOL(MPI_Op_create, mpi_op_create);
+
+    // Objects
+    LOAD_SYMBOL(ompi_mpi_comm_world, comm_world_);
+
+    // Ops
+    LOAD_SYMBOL(ompi_mpi_op_sum, op_sum_);
+    LOAD_SYMBOL(ompi_mpi_op_max, op_max_);
+    LOAD_SYMBOL(ompi_mpi_op_min, op_min_);
+
+    // Datatypes
+    LOAD_SYMBOL(ompi_mpi_c_bool, mpi_bool_);
+    LOAD_SYMBOL(ompi_mpi_int8_t, mpi_int8_);
+    LOAD_SYMBOL(ompi_mpi_uint8_t, mpi_uint8_);
+    LOAD_SYMBOL(ompi_mpi_int16_t, mpi_int16_);
+    LOAD_SYMBOL(ompi_mpi_uint16_t, mpi_uint16_);
+    LOAD_SYMBOL(ompi_mpi_int32_t, mpi_int32_);
+    LOAD_SYMBOL(ompi_mpi_uint32_t, mpi_uint32_);
+    LOAD_SYMBOL(ompi_mpi_int64_t, mpi_int64_);
+    LOAD_SYMBOL(ompi_mpi_uint64_t, mpi_uint64_);
+    LOAD_SYMBOL(ompi_mpi_float, mpi_float_);
+    LOAD_SYMBOL(ompi_mpi_double, mpi_double_);
+    LOAD_SYMBOL(ompi_mpi_c_complex, mpi_complex_);
+  }
+
+  bool is_available() {
+    return libmpi_handle_ != nullptr;
+  }
+
+  bool init_safe() {
+    if (!is_available()) {
+      return false;
+    }
+    bool success = init(nullptr, nullptr) == MPI_SUCCESS;
+
+    // Initialize custom types and ops
+    if (success && !initialized_) {
+      // Custom float16 dtypes
+      mpi_type_contiguous(2, mpi_uint8_, &mpi_float16_);
+      mpi_type_commit(&mpi_float16_);
+      mpi_type_contiguous(2, mpi_uint8_, &mpi_bfloat16_);
+      mpi_type_commit(&mpi_bfloat16_);
+
+      // Custom reduction ops
+      mpi_op_create(&simple_sum<float16_t>, 1, &op_sum_f16_);
+      mpi_op_create(&simple_sum<bfloat16_t>, 1, &op_sum_bf16_);
+      mpi_op_create(&simple_max<float16_t>, 1, &op_max_f16_);
+      mpi_op_create(&simple_max<bfloat16_t>, 1, &op_max_bf16_);
+      mpi_op_create(&simple_max<complex64_t>, 1, &op_max_c64_);
+      mpi_op_create(&simple_min<float16_t>, 1, &op_min_f16_);
+      mpi_op_create(&simple_min<bfloat16_t>, 1, &op_min_bf16_);
+      mpi_op_create(&simple_min<complex64_t>, 1, &op_min_c64_);
+
+      initialized_ = true;
+    }
+
+    return success;
+  }
+
+  void finalize_safe() {
+    if (is_available()) {
+      finalize();
+    }
+  }
+
+  MPI_Comm world() {
+    return comm_world_;
+  }
+
+  MPI_Datatype datatype(const array& arr) {
+    switch (arr.dtype()) {
+      case bool_:
+        return mpi_bool_;
+      case int8:
+        return mpi_int8_;
+      case uint8:
+        return mpi_uint8_;
+      case int16:
+        return mpi_int16_;
+      case uint16:
+        return mpi_uint16_;
+      case int32:
+        return mpi_int32_;
+      case uint32:
+        return mpi_uint32_;
+      case int64:
+        return mpi_int64_;
+      case uint64:
+        return mpi_uint64_;
+      case float32:
+        return mpi_float_;
+      case complex64:
+        return mpi_complex_;
+      case float16:
+        return mpi_float16_;
+      case bfloat16:
+        return mpi_bfloat16_;
+      case float64:
+        return mpi_double_;
+      default:
+        throw std::runtime_error("Invalid type");
+    }
+  }
+
+  MPI_Op op_sum(const array& arr) {
+    switch (arr.dtype()) {
+      case float16:
+        return op_sum_f16_;
+      case bfloat16:
+        return op_sum_bf16_;
+      default:
+        return op_sum_;
+    }
+  }
+
+  MPI_Op op_max(const array& arr) {
+    switch (arr.dtype()) {
+      case float16:
+        return op_max_f16_;
+      case bfloat16:
+        return op_max_bf16_;
+      case complex64:
+        return op_max_c64_;
+      default:
+        return op_max_;
+    }
+  }
+
+  MPI_Op op_min(const array& arr) {
+    switch (arr.dtype()) {
+      case float16:
+        return op_min_f16_;
+      case bfloat16:
+        return op_min_bf16_;
+      case complex64:
+        return op_min_c64_;
+      default:
+        return op_min_;
+    }
+  }
+
+  void* libmpi_handle_;
+
+  // API
+  int (*init)(int*, char***);
+  int (*finalize)();
+  int (*rank)(MPI_Comm, int*);
+  int (*size)(MPI_Comm, int*);
+  int (*all_reduce)(const void*, void*, int, MPI_Datatype, MPI_Op, MPI_Comm);
+  int (*all_gather)(
+      const void*,
+      int,
+      MPI_Datatype,
+      void*,
+      int,
+      MPI_Datatype,
+      MPI_Comm);
+  int (*comm_split)(MPI_Comm, int, int, MPI_Comm*);
+  int (*comm_free)(MPI_Comm*);
+  int (*send)(const void*, int, MPI_Datatype, int, int, MPI_Comm);
+  int (*recv)(void*, int, MPI_Datatype, int, int, MPI_Comm, MPI_Status*);
+
+  // Objects
+  MPI_Comm comm_world_;
+
+  // Ops
+  MPI_Op op_sum_;
+  MPI_Op op_sum_f16_;
+  MPI_Op op_sum_bf16_;
+  MPI_Op op_max_;
+  MPI_Op op_max_f16_;
+  MPI_Op op_max_bf16_;
+  MPI_Op op_max_c64_;
+  MPI_Op op_min_;
+  MPI_Op op_min_f16_;
+  MPI_Op op_min_bf16_;
+  MPI_Op op_min_c64_;
+
+  // Datatypes
+  MPI_Datatype mpi_bool_;
+  MPI_Datatype mpi_int8_;
+  MPI_Datatype mpi_uint8_;
+  MPI_Datatype mpi_int16_;
+  MPI_Datatype mpi_uint16_;
+  MPI_Datatype mpi_int32_;
+  MPI_Datatype mpi_uint32_;
+  MPI_Datatype mpi_int64_;
+  MPI_Datatype mpi_uint64_;
+  MPI_Datatype mpi_float_;
+  MPI_Datatype mpi_double_;
+  MPI_Datatype mpi_complex_;
+  MPI_Datatype mpi_float16_;
+  MPI_Datatype mpi_bfloat16_;
+
+ private:
+  bool initialized_;
+
+  // Private API
+  int (*mpi_type_contiguous)(int, MPI_Datatype, MPI_Datatype*);
+  int (*mpi_type_commit)(MPI_Datatype*);
+  int (*mpi_op_create)(MPI_User_function*, int, MPI_Op*);
+};
+
+MPIWrapper& mpi() {
+  static MPIWrapper wrapper;
+  return wrapper;
+}
+
+} // namespace
+
+class MPIGroup : public GroupImpl {
+ public:
+  MPIGroup(MPI_Comm comm, bool global)
+      : comm_(comm), global_(global), rank_(-1), size_(-1) {}
+
+  virtual ~MPIGroup() {
+    if (global_) {
+      mpi().finalize_safe();
+    } else {
+      mpi().comm_free(&comm_);
+    }
+  }
+
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s, Device::cpu);
+  }
+
+  int rank() override {
+    if (rank_ < 0) {
+      mpi().rank(comm_, &rank_);
+    }
+    return rank_;
+  }
+
+  int size() override {
+    if (size_ < 0) {
+      mpi().size(comm_, &size_);
+    }
+    return size_;
+  }
+
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    key = (key < 0) ? rank() : key;
+
+    MPI_Comm new_comm;
+    int result = mpi().comm_split(comm_, color, key, &new_comm);
+    if (result != MPI_SUCCESS) {
+      throw std::runtime_error("MPI could not split this group");
+    }
+
+    return std::make_shared<MPIGroup>(new_comm, false);
+  }
+
+  void all_sum(const array& input, array& output, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.set_output_array(output);
+    encoder.dispatch(
+        mpi().all_reduce,
+        (input.data<void>() == output.data<void>()) ? MPI_IN_PLACE
+                                                    : input.data<void>(),
+        output.data<void>(),
+        input.size(),
+        mpi().datatype(input),
+        mpi().op_sum(input),
+        comm_);
+  }
+
+  void all_max(const array& input, array& output, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.set_output_array(output);
+    encoder.dispatch(
+        mpi().all_reduce,
+        (input.data<void>() == output.data<void>()) ? MPI_IN_PLACE
+                                                    : input.data<void>(),
+        output.data<void>(),
+        input.size(),
+        mpi().datatype(input),
+        mpi().op_max(input),
+        comm_);
+  }
+
+  void all_min(const array& input, array& output, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.set_output_array(output);
+    encoder.dispatch(
+        mpi().all_reduce,
+        (input.data<void>() == output.data<void>()) ? MPI_IN_PLACE
+                                                    : input.data<void>(),
+        output.data<void>(),
+        input.size(),
+        mpi().datatype(input),
+        mpi().op_min(input),
+        comm_);
+  }
+
+  void all_gather(const array& input, array& output, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.set_output_array(output);
+    encoder.dispatch(
+        mpi().all_gather,
+        input.data<void>(),
+        input.size(),
+        mpi().datatype(input),
+        output.data<void>(),
+        input.size(),
+        mpi().datatype(output),
+        comm_);
+  }
+
+  void send(const array& input, int dst, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.dispatch(
+        mpi().send,
+        input.data<void>(),
+        input.size(),
+        mpi().datatype(input),
+        dst,
+        0,
+        comm_);
+  }
+
+  void recv(array& out, int src, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_output_array(out);
+    encoder.dispatch([out_ptr = out.data<void>(),
+                      out_size = out.size(),
+                      out_type = mpi().datatype(out),
+                      src,
+                      comm = comm_]() {
+      MPI_Status status;
+      mpi().recv(out_ptr, out_size, out_type, src, MPI_ANY_TAG, comm, &status);
+    });
+  }
+
+  void sum_scatter(const array& input, array& output, Stream stream) override {
+    throw std::runtime_error("[mpi] sum_scatter not yet implemented.");
+  }
+
+ private:
+  MPI_Comm comm_;
+  bool global_;
+  int rank_;
+  int size_;
+};
+
+bool is_available() {
+  return mpi().is_available();
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  if (!mpi().init_safe()) {
+    if (strict) {
+      throw std::runtime_error("Cannot initialize MPI");
+    }
+    return nullptr;
+  }
+
+  return std::make_shared<MPIGroup>(mpi().world(), true);
+}
+
+} // namespace mlx::core::distributed::mpi
diff --git a/macos/third_party/mlx/mlx/distributed/mpi/mpi.h b/macos/third_party/mlx/mlx/distributed/mpi/mpi.h
new file mode 100644
index 0000000..cd11a47
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/mpi/mpi.h
@@ -0,0 +1,12 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::mpi {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available();
+std::shared_ptr<GroupImpl> init(bool strict = false);
+
+} // namespace mlx::core::distributed::mpi
diff --git a/macos/third_party/mlx/mlx/distributed/mpi/mpi_declarations.h b/macos/third_party/mlx/mlx/distributed/mpi/mpi_declarations.h
new file mode 100644
index 0000000..99c1a9c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/mpi/mpi_declarations.h
@@ -0,0 +1,28 @@
+// Copyright © 2024 Apple Inc.
+
+// Constants
+
+#define MPI_SUCCESS 0
+#define MPI_ANY_SOURCE -1
+#define MPI_ANY_TAG -1
+#define MPI_IN_PLACE ((void*)1)
+#define MPI_MAX_LIBRARY_VERSION_STRING 256
+
+// Define all the types that we use so that we don't include <mpi.h> which
+// causes linker errors on some platforms.
+//
+// NOTE: We define everything for openmpi.
+
+typedef void* MPI_Comm;
+typedef void* MPI_Datatype;
+typedef void* MPI_Op;
+
+typedef void(MPI_User_function)(void*, void*, int*, MPI_Datatype*);
+
+typedef struct ompi_status_public_t {
+  int MPI_SOURCE;
+  int MPI_TAG;
+  int MPI_ERROR;
+  int _cancelled;
+  size_t _ucount;
+} MPI_Status;
diff --git a/macos/third_party/mlx/mlx/distributed/mpi/no_mpi.cpp b/macos/third_party/mlx/mlx/distributed/mpi/no_mpi.cpp
new file mode 100644
index 0000000..4a7fc66
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/mpi/no_mpi.cpp
@@ -0,0 +1,20 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/distributed/mpi/mpi.h"
+
+namespace mlx::core::distributed::mpi {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available() {
+  return false;
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  if (strict) {
+    throw std::runtime_error("Cannot initialize MPI");
+  }
+  return nullptr;
+}
+
+} // namespace mlx::core::distributed::mpi
diff --git a/macos/third_party/mlx/mlx/distributed/nccl/CMakeLists.txt b/macos/third_party/mlx/mlx/distributed/nccl/CMakeLists.txt
new file mode 100644
index 0000000..644d060
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/nccl/CMakeLists.txt
@@ -0,0 +1,26 @@
+if(MLX_BUILD_CUDA)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp)
+  find_package(NCCL)
+  if(NCCL_FOUND)
+    target_link_libraries(mlx PRIVATE ${NCCL_LIBRARIES})
+    target_include_directories(mlx PRIVATE ${NCCL_INCLUDE_DIRS})
+  else()
+    message(
+      STATUS
+        "NCCL not found, using stubs. To run distributed with NCCL backend, install NCCL."
+    )
+    file(
+      DOWNLOAD
+      "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in"
+      "${CMAKE_CURRENT_BINARY_DIR}/nccl.h")
+    add_library(nccl_stub OBJECT
+                ${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub/nccl_stubs.cpp)
+    target_include_directories(
+      nccl_stub
+      PRIVATE ${CUDAToolkit_INCLUDE_DIRS}
+      PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
+    target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:nccl_stub>)
+  endif()
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp)
+endif()
diff --git a/macos/third_party/mlx/mlx/distributed/nccl/nccl.cpp b/macos/third_party/mlx/mlx/distributed/nccl/nccl.cpp
new file mode 100644
index 0000000..d8244bf
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/nccl/nccl.cpp
@@ -0,0 +1,443 @@
+// NCCL distributed support currently requires Unix socket APIs
+// TODO: Add Windows Winsock2 support for Windows builds
+#ifndef _WIN32
+#include <arpa/inet.h>
+#include <netdb.h>
+#include <sys/socket.h>
+#include <unistd.h>
+#endif
+
+#include <cuda_runtime.h>
+#include <nccl.h>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <mutex>
+#include <stdexcept>
+#include <string>
+#include <type_traits>
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::distributed::nccl {
+
+// Can be tuned with MLX_NCCL_TIMEOUT
+constexpr int nccl_timeout = 300000; // miliseconds
+
+#define CHECK_CUDA(cmd)              \
+  do {                               \
+    cudaError_t e = cmd;             \
+    if (e != cudaSuccess) {          \
+      fprintf(                       \
+          stderr,                    \
+          "CUDA error %s:%d '%s'\n", \
+          __FILE__,                  \
+          __LINE__,                  \
+          cudaGetErrorString(e));    \
+      exit(1);                       \
+    }                                \
+  } while (0)
+
+#define CHECK_NCCL(cmd)              \
+  do {                               \
+    ncclResult_t r = cmd;            \
+    if (r != ncclSuccess) {          \
+      fprintf(                       \
+          stderr,                    \
+          "NCCL error %s:%d '%s'\n", \
+          __FILE__,                  \
+          __LINE__,                  \
+          ncclGetErrorString(r));    \
+      exit(1);                       \
+    }                                \
+  } while (0)
+
+#define MLX_NCCL_TYPE_LIST(X) \
+  X(int8_t, ncclChar)         \
+  X(uint8_t, ncclUint8)       \
+  X(int32_t, ncclInt)         \
+  X(uint32_t, ncclUint32)     \
+  X(int64_t, ncclInt64)       \
+  X(uint64_t, ncclUint64)     \
+  X(float16_t, ncclHalf)      \
+  X(bfloat16_t, ncclBfloat16) \
+  X(float, ncclFloat)         \
+  X(double, ncclDouble)
+
+template <class>
+struct nccl_map {
+  static constexpr bool ok = false; // default: unsupported
+};
+
+#define MLX_DEF_NCCL_MAP(T, E)                 \
+  template <>                                  \
+  struct nccl_map<T> {                         \
+    static constexpr bool ok = true;           \
+    static constexpr ncclDataType_t value = E; \
+  };
+
+MLX_NCCL_TYPE_LIST(MLX_DEF_NCCL_MAP)
+#undef MLX_DEF_NCCL_MAP
+
+namespace detail {
+
+template <typename F>
+void dispatch_dtype(const array& arr, F&& f) {
+  dispatch_all_types(arr.dtype(), [&](auto type_tag) {
+    using T = MLX_GET_TYPE(type_tag);
+    if constexpr (nccl_map<T>::ok) {
+      f(type_tag, nccl_map<T>::value);
+    } else {
+      throw std::invalid_argument("[nccl] Unknown or unsupported dtype");
+    }
+  });
+}
+
+#ifndef _WIN32
+inline void sendAll(int sock, const void* buf, size_t len) {
+  const char* ptr = reinterpret_cast<const char*>(buf);
+  while (len > 0) {
+    ssize_t sent = send(sock, ptr, len, 0);
+    if (sent <= 0) {
+      perror("send");
+      exit(1);
+    }
+    ptr += sent;
+    len -= sent;
+  }
+}
+
+inline void recvAll(int sock, void* buf, size_t len) {
+  char* ptr = reinterpret_cast<char*>(buf);
+  while (len > 0) {
+    ssize_t rec = recv(sock, ptr, len, 0);
+    if (rec <= 0) {
+      perror("recv");
+      exit(1);
+    }
+    ptr += rec;
+    len -= rec;
+  }
+}
+#endif // _WIN32
+
+#ifndef _WIN32
+inline void bootstrap_unique_id(
+    ncclUniqueId& id,
+    int rank,
+    int size,
+    const std::string& initMethod) {
+  // Parse the init method to extract the host and port
+  if (initMethod.rfind("tcp://", 0) != 0)
+    throw;
+  auto hostport = initMethod.substr(6);
+  auto colon = hostport.find(':');
+  std::string host = hostport.substr(0, colon);
+  int port = std::stoi(hostport.substr(colon + 1));
+
+  if (rank == 0) {
+    // create a unique id on the rank 0
+    CHECK_NCCL(ncclGetUniqueId(&id));
+
+    // create a socket to send the unique id to all other ranks
+    int sock = socket(AF_INET, SOCK_STREAM, 0);
+
+    if (sock < 0) {
+      std::ostringstream msg;
+      msg << "[nccl] Couldn't create socket (error: " << errno << ")";
+      throw std::runtime_error(msg.str());
+    }
+
+    sockaddr_in serv = {};
+    serv.sin_family = AF_INET;
+    serv.sin_addr.s_addr = htonl(INADDR_ANY);
+    serv.sin_port = htons(port);
+
+    int reuse = 1;
+    // Without this, if rank-0 crashes or restarts process quickly,
+    // the OS might refuse to let binding to the same port, so reuse
+
+    if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0) {
+      std::ostringstream msg;
+      msg << "[nccl] setsockopt() failed: " << strerror(errno);
+      throw std::runtime_error(msg.str());
+    }
+
+    if (bind(sock, reinterpret_cast<sockaddr*>(&serv), sizeof(serv)) < 0) {
+      std::ostringstream msg;
+      msg << "[nccl] bind() failed: " << strerror(errno);
+      throw std::runtime_error(msg.str());
+    }
+    if (listen(sock, size - 1) < 0) {
+      std::ostringstream msg;
+      msg << "[nccl] listen() failed: " << strerror(errno);
+      throw std::runtime_error(msg.str());
+    }
+
+    for (int peer = 1; peer < size; ++peer) {
+      int conn = accept(sock, nullptr, nullptr);
+      if (conn < 0) {
+        std::ostringstream msg;
+        msg << "[nccl] accept() failed: " << strerror(errno);
+        throw std::runtime_error(msg.str());
+      }
+      sendAll(conn, &id, sizeof(id));
+      close(conn);
+    }
+    close(sock);
+
+  } else {
+    // Here we want to make sure that rank 0 has enough time to bind
+    // so we will retry to connect until elapsed time exceeds nccl_timeout
+    // this is particularity important for multinode setup
+
+    int sock = socket(AF_INET, SOCK_STREAM, 0);
+    if (sock < 0) {
+      std::ostringstream msg;
+      msg << "[nccl] socket() failed: " << strerror(errno);
+      throw std::runtime_error(msg.str());
+    }
+
+    hostent* he = gethostbyname(host.c_str());
+    if (!he) {
+      throw std::runtime_error("[nccl] lookup failed for host: " + host);
+    }
+    sockaddr_in serv = {};
+    serv.sin_family = AF_INET;
+    memcpy(&serv.sin_addr, he->h_addr_list[0], he->h_length);
+    serv.sin_port = htons(port);
+
+    const int timeout_ms = env::nccl_timeout(nccl_timeout);
+    bool connected = false;
+
+    const char* dbg = std::getenv("NCCL_DEBUG");
+    bool do_log = (dbg && std::string(dbg) == "INFO");
+
+    auto start = std::chrono::steady_clock::now();
+    int attempt = 0;
+
+    while (true) {
+      auto elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
+                            std::chrono::steady_clock::now() - start)
+                            .count();
+      if (elapsed_ms > timeout_ms)
+        break;
+      if (connect(sock, reinterpret_cast<sockaddr*>(&serv), sizeof(serv)) ==
+          0) {
+        connected = true;
+        if (do_log) {
+          std::cout << "[Rank " << rank << "] Connected successfully after "
+                    << elapsed_ms << " miliseconds" << std::endl;
+          break;
+        }
+      }
+      if (errno != ECONNREFUSED) {
+        break;
+      }
+      ++attempt;
+      std::this_thread::sleep_for(std::chrono::milliseconds(500));
+    }
+
+    if (!connected) {
+      std::ostringstream msg;
+      msg << "[Rank " << rank << "] connect() failed after " << timeout_ms
+          << " milliseconds and " << attempt << " retries: " << strerror(errno);
+      close(sock);
+      throw std::runtime_error(msg.str());
+    }
+    recvAll(sock, &id, sizeof(id));
+    close(sock);
+  }
+}
+#else // _WIN32
+inline void bootstrap_unique_id(
+    ncclUniqueId& id,
+    int rank,
+    int size,
+    const std::string& initMethod) {
+  throw std::runtime_error(
+      "[nccl] Distributed NCCL is not yet supported on Windows");
+}
+#endif // _WIN32
+
+} // namespace detail
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+class NCCLGroup : public GroupImpl {
+ public:
+  NCCLGroup(int worldRank, int worldSize, const std::string initMethod)
+      : rank_(worldRank),
+        size_(worldSize),
+        comm_(nullptr),
+        initMethod_(initMethod) {
+    if (initialized_)
+      return;
+    int ndev;
+    CHECK_CUDA(cudaGetDeviceCount(&ndev));
+    CHECK_CUDA(cudaSetDevice(rank_ % ndev));
+    detail::bootstrap_unique_id(uniqueId_, rank_, size_, initMethod_);
+    CHECK_NCCL(ncclCommInitRank(&comm_, size_, uniqueId_, rank_));
+    initialized_ = true;
+  }
+
+  ~NCCLGroup() {
+    ncclCommDestroy(comm_);
+    initialized_ = false;
+  }
+
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s, Device::gpu);
+  }
+
+  int rank() override {
+    return rank_;
+  }
+
+  int size() override {
+    return size_;
+  }
+
+  void all_sum(const array& input, array& output, Stream stream) override {
+    detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) {
+      using T = typename decltype(type_tag)::type;
+      all_reduce_impl<T>(input, output, stream, dt, ncclSum);
+    });
+  }
+
+  virtual std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    throw std::runtime_error("[nccl] Group split not supported.");
+  }
+
+  void all_gather(const array& input, array& output, Stream stream) override {
+    detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) {
+      using T = typename decltype(type_tag)::type;
+      auto& encoder = cu::get_command_encoder(stream);
+      CHECK_NCCL(ncclAllGather(
+          gpu_ptr<T>(input),
+          gpu_ptr<T>(output),
+          input.size(),
+          dt,
+          comm_,
+          encoder.stream()));
+    });
+  }
+
+  void send(const array& input, int dst, Stream stream) override {
+    throw std::runtime_error("[nccl] Send not supported in NCCL backend.");
+  }
+
+  void recv(array& output, int src, Stream stream) override {
+    throw std::runtime_error("[nccl] Recv not supported in NCCL backend.");
+  }
+
+  void all_max(const array& input, array& output, Stream stream) override {
+    detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) {
+      using T = typename decltype(type_tag)::type;
+      all_reduce_impl<T>(input, output, stream, dt, ncclMax);
+    });
+  }
+
+  void all_min(const array& input, array& output, Stream stream) override {
+    detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) {
+      using T = typename decltype(type_tag)::type;
+      all_reduce_impl<T>(input, output, stream, dt, ncclMin);
+    });
+  }
+
+  void sum_scatter(const array& input, array& output, Stream stream) override {
+    detail::dispatch_dtype(input, [&](auto type_tag, ncclDataType_t dt) {
+      using T = typename decltype(type_tag)::type;
+      reduce_scatter_impl<T>(input, output, stream, dt, ncclSum);
+    });
+  }
+
+  template <typename T>
+  void all_reduce_impl(
+      const array& input,
+      array& output,
+      Stream stream,
+      ncclDataType_t dt,
+      ncclRedOp_t op) {
+    auto& encoder = cu::get_command_encoder(stream);
+
+    CHECK_NCCL(ncclAllReduce(
+        gpu_ptr<T>(input),
+        gpu_ptr<T>(output),
+        input.size(),
+        dt,
+        op,
+        comm_,
+        encoder.stream()));
+  }
+
+  template <typename T>
+  void reduce_scatter_impl(
+      const array& input,
+      array& output,
+      Stream stream,
+      ncclDataType_t dt,
+      ncclRedOp_t op) {
+    auto& encoder = cu::get_command_encoder(stream);
+
+    CHECK_NCCL(ncclReduceScatter(
+        gpu_ptr<T>(input),
+        gpu_ptr<T>(output),
+        output.size(),
+        dt,
+        op,
+        comm_,
+        encoder.stream()));
+  }
+
+  int rank_, size_;
+  std::string initMethod_;
+  ncclUniqueId uniqueId_;
+  ncclComm_t comm_;
+  bool initialized_ = false;
+};
+
+bool is_available() {
+  return true;
+}
+
+namespace detail {
+std::string get_env_var_or_throw(const char* env_var_name, bool strict) {
+  const char* value = std::getenv(env_var_name);
+  if (value == nullptr && strict) {
+    std::ostringstream msg;
+    msg << "[nccl] Required environment variable '" << env_var_name
+        << "' is not set. "
+        << "Please set it before initializing the distributed backend.";
+    throw std::runtime_error(msg.str());
+  }
+  if (value == nullptr) {
+    return "";
+  }
+  return std::string(value);
+}
+} // namespace detail
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  std::string host = detail::get_env_var_or_throw("NCCL_HOST_IP", strict);
+  std::string port = detail::get_env_var_or_throw("NCCL_PORT", strict);
+  std::string rank_str = detail::get_env_var_or_throw("MLX_RANK", strict);
+  std::string n_nodes_str =
+      detail::get_env_var_or_throw("MLX_WORLD_SIZE", strict);
+  if (!strict &&
+      (host.empty() || port.empty() || rank_str.empty() ||
+       n_nodes_str.empty())) {
+    return nullptr;
+  }
+
+  int rank = std::stoi(rank_str);
+  int n_nodes = std::stoi(n_nodes_str);
+  std::string init_method = "tcp://" + host + ":" + port;
+
+  return std::make_shared<NCCLGroup>(rank, n_nodes, init_method);
+}
+} // namespace mlx::core::distributed::nccl
diff --git a/macos/third_party/mlx/mlx/distributed/nccl/nccl.h b/macos/third_party/mlx/mlx/distributed/nccl/nccl.h
new file mode 100644
index 0000000..5370d2d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/nccl/nccl.h
@@ -0,0 +1,12 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::nccl {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available();
+std::shared_ptr<GroupImpl> init(bool strict = false);
+
+} // namespace mlx::core::distributed::nccl
diff --git a/macos/third_party/mlx/mlx/distributed/nccl/nccl_stub/CMakeLists.txt b/macos/third_party/mlx/mlx/distributed/nccl/nccl_stub/CMakeLists.txt
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/nccl/nccl_stub/CMakeLists.txt
@@ -0,0 +1 @@
+
diff --git a/macos/third_party/mlx/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp b/macos/third_party/mlx/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp
new file mode 100644
index 0000000..cd1b522
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp
@@ -0,0 +1,54 @@
+#if !defined(_WIN32)
+#include <arpa/inet.h>
+#endif
+#include <cuda_runtime.h>
+#include <nccl.h>
+
+ncclResult_t ncclGetUniqueId(ncclUniqueId*) {
+  return ncclSuccess;
+}
+
+const char* ncclGetErrorString(ncclResult_t result) {
+  return nullptr;
+}
+
+ncclResult_t
+ncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank) {
+  return ncclSuccess;
+}
+
+ncclResult_t ncclCommDestroy(ncclComm_t comm) {
+  return ncclSuccess;
+}
+
+ncclResult_t ncclAllGather(
+    const void* sendbuff,
+    void* recvbuff,
+    size_t sendcount,
+    ncclDataType_t datatype,
+    ncclComm_t comm,
+    cudaStream_t stream) {
+  return ncclSuccess;
+}
+
+ncclResult_t ncclAllReduce(
+    const void* sendbuff,
+    void* recvbuff,
+    size_t count,
+    ncclDataType_t datatype,
+    ncclRedOp_t op,
+    ncclComm_t comm,
+    cudaStream_t stream) {
+  return ncclSuccess;
+}
+
+ncclResult_t ncclReduceScatter(
+    const void* sendbuff,
+    void* recvbuff,
+    size_t recvcount,
+    ncclDataType_t datatype,
+    ncclRedOp_t op,
+    ncclComm_t comm,
+    cudaStream_t stream) {
+  return ncclSuccess;
+}
diff --git a/macos/third_party/mlx/mlx/distributed/nccl/no_nccl.cpp b/macos/third_party/mlx/mlx/distributed/nccl/no_nccl.cpp
new file mode 100644
index 0000000..1be256a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/nccl/no_nccl.cpp
@@ -0,0 +1,20 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/distributed/nccl/nccl.h"
+
+namespace mlx::core::distributed::nccl {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available() {
+  return false;
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  if (strict) {
+    throw std::runtime_error("Cannot initialize nccl distributed backend.");
+  }
+  return nullptr;
+}
+
+} // namespace mlx::core::distributed::nccl
diff --git a/macos/third_party/mlx/mlx/distributed/ops.cpp b/macos/third_party/mlx/mlx/distributed/ops.cpp
new file mode 100644
index 0000000..1762f0e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/ops.cpp
@@ -0,0 +1,186 @@
+// Copyright © 2024 Apple Inc.
+
+#include <sstream>
+
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/ops.h"
+#include "mlx/distributed/primitives.h"
+
+namespace mlx::core::distributed {
+
+namespace {
+
+Group to_group(std::optional<Group> group) {
+  if (group.has_value()) {
+    return group.value();
+  } else {
+    return distributed::init();
+  }
+}
+
+} // namespace
+
+array all_sum(
+    const array& x,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    return x;
+  }
+  auto stream = detail::communication_stream(group, s);
+
+  return array(
+      x.shape(),
+      x.dtype(),
+      std::make_shared<AllReduce>(stream, group, AllReduce::Sum),
+      {x});
+}
+
+array all_max(
+    const array& x,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    return x;
+  }
+  auto stream = detail::communication_stream(group, s);
+
+  return array(
+      x.shape(),
+      x.dtype(),
+      std::make_shared<AllReduce>(stream, group, AllReduce::Max),
+      {x});
+}
+
+array all_min(
+    const array& x,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    return x;
+  }
+  auto stream = detail::communication_stream(group, s);
+
+  return array(
+      x.shape(),
+      x.dtype(),
+      std::make_shared<AllReduce>(stream, group, AllReduce::Min),
+      {x});
+}
+
+array all_gather(
+    const array& x,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    return x;
+  }
+  auto stream = detail::communication_stream(group, s);
+
+  auto result_shape = x.shape();
+  if (result_shape.size() == 0) {
+    result_shape.push_back(group.size());
+  } else {
+    result_shape[0] *= group.size();
+  }
+  return array(
+      std::move(result_shape),
+      x.dtype(),
+      std::make_shared<AllGather>(stream, group),
+      {x});
+}
+
+array send(
+    const array& x,
+    int dst,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    throw std::invalid_argument("Cannot send to a singleton group");
+  }
+  auto stream = detail::communication_stream(group, s);
+
+  if (dst < 0 || dst >= group.size()) {
+    std::ostringstream msg;
+    msg << "Invalid destination=" << dst << " for a group of size "
+        << group.size();
+    throw std::invalid_argument(msg.str());
+  }
+
+  return array(
+      x.shape(), x.dtype(), std::make_shared<Send>(stream, group, dst), {x});
+}
+
+array recv(
+    Shape shape,
+    Dtype dtype,
+    int src,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    throw std::invalid_argument("Cannot recv from a singleton group");
+  }
+  auto stream = detail::communication_stream(group, s);
+
+  if (src < 0 || src >= group.size()) {
+    std::ostringstream msg;
+    msg << "Invalid source=" << src << " for a group of size " << group.size();
+    throw std::invalid_argument(msg.str());
+  }
+
+  return array(
+      std::move(shape),
+      std::move(dtype),
+      std::make_shared<Recv>(stream, group, src),
+      std::vector<array>{});
+}
+
+array recv_like(
+    const array& x,
+    int src,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return recv(x.shape(), x.dtype(), src, group_, s);
+}
+
+array sum_scatter(
+    const array& x,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+  if (group.size() == 1) {
+    return x;
+  }
+  if (x.shape()[0] % group.size() != 0) {
+    std::ostringstream msg;
+    msg << "[sum_scatter] Invalid shape=" << x.shape()
+        << " for a group of size " << group.size()
+        << ". The first dimension (axis 0) must be divisible by the group size.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto result_shape = x.shape();
+  result_shape[0] /= group.size();
+  auto stream = detail::communication_stream(group, s);
+
+  return array(
+      std::move(result_shape),
+      x.dtype(),
+      std::make_shared<ReduceScatter>(stream, group, ReduceScatter::Sum),
+      {x});
+}
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/distributed/ops.h b/macos/third_party/mlx/mlx/distributed/ops.h
new file mode 100644
index 0000000..e223c5b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/ops.h
@@ -0,0 +1,57 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <optional>
+
+#include "mlx/api.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::distributed {
+
+MLX_API array all_sum(
+    const array& x,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array all_gather(
+    const array& x,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice S = {});
+
+MLX_API array send(
+    const array& x,
+    int dst,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array recv(
+    Shape shape,
+    Dtype dtype,
+    int src,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array recv_like(
+    const array& x,
+    int src,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array all_max(
+    const array& x,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array all_min(
+    const array& x,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array sum_scatter(
+    const array& x,
+    std::optional<Group> group = std::nullopt,
+    StreamOrDevice s = {});
+
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/distributed/primitives.cpp b/macos/third_party/mlx/mlx/distributed/primitives.cpp
new file mode 100644
index 0000000..5e8d532
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/primitives.cpp
@@ -0,0 +1,95 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/allocator.h"
+#include "mlx/distributed/ops.h"
+#include "mlx/distributed/primitives.h"
+#include "mlx/ops.h"
+
+namespace mlx::core::distributed {
+
+std::pair<std::vector<array>, std::vector<int>> AllReduce::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  switch (reduce_type_) {
+    case Sum:
+      return {{all_sum(inputs[0], group(), stream())}, axes};
+    case Max:
+      return {{all_max(inputs[0], group(), stream())}, axes};
+    case Min:
+      return {{all_min(inputs[0], group(), stream())}, axes};
+    default:
+
+      throw std::runtime_error(
+          "Only all reduce sum, max and min are supported for now");
+  }
+}
+
+std::vector<array> AllReduce::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  switch (reduce_type_) {
+    case Sum:
+      return {all_sum(tangents[0], group(), stream())};
+    case Max:
+      return {all_max(tangents[0], group(), stream())};
+    case Min:
+      return {all_min(tangents[0], group(), stream())};
+    default:
+      throw std::runtime_error(
+          "Only all reduce sum, max and min are supported for now");
+  }
+}
+
+std::vector<array> AllReduce::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>& outputs) {
+  return cotangents;
+}
+
+std::pair<std::vector<array>, std::vector<int>> AllGather::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  return {{all_gather(inputs[0], group(), stream())}, axes};
+}
+
+std::vector<array> AllGather::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {all_gather(tangents[0], group(), stream())};
+}
+
+std::vector<array> AllGather::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  auto g = group();
+  auto ndim = primals[0].ndim();
+  Shape starts(primals[0].ndim(), 0);
+  auto stops = primals[0].shape();
+  if (ndim == 0) {
+    starts.push_back(0);
+    stops.push_back(1);
+  }
+  starts[0] = g.rank() * stops[0];
+  stops[0] += starts[0];
+  auto out = slice(cotangents[0], starts, stops);
+  if (ndim == 0) {
+    out = squeeze(out, 0);
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Send::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  return {{send(inputs[0], dst_, group(), stream())}, axes};
+}
+
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/distributed/primitives.h b/macos/third_party/mlx/mlx/distributed/primitives.h
new file mode 100644
index 0000000..18a0d65
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/primitives.h
@@ -0,0 +1,156 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core::distributed {
+
+class DistPrimitive : public Primitive {
+ public:
+  DistPrimitive(Stream stream, Group group)
+      : Primitive(stream), group_(group) {}
+
+  const Group& group() const {
+    return group_;
+  }
+
+ private:
+  Group group_;
+};
+
+class AllReduce : public DistPrimitive {
+ public:
+  enum ReduceType { And, Or, Sum, Prod, Min, Max };
+
+  AllReduce(Stream stream, Group group, ReduceType reduce_type)
+      : DistPrimitive(stream, group), reduce_type_(reduce_type) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  std::pair<std::vector<array>, std::vector<int>> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+  std::vector<array> jvp(
+      const std::vector<array>& primals,
+      const std::vector<array>& tangents,
+      const std::vector<int>& argnums) override;
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  const char* name() const override {
+    switch (reduce_type_) {
+      case And:
+        return "And AllReduce";
+      case Or:
+        return "Or AllReduce";
+      case Sum:
+        return "Sum AllReduce";
+      case Prod:
+        return "Prod AllReduce";
+      case Min:
+        return "Min AllReduce";
+      case Max:
+        return "Max AllReduce";
+    }
+    return "<unknwon AllReduce>";
+  }
+
+ private:
+  ReduceType reduce_type_;
+};
+
+class AllGather : public DistPrimitive {
+ public:
+  AllGather(Stream stream, Group group) : DistPrimitive(stream, group) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::pair<std::vector<array>, std::vector<int>> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+  std::vector<array> jvp(
+      const std::vector<array>& primals,
+      const std::vector<array>& tangents,
+      const std::vector<int>& argnums) override;
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(AllGather);
+};
+
+class Send : public DistPrimitive {
+ public:
+  Send(Stream stream, Group group, int dst)
+      : DistPrimitive(stream, group), dst_(dst) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  std::pair<std::vector<array>, std::vector<int>> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  DEFINE_NAME(Send);
+
+ private:
+  int dst_;
+};
+
+class Recv : public DistPrimitive {
+ public:
+  Recv(Stream stream, Group group, int src)
+      : DistPrimitive(stream, group), src_(src) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(Recv);
+
+ private:
+  int src_;
+};
+
+class ReduceScatter : public DistPrimitive {
+ public:
+  enum ReduceType { Sum, Min, Max };
+  ReduceScatter(Stream stream, Group group, ReduceType reduce_type)
+      : DistPrimitive(stream, group), reduce_type_(reduce_type) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  const char* name() const override {
+    switch (reduce_type_) {
+      case Sum:
+        return "Sum ReduceScatter";
+      case Min:
+        return "Min ReduceScatter";
+      case Max:
+        return "Max ReduceScatter";
+    }
+    return "<unknwon ReduceScatter>";
+  }
+
+ private:
+  ReduceType reduce_type_;
+};
+} // namespace mlx::core::distributed
diff --git a/macos/third_party/mlx/mlx/distributed/reduction_ops.h b/macos/third_party/mlx/mlx/distributed/reduction_ops.h
new file mode 100644
index 0000000..02777be
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/reduction_ops.h
@@ -0,0 +1,38 @@
+// Copyright © 2025 Apple Inc.
+
+namespace mlx::core::distributed::detail {
+
+template <typename T>
+struct SumOp {
+  void operator()(const T* input, T* output, size_t N) const {
+    while (N-- > 0) {
+      *output += *input;
+      input++;
+      output++;
+    }
+  }
+};
+
+template <typename T>
+struct MaxOp {
+  void operator()(const T* input, T* output, size_t N) const {
+    while (N-- > 0) {
+      *output = std::max(*output, *input);
+      input++;
+      output++;
+    }
+  }
+};
+
+template <typename T>
+struct MinOp {
+  void operator()(const T* input, T* output, size_t N) const {
+    while (N-- > 0) {
+      *output = std::min(*output, *input);
+      input++;
+      output++;
+    }
+  }
+};
+
+} // namespace mlx::core::distributed::detail
diff --git a/macos/third_party/mlx/mlx/distributed/ring/CMakeLists.txt b/macos/third_party/mlx/mlx/distributed/ring/CMakeLists.txt
new file mode 100644
index 0000000..346f7cd
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/ring/CMakeLists.txt
@@ -0,0 +1,5 @@
+if(MLX_BUILD_CPU AND NOT WIN32)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/ring.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_ring.cpp)
+endif()
diff --git a/macos/third_party/mlx/mlx/distributed/ring/no_ring.cpp b/macos/third_party/mlx/mlx/distributed/ring/no_ring.cpp
new file mode 100644
index 0000000..0c31d28
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/ring/no_ring.cpp
@@ -0,0 +1,20 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/distributed/ring/ring.h"
+
+namespace mlx::core::distributed::ring {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available() {
+  return false;
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  if (strict) {
+    throw std::runtime_error("Cannot initialize ring distributed backend.");
+  }
+  return nullptr;
+}
+
+} // namespace mlx::core::distributed::ring
diff --git a/macos/third_party/mlx/mlx/distributed/ring/ring.cpp b/macos/third_party/mlx/mlx/distributed/ring/ring.cpp
new file mode 100644
index 0000000..ea40042
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/ring/ring.cpp
@@ -0,0 +1,870 @@
+// Copyright © 2024 Apple Inc.
+
+#include <fcntl.h>
+#include <netinet/tcp.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+#include <chrono>
+#include <fstream>
+#include <future>
+#include <iostream>
+#include <list>
+#include <sstream>
+#include <thread>
+#include <unordered_map>
+
+#include <json.hpp>
+
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/distributed_impl.h"
+#include "mlx/distributed/reduction_ops.h"
+#include "mlx/distributed/utils.h"
+#include "mlx/threadpool.h"
+
+#ifndef SOL_TCP
+#define SOL_TCP IPPROTO_TCP
+#endif
+
+#define SWITCH_TYPE(x, ...)  \
+  switch ((x).dtype()) {     \
+    case bool_: {            \
+      using T = bool;        \
+      __VA_ARGS__;           \
+    } break;                 \
+    case int8: {             \
+      using T = int8_t;      \
+      __VA_ARGS__;           \
+    } break;                 \
+    case int16: {            \
+      using T = int16_t;     \
+      __VA_ARGS__;           \
+    } break;                 \
+    case int32: {            \
+      using T = int32_t;     \
+      __VA_ARGS__;           \
+    } break;                 \
+    case int64: {            \
+      using T = int64_t;     \
+      __VA_ARGS__;           \
+    } break;                 \
+    case uint8: {            \
+      using T = uint8_t;     \
+      __VA_ARGS__;           \
+    } break;                 \
+    case uint16: {           \
+      using T = uint16_t;    \
+      __VA_ARGS__;           \
+    } break;                 \
+    case uint32: {           \
+      using T = uint32_t;    \
+      __VA_ARGS__;           \
+    } break;                 \
+    case uint64: {           \
+      using T = uint64_t;    \
+      __VA_ARGS__;           \
+    } break;                 \
+    case bfloat16: {         \
+      using T = bfloat16_t;  \
+      __VA_ARGS__;           \
+    } break;                 \
+    case float16: {          \
+      using T = float16_t;   \
+      __VA_ARGS__;           \
+    } break;                 \
+    case float32: {          \
+      using T = float;       \
+      __VA_ARGS__;           \
+    } break;                 \
+    case float64: {          \
+      using T = double;      \
+      __VA_ARGS__;           \
+    } break;                 \
+    case complex64: {        \
+      using T = complex64_t; \
+      __VA_ARGS__;           \
+    } break;                 \
+  }
+
+namespace mlx::core::distributed::ring {
+
+constexpr const size_t ALL_SUM_SIZE = 8 * 1024 * 1024;
+constexpr const size_t ALL_SUM_BUFFERS = 2;
+constexpr const int CONN_ATTEMPTS = 5;
+constexpr const int CONN_WAIT = 1000;
+constexpr const char* RING_TAG = "[ring]";
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+using json = nlohmann::json;
+using namespace std::chrono_literals;
+
+namespace {
+
+template <typename T>
+void log(std::ostream& os, T first) {
+  os << first << std::endl;
+}
+
+template <typename T, typename... Args>
+void log(std::ostream& os, T first, Args... args) {
+  log(os << first << " ", args...);
+}
+
+template <typename... Args>
+void log_info(bool verbose, Args... args) {
+  if (!verbose) {
+    return;
+  }
+
+  log(std::cerr, "[ring]", args...);
+}
+
+template <typename T, typename U>
+decltype(T() * U()) ceildiv(T a, U b) {
+  return (a + b - 1) / b;
+}
+
+class SocketThread {
+ public:
+  SocketThread(int fd) : fd_(fd), stop_(false) {
+    worker_ = std::thread(&SocketThread::worker, this);
+    int flags = fcntl(fd, F_GETFL, 0);
+    fcntl(fd, F_SETFL, flags | O_NONBLOCK);
+  }
+  ~SocketThread() {
+    stop_ = true;
+    condition_.notify_all();
+    worker_.join();
+    int flags = fcntl(fd_, F_GETFL, 0);
+    fcntl(fd_, F_SETFL, flags & ~O_NONBLOCK);
+  }
+
+  template <typename T>
+  std::future<void> send(const T* buffer, size_t size) {
+    return send_impl(reinterpret_cast<const char*>(buffer), size * sizeof(T));
+  }
+
+  template <typename T>
+  std::future<void> recv(T* buffer, size_t size) {
+    return recv_impl(reinterpret_cast<char*>(buffer), size * sizeof(T));
+  }
+
+ private:
+  struct SocketTask {
+    SocketTask(void* b, size_t s, std::promise<void>&& p)
+        : buffer(b), size(s), promise(std::move(p)) {}
+    SocketTask(SocketTask&& t)
+        : buffer(t.buffer), size(t.size), promise(std::move(t.promise)) {}
+    void* buffer;
+    size_t size;
+    std::promise<void> promise;
+  };
+
+  std::future<void> send_impl(const char* buffer, size_t size) {
+    std::promise<void> send_completed_promise;
+    auto send_completed_future = send_completed_promise.get_future();
+    if (size == 0) {
+      send_completed_promise.set_value();
+      return send_completed_future;
+    }
+
+    {
+      std::unique_lock lock(queue_mutex_);
+      sends_.emplace_back(SocketTask(
+          const_cast<char*>(buffer), size, std::move(send_completed_promise)));
+    }
+    condition_.notify_one();
+    return send_completed_future;
+  }
+
+  std::future<void> recv_impl(char* buffer, size_t size) {
+    std::promise<void> recv_completed_promise;
+    auto recv_completed_future = recv_completed_promise.get_future();
+    if (size == 0) {
+      recv_completed_promise.set_value();
+      return recv_completed_future;
+    }
+
+    {
+      std::unique_lock lock(queue_mutex_);
+      recvs_.emplace_back(
+          SocketTask(buffer, size, std::move(recv_completed_promise)));
+    }
+    condition_.notify_one();
+    return recv_completed_future;
+  }
+
+  bool have_tasks() {
+    return !(sends_.empty() && recvs_.empty());
+  }
+
+  void worker() {
+    int error_count = 0;
+    bool delete_recv = false;
+    bool delete_send = false;
+    while (true) {
+      {
+        std::unique_lock lock(queue_mutex_);
+
+        if (delete_recv) {
+          recvs_.front().promise.set_value();
+          recvs_.pop_front();
+          delete_recv = false;
+        }
+        if (delete_send) {
+          sends_.front().promise.set_value();
+          sends_.pop_front();
+          delete_send = false;
+        }
+
+        if (stop_) {
+          return;
+        }
+
+        if (!have_tasks()) {
+          condition_.wait(lock, [this] { return stop_ || have_tasks(); });
+          if (stop_) {
+            return;
+          }
+        }
+      }
+
+      if (!recvs_.empty()) {
+        auto& task = recvs_.front();
+        ssize_t r = ::recv(fd_, task.buffer, task.size, 0);
+        if (r > 0) {
+          task.buffer = static_cast<char*>(task.buffer) + r;
+          task.size -= r;
+          delete_recv = task.size == 0;
+          error_count = 0;
+        } else if (errno != EAGAIN) {
+          error_count++;
+          log_info(
+              true, "Receiving from socket", fd_, "failed with errno", errno);
+        }
+      }
+      if (!sends_.empty()) {
+        auto& task = sends_.front();
+        ssize_t r = ::send(fd_, task.buffer, task.size, 0);
+        if (r > 0) {
+          task.buffer = static_cast<char*>(task.buffer) + r;
+          task.size -= r;
+          delete_send = task.size == 0;
+          error_count = 0;
+        } else if (errno != EAGAIN) {
+          error_count++;
+          log_info(true, "Sending to socket", fd_, "failed with errno", errno);
+        }
+      }
+
+      if (error_count >= 10) {
+        log_info(true, "Too many send/recv errors. Aborting...");
+        return;
+      }
+    }
+  }
+
+  int fd_;
+  bool stop_;
+  std::thread worker_;
+  std::mutex queue_mutex_;
+  std::condition_variable condition_;
+  std::list<SocketTask> sends_;
+  std::list<SocketTask> recvs_;
+};
+
+class CommunicationThreads {
+ public:
+  void add(const std::vector<int>& sockets) {
+    for (int sock : sockets) {
+      threads_.emplace(sock, sock);
+    }
+  }
+
+  template <typename T>
+  std::future<void> send(int socket, T* buffer, size_t size) {
+    return threads_.at(socket).send<T>(buffer, size);
+  }
+
+  template <typename T>
+  std::future<void> recv(int socket, T* buffer, size_t size) {
+    return threads_.at(socket).recv<T>(buffer, size);
+  }
+
+ private:
+  std::unordered_map<int, SocketThread> threads_;
+};
+
+/**
+ * Load all addresses from the json hostfile. The hostfile is a list of
+ * addresses in order of rank. For each rank there can be many addresses so
+ * that we can have multiple connections between peers.
+ *
+ * For example:
+ *  [
+ *    ["ip1:5000", "ip1:5001"],
+ *    ["ip2:5000", "ip2:5001"],
+ *    ["ip3:5000", "ip3:5001"],
+ *  ]
+ */
+std::vector<std::vector<detail::address_t>> load_nodes(const char* hostfile) {
+  std::vector<std::vector<detail::address_t>> nodes;
+  std::ifstream f(hostfile);
+
+  json hosts = json::parse(f);
+  for (auto& h : hosts) {
+    std::vector<detail::address_t> host;
+    for (auto& ips : h) {
+      host.push_back(std::move(detail::parse_address(ips.get<std::string>())));
+    }
+    nodes.push_back(std::move(host));
+  }
+
+  return nodes;
+}
+
+/**
+ * Create a socket and accept one connection for each of the provided
+ * addresses.
+ */
+std::vector<int> accept_connections(
+    const std::vector<detail::address_t>& addresses) {
+  std::vector<int> sockets;
+  int success;
+
+  for (auto& address : addresses) {
+    detail::TCPSocket socket(RING_TAG);
+    socket.listen(RING_TAG, address);
+    sockets.push_back(socket.accept(RING_TAG).detach());
+  }
+
+  return sockets;
+}
+
+/**
+ * The counterpoint of `accept_connections`. Basically connect to each of the
+ * provided addresses.
+ */
+std::vector<int> make_connections(
+    const std::vector<detail::address_t>& addresses,
+    bool verbose) {
+  std::vector<int> sockets;
+  int success;
+
+  for (auto& address : addresses) {
+    sockets.push_back(
+        detail::TCPSocket::connect(
+            RING_TAG,
+            address,
+            CONN_ATTEMPTS,
+            CONN_WAIT,
+            [verbose](int attempt, int wait) {
+              log_info(
+                  verbose,
+                  "Attempt",
+                  attempt,
+                  "waiting",
+                  wait,
+                  "ms (error:",
+                  errno,
+                  ")");
+            })
+            .detach());
+  }
+
+  return sockets;
+}
+
+} // namespace
+
+class RingGroup : public GroupImpl {
+ public:
+  RingGroup(
+      int rank,
+      std::vector<std::vector<detail::address_t>> nodes,
+      bool verbose)
+      : rank_(rank), verbose_(verbose), pool_(0) {
+    if (rank_ > 0 && rank_ >= nodes.size()) {
+      throw std::runtime_error(
+          "[ring] Rank cannot be larger than the size of the group");
+    }
+
+    size_ = nodes.size();
+    int connect_to = (rank_ + 1) % size_;
+
+    // We define the connection order by having the rank_ == size_ - 1 connect
+    // first and accept after.
+    if (rank_ < connect_to) {
+      log_info(verbose_, "Rank", rank_, "accepting");
+      sockets_left_ = accept_connections(nodes[rank_]);
+      log_info(verbose_, "Rank", rank_, "connecting to", connect_to);
+      sockets_right_ = make_connections(nodes[connect_to], verbose);
+    } else {
+      log_info(verbose_, "Rank", rank_, "connecting to", connect_to);
+      sockets_right_ = make_connections(nodes[connect_to], verbose);
+      log_info(verbose_, "Rank", rank_, "accepting");
+      sockets_left_ = accept_connections(nodes[rank_]);
+    }
+
+    // Failure if we couldn't make right or left sockets
+    if (sockets_right_.empty()) {
+      std::ostringstream msg;
+      msg << "[ring] Rank " << rank_ << " has no sockets to the right.";
+      throw std::invalid_argument(msg.str());
+    }
+    if (sockets_left_.empty()) {
+      std::ostringstream msg;
+      msg << "[ring] Rank " << rank_ << " has no sockets to the left.";
+      throw std::invalid_argument(msg.str());
+    }
+
+    // The following could be relaxed since we can define non-homogeneous rings
+    // but it makes things a bit simpler for now.
+    if (sockets_right_.size() != sockets_left_.size()) {
+      std::ostringstream msg;
+      msg << "[ring] It is required to have as many connections to the left as "
+          << "to the right but rank " << rank_ << " has "
+          << sockets_right_.size() << " connections to the right and "
+          << sockets_left_.size() << " to the left.";
+      throw std::invalid_argument(msg.str());
+    }
+
+    // Configure all sockets to use TCP no delay.
+    int one = 1;
+    for (int i = 0; i < sockets_right_.size(); i++) {
+      setsockopt(sockets_right_[i], SOL_TCP, TCP_NODELAY, &one, sizeof(one));
+      setsockopt(sockets_left_[i], SOL_TCP, TCP_NODELAY, &one, sizeof(one));
+    }
+
+    // Start the all reduce threads. One all reduce per direction per ring.
+    pool_.resize(sockets_right_.size() + sockets_left_.size());
+
+    // Create a communication thread per socket. This also converts them to
+    // non-blocking.
+    comm_.add(sockets_right_);
+    comm_.add(sockets_left_);
+
+    // Allocate buffers for the all sum
+    buffers_.resize(
+        (sockets_right_.size() + sockets_left_.size()) * ALL_SUM_BUFFERS *
+        ALL_SUM_SIZE);
+  }
+
+  ~RingGroup() {
+    for (auto s : sockets_right_) {
+      shutdown(s, 2);
+      close(s);
+    }
+    for (auto s : sockets_left_) {
+      shutdown(s, 2);
+      close(s);
+    }
+  }
+
+  Stream communication_stream(StreamOrDevice s) override {
+    return to_stream(s, Device::cpu);
+  }
+
+  int rank() override {
+    return rank_;
+  }
+
+  int size() override {
+    return size_;
+  }
+
+  void all_sum(const array& input, array& output, Stream stream) override {
+    SWITCH_TYPE(
+        output, all_reduce<T>(input, output, stream, detail::SumOp<T>()));
+  }
+
+  void all_max(const array& input, array& output, Stream stream) override {
+    SWITCH_TYPE(
+        output, all_reduce<T>(input, output, stream, detail::MaxOp<T>()));
+  }
+
+  void all_min(const array& input, array& output, Stream stream) override {
+    SWITCH_TYPE(
+        output, all_reduce<T>(input, output, stream, detail::MinOp<T>()));
+  }
+
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    throw std::runtime_error("[ring] Group split not supported.");
+  }
+
+  void all_gather(const array& input, array& output, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.set_output_array(output);
+    encoder.dispatch([input_ptr = input.data<char>(),
+                      nbytes = input.nbytes(),
+                      output_ptr = output.data<char>(),
+                      this]() {
+      constexpr size_t min_send_size = 262144;
+      size_t n_gathers = std::max(
+          std::min(
+              sockets_right_.size() + sockets_left_.size(),
+              nbytes / min_send_size),
+          size_t(1));
+      size_t bytes_per_gather = ceildiv(nbytes, n_gathers);
+      std::vector<std::future<void>> all_gathers;
+      for (int i = 0; i < n_gathers; i++) {
+        auto offset = i * bytes_per_gather;
+        all_gathers.emplace_back(pool_.enqueue(
+            std::bind(
+                &RingGroup::all_gather_impl,
+                this,
+                input_ptr + offset,
+                output_ptr + offset,
+                nbytes,
+                offset + bytes_per_gather > nbytes ? nbytes - offset
+                                                   : bytes_per_gather,
+                sockets_right_[i / 2],
+                sockets_left_[i / 2],
+                (i % 2) ? -1 : 1)));
+      }
+      for (auto& f : all_gathers) {
+        f.wait();
+      }
+    });
+  }
+
+  void send(const array& input, int dst, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_input_array(input);
+    encoder.dispatch(
+        [input_ptr = input.data<char>(), nbytes = input.nbytes(), dst, this]() {
+          int right = (rank_ + 1) % size_;
+          int left = (rank_ + size_ - 1) % size_;
+          if (dst == right) {
+            send(sockets_right_, input_ptr, nbytes);
+          } else if (dst == left) {
+            send(sockets_left_, input_ptr, nbytes);
+          } else {
+            std::ostringstream msg;
+            msg << "[ring] Send only supported to direct neighbors "
+                << "but tried to send to " << dst << " from " << rank_
+                << std::endl;
+            throw std::runtime_error(msg.str());
+          }
+        });
+  }
+
+  void recv(array& out, int src, Stream stream) override {
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_output_array(out);
+    encoder.dispatch(
+        [out_ptr = out.data<char>(), nbytes = out.nbytes(), src, this]() {
+          // NOTE: We 'll check the sockets with the opposite order of send so
+          // that they work even with 2 nodes where left and right is the same
+          // neighbor.
+          int right = (rank_ + 1) % size_;
+          int left = (rank_ + size_ - 1) % size_;
+          if (src == left) {
+            recv(sockets_left_, out_ptr, nbytes);
+          } else if (src == right) {
+            recv(sockets_right_, out_ptr, nbytes);
+          } else {
+            std::ostringstream msg;
+            msg << "[ring] Recv only supported from direct neighbors "
+                << "but tried to recv from " << src << " to " << rank_
+                << std::endl;
+            throw std::runtime_error(msg.str());
+          }
+        });
+  }
+
+  void sum_scatter(const array& input, array& output, Stream stream) override {
+    throw std::runtime_error("[ring] sum_scatter not supported.");
+  }
+
+ private:
+  template <typename T, typename ReduceOp>
+  void all_reduce(
+      const array& input,
+      array& output,
+      Stream stream,
+      ReduceOp reduce_op) {
+    auto in_ptr = input.data<char>();
+    auto out_ptr = output.data<char>();
+    auto& encoder = cpu::get_command_encoder(stream);
+    encoder.set_output_array(output);
+    encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() {
+      // If the input data cannot be split into size_ segments then copy it and
+      // all reduce a local buffer prefilled with 0s.
+      size_t nbytes = size * sizeof(T);
+      if (size < size_) {
+        // TODO: Maybe allocate dynamically so we don't have the constraint
+        // below?
+        if (sizeof(T) * size_ > 1024) {
+          std::ostringstream msg;
+          msg << "Can't perform the ring all reduce of " << size
+              << " elements with a ring of size " << size_;
+          throw std::runtime_error(msg.str());
+        }
+
+        char buffer[1024];
+        std::memset(buffer, 0, size_ * sizeof(T));
+        std::memcpy(buffer, in_ptr, nbytes);
+        all_reduce_impl<T, ReduceOp>(
+            reinterpret_cast<T*>(buffers_.data()),
+            reinterpret_cast<T*>(buffer),
+            size_,
+            sockets_right_[0],
+            sockets_left_[0],
+            -1,
+            reduce_op);
+        std::memcpy(out_ptr, buffer, nbytes);
+        return;
+      }
+
+      // If not inplace all reduce then copy the input to the output first
+      if (in_ptr != out_ptr) {
+        std::memcpy(out_ptr, in_ptr, nbytes);
+      }
+
+      // Split the all reduces so that each member has at least 1 buffer to
+      // send/recv per segment.
+      constexpr size_t min_send_size = 262144;
+      size_t n_reduces = std::max(
+          std::min(
+              sockets_right_.size() + sockets_left_.size(),
+              nbytes / (size_ * min_send_size)),
+          size_t(1));
+      size_t step = ceildiv(size, n_reduces);
+      std::vector<std::future<void>> all_sums;
+
+      for (int i = 0; i < n_reduces; i++) {
+        all_sums.emplace_back(pool_.enqueue(
+            std::bind(
+                &RingGroup::all_reduce_impl<T, ReduceOp>,
+                this,
+                reinterpret_cast<T*>(
+                    buffers_.data() + i * ALL_SUM_SIZE * ALL_SUM_BUFFERS),
+                reinterpret_cast<T*>(out_ptr) + i * step,
+                std::min(size, (i + 1) * step) - i * step,
+                sockets_right_[i / 2],
+                sockets_left_[i / 2],
+                (i % 2) ? -1 : 1,
+                reduce_op)));
+      }
+      for (auto& f : all_sums) {
+        f.wait();
+      }
+    });
+  }
+
+  template <typename T, typename ReduceOp>
+  void all_reduce_impl(
+      T* buffer,
+      T* data,
+      size_t data_size,
+      int socket_right,
+      int socket_left,
+      int direction,
+      ReduceOp reduce_op) {
+    // Choose which socket we send to and recv from
+    int socket_send = (direction < 0) ? socket_right : socket_left;
+    int socket_recv = (direction < 0) ? socket_left : socket_right;
+
+    // We split the data into `size_` segments of size `segment_size` and each
+    // of these in smaller segments of ALL_SUM_SIZE which we 'll call packets.
+    size_t segment_size = ceildiv(data_size, size_);
+    size_t BUFFER_SIZE = std::max(
+        size_t(32768), std::min(ALL_SUM_SIZE / sizeof(T), segment_size / 2));
+    size_t n_packets = ceildiv(segment_size, BUFFER_SIZE);
+
+    // Initial segments
+    int send_segment = rank_;
+    int recv_segment = (rank_ + direction + size_) % size_;
+
+    // Plan the whole reduce in terms of sends and recvs as indices in data.
+    // It makes the actual async send and recv a bit simpler to follow when
+    // there are less offset calculations around.
+    std::vector<std::pair<size_t, size_t>> send_plan;
+    std::vector<std::pair<size_t, size_t>> recv_plan;
+
+    // Two times the same send/recv operations, first scatter reduce and then
+    // gather.
+    for (int k = 0; k < 2; k++) {
+      for (int i = 0; i < size_ - 1; i++) {
+        size_t send_start = send_segment * segment_size;
+        size_t send_stop =
+            std::min((send_segment + 1) * segment_size, data_size);
+        size_t recv_start = recv_segment * segment_size;
+        size_t recv_stop =
+            std::min((recv_segment + 1) * segment_size, data_size);
+
+        for (size_t j = 0; j < n_packets; j++) {
+          send_plan.emplace_back(
+              std::min(send_start + j * BUFFER_SIZE, send_stop),
+              std::min(send_start + (j + 1) * BUFFER_SIZE, send_stop));
+          recv_plan.emplace_back(
+              std::min(recv_start + j * BUFFER_SIZE, recv_stop),
+              std::min(recv_start + (j + 1) * BUFFER_SIZE, recv_stop));
+        }
+
+        send_segment = (send_segment + size_ + direction) % size_;
+        recv_segment = (recv_segment + size_ + direction) % size_;
+      }
+    }
+
+    // Running the plan is fairly simple, we keep a send and a recv in flight
+    // while doing the summation.
+    T* recv_buffers[ALL_SUM_BUFFERS];
+    for (int i = 0; i < ALL_SUM_BUFFERS; i++) {
+      recv_buffers[i] = buffer + i * BUFFER_SIZE;
+    }
+    std::future<void> sends[2], recvs[2];
+    int a = 0;
+    int b = (n_packets > 1) ? 1 : 0;
+    for (int i = 0, j = -b; i < send_plan.size(); j++, i++) {
+      sends[a] = comm_.send(
+          socket_send,
+          data + send_plan[i].first,
+          send_plan[i].second - send_plan[i].first);
+      if (2 * i < send_plan.size()) {
+        recvs[a] = comm_.recv(
+            socket_recv,
+            recv_buffers[i % ALL_SUM_BUFFERS],
+            recv_plan[i].second - recv_plan[i].first);
+      } else {
+        recvs[a] = comm_.recv(
+            socket_recv,
+            data + recv_plan[i].first,
+            recv_plan[i].second - recv_plan[i].first);
+      }
+
+      if (j >= 0) {
+        sends[b].wait();
+        recvs[b].wait();
+        if (2 * j < send_plan.size()) {
+          reduce_op(
+              recv_buffers[j % ALL_SUM_BUFFERS],
+              data + recv_plan[j].first,
+              recv_plan[j].second - recv_plan[j].first);
+        }
+      }
+
+      std::swap(a, b);
+    }
+    sends[b].wait();
+    recvs[b].wait();
+  }
+
+  void all_gather_impl(
+      const char* input,
+      char* output,
+      size_t input_size,
+      size_t data_size,
+      int socket_right,
+      int socket_left,
+      int direction) {
+    // Choose which socket we send to and recv from
+    int socket_send = (direction < 0) ? socket_right : socket_left;
+    int socket_recv = (direction < 0) ? socket_left : socket_right;
+
+    // Initial segments
+    int send_segment = rank_;
+    int recv_segment = (rank_ + direction + size_) % size_;
+
+    // Copy our own segment in the output
+    std::memcpy(output + rank_ * input_size, input, data_size);
+
+    // Simple send/recv all gather. Possible performance improvement by
+    // splitting to multiple chunks and allowing send/recv to run a bit ahead.
+    // See all_sum_impl for an example.
+    for (int i = 0; i < size_ - 1; i++) {
+      auto sent = comm_.send(
+          socket_send, output + send_segment * input_size, data_size);
+      auto recvd = comm_.recv(
+          socket_recv, output + recv_segment * input_size, data_size);
+
+      send_segment = (send_segment + size_ + direction) % size_;
+      recv_segment = (recv_segment + size_ + direction) % size_;
+
+      sent.wait();
+      recvd.wait();
+    }
+  }
+
+  void
+  send(const std::vector<int>& sockets, const char* data, size_t data_size) {
+    size_t segment_size =
+        std::max(size_t(1024), ceildiv(data_size, sockets.size()));
+    std::vector<std::future<void>> sends;
+    for (int i = 0; i < sockets.size(); i++) {
+      if (i * segment_size >= data_size) {
+        break;
+      }
+      sends.emplace_back(comm_.send(
+          sockets[i],
+          data + i * segment_size,
+          std::min(data_size, (i + 1) * segment_size) - i * segment_size));
+    }
+    for (auto& f : sends) {
+      f.wait();
+    }
+  }
+
+  void recv(const std::vector<int>& sockets, char* data, size_t data_size) {
+    size_t segment_size =
+        std::max(size_t(1024), ceildiv(data_size, sockets.size()));
+    std::vector<std::future<void>> recvs;
+    for (int i = 0; i < sockets.size(); i++) {
+      if (i * segment_size >= data_size) {
+        break;
+      }
+      recvs.emplace_back(comm_.recv(
+          sockets[i],
+          data + i * segment_size,
+          std::min(data_size, (i + 1) * segment_size) - i * segment_size));
+    }
+    for (auto& f : recvs) {
+      f.wait();
+    }
+  }
+
+  int rank_;
+  int size_;
+
+  bool verbose_;
+
+  ThreadPool pool_;
+  CommunicationThreads comm_;
+
+  std::vector<int> sockets_right_;
+  std::vector<int> sockets_left_;
+
+  std::vector<char> buffers_;
+};
+
+bool is_available() {
+  return true;
+}
+
+std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
+  const char* hostfile = std::getenv("MLX_HOSTFILE");
+  const char* rank_str = std::getenv("MLX_RANK");
+  const char* ring_verbose = std::getenv("MLX_RING_VERBOSE");
+
+  if (!hostfile || !rank_str) {
+    if (strict) {
+      std::ostringstream msg;
+      msg << "[ring] You need to provide via environment variables both a rank (MLX_RANK) "
+          << "and a hostfile (MLX_HOSTFILE) but provided MLX_RANK=\""
+          << ((rank_str) ? rank_str : "") << "\" and MLX_HOSTFILE=\""
+          << ((hostfile) ? hostfile : "") << "\"";
+      throw std::runtime_error(msg.str());
+    }
+    return nullptr;
+  }
+
+  auto nodes = load_nodes(hostfile);
+  int rank = std::atoi(rank_str);
+
+  return std::make_shared<RingGroup>(rank, nodes, ring_verbose != nullptr);
+}
+
+} // namespace mlx::core::distributed::ring
diff --git a/macos/third_party/mlx/mlx/distributed/ring/ring.h b/macos/third_party/mlx/mlx/distributed/ring/ring.h
new file mode 100644
index 0000000..e0b3fd0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/ring/ring.h
@@ -0,0 +1,12 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/distributed/distributed.h"
+
+namespace mlx::core::distributed::ring {
+
+using GroupImpl = mlx::core::distributed::detail::GroupImpl;
+
+bool is_available();
+std::shared_ptr<GroupImpl> init(bool strict = false);
+
+} // namespace mlx::core::distributed::ring
diff --git a/macos/third_party/mlx/mlx/distributed/utils.cpp b/macos/third_party/mlx/mlx/distributed/utils.cpp
new file mode 100644
index 0000000..2de994c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/utils.cpp
@@ -0,0 +1,206 @@
+// Copyright © 2025 Apple Inc.
+
+#include <netdb.h>
+#include <unistd.h>
+#include <cstring>
+#include <sstream>
+#include <thread>
+
+#include "mlx/distributed/utils.h"
+
+namespace mlx::core::distributed::detail {
+
+/**
+ * Parse a sockaddr from an ip and port provided as strings.
+ */
+address_t parse_address(const std::string& ip, const std::string& port) {
+  struct addrinfo hints, *res;
+  std::memset(&hints, 0, sizeof(hints));
+  hints.ai_family = AF_UNSPEC;
+  hints.ai_socktype = SOCK_STREAM;
+
+  int status = getaddrinfo(ip.c_str(), port.c_str(), &hints, &res);
+  if (status != 0) {
+    std::ostringstream msg;
+    msg << "Can't parse address " << ip << ":" << port;
+    throw std::runtime_error(msg.str());
+  }
+
+  address_t result;
+  memcpy(&result.addr, res->ai_addr, res->ai_addrlen);
+  result.len = res->ai_addrlen;
+  freeaddrinfo(res);
+
+  return result;
+}
+
+/**
+ * Parse a sockaddr provided as an <ip>:<port> string.
+ */
+address_t parse_address(const std::string& ip_port) {
+  auto colon = ip_port.find(":");
+  if (colon == std::string::npos) {
+    std::ostringstream msg;
+    msg << "Can't parse address " << ip_port;
+    throw std::runtime_error(msg.str());
+  }
+  std::string ip(ip_port.begin(), ip_port.begin() + colon);
+  std::string port(ip_port.begin() + colon + 1, ip_port.end());
+
+  return parse_address(ip, port);
+}
+
+TCPSocket::TCPSocket(const char* tag) {
+  sock_ = socket(AF_INET, SOCK_STREAM, 0);
+  if (sock_ < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't create socket (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+TCPSocket::TCPSocket(TCPSocket&& s) {
+  sock_ = s.sock_;
+  s.sock_ = -1;
+}
+
+TCPSocket& TCPSocket::operator=(TCPSocket&& s) {
+  if (this != &s) {
+    sock_ = s.sock_;
+    s.sock_ = -1;
+  }
+  return *this;
+}
+
+TCPSocket::TCPSocket(int s) : sock_(s) {}
+
+TCPSocket::~TCPSocket() {
+  if (sock_ > 0) {
+    shutdown(sock_, 2);
+    close(sock_);
+  }
+}
+
+int TCPSocket::detach() {
+  int s = sock_;
+  sock_ = -1;
+  return s;
+}
+
+void TCPSocket::listen(const char* tag, const address_t& addr) {
+  int success;
+
+  // Make sure we can launch immediately after shutdown by setting the
+  // reuseaddr option so that we don't get address already in use errors
+  int enable = 1;
+  success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int));
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+  success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int));
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't enable reuseport (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Bind the socket to the address and port
+  success = bind(sock_, addr.get(), addr.len);
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't bind socket (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  // Prepare waiting for connections
+  success = ::listen(sock_, 0);
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't listen (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+}
+
+TCPSocket TCPSocket::accept(const char* tag) {
+  int peer = ::accept(sock_, nullptr, nullptr);
+  if (peer < 0) {
+    std::ostringstream msg;
+    msg << tag << " Accept failed (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  return TCPSocket(peer);
+}
+
+void TCPSocket::send(const char* tag, const void* data, size_t len) {
+  while (len > 0) {
+    auto n = ::send(sock_, data, len, 0);
+    if (n <= 0) {
+      std::ostringstream msg;
+      msg << tag << " Send failed with errno=" << errno;
+      throw std::runtime_error(msg.str());
+    }
+    len -= n;
+    data = static_cast<const char*>(data) + n;
+  }
+}
+
+void TCPSocket::recv(const char* tag, void* data, size_t len) {
+  while (len > 0) {
+    auto n = ::recv(sock_, data, len, 0);
+    if (n <= 0) {
+      std::ostringstream msg;
+      msg << tag << " Recv failed with errno=" << errno;
+      throw std::runtime_error(msg.str());
+    }
+    len -= n;
+    data = static_cast<char*>(data) + n;
+  }
+}
+
+TCPSocket TCPSocket::connect(
+    const char* tag,
+    const address_t& addr,
+    int num_retries,
+    int wait,
+    std::function<void(int, int)> cb) {
+  int sock, success;
+
+  // Attempt to connect `num_retries` times with exponential backoff.
+  for (int attempt = 0; attempt < num_retries; attempt++) {
+    // Create the socket
+    sock = socket(AF_INET, SOCK_STREAM, 0);
+    if (sock < 0) {
+      std::ostringstream msg;
+      msg << tag << " Couldn't create socket to connect (error: " << errno
+          << ")";
+      throw std::runtime_error(msg.str());
+    }
+
+    success = ::connect(sock, addr.get(), addr.len);
+    if (success == 0) {
+      break;
+    }
+
+    if (cb != nullptr) {
+      cb(attempt, wait);
+    }
+    if (wait > 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(wait));
+    }
+
+    wait <<= 1;
+  }
+
+  if (success < 0) {
+    std::ostringstream msg;
+    msg << tag << " Couldn't connect (error: " << errno << ")";
+    throw std::runtime_error(msg.str());
+  }
+
+  return TCPSocket(sock);
+}
+
+} // namespace mlx::core::distributed::detail
diff --git a/macos/third_party/mlx/mlx/distributed/utils.h b/macos/third_party/mlx/mlx/distributed/utils.h
new file mode 100644
index 0000000..213dd59
--- /dev/null
+++ b/macos/third_party/mlx/mlx/distributed/utils.h
@@ -0,0 +1,67 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <sys/socket.h>
+#include <functional>
+#include <string>
+
+namespace mlx::core::distributed::detail {
+
+struct address_t {
+  sockaddr_storage addr;
+  socklen_t len;
+
+  const sockaddr* get() const {
+    return (struct sockaddr*)&addr;
+  }
+};
+
+/**
+ * Parse a sockaddr from an ip and port provided as strings.
+ */
+address_t parse_address(const std::string& ip, const std::string& port);
+
+/**
+ * Parse a sockaddr provided as an <ip>:<port> string.
+ */
+address_t parse_address(const std::string& ip_port);
+
+/**
+ * Small wrapper over a TCP socket to simplify initiating connections.
+ */
+class TCPSocket {
+ public:
+  TCPSocket(const char* tag);
+  TCPSocket(const TCPSocket&) = delete;
+  TCPSocket& operator=(const TCPSocket&) = delete;
+  TCPSocket(TCPSocket&& s);
+  TCPSocket& operator=(TCPSocket&&);
+  ~TCPSocket();
+
+  void listen(const char* tag, const address_t& addr);
+  TCPSocket accept(const char* tag);
+
+  void send(const char* tag, const void* data, size_t len);
+  void recv(const char* tag, void* data, size_t len);
+
+  int detach();
+
+  operator int() const {
+    return sock_;
+  }
+
+  static TCPSocket connect(
+      const char* tag,
+      const address_t& addr,
+      int num_retries = 1,
+      int wait = 0,
+      std::function<void(int, int)> cb = nullptr);
+
+ private:
+  TCPSocket(int sock);
+
+  int sock_;
+};
+
+} // namespace mlx::core::distributed::detail
diff --git a/macos/third_party/mlx/mlx/dtype.cpp b/macos/third_party/mlx/mlx/dtype.cpp
new file mode 100644
index 0000000..380a936
--- /dev/null
+++ b/macos/third_party/mlx/mlx/dtype.cpp
@@ -0,0 +1,197 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstdint>
+
+#include "mlx/dtype.h"
+
+namespace mlx::core {
+
+namespace {
+
+constexpr int num_types = 14;
+constexpr int num_cats = 8;
+
+constexpr Dtype::Kind type_kinds[num_types] = {
+    Dtype::Kind::b, // bool_,
+    Dtype::Kind::u, // uint8,
+    Dtype::Kind::u, // uint16,
+    Dtype::Kind::u, // uint32,
+    Dtype::Kind::u, // uint64,
+    Dtype::Kind::i, // int8,
+    Dtype::Kind::i, // int16,
+    Dtype::Kind::i, // int32,
+    Dtype::Kind::i, // int64,
+    Dtype::Kind::f, // float16,
+    Dtype::Kind::f, // float32,
+    Dtype::Kind::f, // float64,
+    Dtype::Kind::V, // bfloat16,
+    Dtype::Kind::c // complex64,
+};
+
+// Following Jax type promotion rules:
+// https://jax.readthedocs.io/en/latest/type_promotion.html
+// clang-format off
+constexpr Dtype type_rules[num_types][num_types] = {
+// bool       uint8      uint16     uint32     uint64     int8       int16      int32      int64      float16    float32   float64    bfloat16   complex64
+  {bool_,     uint8,     uint16,    uint32,    uint64,    int8,      int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // bool
+  {uint8,     uint8,     uint16,    uint32,    uint64,    int16,     int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // uint8
+  {uint16,    uint16,    uint16,    uint32,    uint64,    int32,     int32,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // uint16
+  {uint32,    uint32,    uint32,    uint32,    uint64,    int64,     int64,     int64,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // uint32
+  {uint64,    uint64,    uint64,    uint64,    uint64,    float32,   float32,   float32,   float32,   float16,   float32,  float64,   bfloat16,  complex64}, // uint64
+  {int8,      int16,     int32,     int64,     float32,   int8,      int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int8
+  {int16,     int16,     int32,     int64,     float32,   int16,     int16,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int16
+  {int32,     int32,     int32,     int64,     float32,   int32,     int32,     int32,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int32
+  {int64,     int64,     int64,     int64,     float32,   int64,     int64,     int64,     int64,     float16,   float32,  float64,   bfloat16,  complex64}, // int64
+  {float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float16,   float32,  float64,   float32,   complex64}, // float16
+  {float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,   float32,  float64,   float32,   complex64}, // float32
+  {float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,   float64,  float64,   float64,   complex64}, // float64
+  {bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  bfloat16,  float32,   float32,  float64,   bfloat16,  complex64}, // bfloat16
+  {complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64, complex64,complex64, complex64, complex64}, // complex64
+};
+
+
+constexpr bool subcategory_to_category[num_cats][num_cats] = {
+// complexfloating floating inexact signedinteger unsignedinteger integer number generic
+  {true,           false,   true,   false,        false,          false,  true,  true}, // complexfloating
+  {false,          true,    true,   false,        false,          false,  true,  true}, // floating
+  {false,          false,   true,   false,        false,          false,  true,  true}, // inexact
+  {false,          false,   false,  true,         false,          true,   true,  true}, // signedinteger
+  {false,          false,   false,  false,        true,           true,   true,  true}, // unsignedinteger
+  {false,          false,   false,  false,        false,          true,   true,  true}, // integer
+  {false,          false,   false,  false,        false,          false,  true,  true}, // number
+  {false,          false,   false,  false,        false,          false,  false, true}, // generic
+};
+
+constexpr Dtype::Category type_to_category[num_types] = {
+    Dtype::Category::generic, // bool_,
+    Dtype::Category::unsignedinteger, // uint8,
+    Dtype::Category::unsignedinteger, // uint16,
+    Dtype::Category::unsignedinteger, // uint32,
+    Dtype::Category::unsignedinteger, // uint64,
+    Dtype::Category::signedinteger, // int8,
+    Dtype::Category::signedinteger, // int16,
+    Dtype::Category::signedinteger, // int32,
+    Dtype::Category::signedinteger, // int64,
+    Dtype::Category::floating, // float16,
+    Dtype::Category::floating, // float32,
+    Dtype::Category::floating, // float64,
+    Dtype::Category::floating, // bfloat16,
+    Dtype::Category::complexfloating, // complex64,
+};
+
+// clang-format on
+
+} // namespace
+
+Dtype promote_types(const Dtype& t1, const Dtype& t2) {
+  return Dtype(
+      type_rules[static_cast<int>(t1.val())][static_cast<int>(t2.val())]);
+}
+
+Dtype::Kind kindof(const Dtype& t) {
+  return type_kinds[static_cast<int>(t.val())];
+}
+
+template class MLX_API TypeToDtype<bool>;
+template class MLX_API TypeToDtype<uint8_t>;
+template class MLX_API TypeToDtype<uint16_t>;
+template class MLX_API TypeToDtype<uint32_t>;
+template class MLX_API TypeToDtype<uint64_t>;
+template class MLX_API TypeToDtype<int8_t>;
+template class MLX_API TypeToDtype<int16_t>;
+template class MLX_API TypeToDtype<int32_t>;
+template class MLX_API TypeToDtype<int64_t>;
+template class MLX_API TypeToDtype<float16_t>;
+template class MLX_API TypeToDtype<float>;
+template class MLX_API TypeToDtype<double>;
+template class MLX_API TypeToDtype<bfloat16_t>;
+template class MLX_API TypeToDtype<complex64_t>;
+
+template <>
+TypeToDtype<bool>::operator Dtype() {
+  return bool_;
+}
+
+template <>
+TypeToDtype<uint8_t>::operator Dtype() {
+  return uint8;
+}
+
+template <>
+TypeToDtype<uint16_t>::operator Dtype() {
+  return uint16;
+}
+
+template <>
+TypeToDtype<uint32_t>::operator Dtype() {
+  return uint32;
+}
+
+template <>
+TypeToDtype<uint64_t>::operator Dtype() {
+  return uint64;
+}
+
+template <>
+TypeToDtype<int8_t>::operator Dtype() {
+  return int8;
+}
+
+template <>
+TypeToDtype<int16_t>::operator Dtype() {
+  return int16;
+}
+
+template <>
+TypeToDtype<int32_t>::operator Dtype() {
+  return int32;
+}
+
+template <>
+TypeToDtype<int64_t>::operator Dtype() {
+  return int64;
+}
+
+template <>
+TypeToDtype<float16_t>::operator Dtype() {
+  return float16;
+}
+
+template <>
+TypeToDtype<float>::operator Dtype() {
+  return float32;
+}
+
+template <>
+TypeToDtype<double>::operator Dtype() {
+  return float32;
+}
+
+template <>
+TypeToDtype<bfloat16_t>::operator Dtype() {
+  return bfloat16;
+}
+
+template <>
+TypeToDtype<complex64_t>::operator Dtype() {
+  return complex64;
+}
+
+bool issubdtype(const Dtype& a, const Dtype& b) {
+  return a == b;
+}
+
+bool issubdtype(const Dtype::Category& cat, const Dtype& type) {
+  return false;
+}
+
+bool issubdtype(const Dtype& type, const Dtype::Category& cat) {
+  return issubdtype(type_to_category[static_cast<uint32_t>(type.val())], cat);
+}
+
+bool issubdtype(const Dtype::Category& a, const Dtype::Category& b) {
+  return subcategory_to_category[static_cast<uint32_t>(a)]
+                                [static_cast<uint32_t>(b)];
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/dtype.h b/macos/third_party/mlx/mlx/dtype.h
new file mode 100644
index 0000000..744ca58
--- /dev/null
+++ b/macos/third_party/mlx/mlx/dtype.h
@@ -0,0 +1,116 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <complex>
+#include <cstdint>
+
+#include "mlx/api.h"
+#include "mlx/types/complex.h"
+#include "mlx/types/half_types.h"
+
+namespace mlx::core {
+
+struct Dtype {
+  enum class Val {
+    bool_,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    int8,
+    int16,
+    int32,
+    int64,
+    float16,
+    float32,
+    float64,
+    bfloat16,
+    complex64,
+  };
+
+  enum class Kind {
+    b, /* bool */
+    u, /* unsigned int */
+    i, /* signed int */
+    f, /* float */
+    c, /* complex */
+    V, /* void - used for brain float */
+  };
+
+  enum class Category {
+    complexfloating,
+    floating,
+    inexact,
+    signedinteger,
+    unsignedinteger,
+    integer,
+    number,
+    generic
+  };
+
+  constexpr explicit Dtype(Val val, uint8_t size) : val_(val), size_(size) {}
+
+  constexpr operator Val() const {
+    return val_;
+  }
+  constexpr Val val() const {
+    return val_;
+  }
+  constexpr uint8_t size() const {
+    return size_;
+  }
+
+ private:
+  Val val_;
+  uint8_t size_;
+};
+
+inline constexpr Dtype bool_{Dtype::Val::bool_, sizeof(bool)};
+
+inline constexpr Dtype uint8{Dtype::Val::uint8, sizeof(uint8_t)};
+inline constexpr Dtype uint16{Dtype::Val::uint16, sizeof(uint16_t)};
+inline constexpr Dtype uint32{Dtype::Val::uint32, sizeof(uint32_t)};
+inline constexpr Dtype uint64{Dtype::Val::uint64, sizeof(uint64_t)};
+
+inline constexpr Dtype int8{Dtype::Val::int8, sizeof(int8_t)};
+inline constexpr Dtype int16{Dtype::Val::int16, sizeof(int16_t)};
+inline constexpr Dtype int32{Dtype::Val::int32, sizeof(int32_t)};
+inline constexpr Dtype int64{Dtype::Val::int64, sizeof(int64_t)};
+
+inline constexpr Dtype float16{Dtype::Val::float16, sizeof(uint16_t)};
+inline constexpr Dtype float32{Dtype::Val::float32, sizeof(float)};
+inline constexpr Dtype float64{Dtype::Val::float64, sizeof(double)};
+inline constexpr Dtype bfloat16{Dtype::Val::bfloat16, sizeof(uint16_t)};
+inline constexpr Dtype complex64{Dtype::Val::complex64, sizeof(complex64_t)};
+
+inline constexpr Dtype::Category complexfloating =
+    Dtype::Category::complexfloating;
+inline constexpr Dtype::Category floating = Dtype::Category::floating;
+inline constexpr Dtype::Category inexact = Dtype::Category::inexact;
+inline constexpr Dtype::Category signedinteger = Dtype::Category::signedinteger;
+inline constexpr Dtype::Category unsignedinteger =
+    Dtype::Category::unsignedinteger;
+inline constexpr Dtype::Category integer = Dtype::Category::integer;
+inline constexpr Dtype::Category number = Dtype::Category::number;
+inline constexpr Dtype::Category generic = Dtype::Category::generic;
+
+MLX_API bool issubdtype(const Dtype& a, const Dtype& b);
+MLX_API bool issubdtype(const Dtype::Category& a, const Dtype& b);
+MLX_API bool issubdtype(const Dtype& a, const Dtype::Category& b);
+MLX_API bool issubdtype(const Dtype::Category& a, const Dtype::Category& b);
+
+MLX_API Dtype promote_types(const Dtype& t1, const Dtype& t2);
+
+inline uint8_t size_of(const Dtype& t) {
+  return t.size();
+}
+
+MLX_API Dtype::Kind kindof(const Dtype& t);
+
+template <typename T>
+struct MLX_API TypeToDtype {
+  operator Dtype();
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/dtype_utils.cpp b/macos/third_party/mlx/mlx/dtype_utils.cpp
new file mode 100644
index 0000000..9f10e6a
--- /dev/null
+++ b/macos/third_party/mlx/mlx/dtype_utils.cpp
@@ -0,0 +1,42 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core {
+
+const char* dtype_to_string(Dtype arg) {
+  switch (arg) {
+    case bool_:
+      return "bool";
+    case int8:
+      return "int8";
+    case int16:
+      return "int16";
+    case int32:
+      return "int32";
+    case int64:
+      return "int64";
+    case uint8:
+      return "uint8";
+    case uint16:
+      return "uint16";
+    case uint32:
+      return "uint32";
+    case uint64:
+      return "uint64";
+    case float16:
+      return "float16";
+    case bfloat16:
+      return "bfloat16";
+    case float32:
+      return "float32";
+    case float64:
+      return "float64";
+    case complex64:
+      return "complex64";
+    default:
+      return "unknown";
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/dtype_utils.h b/macos/third_party/mlx/mlx/dtype_utils.h
new file mode 100644
index 0000000..47c6ed6
--- /dev/null
+++ b/macos/third_party/mlx/mlx/dtype_utils.h
@@ -0,0 +1,119 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <sstream>
+
+#include "mlx/dtype.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+// Return string representation of dtype.
+const char* dtype_to_string(Dtype arg);
+
+#define MLX_INTERNAL_DTYPE_SWITCH_CASE(DTYPE, TYPE) \
+  case DTYPE:                                       \
+    f(type_identity<TYPE>{});                       \
+    break
+
+#define MLX_INTERNAL_DTYPE_SWITCH_INTS()            \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(int8, int8_t);     \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(int16, int16_t);   \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(int32, int32_t);   \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(int64, int64_t);   \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(uint8, uint8_t);   \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(uint16, uint16_t); \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(uint32, uint32_t); \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(uint64, uint64_t)
+
+#define MLX_INTERNAL_DTYPE_SWITCH_FLOATS()              \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(float16, float16_t);   \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(bfloat16, bfloat16_t); \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(float32, float);       \
+  MLX_INTERNAL_DTYPE_SWITCH_CASE(float64, double)
+
+// This already exists in C++20 but in C++20 we can also just use templated
+// lambdas which will make this so much nicer.
+template <typename T>
+struct type_identity {
+  using type = T;
+};
+
+#define MLX_GET_TYPE(x) typename decltype(x)::type
+#define MLX_GET_VALUE(x) decltype(x)::value
+
+template <typename F>
+void dispatch_all_types(Dtype dt, F&& f) {
+  switch (dt) {
+    MLX_INTERNAL_DTYPE_SWITCH_CASE(bool_, bool);
+    MLX_INTERNAL_DTYPE_SWITCH_INTS();
+    MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
+    MLX_INTERNAL_DTYPE_SWITCH_CASE(complex64, complex64_t);
+  }
+}
+
+template <typename F>
+void dispatch_int_types(Dtype dt, std::string_view tag, F&& f) {
+  switch (dt) {
+    MLX_INTERNAL_DTYPE_SWITCH_INTS();
+    default:
+      std::ostringstream msg;
+      msg << tag << " Only integer types supported but " << dt
+          << " was provided";
+      throw std::invalid_argument(msg.str());
+  }
+}
+
+template <typename F>
+void dispatch_float_types(Dtype dt, std::string_view tag, F&& f) {
+  switch (dt) {
+    MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
+    default:
+      std::ostringstream msg;
+      msg << tag << " Only float types supported but " << dt << " was provided";
+      throw std::invalid_argument(msg.str());
+  }
+}
+
+template <typename F>
+void dispatch_inexact_types(Dtype dt, std::string_view tag, F&& f) {
+  switch (dt) {
+    MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
+    MLX_INTERNAL_DTYPE_SWITCH_CASE(complex64, complex64_t);
+    default:
+      std::ostringstream msg;
+      msg << tag << " Only inexact (float/complex) types supported but " << dt
+          << " was provided";
+      throw std::invalid_argument(msg.str());
+  }
+}
+
+template <typename F>
+void dispatch_int_float_types(Dtype dt, std::string_view tag, F&& f) {
+  switch (dt) {
+    MLX_INTERNAL_DTYPE_SWITCH_INTS();
+    MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
+    default:
+      std::ostringstream msg;
+      msg << tag << " Only integer and float types supported but " << dt
+          << " was provided";
+      throw std::invalid_argument(msg.str());
+  }
+}
+
+template <typename F>
+void dispatch_real_types(Dtype dt, std::string_view tag, F&& f) {
+  switch (dt) {
+    MLX_INTERNAL_DTYPE_SWITCH_CASE(bool_, bool);
+    MLX_INTERNAL_DTYPE_SWITCH_INTS();
+    MLX_INTERNAL_DTYPE_SWITCH_FLOATS();
+    default:
+      std::ostringstream msg;
+      msg << tag << " Only real numbers supported but " << dt
+          << " was provided";
+      throw std::invalid_argument(msg.str());
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/einsum.cpp b/macos/third_party/mlx/mlx/einsum.cpp
new file mode 100644
index 0000000..6290887
--- /dev/null
+++ b/macos/third_party/mlx/mlx/einsum.cpp
@@ -0,0 +1,941 @@
+// Copyright © 2024 Apple Inc.
+#include <numeric>
+#include <sstream>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/einsum.h"
+#include "mlx/ops.h"
+
+namespace mlx::core {
+
+namespace {
+
+// The MLX einsum implementation is based on NumPy (which is based on
+// opt_einsum):
+// https://github.com/numpy/numpy/blob/1d49c7f7ff527c696fc26ab2278ad51632a66660/numpy/_core/einsumfunc.py#L743
+// https://github.com/dgasmith/opt_einsum
+
+using CharSet = std::unordered_set<char>;
+
+// A helper struct to hold the string and set
+// representation of a subscript to avoid needing
+// to recompute the set
+struct Subscript {
+  Subscript(std::string str, CharSet set)
+      : str(std::move(str)), set(std::move(set)) {};
+  std::string str;
+  CharSet set;
+};
+
+struct PathInfo {
+  size_t naive_cost;
+  size_t naive_scaling;
+  size_t optimized_cost;
+  size_t optimized_scaling;
+  size_t largest_term;
+};
+
+struct PathNode {
+  PathNode(
+      std::vector<Subscript> inputs,
+      Subscript output,
+      std::vector<int> positions)
+      : inputs(std::move(inputs)),
+        output(std::move(output)),
+        positions(std::move(positions)) {};
+
+  std::vector<Subscript> inputs;
+  Subscript output;
+
+  std::vector<int> positions;
+};
+
+// Parse the comma separated subscripts into a vector of strings. If the
+// output subscripts are missing they are inferred.
+//
+// For example:
+//  "ij,jk -> ik" becomes {{"ij", "jk"}, "ik"}
+//  "ij,jk" becomes {{"ij", "jk"}, "ik"}
+std::pair<std::vector<std::string>, std::string> parse(std::string subscripts) {
+  std::string lhs, rhs;
+
+  // Start by removing all white space
+  subscripts.erase(
+      std::remove(subscripts.begin(), subscripts.end(), ' '), subscripts.end());
+
+  if (auto pos = subscripts.find("->"); pos != std::string::npos) {
+    // Explicit mode
+    lhs = subscripts.substr(0, pos);
+    rhs = subscripts.substr(pos + 2);
+  } else {
+    // Implicit mode:
+    // - repeats are summed
+    // - ellipses are placed in the beginning of the output
+    // - remaining output axes are ordered alphabetically
+    lhs = subscripts;
+    std::unordered_map<char, int> temp;
+    for (auto& c : subscripts) {
+      if (c == ',') {
+        continue;
+      }
+      if (c == '.' && rhs.empty()) {
+        rhs += "...";
+        continue;
+      }
+
+      auto inserted = temp.insert({c, 0});
+      inserted.first->second++;
+    }
+    for (auto& k : temp) {
+      if (k.second == 1) {
+        rhs += k.first;
+      }
+    }
+    std::sort(rhs.begin(), rhs.end());
+  }
+  std::vector<std::string> input_list;
+  std::stringstream ss(lhs);
+  std::string token;
+  while (getline(ss, token, ',')) {
+    input_list.push_back(token);
+  }
+  return {input_list, rhs};
+}
+
+// Check if two sets are disjoint
+bool disjoint(const CharSet& x, const CharSet& y) {
+  for (auto& c : x) {
+    if (y.find(c) != y.end()) {
+      return false;
+    }
+  }
+  return true;
+}
+
+template <typename T>
+size_t term_size(const T& term, std::unordered_map<char, ShapeElem> dict) {
+  size_t size = 1;
+  for (auto c : term) {
+    size *= dict[c];
+  }
+  return size;
+}
+
+size_t flop_count(
+    const CharSet& term,
+    bool inner,
+    int num_terms,
+    std::unordered_map<char, ShapeElem> dict) {
+  size_t size = term_size(term, dict);
+  auto op_factor = 1;
+  if ((num_terms - 1) > op_factor) {
+    op_factor = num_terms - 1;
+  }
+  if (inner) {
+    op_factor += 1;
+  }
+  return size * op_factor;
+}
+
+std::pair<size_t, int> compute_cost_and_scaling(
+    const std::vector<Subscript>& inputs,
+    const Subscript& output,
+    std::unordered_map<char, ShapeElem> dim_map) {
+  CharSet contractions;
+  for (auto& in : inputs) {
+    contractions.insert(in.set.begin(), in.set.end());
+  }
+
+  bool inner = false;
+  for (auto c : contractions) {
+    if (output.set.find(c) == output.set.end()) {
+      inner = true;
+      break;
+    }
+  }
+  auto cost = flop_count(contractions, inner, inputs.size(), dim_map);
+  return {cost, contractions.size()};
+}
+
+std::tuple<std::vector<PathNode>, size_t, int> greedy_path(
+    std::vector<Subscript> inputs,
+    const Subscript& output,
+    std::unordered_map<char, ShapeElem> dim_map,
+    size_t cost_limit,
+    size_t memory_limit) {
+  // Helper struct for building the greedy path
+  struct Contraction {
+    Contraction(
+        size_t size,
+        size_t cost,
+        CharSet output,
+        int dims,
+        int x,
+        int y)
+        : size(size),
+          cost(cost),
+          output(std::move(output)),
+          dims(dims),
+          x(x),
+          y(y) {};
+
+    int64_t size; // Size difference, can be negative
+    size_t cost;
+    CharSet output;
+    int dims; // Number of dimensions in the contraction
+    int x;
+    int y;
+  };
+
+  // Start by iterating over all possible combinations
+  std::vector<std::pair<int, int>> pos_pairs;
+  for (int i = 0; i < inputs.size(); ++i) {
+    for (int j = i + 1; j < inputs.size(); ++j) {
+      pos_pairs.emplace_back(i, j);
+    }
+  }
+
+  std::vector<PathNode> path;
+  std::vector<Contraction> possible_contractions;
+  size_t path_cost = 0;
+  int path_scaling = 0;
+  auto num_in = inputs.size();
+  for (int i = 0; i < num_in - 1; ++i) {
+    auto add_contraction = [&](int p1, int p2) {
+      CharSet new_term;
+      CharSet contractions(inputs[p1].set.begin(), inputs[p1].set.end());
+      contractions.insert(inputs[p2].set.begin(), inputs[p2].set.end());
+      for (int i = 0; i < inputs.size(); i++) {
+        if (i == p1 || i == p2) {
+          continue;
+        }
+        auto& in = inputs[i].set;
+        for (auto c : in) {
+          if (contractions.find(c) != contractions.end()) {
+            new_term.insert(c);
+          }
+        }
+      }
+      for (auto c : output.set) {
+        if (contractions.find(c) != contractions.end()) {
+          new_term.insert(c);
+        }
+      }
+
+      // Ignore if:
+      // - The size of the new result is greater than the memory limit
+      // - The cost is larger than the naive cost
+      auto new_size = term_size(new_term, dim_map);
+      if (new_size > memory_limit) {
+        return;
+      }
+      int64_t removed_size = term_size(inputs[p1].set, dim_map) +
+          term_size(inputs[p2].set, dim_map) - new_size;
+
+      bool inner = contractions.size() > new_term.size();
+      auto cost = flop_count(contractions, inner, 2, dim_map);
+      if (path_cost + cost > cost_limit) {
+        return;
+      }
+      possible_contractions.emplace_back(
+          removed_size, cost, std::move(new_term), contractions.size(), p1, p2);
+    };
+
+    for (auto& [p1, p2] : pos_pairs) {
+      // Ignore outer products
+      if (!disjoint(inputs[p1].set, inputs[p2].set)) {
+        add_contraction(p1, p2);
+      }
+    }
+
+    // If there's nothing in the contraction list,
+    // go over the pairs again without ignoring outer products
+    if (possible_contractions.empty()) {
+      for (auto& [p1, p2] : pos_pairs) {
+        add_contraction(p1, p2);
+      }
+    }
+
+    if (possible_contractions.empty()) {
+      // Default to naive einsum for the remaining inputs
+      std::vector<int> positions(inputs.size());
+      std::iota(positions.begin(), positions.end(), 0);
+      auto [cost, scale] = compute_cost_and_scaling(inputs, output, dim_map);
+      path.emplace_back(std::move(inputs), output, std::move(positions));
+
+      path_cost += cost;
+      path_scaling = std::max(scale, path_scaling);
+      break;
+    }
+
+    // Find the best contraction
+    auto& best = *std::min_element(
+        possible_contractions.begin(),
+        possible_contractions.end(),
+        [](const auto& x, const auto& y) {
+          return x.size > y.size || (x.size == y.size && x.cost < y.cost);
+        });
+    path_scaling = std::max(best.dims, path_scaling);
+
+    // Construct the output subscripts
+    std::string out_str(best.output.begin(), best.output.end());
+    // TODO, sorting by dimension size seems suboptimal?
+    std::sort(out_str.begin(), out_str.end(), [&dim_map](auto x, auto y) {
+      return dim_map[x] < dim_map[y];
+    });
+    Subscript new_output(std::move(out_str), std::move(best.output));
+
+    // Add the chosen contraction to the path
+    {
+      std::vector<Subscript> in_terms;
+      in_terms.push_back(std::move(inputs[best.x]));
+      in_terms.push_back(std::move(inputs[best.y]));
+      path.emplace_back(
+          std::move(in_terms), new_output, std::vector<int>{best.x, best.y});
+    }
+    // Remove used terms
+    inputs.erase(inputs.begin() + best.y);
+    inputs.erase(inputs.begin() + best.x);
+
+    // Add the new result
+    inputs.push_back(std::move(new_output));
+
+    // Update the existing contractions based on the selected one
+    std::vector<Contraction> updated_contractions;
+    for (auto& contraction : possible_contractions) {
+      // Drop contractions which contain either selected term
+      if (contraction.x == best.x || contraction.x == best.y ||
+          contraction.y == best.x || contraction.y == best.y) {
+        continue;
+      }
+
+      // Update the positions of other contractions
+      int x =
+          contraction.x - (contraction.x > best.x) - (contraction.x > best.y);
+      int y =
+          contraction.y - (contraction.y > best.x) - (contraction.y > best.y);
+      contraction.x = x;
+      contraction.y = y;
+      updated_contractions.push_back(std::move(contraction));
+    }
+
+    pos_pairs.clear();
+    for (int i = 0; i < inputs.size() - 1; ++i) {
+      pos_pairs.emplace_back(i, inputs.size() - 1);
+    }
+    path_cost += best.cost;
+
+    possible_contractions = std::move(updated_contractions);
+  }
+  return {path, path_cost, path_scaling};
+}
+
+// Assumes inputs have already have had repeats and single axis sums collapsed
+bool can_dot(const std::vector<Subscript>& inputs, const Subscript& output) {
+  if (inputs.size() != 2) {
+    return false;
+  }
+
+  for (auto c : inputs[0].set) {
+    // Use batched tensordot if anything is being contracted
+    if (output.set.find(c) == output.set.end()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+array batch_tensordot(
+    array a,
+    array b,
+    std::vector<int> a_contract,
+    std::vector<int> a_batch,
+    std::vector<int> a_concat,
+    std::vector<int> b_contract,
+    std::vector<int> b_batch,
+    std::vector<int> b_concat,
+    StreamOrDevice s) {
+  // Broadcast contracting dimensions
+  {
+    auto a_shape = a.shape();
+    auto b_shape = b.shape();
+    for (int i = 0; i < a_contract.size(); ++i) {
+      auto d = std::max(a.shape(a_contract[i]), b.shape(b_contract[i]));
+      a_shape[a_contract[i]] = d;
+      b_shape[b_contract[i]] = d;
+    }
+    a = broadcast_to(a, a_shape, s);
+    b = broadcast_to(b, b_shape, s);
+  }
+  auto transpose_reshape = [&s](
+                               const array& x,
+                               const std::vector<int>& i,
+                               const std::vector<int>& j,
+                               const std::vector<int>& k) {
+    std::vector<int> reorder(i.begin(), i.end());
+    reorder.insert(reorder.end(), j.begin(), j.end());
+    reorder.insert(reorder.end(), k.begin(), k.end());
+
+    int size1 = 1;
+    for (auto s : j) {
+      size1 *= x.shape(s);
+    }
+
+    int size2 = 1;
+    for (auto s : k) {
+      size2 *= x.shape(s);
+    }
+
+    Shape shape;
+    for (auto ax : i) {
+      shape.push_back(x.shape(ax));
+    }
+    shape.push_back(size1);
+    shape.push_back(size2);
+
+    return reshape(transpose(x, reorder, s), std::move(shape), s);
+  };
+
+  Shape out_shape;
+  for (auto ax : a_batch) {
+    out_shape.push_back(a.shape(ax));
+  }
+  for (auto ax : a_concat) {
+    out_shape.push_back(a.shape(ax));
+  }
+  for (auto ax : b_concat) {
+    out_shape.push_back(b.shape(ax));
+  }
+
+  a = transpose_reshape(a, a_batch, a_concat, a_contract);
+  b = transpose_reshape(b, b_batch, b_contract, b_concat);
+
+  return reshape(matmul(a, b, s), std::move(out_shape), s);
+}
+
+// Collapse repeated subscripts and return the resulting array. The subscript
+// is also updated in place. For example:
+// - Given an input with shape (4, 4) and subscript "ii", returns
+//   the diagonal of shape (4,) and updates the subscript to "i".
+// - Given an input with shape (4, 2, 4, 2) and subscript "ijij",
+//   returns an output with shape (4, 2) and updates the subscript
+//   to "ij".
+array collapse_repeats(array in, Subscript& subscript, StreamOrDevice s) {
+  // Build a list of (repeat chars, num repeats)
+  auto& str = subscript.str;
+  std::vector<std::pair<char, int>> repeats;
+  std::string new_str;
+  {
+    std::string repeat_str;
+    std::string no_repeat_str;
+    std::unordered_map<char, int> counts;
+    for (int i = 0; i < str.size(); ++i) {
+      auto [it, _] = counts.insert({str[i], 0});
+      it->second++;
+    }
+
+    for (auto& v : counts) {
+      if (v.second > 1) {
+        repeats.emplace_back(v.first, v.second);
+        repeat_str += v.first;
+      }
+    }
+    for (auto& c : str) {
+      if (counts[c] == 1) {
+        no_repeat_str += c;
+      }
+    }
+    new_str = repeat_str + no_repeat_str;
+  }
+
+  // Build the inputs for gather
+  auto slice_sizes = in.shape();
+  std::vector<int> axes;
+  std::vector<array> indices;
+  int n_expand = repeats.size();
+  for (auto [c, v] : repeats) {
+    for (int i = 0; i < str.size(); ++i) {
+      if (str[i] == c) {
+        slice_sizes[i] = 1;
+        axes.push_back(i);
+      }
+    }
+    Shape idx_shape(n_expand--, 1);
+    idx_shape[0] = in.shape(axes.back());
+    auto idx = reshape(
+        arange(static_cast<ShapeElem>(in.shape(axes.back())), s), idx_shape, s);
+    for (int i = 0; i < v; ++i) {
+      indices.push_back(idx);
+    }
+  }
+
+  in = gather(in, indices, axes, slice_sizes, s);
+
+  // Update subscript string with removed dups
+  str = new_str;
+
+  // Squeeze singleton dimensions left over from the gather
+  for (auto& ax : axes) {
+    ax += indices[0].ndim();
+  }
+
+  return squeeze(in, axes, s);
+}
+
+// Collapse repeat indices and sum single dimensions.
+// For example:
+// - "aa" becomes "a"
+// - "ij,jk->k" becoms "j,jk->k"
+void preprocess_einsum_inputs(
+    std::vector<Subscript>& inputs,
+    const Subscript& output,
+    const std::vector<int>& positions,
+    std::vector<array>& operands,
+    StreamOrDevice s) {
+  // Collapse repeat indices
+  for (int i = 0; i < inputs.size(); ++i) {
+    auto& in = inputs[i];
+    if (in.set.size() < in.str.size()) {
+      operands[positions[i]] = collapse_repeats(operands[positions[i]], in, s);
+    }
+  }
+
+  // Sum indices that are only in a single input
+  {
+    std::unordered_map<char, int> counts;
+    for (auto& in : inputs) {
+      for (auto c : in.set) {
+        auto inserted = counts.insert({c, 0});
+        inserted.first->second++;
+      }
+    }
+    for (auto c : output.set) {
+      auto inserted = counts.insert({c, 0});
+      inserted.first->second++;
+    }
+    for (int i = 0; i < inputs.size(); ++i) {
+      auto& in = inputs[i];
+      std::vector<int> sum_axes;
+      for (int ax = 0; ax < in.str.size(); ++ax) {
+        if (counts[in.str[ax]] == 1) {
+          sum_axes.push_back(ax);
+        }
+      }
+      if (!sum_axes.empty()) {
+        operands[positions[i]] =
+            sum(operands[positions[i]], sum_axes, false, s);
+      }
+      for (auto it = sum_axes.rbegin(); it != sum_axes.rend(); ++it) {
+        in.set.erase(in.str[*it]);
+        in.str.erase(in.str.begin() + *it);
+      }
+    }
+  }
+}
+
+array einsum_naive(
+    std::vector<Subscript> inputs,
+    const Subscript& output,
+    const std::vector<int>& positions,
+    std::vector<array> operands,
+    StreamOrDevice s) {
+  // Map each character to an axis
+  std::unordered_map<char, int> char_to_ax;
+  for (auto& in : inputs) {
+    for (auto c : in.str) {
+      char_to_ax.insert({c, char_to_ax.size()});
+    }
+  }
+
+  // Expand and transpose inputs as needed
+  for (int i = 0; i < inputs.size(); ++i) {
+    int pos = positions[i];
+    auto& op = operands[pos];
+
+    // Add missing dimensions at the end
+    if (op.ndim() != char_to_ax.size()) {
+      auto shape = op.shape();
+      shape.insert(shape.end(), char_to_ax.size() - shape.size(), 1);
+      op = reshape(op, std::move(shape), s);
+    }
+
+    // Transpose:
+    // - Build a vector of (char, ax) pairs for the current input
+    // - Sort the vector by the canonical axis in char_to_ax
+    // - Extract the sorted axis to get transpose order
+    std::vector<std::pair<char, int>> str_ax;
+    for (auto c : inputs[i].str) {
+      str_ax.emplace_back(c, str_ax.size());
+    }
+    for (auto [c, ax] : char_to_ax) {
+      if (inputs[i].set.find(c) == inputs[i].set.end()) {
+        str_ax.emplace_back(c, str_ax.size());
+      }
+    }
+    std::sort(
+        str_ax.begin(),
+        str_ax.end(),
+        [&char_to_ax](const auto& x, const auto& y) {
+          return char_to_ax[x.first] < char_to_ax[y.first];
+        });
+
+    // Skip the transpose if not needed
+    if (std::is_sorted(
+            str_ax.begin(), str_ax.end(), [](const auto& x, const auto& y) {
+              return x.second < y.second;
+            })) {
+      continue;
+    }
+
+    std::vector<int> reorder;
+    for (auto [c, ax] : str_ax) {
+      reorder.push_back(ax);
+    }
+    op = transpose(op, reorder, s);
+  }
+
+  // Multiply and sum
+  auto out = operands[positions[0]];
+  for (int i = 1; i < positions.size(); ++i) {
+    out = multiply(out, operands[positions[i]], s);
+  }
+  std::vector<int> sum_axes;
+  for (auto [c, ax] : char_to_ax) {
+    if (output.set.find(c) == output.set.end()) {
+      sum_axes.push_back(ax);
+    }
+  }
+  if (!sum_axes.empty()) {
+    out = sum(out, sum_axes, false, s);
+  }
+
+  // Transpose output if needed
+  std::vector<int> reorder;
+  for (auto c : output.str) {
+    reorder.push_back(char_to_ax[c]);
+  }
+  for (auto& r : reorder) {
+    int offset = 0;
+    for (auto s : sum_axes) {
+      if (r > s) {
+        offset++;
+      }
+    }
+    r -= offset;
+  }
+  return transpose(out, reorder, s);
+}
+
+std::pair<std::vector<PathNode>, PathInfo> einsum_path_helper(
+    const std::string& subscripts,
+    const std::vector<array>& operands,
+    const std::string& fn_name) {
+  if (operands.size() == 0) {
+    std::ostringstream msg;
+    msg << "[" << fn_name << "] At least one operand is required.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto [in_subscripts, out_subscript] = parse(subscripts);
+
+  if (operands.size() != in_subscripts.size()) {
+    std::ostringstream msg;
+    msg << "[" << fn_name << "] Number of operands, " << operands.size()
+        << ", does not match number of input subscripts, "
+        << in_subscripts.size();
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Expand ellipses
+  // 1. Collect all the characters we can use for the missing axes.
+  // 2. Go over each subscript and check if all the characters are either
+  //    alphanumeric or an ellipsis.
+  // 3. Expand the ellipsis with as many characters from the unused ones as
+  //    necessary. We use the last N characters effectively prepending with
+  //    singleton dims for inputs with fewer dimensions.
+  // 4. For the output use the maximum size of ellipsis that we encountered in
+  //    the input.
+  CharSet used_chars(subscripts.begin(), subscripts.end());
+  std::string remaining_chars;
+  remaining_chars.reserve(52 - used_chars.size());
+  for (char c = 'a'; c <= 'z'; c++) {
+    if (used_chars.find(c) == used_chars.end()) {
+      remaining_chars += c;
+    }
+  }
+  for (char c = 'A'; c <= 'Z'; c++) {
+    if (used_chars.find(c) == used_chars.end()) {
+      remaining_chars += c;
+    }
+  }
+  int max_ellipsis_length = 0;
+  auto check_letters_and_expand_ellipsis = [&](auto& subscript,
+                                               const array* operand,
+                                               int operand_idx) {
+    bool have_ellipsis = false;
+    int cnt_before = 0, cnt_after = 0;
+    for (int i = 0; i < subscript.size(); i++) {
+      if (!isalpha(subscript[i])) {
+        if (i + 2 >= subscript.size() || subscript[i] != '.' ||
+            subscript[i + 1] != '.' || subscript[i + 2] != '.') {
+          std::ostringstream msg;
+          msg << "[" << fn_name << "] Subscripts must be letters, but got '"
+              << subscript[i] << "'.";
+          throw std::invalid_argument(msg.str());
+        }
+
+        if (have_ellipsis) {
+          std::ostringstream msg;
+          msg << "[" << fn_name
+              << "] Only one ellipsis per subscript is allowed but found more in '"
+              << subscript << "'.";
+          throw std::invalid_argument(msg.str());
+        }
+
+        have_ellipsis = true;
+        i += 2;
+        continue;
+      }
+
+      if (have_ellipsis) {
+        cnt_after++;
+      } else {
+        cnt_before++;
+      }
+    }
+
+    if (have_ellipsis) {
+      int ellipsis_length;
+      if (operand != nullptr) {
+        ellipsis_length = operand->ndim() - cnt_before - cnt_after;
+        if (ellipsis_length < 0) {
+          std::ostringstream msg;
+          msg << "[" << fn_name << "] Operand " << operand_idx << " with shape "
+              << operand->shape()
+              << " has insufficient dimensions for subscript '" << subscript
+              << "'. The ellipsis requires at least "
+              << (cnt_before + cnt_after) << " dimensions but the operand has "
+              << operand->ndim() << " dimensions.";
+          throw std::invalid_argument(msg.str());
+        }
+        max_ellipsis_length = std::max(ellipsis_length, max_ellipsis_length);
+      } else {
+        ellipsis_length = max_ellipsis_length;
+      }
+
+      subscript.replace(
+          subscript.begin() + cnt_before,
+          subscript.begin() + cnt_before + 3,
+          remaining_chars.end() - ellipsis_length,
+          remaining_chars.end());
+    }
+  };
+
+  for (int i = 0; i < operands.size(); i++) {
+    check_letters_and_expand_ellipsis(in_subscripts[i], &operands[i], i);
+  }
+  check_letters_and_expand_ellipsis(out_subscript, nullptr, -1);
+
+  CharSet out_set(out_subscript.begin(), out_subscript.end());
+  if (out_set.size() != out_subscript.size()) {
+    std::ostringstream msg;
+    msg << "[" << fn_name << "] Repeat indices not allowed in output.";
+    throw std::invalid_argument(msg.str());
+  }
+  Subscript output(out_subscript, std::move(out_set));
+
+  std::unordered_map<char, ShapeElem> dim_map;
+  std::vector<Subscript> inputs;
+  for (int i = 0; i < in_subscripts.size(); ++i) {
+    auto& in = in_subscripts[i];
+    CharSet in_set(in.begin(), in.end());
+    inputs.emplace_back(in, in_set);
+
+    if (in.size() != operands[i].ndim()) {
+      std::ostringstream msg;
+      msg << "[" << fn_name << "] Invalid number of subscripts " << in.size()
+          << " for input " << i << " with " << operands[i].ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+
+    // Check repeat subscripts are valid
+    if (in_set.size() < in.size()) {
+      std::unordered_map<char, ShapeElem> local_dims;
+      for (int j = 0; j < in.size(); ++j) {
+        auto dim = operands[i].shape(j);
+        auto inserted = local_dims.insert({in[j], dim});
+        if (!inserted.second) {
+          if (inserted.first->second != dim) {
+            std::ostringstream msg;
+            msg << "[" << fn_name << "] Dimensions of repeated subscripts "
+                << "do not have the same size (" << inserted.first->second
+                << " != " << dim << ").";
+            throw std::invalid_argument(msg.str());
+          }
+        }
+      }
+    }
+
+    for (int j = 0; j < in.size(); j++) {
+      auto c = in[j];
+      auto dim = operands[i].shape(j);
+      auto inserted = dim_map.insert({c, dim});
+      auto& in_dim = inserted.first->second;
+      if (dim != 1 && in_dim != 1 && in_dim != dim) {
+        std::ostringstream msg;
+        msg << "[" << fn_name << "] Cannot broadcast dimension " << j
+            << " of input " << i << " with shape " << operands[i].shape()
+            << " to size " << in_dim << ".";
+        throw std::invalid_argument(msg.str());
+      }
+      // Ensure the broadcasted size is used
+      in_dim = std::max(in_dim, dim);
+    }
+  }
+
+  size_t max_size = term_size(out_subscript, dim_map);
+  for (auto& in : in_subscripts) {
+    max_size = std::max(max_size, term_size(in, dim_map));
+  }
+
+  PathInfo path_info;
+
+  // Get the full naive cost
+  std::tie(path_info.naive_cost, path_info.naive_scaling) =
+      compute_cost_and_scaling(inputs, output, dim_map);
+
+  // Calculate the path
+  std::vector<PathNode> path;
+  if (inputs.size() <= 2) {
+    std::vector<int> positions(in_subscripts.size());
+    std::iota(positions.begin(), positions.end(), 0);
+    path.emplace_back(
+        std::move(inputs), std::move(output), std::move(positions));
+  } else {
+    std::tie(path, path_info.optimized_cost, path_info.optimized_scaling) =
+        greedy_path(inputs, output, dim_map, path_info.naive_cost, max_size);
+    // Set the final output subscript to the actual output
+    path.back().output = std::move(output);
+  }
+  return {path, path_info};
+}
+
+} // namespace
+
+std::pair<std::vector<std::vector<int>>, std::string> einsum_path(
+    const std::string& subscripts,
+    const std::vector<array>& operands) {
+  auto [path, path_info] =
+      einsum_path_helper(subscripts, operands, "einsum_path");
+
+  std::vector<std::vector<int>> pos_path;
+  for (auto& p : path) {
+    pos_path.push_back(p.positions);
+  }
+
+  std::ostringstream path_print;
+  path_print << "  Complete contraction:  " << subscripts << "\n"
+             << "         Naive scaling:  " << path_info.naive_scaling << "\n"
+             << "     Optimized scaling:  " << path_info.optimized_scaling
+             << "\n"
+             << "      Naive FLOP count:  " << path_info.naive_cost << "\n"
+             << "  Optimized FLOP count:  " << path_info.optimized_cost << "\n";
+  // TODO add more info here
+  return {pos_path, path_print.str()};
+}
+
+array einsum(
+    const std::string& subscripts,
+    const std::vector<array>& operands,
+    StreamOrDevice s /* = {} */) {
+  auto [path, path_info] = einsum_path_helper(subscripts, operands, "einsum");
+  auto inputs = operands;
+  for (auto& node : path) {
+    preprocess_einsum_inputs(
+        node.inputs, node.output, node.positions, inputs, s);
+
+    if (can_dot(node.inputs, node.output)) {
+      auto& in_a = node.inputs[0];
+      auto& in_b = node.inputs[1];
+      auto& out = node.output;
+
+      std::vector<int> a_contract;
+      std::vector<int> a_batch;
+      std::vector<int> a_concat;
+      for (int i = 0; i < in_a.str.size(); ++i) {
+        auto c = in_a.str[i];
+        if (out.set.find(c) == out.set.end()) {
+          // Not in the output, contraction
+          a_contract.push_back(i);
+        } else if (in_b.set.find(c) != in_b.set.end()) {
+          // Not a contraction but in both inputs, batch dim
+          a_batch.push_back(i);
+        } else {
+          // Not a batch dim or contract dim, so concat dim
+          a_concat.push_back(i);
+        }
+      }
+
+      std::vector<int> b_contract;
+      std::vector<int> b_batch;
+      std::vector<int> b_concat;
+      for (auto a_i : a_contract) {
+        b_contract.push_back(in_b.str.find(in_a.str[a_i]));
+      }
+      for (auto a_i : a_batch) {
+        b_batch.push_back(in_b.str.find(in_a.str[a_i]));
+      }
+      for (int i = 0; i < in_b.str.size(); ++i) {
+        auto c = in_b.str[i];
+        if (out.set.find(c) != out.set.end() &&
+            in_a.set.find(c) == in_a.set.end()) {
+          b_concat.push_back(i);
+        }
+      }
+
+      auto& a = inputs[node.positions[0]];
+      auto& b = inputs[node.positions[1]];
+
+      std::unordered_map<char, int> char_map;
+      for (auto i : a_batch) {
+        char_map.insert({in_a.str[i], char_map.size()});
+      }
+      for (auto i : a_concat) {
+        char_map.insert({in_a.str[i], char_map.size()});
+      }
+      for (auto i : b_concat) {
+        char_map.insert({in_b.str[i], char_map.size()});
+      }
+      inputs.emplace_back(batch_tensordot(
+          a,
+          b,
+          std::move(a_contract),
+          std::move(a_batch),
+          std::move(a_concat),
+          std::move(b_contract),
+          std::move(b_batch),
+          std::move(b_concat),
+          s));
+
+      std::vector<int> reorder;
+      for (auto c : node.output.str) {
+        reorder.push_back(char_map[c]);
+      }
+      inputs.back() = transpose(inputs.back(), reorder, s);
+
+    } else {
+      inputs.emplace_back(
+          einsum_naive(node.inputs, node.output, node.positions, inputs, s));
+    }
+
+    // Positions are always sorted increasing, so start from the back
+    for (auto it = node.positions.rbegin(); it != node.positions.rend(); ++it) {
+      inputs.erase(inputs.begin() + *it);
+    }
+  }
+  return inputs.front();
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/einsum.h b/macos/third_party/mlx/mlx/einsum.h
new file mode 100644
index 0000000..05588f8
--- /dev/null
+++ b/macos/third_party/mlx/mlx/einsum.h
@@ -0,0 +1,23 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+MLX_API std::pair<std::vector<std::vector<int>>, std::string> einsum_path(
+    const std::string& subscripts,
+    const std::vector<array>& operands);
+
+MLX_API array einsum(
+    const std::string& subscripts,
+    const std::vector<array>& operands,
+    StreamOrDevice s = {});
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/event.h b/macos/third_party/mlx/mlx/event.h
new file mode 100644
index 0000000..66a6a75
--- /dev/null
+++ b/macos/third_party/mlx/mlx/event.h
@@ -0,0 +1,58 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <stdexcept>
+
+#include "mlx/stream.h"
+
+namespace mlx::core {
+
+class Event {
+ public:
+  Event() {};
+  explicit Event(Stream stream);
+
+  // Wait for the event to be signaled at its current value
+  void wait();
+
+  // Wait in the given stream for the event to be signaled at its current value
+  void wait(Stream stream);
+
+  // Signal the event at its current value in the given stream
+  void signal(Stream stream);
+
+  // Check if the event has been signaled at its current value
+  bool is_signaled() const;
+
+  // Check if the event is valid
+  bool valid() const {
+    return event_ != nullptr;
+  }
+
+  uint64_t value() const {
+    return value_;
+  }
+
+  void set_value(uint64_t v) {
+    value_ = v;
+  }
+
+  const Stream& stream() const {
+    if (!valid()) {
+      throw std::runtime_error(
+          "[Event::stream] Cannot access stream on invalid event.");
+    }
+    return stream_;
+  }
+
+ private:
+  // Default constructed stream should never be used
+  // since the event is not yet valid
+  Stream stream_{0, Device::cpu};
+  std::shared_ptr<void> event_{nullptr};
+  uint64_t value_{0};
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/export.cpp b/macos/third_party/mlx/mlx/export.cpp
new file mode 100644
index 0000000..9a773cd
--- /dev/null
+++ b/macos/third_party/mlx/mlx/export.cpp
@@ -0,0 +1,1130 @@
+// Copyright © 2024 Apple Inc.
+#include "mlx/export.h"
+#include <map>
+#include "mlx/compile_impl.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/graph_utils.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+#include "mlx/version.h"
+
+// clang-format off
+#define SERIALIZE_PRIMITIVE(primitive, ...)  \
+  {                                          \
+    #primitive, {                            \
+      serialize_primitive<primitive>,        \
+      deserialize_primitive<primitive>,      \
+      primitive_state<primitive>,            \
+      {__VA_ARGS__}                          \
+    }                                        \
+  }
+// clang-format on
+
+bool is_big_endian() {
+  int num = 1;
+  return *reinterpret_cast<char*>(&num) != 1;
+}
+
+namespace mlx::core {
+
+using namespace mlx::core::fast;
+
+using Reader = io::ParallelFileReader;
+using Writer = io::FileWriter;
+
+struct PrimitiveSerializer {
+  using Serializer = std::function<void(Writer&, const Primitive&)>;
+  using Deserializer =
+      std::function<std::shared_ptr<Primitive>(Reader&, Stream s)>;
+  using StateExtractor = std::function<std::vector<StateT>(const Primitive&)>;
+
+  PrimitiveSerializer(
+      Serializer serialize,
+      Deserializer deserialize,
+      StateExtractor extract_state,
+      std::vector<std::string> keys = {})
+      : serialize(std::move(serialize)),
+        deserialize(std::move(deserialize)),
+        extract_state(std::move(extract_state)),
+        keys(std::move(keys)) {};
+  Serializer serialize;
+  Deserializer deserialize;
+  StateExtractor extract_state;
+  std::vector<std::string> keys;
+};
+
+template <typename, typename = void>
+constexpr bool is_iterable = false;
+
+template <typename T>
+constexpr bool is_iterable<
+    T,
+    std::void_t<
+        decltype(std::declval<T>().begin()),
+        decltype(std::declval<T>().end())>> = true;
+
+template <template <typename...> class T, typename U>
+constexpr bool is_specialization_of = false;
+
+template <template <typename...> class T, typename... Us>
+constexpr bool is_specialization_of<T, T<Us...>> = true;
+
+template <typename T>
+constexpr bool is_pair = is_specialization_of<std::pair, std::decay_t<T>>;
+
+template <typename T>
+constexpr bool is_tuple = is_specialization_of<std::tuple, std::decay_t<T>>;
+
+template <typename T>
+inline constexpr bool is_optional =
+    is_specialization_of<std::optional, std::decay_t<T>>;
+
+template <typename T>
+inline constexpr bool is_variant =
+    is_specialization_of<std::variant, std::decay_t<T>>;
+
+template <typename>
+constexpr bool dependent_false = false;
+
+template <typename T>
+struct NotSerializable {
+  static_assert(dependent_false<T>, "Type is not serializable.");
+};
+
+template <typename T>
+struct NotDeserializable {
+  static_assert(dependent_false<T>, "Type is not deserializable.");
+};
+
+template <typename T>
+void reverse_bytes(T& data) {
+  auto* bytes = reinterpret_cast<uint8_t*>(&data);
+  for (size_t j = 0; j < (sizeof(T) / 2); j++) {
+    std::swap(bytes[j], bytes[sizeof(T) - j - 1]);
+  }
+}
+
+template <typename T>
+void serialize_variant(Writer& os, T v);
+
+template <typename T>
+T deserialize_variant(Reader& is);
+
+template <typename T>
+void serialize(Writer& os, T v) {
+  if constexpr (std::is_arithmetic_v<T>) {
+    if (is_big_endian()) {
+      reverse_bytes(v);
+    }
+    os.write(reinterpret_cast<const char*>(&v), sizeof(T));
+  } else if constexpr (std::is_enum_v<T>) {
+    serialize(os, static_cast<int>(v));
+  } else if constexpr (std::is_same_v<T, std::nullptr_t>) {
+  } else if constexpr (is_iterable<T>) {
+    serialize(os, static_cast<uint64_t>(v.size()));
+    for (const auto& t : v) {
+      serialize(os, t);
+    }
+  } else if constexpr (is_pair<T> || is_tuple<T>) {
+    std::apply([&os](auto&... x) { (..., serialize(os, x)); }, v);
+  } else if constexpr (is_variant<T>) {
+    serialize_variant(os, v);
+  } else if constexpr (is_optional<T>) {
+    serialize(os, v.has_value());
+    if (v.has_value()) {
+      serialize(os, *v);
+    }
+  } else {
+    NotSerializable<T>();
+  }
+}
+
+template <typename T, std::size_t... I>
+decltype(auto) deserialize_tuple(Reader& is, std::index_sequence<I...>);
+
+template <typename T>
+T deserialize(Reader& is) {
+  if constexpr (std::is_arithmetic_v<T>) {
+    T v;
+    is.read(reinterpret_cast<char*>(&v), sizeof(T));
+    if (is_big_endian()) {
+      reverse_bytes(v);
+    }
+    return v;
+  } else if constexpr (std::is_enum_v<T>) {
+    return static_cast<T>(deserialize<int>(is));
+  } else if constexpr (std::is_same_v<T, std::nullptr_t>) {
+    return nullptr;
+  } else if constexpr (is_iterable<T>) {
+    T v;
+    auto size = deserialize<uint64_t>(is);
+    v.reserve(size);
+    for (int i = 0; i < size; ++i) {
+      v.push_back(deserialize<typename T::value_type>(is));
+    }
+    return v;
+  } else if constexpr (is_pair<T> || is_tuple<T>) {
+    return deserialize_tuple<T>(
+        is, std::make_index_sequence<std::tuple_size_v<std::decay_t<T>>>{});
+  } else if constexpr (is_optional<T>) {
+    auto has_value = deserialize<bool>(is);
+    if (has_value) {
+      return deserialize<T>(is);
+    } else {
+      return std::nullopt;
+    }
+  } else if constexpr (is_variant<T>) {
+    return deserialize_variant<T>(is);
+  } else {
+    NotDeserializable<T>();
+  }
+}
+
+enum class VariantType { Int = 0, Float = 1, Bool = 2 };
+
+template <typename T>
+void serialize_variant(Writer& os, T v) {
+  std::visit(
+      [&](auto&& x) {
+        using ElemT = std::decay_t<decltype(x)>;
+        if constexpr (std::is_same_v<ElemT, int>) {
+          serialize(os, VariantType::Int);
+        } else if constexpr (std::is_same_v<ElemT, float>) {
+          serialize(os, VariantType::Float);
+        } else if constexpr (std::is_same_v<ElemT, bool>) {
+          serialize(os, VariantType::Bool);
+        } else {
+          static_assert(
+              std::is_same_v<ElemT, void>, "Can't serialize variant type.");
+        }
+        serialize(os, x);
+      },
+      v);
+}
+
+template <typename T>
+T deserialize_variant(Reader& is) {
+  auto vt = deserialize<VariantType>(is);
+  switch (vt) {
+    case VariantType::Int:
+      return deserialize<int>(is);
+    case VariantType::Float:
+      return deserialize<float>(is);
+    case VariantType::Bool:
+      return deserialize<bool>(is);
+    default:
+      throw std::runtime_error(
+          "[deserialize_variant] Unknonw variant type tag.");
+  }
+}
+
+template <typename T, std::size_t... I>
+decltype(auto) deserialize_tuple(Reader& is, std::index_sequence<I...>) {
+  return T{deserialize<std::tuple_element_t<I, T>>(is)...};
+};
+
+void serialize(Writer& os, const Stream& s) {
+  serialize(os, s.index);
+  serialize(os, s.device.type);
+  serialize(os, s.device.index);
+}
+template <>
+Stream deserialize(Reader& is) {
+  auto stream_index = deserialize<int>(is);
+  auto device_type = deserialize<Device::DeviceType>(is);
+  auto device_index = deserialize<int>(is);
+  return Stream(stream_index, Device(device_type, device_index));
+}
+
+void serialize(Writer& os, const Dtype& t) {
+  serialize(os, t.val());
+  serialize(os, t.size());
+}
+
+template <>
+Dtype deserialize(Reader& is) {
+  auto val = deserialize<Dtype::Val>(is);
+  auto size = deserialize<uint8_t>(is);
+  return Dtype(val, size);
+};
+
+void serialize(Writer& os, const array& arr) {
+  serialize(os, arr.shape());
+  serialize(os, arr.dtype());
+}
+template <>
+array deserialize(Reader& is) {
+  auto shape = deserialize<Shape>(is);
+  auto type = deserialize<Dtype>(is);
+  return array(std::move(shape), type, nullptr, std::vector<array>{});
+}
+
+template <typename, typename = void>
+constexpr bool has_state = false;
+
+template <typename T>
+constexpr bool has_state<T, std::void_t<decltype(std::declval<T>().state())>> =
+    true;
+
+template <typename T>
+void serialize_primitive(Writer& os, const Primitive& p) {
+  if constexpr (has_state<T>) {
+    serialize(os, static_cast<const T&>(p).state());
+  }
+}
+
+template <typename T>
+void extract_state(const T state, std::vector<StateT>& unpacked_state) {
+  if constexpr (std::is_arithmetic_v<T>) {
+    unpacked_state.push_back(state);
+  } else if constexpr (std::is_enum_v<T>) {
+    unpacked_state.push_back(static_cast<int>(state));
+  } else if constexpr (is_iterable<T>) {
+    unpacked_state.push_back(state);
+  } else if constexpr (is_pair<T> || is_tuple<T>) {
+    std::apply(
+        [&unpacked_state](auto&... x) {
+          (..., extract_state(x, unpacked_state));
+        },
+        state);
+  }
+}
+
+template <typename T>
+std::vector<StateT> primitive_state(const Primitive& p) {
+  std::vector<StateT> state;
+  if constexpr (has_state<T>) {
+    extract_state(static_cast<const T&>(p).state(), state);
+  }
+  return state;
+}
+
+template <typename T>
+std::shared_ptr<T> deserialize_primitive(Reader& is, Stream s) {
+  if constexpr (has_state<T>) {
+    auto args = deserialize<decltype(std::declval<T>().state())>(is);
+    if constexpr (is_pair<decltype(args)> || is_tuple<decltype(args)>) {
+      auto fn = [s](auto&&... args) {
+        return std::make_shared<T>(s, std::move(args)...);
+      };
+      return std::apply(fn, std::move(args));
+    } else {
+      return std::make_shared<T>(s, std::move(args));
+    }
+  } else {
+    return std::make_shared<T>(s);
+  }
+}
+
+struct PrimitiveFactory {
+  std::unordered_map<std::string, PrimitiveSerializer> factory = {
+      SERIALIZE_PRIMITIVE(Abs),
+      SERIALIZE_PRIMITIVE(Add),
+      SERIALIZE_PRIMITIVE(AddMM),
+      SERIALIZE_PRIMITIVE(Arange),
+      SERIALIZE_PRIMITIVE(ArcCos),
+      SERIALIZE_PRIMITIVE(ArcCosh),
+      SERIALIZE_PRIMITIVE(ArcSin),
+      SERIALIZE_PRIMITIVE(ArcSinh),
+      SERIALIZE_PRIMITIVE(ArcTan),
+      SERIALIZE_PRIMITIVE(ArcTan2),
+      SERIALIZE_PRIMITIVE(ArcTanh),
+      SERIALIZE_PRIMITIVE(ArgPartition),
+      SERIALIZE_PRIMITIVE(ArgReduce),
+      SERIALIZE_PRIMITIVE(ArgSort),
+      SERIALIZE_PRIMITIVE(AsType),
+      SERIALIZE_PRIMITIVE(AsStrided),
+      SERIALIZE_PRIMITIVE(
+          BitwiseBinary,
+          "BitwiseAnd",
+          "BitwiseOr",
+          "BitwiseXor",
+          "LeftShift",
+          "RightShift"),
+      SERIALIZE_PRIMITIVE(BlockMaskedMM),
+      SERIALIZE_PRIMITIVE(Broadcast),
+      SERIALIZE_PRIMITIVE(BroadcastAxes),
+      SERIALIZE_PRIMITIVE(Ceil),
+      SERIALIZE_PRIMITIVE(Concatenate),
+      SERIALIZE_PRIMITIVE(Conjugate),
+      SERIALIZE_PRIMITIVE(Convolution),
+      SERIALIZE_PRIMITIVE(Copy),
+      SERIALIZE_PRIMITIVE(Cos),
+      SERIALIZE_PRIMITIVE(Cosh),
+      SERIALIZE_PRIMITIVE(Depends),
+      SERIALIZE_PRIMITIVE(Divide),
+      SERIALIZE_PRIMITIVE(DivMod),
+      SERIALIZE_PRIMITIVE(DynamicSlice),
+      SERIALIZE_PRIMITIVE(DynamicSliceUpdate),
+      SERIALIZE_PRIMITIVE(Equal, "NaNEqual"),
+      SERIALIZE_PRIMITIVE(Erf),
+      SERIALIZE_PRIMITIVE(ErfInv),
+      SERIALIZE_PRIMITIVE(Exp),
+      SERIALIZE_PRIMITIVE(Expm1),
+      SERIALIZE_PRIMITIVE(ExpandDims),
+      SERIALIZE_PRIMITIVE(FFT),
+      SERIALIZE_PRIMITIVE(Flatten),
+      SERIALIZE_PRIMITIVE(Floor),
+      SERIALIZE_PRIMITIVE(Full),
+      SERIALIZE_PRIMITIVE(Gather),
+      SERIALIZE_PRIMITIVE(GatherAxis),
+      SERIALIZE_PRIMITIVE(GatherMM),
+      SERIALIZE_PRIMITIVE(Greater),
+      SERIALIZE_PRIMITIVE(GreaterEqual),
+      SERIALIZE_PRIMITIVE(Hadamard),
+      SERIALIZE_PRIMITIVE(Imag),
+      SERIALIZE_PRIMITIVE(Less),
+      SERIALIZE_PRIMITIVE(LessEqual),
+      SERIALIZE_PRIMITIVE(Log, "Log2", "Log10"),
+      SERIALIZE_PRIMITIVE(Log1p),
+      SERIALIZE_PRIMITIVE(LogicalNot),
+      SERIALIZE_PRIMITIVE(LogicalAnd),
+      SERIALIZE_PRIMITIVE(LogicalOr),
+      SERIALIZE_PRIMITIVE(LogAddExp),
+      SERIALIZE_PRIMITIVE(LogSumExp),
+      SERIALIZE_PRIMITIVE(MaskedScatter),
+      SERIALIZE_PRIMITIVE(Matmul),
+      SERIALIZE_PRIMITIVE(Maximum),
+      SERIALIZE_PRIMITIVE(Minimum),
+      SERIALIZE_PRIMITIVE(Multiply),
+      SERIALIZE_PRIMITIVE(Negative),
+      SERIALIZE_PRIMITIVE(NotEqual),
+      SERIALIZE_PRIMITIVE(Reshape),
+      SERIALIZE_PRIMITIVE(NumberOfElements),
+      SERIALIZE_PRIMITIVE(Pad),
+      SERIALIZE_PRIMITIVE(Partition),
+      SERIALIZE_PRIMITIVE(Power),
+      SERIALIZE_PRIMITIVE(QuantizedMatmul),
+      SERIALIZE_PRIMITIVE(GatherQMM),
+      SERIALIZE_PRIMITIVE(RandomBits),
+      SERIALIZE_PRIMITIVE(Real),
+      SERIALIZE_PRIMITIVE(Remainder),
+      SERIALIZE_PRIMITIVE(Reshape),
+      SERIALIZE_PRIMITIVE(Reduce, "And", "Or", "Sum", "Prod", "Min", "Max"),
+      SERIALIZE_PRIMITIVE(Round),
+      SERIALIZE_PRIMITIVE(
+          Scan,
+          "CumSum",
+          "CumProd",
+          "CumMin",
+          "CumMax",
+          "CumLogaddexp"),
+      SERIALIZE_PRIMITIVE(Scatter),
+      SERIALIZE_PRIMITIVE(ScatterAxis),
+      SERIALIZE_PRIMITIVE(Select),
+      SERIALIZE_PRIMITIVE(Sigmoid),
+      SERIALIZE_PRIMITIVE(Sign),
+      SERIALIZE_PRIMITIVE(Sin),
+      SERIALIZE_PRIMITIVE(Sinh),
+      SERIALIZE_PRIMITIVE(Slice),
+      SERIALIZE_PRIMITIVE(SliceUpdate),
+      SERIALIZE_PRIMITIVE(Softmax),
+      SERIALIZE_PRIMITIVE(Sort),
+      SERIALIZE_PRIMITIVE(Split),
+      SERIALIZE_PRIMITIVE(Square),
+      SERIALIZE_PRIMITIVE(Squeeze),
+      SERIALIZE_PRIMITIVE(Sqrt, "Rsqrt", "Sqrt"),
+      SERIALIZE_PRIMITIVE(StopGradient),
+      SERIALIZE_PRIMITIVE(Subtract),
+      SERIALIZE_PRIMITIVE(Tan),
+      SERIALIZE_PRIMITIVE(Tanh),
+      SERIALIZE_PRIMITIVE(View),
+      SERIALIZE_PRIMITIVE(Transpose),
+      SERIALIZE_PRIMITIVE(Unflatten),
+      SERIALIZE_PRIMITIVE(QRF),
+      SERIALIZE_PRIMITIVE(SVD),
+      SERIALIZE_PRIMITIVE(Inverse),
+      SERIALIZE_PRIMITIVE(Cholesky),
+      SERIALIZE_PRIMITIVE(Eig),
+      SERIALIZE_PRIMITIVE(Eigh),
+      SERIALIZE_PRIMITIVE(Quantize),
+      SERIALIZE_PRIMITIVE(RMSNorm),
+      SERIALIZE_PRIMITIVE(RMSNormVJP),
+      SERIALIZE_PRIMITIVE(LayerNorm),
+      SERIALIZE_PRIMITIVE(LayerNormVJP),
+      SERIALIZE_PRIMITIVE(RoPE),
+      SERIALIZE_PRIMITIVE(ScaledDotProductAttention),
+      SERIALIZE_PRIMITIVE(CustomKernel)};
+  std::unordered_map<std::string, std::string> name_remap;
+
+  PrimitiveFactory() {
+    for (auto& [n, f] : factory) {
+      for (auto& k : f.keys) {
+        name_remap[k] = n;
+      }
+    }
+  }
+
+  void save(Writer& os, const std::shared_ptr<Primitive>& p) {
+    serialize(os, p->stream());
+    std::string name = p->name();
+    name = name.substr(0, name.find(' '));
+    if (auto it = name_remap.find(name); it != name_remap.end()) {
+      name = it->second;
+    }
+    serialize(os, name);
+    if (auto it = factory.find(name); it != factory.end()) {
+      it->second.serialize(os, *p);
+    } else {
+      throw std::invalid_argument(
+          "[export_function] Unable to serialize primitive " + name);
+    }
+  };
+
+  std::shared_ptr<Primitive> load(Reader& is) {
+    auto stream = deserialize<Stream>(is);
+    if (get_stream(stream.index) != stream) {
+      std::ostringstream msg;
+      msg << "[import_function] Invalid stream encountered " << stream << ".";
+      throw std::invalid_argument(msg.str());
+    }
+    auto name = deserialize<std::string>(is);
+    if (auto it = factory.find(name); it != factory.end()) {
+      return it->second.deserialize(is, stream);
+    } else {
+      throw std::invalid_argument(
+          "[import_function] Unable to deserialize primitive " + name);
+    }
+  };
+
+  std::pair<std::string, std::vector<StateT>> extract_state(
+      const std::shared_ptr<Primitive>& p) {
+    std::string name = p->name();
+    name = name.substr(0, name.find(' '));
+    if (auto it = name_remap.find(name); it != name_remap.end()) {
+      name = it->second;
+    }
+
+    if (auto it = factory.find(name); it != factory.end()) {
+      auto state = it->second.extract_state(*p);
+      return {name, state};
+    } else {
+      throw std::invalid_argument(
+          "[export_function] Unable to get state for primitive " + name);
+    }
+  };
+};
+
+void write_header(Writer& os, int count, bool shapeless) {
+  serialize(os, std::string(version()));
+  serialize(os, count);
+  serialize(os, shapeless);
+}
+
+// A struct to hold and retrieve the graphs that are exported / imported
+struct FunctionTable {
+  FunctionTable(bool shapeless = false) : shapeless(shapeless) {};
+  struct Function {
+    Function(
+        std::vector<std::string> kwarg_keys,
+        std::vector<array> inputs,
+        std::vector<array> outputs,
+        std::vector<array> tape)
+        : kwarg_keys(std::move(kwarg_keys)),
+          inputs(std::move(inputs)),
+          outputs(std::move(outputs)),
+          tape(std::move(tape)) {}
+
+    std::vector<std::string> kwarg_keys;
+    std::vector<array> inputs;
+    std::vector<array> outputs;
+    std::vector<array> tape;
+    Function(const Function&) = delete;
+    Function& operator=(const Function&) = delete;
+    Function(Function&&) = default;
+    Function() = default;
+  };
+  bool shapeless;
+  std::unordered_map<int, std::vector<Function>> table;
+  Function* find(const Args& args, const std::map<std::string, array>& kwargs);
+  std::pair<Function&, bool> emplace(
+      const Args& args,
+      const std::map<std::string, array>& kwargs);
+  void insert(
+      std::vector<std::string> kwarg_keys,
+      std::vector<array> inputs,
+      std::vector<array> outputs,
+      std::vector<array> tape) {
+    auto [it, _] = table.emplace(inputs.size(), std::vector<Function>{});
+    it->second.emplace_back(
+        std::move(kwarg_keys),
+        std::move(inputs),
+        std::move(outputs),
+        std::move(tape));
+  }
+
+  void print_functions(std::ostream& os) {
+    int n = 1;
+    for (auto& [_, vec] : table) {
+      for (auto& fun : vec) {
+        auto npos = fun.inputs.size() - fun.kwarg_keys.size();
+        os << " " << n++ << ". Function with " << npos
+           << " positional inputs and " << fun.kwarg_keys.size()
+           << " keyword inputs:\n";
+        for (int j = 0; j < fun.inputs.size(); ++j) {
+          auto& in = fun.inputs[j];
+          if (j < npos) {
+            os << "   " << j + 1 << ": ";
+          } else {
+            os << "   \"" << fun.kwarg_keys[j - npos] << "\": ";
+          }
+          os << in.shape() << " " << in.dtype() << "\n";
+        }
+      }
+    }
+  }
+
+ private:
+  bool match(
+      const Args& args,
+      const std::map<std::string, array>& kwargs,
+      const Function& fun);
+};
+
+bool FunctionTable::match(
+    const Args& args,
+    const std::map<std::string, array>& kwargs,
+    const Function& fun) {
+  for (auto& k : fun.kwarg_keys) {
+    if (kwargs.find(k) == kwargs.end()) {
+      return false;
+    }
+  }
+
+  auto match_inputs = [shapeless = this->shapeless](
+                          const array& x, const array& y) {
+    if (x.dtype() != y.dtype()) {
+      return false;
+    }
+    if (x.ndim() != y.ndim()) {
+      return false;
+    }
+    if (!shapeless && x.shape() != y.shape()) {
+      return false;
+    }
+    return true;
+  };
+
+  int i = 0;
+  for (; i < args.size(); ++i) {
+    if (!match_inputs(args[i], fun.inputs[i])) {
+      return false;
+    }
+  }
+  for (auto& [_, in] : kwargs) {
+    if (!match_inputs(in, fun.inputs[i++])) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+std::pair<FunctionTable::Function&, bool> FunctionTable::emplace(
+    const Args& args,
+    const std::map<std::string, array>& kwargs) {
+  auto n_inputs = args.size() + kwargs.size();
+  auto [it, _] = table.emplace(n_inputs, std::vector<Function>{});
+  auto& funs_vec = it->second;
+
+  for (auto& fun : funs_vec) {
+    if (match(args, kwargs, fun)) {
+      return {fun, false};
+    }
+  }
+
+  funs_vec.emplace_back();
+  return {funs_vec.back(), true};
+}
+
+FunctionTable::Function* FunctionTable::find(
+    const Args& args,
+    const std::map<std::string, array>& kwargs) {
+  auto n_inputs = args.size() + kwargs.size();
+  auto it = table.find(n_inputs);
+  if (it == table.end()) {
+    return nullptr;
+  }
+
+  for (auto& fun : it->second) {
+    if (match(args, kwargs, fun)) {
+      return &fun;
+    }
+  }
+
+  return nullptr;
+}
+
+FunctionExporter::FunctionExporter(
+    const std::string& file,
+    std::function<std::vector<array>(const Args&, const Kwargs&)> fun,
+    bool shapeless)
+    : os(file),
+      fun(std::move(fun)),
+      ftable(std::make_shared<FunctionTable>(shapeless)) {
+  if (!os.is_open()) {
+    throw std::runtime_error("[export_function] Failed to open " + file);
+  }
+  write_header(os, count, shapeless);
+}
+
+FunctionExporter::FunctionExporter(
+    const ExportCallback& callback,
+    std::function<std::vector<array>(const Args&, const Kwargs&)> fun,
+    bool shapeless)
+    : callback(callback),
+      fun(std::move(fun)),
+      ftable(std::make_shared<FunctionTable>(shapeless)) {}
+
+void FunctionExporter::close() {
+  closed = true;
+};
+
+void FunctionExporter::export_with_callback(
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    const std::vector<array>& tape,
+    const std::vector<std::string>& kwarg_keys) {
+  NodeNamer namer{};
+  auto to_vector_data = [&namer](const auto& arrays) {
+    std::vector<std::tuple<std::string, Shape, Dtype>> data;
+    for (auto& a : arrays) {
+      data.emplace_back(namer.get_name(a), a.shape(), a.dtype());
+    }
+    return data;
+  };
+
+  // Callback on the inputs
+  callback({{"type", "inputs"}, {"inputs", to_vector_data(inputs)}});
+  std::vector<std::pair<std::string, std::string>> keyword_inputs;
+  for (int i = inputs.size() - kwarg_keys.size(), j = 0; i < inputs.size();
+       ++i, ++j) {
+    keyword_inputs.emplace_back(kwarg_keys[j], namer.get_name(inputs[i]));
+  }
+  callback({{"type", "keyword_inputs"}, {"keywords", keyword_inputs}});
+
+  // Callback on the outputs
+  callback({{"type", "outputs"}, {"outputs", to_vector_data(outputs)}});
+
+  // Callback on the constants
+  {
+    std::unordered_set<std::uintptr_t> input_set;
+    for (auto& in : inputs) {
+      input_set.insert(in.id());
+    }
+    std::vector<std::pair<std::string, array>> new_constants;
+    for (auto& arr : tape) {
+      if (arr.has_primitive() || input_set.find(arr.id()) != input_set.end()) {
+        continue;
+      }
+      if (constants.insert({arr.id(), arr}).second) {
+        new_constants.emplace_back(namer.get_name(arr), arr);
+      }
+    }
+    callback({{"type", "constants"}, {"constants", new_constants}});
+  }
+  auto factory = PrimitiveFactory();
+
+  // Callback for each primitive in the tape
+  for (auto& arr : tape) {
+    if (!arr.has_primitive()) {
+      continue;
+    }
+    auto [name, state] = factory.extract_state(arr.primitive_ptr());
+    callback(
+        {{"type", "primitive"},
+         {"inputs", to_vector_data(arr.inputs())},
+         {"outputs", to_vector_data(arr.outputs())},
+         {"name", name},
+         {"arguments", state}});
+  }
+}
+
+void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) {
+  if (closed) {
+    throw std::runtime_error(
+        "[export_function] Attempting to write after exporting is closed.");
+  }
+  auto sorted_kwargs =
+      std::map<std::string, array>(kwargs.begin(), kwargs.end());
+  auto [fentry, inserted] = ftable->emplace(args, sorted_kwargs);
+  if (!inserted) {
+    throw std::runtime_error(
+        "[export_function] Attempting to export a function twice with "
+        "the same signature is not allowed.");
+  }
+
+  // Flatten the inputs to the function for tracing
+  std::vector<std::string> kwarg_keys;
+  auto inputs = args;
+  for (auto& [k, v] : sorted_kwargs) {
+    kwarg_keys.push_back(k);
+    inputs.push_back(v);
+  }
+
+  auto flat_fun = [this, &kwarg_keys](const Args& flat_args) {
+    auto args = Args(flat_args.begin(), flat_args.end() - kwarg_keys.size());
+    Kwargs kwargs;
+    auto it = flat_args.end() - kwarg_keys.size();
+    ;
+    for (auto& k : kwarg_keys) {
+      kwargs.insert({k, *it++});
+    }
+    return detail::ArraysAndExtra{fun(args, kwargs), nullptr};
+  };
+
+  // Trace to build the graph
+  auto [trace_inputs, trace_outputs, extra] =
+      detail::compile_trace(flat_fun, inputs, ftable->shapeless);
+
+  // DFS the graph and get the tape
+  auto [tape, parents_map] =
+      detail::compile_dfs(trace_inputs, trace_outputs, inputs);
+
+  detail::compile_simplify(tape, parents_map, trace_outputs, /* passes */ 3);
+
+  // Update the table entry
+  fentry.kwarg_keys = kwarg_keys;
+  fentry.inputs = trace_inputs;
+
+  count++;
+
+  if (callback) {
+    export_with_callback(trace_inputs, trace_outputs, tape, kwarg_keys);
+    return;
+  }
+
+  // Update the header
+  auto pos = os.tell();
+  os.seek(0);
+  write_header(os, count, ftable->shapeless);
+  os.seek(pos);
+  serialize(os, kwarg_keys);
+
+  auto arrays_to_ids = [](const std::vector<array>& arrs) {
+    std::vector<uint64_t> ids;
+    for (auto& arr : arrs) {
+      ids.push_back(arr.id());
+    }
+    return ids;
+  };
+
+  // Inputs and outputs
+  auto trace_input_ids = arrays_to_ids(trace_inputs);
+  serialize(os, trace_input_ids);
+  serialize(os, trace_inputs);
+  serialize(os, arrays_to_ids(trace_outputs));
+
+  std::unordered_set<std::uintptr_t> input_set(
+      trace_input_ids.begin(), trace_input_ids.end());
+
+  // Tape
+  auto factory = PrimitiveFactory();
+  serialize(os, static_cast<uint64_t>(tape.size()));
+  for (auto& arr : tape) {
+    serialize(os, static_cast<uint64_t>(arr.id()));
+    if (arr.has_primitive()) {
+      serialize(os, true);
+      serialize(os, arrays_to_ids(arr.inputs()));
+      factory.save(os, arr.primitive_ptr());
+      serialize(os, static_cast<uint64_t>(arr.siblings().size()));
+      if (arr.siblings().empty()) {
+        serialize(os, arr.shape());
+        serialize(os, arr.dtype());
+      } else {
+        auto outputs = arr.outputs();
+        serialize(os, arrays_to_ids(outputs));
+
+        std::vector<Shape> shapes;
+        std::vector<Dtype> dtypes;
+        for (auto& o : outputs) {
+          shapes.push_back(o.shape());
+          dtypes.push_back(o.dtype());
+        }
+        serialize(os, shapes);
+        serialize(os, dtypes);
+      }
+    } else {
+      serialize(os, false);
+      if (input_set.find(arr.id()) == input_set.end()) {
+        serialize(os, true);
+        // Save constant data if not already saved
+        if (constants.insert({arr.id(), arr}).second) {
+          serialize(os, arr.shape());
+          serialize(os, arr.dtype());
+          os.write(arr.data<char>(), arr.nbytes());
+        }
+      } else {
+        serialize(os, false);
+      }
+    }
+  }
+}
+
+void FunctionExporter::operator()(const Args& args) {
+  export_function(args, {});
+}
+
+void FunctionExporter::operator()(const Kwargs& kwargs) {
+  export_function({}, kwargs);
+}
+
+void FunctionExporter::operator()(const Args& args, const Kwargs& kwargs) {
+  export_function(args, kwargs);
+}
+
+FunctionExporter exporter(
+    const std::string& file,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    bool shapeless /* = false */) {
+  return FunctionExporter{
+      file,
+      [fun](const Args& args, const Kwargs&) { return fun(args); },
+      shapeless};
+}
+
+FunctionExporter exporter(
+    const std::string& file,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    bool shapeless /* = false */) {
+  return exporter(
+      file,
+      [fun](const Args&, const Kwargs kwargs) { return fun(kwargs); },
+      shapeless);
+}
+
+FunctionExporter exporter(
+    const std::string& file,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    bool shapeless /* = false */) {
+  return FunctionExporter{file, fun, shapeless};
+}
+
+void export_function(
+    const std::string& file,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    const Args& args,
+    bool shapeless /* = false */) {
+  exporter(file, fun, shapeless)(args);
+}
+
+void export_function(
+    const std::string& file,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    const Kwargs& kwargs,
+    bool shapeless /* = false */) {
+  exporter(file, fun, shapeless)(kwargs);
+}
+
+void export_function(
+    const std::string& file,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    const Args& args,
+    const Kwargs& kwargs,
+    bool shapeless /* = false */) {
+  exporter(file, fun, shapeless)(args, kwargs);
+}
+
+FunctionExporter exporter(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    bool shapeless /* = false */) {
+  return FunctionExporter{
+      callback,
+      [fun](const Args& args, const Kwargs&) { return fun(args); },
+      shapeless};
+}
+
+FunctionExporter exporter(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    bool shapeless /* = false */) {
+  return exporter(
+      callback,
+      [fun](const Args&, const Kwargs kwargs) { return fun(kwargs); },
+      shapeless);
+}
+
+FunctionExporter exporter(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    bool shapeless /* = false */) {
+  return FunctionExporter{callback, fun, shapeless};
+}
+
+void export_function(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    const Args& args,
+    bool shapeless /* = false */) {
+  exporter(callback, fun, shapeless)(args);
+}
+
+void export_function(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    const Kwargs& kwargs,
+    bool shapeless /* = false */) {
+  exporter(callback, fun, shapeless)(kwargs);
+}
+
+void export_function(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    const Args& args,
+    const Kwargs& kwargs,
+    bool shapeless /* = false */) {
+  exporter(callback, fun, shapeless)(args, kwargs);
+}
+
+std::vector<array> ImportedFunction::operator()(const Kwargs& kwargs) const {
+  return this->operator()({}, kwargs);
+}
+
+std::vector<array> ImportedFunction::operator()(const Args& args) const {
+  return this->operator()(args, {});
+}
+
+std::vector<array> ImportedFunction::operator()(
+    const Args& args,
+    const Kwargs& kwargs) const {
+  auto sorted_kwargs =
+      std::map<std::string, array>(kwargs.begin(), kwargs.end());
+  auto* fun = ftable->find(args, sorted_kwargs);
+  if (fun == nullptr) {
+    std::ostringstream msg;
+    msg << "[import_function::call] No imported function found which matches "
+        << "the given positional and keyword arguments. Possible functions include:\n";
+    ftable->print_functions(msg);
+    msg << "\nCalled with " << args.size() << " positional inputs and "
+        << kwargs.size() << " keyword inputs:\n";
+    for (int i = 0; i < args.size(); ++i) {
+      auto& in = args[i];
+      msg << "  " << i + 1 << ": " << in.shape() << " " << in.dtype() << "\n";
+    }
+    for (auto& [k, in] : kwargs) {
+      msg << "  \"" << k << "\": " << in.shape() << " " << in.dtype() << "\n";
+    }
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto inputs = args;
+  for (auto& [_, v] : sorted_kwargs) {
+    inputs.push_back(v);
+  }
+  return detail::compile_replace(
+      fun->tape, fun->inputs, fun->outputs, inputs, ftable->shapeless);
+}
+
+ImportedFunction import_function(const std::string& file) {
+  return ImportedFunction{file};
+}
+
+ImportedFunction::ImportedFunction(const std::string& file)
+    : ftable(std::make_shared<FunctionTable>()) {
+  auto is_ptr = std::make_shared<Reader>(file);
+  auto& is = *is_ptr;
+  if (!is.is_open()) {
+    throw std::runtime_error("[import_function] Failed to open " + file);
+  }
+
+  // Parse header
+  auto mlx_version = deserialize<std::string>(is);
+  auto function_count = deserialize<int>(is);
+  ftable->shapeless = deserialize<bool>(is);
+  std::unordered_map<std::uintptr_t, array> constants;
+
+  auto import_one = [&]() {
+    auto kwarg_keys = deserialize<std::vector<std::string>>(is);
+
+    std::unordered_map<uint64_t, array> array_map;
+    auto trace_input_ids = deserialize<std::vector<uint64_t>>(is);
+    auto trace_inputs = deserialize<std::vector<array>>(is);
+    for (int i = 0; i < trace_inputs.size(); ++i) {
+      array_map.emplace(trace_input_ids[i], trace_inputs[i]);
+    }
+    auto trace_output_ids = deserialize<std::vector<uint64_t>>(is);
+
+    std::vector<array> tape;
+    auto tape_size = deserialize<uint64_t>(is);
+    tape.reserve(tape_size);
+
+    auto factory = PrimitiveFactory();
+    for (size_t i = 0; i < tape_size; ++i) {
+      auto id = deserialize<uint64_t>(is);
+      if (deserialize<bool>(is)) {
+        auto input_ids = deserialize<std::vector<uint64_t>>(is);
+        std::vector<array> inputs;
+        inputs.reserve(input_ids.size());
+        for (auto id : input_ids) {
+          inputs.push_back(array_map.at(id));
+        }
+        std::shared_ptr<Primitive> prim = factory.load(is);
+        auto num_siblings = deserialize<uint64_t>(is);
+        if (num_siblings == 0) {
+          auto shape = deserialize<Shape>(is);
+          auto type = deserialize<Dtype>(is);
+          tape.emplace_back(
+              std::move(shape), type, std::move(prim), std::move(inputs));
+          array_map.emplace(id, tape.back());
+        } else {
+          auto ids = deserialize<std::vector<uint64_t>>(is);
+          auto shapes = deserialize<std::vector<Shape>>(is);
+          auto types = deserialize<std::vector<Dtype>>(is);
+          auto arrays = array::make_arrays(
+              std::move(shapes),
+              std::move(types),
+              std::move(prim),
+              std::move(inputs));
+          for (int i = 0; i < arrays.size(); ++i) {
+            auto sid = ids[i];
+            if (sid == id) {
+              tape.push_back(arrays[i]);
+            }
+            array_map.emplace(sid, arrays[i]);
+          }
+        }
+      } else {
+        if (deserialize<bool>(is)) {
+          // Load constant
+          if (auto it = constants.find(id); it != constants.end()) {
+            tape.push_back(it->second);
+          } else {
+            auto shape = deserialize<Shape>(is);
+            auto type = deserialize<Dtype>(is);
+            size_t offset = is.tell();
+            tape.push_back(array(
+                std::move(shape),
+                type,
+                std::make_shared<Load>(
+                    default_stream(Device::cpu), is_ptr, offset),
+                {}));
+            is.seek(offset + tape.back().nbytes());
+            constants.insert({id, tape.back()});
+          }
+          array_map.emplace(id, tape.back());
+        } else {
+          // Function inputs are in the map
+          tape.push_back(array_map.at(id));
+        }
+      }
+    }
+
+    std::vector<array> trace_outputs;
+    trace_outputs.reserve(trace_output_ids.size());
+    for (auto id : trace_output_ids) {
+      trace_outputs.push_back(array_map.at(id));
+    }
+    ftable->insert(
+        std::move(kwarg_keys),
+        std::move(trace_inputs),
+        std::move(trace_outputs),
+        std::move(tape));
+  };
+
+  for (int i = 0; i < function_count; ++i) {
+    import_one();
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/export.h b/macos/third_party/mlx/mlx/export.h
new file mode 100644
index 0000000..5532f7c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/export.h
@@ -0,0 +1,137 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <optional>
+#include <set>
+#include <unordered_map>
+#include <variant>
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+using Args = std::vector<array>;
+using Kwargs = std::unordered_map<std::string, array>;
+
+// Possible types for a Primitive's state
+using StateT = std::variant<
+    bool,
+    int,
+    size_t,
+    float,
+    double,
+    Dtype,
+    Shape,
+    Strides,
+    std::vector<int>,
+    std::vector<size_t>,
+    std::vector<std::tuple<bool, bool, bool>>,
+    std::vector<std::variant<bool, int, float>>,
+    std::optional<float>,
+    std::string>;
+
+using ExportCallbackInput = std::unordered_map<
+    std::string,
+    std::variant<
+        std::vector<std::tuple<std::string, Shape, Dtype>>,
+        std::vector<std::pair<std::string, array>>,
+        std::vector<std::pair<std::string, std::string>>,
+        std::vector<StateT>,
+        std::string>>;
+using ExportCallback = std::function<void(const ExportCallbackInput&)>;
+
+struct FunctionExporter;
+
+/**
+ * Make an exporter to save multiple traces of a given function to
+ * the same file.
+ */
+MLX_API FunctionExporter exporter(
+    const std::string& file,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    bool shapeless = false);
+
+MLX_API FunctionExporter exporter(
+    const std::string& file,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    bool shapeless = false);
+
+MLX_API FunctionExporter exporter(
+    const std::string& path,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    bool shapeless = false);
+
+/**
+ * Export a function to a file.
+ */
+MLX_API void export_function(
+    const std::string& file,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    const Args& args,
+    bool shapeless = false);
+
+MLX_API void export_function(
+    const std::string& file,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    const Kwargs& kwargs,
+    bool shapeless = false);
+
+MLX_API void export_function(
+    const std::string& file,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    const Args& args,
+    const Kwargs& kwargs,
+    bool shapeless = false);
+
+struct ImportedFunction;
+
+/**
+ * Import a function from a file.
+ */
+MLX_API ImportedFunction import_function(const std::string& file);
+
+/**
+ * Make an exporter to export multiple traces of a given function with the same
+ * callback.
+ */
+MLX_API FunctionExporter exporter(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    bool shapeless = false);
+
+MLX_API FunctionExporter exporter(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    bool shapeless = false);
+
+MLX_API FunctionExporter exporter(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    bool shapeless = false);
+
+/**
+ * Export a function with a callback.
+ */
+MLX_API void export_function(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&)>& fun,
+    const Args& args,
+    bool shapeless = false);
+
+MLX_API void export_function(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Kwargs&)>& fun,
+    const Kwargs& kwargs,
+    bool shapeless = false);
+
+MLX_API void export_function(
+    const ExportCallback& callback,
+    const std::function<std::vector<array>(const Args&, const Kwargs&)>& fun,
+    const Args& args,
+    const Kwargs& kwargs,
+    bool shapeless = false);
+
+} // namespace mlx::core
+
+#include "mlx/export_impl.h"
diff --git a/macos/third_party/mlx/mlx/export_impl.h b/macos/third_party/mlx/mlx/export_impl.h
new file mode 100644
index 0000000..467a5f0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/export_impl.h
@@ -0,0 +1,99 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/api.h"
+#include "mlx/io/load.h"
+
+#pragma once
+
+namespace mlx::core {
+
+struct FunctionTable;
+
+struct MLX_API FunctionExporter {
+  void operator()(const std::initializer_list<array>& args) {
+    this->operator()(Args(args));
+  }
+  void operator()(const Args& args);
+  void operator()(const Kwargs& kwargs);
+  void operator()(const Args& args, const Kwargs& kwargs);
+
+  void close();
+
+  FunctionExporter(const FunctionExporter&) = delete;
+  FunctionExporter& operator=(const FunctionExporter&) = delete;
+  FunctionExporter(FunctionExporter&& other) = default;
+
+ private:
+  friend MLX_API FunctionExporter exporter(
+      const std::string&,
+      const std::function<std::vector<array>(const Args&)>&,
+      bool shapeless);
+
+  friend MLX_API FunctionExporter exporter(
+      const std::string&,
+      const std::function<std::vector<array>(const Kwargs&)>&,
+      bool shapeless);
+
+  friend MLX_API FunctionExporter exporter(
+      const std::string&,
+      const std::function<std::vector<array>(const Args&, const Kwargs&)>&,
+      bool shapeless);
+
+  friend MLX_API FunctionExporter exporter(
+      const ExportCallback&,
+      const std::function<std::vector<array>(const Args&)>&,
+      bool shapeless);
+
+  friend MLX_API FunctionExporter exporter(
+      const ExportCallback&,
+      const std::function<std::vector<array>(const Kwargs&)>&,
+      bool shapeless);
+
+  friend MLX_API FunctionExporter exporter(
+      const ExportCallback&,
+      const std::function<std::vector<array>(const Args&, const Kwargs&)>&,
+      bool shapeless);
+
+  FunctionExporter(
+      const std::string& file,
+      std::function<std::vector<array>(const Args&, const Kwargs&)> fun,
+      bool shapeless);
+
+  FunctionExporter(
+      const ExportCallback& callback,
+      std::function<std::vector<array>(const Args&, const Kwargs&)> fun,
+      bool shapeless);
+
+  io::FileWriter os;
+  ExportCallback callback;
+  std::function<std::vector<array>(const Args&, const Kwargs& kwargs)> fun;
+  void export_function(const Args& args, const Kwargs& kwargs);
+  void export_with_callback(
+      const std::vector<array>& inputs,
+      const std::vector<array>& outputs,
+      const std::vector<array>& tape,
+      const std::vector<std::string>& kwarg_keys);
+  std::unordered_map<std::uintptr_t, array> constants;
+  int count{0};
+  bool closed{false};
+  std::shared_ptr<FunctionTable> ftable;
+};
+
+struct MLX_API ImportedFunction {
+  std::vector<array> operator()(
+      const std::initializer_list<array>& args) const {
+    return this->operator()(Args(args));
+  }
+  std::vector<array> operator()(const Args& args) const;
+  std::vector<array> operator()(const Kwargs& kwargs) const;
+  std::vector<array> operator()(const Args& args, const Kwargs& kwargs) const;
+
+ private:
+  ImportedFunction(const std::string& file);
+  friend MLX_API ImportedFunction import_function(const std::string&);
+  ImportedFunction();
+
+  std::shared_ptr<FunctionTable> ftable;
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/fast.cpp b/macos/third_party/mlx/mlx/fast.cpp
new file mode 100644
index 0000000..bf140b7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/fast.cpp
@@ -0,0 +1,941 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <cassert>
+#include <numeric>
+
+#include "mlx/fast.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/ops.h"
+#include "mlx/transforms.h"
+#include "mlx/transforms_impl.h"
+
+namespace mlx::core::fast {
+
+std::vector<array> Custom::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  auto [_, vjps] = mlx::core::vjp(fallback_, primals, cotangents);
+  std::vector<array> vjp_outs;
+  for (int i = 0, j = 0; i < vjps.size(); ++i) {
+    if (j < argnums.size() && i == argnums[j]) {
+      vjp_outs.push_back(vjps[i]);
+      j++;
+    }
+  }
+  return vjp_outs;
+}
+
+std::vector<array> Custom::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  std::vector<array> all_tangents;
+  for (int i = 0, j = 0; i < primals.size(); i++) {
+    if (j < argnums.size() && i == argnums[j]) {
+      all_tangents.emplace_back(tangents[j++]);
+    } else {
+      all_tangents.emplace_back(zeros_like(primals[i]));
+    }
+  }
+  auto [_, jvps] = mlx::core::jvp(fallback_, primals, all_tangents);
+  return jvps;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Custom::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto outputs = mlx::core::vmap(fallback_, axes)(inputs);
+  auto out_axes = std::vector<int>(outputs.size(), 0);
+  return {outputs, out_axes};
+}
+
+array rms_norm(
+    const array& x,
+    const std::optional<array>& weight,
+    float eps,
+    StreamOrDevice s_ /* = {} */) {
+  bool has_weight = weight.has_value();
+
+  if (x.ndim() == 0) {
+    std::ostringstream msg;
+    msg << "[rms_norm] Input must have at least 1 dimension but got input with "
+           "0 dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (has_weight) {
+    if ((*weight).ndim() != 1) {
+      std::ostringstream msg;
+      msg << "[rms_norm] (*weight) must have 1 dimension but has "
+          << (*weight).ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if ((*weight).size() != x.shape(-1)) {
+      std::ostringstream msg;
+      msg << "[rms_norm] (*weight) must have the same size as the last dimension of"
+             " x but has "
+          << (*weight).size() << " elements.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  auto out_type = (weight.has_value()) ? result_type(x, (*weight)) : x.dtype();
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[rms_norm] Received unsupported type " << out_type << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto s = to_stream(s_);
+  auto fallback =
+      [has_weight, eps, out_type, s](const std::vector<array>& inputs) {
+        auto x = astype(inputs[0], float32, s);
+        x = multiply(
+            x,
+            rsqrt(
+                add(mean(square(x, s), -1, /* keepdims */ true, s),
+                    array(eps, float32),
+                    s),
+                s),
+            s);
+        x = astype(x, out_type, s);
+
+        if (has_weight) {
+          x = multiply(x, inputs[1], s);
+        }
+
+        return std::vector<array>{x};
+      };
+
+  auto passed_weight =
+      (has_weight) ? astype(*weight, out_type, s) : array(1, out_type);
+
+  if (!RMSNorm::use_fallback(s)) {
+    return array(
+        x.shape(),
+        out_type,
+        std::make_shared<RMSNorm>(s, fallback, eps),
+        {astype(x, out_type, s), passed_weight});
+  }
+  return fallback({x, passed_weight})[0];
+}
+
+std::vector<array> RMSNorm::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  assert(primals.size() == 2);
+  assert(outputs.size() == 1);
+  assert(cotangents.size() == 1);
+
+  auto s = stream();
+  auto fallback = [eps = eps_, s](const std::vector<array>& inputs) {
+    auto& x = inputs[0];
+    auto& w = inputs[1];
+    auto& g = inputs[2];
+
+    std::vector<array> vjps;
+
+    auto n = rsqrt(
+        add(mean(square(x, s), /* axis= */ -1, /* keepdims= */ true, s),
+            array(eps, x.dtype()),
+            s),
+        s);
+    auto n3 = power(n, array(3, x.dtype()), s);
+
+    // df/dx
+    auto gw = multiply(g, w, s);
+    auto t = mean(multiply(gw, x, s), /* axis= */ -1, /* keepdims= */ true, s);
+    t = multiply(multiply(x, t, s), n3, s);
+    vjps.push_back(subtract(multiply(gw, n, s), t, s));
+
+    // df/dw
+    std::vector<int> axes(g.ndim() - 1);
+    std::iota(axes.begin(), axes.end(), 0);
+    if (w.ndim() == 0) {
+      vjps.push_back(zeros_like(w, s));
+    } else {
+      vjps.push_back(sum(
+          multiply(g, multiply(x, n, s), s), axes, /* keepdims= */ false, s));
+    }
+
+    return vjps;
+  };
+
+  auto vjps = array::make_arrays(
+      {primals[0].shape(), primals[1].shape()},
+      {primals[0].dtype(), primals[1].dtype()},
+      std::make_shared<RMSNormVJP>(s, fallback, eps_),
+      {primals[0], primals[1], cotangents[0]});
+
+  std::vector<array> returned_vjps;
+  for (auto& arg : argnums) {
+    returned_vjps.push_back(std::move(vjps[arg]));
+  }
+
+  return returned_vjps;
+}
+
+bool RMSNorm::is_equivalent(const Primitive& other) const {
+  const RMSNorm& a_other = static_cast<const RMSNorm&>(other);
+  return eps_ == a_other.eps_;
+}
+
+bool RMSNormVJP::is_equivalent(const Primitive& other) const {
+  const RMSNormVJP& a_other = static_cast<const RMSNormVJP&>(other);
+  return eps_ == a_other.eps_;
+}
+
+array layer_norm(
+    const array& x,
+    const std::optional<array>& weight,
+    const std::optional<array>& bias,
+    float eps,
+    StreamOrDevice s_ /* = {} */) {
+  bool has_weight = weight.has_value();
+  bool has_bias = bias.has_value();
+
+  if (x.ndim() == 0) {
+    std::ostringstream msg;
+    msg << "[layer_norm] Input must have at least 1 dimension but got input with "
+           "0 dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (has_weight) {
+    if ((*weight).ndim() != 1) {
+      std::ostringstream msg;
+      msg << "[layer_norm] weight must have 1 dimension but has "
+          << (*weight).ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if ((*weight).size() != x.shape(-1)) {
+      std::ostringstream msg;
+      msg << "[layer_norm] weight must have the same size as the last dimension of"
+             " x but has "
+          << (*weight).size() << " elements.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  if (has_bias) {
+    if ((*bias).ndim() != 1) {
+      std::ostringstream msg;
+      msg << "[layer_norm] bias must have 1 dimension but has "
+          << (*bias).ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if ((*bias).size() != x.shape(-1)) {
+      std::ostringstream msg;
+      msg << "[layer_norm] bias must have the same size as the last dimension of"
+             " x but has "
+          << (*bias).size() << " elements.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  auto out_type = (has_weight)
+      ? ((has_bias) ? result_type(x, *weight, *bias) : result_type(x, *weight))
+      : x.dtype();
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[layer_norm] Received unsupported type " << out_type << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto s = to_stream(s_);
+  auto fallback = [has_weight, has_bias, eps, out_type, s](
+                      const std::vector<array>& inputs) {
+    auto x = astype(inputs[0], float32, s);
+
+    auto mu = mean(x, /* axis= */ -1, /* keepdims= */ true, s);
+    auto xc = subtract(x, mu, s);
+    auto v = mean(square(xc, s), /* axis= */ -1, /* keepdims= */ true, s);
+
+    x = multiply(xc, rsqrt(add(v, array(eps, float32), s), s));
+    x = astype(x, out_type, s);
+
+    // If the LN is affine then transform x according to the weight and bias
+    if (has_weight) {
+      x = multiply(x, inputs[1], s);
+    }
+    if (has_bias) {
+      x = add(x, inputs[2], s);
+    }
+
+    return std::vector<array>{x};
+  };
+
+  auto passed_weight =
+      (has_weight) ? astype(*weight, out_type, s) : array(1, out_type);
+  auto passed_bias =
+      (has_bias) ? astype(*bias, out_type, s) : array(0, out_type);
+
+  if (!LayerNorm::use_fallback(s)) {
+    return array(
+        x.shape(),
+        out_type,
+        std::make_shared<LayerNorm>(s, fallback, eps),
+        {astype(x, out_type, s), passed_weight, passed_bias});
+  }
+  return fallback({x, passed_weight, passed_bias})[0];
+}
+
+std::vector<array> LayerNorm::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  assert(primals.size() == 3);
+  assert(outputs.size() == 1);
+  assert(cotangents.size() == 1);
+
+  auto s = stream();
+  auto fallback = [eps = eps_, s](const std::vector<array>& inputs) {
+    auto& x = inputs[0];
+    auto& w = inputs[1];
+    auto& b = inputs[2];
+    auto& g = inputs[3];
+
+    std::vector<array> vjps;
+
+    auto norm = number_of_elements(x, {-1}, true, x.dtype(), s);
+    auto sumx = sum(x, /* axis= */ -1, /* keepdims= */ true, s);
+    auto sumx2 = sum(square(x, s), /* axis= */ -1, /* keepdims= */ true, s);
+    auto mu = multiply(sumx, norm, s);
+    auto mu2 = multiply(sumx2, norm, s);
+    auto var = subtract(mu2, square(mu, s), s);
+    auto n = rsqrt(add(var, array(eps, x.dtype()), s));
+    auto n3 = power(n, array(3, x.dtype()), s);
+    auto x_c = subtract(x, mu, s);
+
+    // df/dx
+    auto wg = multiply(w, g, s);
+    auto sumwg =
+        multiply(sum(wg, /* axis= */ -1, /* keepdims= */ true, s), norm, s);
+    auto sumwgxc = multiply(
+        sum(multiply(wg, x_c, s), /* axis= */ -1, /* keepdims= */ true, s),
+        norm,
+        s);
+    auto t1 = multiply(multiply(x_c, sumwgxc, s), n3, s);
+    auto t2 = multiply(subtract(wg, sumwg, s), n, s);
+    vjps.push_back(subtract(t2, t1, s));
+
+    // df/dw
+    std::vector<int> axes(g.ndim() - 1);
+    std::iota(axes.begin(), axes.end(), 0);
+    if (w.ndim() == 0) {
+      vjps.push_back(zeros_like(w, s));
+    } else {
+      vjps.push_back(sum(
+          multiply(g, multiply(x_c, n, s), s), axes, /* keepdims= */ false, s));
+    }
+
+    // df/db
+    if (b.ndim() == 0) {
+      vjps.push_back(zeros_like(w, s));
+    } else {
+      vjps.push_back(sum(g, axes, /* keepdims= */ false, s));
+    }
+
+    return vjps;
+  };
+
+  auto vjps = array::make_arrays(
+      {primals[0].shape(), primals[1].shape(), primals[2].shape()},
+      {primals[0].dtype(), primals[1].dtype(), primals[2].dtype()},
+      std::make_shared<LayerNormVJP>(s, fallback, eps_),
+      {primals[0], primals[1], primals[2], cotangents[0]});
+
+  std::vector<array> returned_vjps;
+  for (auto& arg : argnums) {
+    returned_vjps.push_back(std::move(vjps[arg]));
+  }
+
+  return returned_vjps;
+}
+
+bool LayerNorm::is_equivalent(const Primitive& other) const {
+  const LayerNorm& a_other = static_cast<const LayerNorm&>(other);
+  return eps_ == a_other.eps_;
+}
+
+bool LayerNormVJP::is_equivalent(const Primitive& other) const {
+  const LayerNormVJP& a_other = static_cast<const LayerNormVJP&>(other);
+  return eps_ == a_other.eps_;
+}
+
+array rope(
+    std::vector<array> inputs,
+    int dims,
+    bool traditional,
+    float base,
+    float scale,
+    bool forward,
+    StreamOrDevice s) {
+  auto& x = inputs[0];
+  auto& offset = inputs[1];
+  if (x.ndim() < 3) {
+    std::ostringstream msg;
+    msg << "[rope] Input must have at least 3 dimensions but got input with "
+        << x.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (!issubdtype(x.dtype(), floating)) {
+    std::ostringstream msg;
+    msg << "[rope] Input must be a floating type but got " << x.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (offset.ndim() > 1) {
+    std::ostringstream msg;
+    msg << "[rope] offset must have at most one dimension but has shape "
+        << offset.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (offset.size() != 1 && offset.size() != x.shape(0)) {
+    std::ostringstream msg;
+    msg << "[rope] offset must be a scalar or vector with " << x.shape(0)
+        << " elements but has shape " << offset.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (!issubdtype(offset.dtype(), integer)) {
+    std::ostringstream msg;
+    msg << "[rope] offset must be an integer but got type " << offset.dtype()
+        << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (offset.dtype().size() != 4) {
+    inputs[1] = astype(offset, int32, s);
+  }
+  if (inputs.size() == 3 &&
+      (inputs[2].ndim() != 1 || inputs[2].shape(0) != dims / 2)) {
+    std::ostringstream msg;
+    msg << "[rope] freqs must be one dimensional with size " << dims / 2
+        << " but got shape " << inputs[2].shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto fallback = [dims, traditional, base, scale, forward, s](
+                      std::vector<array> inputs) {
+    auto x = inputs[0];
+    auto shape = x.shape();
+    if (x.ndim() == 3) {
+      x = expand_dims(x, 1, s);
+    } else if (x.ndim() > 4) {
+      x = flatten(x, 1, 1 + (x.ndim() - 4), s);
+    }
+
+    auto B = x.shape(0);
+    auto N = x.shape(1);
+    auto T = x.shape(2);
+    auto t = x.dtype();
+    // Compute sines and cosines
+    auto half_dims = dims / 2;
+    auto offset = inputs[1];
+    if (offset.size() > 1) {
+      offset = expand_dims(offset, {-1, -2}, s);
+    }
+    auto positions = multiply(
+        add(arange(x.shape(2), float32, s), offset, s),
+        array(scale, float32),
+        s);
+
+    auto default_inv_freqs = [&s, base, half_dims]() {
+      return exp(
+          multiply(
+              arange(0, -half_dims, -1, float32, s),
+              array(std::log(base) / half_dims, float32),
+              s),
+          s);
+    };
+
+    auto inv_freqs =
+        inputs.size() == 3 ? reciprocal(inputs[2], s) : default_inv_freqs();
+    auto theta = multiply(expand_dims(positions, -1, s), inv_freqs, s);
+    auto coss = astype(cos(theta, s), t, s);
+    auto sins = astype(sin(theta, s), t, s);
+
+    auto apply_rope = [forward, s](
+                          const array& x1,
+                          const array& x2,
+                          const array& coss,
+                          const array& sins) {
+      std::vector<array> outs;
+      if (forward) {
+        outs.push_back(
+            subtract(multiply(x1, coss, s), multiply(x2, sins, s), s));
+        outs.push_back(add(multiply(x1, sins, s), multiply(x2, coss, s), s));
+      } else {
+        outs.push_back(add(multiply(x2, sins, s), multiply(x1, coss, s), s));
+        outs.push_back(
+            subtract(multiply(x2, coss, s), multiply(x1, sins, s), s));
+      }
+      return outs;
+    };
+
+    if (traditional) {
+      auto x1 = slice(x, {0, 0, 0, 0}, {B, N, T, dims}, {1, 1, 1, 2}, s);
+      auto x2 = slice(x, {0, 0, 0, 1}, {B, N, T, dims}, {1, 1, 1, 2}, s);
+      auto outs = apply_rope(x1, x2, coss, sins);
+      for (auto& o : outs) {
+        o = expand_dims(o, -1, s);
+      }
+      auto out = reshape(concatenate(outs, -1, s), {B, N, T, dims}, s);
+      if (dims < x.shape(-1)) {
+        out =
+            concatenate({out, slice(x, {0, 0, 0, dims}, x.shape(), s)}, -1, s);
+      }
+      return std::vector<array>{reshape(out, shape, s)};
+    } else {
+      auto out_s = x.shape();
+      out_s.back() = half_dims;
+      auto x1 = slice(x, {0, 0, 0, 0}, out_s, s);
+      out_s.back() = dims;
+      auto x2 = slice(x, {0, 0, 0, half_dims}, out_s, s);
+
+      auto outs = apply_rope(x1, x2, coss, sins);
+      if (dims < x.shape(-1)) {
+        outs.push_back(slice(x, {0, 0, 0, dims}, x.shape(), s));
+      }
+      return std::vector<array>{reshape(concatenate(outs, -1, s), shape, s)};
+    }
+  };
+  auto stream = to_stream(s);
+  if (!RoPE::use_fallback(stream)) {
+    return array(
+        x.shape(),
+        x.dtype(),
+        std::make_shared<RoPE>(
+            stream, fallback, dims, traditional, base, scale, forward),
+        std::move(inputs));
+  }
+  return fallback(std::move(inputs))[0];
+}
+
+array rope(
+    const array& x,
+    int dims,
+    bool traditional,
+    std::optional<float> base,
+    float scale,
+    const array& offset,
+    const std::optional<array>& freqs /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  std::vector<array> inputs = {x, offset};
+  if (freqs) {
+    inputs.push_back(astype(*freqs, float32, s));
+    if (base) {
+      throw std::invalid_argument(
+          "[rope] Only one of base or freqs can have a value.");
+    }
+  } else if (!base) {
+    throw std::invalid_argument("[rope] Neither base nor freqs has a value.");
+  }
+  return rope(
+      std::move(inputs),
+      dims,
+      traditional,
+      base.has_value() ? *base : 1.0,
+      scale,
+      true,
+      s);
+}
+
+array rope(
+    const array& x,
+    int dims,
+    bool traditional,
+    std::optional<float> base,
+    float scale,
+    int offset,
+    const std::optional<array>& freqs /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return rope(
+      x, dims, traditional, base, scale, array(offset, int32), freqs, s);
+}
+
+std::vector<array> RoPE::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  auto s = stream();
+  auto fallback = [dims = dims_,
+                   traditional = traditional_,
+                   base = base_,
+                   scale = scale_,
+                   forward = forward_,
+                   s](std::vector<array> inputs) {
+    return std::vector<array>{
+        rope(std::move(inputs), dims, traditional, base, scale, !forward, s)};
+  };
+  if (argnums.size() > 1 || argnums[0] != 0) {
+    throw std::invalid_argument(
+        "[RoPE::vjp] vjp for offset or frequencies not supported");
+  }
+  auto inputs = std::vector<array>{cotangents[0], primals[1]};
+  if (primals.size() == 3) {
+    inputs.push_back(primals[2]);
+  }
+  return {array(
+      cotangents[0].shape(),
+      cotangents[0].dtype(),
+      std::make_shared<RoPE>(
+          s, fallback, dims_, traditional_, base_, scale_, !forward_),
+      std::move(inputs))};
+}
+
+bool RoPE::is_equivalent(const Primitive& other) const {
+  const RoPE& a_other = static_cast<const RoPE&>(other);
+  return (
+      dims_ == a_other.dims_ && base_ == a_other.base_ &&
+      scale_ == a_other.scale_ && traditional_ == a_other.traditional_ &&
+      forward_ == a_other.forward_);
+}
+
+/** Computes: O = softmax(Q @ K.T) @ V **/
+array scaled_dot_product_attention(
+    const array& queries,
+    const array& keys,
+    const array& values,
+    const float scale,
+    const std::string& mask_mode /* = "" */,
+    std::optional<array> mask_arr /* = {} */,
+    const std::optional<array>& sinks /* = {} */,
+    StreamOrDevice s /* = {}*/) {
+  for (const auto& tensor : {queries, keys, values}) {
+    if (tensor.ndim() != 4) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] input with shape "
+          << tensor.shape() << " expected to be rank 4";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  // Check valid mask
+  if (mask_mode != "" && mask_mode != "causal" && mask_mode != "array") {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] Invalid mask_mode " << mask_mode
+        << ". mask_mode must be 'causal', 'array' or ''.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  bool do_causal = false;
+  bool has_mask = false;
+  bool has_arr_mask = false;
+  bool has_bool_mask = false;
+
+  if (mask_mode == "causal") {
+    has_mask = true;
+    do_causal = true;
+
+    if (mask_arr) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] Invalid mask_arr for mask_mode "
+          << "'casusal'. No array mask should be passed.";
+      throw std::invalid_argument(msg.str());
+    }
+  } else if (mask_arr) {
+    has_mask = true;
+    has_arr_mask = true;
+    has_bool_mask = mask_arr->dtype() == bool_;
+  }
+
+  if (has_arr_mask && mask_arr->ndim() > 4) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] the mask with shape "
+        << mask_arr->shape() << " expected to have at most rank 4.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  const size_t batch_dim = queries.shape(0);
+  for (const auto& tensor : {keys, values}) {
+    if (tensor.shape(0) != batch_dim) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] mismatching batch dimension for input with shape "
+          << tensor.shape() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  // Q, K must have matching last dims (d_k aka 'head_dim');
+  if (queries.shape(-1) != keys.shape(-1)) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] query, keys expected to have matching last dimension; found query shape "
+        << queries.shape() << " for keys shape " << keys.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // K, V must have matching number of heads (n_kv_heads);
+  auto n_q_heads = queries.shape(-3);
+  auto n_kv_heads = keys.shape(-3);
+
+  if (keys.shape(-3) != values.shape(-3)) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] keys, values expected to have matching n_kv_heads; found keys with n_heads "
+        << keys.shape(-3) << " for values with n_heads " << values.shape(-3)
+        << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // n_heads % n_kv_heads == 0; n_heads >= 1, n_kv_heads >= 1.
+  if (n_q_heads % n_kv_heads != 0) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] n_heads must be a multiple of n_kv_heads, found n_heads "
+        << n_q_heads << " for n_kv_heads " << n_kv_heads << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto final_type = result_type(queries, keys, values);
+  if (!issubdtype(final_type, floating)) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] Received unsupported type "
+        << final_type << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  bool has_sinks = sinks.has_value();
+
+  auto q = astype(queries, final_type, s);
+  auto k = astype(keys, final_type, s);
+  auto v = astype(values, final_type, s);
+
+  auto fallback = [scale,
+                   n_q_heads,
+                   n_kv_heads,
+                   do_causal,
+                   has_sinks,
+                   has_arr_mask,
+                   s](const std::vector<array>& inputs) {
+    auto q = multiply(array(scale, inputs[0].dtype()), inputs[0], s);
+    int n_repeats = n_q_heads / n_kv_heads;
+    auto k = inputs[1];
+    auto v = inputs[2];
+    if (n_repeats > 1) {
+      q = unflatten(q, 1, {n_kv_heads, n_repeats}, s);
+      k = expand_dims(k, 2, s);
+      v = expand_dims(v, 2, s);
+    }
+    auto scores = matmul(q, swapaxes(k, -1, -2, s), s);
+    if (has_arr_mask || do_causal) {
+      // Mask must be broadcast-compatible with [B, n_q_heads, L_q, L_kv]
+      auto make_or_fetch_mask = [&]() {
+        if (do_causal) {
+          int kL = k.shape(-2);
+          int qL = q.shape(-2);
+          int offset = kL - qL;
+          auto q_idx = arange(offset, qL + offset, s);
+          auto k_idx = arange(0, kL, s);
+          q_idx = expand_dims(q_idx, 1, s);
+          k_idx = expand_dims(k_idx, 0, s);
+          return greater_equal(q_idx, k_idx, s);
+        }
+        return inputs[3];
+      };
+      auto mask = make_or_fetch_mask();
+
+      if (n_repeats > 1 && mask.ndim() >= 3) {
+        if (mask.shape(-3) == 1) {
+          mask = expand_dims(mask, -3, s);
+        } else {
+          mask = unflatten(mask, -3, {n_kv_heads, n_repeats}, s);
+        }
+      }
+      if (mask.dtype() == bool_) {
+        scores = where(
+            mask, scores, array(finfo(scores.dtype()).min, scores.dtype()), s);
+      } else {
+        scores = add(scores, mask, s);
+      }
+    }
+    if (has_sinks) {
+      auto sinks = inputs.back();
+      // scores has shape B N_q N_k L_q L_k
+      sinks = expand_dims(sinks, {0, 2, 3}, s);
+      if (scores.ndim() == 5) {
+        sinks = unflatten(sinks, 1, {n_kv_heads, n_repeats}, s);
+      }
+      auto bsx_shape = scores.shape();
+      bsx_shape.back() = 1;
+      scores = concatenate({broadcast_to(sinks, bsx_shape, s), scores}, -1, s);
+    }
+    scores = softmax(scores, std::vector<int>{-1}, true, s);
+    if (has_sinks) {
+      // Slice off scores
+      auto start = Shape(scores.ndim(), 0);
+      start.back() = 1;
+      auto stop = scores.shape();
+      scores = slice(scores, std::move(start), std::move(stop), s);
+    }
+    auto out = matmul(scores, v, s);
+    if (n_repeats > 1) {
+      out = flatten(out, 1, 2, s);
+    }
+    return std::vector<array>{out};
+  };
+
+  auto stream = to_stream(s);
+  std::vector<array> inputs = {q, k, v};
+  if (has_arr_mask) {
+    // Check type
+    has_bool_mask = mask_arr->dtype() == bool_;
+    if (promote_types(mask_arr->dtype(), final_type) != final_type) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] Mask type must promote to output type "
+          << final_type << ".";
+      throw std::invalid_argument(msg.str());
+    } else if (!has_bool_mask) {
+      mask_arr = astype(*mask_arr, final_type, stream);
+    }
+    // Broadcast mask
+    auto mask_shape = queries.shape();
+    mask_shape.back() = keys.shape(-2);
+    inputs.push_back(broadcast_to(*mask_arr, mask_shape, stream));
+  }
+  if (has_sinks) {
+    if (promote_types(sinks->dtype(), final_type) != final_type) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] Type of sinks must promote to output type "
+          << final_type << ".";
+      throw std::invalid_argument(msg.str());
+    }
+    if (sinks->ndim() != 1 || sinks->shape(0) != n_q_heads) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] Received invalid shape for sinks "
+          << sinks->shape() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+    inputs.push_back(astype(*sinks, final_type, stream));
+  }
+
+  bool is_training = detail::in_grad_tracing();
+  bool has_fast_vjp = !ScaledDotProductAttentionVJP::use_fallback(q, stream);
+  bool output_logsumexp = is_training && has_fast_vjp;
+  if (!ScaledDotProductAttention::use_fallback(
+          q,
+          k,
+          v,
+          has_mask,
+          has_arr_mask,
+          do_causal,
+          is_training,
+          output_logsumexp,
+          stream)) {
+    if (has_bool_mask && !ScaledDotProductAttention::supports_bool_mask()) {
+      // Convert bool mask to additive mask.
+      float inf = std::numeric_limits<float>::infinity();
+      array& mask = inputs[3];
+      mask = where(
+          mask,
+          full_like(mask, 0, final_type, s),
+          full_like(mask, -inf, final_type, s));
+    }
+    Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)};
+    auto primitive = std::make_shared<ScaledDotProductAttention>(
+        stream, fallback, scale, do_causal, has_sinks, output_logsumexp);
+    if (output_logsumexp) {
+      return array::make_arrays(
+          {std::move(out_shape), Shape{q.shape(0), q.shape(1), q.shape(2), 1}},
+          {final_type, float32},
+          primitive,
+          std::move(inputs))[0];
+    } else {
+      return array(
+          std::move(out_shape), final_type, primitive, std::move(inputs));
+    }
+  }
+  return fallback(std::move(inputs))[0];
+}
+
+std::vector<array> ScaledDotProductAttention::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  assert(primals.size() >= 3);
+  assert(cotangents.size() == outputs.size());
+
+  auto s = stream();
+  if (ScaledDotProductAttentionVJP::use_fallback(primals[0], s)) {
+    assert(outputs.size() == 1);
+    return Custom::vjp(primals, cotangents, argnums, outputs);
+  }
+
+  auto fallback = [sdpa = fallback_, s](const std::vector<array>& inputs) {
+    std::vector<array> primals(inputs.begin(), std::prev(inputs.end()));
+    auto [_, vjps] = mlx::core::vjp(sdpa, primals, {inputs.back()});
+    return vjps;
+  };
+
+  std::vector<Shape> shapes;
+  std::vector<Dtype> dtypes;
+  for (int i = 0; i < /* outputs size */ 3; ++i) {
+    shapes.push_back(primals[i].shape());
+    dtypes.push_back(primals[i].dtype());
+  }
+  auto primitive = std::make_shared<ScaledDotProductAttentionVJP>(
+      s, fallback, scale_, do_causal_, has_sinks_);
+  std::vector<array> inputs = primals;
+  inputs.push_back(outputs[0]);
+  inputs.push_back(outputs[1]);
+  inputs.push_back(cotangents[0]);
+  auto vjps = array::make_arrays(std::move(shapes), dtypes, primitive, inputs);
+
+  std::vector<array> returned_vjps;
+  for (int arg : argnums) {
+    if (arg >= 3) {
+      throw std::invalid_argument(
+          "[scale_dot_product_attention] Does not support VJP with respect "
+          " to mask or attention sinks.");
+    }
+    returned_vjps.push_back(std::move(vjps[arg]));
+  }
+  return returned_vjps;
+}
+
+bool ScaledDotProductAttention::is_equivalent(const Primitive& other) const {
+  const ScaledDotProductAttention& a_other =
+      static_cast<const ScaledDotProductAttention&>(other);
+  return scale_ == a_other.scale_ && do_causal_ == a_other.do_causal_ &&
+      has_sinks_ == a_other.has_sinks_ &&
+      output_logsumexp_ == a_other.output_logsumexp_;
+}
+
+bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const {
+  const ScaledDotProductAttentionVJP& a_other =
+      static_cast<const ScaledDotProductAttentionVJP&>(other);
+  return scale_ == a_other.scale_ && do_causal_ == a_other.do_causal_ &&
+      has_sinks_ == a_other.has_sinks_;
+}
+
+bool Quantize::is_equivalent(const Primitive& other) const {
+  const Quantize& p_other = static_cast<const Quantize&>(other);
+  return (
+      p_other.group_size_ == group_size_ && p_other.bits_ == bits_ &&
+      p_other.mode_ == mode_ && p_other.dequantize_ == dequantize_);
+}
+
+std::vector<Shape> Quantize::output_shapes(const std::vector<array>& inputs) {
+  auto& w = inputs[0];
+  if (dequantize_) {
+    auto out_size = w.shape(-1) * 32 / bits_;
+    auto out_shape = w.shape();
+    out_shape.back() = out_size;
+    return {std::move(out_shape)};
+  } else {
+    auto wq_shape = w.shape();
+    wq_shape.back() = w.shape(-1) * bits_ / 32;
+    auto sshape = w.shape();
+    sshape.back() = w.shape(-1) / group_size_;
+    if (inputs.size() == 2) {
+      return {std::move(wq_shape), std::move(sshape)};
+    } else {
+      auto bshape = sshape;
+      return {std::move(wq_shape), std::move(sshape), std::move(bshape)};
+    }
+  }
+}
+
+bool ConvertFP8::is_equivalent(const Primitive& other) const {
+  const ConvertFP8& a_other = static_cast<const ConvertFP8&>(other);
+  return to_fp8_ == a_other.to_fp8_;
+}
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/fast.h b/macos/third_party/mlx/mlx/fast.h
new file mode 100644
index 0000000..1183aba
--- /dev/null
+++ b/macos/third_party/mlx/mlx/fast.h
@@ -0,0 +1,103 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <optional>
+#include <variant>
+
+#include "mlx/api.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fast {
+
+MLX_API array rms_norm(
+    const array& x,
+    const std::optional<array>& weight,
+    float eps,
+    StreamOrDevice s = {});
+
+MLX_API array layer_norm(
+    const array& x,
+    const std::optional<array>& weight,
+    const std::optional<array>& bias,
+    float eps,
+    StreamOrDevice s = {});
+
+MLX_API array rope(
+    const array& x,
+    int dims,
+    bool traditional,
+    std::optional<float> base,
+    float scale,
+    int offset,
+    const std::optional<array>& freqs = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array rope(
+    const array& x,
+    int dims,
+    bool traditional,
+    std::optional<float> base,
+    float scale,
+    const array& offset,
+    const std::optional<array>& freqs = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Computes: O = softmax(Q @ K.T) @ V **/
+MLX_API array scaled_dot_product_attention(
+    const array& queries,
+    const array& keys,
+    const array& values,
+    const float scale,
+    const std::string& mask_mode = "",
+    std::optional<array> mask_arr = {},
+    const std::optional<array>& sinks = {},
+    StreamOrDevice s = {});
+
+using TemplateArg = std::variant<int, bool, Dtype>;
+using ScalarArg = std::variant<bool, int, float>;
+
+using CustomKernelFunction = std::function<std::vector<array>(
+    const std::vector<array>&,
+    const std::vector<Shape>&,
+    const std::vector<Dtype>&,
+    std::tuple<int, int, int>,
+    std::tuple<int, int, int>,
+    std::vector<std::pair<std::string, TemplateArg>>,
+    std::optional<float>,
+    bool,
+    StreamOrDevice)>;
+
+MLX_API CustomKernelFunction metal_kernel(
+    const std::string& name,
+    const std::vector<std::string>& input_names,
+    const std::vector<std::string>& output_names,
+    const std::string& source,
+    const std::string& header = "",
+    bool ensure_row_contiguous = true,
+    bool atomic_outputs = false);
+
+MLX_API CustomKernelFunction cuda_kernel(
+    const std::string& name,
+    const std::vector<std::string>& input_names,
+    const std::vector<std::string>& output_names,
+    const std::string& source,
+    const std::string& header = "",
+    bool ensure_row_contiguous = true,
+    int shared_memory = 0);
+
+MLX_API std::vector<array> precompiled_cuda_kernel(
+    const std::string& name,
+    const std::string& compiled_source,
+    const std::vector<array>& inputs,
+    const std::vector<Shape>& output_shapes,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<ScalarArg>& scalars,
+    std::tuple<int, int, int> grid,
+    std::tuple<int, int, int> threadgroup,
+    int shared_memory = 0,
+    std::optional<float> init_value = std::nullopt,
+    bool ensure_row_contiguous = false,
+    StreamOrDevice s = {});
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/fast_primitives.h b/macos/third_party/mlx/mlx/fast_primitives.h
new file mode 100644
index 0000000..4434830
--- /dev/null
+++ b/macos/third_party/mlx/mlx/fast_primitives.h
@@ -0,0 +1,427 @@
+// Copyright © 2024 Apple Inc.
+
+#include <optional>
+#include <variant>
+
+#include "mlx/primitives.h"
+
+namespace mlx::core::fast {
+
+// Custom primitive accepts a fallback function which it uses for
+// transformations. Transformations are virtual so that derived classes may
+// override the default behavior.
+class Custom : public Primitive {
+ public:
+  explicit Custom(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback)
+      : Primitive(stream), fallback_(std::move(fallback)) {}
+
+  virtual std::pair<std::vector<array>, std::vector<int>> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes) override;
+
+  virtual std::vector<array> jvp(
+      const std::vector<array>& primals,
+      const std::vector<array>& tangents,
+      const std::vector<int>& argnums) override;
+
+  virtual std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+ protected:
+  std::function<std::vector<array>(std::vector<array>)> fallback_;
+};
+
+class RMSNorm : public Custom {
+ public:
+  RMSNorm(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  static bool use_fallback(Stream stream);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(RMSNorm)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class RMSNormVJP : public Custom {
+ public:
+  RMSNormVJP(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(RMSNormVJP)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class LayerNorm : public Custom {
+ public:
+  LayerNorm(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  static bool use_fallback(Stream s);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(LayerNorm)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class LayerNormVJP : public Custom {
+ public:
+  LayerNormVJP(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float eps)
+      : Custom(stream, std::move(fallback)), eps_(eps) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(LayerNormVJP)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_pair(nullptr, eps_);
+  }
+
+ private:
+  float eps_;
+};
+
+class RoPE : public Custom {
+ public:
+  RoPE(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      int dims,
+      bool traditional,
+      float base,
+      float scale,
+      bool forward)
+      : Custom(stream, std::move(fallback)),
+        dims_(dims),
+        traditional_(traditional),
+        base_(base),
+        scale_(scale),
+        forward_(forward) {}
+
+  static bool use_fallback(Stream s);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(RoPE)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return std::make_tuple(
+        nullptr, dims_, traditional_, base_, scale_, forward_);
+  }
+
+ private:
+  int dims_;
+  bool traditional_;
+  float base_;
+  float scale_;
+  bool forward_;
+};
+
+class ScaledDotProductAttention : public Custom {
+ public:
+  ScaledDotProductAttention(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float scale,
+      bool do_causal,
+      bool has_sinks,
+      bool output_logsumexp)
+      : Custom(stream, std::move(fallback)),
+        scale_(scale),
+        do_causal_(do_causal),
+        has_sinks_(has_sinks),
+        output_logsumexp_(output_logsumexp) {}
+
+  static bool use_fallback(
+      const array& q,
+      const array& k,
+      const array& v,
+      bool has_mask,
+      bool has_arr_mask,
+      bool do_causal,
+      bool is_training,
+      bool output_logsumexp,
+      Stream s);
+  static bool supports_bool_mask();
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  bool is_equivalent(const Primitive& other) const override;
+
+  DEFINE_NAME(ScaledDotProductAttention);
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return std::make_tuple(
+        nullptr, scale_, do_causal_, has_sinks_, output_logsumexp_);
+  }
+
+ private:
+  float scale_;
+  bool do_causal_;
+  bool has_sinks_;
+  bool output_logsumexp_;
+};
+
+class ScaledDotProductAttentionVJP : public Custom {
+ public:
+  ScaledDotProductAttentionVJP(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      float scale,
+      bool do_causal,
+      bool has_sinks)
+      : Custom(stream, std::move(fallback)),
+        scale_(scale),
+        do_causal_(do_causal),
+        has_sinks_(has_sinks) {}
+
+  static bool use_fallback(const array& q, Stream s);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("NYI");
+  }
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(ScaledDotProductAttentionVJP);
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(nullptr, scale_, do_causal_, has_sinks_);
+  }
+
+ private:
+  float scale_;
+  bool do_causal_;
+  bool has_sinks_;
+};
+
+class ConvertFP8 : public Primitive {
+ public:
+  explicit ConvertFP8(Stream stream, bool to_fp8)
+      : Primitive(stream), to_fp8_(to_fp8) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  const char* name() const override {
+    if (to_fp8_) {
+      return "ToFP8";
+    } else {
+      return "FromFP8";
+    }
+  }
+  bool state() const {
+    return to_fp8_;
+  };
+
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE();
+
+ private:
+  bool to_fp8_;
+};
+
+class Quantize : public Custom {
+ public:
+  explicit Quantize(
+      Stream stream,
+      std::function<std::vector<array>(std::vector<array>)> fallback,
+      int group_size,
+      int bits,
+      QuantizationMode mode,
+      bool dequantize)
+      : Custom(stream, std::move(fallback)),
+        group_size_(group_size),
+        bits_(bits),
+        mode_(mode),
+        dequantize_(dequantize) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(Quantize);
+
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return std::make_tuple(nullptr, group_size_, bits_, mode_, dequantize_);
+  }
+
+ private:
+  int group_size_;
+  int bits_;
+  QuantizationMode mode_;
+  bool dequantize_;
+};
+
+using ScalarArg = std::variant<bool, int, float>;
+
+class CustomKernel : public Primitive {
+ public:
+  CustomKernel(
+      Stream stream,
+      std::string name,
+      std::string source,
+      std::tuple<int, int, int> grid,
+      std::tuple<int, int, int> threadgroup,
+      std::vector<std::tuple<bool, bool, bool>> shape_infos,
+      bool ensure_row_contiguous,
+      std::optional<float> init_value,
+      std::vector<ScalarArg> scalar_arguments,
+      bool is_precompiled,
+      int shared_memory)
+      : Primitive(stream),
+        name_(std::move(name)),
+        source_(std::move(source)),
+        grid_(grid),
+        threadgroup_(threadgroup),
+        shape_infos_(std::move(shape_infos)),
+        ensure_row_contiguous_(ensure_row_contiguous),
+        init_value_(init_value),
+        scalar_arguments_(std::move(scalar_arguments)),
+        is_precompiled_(is_precompiled),
+        shared_memory_(shared_memory) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override {
+    throw std::runtime_error("Custom kernels only run on GPU.");
+  }
+
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(CustomKernel);
+  auto state() const {
+    return std::make_tuple(
+        name_,
+        source_,
+        grid_,
+        threadgroup_,
+        shape_infos_,
+        ensure_row_contiguous_,
+        init_value_,
+        scalar_arguments_,
+        is_precompiled_,
+        shared_memory_);
+  }
+
+ private:
+  std::string name_;
+  std::string source_;
+  std::tuple<int, int, int> grid_;
+  std::tuple<int, int, int> threadgroup_;
+  std::vector<std::tuple<bool, bool, bool>> shape_infos_;
+  bool ensure_row_contiguous_;
+  std::optional<float> init_value_;
+  std::vector<ScalarArg> scalar_arguments_;
+  bool is_precompiled_;
+  int shared_memory_;
+};
+
+} // namespace mlx::core::fast
diff --git a/macos/third_party/mlx/mlx/fence.h b/macos/third_party/mlx/mlx/fence.h
new file mode 100644
index 0000000..0ececdb
--- /dev/null
+++ b/macos/third_party/mlx/mlx/fence.h
@@ -0,0 +1,39 @@
+// Copyright © 2024 Apple Inc.
+
+#include <vector>
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+/* A fence to be used for synchronizing work between streams.
+ *
+ * Calls to `wait` wait in the given stream until all previous calls to update
+ * are complete on their given stream.
+ *
+ * The array passed to `update` is computed and visible after the call to
+ * `wait` returns. The array passed to `wait` will not be read until all
+ * previous calls to `update` have completed.
+ *
+ * Note, calls to `update` should always be from the same thread or explicitly
+ * synchronized so that they occur in sequence. Calls to `wait` can be on any
+ * thread.
+ *
+ * For the Metal back-end the fence supports slow (default) and fast mode.
+ * Fast mode requires setting the environment variable
+ * `MLX_METAL_FAST_SYNCH=1`. Fast mode also requires Metal 3.2+ (macOS 15+,
+ * iOS 18+).
+ */
+class Fence {
+ public:
+  Fence() {};
+  explicit Fence(Stream stream);
+
+  void update(Stream stream, const array& x, bool cross_device);
+  void wait(Stream stream, const array& x);
+
+ private:
+  std::shared_ptr<void> fence_{nullptr};
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/fft.cpp b/macos/third_party/mlx/mlx/fft.cpp
new file mode 100644
index 0000000..6510fae
--- /dev/null
+++ b/macos/third_party/mlx/mlx/fft.cpp
@@ -0,0 +1,262 @@
+// Copyright © 2023 Apple Inc.
+#include <numeric>
+#include <set>
+
+#include "mlx/fft.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fft {
+
+array fft_impl(
+    const array& a,
+    Shape n,
+    const std::vector<int>& axes,
+    bool real,
+    bool inverse,
+    StreamOrDevice s) {
+  if (a.ndim() < 1) {
+    throw std::invalid_argument(
+        "[fftn] Requires array with at least one dimension.");
+  }
+  if (n.size() != axes.size()) {
+    throw std::invalid_argument("[fftn] Shape and axes have different sizes.");
+  }
+  if (axes.empty()) {
+    return a;
+  }
+
+  std::vector<size_t> valid_axes;
+  for (int ax : axes) {
+    valid_axes.push_back(ax < 0 ? ax + a.ndim() : ax);
+  }
+  std::set<int> unique_axes(valid_axes.begin(), valid_axes.end());
+  if (unique_axes.size() != axes.size()) {
+    std::ostringstream msg;
+    msg << "[fftn] Duplicated axis received " << axes;
+    throw std::invalid_argument(msg.str());
+  }
+  if (*unique_axes.begin() < 0 || *unique_axes.rbegin() >= a.ndim()) {
+    std::ostringstream msg;
+    msg << "[fftn] Invalid axis received for array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // In the following shape manipulations there are three cases to consider:
+  // 1. In a complex to complex transform (fftn / ifftn) the output
+  //    and input shapes are the same.
+  // 2. In a real to complex transform (rfftn) n specifies the input dims
+  //    and the output dims are n[i] / 2 + 1
+  // 3  In a complex to real transform (irfftn) n specifies the output dims
+  //    and the input dims are n[i] / 2 + 1
+
+  if (std::any_of(n.begin(), n.end(), [](auto i) { return i <= 0; })) {
+    std::ostringstream msg;
+    msg << "[fftn] Invalid FFT output size requested " << n;
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto in_shape = a.shape();
+  for (int i = 0; i < valid_axes.size(); ++i) {
+    in_shape[valid_axes[i]] = n[i];
+  }
+  if (real && inverse) {
+    in_shape[valid_axes.back()] = n.back() / 2 + 1;
+  }
+
+  bool any_greater = false;
+  bool any_less = false;
+  for (int i = 0; i < in_shape.size(); ++i) {
+    any_greater |= in_shape[i] > a.shape()[i];
+    any_less |= in_shape[i] < a.shape()[i];
+  }
+
+  auto in = a;
+  if (any_less) {
+    in = slice(in, Shape(in.ndim(), 0), in_shape, s);
+  }
+  if (any_greater) {
+    // Pad with zeros
+    auto tmp = zeros(in_shape, a.dtype(), s);
+    in = slice_update(tmp, in, Shape(in.ndim(), 0), in.shape());
+  }
+
+  auto out_shape = in_shape;
+  if (real) {
+    auto ax = valid_axes.back();
+    out_shape[ax] = inverse ? n.back() : out_shape[ax] / 2 + 1;
+  }
+
+  auto in_type = real && !inverse ? float32 : complex64;
+  auto out_type = real && inverse ? float32 : complex64;
+  return array(
+      out_shape,
+      out_type,
+      std::make_shared<FFT>(to_stream(s), valid_axes, inverse, real),
+      {astype(in, in_type, s)});
+}
+
+array fft_impl(
+    const array& a,
+    const std::vector<int>& axes,
+    bool real,
+    bool inverse,
+    StreamOrDevice s) {
+  Shape n;
+  for (auto ax : axes) {
+    n.push_back(a.shape(ax));
+  }
+  if (real && inverse && a.ndim() > 0) {
+    n.back() = (n.back() - 1) * 2;
+  }
+  return fft_impl(a, n, axes, real, inverse, s);
+}
+
+array fft_impl(const array& a, bool real, bool inverse, StreamOrDevice s) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return fft_impl(a, axes, real, inverse, s);
+}
+
+array fftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, false, false, s);
+}
+array fftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, false, false, s);
+}
+array fftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, false, false, s);
+}
+
+array ifftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, false, true, s);
+}
+array ifftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, false, true, s);
+}
+array ifftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, false, true, s);
+}
+
+array rfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, true, false, s);
+}
+array rfftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, true, false, s);
+}
+array rfftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, true, false, s);
+}
+
+array irfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, true, true, s);
+}
+array irfftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, true, true, s);
+}
+
+array irfftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, true, true, s);
+}
+
+array fftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+
+  Shape shifts;
+  for (int ax : axes) {
+    // Convert negative axes to positive
+    int axis = ax < 0 ? ax + a.ndim() : ax;
+    if (axis < 0 || axis >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[fftshift] Invalid axis " << ax << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    // Match NumPy's implementation
+    shifts.push_back(a.shape(axis) / 2);
+  }
+
+  return roll(a, shifts, axes, s);
+}
+
+array ifftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+
+  Shape shifts;
+  for (int ax : axes) {
+    // Convert negative axes to positive
+    int axis = ax < 0 ? ax + a.ndim() : ax;
+    if (axis < 0 || axis >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[ifftshift] Invalid axis " << ax << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    // Match NumPy's implementation
+    int size = a.shape(axis);
+    shifts.push_back(-(size / 2));
+  }
+
+  return roll(a, shifts, axes, s);
+}
+
+// Default versions that operate on all axes
+array fftshift(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.ndim() < 1) {
+    return a;
+  }
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return fftshift(a, axes, s);
+}
+
+array ifftshift(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.ndim() < 1) {
+    return a;
+  }
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return ifftshift(a, axes, s);
+}
+
+} // namespace mlx::core::fft
diff --git a/macos/third_party/mlx/mlx/fft.h b/macos/third_party/mlx/mlx/fft.h
new file mode 100644
index 0000000..9abf2b1
--- /dev/null
+++ b/macos/third_party/mlx/mlx/fft.h
@@ -0,0 +1,159 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <variant>
+
+#include "array.h"
+#include "device.h"
+#include "mlx/api.h"
+#include "utils.h"
+
+namespace mlx::core::fft {
+
+/** Compute the n-dimensional Fourier Transform. */
+MLX_API array fftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+fftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array fftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the n-dimensional inverse Fourier Transform. */
+MLX_API array ifftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+ifftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array ifftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the one-dimensional Fourier Transform. */
+inline array fft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return fftn(a, {n}, {axis}, s);
+}
+inline array fft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return fftn(a, {axis}, s);
+}
+
+/** Compute the one-dimensional inverse Fourier Transform. */
+inline array ifft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return ifftn(a, {n}, {axis}, s);
+}
+inline array ifft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return ifftn(a, {axis}, s);
+}
+
+/** Compute the two-dimensional Fourier Transform. */
+inline array fft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return fftn(a, n, axes, s);
+}
+inline array fft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return fftn(a, axes, s);
+}
+
+/** Compute the two-dimensional inverse Fourier Transform. */
+inline array ifft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return ifftn(a, n, axes, s);
+}
+inline array ifft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return ifftn(a, axes, s);
+}
+
+/** Compute the n-dimensional Fourier Transform on a real input. */
+MLX_API array rfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+rfftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array rfftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the n-dimensional inverse of `rfftn`. */
+MLX_API array irfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+irfftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array irfftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the one-dimensional Fourier Transform on a real input. */
+inline array rfft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return rfftn(a, {n}, {axis}, s);
+}
+inline array rfft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return rfftn(a, {axis}, s);
+}
+/** Compute the one-dimensional inverse of `rfft`. */
+inline array irfft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return irfftn(a, {n}, {axis}, s);
+}
+inline array irfft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return irfftn(a, {axis}, s);
+}
+
+/** Compute the two-dimensional Fourier Transform on a real input. */
+inline array rfft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return rfftn(a, n, axes, s);
+}
+inline array rfft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return rfftn(a, axes, s);
+}
+
+/** Compute the two-dimensional inverse of `rfft2`. */
+inline array irfft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return irfftn(a, n, axes, s);
+}
+inline array irfft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return irfftn(a, axes, s);
+}
+/** Shift the zero-frequency component to the center of the spectrum. */
+MLX_API array fftshift(const array& a, StreamOrDevice s = {});
+
+/** Shift the zero-frequency component to the center of the spectrum along
+ * specified axes. */
+MLX_API array
+fftshift(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+
+/** The inverse of fftshift. */
+MLX_API array ifftshift(const array& a, StreamOrDevice s = {});
+
+/** The inverse of fftshift along specified axes. */
+MLX_API array
+ifftshift(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+
+} // namespace mlx::core::fft
diff --git a/macos/third_party/mlx/mlx/graph_utils.cpp b/macos/third_party/mlx/mlx/graph_utils.cpp
new file mode 100644
index 0000000..854881b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/graph_utils.cpp
@@ -0,0 +1,175 @@
+// Copyright © 2023 Apple Inc.
+
+#include <functional>
+#include <optional>
+#include <sstream>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/graph_utils.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+const std::string& NodeNamer::get_name(const array& x) {
+  auto it = names.find(x.id());
+  if (it == names.end()) {
+    // Get the next name in the sequence
+    // [A, B, ..., Z, AA, AB, ...]
+    std::vector<char> letters;
+    auto var_num = names.size() + 1;
+    while (var_num > 0) {
+      letters.push_back('A' + (var_num - 1) % 26);
+      var_num = (var_num - 1) / 26;
+    }
+    names.emplace(x.id(), std::string(letters.rbegin(), letters.rend()));
+
+    return get_name(x);
+  }
+  return it->second;
+}
+
+void NodeNamer::set_name(const array& x, std::string n) {
+  names[x.id()] = std::move(n);
+}
+
+void depth_first_traversal(
+    std::function<void(array)> callback,
+    const std::vector<array>& outputs) {
+  std::function<void(const array&)> recurse;
+  std::unordered_set<std::uintptr_t> cache;
+  recurse = [&](const array& x) {
+    auto id = x.id();
+    if (cache.find(id) != cache.end()) {
+      return;
+    }
+    cache.insert(id);
+    for (auto& s : x.siblings()) {
+      cache.insert(s.id());
+    }
+    for (auto& in : x.inputs()) {
+      recurse(in);
+    }
+    callback(x);
+  };
+
+  for (auto& o : outputs) {
+    recurse(o);
+  }
+}
+
+void print_graph(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs) {
+  std::vector<array> tape;
+  std::vector<array> inputs;
+
+  depth_first_traversal(
+      [&](const array& x) {
+        if (x.has_primitive()) {
+          tape.push_back(x);
+        } else {
+          inputs.push_back(x);
+        }
+      },
+      outputs);
+
+  auto print_arrs = [&namer, &os](std::vector<array> arrs) {
+    for (auto& arr : arrs) {
+      os << namer.get_name(arr);
+      os << " [" << arr.shape() << ", " << arr.dtype() << "]";
+      if (&arr != &arrs.back()) {
+        os << ", ";
+      }
+    }
+  };
+
+  os << "Inputs: ";
+  print_arrs(inputs);
+  os << "\nOutputs: ";
+  print_arrs(outputs);
+  os << "\n";
+
+  for (auto& arr : tape) {
+    os << arr.primitive().name();
+    os << " ";
+    print_arrs(arr.inputs());
+    os << " -> ";
+    print_arrs(arr.outputs());
+    os << "\n";
+  }
+}
+
+void export_to_dot(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& nodes) {
+  // Perform one DFS to mark arrays as intermediate if they are used as inputs
+  // to other arrays.
+  std::unordered_set<std::uintptr_t> intermediate_set;
+  depth_first_traversal(
+      [&](const array& x) {
+        // No primitive so it is an input
+        if (!x.has_primitive()) {
+          return;
+        }
+
+        for (auto& a : x.inputs()) {
+          intermediate_set.insert(a.id());
+        }
+      },
+      nodes);
+
+  // Now we got everything we need to make the graph. Arrays can be one of 3
+  // things:
+  //  1. Inputs, when they have no primitive ie are evaluated
+  //  2. Intermediates, when they are the intermediate set
+  //  3. Outputs, if they are not inputs and not intermediates
+
+  os << "digraph {" << std::endl;
+
+  depth_first_traversal(
+      [&](const array& x) {
+        if (!x.has_primitive()) {
+          os << "{ rank=source; \"" << namer.get_name(x) << "\"; }"
+             << std::endl;
+          return;
+        }
+
+        // Node for primitive
+        if (x.has_primitive()) {
+          os << "{ ";
+          os << x.primitive_id();
+          os << " [label =\"";
+          os << x.primitive().name();
+          os << "\", shape=rectangle]";
+          os << "; }" << std::endl;
+          // Arrows to primitive's inputs
+          for (auto& a : x.inputs()) {
+            os << '"' << namer.get_name(a) << "\" -> " << x.primitive_id()
+               << std::endl;
+          }
+        }
+
+        // Point outputs to their primitive
+        for (auto& a : x.outputs()) {
+          os << "{ ";
+          if (intermediate_set.find(a.id()) == intermediate_set.end()) {
+            os << "rank=sink; ";
+          }
+          os << '"' << namer.get_name(a);
+          os << "\"; }" << std::endl;
+          if (x.has_primitive()) {
+            os << x.primitive_id() << " -> \"" << namer.get_name(a) << '"'
+               << std::endl;
+          }
+        }
+      },
+      nodes);
+
+  os << "}";
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/graph_utils.h b/macos/third_party/mlx/mlx/graph_utils.h
new file mode 100644
index 0000000..54297c2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/graph_utils.h
@@ -0,0 +1,67 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+struct MLX_API NodeNamer {
+  std::unordered_map<std::uintptr_t, std::string> names;
+
+  const std::string& get_name(const array& x);
+  void set_name(const array& x, std::string n);
+};
+
+MLX_API void print_graph(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs);
+
+inline void print_graph(std::ostream& os, const std::vector<array>& outputs) {
+  print_graph(os, NodeNamer{}, outputs);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void print_graph(std::ostream& os, Arrays&&... outputs) {
+  print_graph(
+      os, NodeNamer{}, std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void
+print_graph(std::ostream& os, NodeNamer namer, Arrays&&... outputs) {
+  print_graph(
+      os,
+      std::move(namer),
+      std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+MLX_API void export_to_dot(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs);
+
+inline void export_to_dot(std::ostream& os, const std::vector<array>& outputs) {
+  export_to_dot(os, NodeNamer{}, outputs);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void export_to_dot(std::ostream& os, Arrays&&... outputs) {
+  export_to_dot(
+      os, NodeNamer{}, std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void
+export_to_dot(std::ostream& os, NodeNamer namer, Arrays&&... outputs) {
+  export_to_dot(
+      os,
+      std::move(namer),
+      std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io.h b/macos/third_party/mlx/mlx/io.h
new file mode 100644
index 0000000..760f298
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io.h
@@ -0,0 +1,61 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/io/load.h"
+#include "mlx/stream.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+using GGUFMetaData =
+    std::variant<std::monostate, array, std::string, std::vector<std::string>>;
+using GGUFLoad = std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, GGUFMetaData>>;
+using SafetensorsLoad = std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string>>;
+
+/** Save array to out stream in .npy format */
+MLX_API void save(std::shared_ptr<io::Writer> out_stream, array a);
+
+/** Save array to file in .npy format */
+MLX_API void save(std::string file, array a);
+
+/** Load array from reader in .npy format */
+MLX_API array
+load(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s = {});
+
+/** Load array from file in .npy format */
+MLX_API array load(std::string file, StreamOrDevice s = {});
+
+/** Load array map from .safetensors file format */
+MLX_API SafetensorsLoad
+load_safetensors(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s = {});
+MLX_API SafetensorsLoad
+load_safetensors(const std::string& file, StreamOrDevice s = {});
+
+MLX_API void save_safetensors(
+    std::shared_ptr<io::Writer> in_stream,
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string> metadata = {});
+MLX_API void save_safetensors(
+    std::string file,
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string> metadata = {});
+
+/** Load array map and metadata from .gguf file format */
+
+MLX_API GGUFLoad load_gguf(const std::string& file, StreamOrDevice s = {});
+
+MLX_API void save_gguf(
+    std::string file,
+    std::unordered_map<std::string, array> array_map,
+    std::unordered_map<std::string, GGUFMetaData> meta_data = {});
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/CMakeLists.txt b/macos/third_party/mlx/mlx/io/CMakeLists.txt
new file mode 100644
index 0000000..eebd761
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/CMakeLists.txt
@@ -0,0 +1,25 @@
+target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp)
+
+if(MLX_BUILD_SAFETENSORS)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/safetensors.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_safetensors.cpp)
+endif()
+
+if(MLX_BUILD_GGUF)
+  message(STATUS "Downloading gguflib")
+  FetchContent_Declare(
+    gguflib
+    GIT_REPOSITORY https://github.com/antirez/gguf-tools/
+    GIT_TAG 8fa6eb65236618e28fd7710a0fba565f7faa1848)
+  FetchContent_MakeAvailable(gguflib)
+  target_include_directories(mlx
+                             PRIVATE $<BUILD_INTERFACE:${gguflib_SOURCE_DIR}>)
+  add_library(gguflib STATIC ${gguflib_SOURCE_DIR}/fp16.c
+                             ${gguflib_SOURCE_DIR}/gguflib.c)
+  target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:gguflib>)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gguf.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/gguf_quants.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_gguf.cpp)
+endif()
diff --git a/macos/third_party/mlx/mlx/io/gguf.cpp b/macos/third_party/mlx/mlx/io/gguf.cpp
new file mode 100644
index 0000000..206f6fb
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/gguf.cpp
@@ -0,0 +1,470 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstdint>
+#include <cstring>
+#include <fstream>
+#include <numeric>
+
+#include "mlx/io/gguf.h"
+#include "mlx/ops.h"
+
+namespace mlx::core {
+
+// https://github.com/antirez/gguf-tools/blob/af7d88d808a7608a33723fba067036202910acb3/gguflib.h#L102-L108
+constexpr int gguf_array_header_size = 12;
+
+std::optional<uint32_t> dtype_to_gguf_tensor_type(const Dtype& dtype) {
+  switch (dtype) {
+    case float32:
+      return GGUF_TYPE_F32;
+    case float16:
+      return GGUF_TYPE_F16;
+    case int8:
+      return GGUF_TYPE_I8;
+    case int16:
+      return GGUF_TYPE_I16;
+    case int32:
+      return GGUF_TYPE_I32;
+    default:
+      return {};
+  }
+}
+
+std::optional<Dtype> gguf_type_to_dtype(const uint32_t& gguf_type) {
+  switch (gguf_type) {
+    case GGUF_TYPE_F32:
+      return float32;
+    case GGUF_TYPE_F16:
+      return float16;
+    case GGUF_TYPE_I8:
+      return int8;
+    case GGUF_TYPE_I16:
+      return int16;
+    case GGUF_TYPE_I32:
+      return int32;
+    default:
+      return {};
+  }
+}
+
+Shape get_shape(const gguf_tensor& tensor) {
+  Shape shape;
+  // The dimension order in GGML is the reverse of the order used in MLX.
+  for (int i = tensor.ndim - 1; i >= 0; i--) {
+    shape.push_back(tensor.dim[i]);
+  }
+  return shape;
+}
+
+std::tuple<allocator::Buffer, Dtype> extract_tensor_data(gguf_tensor* tensor) {
+  if (tensor == nullptr) {
+    throw std::invalid_argument(
+        "[extract_tensor_data] Input tensor pointer is null.");
+  }
+  std::optional<Dtype> equivalent_dtype = gguf_type_to_dtype(tensor->type);
+  // If there's an equivalent type, we can simply copy.
+  if (equivalent_dtype.has_value()) {
+    if (tensor->weights_data == nullptr) {
+      throw std::runtime_error("[load_gguf] NULL tensor data pointer");
+    }
+    allocator::Buffer buffer = allocator::malloc(tensor->bsize);
+    memcpy(
+        buffer.raw_ptr(),
+        tensor->weights_data,
+        tensor->num_weights * equivalent_dtype.value().size());
+    return {buffer, equivalent_dtype.value()};
+  }
+  // Otherwise, we convert to float16.
+  // TODO: Add other dequantization options.
+  int16_t* data = gguf_tensor_to_f16(tensor);
+  if (data == NULL) {
+    throw std::runtime_error("[load_gguf] gguf_tensor_to_f16 failed");
+  }
+  const size_t new_size = tensor->num_weights * sizeof(int16_t);
+  allocator::Buffer buffer = allocator::malloc(new_size);
+  memcpy(buffer.raw_ptr(), data, new_size);
+  free(data);
+  return {buffer, float16};
+}
+
+void set_mx_value_from_gguf(
+    gguf_ctx* ctx,
+    uint32_t type,
+    gguf_value* val,
+    GGUFMetaData& value) {
+  switch (type) {
+    case GGUF_VALUE_TYPE_UINT8:
+      value = array(val->uint8, uint8);
+      break;
+    case GGUF_VALUE_TYPE_INT8:
+      value = array(val->int8, int8);
+      break;
+    case GGUF_VALUE_TYPE_UINT16:
+      value = array(val->uint16, uint16);
+      break;
+    case GGUF_VALUE_TYPE_INT16:
+      value = array(val->int16, int16);
+      break;
+    case GGUF_VALUE_TYPE_UINT32:
+      value = array(val->uint32, uint32);
+      break;
+    case GGUF_VALUE_TYPE_INT32:
+      value = array(val->int32, int32);
+      break;
+    case GGUF_VALUE_TYPE_UINT64:
+      value = array(val->uint64, uint64);
+      break;
+    case GGUF_VALUE_TYPE_INT64:
+      value = array(val->int64, int64);
+      break;
+    case GGUF_VALUE_TYPE_FLOAT32:
+      value = array(val->float32, float32);
+      break;
+    case GGUF_VALUE_TYPE_BOOL:
+      value = array(val->boolval, bool_);
+      break;
+    case GGUF_VALUE_TYPE_STRING:
+      value =
+          std::string(val->string.string, static_cast<int>(val->string.len));
+      break;
+    case GGUF_VALUE_TYPE_FLOAT64:
+      value = array(val->float64, float32);
+      break;
+    case GGUF_VALUE_TYPE_ARRAY: {
+      ctx->off += gguf_array_header_size; // Skip header
+      char* data = reinterpret_cast<char*>(val) + gguf_array_header_size;
+      auto size = static_cast<int>(val->array.len);
+      if (val->array.type == GGUF_VALUE_TYPE_ARRAY) {
+        throw std::invalid_argument(
+            "[load_gguf] Only supports loading 1-layer of nested arrays.");
+      }
+      switch (val->array.type) {
+        case GGUF_VALUE_TYPE_UINT8:
+          value = array(reinterpret_cast<uint8_t*>(data), {size}, uint8);
+          break;
+        case GGUF_VALUE_TYPE_INT8:
+          value = array(reinterpret_cast<int8_t*>(data), {size}, int8);
+          break;
+        case GGUF_VALUE_TYPE_UINT16:
+          value = array(reinterpret_cast<uint16_t*>(data), {size}, uint16);
+          break;
+        case GGUF_VALUE_TYPE_INT16:
+          value = array(reinterpret_cast<int16_t*>(data), {size}, int16);
+          break;
+        case GGUF_VALUE_TYPE_UINT32:
+          value = array(reinterpret_cast<uint32_t*>(data), {size}, uint32);
+          break;
+        case GGUF_VALUE_TYPE_INT32:
+          value = array(reinterpret_cast<int32_t*>(data), {size}, int32);
+          break;
+        case GGUF_VALUE_TYPE_UINT64:
+          value = array(reinterpret_cast<uint64_t*>(data), {size}, uint64);
+          break;
+        case GGUF_VALUE_TYPE_INT64:
+          value = array(reinterpret_cast<uint64_t*>(data), {size}, int64);
+          break;
+        case GGUF_VALUE_TYPE_FLOAT32:
+          value = array(reinterpret_cast<float*>(data), {size}, float32);
+          break;
+        case GGUF_VALUE_TYPE_BOOL:
+          value = array(reinterpret_cast<bool*>(data), {size}, bool_);
+          break;
+        case GGUF_VALUE_TYPE_STRING: {
+          std::vector<std::string> strs(size);
+          for (auto& str : strs) {
+            auto str_val = reinterpret_cast<gguf_string*>(data);
+            data += (str_val->len + sizeof(gguf_string));
+            str = std::string(str_val->string, static_cast<int>(str_val->len));
+            ctx->off += (str_val->len + sizeof(gguf_string));
+          }
+          value = std::move(strs);
+          break;
+        }
+        case GGUF_VALUE_TYPE_FLOAT64:
+          value = array(reinterpret_cast<double*>(data), {size}, float32);
+          break;
+        default:
+          throw std::runtime_error(
+              "[load_gguf] Multiple levels of nested arrays are not supported.");
+      }
+      break;
+    }
+    default:
+      throw std::runtime_error("[load_gguf] Received unexpected type.");
+      break;
+  }
+  if (type == GGUF_VALUE_TYPE_STRING) {
+    ctx->off += (sizeof(gguf_string) + std::get<std::string>(value).size());
+  } else if (auto pv = std::get_if<array>(&value); pv) {
+    ctx->off += pv->nbytes();
+  }
+}
+
+std::unordered_map<std::string, GGUFMetaData> load_metadata(gguf_ctx* ctx) {
+  std::unordered_map<std::string, GGUFMetaData> metadata;
+  gguf_key key;
+  while (gguf_get_key(ctx, &key)) {
+    std::string key_name = std::string(key.name, key.namelen);
+    auto& val = metadata.insert({key_name, GGUFMetaData{}}).first->second;
+    set_mx_value_from_gguf(ctx, key.type, key.val, val);
+  }
+  return metadata;
+}
+
+std::unordered_map<std::string, array> load_arrays(gguf_ctx* ctx) {
+  std::unordered_map<std::string, array> array_map;
+  gguf_tensor tensor;
+
+  auto check_insert = [](const auto& inserted) {
+    if (!inserted.second) {
+      std::ostringstream msg;
+      msg << "[load_gguf] Duplicate parameter name " << inserted.first->second
+          << " this can happend when loading quantized tensors.";
+      throw std::runtime_error(msg.str());
+    }
+  };
+
+  while (gguf_get_tensor(ctx, &tensor)) {
+    if (tensor.type == GGUF_TYPE_Q4_0 || tensor.type == GGUF_TYPE_Q4_1 ||
+        tensor.type == GGUF_TYPE_Q8_0) {
+      gguf_load_quantized(array_map, tensor);
+    } else {
+      std::string name(tensor.name, tensor.namelen);
+      const auto& [data, dtype] = extract_tensor_data(&tensor);
+      array loaded_array = array(data, get_shape(tensor), dtype);
+      check_insert(array_map.insert({name, loaded_array}));
+    }
+  }
+  return array_map;
+}
+
+GGUFLoad load_gguf(const std::string& file, StreamOrDevice s) {
+  bool exists;
+  {
+    std::ifstream f(file.c_str());
+    exists = f.good();
+  }
+  if (!exists) {
+    throw std::invalid_argument("[load_gguf] Failed to open " + file);
+  }
+
+  std::unique_ptr<gguf_ctx, decltype(&gguf_close)> ctx(
+      gguf_open(file.data()), gguf_close);
+  if (!ctx) {
+    throw std::runtime_error("[load_gguf] gguf_init failed");
+  }
+  auto metadata = load_metadata(ctx.get());
+  auto arrays = load_arrays(ctx.get());
+  return {arrays, metadata};
+}
+
+void append_kv_array(
+    gguf_ctx* ctx,
+    const std::string& key,
+    array& val,
+    uint32_t gguf_type) {
+  if (val.ndim() == 1) {
+    size_t gguf_size = val.nbytes() + gguf_array_header_size;
+    std::vector<char> val_vec(gguf_size);
+    gguf_value* gguf_val = reinterpret_cast<gguf_value*>(val_vec.data());
+    gguf_val->array.type = gguf_type;
+    gguf_val->array.len = val.size();
+    memcpy(
+        val_vec.data() + gguf_array_header_size,
+        val.data<char>(),
+        val.nbytes());
+    gguf_append_kv(
+        ctx,
+        key.c_str(),
+        key.length(),
+        GGUF_VALUE_TYPE_ARRAY,
+        reinterpret_cast<void*>(val_vec.data()),
+        gguf_size);
+  } else {
+    gguf_append_kv(
+        ctx,
+        key.c_str(),
+        key.length(),
+        gguf_type,
+        reinterpret_cast<void*>(val.data<char>()),
+        val.nbytes());
+  }
+}
+
+void save_gguf(
+    std::string file,
+    std::unordered_map<std::string, array> array_map,
+    std::unordered_map<std::string, GGUFMetaData> metadata /* = {} */) {
+  // Add .gguf to file name if it is not there
+  if (file.length() < 5 || file.substr(file.length() - 5, 5) != ".gguf") {
+    file += ".gguf";
+  }
+
+  std::unique_ptr<gguf_ctx, decltype(&gguf_close)> ctx(
+      gguf_create(file.c_str(), GGUF_OVERWRITE), gguf_close);
+  if (!ctx) {
+    throw std::runtime_error("[save_gguf] gguf_create failed");
+  }
+
+  auto string_to_gguf = [](char* dst, const std::string& src) {
+    gguf_string* val = reinterpret_cast<gguf_string*>(dst);
+    val->len = src.length();
+    memcpy(val->string, src.c_str(), src.length());
+  };
+
+  // Save any meta data
+  for (auto& [key, value] : metadata) {
+    if (auto pv = std::get_if<std::string>(&value); pv) {
+      const std::string& str = *pv;
+      size_t size = sizeof(gguf_string) + str.length();
+      std::vector<char> val_vec(size);
+      string_to_gguf(val_vec.data(), str);
+      gguf_append_kv(
+          ctx.get(),
+          key.c_str(),
+          key.length(),
+          GGUF_VALUE_TYPE_STRING,
+          static_cast<void*>(val_vec.data()),
+          size);
+    } else if (auto pv = std::get_if<std::vector<std::string>>(&value); pv) {
+      const auto& str_vec = *pv;
+      auto mem_size = std::accumulate(
+          str_vec.begin(), str_vec.end(), 0, [](size_t accum, const auto& s) {
+            return accum + s.size();
+          });
+      mem_size += str_vec.size() * sizeof(gguf_string) + gguf_array_header_size;
+      std::vector<char> val_vec(mem_size);
+      gguf_value* val = reinterpret_cast<gguf_value*>(val_vec.data());
+      val->array.type = GGUF_VALUE_TYPE_STRING;
+      val->array.len = str_vec.size();
+      auto str_ptr = val_vec.data() + gguf_array_header_size;
+      for (auto& str : str_vec) {
+        string_to_gguf(str_ptr, str);
+        str_ptr += str.length() + sizeof(gguf_string);
+      }
+      gguf_append_kv(
+          ctx.get(),
+          key.c_str(),
+          key.length(),
+          GGUF_VALUE_TYPE_ARRAY,
+          static_cast<void*>(val),
+          mem_size);
+    } else if (auto pv = std::get_if<array>(&value); pv) {
+      array v = *pv;
+      if (v.ndim() > 1) {
+        throw std::runtime_error(
+            "[save_gguf] Cannot save arrays with more than one dimension.");
+      }
+      if (v.size() == 0) {
+        throw std::runtime_error("[save_gguf] Cannot save empty arrays.");
+      }
+
+      eval(v);
+      if (!v.flags().row_contiguous) {
+        v = reshape(flatten(v), v.shape());
+      }
+      if (!v.flags().row_contiguous) {
+        throw std::runtime_error(
+            "[save_gguf] Cannot save non contiguous arrays.");
+      }
+      switch (v.dtype()) {
+        case float32:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_FLOAT32);
+          break;
+        case int64:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_INT64);
+          break;
+        case int32:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_INT32);
+          break;
+        case int16:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_INT16);
+          break;
+        case int8:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_INT8);
+          break;
+        case uint64:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_UINT64);
+          break;
+        case uint32:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_UINT32);
+          break;
+        case uint16:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_UINT16);
+          break;
+        case uint8:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_UINT8);
+          break;
+        case bool_:
+          append_kv_array(ctx.get(), key, v, GGUF_VALUE_TYPE_BOOL);
+          break;
+        default:
+          std::ostringstream msg;
+          msg << "[save_gguf] array type " << v.dtype()
+              << " not support for metadata.";
+          throw std::invalid_argument(msg.str());
+      }
+    } else {
+      throw std::runtime_error(
+          "[save_gguf] Received unexpected type in metadata");
+    }
+  }
+
+  // Tensor offsets are relative to data section, so we start at offset 0.
+  uint64_t tensor_offset = 0;
+
+  // First, append the tensor info
+  for (auto& [key, arr] : array_map) {
+    arr.eval();
+
+    // Try to make it row contiguous
+    if (!arr.flags().row_contiguous) {
+      arr = reshape(flatten(arr), arr.shape());
+      arr.eval();
+    }
+
+    // Has to be row-major now but, check one more time in case
+    // any of the above change in the future
+    if (!arr.flags().row_contiguous) {
+      throw std::invalid_argument(
+          "[save_gguf] can only serialize row-major arrays");
+    }
+
+    tensor_offset += gguf_get_alignment_padding(ctx->alignment, tensor_offset);
+    const std::optional<uint32_t> gguf_type =
+        dtype_to_gguf_tensor_type(arr.dtype());
+    if (!gguf_type.has_value()) {
+      std::ostringstream msg;
+      msg << "[save_gguf] dtype " << arr.dtype() << " is not supported";
+      throw std::runtime_error(msg.str());
+    }
+    const char* tensorname = key.c_str();
+    const uint64_t namelen = key.length();
+    const uint32_t num_dim = arr.ndim();
+    std::vector<uint64_t> dim(num_dim);
+    for (int i = 0; i < num_dim; i++) {
+      dim[i] = arr.shape()[num_dim - 1 - i];
+    }
+    if (!gguf_append_tensor_info(
+            ctx.get(),
+            tensorname,
+            namelen,
+            num_dim,
+            dim.data(),
+            gguf_type.value(),
+            tensor_offset)) {
+      throw std::runtime_error("[save_gguf] gguf_append_tensor_info failed");
+    }
+    tensor_offset += arr.nbytes();
+  }
+
+  // Then, append the tensor weights
+  for (const auto& [key, arr] : array_map) {
+    if (!gguf_append_tensor_data(
+            ctx.get(), (void*)arr.data<void>(), arr.nbytes())) {
+      throw std::runtime_error("[save_gguf] gguf_append_tensor_data failed");
+    }
+  }
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/gguf.h b/macos/third_party/mlx/mlx/io/gguf.h
new file mode 100644
index 0000000..fa5bc45
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/gguf.h
@@ -0,0 +1,20 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+
+#include "mlx/io.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
+
+extern "C" {
+#include <gguflib.h>
+}
+
+namespace mlx::core {
+
+Shape get_shape(const gguf_tensor& tensor);
+void gguf_load_quantized(
+    std::unordered_map<std::string, array>& a,
+    const gguf_tensor& tensor);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/gguf_quants.cpp b/macos/third_party/mlx/mlx/io/gguf_quants.cpp
new file mode 100644
index 0000000..148ed6c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/gguf_quants.cpp
@@ -0,0 +1,164 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cstdint>
+#include <cstring>
+#include <numeric>
+
+#include "mlx/io/gguf.h"
+
+namespace mlx::core {
+
+void unpack_32_4(uint8_t* data, int8_t* dst) {
+  std::fill_n(dst, 16, 0);
+  for (int j = 0; j < 16; ++j) {
+    uint8_t x = (data[j + 2] & 0x0F); // j+2 to skip scale bytes.
+    if (j % 2 != 0) {
+      x <<= 4;
+    }
+    dst[j / 2] += x;
+  }
+  // Last 16 weights are in the higher bits
+  for (int j = 0; j < 16; ++j) {
+    uint8_t x = (data[j + 2] >> 4);
+    if (j % 2 != 0) {
+      x <<= 4;
+    }
+    dst[8 + j / 2] += x;
+  }
+}
+
+// Extracts (weight, scales, biases) from Q4_0 tensors.
+// Data layout is: |16 bit scale|32 x 4bit weights|.
+void extract_q4_0_data(
+    const gguf_tensor& tensor,
+    array& weights_arr,
+    array& scales_arr,
+    array& biases_arr) {
+  const uint64_t bytes_per_block = 18; // 2 bytes scale, 32x0.5 byte weights
+  auto data = static_cast<uint8_t*>(tensor.weights_data);
+  auto weights = weights_arr.data<int8_t>();
+  auto scales = scales_arr.data<float16_t>();
+  auto biases = biases_arr.data<float16_t>();
+  for (int64_t i = 0; i < scales_arr.size(); i++) {
+    scales[i] = *((float16_t*)data);
+    biases[i] = -8 * scales[i];
+    unpack_32_4(data, weights);
+    weights += 16;
+    data += bytes_per_block;
+  }
+}
+
+// Extracts (weight, scales, biases) from Q4_1 tensors.
+// Data layout is: |16 bit scale|16 bit bias|32 x 4bit weights|.
+void extract_q4_1_data(
+    const gguf_tensor& tensor,
+    array& weights_arr,
+    array& scales_arr,
+    array& biases_arr) {
+  const uint64_t bytes_per_block =
+      20; // 2 bytes scale, 2 bytes bias, 32x0.5 byte weights
+  auto data = static_cast<uint8_t*>(tensor.weights_data);
+  auto weights = weights_arr.data<int8_t>();
+  auto scales = scales_arr.data<float16_t>();
+  auto biases = biases_arr.data<float16_t>();
+  for (int64_t i = 0; i < scales_arr.size(); i++) {
+    scales[i] = *((float16_t*)data);
+    biases[i] = *((float16_t*)(data) + 1);
+    unpack_32_4(data, weights);
+    weights += 16;
+    data += bytes_per_block;
+  }
+}
+
+// Extracts (weight, scales, biases) from Q8_0 tensors.
+// Data layout is: |16 bit scale|32 x 8bit weights|.
+void extract_q8_0_data(
+    const gguf_tensor& tensor,
+    array& weights_arr,
+    array& scales_arr,
+    array& biases_arr) {
+  const uint64_t weights_per_block = 32;
+  const uint64_t bytes_per_block = 34; // 2 bytes scale, 32x1 byte weights
+  auto data = static_cast<uint8_t*>(tensor.weights_data);
+  auto weights = weights_arr.data<int8_t>();
+  auto scales = scales_arr.data<float16_t>();
+  auto biases = biases_arr.data<float16_t>();
+  for (int64_t i = 0; i < scales_arr.size(); i++) {
+    uint8_t* block_data = data + i * bytes_per_block;
+    scales[i] = *((float16_t*)block_data);
+    biases[i] = -128 * scales[i];
+    for (int64_t j = 0; j < weights_per_block; ++j) {
+      uint8_t x = block_data[j + 2]; // j+2 to skip the scale bytes.
+      // Original data is in int8_t, so we add a bias of -128 and invert the
+      // first bit.
+      x ^= 1 << 7;
+      weights[i * weights_per_block + j] = x;
+    }
+  }
+}
+
+void gguf_load_quantized(
+    std::unordered_map<std::string, array>& a,
+    const gguf_tensor& tensor) {
+  uint64_t weights_per_byte;
+  if (tensor.type == GGUF_TYPE_Q4_0 || tensor.type == GGUF_TYPE_Q4_1) {
+    weights_per_byte = 2;
+  } else { // tensor.type == GGUF_TYPE_Q8_0
+    weights_per_byte = 1;
+  }
+
+  std::string name(tensor.name, tensor.namelen);
+
+  auto shape = get_shape(tensor);
+  const uint64_t weights_per_block = 32;
+  if (shape[shape.size() - 1] % weights_per_block != 0) {
+    std::ostringstream msg;
+    msg << "[load_gguf] tensor " << name
+        << "has incompatible last dim shape: " << shape[shape.size() - 1];
+    throw std::runtime_error(msg.str());
+  }
+
+  auto weights_shape = shape;
+  weights_shape.back() /= (weights_per_byte * 4);
+  auto w_nbytes = uint32.size() *
+      std::accumulate(weights_shape.begin(),
+                      weights_shape.end(),
+                      1,
+                      std::multiplies<size_t>());
+
+  array weights(allocator::malloc(w_nbytes), std::move(weights_shape), uint32);
+
+  // For scales and bias
+  shape[shape.size() - 1] = shape[shape.size() - 1] / weights_per_block;
+  auto sb_nbytes = float16.size() *
+      std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<size_t>());
+
+  array scales(allocator::malloc(sb_nbytes), shape, float16);
+  array biases(allocator::malloc(sb_nbytes), std::move(shape), float16);
+  if (tensor.type == GGUF_TYPE_Q4_0) {
+    extract_q4_0_data(tensor, weights, scales, biases);
+  } else if (tensor.type == GGUF_TYPE_Q4_1) {
+    extract_q4_1_data(tensor, weights, scales, biases);
+  } else if (tensor.type == GGUF_TYPE_Q8_0) {
+    extract_q8_0_data(tensor, weights, scales, biases);
+  }
+
+  a.emplace(name, std::move(weights));
+
+  auto check_insert = [](const auto& inserted) {
+    if (!inserted.second) {
+      std::ostringstream msg;
+      msg << "[load_gguf] Duplicate parameter name " << inserted.first->second
+          << " this can happend when loading quantized tensors.";
+      throw std::runtime_error(msg.str());
+    }
+  };
+
+  constexpr std::string_view weight_suffix = ".weight";
+  const std::string name_prefix =
+      name.substr(0, name.length() - weight_suffix.length());
+  check_insert(a.emplace(name_prefix + ".scales", std::move(scales)));
+  check_insert(a.emplace(name_prefix + ".biases", std::move(biases)));
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/load.cpp b/macos/third_party/mlx/mlx/io/load.cpp
new file mode 100644
index 0000000..15fc4bc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/load.cpp
@@ -0,0 +1,397 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <algorithm>
+#include <cstring>
+#include <fstream>
+#include <limits>
+#include <sstream>
+
+// Used by pread implementation.
+#ifdef _WIN32
+#include <windows.h>
+#endif // _WIN32
+
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/io.h"
+#include "mlx/io/load.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+// Adapted from
+// https://github.com/angeloskath/supervised-lda/blob/master/include/ldaplusplus/NumpyFormat.hpp
+
+namespace mlx::core {
+
+namespace {
+
+constexpr uint8_t MAGIC[] = {
+    0x93,
+    0x4e,
+    0x55,
+    0x4d,
+    0x50,
+    0x59,
+};
+
+inline bool is_big_endian() {
+  union ByteOrder {
+    int32_t i;
+    uint8_t c[4];
+  };
+  ByteOrder b = {0x01234567};
+
+  return b.c[0] == 0x01;
+}
+
+// Array protocol typestring for Dtype
+std::string dtype_to_array_protocol(const Dtype& t) {
+  std::ostringstream r;
+  if (size_of(t) > 1) {
+    r << (is_big_endian() ? ">" : "<");
+  } else {
+    r << "|";
+  }
+  r << kindof(t) << (int)size_of(t);
+  return r.str();
+}
+
+// Dtype from array protocol type string
+Dtype dtype_from_array_protocol(std::string_view t) {
+  if (t.length() == 2 || t.length() == 3) {
+    std::string_view r = t.length() == 3 ? t.substr(1, 2) : t;
+
+    if (r == "V2") {
+      return bfloat16;
+    }
+
+    uint8_t size = r[1] - '0';
+
+    switch (r[0]) {
+      case 'b': {
+        if (size == 1)
+          return bool_;
+        break;
+      }
+      case 'i': {
+        if (size == 1)
+          return int8;
+        else if (size == 2)
+          return int16;
+        else if (size == 4)
+          return int32;
+        else if (size == 8)
+          return int64;
+        break;
+      }
+      case 'u': {
+        if (size == 1)
+          return uint8;
+        else if (size == 2)
+          return uint16;
+        else if (size == 4)
+          return uint32;
+        else if (size == 8)
+          return uint64;
+        break;
+      }
+      case 'f': {
+        if (size == 2)
+          return float16;
+        else if (size == 4)
+          return float32;
+        else if (size == 8)
+          return float64;
+        break;
+      }
+      case 'c': {
+        if (size == 8)
+          return complex64;
+        break;
+      }
+    }
+  }
+
+  throw std::invalid_argument(
+      "[from_str] Unsupported array protocol type-string: " + std::string(t));
+}
+
+#ifdef _WIN32
+// There is no pread on Windows, emulate it with ReadFile.
+int64_t pread(int fd, void* buf, uint64_t size, uint64_t offset) {
+  HANDLE file = reinterpret_cast<HANDLE>(_get_osfhandle(fd));
+  if (file == INVALID_HANDLE_VALUE) {
+    return -1;
+  }
+
+  OVERLAPPED overlapped = {0};
+  overlapped.Offset = offset & 0xFFFFFFFF;
+  overlapped.OffsetHigh = (offset >> 32) & 0xFFFFFFFF;
+
+  DWORD bytes_read;
+  if (!ReadFile(file, buf, size, &bytes_read, &overlapped)) {
+    if (GetLastError() != ERROR_HANDLE_EOF) {
+      return -1;
+    }
+  }
+
+  return bytes_read;
+}
+#endif
+
+} // namespace
+
+/** Save array to out stream in .npy format */
+void save(std::shared_ptr<io::Writer> out_stream, array a) {
+  ////////////////////////////////////////////////////////
+  // Check array
+
+  a = contiguous(a, true);
+  a.eval();
+
+  if (a.nbytes() == 0) {
+    throw std::invalid_argument("[save] cannot serialize an empty array");
+  }
+
+  ////////////////////////////////////////////////////////
+  // Check file
+  if (!out_stream->good() || !out_stream->is_open()) {
+    throw std::runtime_error("[save] Failed to open " + out_stream->label());
+  }
+
+  ////////////////////////////////////////////////////////
+  // Prepare header
+  std::ostringstream magic_ver_len;
+  magic_ver_len.write(reinterpret_cast<const char*>(MAGIC), 6);
+
+  std::string fortran_order = a.flags().col_contiguous ? "True" : "False";
+  std::ostringstream header;
+  header << "{'descr': '" << dtype_to_array_protocol(a.dtype()) << "',"
+         << " 'fortran_order': " << fortran_order << "," << " 'shape': (";
+  for (auto i : a.shape()) {
+    header << i << ", ";
+  }
+  header << ")}";
+
+  size_t header_len = static_cast<size_t>(header.tellp());
+  bool is_v1 = header_len + 15 < std::numeric_limits<uint16_t>::max();
+
+  // Pad out magic + version + header_len + header + \n to be divisible by 16
+  size_t padding = (6 + 2 + (2 + 2 * is_v1) + header_len + 1) % 16;
+
+  header << std::string(padding, ' ') << '\n';
+
+  if (is_v1) {
+    magic_ver_len << (char)0x01 << (char)0x00;
+
+    uint16_t v1_header_len = header.tellp();
+    const char* len_bytes = reinterpret_cast<const char*>(&v1_header_len);
+
+    if (!is_big_endian()) {
+      magic_ver_len.write(len_bytes, 2);
+    } else {
+      magic_ver_len.write(len_bytes + 1, 1);
+      magic_ver_len.write(len_bytes, 1);
+    }
+  } else {
+    magic_ver_len << (char)0x02 << (char)0x00;
+
+    uint32_t v2_header_len = header.tellp();
+    const char* len_bytes = reinterpret_cast<const char*>(&v2_header_len);
+
+    if (!is_big_endian()) {
+      magic_ver_len.write(len_bytes, 4);
+    } else {
+      magic_ver_len.write(len_bytes + 3, 1);
+      magic_ver_len.write(len_bytes + 2, 1);
+      magic_ver_len.write(len_bytes + 1, 1);
+      magic_ver_len.write(len_bytes, 1);
+    }
+  }
+  ////////////////////////////////////////////////////////
+  // Serialize array
+
+  out_stream->write(magic_ver_len.str().c_str(), magic_ver_len.str().length());
+  out_stream->write(header.str().c_str(), header.str().length());
+  out_stream->write(a.data<char>(), a.nbytes());
+}
+
+/** Save array to file in .npy format */
+void save(std::string file, array a) {
+  // Add .npy to file name if it is not there
+  if (file.length() < 4 || file.substr(file.length() - 4, 4) != ".npy")
+    file += ".npy";
+
+  // Serialize array
+  save(std::make_shared<io::FileWriter>(std::move(file)), a);
+}
+
+/** Load array from reader in .npy format */
+array load(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s) {
+  ////////////////////////////////////////////////////////
+  // Open and check file
+  if (!in_stream->good() || !in_stream->is_open()) {
+    throw std::runtime_error("[load] Failed to open " + in_stream->label());
+  }
+
+  auto stream = cu::is_available() ? to_stream(s) : to_stream(s, Device::cpu);
+
+  ////////////////////////////////////////////////////////
+  // Read header and prepare array details
+
+  // Read and check magic
+  char read_magic_and_ver[8];
+  in_stream->read(read_magic_and_ver, 8);
+  if (std::memcmp(read_magic_and_ver, MAGIC, 6) != 0) {
+    throw std::runtime_error("[load] Invalid header in " + in_stream->label());
+  }
+
+  // Read and check version
+  if (read_magic_and_ver[6] != 1 && read_magic_and_ver[6] != 2) {
+    throw std::runtime_error(
+        "[load] Unsupported npy format version in " + in_stream->label());
+  }
+
+  // Read header len and header
+  int header_len_size = read_magic_and_ver[6] == 1 ? 2 : 4;
+  size_t header_len;
+
+  if (header_len_size == 2) {
+    uint16_t v1_header_len;
+    in_stream->read(reinterpret_cast<char*>(&v1_header_len), header_len_size);
+    header_len = v1_header_len;
+  } else {
+    uint32_t v2_header_len;
+    in_stream->read(reinterpret_cast<char*>(&v2_header_len), header_len_size);
+    header_len = v2_header_len;
+  }
+
+  // Read the header
+  std::vector<char> buffer(header_len + 1);
+  in_stream->read(&buffer[0], header_len);
+  buffer[header_len] = 0;
+  std::string header(buffer.data(), header_len);
+
+  // Read data type from header
+  std::string dtype_str = header.substr(11, 3);
+  bool read_is_big_endian = dtype_str[0] == '>';
+  Dtype dtype = dtype_from_array_protocol(dtype_str);
+
+  // Read contiguity order
+  bool col_contiguous = header.at(34) == 'T';
+
+  // Read array shape from header
+  Shape shape;
+
+  size_t st = header.find_last_of('(') + 1;
+  size_t ed = header.find_last_of(')');
+  std::string shape_str = header.substr(st, ed - st);
+
+  while (!shape_str.empty()) {
+    // Read current number and get position of comma
+    size_t pos;
+    int dim = std::stoi(shape_str, &pos);
+    shape.push_back(dim);
+
+    // Skip the comma and space and read the next number
+    if (pos + 2 <= shape_str.length())
+      shape_str = shape_str.substr(pos + 2);
+    else {
+      shape_str = shape_str.substr(pos);
+      if (!shape_str.empty() && shape_str != " " && shape_str != ",") {
+        throw std::runtime_error(
+            "[load] Unknown error while parsing header in " +
+            in_stream->label());
+      }
+      shape_str = "";
+    }
+  }
+
+  ////////////////////////////////////////////////////////
+  // Build primitive
+
+  size_t offset = 8 + header_len_size + header.length();
+  bool swap_endianness = read_is_big_endian != is_big_endian();
+
+  if (col_contiguous) {
+    std::reverse(shape.begin(), shape.end());
+  }
+  auto loaded_array = array(
+      shape,
+      dtype,
+      std::make_shared<Load>(stream, in_stream, offset, swap_endianness),
+      std::vector<array>{});
+  if (col_contiguous) {
+    loaded_array = transpose(loaded_array, s);
+  }
+
+  return loaded_array;
+}
+
+/** Load array from file in .npy format */
+array load(std::string file, StreamOrDevice s) {
+  return load(std::make_shared<io::ParallelFileReader>(std::move(file)), s);
+}
+
+namespace io {
+
+ThreadPool& thread_pool() {
+  static ThreadPool pool_{4};
+  return pool_;
+}
+
+ThreadPool& ParallelFileReader::thread_pool() {
+  static ThreadPool thread_pool{4};
+  return thread_pool;
+}
+
+void ParallelFileReader::read(char* data, size_t n) {
+  while (n != 0) {
+    auto m = ::read(fd_, data, std::min(n, static_cast<size_t>(INT32_MAX)));
+    if (m <= 0) {
+      std::ostringstream msg;
+      msg << "[read] Unable to read " << n << " bytes from file.";
+      throw std::runtime_error(msg.str());
+    }
+    data += m;
+    n -= m;
+  }
+}
+
+void ParallelFileReader::read(char* data, size_t n, size_t offset) {
+  auto readfn = [fd = fd_](size_t offset, size_t size, char* buffer) -> bool {
+    while (size != 0) {
+      auto m = pread(fd, buffer, size, offset);
+      if (m <= 0) {
+        return false;
+      }
+      buffer += m;
+      size -= m;
+    }
+    return true;
+  };
+  std::vector<std::future<bool>> futs;
+  while (n != 0) {
+    if (n < batch_size_) {
+      if (!readfn(offset, n, data)) {
+        throw std::runtime_error("[read] Unable to read from file.");
+      }
+      break;
+    } else {
+      size_t m = batch_size_;
+      futs.emplace_back(
+          ParallelFileReader::thread_pool().enqueue(readfn, offset, m, data));
+      data += m;
+      n -= m;
+      offset += m;
+    }
+  }
+  for (auto& f : futs) {
+    if (!f.get()) {
+      throw std::runtime_error("[read] Unable to read from file.");
+    }
+  }
+}
+
+} // namespace io
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/load.h b/macos/third_party/mlx/mlx/io/load.h
new file mode 100644
index 0000000..0efcb36
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/load.h
@@ -0,0 +1,175 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <memory>
+#include <sstream>
+
+#include <fcntl.h>
+#ifdef _MSC_VER
+#include <io.h>
+#else
+#include <sys/stat.h>
+#include <unistd.h>
+#endif
+
+#include "mlx/threadpool.h"
+
+// Strictly we need to operate on files in binary mode (to avoid \r getting
+// automatically inserted), but every modern system except for Windows no
+// longer differentiates between binary and text files and for them define
+// the flag as no-op.
+#ifndef O_BINARY
+#define O_BINARY 0
+#endif
+
+namespace mlx::core {
+
+namespace io {
+
+ThreadPool& thread_pool();
+
+class Reader {
+ public:
+  virtual bool is_open() const = 0;
+  virtual bool good() const = 0;
+  virtual size_t tell() = 0; // tellp is non-const in iostream
+  virtual void seek(
+      int64_t off,
+      std::ios_base::seekdir way = std::ios_base::beg) = 0;
+  virtual void read(char* data, size_t n) = 0;
+  virtual void read(char* data, size_t n, size_t offset) = 0;
+  virtual std::string label() const = 0;
+  virtual ~Reader() = default;
+};
+
+class Writer {
+ public:
+  virtual bool is_open() const = 0;
+  virtual bool good() const = 0;
+  virtual size_t tell() = 0;
+  virtual void seek(
+      int64_t off,
+      std::ios_base::seekdir way = std::ios_base::beg) = 0;
+  virtual void write(const char* data, size_t n) = 0;
+  virtual std::string label() const = 0;
+  virtual ~Writer() = default;
+};
+
+class ParallelFileReader : public Reader {
+ public:
+  explicit ParallelFileReader(std::string file_path)
+      : fd_(open(file_path.c_str(), O_RDONLY | O_BINARY)),
+        label_(std::move(file_path)) {}
+
+  ~ParallelFileReader() override {
+    close(fd_);
+  }
+
+  bool is_open() const override {
+    return fd_ > 0;
+  }
+
+  bool good() const override {
+    return is_open();
+  }
+
+  size_t tell() override {
+    return lseek(fd_, 0, SEEK_CUR);
+  }
+
+  // Warning: do not use this function from multiple threads as
+  // it advances the file descriptor
+  void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
+      override {
+    if (way == std::ios_base::beg) {
+      lseek(fd_, off, 0);
+    } else {
+      lseek(fd_, off, SEEK_CUR);
+    }
+  }
+
+  // Warning: do not use this function from multiple threads as
+  // it advances the file descriptor
+  void read(char* data, size_t n) override;
+
+  void read(char* data, size_t n, size_t offset) override;
+
+  std::string label() const override {
+    return "file " + label_;
+  }
+
+ private:
+  static constexpr size_t batch_size_ = 1 << 25;
+  static ThreadPool& thread_pool();
+  int fd_;
+  std::string label_;
+};
+
+class FileWriter : public Writer {
+ public:
+  explicit FileWriter() {}
+  explicit FileWriter(std::string file_path)
+      : fd_(open(
+            file_path.c_str(),
+            O_CREAT | O_WRONLY | O_TRUNC | O_BINARY,
+            0644)),
+        label_(std::move(file_path)) {}
+
+  FileWriter(const FileWriter&) = delete;
+  FileWriter& operator=(const FileWriter&) = delete;
+  FileWriter(FileWriter&& other) {
+    std::swap(fd_, other.fd_);
+  }
+
+  ~FileWriter() override {
+    if (fd_ != 0) {
+      close(fd_);
+    }
+  }
+
+  bool is_open() const override {
+    return fd_ >= 0;
+  }
+
+  bool good() const override {
+    return is_open();
+  }
+
+  size_t tell() override {
+    return lseek(fd_, 0, SEEK_CUR);
+  }
+
+  void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
+      override {
+    if (way == std::ios_base::beg) {
+      lseek(fd_, off, 0);
+    } else {
+      lseek(fd_, off, SEEK_CUR);
+    }
+  }
+
+  void write(const char* data, size_t n) override {
+    while (n != 0) {
+      auto m = ::write(fd_, data, std::min(n, static_cast<size_t>(INT32_MAX)));
+      if (m <= 0) {
+        std::ostringstream msg;
+        msg << "[write] Unable to write " << n << " bytes to file.";
+        throw std::runtime_error(msg.str());
+      }
+      data += m;
+      n -= m;
+    }
+  }
+
+  std::string label() const override {
+    return "file " + label_;
+  }
+
+ private:
+  int fd_{0};
+  std::string label_;
+};
+
+} // namespace io
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/no_gguf.cpp b/macos/third_party/mlx/mlx/io/no_gguf.cpp
new file mode 100644
index 0000000..822d8fc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/no_gguf.cpp
@@ -0,0 +1,20 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/io.h"
+
+namespace mlx::core {
+
+GGUFLoad load_gguf(const std::string&, StreamOrDevice s) {
+  throw std::runtime_error(
+      "[load_gguf] Compile with MLX_BUILD_GGUF=ON to enable GGUF support.");
+}
+
+void save_gguf(
+    std::string,
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, GGUFMetaData>) {
+  throw std::runtime_error(
+      "[save_gguf] Compile with MLX_BUILD_GGUF=ON to enable GGUF support.");
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/no_safetensors.cpp b/macos/third_party/mlx/mlx/io/no_safetensors.cpp
new file mode 100644
index 0000000..949c5d2
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/no_safetensors.cpp
@@ -0,0 +1,37 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/io.h"
+
+namespace mlx::core {
+
+SafetensorsLoad load_safetensors(std::shared_ptr<io::Reader>, StreamOrDevice) {
+  throw std::runtime_error(
+      "[load_safetensors] Compile with MLX_BUILD_SAFETENSORS=ON "
+      "to enable safetensors support.");
+}
+
+SafetensorsLoad load_safetensors(const std::string&, StreamOrDevice) {
+  throw std::runtime_error(
+      "[load_safetensors] Compile with MLX_BUILD_SAFETENSORS=ON "
+      "to enable safetensors support.");
+}
+
+void save_safetensors(
+    std::shared_ptr<io::Writer>,
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string>) {
+  throw std::runtime_error(
+      "[save_safetensors] Compile with MLX_BUILD_SAFETENSORS=ON "
+      "to enable safetensors support.");
+}
+
+void save_safetensors(
+    std::string file,
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string>) {
+  throw std::runtime_error(
+      "[save_safetensors] Compile with MLX_BUILD_SAFETENSORS=ON "
+      "to enable safetensors support.");
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/io/safetensors.cpp b/macos/third_party/mlx/mlx/io/safetensors.cpp
new file mode 100644
index 0000000..b8f91ba
--- /dev/null
+++ b/macos/third_party/mlx/mlx/io/safetensors.cpp
@@ -0,0 +1,234 @@
+// Copyright © 2023 Apple Inc.
+//
+#include <json.hpp>
+#include <memory>
+#include <stack>
+
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/io.h"
+#include "mlx/io/load.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+
+using json = nlohmann::json;
+
+#define ST_F16 "F16"
+#define ST_BF16 "BF16"
+#define ST_F32 "F32"
+
+#define ST_BOOL "BOOL"
+#define ST_I8 "I8"
+#define ST_I16 "I16"
+#define ST_I32 "I32"
+#define ST_I64 "I64"
+#define ST_U8 "U8"
+#define ST_U16 "U16"
+#define ST_U32 "U32"
+#define ST_U64 "U64"
+#define ST_F8_E4M3 "F8_E4M3"
+
+// Note: Complex numbers aren't in the spec yet so this could change -
+// https://github.com/huggingface/safetensors/issues/389
+#define ST_C64 "C64"
+
+namespace mlx::core {
+
+std::string dtype_to_safetensor_str(Dtype t) {
+  switch (t) {
+    case float32:
+      return ST_F32;
+    case bfloat16:
+      return ST_BF16;
+    case float16:
+      return ST_F16;
+    case int64:
+      return ST_I64;
+    case int32:
+      return ST_I32;
+    case int16:
+      return ST_I16;
+    case int8:
+      return ST_I8;
+    case uint64:
+      return ST_U64;
+    case uint32:
+      return ST_U32;
+    case uint16:
+      return ST_U16;
+    case uint8:
+      return ST_U8;
+    case bool_:
+      return ST_BOOL;
+    case complex64:
+      return ST_C64;
+    default:
+      throw std::runtime_error("[save_safetensors] received invalid dtype.");
+  }
+}
+
+Dtype dtype_from_safetensor_str(std::string_view str) {
+  if (str == ST_F32) {
+    return float32;
+  } else if (str == ST_F16) {
+    return float16;
+  } else if (str == ST_BF16) {
+    return bfloat16;
+  } else if (str == ST_I64) {
+    return int64;
+  } else if (str == ST_I32) {
+    return int32;
+  } else if (str == ST_I16) {
+    return int16;
+  } else if (str == ST_I8) {
+    return int8;
+  } else if (str == ST_U64) {
+    return uint64;
+  } else if (str == ST_U32) {
+    return uint32;
+  } else if (str == ST_U16) {
+    return uint16;
+  } else if (str == ST_U8) {
+    return uint8;
+  } else if (str == ST_BOOL) {
+    return bool_;
+  } else if (str == ST_C64) {
+    return complex64;
+  } else if (str == ST_F8_E4M3) {
+    return uint8;
+  } else {
+    throw std::runtime_error(
+        "[safetensor] unsupported dtype " + std::string(str));
+  }
+}
+
+/** Load array from reader in safetensor format */
+SafetensorsLoad load_safetensors(
+    std::shared_ptr<io::Reader> in_stream,
+    StreamOrDevice s) {
+  ////////////////////////////////////////////////////////
+  // Open and check file
+  if (!in_stream->good() || !in_stream->is_open()) {
+    throw std::runtime_error(
+        "[load_safetensors] Failed to open " + in_stream->label());
+  }
+
+  auto stream = cu::is_available() ? to_stream(s) : to_stream(s, Device::cpu);
+
+  uint64_t jsonHeaderLength = 0;
+  // This is the same limit as in the original Rust Safetensors code.
+  constexpr uint64_t kMaxJsonHeaderLength = 100000000;
+  in_stream->read(reinterpret_cast<char*>(&jsonHeaderLength), 8);
+  if (jsonHeaderLength <= 0 || jsonHeaderLength >= kMaxJsonHeaderLength) {
+    throw std::runtime_error(
+        "[load_safetensors] Invalid json header length " + in_stream->label());
+  }
+  // Load the json metadata
+  auto rawJson = std::make_unique<char[]>(jsonHeaderLength);
+  in_stream->read(rawJson.get(), jsonHeaderLength);
+  auto metadata = json::parse(rawJson.get(), rawJson.get() + jsonHeaderLength);
+  // Should always be an object on the top-level
+  if (!metadata.is_object()) {
+    throw std::runtime_error(
+        "[load_safetensors] Invalid json metadata " + in_stream->label());
+  }
+  size_t offset = jsonHeaderLength + 8;
+  // Load the arrays using metadata
+  std::unordered_map<std::string, array> res;
+  std::unordered_map<std::string, std::string> metadata_map;
+  for (const auto& item : metadata.items()) {
+    if (item.key() == "__metadata__") {
+      for (const auto& meta_item : item.value().items()) {
+        metadata_map.insert({meta_item.key(), meta_item.value()});
+      }
+      continue;
+    }
+    const std::string& dtype = item.value().at("dtype");
+    const Shape& shape = item.value().at("shape");
+    const std::vector<size_t>& data_offsets = item.value().at("data_offsets");
+    Dtype type = dtype_from_safetensor_str(dtype);
+    res.insert(
+        {item.key(),
+         array(
+             shape,
+             type,
+             std::make_shared<Load>(
+                 stream, in_stream, offset + data_offsets.at(0), false),
+             std::vector<array>{})});
+  }
+  return {res, metadata_map};
+}
+
+SafetensorsLoad load_safetensors(const std::string& file, StreamOrDevice s) {
+  return load_safetensors(std::make_shared<io::ParallelFileReader>(file), s);
+}
+
+void save_safetensors(
+    std::shared_ptr<io::Writer> out_stream,
+    std::unordered_map<std::string, array> a,
+    std::unordered_map<std::string, std::string> metadata /* = {} */) {
+  ////////////////////////////////////////////////////////
+  // Check file
+  if (!out_stream->good() || !out_stream->is_open()) {
+    throw std::runtime_error(
+        "[save_safetensors] Failed to open " + out_stream->label());
+  }
+
+  ////////////////////////////////////////////////////////
+  // Check array map
+  json parent;
+  json _metadata;
+  for (auto& [key, value] : metadata) {
+    _metadata[key] = value;
+  }
+  parent["__metadata__"] = _metadata;
+
+  {
+    std::vector<array> to_eval;
+    to_eval.reserve(a.size());
+    for (auto& p : a) {
+      p.second = contiguous(p.second);
+      to_eval.push_back(p.second);
+    }
+    eval(std::move(to_eval));
+  }
+
+  size_t offset = 0;
+  for (auto& [key, arr] : a) {
+    if (arr.nbytes() == 0) {
+      throw std::invalid_argument(
+          "[save_safetensors] cannot serialize an empty array key: " + key);
+    }
+
+    json child;
+    child["dtype"] = dtype_to_safetensor_str(arr.dtype());
+    child["shape"] = arr.shape();
+    child["data_offsets"] = std::vector<size_t>{offset, offset + arr.nbytes()};
+    parent[key] = child;
+    offset += arr.nbytes();
+  }
+
+  auto header = parent.dump();
+  uint64_t header_len = header.length();
+  out_stream->write(reinterpret_cast<char*>(&header_len), 8);
+  out_stream->write(header.c_str(), header_len);
+  for (auto& [key, arr] : a) {
+    out_stream->write(arr.data<char>(), arr.nbytes());
+  }
+}
+
+void save_safetensors(
+    std::string file,
+    std::unordered_map<std::string, array> a,
+    std::unordered_map<std::string, std::string> metadata /* = {} */) {
+  // Add .safetensors to file name if it is not there
+  if (file.length() < 12 ||
+      file.substr(file.length() - 12, 12) != ".safetensors")
+    file += ".safetensors";
+
+  // Serialize array
+  save_safetensors(
+      std::make_shared<io::FileWriter>(std::move(file)), a, metadata);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/linalg.cpp b/macos/third_party/mlx/mlx/linalg.cpp
new file mode 100644
index 0000000..7ac080d
--- /dev/null
+++ b/macos/third_party/mlx/mlx/linalg.cpp
@@ -0,0 +1,708 @@
+// Copyright © 2023 Apple Inc.
+
+#include <numeric>
+#include <ostream>
+#include <vector>
+
+#include "mlx/linalg.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::linalg {
+
+void check_cpu_stream(const StreamOrDevice& s, const std::string& prefix) {
+  if (to_stream(s).device == Device::gpu) {
+    throw std::invalid_argument(
+        prefix +
+        " This op is not yet supported on the GPU. "
+        "Explicitly pass a CPU stream to run it.");
+  }
+}
+void check_float(Dtype dtype, const std::string& prefix) {
+  if (dtype != float32 && dtype != float64) {
+    std::ostringstream msg;
+    msg << prefix << " Arrays must have type float32 or float64. "
+        << "Received array with type " << dtype << ".";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+void check_float_or_complex(Dtype dtype, const std::string& prefix) {
+  if (dtype != float32 && dtype != float64 && dtype != complex64) {
+    std::ostringstream msg;
+    msg << prefix << " Arrays must have type float32, float64 or complex64. "
+        << "Received array with type " << dtype << ".";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+Dtype at_least_float(const Dtype& d) {
+  return issubdtype(d, inexact) ? d : promote_types(d, float32);
+}
+
+inline array l2_norm(
+    const array& a,
+    const std::vector<int>& axis,
+    bool keepdims,
+    StreamOrDevice s) {
+  if (issubdtype(a.dtype(), complexfloating)) {
+    return sqrt(sum(abs(a, s) * abs(a, s), axis, keepdims, s), s);
+  } else {
+    return sqrt(sum(square(a, s), axis, keepdims, s), s);
+  }
+}
+
+inline array vector_norm(
+    const array& a,
+    const double ord,
+    const std::vector<int>& axis,
+    bool keepdims,
+    StreamOrDevice s) {
+  auto dtype = at_least_float(a.dtype());
+  if (ord == 0.0) {
+    return astype(sum(not_equal(a, array(0), s), axis, keepdims, s), dtype, s);
+  } else if (ord == 1.0) {
+    return astype(sum(abs(a, s), axis, keepdims, s), dtype, s);
+  } else if (ord == 2.0) {
+    return l2_norm(a, axis, keepdims, s);
+  } else if (ord == std::numeric_limits<double>::infinity()) {
+    return astype(max(abs(a, s), axis, keepdims, s), dtype, s);
+  } else if (ord == -std::numeric_limits<double>::infinity()) {
+    return astype(min(abs(a, s), axis, keepdims, s), dtype, s);
+  } else {
+    return power(
+        sum(power(abs(a, s), array(ord, dtype), s), axis, keepdims, s),
+        array(1.0 / ord, dtype),
+        s);
+  }
+}
+
+inline array matrix_norm(
+    const array& a,
+    const double ord,
+    const std::vector<int>& axis,
+    bool keepdims,
+    StreamOrDevice s) {
+  auto dtype = at_least_float(a.dtype());
+  auto row_axis = axis[0];
+  auto col_axis = axis[1];
+  if (ord == -1.0) {
+    col_axis -= (!keepdims && col_axis > row_axis && col_axis > 0);
+    return astype(
+        min(sum(abs(a, s), row_axis, keepdims, s), col_axis, keepdims, s),
+        dtype,
+        s);
+  } else if (ord == 1.0) {
+    col_axis -= (!keepdims && col_axis > row_axis && col_axis > 0);
+    return astype(
+        max(sum(abs(a, s), row_axis, keepdims, s), col_axis, keepdims, s),
+        dtype,
+        s);
+  } else if (ord == std::numeric_limits<double>::infinity()) {
+    row_axis -= (!keepdims && row_axis > col_axis && row_axis > 0);
+    return astype(
+        max(sum(abs(a, s), col_axis, keepdims, s), row_axis, keepdims, s),
+        dtype,
+        s);
+  } else if (ord == -std::numeric_limits<double>::infinity()) {
+    row_axis -= (!keepdims && row_axis > col_axis && row_axis > 0);
+    return astype(
+        min(sum(abs(a, s), col_axis, keepdims, s), row_axis, keepdims, s),
+        dtype,
+        s);
+  } else if (ord == 2.0 || ord == -2.0) {
+    row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0];
+    col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1];
+    auto a_matrix = (row_axis > col_axis)
+        ? moveaxis(moveaxis(a, row_axis, -1, s), col_axis, -1, s)
+        : moveaxis(moveaxis(a, col_axis, -1, s), row_axis, -2, s);
+    a_matrix = svd(a_matrix, false, s).at(0);
+    a_matrix = (ord == 2.0) ? max(a_matrix, -1, false, s)
+                            : min(a_matrix, -1, false, s);
+    if (keepdims) {
+      std::vector<int> sorted_axes = (row_axis < col_axis)
+          ? std::vector<int>{row_axis, col_axis}
+          : std::vector<int>{col_axis, row_axis};
+      a_matrix = expand_dims(a_matrix, sorted_axes, s);
+    }
+    return astype(a_matrix, dtype, s);
+  } else {
+    std::ostringstream msg;
+    msg << "[linalg::norm] Invalid ord " << ord << " for matrix norm.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+inline array matrix_norm(
+    const array& a,
+    const std::string& ord,
+    const std::vector<int>& axis,
+    bool keepdims,
+    StreamOrDevice s) {
+  if (ord == "f" || ord == "fro") {
+    return l2_norm(a, axis, keepdims, s);
+  } else if (ord == "nuc") {
+    int row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0];
+    int col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1];
+    auto a_matrix = (row_axis > col_axis)
+        ? moveaxis(moveaxis(a, row_axis, -1, s), col_axis, -1, s)
+        : moveaxis(moveaxis(a, col_axis, -1, s), row_axis, -2, s);
+    a_matrix = sum(svd(a_matrix, false, s).at(0), -1, false, s);
+    if (keepdims) {
+      std::vector<int> sorted_axes = (row_axis < col_axis)
+          ? std::vector<int>{row_axis, col_axis}
+          : std::vector<int>{col_axis, row_axis};
+      a_matrix = expand_dims(a_matrix, sorted_axes, s);
+    }
+    return a_matrix;
+  } else {
+    std::ostringstream msg;
+    msg << "[linalg::norm] Invalid ord value '" << ord << "' for matrix norm.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+array norm(
+    const array& a,
+    const std::optional<std::vector<int>>& axis /* = std::nullopt */,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  if (!axis) {
+    return norm(flatten(a, s), std::vector<int>{0}, keepdims, s);
+  }
+
+  if (axis.value().size() > 2) {
+    throw std::invalid_argument(
+        "[linalg::norm] Received too many axes for norm.");
+  }
+  return l2_norm(a, axis.value(), keepdims, s);
+}
+
+array norm(
+    const array& a,
+    const double ord,
+    const std::optional<std::vector<int>>& axis /* = std::nullopt */,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  std::vector<int> ax;
+  if (!axis) {
+    ax.resize(a.ndim());
+    std::iota(ax.begin(), ax.end(), 0);
+  } else {
+    ax = axis.value();
+  }
+  if (ax.size() == 1) {
+    return vector_norm(a, ord, ax, keepdims, s);
+  } else if (ax.size() == 2) {
+    return matrix_norm(a, ord, ax, keepdims, s);
+  } else {
+    throw std::invalid_argument(
+        "[linalg::norm] Received too many axes for norm.");
+  }
+}
+
+array norm(
+    const array& a,
+    const std::string& ord,
+    const std::optional<std::vector<int>>& axis /* = std::nullopt */,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  std::vector<int> ax;
+  if (!axis) {
+    ax.resize(a.ndim());
+    std::iota(ax.begin(), ax.end(), 0);
+  } else {
+    ax = axis.value();
+  }
+  if (ax.size() != 2) {
+    std::ostringstream msg;
+    msg << "[linalg::norm] Norm '" << ord << "' only supported for matrices,"
+        << " but received " << ax.size() << " axis/axes.";
+    throw std::invalid_argument(msg.str());
+  }
+  return matrix_norm(a, ord, ax, keepdims, s);
+}
+
+std::pair<array, array> qr(const array& a, StreamOrDevice s /* = {} */) {
+  check_cpu_stream(s, "[linalg::qr]");
+  check_float(a.dtype(), "[linalg::qr]");
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::qr] Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  int k = std::min(a.shape(-2), a.shape(-1));
+  auto q_shape = a.shape();
+  q_shape.back() = k;
+  auto r_shape = a.shape();
+  r_shape[r_shape.size() - 2] = k;
+  auto out = array::make_arrays(
+      {std::move(q_shape), std::move(r_shape)},
+      {a.dtype(), a.dtype()},
+      std::make_shared<QRF>(to_stream(s)),
+      {astype(a, a.dtype(), s)});
+  return std::make_pair(out[0], out[1]);
+}
+
+std::vector<array>
+svd(const array& a, bool compute_uv, StreamOrDevice s /* = {} */) {
+  check_cpu_stream(s, "[linalg::svd]");
+  check_float_or_complex(a.dtype(), "[linalg::svd]");
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::svd] Input array must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  const auto m = a.shape(-2);
+  const auto n = a.shape(-1);
+  const auto rank = a.ndim();
+
+  auto s_shape = a.shape();
+  s_shape.pop_back();
+  s_shape[rank - 2] = std::min(m, n);
+
+  auto s_dtype = a.dtype() == complex64 ? float32 : a.dtype();
+
+  if (!compute_uv) {
+    return {array(
+        std::move(s_shape),
+        s_dtype,
+        std::make_shared<SVD>(to_stream(s), compute_uv),
+        {a})};
+  }
+
+  auto u_shape = a.shape();
+  u_shape[rank - 2] = m;
+  u_shape[rank - 1] = m;
+
+  auto vt_shape = a.shape();
+  vt_shape[rank - 2] = n;
+  vt_shape[rank - 1] = n;
+
+  return array::make_arrays(
+      {u_shape, s_shape, vt_shape},
+      {a.dtype(), s_dtype, a.dtype()},
+      std::make_shared<SVD>(to_stream(s), compute_uv),
+      {a});
+}
+
+array inv_impl(const array& a, bool tri, bool upper, StreamOrDevice s) {
+  check_cpu_stream(s, "[linalg::inv]");
+  check_float(a.dtype(), "[linalg::inv]");
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::inv] Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (a.shape(-1) != a.shape(-2)) {
+    throw std::invalid_argument(
+        "[linalg::inv] Inverses are only defined for square matrices.");
+  }
+
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Inverse>(to_stream(s), tri, upper),
+      {a});
+}
+
+array inv(const array& a, StreamOrDevice s /* = {} */) {
+  return inv_impl(a, /*tri=*/false, /*upper=*/true, s);
+}
+
+array tri_inv(
+    const array& a,
+    bool upper /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return inv_impl(a, /*tri=*/true, upper, s);
+}
+
+array cholesky(
+    const array& a,
+    bool upper /* = false */,
+    StreamOrDevice s /* = {} */) {
+  check_cpu_stream(s, "[linalg::cholesky]");
+  check_float(a.dtype(), "[linalg::cholesky]");
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::cholesky] Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.shape(-1) != a.shape(-2)) {
+    throw std::invalid_argument(
+        "[linalg::cholesky] Cholesky decomposition is only defined for square "
+        "matrices.");
+  }
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Cholesky>(to_stream(s), upper),
+      {a});
+}
+
+array pinv(const array& a, StreamOrDevice s /* = {} */) {
+  check_cpu_stream(s, "[linalg::pinv]");
+  check_float(a.dtype(), "[linalg::pinv]");
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::pinv] Arrays must have >= 2 dimensions. Received array "
+        << "with " << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  int m = a.shape(-2);
+  int n = a.shape(-1);
+  int k = std::min(m, n);
+  auto outs = linalg::svd(a, true, s);
+  array U = outs[0];
+  array S = outs[1];
+  array V = outs[2];
+
+  Shape starts(a.ndim(), 0);
+  auto ends = a.shape();
+  int i = a.ndim() - 2;
+  int j = a.ndim() - 1;
+
+  // Prepare U
+  ends[i] = m;
+  ends[j] = k;
+  U = swapaxes(slice(U, starts, ends, s), -1, -2, s);
+
+  // Prepare V
+  ends[i] = k;
+  ends[j] = n;
+  V = swapaxes(slice(V, starts, ends, s), -1, -2, s);
+
+  // Prepare S
+  S = expand_dims(S, -2, s);
+
+  auto rcond = 10. * std::max(m, n) * finfo(a.dtype()).eps;
+  auto cutoff = multiply(array(rcond, a.dtype()), max(S, -1, true, s), s);
+  auto rS =
+      where(greater(S, cutoff, s), reciprocal(S, s), array(0.0f, a.dtype()), s);
+
+  return matmul(multiply(V, rS, s), U, s);
+}
+
+array cholesky_inv(
+    const array& L,
+    bool upper /* = false */,
+    StreamOrDevice s /* = {} */) {
+  check_cpu_stream(s, "[linalg::cholesky_inv]");
+  check_float(L.dtype(), "[linalg::cholesky_inv]");
+
+  if (L.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[linalg::cholesky_inv] Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << L.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (L.shape(-1) != L.shape(-2)) {
+    throw std::invalid_argument(
+        "[linalg::cholesky_inv] Cholesky inverse is only defined for square "
+        "matrices.");
+  }
+
+  array L_inv = tri_inv(L, upper, s);
+  if (upper) {
+    return matmul(L_inv, swapaxes(L_inv, -1, -2, s), s);
+  } else {
+    return matmul(swapaxes(L_inv, -1, -2, s), L_inv, s);
+  }
+}
+
+array cross(
+    const array& a,
+    const array& b,
+    int axis /* = -1 */,
+    StreamOrDevice s /* = {} */) {
+  auto check_ax = [axis](const array& arr) {
+    if (axis >= static_cast<int>(arr.ndim()) || axis + arr.ndim() < 0) {
+      std::ostringstream msg;
+      msg << "[linalg::cross] axis " << axis << " invalid for array with "
+          << arr.ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if (arr.shape(axis) < 2 || arr.shape(axis) > 3) {
+      throw std::invalid_argument(
+          "[linalg::cross] The specified axis must have size 2 or 3.");
+    }
+  };
+  check_ax(a);
+  check_ax(b);
+
+  bool a_2d = a.shape(axis) == 2;
+  bool b_2d = b.shape(axis) == 2;
+
+  auto out_type = promote_types(a.dtype(), b.dtype());
+  auto ashape = a.shape();
+  auto bshape = b.shape();
+
+  ashape[axis < 0 ? axis + a.ndim() : axis] = 3;
+  bshape[axis < 0 ? axis + b.ndim() : axis] = 3;
+  auto out_shape = broadcast_shapes(ashape, bshape);
+
+  if (axis < 0) {
+    axis += out_shape.size();
+  }
+
+  out_shape[axis] = a_2d ? 2 : 3;
+  auto a_ = broadcast_to(astype(a, out_type, s), out_shape, s);
+
+  out_shape[axis] = b_2d ? 2 : 3;
+  auto b_ = broadcast_to(astype(b, out_type, s), out_shape, s);
+
+  auto a_splits = split(a_, a_2d ? 2 : 3, axis);
+  auto b_splits = split(b_, b_2d ? 2 : 3, axis);
+
+  std::vector<array> outputs;
+  if (a_2d && b_2d) {
+    auto z = zeros_like(a_splits[0], s);
+    outputs.push_back(z);
+    outputs.push_back(z);
+  } else if (b_2d) {
+    outputs.push_back(negative(multiply(a_splits[2], b_splits[1], s), s));
+    outputs.push_back(multiply(a_splits[2], b_splits[0], s));
+  } else if (a_2d) {
+    outputs.push_back(multiply(a_splits[1], b_splits[2], s));
+    outputs.push_back(negative(multiply(a_splits[0], b_splits[2], s), s));
+  } else {
+    outputs.push_back(subtract(
+        multiply(a_splits[1], b_splits[2], s),
+        multiply(a_splits[2], b_splits[1], s),
+        s));
+    outputs.push_back(subtract(
+        multiply(a_splits[2], b_splits[0], s),
+        multiply(a_splits[0], b_splits[2], s),
+        s));
+  }
+  outputs.push_back(subtract(
+      multiply(a_splits[0], b_splits[1], s),
+      multiply(a_splits[1], b_splits[0], s),
+      s));
+  return concatenate(outputs, axis, s);
+}
+
+void validate_eig(
+    const array& a,
+    const StreamOrDevice& stream,
+    const std::string& fname) {
+  check_cpu_stream(stream, fname);
+  check_float_or_complex(a.dtype(), fname);
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << fname << " Arrays must have >= 2 dimensions. Received array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.shape(-1) != a.shape(-2)) {
+    throw std::invalid_argument(fname + " Only defined for square matrices.");
+  }
+}
+
+array eigvalsh(
+    const array& a,
+    std::string UPLO /* = "L" */,
+    StreamOrDevice s /* = {} */) {
+  validate_eig(a, s, "[linalg::eigvalsh]");
+  Shape out_shape(a.shape().begin(), a.shape().end() - 1);
+  Dtype eigval_type = a.dtype() == complex64 ? float32 : a.dtype();
+  return array(
+      std::move(out_shape),
+      eigval_type,
+      std::make_shared<Eigh>(to_stream(s), UPLO, false),
+      {a});
+}
+
+std::pair<array, array> eigh(
+    const array& a,
+    std::string UPLO /* = "L" */,
+    StreamOrDevice s /* = {} */) {
+  validate_eig(a, s, "[linalg::eigh]");
+  Dtype eigval_type = a.dtype() == complex64 ? float32 : a.dtype();
+  auto out = array::make_arrays(
+      {Shape(a.shape().begin(), a.shape().end() - 1), a.shape()},
+      {eigval_type, a.dtype()},
+      std::make_shared<Eigh>(to_stream(s), UPLO, true),
+      {a});
+  return std::make_pair(out[0], out[1]);
+}
+
+array eigvals(const array& a, StreamOrDevice s /* = {} */) {
+  validate_eig(a, s, "[linalg::eigvals]");
+  Shape out_shape(a.shape().begin(), a.shape().end() - 1);
+  return array(
+      std::move(out_shape),
+      complex64,
+      std::make_shared<Eig>(to_stream(s), false),
+      {a});
+}
+
+std::pair<array, array> eig(const array& a, StreamOrDevice s /* = {} */) {
+  validate_eig(a, s, "[linalg::eig]");
+  auto out = array::make_arrays(
+      {Shape(a.shape().begin(), a.shape().end() - 1), a.shape()},
+      {complex64, complex64},
+      std::make_shared<Eig>(to_stream(s), true),
+      {a});
+  return std::make_pair(out[0], out[1]);
+}
+
+void validate_lu(
+    const array& a,
+    const StreamOrDevice& stream,
+    const std::string& fname) {
+  check_cpu_stream(stream, fname);
+  check_float(a.dtype(), fname);
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << fname
+        << " Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+std::vector<array> lu_helper(const array& a, StreamOrDevice s /* = {} */) {
+  int m = a.shape()[a.shape().size() - 2];
+  int n = a.shape()[a.shape().size() - 1];
+
+  Shape pivots_shape(a.shape().begin(), a.shape().end() - 2);
+  pivots_shape.push_back(std::min(m, n));
+
+  Shape row_idx_shape(a.shape().begin(), a.shape().end() - 1);
+
+  return array::make_arrays(
+      {a.shape(), pivots_shape, row_idx_shape},
+      {a.dtype(), uint32, uint32},
+      std::make_shared<LUF>(to_stream(s)),
+      {astype(a, a.dtype(), s)});
+}
+
+std::vector<array> lu(const array& a, StreamOrDevice s /* = {} */) {
+  validate_lu(a, s, "[linalg::lu]");
+
+  auto out = lu_helper(a, s);
+  auto& LU = out[0];
+  auto& row_pivots = out[2];
+  auto L = tril(LU, /* k = */ -1, s);
+  auto U = triu(LU, /* k = */ 0, s);
+
+  int M = a.shape(-2);
+  int N = a.shape(-1);
+  int K = std::min(M, N);
+  if (N != K) {
+    auto start = Shape(L.ndim(), 0);
+    auto stop = L.shape();
+    stop.back() = K;
+    L = slice(L, std::move(start), std::move(stop), s);
+  } else if (M != K) {
+    auto start = Shape(U.ndim(), 0);
+    auto stop = U.shape();
+    stop[U.ndim() - 2] = K;
+    U = slice(U, std::move(start), std::move(stop), s);
+  }
+  L = add(L, eye(M, K, s), s);
+  return {row_pivots, L, U};
+}
+
+std::pair<array, array> lu_factor(const array& a, StreamOrDevice s /* = {} */) {
+  validate_lu(a, s, "[linalg::lu_factor]");
+  auto out = lu_helper(a, s);
+  return std::make_pair(out[0], out[1]);
+}
+
+void validate_solve(
+    const array& a,
+    const array& b,
+    const StreamOrDevice& stream,
+    const std::string& fname) {
+  check_cpu_stream(stream, fname);
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << fname << " First input must have >= 2 dimensions. "
+        << "Received array with " << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (b.ndim() < 1) {
+    std::ostringstream msg;
+    msg << fname << " Second input must have >= 1 dimensions. "
+        << "Received array with " << b.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.shape(-1) != a.shape(-2)) {
+    std::ostringstream msg;
+    msg << fname << " First input must be a square matrix. "
+        << "Received array with shape " << a.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  int lastDim = b.ndim() > 1 ? -2 : -1;
+  if (a.shape(-1) != b.shape(lastDim)) {
+    std::ostringstream msg;
+    msg << fname << " Last dimension of first input with shape " << a.shape()
+        << " must match second to last dimension of"
+        << " second input with shape " << b.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto out_type = promote_types(a.dtype(), b.dtype());
+  if (out_type != float32 && out_type != float64) {
+    std::ostringstream msg;
+    msg << fname
+        << " Input arrays must promote to float32 or float64. "
+           " Received arrays with type "
+        << a.dtype() << " and " << b.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+array solve(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  validate_solve(a, b, s, "[linalg::solve]");
+
+  // P, L, U matrices
+  const auto luf = lu(a, s);
+  auto perm = argsort(luf[0], -1, s);
+  int take_axis = -1;
+  if (b.ndim() >= 2) {
+    perm = expand_dims(perm, -1, s);
+    take_axis -= 1;
+  }
+  auto pb = take_along_axis(b, perm, take_axis, s);
+  auto y = solve_triangular(luf[1], pb, /* upper = */ false, s);
+  return solve_triangular(luf[2], y, /* upper = */ true, s);
+}
+
+array solve_triangular(
+    const array& a,
+    const array& b,
+    bool upper /* = false */,
+    StreamOrDevice s /* = {} */) {
+  validate_solve(a, b, s, "[linalg::solve_triangular]");
+  auto a_inv = tri_inv(a, upper, s);
+  return matmul(a_inv, b, s);
+}
+
+} // namespace mlx::core::linalg
\ No newline at end of file
diff --git a/macos/third_party/mlx/mlx/linalg.h b/macos/third_party/mlx/mlx/linalg.h
new file mode 100644
index 0000000..fe3f83c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/linalg.h
@@ -0,0 +1,115 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <optional>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/device.h"
+#include "mlx/ops.h"
+#include "mlx/stream.h"
+
+namespace mlx::core::linalg {
+
+/**
+ * Compute vector or matrix norms.
+ *
+ * - If axis and ord are both unspecified, computes the 2-norm of flatten(x).
+ * - If axis is not provided but ord is, then x must be either 1D or 2D.
+ * - If axis is provided, but ord is not, then the 2-norm (or Frobenius norm
+ *   for matrices) is computed along the given axes. At most 2 axes can be
+ *   specified.
+ * - If both axis and ord are provided, then the corresponding matrix or vector
+ *   norm is computed. At most 2 axes can be specified.
+ */
+MLX_API array norm(
+    const array& a,
+    const double ord,
+    const std::optional<std::vector<int>>& axis = std::nullopt,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+inline array norm(
+    const array& a,
+    const double ord,
+    int axis,
+    bool keepdims = false,
+    StreamOrDevice s = {}) {
+  return norm(a, ord, std::vector<int>{axis}, keepdims, s);
+}
+MLX_API array norm(
+    const array& a,
+    const std::string& ord,
+    const std::optional<std::vector<int>>& axis = std::nullopt,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+inline array norm(
+    const array& a,
+    const std::string& ord,
+    int axis,
+    bool keepdims = false,
+    StreamOrDevice s = {}) {
+  return norm(a, ord, std::vector<int>{axis}, keepdims, s);
+}
+MLX_API array norm(
+    const array& a,
+    const std::optional<std::vector<int>>& axis = std::nullopt,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+inline array
+norm(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}) {
+  return norm(a, std::vector<int>{axis}, keepdims, s);
+}
+
+MLX_API std::pair<array, array> qr(const array& a, StreamOrDevice s = {});
+
+MLX_API std::vector<array>
+svd(const array& a, bool compute_uv, StreamOrDevice s /* = {} */);
+inline std::vector<array> svd(const array& a, StreamOrDevice s = {}) {
+  return svd(a, true, s);
+}
+
+MLX_API array inv(const array& a, StreamOrDevice s = {});
+
+MLX_API array
+tri_inv(const array& a, bool upper = false, StreamOrDevice s = {});
+
+MLX_API array
+cholesky(const array& a, bool upper = false, StreamOrDevice s = {});
+
+MLX_API array pinv(const array& a, StreamOrDevice s = {});
+
+MLX_API array
+cholesky_inv(const array& a, bool upper = false, StreamOrDevice s = {});
+
+MLX_API std::vector<array> lu(const array& a, StreamOrDevice s = {});
+
+MLX_API std::pair<array, array> lu_factor(
+    const array& a,
+    StreamOrDevice s = {});
+
+MLX_API array solve(const array& a, const array& b, StreamOrDevice s = {});
+
+MLX_API array solve_triangular(
+    const array& a,
+    const array& b,
+    bool upper = false,
+    StreamOrDevice s = {});
+
+/**
+ * Compute the cross product of two arrays along the given axis.
+ */
+MLX_API array
+cross(const array& a, const array& b, int axis = -1, StreamOrDevice s = {});
+
+MLX_API std::pair<array, array> eig(const array& a, StreamOrDevice s = {});
+
+MLX_API array eigvals(const array& a, StreamOrDevice s = {});
+
+MLX_API array
+eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
+
+MLX_API std::pair<array, array>
+eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
+
+} // namespace mlx::core::linalg
diff --git a/macos/third_party/mlx/mlx/memory.h b/macos/third_party/mlx/mlx/memory.h
new file mode 100644
index 0000000..f4eabc9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/memory.h
@@ -0,0 +1,80 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "mlx/api.h"
+
+namespace mlx::core {
+
+/* Get the actively used memory in bytes.
+ *
+ * Note, this will not always match memory use reported by the system because
+ * it does not include cached memory buffers.
+ * */
+MLX_API size_t get_active_memory();
+
+/* Get the peak amount of used memory in bytes.
+ *
+ * The maximum memory used recorded from the beginning of the program
+ * execution or since the last call to reset_peak_memory.
+ * */
+MLX_API size_t get_peak_memory();
+
+/* Reset the peak memory to zero.
+ * */
+MLX_API void reset_peak_memory();
+
+/* Get the cache size in bytes.
+ *
+ * The cache includes memory not currently used that has not been returned
+ * to the system allocator.
+ * */
+MLX_API size_t get_cache_memory();
+
+/* Set the memory limit.
+ * The memory limit is a guideline for the maximum amount of memory to use
+ * during graph evaluation. If the memory limit is exceeded and there is no
+ * more RAM (including swap when available) allocations will result in an
+ * exception.
+ *
+ * When Metal is available the memory limit defaults to 1.5 times the maximum
+ * recommended working set size reported by the device.
+ *
+ * Returns the previous memory limit.
+ * */
+MLX_API size_t set_memory_limit(size_t limit);
+
+/* Get the current memory limit. */
+MLX_API size_t get_memory_limit();
+
+/* Set the cache limit.
+ * If using more than the given limit, free memory will be reclaimed
+ * from the cache on the next allocation. To disable the cache,
+ * set the limit to 0.
+ *
+ * The cache limit defaults to the memory limit.
+ *
+ * Returns the previous cache limit.
+ * */
+MLX_API size_t set_cache_limit(size_t limit);
+
+/* Clear the memory cache. */
+MLX_API void clear_cache();
+
+/* Set the wired size limit.
+ *
+ * Note, this function is only useful when using the Metal backend with
+ * macOS 15.0 or higher.
+ *
+ * The wired limit is the total size in bytes of memory that will be kept
+ * resident. The default value is ``0``.
+ *
+ * Setting a wired limit larger than system wired limit is an error.
+ *
+ * Returns the previous wired limit.
+ * */
+MLX_API size_t set_wired_limit(size_t limit);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/mlx.h b/macos/third_party/mlx/mlx/mlx.h
new file mode 100644
index 0000000..eda7333
--- /dev/null
+++ b/macos/third_party/mlx/mlx/mlx.h
@@ -0,0 +1,25 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/compile.h"
+#include "mlx/device.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/ops.h"
+#include "mlx/einsum.h"
+#include "mlx/export.h"
+#include "mlx/fast.h"
+#include "mlx/fft.h"
+#include "mlx/io.h"
+#include "mlx/linalg.h"
+#include "mlx/memory.h"
+#include "mlx/ops.h"
+#include "mlx/random.h"
+#include "mlx/stream.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
+#include "mlx/version.h"
diff --git a/macos/third_party/mlx/mlx/ops.cpp b/macos/third_party/mlx/mlx/ops.cpp
new file mode 100644
index 0000000..076de69
--- /dev/null
+++ b/macos/third_party/mlx/mlx/ops.cpp
@@ -0,0 +1,6094 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// Required for using M_PI in MSVC.
+#define _USE_MATH_DEFINES
+#include <algorithm>
+#include <climits>
+#include <cmath>
+#include <numeric>
+#include <set>
+#include <sstream>
+
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/transforms_impl.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+std::tuple<Shape, std::vector<int>, bool> compute_reduce_shape(
+    const std::vector<int>& axes,
+    const Shape& shape) {
+  bool is_noop = true;
+  std::set<int> axes_set;
+  auto ndim = shape.size();
+  for (auto ax : axes) {
+    int ax_ = (ax < 0) ? ax + ndim : ax;
+    if (ax_ < 0 || ax_ >= ndim) {
+      std::ostringstream msg;
+      msg << "Invalid axis " << ax << " for array with " << ndim
+          << " dimensions.";
+      throw std::out_of_range(msg.str());
+    }
+    axes_set.insert(ax_);
+  }
+  if (axes_set.size() != axes.size()) {
+    throw std::invalid_argument("Duplicate axes detected in reduction.");
+  }
+  Shape out_shape;
+  for (int i = 0; i < ndim; ++i) {
+    if (axes_set.count(i) == 0) {
+      out_shape.push_back(shape[i]);
+    } else {
+      out_shape.push_back(1);
+    }
+    is_noop &= (out_shape.back() == shape[i]);
+  }
+  std::vector<int> sorted_axes(axes_set.begin(), axes_set.end());
+  return {out_shape, sorted_axes, is_noop};
+}
+
+Dtype at_least_float(const Dtype& d) {
+  return issubdtype(d, inexact) ? d : promote_types(d, float32);
+}
+
+array indices_or_default(
+    std::optional<array> indices,
+    const array& x,
+    StreamOrDevice s) {
+  if (indices.has_value()) {
+    return indices.value();
+  }
+
+  Shape shape(x.shape().begin(), x.shape().end() - 2);
+  int total =
+      std::reduce(shape.begin(), shape.end(), 1, std::multiplies<int>());
+  return reshape(arange(total, uint32, s), std::move(shape), s);
+}
+
+void validate_quantized_input(
+    std::string_view tag,
+    const array& w,
+    const array& scales,
+    int group_size,
+    int bits,
+    const std::optional<array>& biases = std::nullopt) {
+  if (w.dtype() != uint32) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] The weight matrix should be uint32 "
+        << "but received " << w.dtype();
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (biases && scales.shape() != biases->shape()) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] Scales and biases should have the same shape. "
+        << "Received scales with shape " << scales.shape()
+        << " and biases with " << biases->shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (!std::equal(
+          w.shape().begin(), w.shape().end() - 2, scales.shape().begin())) {
+    std::ostringstream msg;
+    msg << "[" << tag
+        << "] Weight and scales should have the same batch shape. "
+        << "Received weight with shape " << w.shape() << ", scales with "
+        << scales.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (w.shape(-1) * 32 / bits != scales.shape(-1) * group_size) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] The shapes of the weight and scales are "
+        << "incompatible based on bits and group_size. w.shape() == "
+        << w.shape() << " and scales.shape() == " << scales.shape()
+        << " with group_size=" << group_size << " and bits=" << bits;
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+std::pair<int, int> extract_quantized_matmul_dims(
+    std::string_view tag,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    bool transpose,
+    int group_size,
+    int bits) {
+  validate_quantized_input(tag, w, scales, group_size, bits, biases);
+
+  int x_inner_dims = x.shape(-1);
+
+  // Calculate the expanded w's dims
+  int w_inner_dims = (transpose) ? w.shape(-1) * 32 / bits : w.shape(-2);
+  int w_outer_dims = (transpose) ? w.shape(-2) : w.shape(-1) * 32 / bits;
+
+  if (w_inner_dims != x_inner_dims) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] Last dimension of first input with "
+        << "shape (..., " << x_inner_dims << ") does not match "
+        << "the expanded quantized matrix (" << w_inner_dims << ", "
+        << w_outer_dims << ") computed from shape " << w.shape()
+        << " with group_size=" << group_size << ", bits=" << bits
+        << " and transpose=" << std::boolalpha << transpose;
+    throw std::invalid_argument(msg.str());
+  }
+
+  return {w_inner_dims, w_outer_dims};
+}
+
+} // namespace
+
+array arange(
+    double start,
+    double stop,
+    double step,
+    Dtype dtype,
+    StreamOrDevice s /* = {} */) {
+  if (dtype == bool_) {
+    std::ostringstream msg;
+    msg << bool_ << " not supported for arange.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (std::isnan(start) || std::isnan(step) || std::isnan(stop)) {
+    throw std::invalid_argument("[arange] Cannot compute length.");
+  }
+
+  if (std::isinf(start) || std::isinf(stop)) {
+    throw std::invalid_argument("[arange] Cannot compute length.");
+  }
+
+  // Check if start and stop specify a valid range because if not, we have to
+  // return an empty array
+  if (std::isinf(step) &&
+      ((step > 0 && start < stop) || (step < 0 && start > stop))) {
+    return array({start}, dtype);
+  }
+
+  double real_size = std::ceil((stop - start) / step);
+
+  if (real_size > INT_MAX) {
+    throw std::invalid_argument("[arange] Maximum size exceeded.");
+  }
+
+  int size = std::max(static_cast<int>(real_size), 0);
+  return array(
+      {size},
+      dtype,
+      std::make_shared<Arange>(to_stream(s), start, stop, step),
+      {});
+}
+array arange(
+    double start,
+    double stop,
+    double step,
+    StreamOrDevice s /* = {} */) {
+  return arange(start, stop, step, float32, to_stream(s));
+}
+array arange(
+    double start,
+    double stop,
+    Dtype dtype,
+    StreamOrDevice s /* = {} */) {
+  return arange(start, stop, 1.0, dtype, to_stream(s));
+}
+array arange(double start, double stop, StreamOrDevice s /* = {} */) {
+  return arange(start, stop, 1.0, float32, to_stream(s));
+}
+array arange(double stop, Dtype dtype, StreamOrDevice s /* = {} */) {
+  return arange(0.0, stop, 1.0, dtype, to_stream(s));
+}
+array arange(double stop, StreamOrDevice s /* = {} */) {
+  return arange(0.0, stop, 1.0, float32, to_stream(s));
+}
+array arange(int start, int stop, int step, StreamOrDevice s /* = {} */) {
+  return arange(
+      static_cast<double>(start),
+      static_cast<double>(stop),
+      static_cast<double>(step),
+      int32,
+      to_stream(s));
+}
+array arange(int start, int stop, StreamOrDevice s /* = {} */) {
+  return arange(
+      static_cast<double>(start),
+      static_cast<double>(stop),
+      1.0,
+      int32,
+      to_stream(s));
+}
+array arange(int stop, StreamOrDevice s /* = {} */) {
+  return arange(0.0, static_cast<double>(stop), 1.0, int32, to_stream(s));
+}
+
+array linspace(
+    double start,
+    double stop,
+    int num /* = 50 */,
+    Dtype dtype /* = float32 */,
+    StreamOrDevice s /* = {} */) {
+  if (num < 0) {
+    std::ostringstream msg;
+    msg << "[linspace] number of samples, " << num << ", must be non-negative.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (num == 1) {
+    return astype(array({start}), dtype, s);
+  }
+  auto inner_type = dtype == float64 ? float64 : float32;
+  array t =
+      divide(arange(0, num, inner_type, s), array(num - 1, inner_type), s);
+  array t_bar = subtract(array(1, inner_type), t, s);
+  return astype(
+      add(multiply(t_bar, array(start, inner_type), s),
+          multiply(t, array(stop, inner_type), s),
+          s),
+      dtype,
+      s);
+}
+
+array astype(array a, Dtype dtype, StreamOrDevice s /* = {} */) {
+  if (dtype == a.dtype()) {
+    return a;
+  }
+  auto copied_shape = a.shape(); // |a| will be moved
+  return array(
+      std::move(copied_shape),
+      dtype,
+      std::make_shared<AsType>(to_stream(s), dtype),
+      {std::move(a)});
+}
+
+array as_strided(
+    array a,
+    Shape shape,
+    Strides strides,
+    size_t offset,
+    StreamOrDevice s /* = {} */) {
+  auto copied_shape = shape; // |shape| will be moved
+  auto dtype = a.dtype(); // |a| will be moved
+  return array(
+      std::move(copied_shape),
+      dtype,
+      std::make_shared<AsStrided>(
+          to_stream(s), std::move(shape), std::move(strides), offset),
+      // Force the input array to be contiguous.
+      {flatten(std::move(a), s)});
+}
+
+array copy(array a, StreamOrDevice s /* = {} */) {
+  auto copied_shape = a.shape(); // |a| will be moved
+  auto dtype = a.dtype();
+  return array(
+      std::move(copied_shape),
+      dtype,
+      std::make_shared<Copy>(to_stream(s)),
+      {std::move(a)});
+}
+
+array full_impl(array vals, Dtype dtype, StreamOrDevice s /* = {} */) {
+  return array(
+      vals.shape(),
+      dtype,
+      std::make_shared<Full>(to_stream(s)),
+      {astype(vals, dtype, s)});
+}
+
+array full(Shape shape, array vals, Dtype dtype, StreamOrDevice s /* = {} */) {
+  if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) {
+    throw std::invalid_argument("[full] Negative dimensions not allowed.");
+  }
+  return full_impl(broadcast_to(vals, std::move(shape), s), dtype, s);
+}
+
+array full(Shape shape, array vals, StreamOrDevice s /* = {} */) {
+  auto dtype = vals.dtype(); // |vals| will be moved
+  return full(std::move(shape), std::move(vals), dtype, to_stream(s));
+}
+
+array full_like(
+    const array& a,
+    array vals,
+    Dtype dtype,
+    StreamOrDevice s /* = {} */) {
+  auto inputs = broadcast_arrays({a, std::move(vals)}, s);
+  return full_impl(std::move(inputs[1]), dtype, s);
+}
+
+array full_like(const array& a, array vals, StreamOrDevice s /* = {} */) {
+  return full_like(a, std::move(vals), a.dtype(), to_stream(s));
+}
+
+array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) {
+  return full(shape, array(0, dtype), to_stream(s));
+}
+
+array zeros_like(const array& a, StreamOrDevice s /* = {} */) {
+  return full_like(a, 0, a.dtype(), to_stream(s));
+}
+
+array ones(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) {
+  return full(shape, array(1, dtype), to_stream(s));
+}
+
+array ones_like(const array& a, StreamOrDevice s /* = {} */) {
+  return full_like(a, 1, a.dtype(), to_stream(s));
+}
+
+array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s /* = {} */) {
+  if (n <= 0 || m <= 0) {
+    throw std::invalid_argument("[eye] N and M must be positive integers.");
+  }
+  array result = zeros({n, m}, dtype, s);
+  if (k >= m || -k >= n) {
+    return result;
+  }
+
+  int diagonal_length = k >= 0 ? std::min(n, m - k) : std::min(n + k, m);
+
+  std::vector<array> indices;
+  auto s1 = std::max(0, -k);
+  auto s2 = std::max(0, k);
+  indices.push_back(arange(s1, diagonal_length + s1, int32, s));
+  indices.push_back(arange(s2, diagonal_length + s2, int32, s));
+  array ones_array = ones({diagonal_length, 1, 1}, dtype, s);
+  return scatter(result, indices, ones_array, {0, 1}, s);
+}
+
+array identity(int n, Dtype dtype, StreamOrDevice s /* = {} */) {
+  return eye(n, n, 0, dtype, s);
+}
+
+array tri(int n, int m, int k, Dtype type, StreamOrDevice s /* = {} */) {
+  auto l = expand_dims(arange(n, s), 1, s);
+  auto r = expand_dims(arange(-k, m - k, s), 0, s);
+  return astype(greater_equal(l, r, s), type, s);
+}
+
+array tril(array x, int k /* = 0 */, StreamOrDevice s /* = {} */) {
+  if (x.ndim() < 2) {
+    throw std::invalid_argument("[tril] array must be at least 2-D");
+  }
+  auto mask = tri(x.shape(-2), x.shape(-1), k, bool_, s);
+  return where(mask, x, array(0, x.dtype()), s);
+}
+
+array triu(array x, int k /* = 0 */, StreamOrDevice s /* = {} */) {
+  if (x.ndim() < 2) {
+    throw std::invalid_argument("[triu] array must be at least 2-D");
+  }
+  auto mask = tri(x.shape(-2), x.shape(-1), k - 1, bool_, s);
+  return where(mask, array(0, x.dtype()), x, s);
+}
+
+array reshape(const array& a, Shape shape, StreamOrDevice s /* = {} */) {
+  if (a.shape() == shape) {
+    return a;
+  }
+  auto out_shape = Reshape::output_shape(a, shape);
+  return array(
+      std::move(out_shape),
+      a.dtype(),
+      std::make_shared<Reshape>(to_stream(s), std::move(shape)),
+      {a});
+}
+
+array unflatten(
+    const array& a,
+    int axis,
+    Shape shape,
+    StreamOrDevice s /* = {} */) {
+  if (shape.empty()) {
+    throw std::invalid_argument(
+        "[unflatten] Shape to unflatten to cannot be empty.");
+  }
+  auto ndim = static_cast<int>(a.ndim());
+  auto ax = axis < 0 ? axis + ndim : axis;
+  if (ax < 0 || ax >= ndim) {
+    std::ostringstream msg;
+    msg << "[unflatten] Invalid axes " << ax << " for array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  size_t size = 1;
+  int infer_idx = -1;
+  for (int i = 0; i < shape.size(); ++i) {
+    if (shape[i] == -1) {
+      if (infer_idx >= 0) {
+        throw std::invalid_argument(
+            "[Unflatten] Can only infer one dimension.");
+      }
+      infer_idx = i;
+    } else {
+      size *= shape[i];
+    }
+  }
+  if (infer_idx >= 0) {
+    shape[infer_idx] = a.shape(ax) / size;
+    size *= shape[infer_idx];
+  }
+  if (size != a.shape(ax)) {
+    std::ostringstream msg;
+    msg << "[Unflatten] Cannot unflatten axis " << axis << " with size "
+        << a.shape(ax) << " into shape " << shape << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (shape.size() == 1) {
+    return a;
+  }
+
+  auto out_shape = Unflatten::output_shape(a, ax, shape);
+  return array(
+      std::move(out_shape),
+      a.dtype(),
+      std::make_shared<Unflatten>(to_stream(s), ax, std::move(shape)),
+      {a});
+}
+
+array flatten(
+    const array& a,
+    int start_axis,
+    int end_axis /* = -1 */,
+    StreamOrDevice s /* = {} */) {
+  auto ndim = static_cast<int>(a.ndim());
+  auto start_ax = start_axis + (start_axis < 0 ? ndim : 0);
+  auto end_ax = end_axis + (end_axis < 0 ? ndim : 0);
+  start_ax = std::max(0, start_ax);
+  end_ax = std::min(ndim - 1, end_ax);
+  if (a.ndim() == 0) {
+    return reshape(a, {1}, s);
+  }
+  if (end_ax < start_ax) {
+    throw std::invalid_argument(
+        "[flatten] start_axis must be less than or equal to end_axis");
+  }
+  if (start_ax >= ndim) {
+    std::ostringstream msg;
+    msg << "[flatten] Invalid start_axis " << start_axis << " for array with "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (end_ax < 0) {
+    std::ostringstream msg;
+    msg << "[flatten] Invalid end_axis " << end_axis << " for array with "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (start_ax == end_ax) {
+    return a;
+  }
+  return array(
+      Flatten::output_shape(a, start_ax, end_ax),
+      a.dtype(),
+      std::make_shared<Flatten>(to_stream(s), start_ax, end_ax),
+      {a});
+}
+
+array flatten(const array& a, StreamOrDevice s /* = {} */) {
+  return flatten(a, 0, a.ndim() - 1, s);
+}
+
+array hadamard_transform(
+    const array& a,
+    std::optional<float> scale_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  if (a.size() == 0) {
+    throw std::invalid_argument(
+        "[hadamard_transform] Does not support empty arrays.");
+  }
+  // Default to an orthonormal Hadamard matrix scaled by 1/sqrt(N)
+  int n = a.ndim() > 0 ? a.shape(-1) : 1;
+  float scale = scale_.has_value() ? *scale_ : 1.0f / std::sqrt(n);
+  auto dtype = issubdtype(a.dtype(), floating) ? a.dtype() : float32;
+
+  // Nothing to do for a scalar
+  if (n == 1) {
+    if (scale == 1) {
+      return a;
+    }
+
+    return multiply(a, array(scale, dtype), s);
+  }
+
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<Hadamard>(to_stream(s), scale),
+      {astype(a, dtype, s)});
+}
+
+array squeeze_impl(
+    const array& a,
+    std::vector<int> axes,
+    StreamOrDevice s /* = {} */) {
+  for (auto& ax : axes) {
+    auto new_ax = ax < 0 ? ax + a.ndim() : ax;
+    if (new_ax < 0 || new_ax >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[squeeze] Invalid axes " << ax << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if (a.shape(new_ax) != 1) {
+      std::ostringstream msg;
+      msg << "[squeeze] Cannot squeeze axis " << ax << " with size "
+          << a.shape(ax) << " which is not equal to 1.";
+      throw std::invalid_argument(msg.str());
+    }
+    ax = new_ax;
+  }
+  auto shape = Squeeze::output_shape(a, axes);
+  return array(
+      std::move(shape),
+      a.dtype(),
+      std::make_shared<Squeeze>(to_stream(s), std::move(axes)),
+      {a});
+}
+
+array squeeze(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+  std::set<int> unique_axes;
+  for (auto ax : axes) {
+    unique_axes.insert(ax < 0 ? ax + a.ndim() : ax);
+  }
+  if (unique_axes.size() != axes.size()) {
+    throw std::invalid_argument("[squeeze] Received duplicate axes.");
+  }
+  std::vector<int> sorted_axes(unique_axes.begin(), unique_axes.end());
+  return squeeze_impl(a, std::move(sorted_axes), s);
+}
+
+array squeeze(const array& a, int axis, StreamOrDevice s /* = {} */) {
+  return squeeze_impl(a, {axis}, s);
+}
+
+array squeeze(const array& a, StreamOrDevice s /* = {} */) {
+  std::vector<int> axes;
+  for (int i = 0; i < a.ndim(); ++i) {
+    if (a.shape(i) == 1) {
+      axes.push_back(i);
+    }
+  }
+  return squeeze_impl(a, std::move(axes), s);
+}
+
+array expand_dims_impl(
+    const array& a,
+    std::vector<int> axes,
+    StreamOrDevice s /* = {} */) {
+  auto out_ndim = a.ndim() + axes.size();
+  for (auto& ax : axes) {
+    auto new_ax = ax < 0 ? ax + out_ndim : ax;
+    if (new_ax < 0 || new_ax >= out_ndim) {
+      std::ostringstream msg;
+      msg << "[expand_dims] Invalid axis " << ax << " for output array with "
+          << a.ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    ax = new_ax;
+  }
+  auto shape = ExpandDims::output_shape(a, axes);
+  return array(
+      std::move(shape),
+      a.dtype(),
+      std::make_shared<ExpandDims>(to_stream(s), std::move(axes)),
+      {a});
+}
+
+array expand_dims(const array& a, int axis, StreamOrDevice s /* = {} */) {
+  return expand_dims_impl(a, {axis}, s);
+}
+
+array expand_dims(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+  { // Check for repeats
+    std::set<int> unique_axes(axes.begin(), axes.end());
+    if (unique_axes.size() != axes.size()) {
+      throw std::invalid_argument("[expand_dims] Received duplicate axes.");
+    }
+  }
+  // Check for repeats again
+  auto out_ndim = a.ndim() + axes.size();
+  std::set<int> unique_axes;
+  for (auto ax : axes) {
+    unique_axes.insert(ax < 0 ? ax + out_ndim : ax);
+  }
+  if (unique_axes.size() != axes.size()) {
+    throw std::invalid_argument("[expand_dims] Received duplicate axes.");
+  }
+  std::vector<int> sorted_axes(unique_axes.begin(), unique_axes.end());
+  return expand_dims_impl(a, std::move(sorted_axes), s);
+}
+
+// Slice helper
+namespace {
+
+inline auto
+normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) {
+  // - Start indices are normalized
+  // - End indices are unchanged as -1 means something different
+  //   pre-normalization (the end of the axis) versus post normalization (the
+  //   position left of 0).
+  // - Any strides corresponding to singleton dimension are set to 1
+
+  Shape out_shape(shape.size());
+  bool has_neg_strides = false;
+
+  for (int i = 0; i < shape.size(); ++i) {
+    // Following numpy docs
+    //  Negative i and j are interpreted as n + i and n + j where n is
+    //  the number of elements in the corresponding dimension. Negative
+    //  k makes stepping go towards smaller indices
+
+    auto n = shape[i];
+    auto s = start[i];
+    s = s < 0 ? s + n : s;
+    auto e = stop[i];
+    e = e < 0 ? e + n : e;
+
+    // Note: -ve strides require start >= stop
+    if (strides[i] < 0) {
+      has_neg_strides = true;
+
+      // Clamp to bounds
+      auto st = std::min(s, n - 1);
+      auto ed = e > -1 ? e : -1;
+
+      start[i] = st;
+      ed = ed > st ? st : ed;
+
+      auto str = -strides[i];
+      out_shape[i] = (start[i] - ed + str - 1) / str;
+
+    } else {
+      // Clamp to bounds
+      auto st = std::max(static_cast<ShapeElem>(0), std::min(s, n));
+      auto ed = std::max(static_cast<ShapeElem>(0), std::min(e, n));
+
+      start[i] = st;
+      ed = ed < st ? st : ed;
+
+      out_shape[i] = (ed - start[i] + strides[i] - 1) / strides[i];
+    }
+    // Simplify the stride if it's unused
+    if (out_shape[i] == 1) {
+      strides[i] = 1;
+    }
+  }
+
+  return std::make_pair(has_neg_strides, out_shape);
+}
+
+void normalize_dynamic_slice_inputs(
+    const array& a,
+    const array& start,
+    std::vector<int>& axes,
+    std::string_view prefix) {
+  if (start.size() > a.ndim()) {
+    std::ostringstream msg;
+    msg << prefix << " Invalid number of starting positions for "
+        << "array with dimension " << a.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (start.ndim() > 1) {
+    std::ostringstream msg;
+    msg << prefix << " array of starting indices "
+        << "must be zero or one dimensional but has dimension " << start.ndim()
+        << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (start.size() != axes.size()) {
+    std::ostringstream msg;
+    msg << prefix << " Number of starting indices " << start.size()
+        << " does not match number of axes " << axes.size() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (!issubdtype(start.dtype(), integer)) {
+    std::ostringstream msg;
+    msg << prefix << " Start indices must be integers, got type "
+        << start.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  for (auto& ax : axes) {
+    auto new_ax = ax < 0 ? ax + a.ndim() : ax;
+    if (new_ax < 0 || new_ax >= a.ndim()) {
+      std::ostringstream msg;
+      msg << prefix << " Invalid axis " << ax << " for array with dimension "
+          << a.ndim() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+    ax = new_ax;
+  }
+  std::set dims(axes.begin(), axes.end());
+  if (dims.size() != axes.size()) {
+    std::ostringstream msg;
+    msg << prefix << " Repeat axes not allowed.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+} // namespace
+
+array slice(
+    const array& a,
+    Shape start,
+    Shape stop,
+    Shape strides,
+    StreamOrDevice s /* = {} */) {
+  if (start.size() != a.ndim() || stop.size() != a.ndim() ||
+      strides.size() != a.ndim()) {
+    std::ostringstream msg;
+    msg << "[slice] Invalid number of indices or strides for "
+        << "array with dimension " << a.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto [has_neg_strides, out_shape] =
+      normalize_slice(a.shape(), start, stop, strides);
+
+  if (!has_neg_strides && out_shape == a.shape()) {
+    return a;
+  }
+
+  return array(
+      out_shape,
+      a.dtype(),
+      std::make_shared<Slice>(
+          to_stream(s), std::move(start), std::move(stop), std::move(strides)),
+      {a});
+}
+
+array slice(
+    const array& a,
+    Shape start,
+    Shape stop,
+    StreamOrDevice s /* = {} */) {
+  return slice(
+      a, std::move(start), std::move(stop), Shape(a.ndim(), 1), to_stream(s));
+}
+
+array slice(
+    const array& a,
+    const array& start,
+    std::vector<int> axes,
+    Shape slice_size,
+    StreamOrDevice s /* = {} */) {
+  normalize_dynamic_slice_inputs(a, start, axes, "[slice]");
+
+  // Check the slice_size
+  if (slice_size.size() != a.ndim()) {
+    std::ostringstream msg;
+    msg << "[slice] Invalid slice size for array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  for (int i = 0; i < a.ndim(); ++i) {
+    if (slice_size[i] > a.shape(i)) {
+      std::ostringstream msg;
+      msg << "[slice] Invalid slice size " << slice_size
+          << " for array with shape " << a.shape() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  auto out_shape = slice_size;
+  return array(
+      std::move(out_shape),
+      a.dtype(),
+      std::make_shared<DynamicSlice>(
+          to_stream(s), std::move(axes), std::move(slice_size)),
+      {a, start});
+}
+
+/** Update a slice from the source array */
+array slice_update(
+    const array& src,
+    const array& update,
+    Shape start,
+    Shape stop,
+    Shape strides,
+    StreamOrDevice s /* = {} */) {
+  // Check dimensions
+  if (start.size() != src.ndim() || stop.size() != src.ndim() ||
+      strides.size() != src.ndim()) {
+    std::ostringstream msg;
+    msg << "[slice_update] Invalid number of indices or strides for "
+        << "array with dimension " << src.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Process slice dimensions
+  auto [has_neg_strides, upd_shape] =
+      normalize_slice(src.shape(), start, stop, strides);
+
+  // Cast update to src type and broadcast update shape to slice shape
+  auto upd = broadcast_to(astype(update, src.dtype(), s), upd_shape, s);
+
+  // If the entire src is the slice, just return the update
+  if (!has_neg_strides && upd_shape == src.shape()) {
+    return upd;
+  }
+  return array(
+      src.shape(),
+      src.dtype(),
+      std::make_shared<SliceUpdate>(
+          to_stream(s), std::move(start), std::move(stop), std::move(strides)),
+      {src, upd});
+}
+
+/** Update a slice from the source array with stride 1 in each dimension */
+array slice_update(
+    const array& src,
+    const array& update,
+    Shape start,
+    Shape stop,
+    StreamOrDevice s /* = {} */) {
+  return slice_update(
+      src, update, std::move(start), std::move(stop), Shape(src.ndim(), 1), s);
+}
+
+/** Update a slice from the source array */
+array slice_update(
+    const array& src,
+    const array& update,
+    const array& start,
+    std::vector<int> axes,
+    StreamOrDevice s /* = {} */) {
+  normalize_dynamic_slice_inputs(src, start, axes, "[slice_update]");
+
+  // Broadcast update with unspecified axes
+  auto up_shape = update.shape();
+  auto dim_diff = std::max(src.ndim() - update.ndim(), size_t(0));
+  up_shape.insert(
+      up_shape.begin(), src.shape().begin(), src.shape().begin() + dim_diff);
+  for (int d = dim_diff; d < src.ndim(); ++d) {
+    up_shape[d] = std::min(up_shape[d], src.shape(d));
+  }
+  for (auto ax : axes) {
+    if (ax < dim_diff) {
+      up_shape[ax] = 1;
+    }
+  }
+  auto upd = broadcast_to(astype(update, src.dtype(), s), up_shape, s);
+  return array(
+      src.shape(),
+      src.dtype(),
+      std::make_shared<DynamicSliceUpdate>(to_stream(s), std::move(axes)),
+      {src, upd, start});
+}
+
+std::vector<array> split(
+    const array& a,
+    const Shape& indices,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  auto ax = axis < 0 ? axis + a.ndim() : axis;
+  if (ax < 0 || ax >= a.ndim()) {
+    std::ostringstream msg;
+    msg << "Invalid axis (" << axis << ") passed to split"
+        << " for array with shape " << a.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (indices.empty()) {
+    return {a};
+  }
+
+  if (indices.size() < 10 &&
+      std::is_sorted(indices.begin(), indices.end(), std::less<>{}) &&
+      indices[0] > 0 && indices.back() < a.shape(ax)) {
+    std::vector<Dtype> dtypes(indices.size() + 1, a.dtype());
+    std::vector<Shape> shapes(indices.size() + 1, a.shape());
+    shapes[0][ax] = indices[0];
+    for (int i = 1; i < indices.size(); i++) {
+      shapes[i][ax] = indices[i] - indices[i - 1];
+    }
+    shapes.back()[ax] = a.shape(ax) - indices.back();
+
+    return array::make_arrays(
+        std::move(shapes),
+        dtypes,
+        std::make_shared<Split>(to_stream(s), indices, ax),
+        {a});
+  }
+
+  std::vector<array> res;
+  auto start_indices = Shape(a.ndim(), 0);
+  auto stop_indices = a.shape();
+  for (int i = 0; i < indices.size() + 1; ++i) {
+    stop_indices[ax] = i < indices.size() ? indices[i] : a.shape(ax);
+    res.push_back(slice(a, start_indices, stop_indices, to_stream(s)));
+    start_indices[ax] = stop_indices[ax];
+  }
+  return res;
+}
+
+std::vector<array>
+split(const array& a, const Shape& indices, StreamOrDevice s /* = {} */) {
+  return split(a, indices, 0, s);
+}
+
+std::vector<array>
+split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) {
+  auto ax = axis < 0 ? axis + a.ndim() : axis;
+  if (ax < 0 || ax >= a.ndim()) {
+    std::ostringstream msg;
+    msg << "Invalid axis " << axis << " passed to split"
+        << " for array with shape " << a.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  auto q_and_r = std::ldiv(a.shape(axis), num_splits);
+  if (q_and_r.rem) {
+    std::ostringstream msg;
+    msg << "Array split does not result in sub arrays with equal size:"
+        << " attempting " << num_splits << " splits along axis " << axis
+        << " for shape " << a.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  auto split_size = q_and_r.quot;
+  Shape indices(num_splits - 1);
+  for (int i = 0; i < indices.size(); ++i) {
+    indices[i] = (i + 1) * split_size;
+  }
+  return split(a, indices, axis, s);
+}
+
+std::vector<array>
+split(const array& a, int num_splits, StreamOrDevice s /* = {} */) {
+  return split(a, num_splits, 0, to_stream(s));
+}
+
+std::vector<array> meshgrid(
+    const std::vector<array>& arrays,
+    bool sparse /* = false */,
+    const std::string& indexing /* = "xy" */,
+    StreamOrDevice s /* = {} */) {
+  if (indexing != "xy" && indexing != "ij") {
+    throw std::invalid_argument(
+        "[meshgrid] Invalid indexing value. Valid values are 'xy' and 'ij'.");
+  }
+
+  auto ndim = arrays.size();
+  std::vector<array> outputs;
+  for (int i = 0; i < ndim; ++i) {
+    Shape shape(ndim, 1);
+    shape[i] = -1;
+    outputs.push_back(reshape(arrays[i], std::move(shape), s));
+  }
+
+  if (indexing == "xy" && ndim > 1) {
+    Shape shape(ndim, 1);
+
+    shape[1] = arrays[0].size();
+    outputs[0] = reshape(arrays[0], shape, s);
+    shape[1] = 1;
+    shape[0] = arrays[1].size();
+    outputs[1] = reshape(arrays[1], std::move(shape), s);
+  }
+
+  if (!sparse) {
+    outputs = broadcast_arrays(outputs, s);
+  }
+
+  return outputs;
+}
+
+array clip(
+    const array& a,
+    const std::optional<array>& a_min,
+    const std::optional<array>& a_max,
+    StreamOrDevice s /* = {} */) {
+  if (!a_min.has_value() && !a_max.has_value()) {
+    throw std::invalid_argument("At most one of a_min and a_max may be None");
+  }
+  array result = a;
+  if (a_min.has_value()) {
+    result = maximum(result, a_min.value(), s);
+  }
+  if (a_max.has_value()) {
+    result = minimum(result, a_max.value(), s);
+  }
+  return result;
+}
+
+array concatenate(
+    std::vector<array> arrays,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  if (arrays.size() == 0) {
+    throw std::invalid_argument(
+        "[concatenate] No arrays provided for concatenation");
+  }
+  if (arrays.size() == 1) {
+    return arrays[0];
+  }
+
+  auto ax = normalize_axis_index(axis, arrays[0].ndim(), "[concatenate] ");
+
+  auto throw_invalid_shapes = [&]() {
+    std::ostringstream msg;
+    msg << "[concatenate] All the input array dimensions must match exactly "
+        << "except for the concatenation axis. However, the provided shapes are ";
+    for (auto& a : arrays) {
+      msg << a.shape() << ", ";
+    }
+    msg << "and the concatenation axis is " << axis << ".";
+    throw std::invalid_argument(msg.str());
+  };
+
+  auto shape = arrays[0].shape();
+  shape[ax] = 0;
+  // Make the output shape and validate that all arrays have the same shape
+  // except for the concatenation axis.
+  for (auto& a : arrays) {
+    if (a.ndim() != shape.size()) {
+      std::ostringstream msg;
+      msg << "[concatenate] All the input arrays must have the same number of "
+          << "dimensions. However, got arrays with dimensions " << shape.size()
+          << " and " << a.ndim() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+    for (int i = 0; i < a.ndim(); i++) {
+      if (i == ax) {
+        continue;
+      }
+      if (a.shape(i) != shape[i]) {
+        throw_invalid_shapes();
+      }
+    }
+    shape[ax] += a.shape(ax);
+  }
+
+  // Promote all the arrays to the same type
+  auto dtype = result_type(arrays);
+  for (auto& a : arrays) {
+    a = astype(a, dtype, s);
+  }
+
+  return array(
+      std::move(shape),
+      dtype,
+      std::make_shared<Concatenate>(to_stream(s), ax),
+      std::move(arrays));
+}
+
+array concatenate(std::vector<array> arrays, StreamOrDevice s /* = {} */) {
+  for (auto& a : arrays) {
+    a = flatten(a, s);
+  }
+  return concatenate(std::move(arrays), 0, s);
+}
+
+/** Stack arrays along a new axis */
+array stack(
+    const std::vector<array>& arrays,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  if (arrays.empty()) {
+    throw std::invalid_argument("[stack] No arrays provided for stacking");
+  }
+  if (!std::all_of(arrays.begin(), arrays.end(), [&](const auto& a) {
+        return arrays[0].shape() == a.shape();
+      })) {
+    throw std::invalid_argument("[stack] All arrays must have the same shape");
+  }
+  auto normalized_axis =
+      normalize_axis_index(axis, arrays[0].ndim() + 1, "[stack] ");
+  std::vector<array> new_arrays;
+  new_arrays.reserve(arrays.size());
+  for (auto& a : arrays) {
+    new_arrays.emplace_back(expand_dims(a, normalized_axis, s));
+  }
+  return concatenate(new_arrays, axis, s);
+}
+
+array stack(const std::vector<array>& arrays, StreamOrDevice s /* = {} */) {
+  return stack(arrays, 0, s);
+}
+
+/** array repeat with axis */
+array repeat(const array& arr, int repeats, int axis, StreamOrDevice s) {
+  axis = normalize_axis_index(axis, arr.ndim(), "[repeat] ");
+
+  if (repeats < 0) {
+    throw std::invalid_argument(
+        "[repeat] Number of repeats cannot be negative");
+  }
+
+  if (repeats == 0) {
+    return array({}, arr.dtype());
+  }
+
+  if (repeats == 1) {
+    return arr;
+  }
+
+  // Broadcast to (S_1, S_2, ..., S_axis, repeats, S_axis+1, ...)
+  auto shape = arr.shape();
+  shape.insert(shape.begin() + axis + 1, repeats);
+  array out = expand_dims(arr, axis + 1, s);
+  out = broadcast_to(out, shape, s);
+
+  // Reshape back into a contiguous array where S_axis is now S_axis * repeats
+  shape.erase(shape.begin() + axis + 1);
+  shape[axis] *= repeats;
+  out = reshape(out, shape, s);
+
+  return out;
+}
+
+array repeat(const array& arr, int repeats, StreamOrDevice s) {
+  return repeat(flatten(arr, s), repeats, 0, s);
+}
+
+array tile(
+    const array& arr,
+    std::vector<int> reps,
+    StreamOrDevice s /* = {} */) {
+  auto shape = arr.shape();
+  if (reps.size() < shape.size()) {
+    reps.insert(reps.begin(), shape.size() - reps.size(), 1);
+  }
+  if (reps.size() > shape.size()) {
+    shape.insert(shape.begin(), reps.size() - shape.size(), 1);
+  }
+
+  Shape expand_shape;
+  Shape broad_shape;
+  Shape final_shape;
+  for (int i = 0; i < shape.size(); i++) {
+    if (reps[i] != 1) {
+      expand_shape.push_back(1);
+      broad_shape.push_back(reps[i]);
+    }
+    expand_shape.push_back(shape[i]);
+    broad_shape.push_back(shape[i]);
+    final_shape.push_back(reps[i] * shape[i]);
+  }
+
+  auto x = reshape(arr, std::move(expand_shape), s);
+  x = broadcast_to(x, std::move(broad_shape), s);
+  return reshape(x, std::move(final_shape), s);
+}
+
+array edge_pad(
+    const array& a,
+    const std::vector<int>& axes,
+    const Shape& low_pad_size,
+    const Shape& high_pad_size,
+    const Shape& out_shape,
+    StreamOrDevice s /* = {}*/) {
+  array out = zeros(out_shape, a.dtype(), s);
+  auto stops = a.shape();
+  for (int i = 0; i < stops.size(); i++) {
+    stops[i] += low_pad_size[i];
+  }
+  // Copy over values from the unpadded array
+  array padded = slice_update(out, a, low_pad_size, stops, s);
+
+  for (int axis = 0; axis < a.ndim(); axis++) {
+    if (low_pad_size[axis] > 0) {
+      Shape starts(a.ndim(), 0);
+      starts[axis] = low_pad_size[axis];
+      auto stops = out.shape();
+      stops[axis] = low_pad_size[axis] + 1;
+      // Fetch edge values
+      array edge_value = slice(padded, starts, stops, s);
+
+      starts[axis] = 0;
+      stops[axis] = low_pad_size[axis];
+      // Update edge values in the padded array
+      padded = slice_update(padded, edge_value, starts, stops, s);
+    }
+
+    if (high_pad_size[axis] > 0) {
+      Shape starts(a.ndim(), 0);
+      starts[axis] = -high_pad_size[axis] - 1;
+      auto stops = out.shape();
+      stops[axis] = -high_pad_size[axis];
+      array edge_value = slice(padded, starts, stops, s);
+
+      starts[axis] = -high_pad_size[axis];
+      stops[axis] = out.shape(axis);
+      padded = slice_update(padded, edge_value, starts, stops, s);
+    }
+  }
+  return padded;
+}
+
+/** Pad an array with a constant value */
+array pad(
+    const array& a,
+    const std::vector<int>& axes,
+    const Shape& low_pad_size,
+    const Shape& high_pad_size,
+    const array& pad_value /*= array(0)*/,
+    const std::string& mode /*= "constant"*/,
+    StreamOrDevice s /* = {}*/) {
+  if (axes.size() != low_pad_size.size() ||
+      axes.size() != high_pad_size.size()) {
+    std::ostringstream msg;
+    msg << "Invalid number of padding sizes passed to pad "
+        << "with axes of size " << axes.size();
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto out_shape = a.shape();
+
+  for (int i = 0; i < axes.size(); i++) {
+    if (low_pad_size[i] < 0) {
+      std::ostringstream msg;
+      msg << "Invalid low padding size (" << low_pad_size[i]
+          << ") passed to pad for axis " << i
+          << ". Padding sizes must be non-negative";
+      throw std::invalid_argument(msg.str());
+    }
+    if (high_pad_size[i] < 0) {
+      std::ostringstream msg;
+      msg << "Invalid high padding size (" << high_pad_size[i]
+          << ") passed to pad for axis " << i
+          << ". Padding sizes must be non-negative";
+      throw std::invalid_argument(msg.str());
+    }
+
+    auto ax = axes[i] < 0 ? a.ndim() + axes[i] : axes[i];
+    out_shape[ax] += low_pad_size[i] + high_pad_size[i];
+  }
+
+  if (mode == "constant") {
+    return array(
+        std::move(out_shape),
+        a.dtype(),
+        std::make_shared<Pad>(to_stream(s), axes, low_pad_size, high_pad_size),
+        {a, astype(pad_value, a.dtype(), s)});
+  } else if (mode == "edge") {
+    return edge_pad(a, axes, low_pad_size, high_pad_size, out_shape, s);
+  } else {
+    std::ostringstream msg;
+    msg << "Invalid padding mode (" << mode << ") passed to pad";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+/** Pad an array with a constant value along all axes */
+array pad(
+    const array& a,
+    const std::vector<std::pair<int, int>>& pad_width,
+    const array& pad_value /*= array(0)*/,
+    const std::string& mode /*= "constant"*/,
+    StreamOrDevice s /*= {}*/) {
+  std::vector<int> axes(a.ndim(), 0);
+  std::iota(axes.begin(), axes.end(), 0);
+
+  Shape lows;
+  Shape highs;
+
+  for (auto& pads : pad_width) {
+    lows.push_back(pads.first);
+    highs.push_back(pads.second);
+  }
+
+  return pad(a, axes, lows, highs, pad_value, mode, s);
+}
+
+array pad(
+    const array& a,
+    const std::pair<int, int>& pad_width,
+    const array& pad_value /*= array(0)*/,
+    const std::string& mode /*= "constant"*/,
+    StreamOrDevice s /*= {}*/) {
+  return pad(
+      a,
+      std::vector<std::pair<int, int>>(a.ndim(), pad_width),
+      pad_value,
+      mode,
+      s);
+}
+
+array pad(
+    const array& a,
+    int pad_width,
+    const array& pad_value /*= array(0)*/,
+    const std::string& mode /*= "constant"*/,
+    StreamOrDevice s /*= {}*/) {
+  return pad(
+      a,
+      std::vector<std::pair<int, int>>(a.ndim(), {pad_width, pad_width}),
+      pad_value,
+      mode,
+      s);
+}
+
+array moveaxis(
+    const array& a,
+    int source,
+    int destination,
+    StreamOrDevice s /* = {} */) {
+  auto check_ax = [&a](int ax) {
+    auto ndim = static_cast<int>(a.ndim());
+    if (ax < -ndim || ax >= ndim) {
+      std::ostringstream msg;
+      msg << "[moveaxis] Invalid axis " << ax << " for array with " << ndim
+          << " dimensions.";
+      throw std::out_of_range(msg.str());
+    }
+    return ax < 0 ? ax + ndim : ax;
+  };
+  source = check_ax(source);
+  destination = check_ax(destination);
+  if (source == destination) {
+    return a;
+  }
+  std::vector<int> reorder(a.ndim());
+  std::iota(reorder.begin(), reorder.end(), 0);
+  reorder.erase(reorder.begin() + source);
+  reorder.insert(reorder.begin() + destination, source);
+  return transpose(a, reorder, s);
+}
+
+array swapaxes(
+    const array& a,
+    int axis1,
+    int axis2,
+    StreamOrDevice s /* = {} */) {
+  auto check_ax = [&a](int ax) {
+    auto ndim = static_cast<int>(a.ndim());
+    if (ax < -ndim || ax >= ndim) {
+      std::ostringstream msg;
+      msg << "[swapaxes] Invalid axis " << ax << " for array with " << ndim
+          << " dimensions.";
+      throw std::out_of_range(msg.str());
+    }
+    return ax < 0 ? ax + ndim : ax;
+  };
+  axis1 = check_ax(axis1);
+  axis2 = check_ax(axis2);
+  std::vector<int> reorder(a.ndim());
+  std::iota(reorder.begin(), reorder.end(), 0);
+  std::swap(reorder[axis1], reorder[axis2]);
+  return transpose(a, std::move(reorder), s);
+}
+
+array transpose(
+    const array& a,
+    std::vector<int> axes,
+    StreamOrDevice s /* = {} */) {
+  for (auto& ax : axes) {
+    ax = ax < 0 ? ax + a.ndim() : ax;
+  }
+  if (axes.size() != a.ndim()) {
+    std::ostringstream msg;
+    msg << "[transpose] Recived " << axes.size() << " axes for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Check in bounds and for duplicates
+  Shape shape(axes.size(), 0);
+  for (auto& ax : axes) {
+    if (ax < 0 || ax >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[transpose] Invalid axis (" << ax << ") for array with "
+          << a.ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    if (shape[ax] != 0) {
+      throw std::invalid_argument("[transpose] Repeat axes not allowed.");
+    }
+    shape[ax] = 1;
+  }
+
+  for (int i = 0; i < axes.size(); ++i) {
+    shape[i] = a.shape()[axes[i]];
+  }
+  return array(
+      std::move(shape),
+      a.dtype(),
+      std::make_shared<Transpose>(to_stream(s), std::move(axes)),
+      {a});
+}
+
+array transpose(const array& a, StreamOrDevice s /* = {} */) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.rbegin(), axes.rend(), 0);
+  return transpose(a, std::move(axes), to_stream(s));
+}
+
+array broadcast_to(
+    const array& a,
+    const Shape& shape,
+    StreamOrDevice s /* = {} */) {
+  if (a.shape() == shape) {
+    return a;
+  }
+
+  // Make sure the shapes are broadcastable
+  auto bxshape = broadcast_shapes(a.shape(), shape);
+  if (bxshape != shape) {
+    std::ostringstream msg;
+    msg << "Cannot broadcast array of shape " << a.shape() << " into shape "
+        << shape << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  return array(
+      std::move(bxshape),
+      a.dtype(),
+      std::make_shared<Broadcast>(to_stream(s), shape),
+      {a});
+}
+
+/** Broadcast the input arrays against one another while ignoring the
+ * axes specified in `ignore_axes`. Note, this API is internal only.
+ * The `ignore_axes` should be:
+ * - negative values indicating axes from the end
+ * - sorted in increasing order
+ */
+std::vector<array> broadcast_arrays(
+    const std::vector<array>& inputs,
+    std::vector<int> ignore_axes,
+    StreamOrDevice s) {
+  if (inputs.size() <= 1) {
+    return inputs;
+  }
+
+  std::vector<array> outputs;
+  auto shape = BroadcastAxes::output_shape(inputs, ignore_axes);
+  auto check_and_get_shape = [&shape, &ignore_axes](const array& in) {
+    auto out_shape = shape;
+    for (int i = 0; i < ignore_axes.size(); ++i) {
+      auto ax = ignore_axes[i];
+      auto pos_ax = in.ndim() + ax;
+      if (pos_ax < 0 || pos_ax > in.ndim() ||
+          (i > 0 && ax <= ignore_axes[i - 1])) {
+        throw std::invalid_argument(
+            "[broadcast_arrays] Received invalid axes to ignore.");
+      }
+      out_shape[out_shape.size() + ax] = in.shape(ax);
+    }
+    return out_shape;
+  };
+
+  if (!detail::in_dynamic_tracing()) {
+    for (auto& in : inputs) {
+      auto out_shape = check_and_get_shape(in);
+      if (in.shape() == out_shape) {
+        outputs.push_back(in);
+      } else {
+        outputs.push_back(array(
+            std::move(out_shape),
+            in.dtype(),
+            std::make_shared<Broadcast>(to_stream(s), out_shape),
+            {in}));
+      }
+    }
+    return outputs;
+  }
+
+  std::vector<array> stop_grad_inputs;
+  for (auto& in : inputs) {
+    stop_grad_inputs.push_back(stop_gradient(in, s));
+  }
+
+  for (int i = 0; i < inputs.size(); ++i) {
+    auto& in = inputs[i];
+    auto out_shape = check_and_get_shape(in);
+    if (in.shape() == out_shape) {
+      outputs.push_back(in);
+    } else {
+      // broadcasted array goes first followed by other stopgrad inputs
+      std::vector<array> p_inputs = {in};
+      for (int j = 0; j < inputs.size(); ++j) {
+        if (j == i) {
+          continue;
+        }
+        p_inputs.push_back(stop_grad_inputs[j]);
+      }
+      outputs.push_back(array(
+          std::move(out_shape),
+          in.dtype(),
+          std::make_shared<BroadcastAxes>(to_stream(s), ignore_axes),
+          std::move(p_inputs)));
+    }
+  }
+  return outputs;
+}
+
+std::vector<array> broadcast_arrays(
+    const std::vector<array>& inputs,
+    StreamOrDevice s /* = {} */) {
+  if (inputs.size() <= 1) {
+    return inputs;
+  }
+  auto shape = Broadcast::output_shape(inputs);
+  std::vector<array> outputs;
+
+  if (!detail::in_dynamic_tracing()) {
+    for (auto& in : inputs) {
+      if (in.shape() == shape) {
+        outputs.push_back(in);
+      } else {
+        outputs.push_back(array(
+            shape,
+            in.dtype(),
+            std::make_shared<Broadcast>(to_stream(s), shape),
+            {in}));
+      }
+    }
+    return outputs;
+  }
+
+  std::vector<array> stop_grad_inputs;
+  for (auto& in : inputs) {
+    stop_grad_inputs.push_back(stop_gradient(in, s));
+  }
+  for (int i = 0; i < inputs.size(); ++i) {
+    auto& in = inputs[i];
+    if (in.shape() == shape) {
+      outputs.push_back(in);
+    } else {
+      // broadcasted array goes first followed by other stopgrad inputs
+      std::vector<array> p_inputs = {in};
+      for (int j = 0; j < inputs.size(); ++j) {
+        if (j == i) {
+          continue;
+        }
+        p_inputs.push_back(stop_grad_inputs[j]);
+      }
+      outputs.push_back(array(
+          shape,
+          in.dtype(),
+          std::make_shared<Broadcast>(to_stream(s), shape),
+          std::move(p_inputs)));
+    }
+  }
+  return outputs;
+}
+
+std::pair<array, array>
+broadcast_arrays(const array& a, const array& b, StreamOrDevice s) {
+  auto out = broadcast_arrays({a, b}, s);
+  return {out[0], out[1]};
+}
+
+std::pair<array, array> broadcast_arrays(
+    const array& a,
+    const array& b,
+    std::vector<int> ignore_axes,
+    StreamOrDevice s) {
+  auto out = broadcast_arrays({a, b}, std::move(ignore_axes), s);
+  return {out[0], out[1]};
+}
+
+array equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape, bool_, std::make_shared<Equal>(to_stream(s)), std::move(inputs));
+}
+
+array not_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      bool_,
+      std::make_shared<NotEqual>(to_stream(s)),
+      std::move(inputs));
+}
+
+array greater(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape, bool_, std::make_shared<Greater>(to_stream(s)), std::move(inputs));
+}
+
+array greater_equal(
+    const array& a,
+    const array& b,
+    StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      bool_,
+      std::make_shared<GreaterEqual>(to_stream(s)),
+      std::move(inputs));
+}
+
+array less(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape, bool_, std::make_shared<Less>(to_stream(s)), std::move(inputs));
+}
+
+array less_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      bool_,
+      std::make_shared<LessEqual>(to_stream(s)),
+      std::move(inputs));
+}
+
+array array_equal(
+    const array& a,
+    const array& b,
+    bool equal_nan,
+    StreamOrDevice s /* = {} */) {
+  if (a.shape() != b.shape()) {
+    return array(false);
+  } else {
+    auto dtype = promote_types(a.dtype(), b.dtype());
+    equal_nan &= issubdtype(dtype, inexact);
+    return all(
+        array(
+            a.shape(),
+            bool_,
+            std::make_shared<Equal>(to_stream(s), equal_nan),
+            {astype(a, dtype, s), astype(b, dtype, s)}),
+        false,
+        s);
+  }
+}
+
+array isnan(const array& a, StreamOrDevice s /* = {} */) {
+  if (issubdtype(a.dtype(), integer) || a.dtype() == bool_) {
+    return full(a.shape(), false, bool_, s);
+  }
+  return not_equal(a, a, s);
+}
+
+array isinf(const array& a, StreamOrDevice s /* = {} */) {
+  if (issubdtype(a.dtype(), integer) || a.dtype() == bool_) {
+    return full(a.shape(), false, bool_, s);
+  }
+  return logical_or(isposinf(a, s), isneginf(a, s), s);
+}
+
+array isfinite(const array& a, StreamOrDevice s /* = {} */) {
+  if (issubdtype(a.dtype(), integer) || a.dtype() == bool_) {
+    return full(a.shape(), true, bool_, s);
+  }
+  return logical_not(logical_or(isinf(a, s), isnan(a, s), s), s);
+}
+
+array isposinf(const array& a, StreamOrDevice s /* = {} */) {
+  if (issubdtype(a.dtype(), integer) || a.dtype() == bool_) {
+    return full(a.shape(), false, bool_, s);
+  }
+  return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
+}
+
+array isneginf(const array& a, StreamOrDevice s /* = {} */) {
+  if (issubdtype(a.dtype(), integer) || a.dtype() == bool_) {
+    return full(a.shape(), false, bool_, s);
+  }
+  return equal(a, array(-std::numeric_limits<float>::infinity(), a.dtype()), s);
+}
+
+array where(
+    const array& a,
+    const array& b,
+    const array& c,
+    StreamOrDevice s /* = {} */) {
+  auto condition = astype(a, bool_, s);
+  Dtype out_dtype = promote_types(b.dtype(), c.dtype());
+  auto inputs = broadcast_arrays(
+      {condition, astype(b, out_dtype, s), astype(c, out_dtype, s)}, s);
+
+  return array(
+      inputs[0].shape(),
+      out_dtype,
+      std::make_shared<Select>(to_stream(s)),
+      inputs);
+}
+
+array nan_to_num(
+    const array& a,
+    float nan /* = 0.0f */,
+    const std::optional<float> posinf_ /* = std::nullopt */,
+    const std::optional<float> neginf_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  Dtype dtype = a.dtype();
+  if (!issubdtype(dtype, inexact)) {
+    return a;
+  }
+
+  auto type_to_max = [](const auto& dtype) -> float {
+    if (dtype == float32) {
+      return std::numeric_limits<float>::max();
+    } else if (dtype == bfloat16) {
+      return std::numeric_limits<bfloat16_t>::max();
+    } else if (dtype == float16) {
+      return std::numeric_limits<float16_t>::max();
+    } else {
+      std::ostringstream msg;
+      msg << "[nan_to_num] Does not yet support given type: " << dtype << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  };
+
+  float posinf = posinf_ ? *posinf_ : type_to_max(dtype);
+  float neginf = neginf_ ? *neginf_ : -type_to_max(dtype);
+
+  auto out = where(isnan(a, s), array(nan, dtype), a, s);
+  out = where(isposinf(a, s), array(posinf, dtype), out, s);
+  out = where(isneginf(a, s), array(neginf, dtype), out, s);
+  return out;
+}
+
+array allclose(
+    const array& a,
+    const array& b,
+    double rtol /* = 1e-5 */,
+    double atol /* = 1e-8 */,
+    bool equal_nan /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  return all(isclose(a, b, rtol, atol, equal_nan, s), s);
+}
+
+array isclose(
+    const array& a,
+    const array& b,
+    double rtol /* = 1e-5 */,
+    double atol /* = 1e-8 */,
+    bool equal_nan /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  // |a - b| <= atol + rtol * |b|
+  auto rhs = add(array(atol), multiply(array(rtol), abs(b, s), s), s);
+  auto lhs = abs(subtract(a, b, s), s);
+  auto out = less_equal(lhs, rhs, s);
+
+  // Correct the result for infinite values.
+  auto a_pos_inf = isposinf(a, s);
+  auto b_pos_inf = isposinf(b, s);
+  auto a_neg_inf = isneginf(a, s);
+  auto b_neg_inf = isneginf(b, s);
+  auto any_inf = logical_or(
+      logical_or(a_pos_inf, a_neg_inf, s),
+      logical_or(b_pos_inf, b_neg_inf, s),
+      s);
+  auto both_inf = logical_or(
+      logical_and(a_pos_inf, b_pos_inf, s),
+      logical_and(a_neg_inf, b_neg_inf, s),
+      s);
+
+  // Convert all elements where either value is infinite to False.
+  out = logical_and(out, logical_not(any_inf, s), s);
+
+  // Convert all the elements where both values are infinite and of the same
+  // sign to True.
+  out = logical_or(out, both_inf, s);
+
+  if (equal_nan) {
+    auto both_nan = logical_and(isnan(a, s), isnan(b, s), s);
+    out = logical_or(out, both_nan, s);
+  }
+
+  return out;
+}
+
+array all(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return all(a, axes, keepdims, s);
+}
+
+array all(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape(axes, a.shape());
+  auto out = (is_noop)
+      ? astype(a, bool_, s)
+      : array(
+            std::move(out_shape),
+            bool_,
+            std::make_shared<Reduce>(to_stream(s), Reduce::And, sorted_axes),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
+}
+
+array all(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return all(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array any(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return any(a, axes, keepdims, s);
+}
+
+array any(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape(axes, a.shape());
+  auto out = (is_noop)
+      ? astype(a, bool_, s)
+      : array(
+            std::move(out_shape),
+            bool_,
+            std::make_shared<Reduce>(to_stream(s), Reduce::Or, sorted_axes),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
+}
+
+array any(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return any(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array sum(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return sum(a, axes, keepdims, s);
+}
+
+array sum(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  if (axes.empty()) {
+    return a;
+  }
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape(axes, a.shape());
+  Dtype out_type = a.dtype();
+  if (issubdtype(a.dtype(), signedinteger)) {
+    out_type = a.dtype().size() <= 4 ? int32 : int64;
+  } else if (issubdtype(a.dtype(), unsignedinteger)) {
+    out_type = a.dtype().size() <= 4 ? uint32 : uint64;
+  } else if (a.dtype() == bool_) {
+    out_type = int32;
+  }
+  auto out = (is_noop)
+      ? astype(a, out_type, s)
+      : array(
+            std::move(out_shape),
+            out_type,
+            std::make_shared<Reduce>(to_stream(s), Reduce::Sum, sorted_axes),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
+}
+
+array sum(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return sum(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array mean(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return mean(a, axes, keepdims, to_stream(s));
+}
+
+array mean(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  int ndim = a.ndim();
+  for (int axis : axes) {
+    if (axis < -ndim || axis >= ndim) {
+      std::ostringstream msg;
+      msg << "[mean] axis " << axis << " is out of bounds for array with "
+          << ndim << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  auto dtype = at_least_float(a.dtype());
+  auto normalizer = number_of_elements(a, axes, true, dtype, s);
+  return multiply(sum(a, axes, keepdims, s), normalizer, s);
+}
+
+array mean(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return mean(a, std::vector<int>{axis}, keepdims, to_stream(s));
+}
+
+array median(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return median(a, axes, keepdims, to_stream(s));
+}
+
+array median(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  int ndim = a.ndim();
+  std::set<int> set_axes;
+  for (int axis : axes) {
+    if (axis < -ndim || axis >= ndim) {
+      std::ostringstream msg;
+      msg << "[median] axis " << axis << " is out of bounds for array with "
+          << ndim << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    set_axes.insert(axis < 0 ? axis + ndim : axis);
+  }
+  if (set_axes.size() != axes.size()) {
+    throw std::invalid_argument("[median] Received duplicate axis.");
+  }
+  std::vector<int> sorted_axes(set_axes.begin(), set_axes.end());
+  auto dtype = at_least_float(a.dtype());
+  std::vector<int> transpose_axes;
+  for (int i = 0, j = 0; i < a.ndim(); ++i) {
+    if (j < sorted_axes.size() && i == sorted_axes[j]) {
+      j++;
+      continue;
+    }
+    transpose_axes.push_back(i);
+  }
+  int flat_start = transpose_axes.size();
+  transpose_axes.insert(
+      transpose_axes.end(), sorted_axes.begin(), sorted_axes.end());
+
+  // Move all the median axes to the back and flatten
+  auto flat_a =
+      flatten(transpose(a, transpose_axes, s), flat_start, a.ndim(), s);
+  int flat_size = flat_a.shape(-1);
+  if (flat_size == 0) {
+    throw std::invalid_argument(
+        "[median] Cannot take median along empty axis.");
+  }
+
+  // Sort the last axis
+  auto sorted_a = sort(flat_a, -1, s);
+
+  // Take the midpoint
+  auto mp = flat_size / 2;
+  auto start = Shape(sorted_a.ndim(), 0);
+  auto stop = sorted_a.shape();
+  start.back() = mp;
+  stop.back() = mp + 1;
+  auto median_a = astype(slice(sorted_a, start, stop, s), dtype, s);
+  if (flat_size % 2 == 0) {
+    start.back() = mp - 1;
+    stop.back() = mp;
+    median_a = multiply(
+        add(median_a, astype(slice(sorted_a, start, stop, s), dtype, s), s),
+        array(0.5, dtype),
+        s);
+  }
+  median_a = squeeze(median_a, -1, s);
+  if (keepdims) {
+    median_a = expand_dims(median_a, sorted_axes, s);
+  }
+  return median_a;
+}
+
+array median(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return median(a, std::vector<int>{axis}, keepdims, to_stream(s));
+}
+
+array var(
+    const array& a,
+    bool keepdims,
+    int ddof /* = 0*/,
+    StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return var(a, axes, keepdims, ddof, to_stream(s));
+}
+
+array var(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    int ddof /* = 0*/,
+    StreamOrDevice s /* = {}*/) {
+  auto dtype = at_least_float(a.dtype());
+  auto mu = mean(a, axes, /* keepdims= */ true, s);
+  auto v = sum(square(subtract(a, mu, s), s), axes, keepdims, s);
+
+  if (ddof != 0) {
+    auto normalizer = maximum(
+        subtract(
+            number_of_elements(a, axes, false, dtype, s),
+            array(ddof, dtype),
+            s),
+        array(0, dtype),
+        s);
+    v = divide(v, normalizer, s);
+  } else {
+    auto normalizer = number_of_elements(a, axes, true, dtype, s);
+    v = multiply(v, normalizer, s);
+  }
+
+  return v;
+}
+
+array var(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    int ddof /* = 0*/,
+    StreamOrDevice s /* = {} */) {
+  return var(a, std::vector<int>{axis}, keepdims, ddof, to_stream(s));
+}
+
+array std(
+    const array& a,
+    bool keepdims,
+    int ddof /* = 0*/,
+    StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return std(a, axes, keepdims, ddof, to_stream(s));
+}
+
+array std(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    int ddof /* = 0*/,
+    StreamOrDevice s /* = {}*/) {
+  return sqrt(var(a, axes, keepdims, ddof, s), s);
+}
+
+array std(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    int ddof /* = 0*/,
+    StreamOrDevice s /* = {} */) {
+  return std(a, std::vector<int>{axis}, keepdims, ddof, to_stream(s));
+}
+
+array prod(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return prod(a, axes, keepdims, s);
+}
+
+array prod(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  if (axes.empty()) {
+    return a;
+  }
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape(axes, a.shape());
+  Dtype out_type = a.dtype();
+  if (issubdtype(a.dtype(), signedinteger)) {
+    out_type = a.dtype().size() <= 4 ? int32 : int64;
+  } else if (issubdtype(a.dtype(), unsignedinteger)) {
+    out_type = a.dtype().size() <= 4 ? uint32 : uint64;
+  } else if (a.dtype() == bool_) {
+    out_type = int32;
+  }
+  auto out = (is_noop)
+      ? a
+      : array(
+            std::move(out_shape),
+            out_type,
+            std::make_shared<Reduce>(to_stream(s), Reduce::Prod, sorted_axes),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
+}
+
+array prod(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return prod(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array max(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return max(a, axes, keepdims, s);
+}
+
+array max(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  if (a.size() == 0) {
+    throw std::invalid_argument("[max] Cannot max reduce zero size array.");
+  }
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape(axes, a.shape());
+  auto out = (is_noop)
+      ? a
+      : array(
+            std::move(out_shape),
+            a.dtype(),
+            std::make_shared<Reduce>(to_stream(s), Reduce::Max, sorted_axes),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
+}
+
+array max(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return max(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array min(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return min(a, axes, keepdims, s);
+}
+
+array min(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  if (a.size() == 0) {
+    throw std::invalid_argument("[min] Cannot min reduce zero size array.");
+  }
+  if (axes.empty()) {
+    return a;
+  }
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape(axes, a.shape());
+  auto out = (is_noop)
+      ? a
+      : array(
+            std::move(out_shape),
+            a.dtype(),
+            std::make_shared<Reduce>(to_stream(s), Reduce::Min, sorted_axes),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes, s);
+  }
+  return out;
+}
+
+array min(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return min(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array argmin(const array& a, bool keepdims, StreamOrDevice s /* = {} */) {
+  auto result = argmin(flatten(a, s), 0, true, s);
+  if (keepdims) {
+    std::vector<int> axes(a.ndim() - 1);
+    std::iota(axes.begin(), axes.end(), 0);
+    result = expand_dims(result, axes, s);
+  } else {
+    result = squeeze(result, s);
+  }
+  return result;
+}
+
+array argmin(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  if (a.size() == 0) {
+    throw std::invalid_argument(
+        "[argmin] Cannot argmin reduce zero size array.");
+  }
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape({axis}, a.shape());
+  auto out = (is_noop)
+      ? zeros(out_shape, uint32, s)
+      : array(
+            std::move(out_shape),
+            uint32,
+            std::make_shared<ArgReduce>(
+                to_stream(s), ArgReduce::ArgMin, sorted_axes[0]),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes[0], s);
+  }
+  return out;
+}
+
+array argmax(const array& a, bool keepdims, StreamOrDevice s /* = {} */) {
+  auto result = argmax(flatten(a, s), 0, true, s);
+  if (keepdims) {
+    std::vector<int> axes(a.ndim() - 1);
+    std::iota(axes.begin(), axes.end(), 0);
+    result = expand_dims(result, axes, s);
+  } else {
+    result = squeeze(result, s);
+  }
+  return result;
+}
+
+array argmax(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  if (a.size() == 0) {
+    throw std::invalid_argument(
+        "[argmax] Cannot argmax reduce zero size array.");
+  }
+  auto [out_shape, sorted_axes, is_noop] =
+      compute_reduce_shape({axis}, a.shape());
+  auto out = (is_noop)
+      ? zeros(out_shape, uint32, s)
+      : array(
+            std::move(out_shape),
+            uint32,
+            std::make_shared<ArgReduce>(
+                to_stream(s), ArgReduce::ArgMax, sorted_axes[0]),
+            {a});
+  if (!keepdims) {
+    out = squeeze(out, sorted_axes[0], s);
+  }
+  return out;
+}
+
+/** Returns a sorted copy of the flattened array. */
+array sort(const array& a, StreamOrDevice s /* = {} */) {
+  int size = a.size();
+  return sort(reshape(a, {size}, s), 0, s);
+}
+
+/** Returns a sorted copy of the array along a given axis. */
+array sort(const array& a, int axis, StreamOrDevice s /* = {} */) {
+  // Check for valid axis
+  if (axis + static_cast<int>(a.ndim()) < 0 ||
+      axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[sort] Received invalid axis " << axis << " for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  return array(
+      a.shape(), a.dtype(), std::make_shared<Sort>(to_stream(s), axis), {a});
+}
+
+/** Returns indices that sort the flattened array. */
+array argsort(const array& a, StreamOrDevice s /* = {} */) {
+  int size = a.size();
+  return argsort(reshape(a, {size}, s), 0, s);
+}
+
+/** Returns indices that sort the array along a given axis. */
+array argsort(const array& a, int axis, StreamOrDevice s /* = {} */) {
+  // Check for valid axis
+  if (axis + static_cast<int>(a.ndim()) < 0 ||
+      axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[argsort] Received invalid axis " << axis << " for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  return array(
+      a.shape(), uint32, std::make_shared<ArgSort>(to_stream(s), axis), {a});
+}
+
+/**
+ * Returns a partitioned copy of the flattened array
+ * such that the smaller kth elements are first.
+ **/
+array partition(const array& a, int kth, StreamOrDevice s /* = {} */) {
+  int size = a.size();
+  return partition(reshape(a, {size}, s), kth, 0, s);
+}
+
+/**
+ * Returns a partitioned copy of the array along a given axis
+ * such that the smaller kth elements are first.
+ **/
+array partition(
+    const array& a,
+    int kth,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  // Check for valid axis
+  if (axis + static_cast<int>(a.ndim()) < 0 ||
+      axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[partition] Received invalid axis " << axis << " for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  int axis_ = axis < 0 ? axis + a.ndim() : axis;
+  int kth_ = kth < 0 ? kth + a.shape(axis) : kth;
+  if (kth_ < 0 || kth_ >= a.shape(axis_)) {
+    std::ostringstream msg;
+    msg << "[partition] Received invalid kth " << kth << "along axis " << axis
+        << " for array with shape: " << a.shape();
+    throw std::invalid_argument(msg.str());
+  }
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Partition>(to_stream(s), kth_, axis_),
+      {a});
+}
+
+/**
+ * Returns indices that partition the flattened array
+ * such that the smaller kth elements are first.
+ **/
+array argpartition(const array& a, int kth, StreamOrDevice s /* = {} */) {
+  int size = a.size();
+  return argpartition(reshape(a, {size}, s), kth, 0, s);
+}
+
+/**
+ * Returns indices that partition the array along a given axis
+ * such that the smaller kth elements are first.
+ **/
+array argpartition(
+    const array& a,
+    int kth,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  // Check for valid axis
+  if (axis + static_cast<int>(a.ndim()) < 0 ||
+      axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[argpartition] Received invalid axis " << axis << " for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  int axis_ = axis < 0 ? axis + a.ndim() : axis;
+  int kth_ = kth < 0 ? kth + a.shape(axis) : kth;
+  if (kth_ < 0 || kth_ >= a.shape(axis_)) {
+    std::ostringstream msg;
+    msg << "[argpartition] Received invalid kth " << kth << " along axis "
+        << axis << " for array with shape: " << a.shape();
+    throw std::invalid_argument(msg.str());
+  }
+  return array(
+      a.shape(),
+      uint32,
+      std::make_shared<ArgPartition>(to_stream(s), kth_, axis_),
+      {a});
+}
+
+/** Returns topk elements of the flattened array. */
+array topk(const array& a, int k, StreamOrDevice s /* = {}*/) {
+  int size = a.size();
+  return topk(reshape(a, {size}, s), k, 0, s);
+}
+
+/** Returns topk elements of the array along a given axis. */
+array topk(const array& a, int k, int axis, StreamOrDevice s /* = {}*/) {
+  // Check for valid axis
+  int axis_ = axis < 0 ? axis + a.ndim() : axis;
+  if (axis_ < 0 || axis_ >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[topk] Received invalid axis " << axis << " for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (k < 0 || k > a.shape(axis_)) {
+    std::ostringstream msg;
+    msg << "[topk] Received invalid k=" << k << " along axis " << axis
+        << " for array with shape: " << a.shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Return early if the whole input was requested.
+  if (k == a.shape(axis_)) {
+    return a;
+  }
+
+  array a_partitioned = partition(a, -k, axis_, s);
+  Shape slice_starts(a.ndim(), 0);
+  auto slice_ends = a.shape();
+  slice_starts[axis_] = a.shape(axis_) - k;
+  return slice(a_partitioned, slice_starts, slice_ends, s);
+}
+
+array logsumexp(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return logsumexp(a, axes, keepdims, s);
+}
+
+array logsumexp(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  if (a.size() == 0) {
+    throw std::invalid_argument("[logsumexp] Received empty array.");
+  }
+  if (a.ndim() == 0 && !axes.empty()) {
+    throw std::invalid_argument(
+        "[logsumexp] Received non-empty axes for array with 0 dimensions.");
+  }
+  bool reduce_last_dim =
+      !axes.empty() && (axes.back() == a.ndim() - 1 || axes.back() == -1);
+  if (reduce_last_dim) {
+    // For more than 2 axes check if axes is [0, 1, ..., NDIM - 1] and shape
+    // is [1, 1, ..., N].
+    for (int i = axes.size() - 2; i >= 0; --i) {
+      if ((axes[i] + 1 != axes[i + 1]) || (a.shape(axes[i]) != 1)) {
+        reduce_last_dim = false;
+        break;
+      }
+    }
+  }
+  bool is_complex = issubdtype(a.dtype(), complexfloating);
+  if (!is_complex && reduce_last_dim) {
+    auto dtype = at_least_float(a.dtype());
+    auto out_shape = a.shape();
+    out_shape.back() = 1;
+    auto out = array(
+        std::move(out_shape),
+        dtype,
+        std::make_shared<LogSumExp>(to_stream(s)),
+        {astype(a, dtype, s)});
+    if (!keepdims) {
+      out = squeeze(out, -1, s);
+    }
+    return out;
+  }
+  auto maxval = stop_gradient(max(a, axes, true, s), s);
+  auto out = log(sum(exp(subtract(a, maxval, s), s), axes, keepdims, s), s);
+  out = add(out, reshape(maxval, out.shape(), s), s);
+  if (!keepdims) {
+    maxval = squeeze(maxval, axes, s);
+  }
+  return where(isinf(maxval, s), maxval, out, s);
+}
+
+array logsumexp(
+    const array& a,
+    int axis,
+    bool keepdims /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return logsumexp(a, std::vector<int>{axis}, keepdims, s);
+}
+
+array abs(const array& a, StreamOrDevice s /* = {} */) {
+  auto out =
+      array(a.shape(), a.dtype(), std::make_shared<Abs>(to_stream(s)), {a});
+  if (a.dtype() == complex64) {
+    out = astype(out, float32, s);
+  }
+  return out;
+}
+
+array negative(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() == bool_) {
+    auto msg = "[negative] Not supported for bool, use logical_not instead.";
+    throw std::invalid_argument(msg);
+  }
+  return array(
+      a.shape(), a.dtype(), std::make_shared<Negative>(to_stream(s)), {a});
+}
+array operator-(const array& a) {
+  return negative(a);
+}
+
+array sign(const array& a, StreamOrDevice s /* = {} */) {
+  return array(a.shape(), a.dtype(), std::make_shared<Sign>(to_stream(s)), {a});
+}
+
+array logical_not(const array& a, StreamOrDevice s /* = {} */) {
+  return array(
+      a.shape(),
+      bool_,
+      std::make_shared<LogicalNot>(to_stream(s)),
+      {astype(a, bool_, s)});
+}
+
+array logical_and(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  // Broadcast arrays to a common shape
+  auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      bool_,
+      std::make_shared<LogicalAnd>(to_stream(s)),
+      std::move(inputs));
+}
+array operator&&(const array& a, const array& b) {
+  return logical_and(a, b);
+}
+
+array logical_or(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  // Broadcast arrays to a common shape
+  auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      bool_,
+      std::make_shared<LogicalOr>(to_stream(s)),
+      std::move(inputs));
+}
+array operator||(const array& a, const array& b) {
+  return logical_or(a, b);
+}
+
+array reciprocal(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  return divide(array(1.0f, dtype), a, to_stream(s));
+}
+
+array add(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto out_type = promote_types(a.dtype(), b.dtype());
+  auto inputs =
+      broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape, out_type, std::make_shared<Add>(to_stream(s)), std::move(inputs));
+}
+
+array operator+(const array& a, const array& b) {
+  return add(a, b);
+}
+
+array subtract(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto out_type = promote_types(a.dtype(), b.dtype());
+  auto inputs =
+      broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      out_type,
+      std::make_shared<Subtract>(to_stream(s)),
+      std::move(inputs));
+}
+
+array operator-(const array& a, const array& b) {
+  return subtract(a, b);
+}
+
+array multiply(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto out_type = promote_types(a.dtype(), b.dtype());
+  auto inputs =
+      broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      out_type,
+      std::make_shared<Multiply>(to_stream(s)),
+      std::move(inputs));
+}
+
+array operator*(const array& a, const array& b) {
+  return multiply(a, b);
+}
+
+array divide(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(promote_types(a.dtype(), b.dtype()));
+  auto inputs = broadcast_arrays(
+      {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape, dtype, std::make_shared<Divide>(to_stream(s)), std::move(inputs));
+}
+array operator/(const array& a, const array& b) {
+  return divide(a, b);
+}
+array operator/(double a, const array& b) {
+  return divide(array(a), b);
+}
+array operator/(const array& a, double b) {
+  return divide(a, array(b));
+}
+
+array floor_divide(
+    const array& a,
+    const array& b,
+    StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  if (issubdtype(dtype, inexact)) {
+    return floor(divide(a, b, s), s);
+  }
+
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape, dtype, std::make_shared<Divide>(to_stream(s)), std::move(inputs));
+}
+
+array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  auto inputs = broadcast_arrays(
+      {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      dtype,
+      std::make_shared<Remainder>(to_stream(s)),
+      std::move(inputs));
+}
+array operator%(const array& a, const array& b) {
+  return remainder(a, b);
+}
+
+std::vector<array>
+divmod(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  if (issubdtype(dtype, complexfloating)) {
+    throw std::invalid_argument("[divmod] Complex type not supported.");
+  }
+  auto inputs = broadcast_arrays(
+      {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s);
+  return array::make_arrays(
+      {inputs[0].shape(), inputs[0].shape()},
+      {inputs[0].dtype(), inputs[0].dtype()},
+      std::make_shared<DivMod>(to_stream(s)),
+      inputs);
+}
+
+array maximum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto out_type = promote_types(a.dtype(), b.dtype());
+  auto inputs =
+      broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      out_type,
+      std::make_shared<Maximum>(to_stream(s)),
+      std::move(inputs));
+}
+
+array minimum(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto out_type = promote_types(a.dtype(), b.dtype());
+  auto inputs =
+      broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      out_type,
+      std::make_shared<Minimum>(to_stream(s)),
+      std::move(inputs));
+}
+
+array floor(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() == complex64) {
+    throw std::invalid_argument("[floor] Not supported for complex64.");
+  }
+  return array(
+      a.shape(), a.dtype(), std::make_shared<Floor>(to_stream(s)), {a});
+}
+
+array ceil(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.dtype() == complex64) {
+    throw std::invalid_argument("[floor] Not supported for complex64.");
+  }
+  return array(a.shape(), a.dtype(), std::make_shared<Ceil>(to_stream(s)), {a});
+}
+
+array square(const array& a, StreamOrDevice s /* = {} */) {
+  return array(
+      a.shape(), a.dtype(), std::make_shared<Square>(to_stream(s)), {a});
+}
+
+array exp(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(a.shape(), dtype, std::make_shared<Exp>(to_stream(s)), {input});
+}
+
+array expm1(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<Expm1>(to_stream(s)), {input});
+}
+
+array sin(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(a.shape(), dtype, std::make_shared<Sin>(to_stream(s)), {input});
+}
+
+array cos(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(a.shape(), dtype, std::make_shared<Cos>(to_stream(s)), {input});
+}
+
+array tan(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(a.shape(), dtype, std::make_shared<Tan>(to_stream(s)), {input});
+}
+
+array arcsin(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<ArcSin>(to_stream(s)), {input});
+}
+
+array arccos(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<ArcCos>(to_stream(s)), {input});
+}
+
+array arctan(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<ArcTan>(to_stream(s)), {input});
+}
+
+array arctan2(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(promote_types(a.dtype(), b.dtype()));
+  auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape, dtype, std::make_shared<ArcTan2>(to_stream(s)), std::move(inputs));
+}
+
+array sinh(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(a.shape(), dtype, std::make_shared<Sinh>(to_stream(s)), {input});
+}
+
+array cosh(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(a.shape(), dtype, std::make_shared<Cosh>(to_stream(s)), {input});
+}
+
+array tanh(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(a.shape(), dtype, std::make_shared<Tanh>(to_stream(s)), {input});
+}
+
+array arcsinh(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<ArcSinh>(to_stream(s)), {input});
+}
+
+array arccosh(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<ArcCosh>(to_stream(s)), {input});
+}
+
+array arctanh(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<ArcTanh>(to_stream(s)), {input});
+}
+
+array degrees(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  return multiply(a, array(180.0 / M_PI, dtype), s);
+}
+
+array radians(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  return multiply(a, array(M_PI / 180.0, dtype), s);
+}
+
+array log(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<Log>(to_stream(s), Log::Base::e),
+      {input});
+}
+
+array log2(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<Log>(to_stream(s), Log::Base::two),
+      {input});
+}
+
+array log10(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<Log>(to_stream(s), Log::Base::ten),
+      {input});
+}
+
+array log1p(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<Log1p>(to_stream(s)), {input});
+}
+
+array logaddexp(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  // Make sure out type is floating point
+  auto out_type = at_least_float(promote_types(a.dtype(), b.dtype()));
+  auto inputs =
+      broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
+  auto& shape = inputs[0].shape();
+  return array(
+      shape,
+      out_type,
+      std::make_shared<LogAddExp>(to_stream(s)),
+      std::move(inputs));
+}
+
+array sigmoid(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return array(
+      a.shape(), dtype, std::make_shared<Sigmoid>(to_stream(s)), {input});
+}
+
+array erf(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<Erf>(to_stream(s)),
+      {astype(a, dtype, s)});
+}
+
+array erfinv(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<ErfInv>(to_stream(s)),
+      {astype(a, dtype, s)});
+}
+
+array stop_gradient(const array& a, StreamOrDevice s /* = {} */) {
+  return array(
+      a.shape(), a.dtype(), std::make_shared<StopGradient>(to_stream(s)), {a});
+}
+
+array round(const array& a, int decimals, StreamOrDevice s /* = {} */) {
+  if (decimals == 0) {
+    return array(
+        a.shape(), a.dtype(), std::make_shared<Round>(to_stream(s)), {a});
+  }
+
+  auto dtype = at_least_float(a.dtype());
+  float scale = std::pow(10, decimals);
+  auto result = multiply(a, array(scale, dtype), s);
+  result = round(result, 0, s);
+  result = multiply(result, array(1 / scale, dtype), s);
+
+  return astype(result, a.dtype(), s);
+}
+
+array matmul(
+    const array& in_a,
+    const array& in_b,
+    StreamOrDevice s /* = {} */) {
+  auto a = in_a;
+  auto b = in_b;
+  if (a.ndim() == 0 || b.ndim() == 0) {
+    throw std::invalid_argument(
+        "[matmul] Got 0 dimension input. Inputs must "
+        "have at least one dimension.");
+  }
+
+  if (a.ndim() == 1) {
+    // Insert a singleton dim in the beginning
+    a = expand_dims(a, 0, s);
+  }
+  if (b.ndim() == 1) {
+    // Insert a singleton dim at the end
+    b = expand_dims(b, 1, s);
+  }
+  if (a.shape(-1) != b.shape(-2)) {
+    std::ostringstream msg;
+    msg << "[matmul] Last dimension of first input with shape " << a.shape()
+        << " must match second to last dimension of"
+        << " second input with shape " << b.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Type promotion
+  auto out_type = promote_types(a.dtype(), b.dtype());
+
+  if (!issubdtype(out_type, inexact)) {
+    std::ostringstream msg;
+    msg << "[matmul] Only inexact types are supported but " << a.dtype()
+        << " and " << b.dtype() << " were provided which results"
+        << " in " << out_type << ", which is not a floating point type.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (a.dtype() != out_type) {
+    a = astype(a, out_type, s);
+  }
+  if (b.dtype() != out_type) {
+    b = astype(b, out_type, s);
+  }
+
+  // We can batch the multiplication by reshaping a
+  if (in_a.ndim() > 2 && in_b.ndim() <= 2) {
+    a = flatten(a, 0, -2, s);
+  } else if (in_b.ndim() > 2) {
+    std::tie(a, b) = broadcast_arrays(a, b, {-2, -1}, s);
+  }
+
+  auto out_shape = a.shape();
+  out_shape.back() = b.shape(-1);
+
+  auto out = array(
+      std::move(out_shape),
+      out_type,
+      std::make_shared<Matmul>(to_stream(s)),
+      {a, b});
+  if (in_a.ndim() > 2 && in_b.ndim() <= 2) {
+    auto orig_shape = in_a.shape();
+    orig_shape.pop_back();
+    out = unflatten(out, 0, std::move(orig_shape), s);
+  }
+
+  // Remove the possibly inserted singleton dimensions
+  std::vector<int> axes;
+  if (in_a.ndim() == 1) {
+    axes.push_back(out.ndim() - 2);
+  }
+  if (in_b.ndim() == 1) {
+    axes.push_back(out.ndim() - 1);
+  }
+  return axes.empty() ? out : squeeze(out, axes, s);
+}
+
+array gather(
+    const array& a,
+    const std::vector<array>& indices,
+    const std::vector<int>& axes,
+    const Shape& slice_sizes,
+    StreamOrDevice s /* = {} */) {
+  // Checks that indices, dimensions, and slice_sizes are all valid
+  if (indices.size() > a.ndim()) {
+    std::ostringstream msg;
+    msg << "[gather] Too many index arrays. Got " << indices.size()
+        << " index arrays for input with " << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  std::set dims(axes.begin(), axes.end());
+  if (dims.size() != axes.size()) {
+    throw std::invalid_argument("[gather] Repeat axes not allowed in gather.");
+  }
+  if (!dims.empty() && (*dims.begin() < 0 || *dims.rbegin() >= a.ndim())) {
+    throw std::invalid_argument("[gather] Axes don't match array dimensions.");
+  }
+  if (indices.size() != axes.size()) {
+    throw std::invalid_argument(
+        "[gather] Number of index arrays does not match number of axes.");
+  }
+  for (auto& x : indices) {
+    if (x.dtype() == bool_) {
+      throw std::invalid_argument("[Gather] Boolean indices not supported.");
+    }
+  }
+
+  if (slice_sizes.size() != a.ndim()) {
+    std::ostringstream msg;
+    msg << "[gather] Got slice_sizes with size " << slice_sizes.size()
+        << " for array with " << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  // Promote indices to the same type
+  auto dtype = result_type(indices);
+  if (issubdtype(dtype, inexact)) {
+    throw std::invalid_argument(
+        "[gather] Got indices with invalid dtype. Indices must be integral.");
+  }
+
+  // Broadcast and cast indices if necessary
+  auto inputs = broadcast_arrays(indices);
+  for (auto& idx : inputs) {
+    idx = astype(idx, dtype, s);
+  }
+
+  if (a.size() == 0) {
+    // Empty input, either the total slice size is 0 or the indices are empty
+    auto total_slice = std::accumulate(
+        slice_sizes.begin(), slice_sizes.end(), 1, std::multiplies<int64_t>{});
+    auto idx_size = !inputs.empty() ? inputs[0].size() : 1;
+    if (idx_size != 0 && total_slice != 0) {
+      std::ostringstream msg;
+      msg << "[gather] If the input is empty, either the indices must be"
+          << " empty or the total slice size must be 0.";
+      throw std::invalid_argument(msg.str());
+    }
+  } else {
+    // Non-empty input, check slice sizes are valid
+    for (int i = 0; i < a.ndim(); ++i) {
+      if (slice_sizes[i] < 0 || slice_sizes[i] > a.shape(i)) {
+        std::ostringstream msg;
+        msg << "[gather] Slice sizes must be in [0, a.shape(i)]. Got "
+            << slice_sizes << " for array with shape " << a.shape() << ".";
+        throw std::invalid_argument(msg.str());
+      }
+    }
+  }
+
+  Shape out_shape;
+  if (!inputs.empty()) {
+    out_shape = inputs[0].shape();
+  }
+  out_shape.insert(out_shape.end(), slice_sizes.begin(), slice_sizes.end());
+
+  inputs.insert(inputs.begin(), a);
+  return array(
+      std::move(out_shape),
+      a.dtype(),
+      std::make_shared<Gather>(
+          to_stream(s), std::move(axes), std::move(slice_sizes)),
+      inputs);
+}
+
+array kron(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  if (a.size() == 0 || b.size() == 0) {
+    throw std::invalid_argument("[kron] Input arrays cannot be empty.");
+  }
+
+  int ndim = std::max(a.ndim(), b.ndim());
+  Shape a_shape(2 * ndim, 1);
+  Shape b_shape(2 * ndim, 1);
+  Shape out_shape(ndim, 1);
+
+  for (int i = ndim - 1, j = a.ndim() - 1; j >= 0; j--, i--) {
+    a_shape[2 * i] = a.shape(j);
+    out_shape[i] *= a.shape(j);
+  }
+  for (int i = ndim - 1, j = b.ndim() - 1; j >= 0; j--, i--) {
+    b_shape[2 * i + 1] = b.shape(j);
+    out_shape[i] *= b.shape(j);
+  }
+
+  return reshape(
+      multiply(
+          reshape(a, std::move(a_shape), s),
+          reshape(b, std::move(b_shape), s),
+          s),
+      std::move(out_shape),
+      s);
+}
+
+array take(
+    const array& a,
+    const array& indices,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  // Check for valid axis
+  if (axis + static_cast<int>(a.ndim()) < 0 ||
+      axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[take] Received invalid axis " << axis << " for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Check for valid take
+  if (a.shape(axis) == 0 && indices.size() != 0) {
+    throw std::invalid_argument(
+        "[take] Cannot do a non-empty take from an empty axis.");
+  }
+
+  // Handle negative axis
+  axis = axis < 0 ? a.ndim() + axis : axis;
+
+  // Make slice sizes to pass to gather
+  Shape slice_sizes = a.shape();
+  slice_sizes[axis] = 1;
+
+  auto out = gather(a, indices, axis, slice_sizes, s);
+
+  // Transpose indices dimensions to axis dimension
+  if (axis != 0) {
+    std::vector<int> t_axes(out.ndim());
+    std::iota(t_axes.begin(), t_axes.begin() + axis, indices.ndim());
+    std::iota(t_axes.begin() + axis, t_axes.begin() + axis + indices.ndim(), 0);
+    std::iota(
+        t_axes.begin() + axis + indices.ndim(),
+        t_axes.end(),
+        indices.ndim() + axis);
+    out = transpose(out, t_axes, s);
+  }
+
+  // Squeeze the axis we take over
+  return squeeze(out, indices.ndim() + axis, s);
+}
+
+array take(const array& a, const array& indices, StreamOrDevice s /* = {} */) {
+  return take(flatten(a, s), indices, 0, s);
+}
+
+array take(const array& a, int index, int axis, StreamOrDevice s /* = {} */) {
+  // Check for valid axis
+  if (axis + static_cast<int>(a.ndim()) < 0 ||
+      axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[take] Received invalid axis " << axis << " for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Check for valid take
+  if (a.size() == 0) {
+    throw std::invalid_argument(
+        "[take] Cannot do a non-empty take from an array with zero elements.");
+  }
+
+  // Handle negative axis
+  axis = axis < 0 ? a.ndim() + axis : axis;
+
+  Shape starts(a.ndim(), 0);
+  Shape stops = a.shape();
+  starts[axis] = index;
+  stops[axis] = index + 1;
+  return squeeze(slice(a, std::move(starts), std::move(stops), s), axis, s);
+}
+
+array take(const array& a, int index, StreamOrDevice s /* = {} */) {
+  return take(flatten(a, s), index, 0, s);
+}
+
+array take_along_axis(
+    const array& a,
+    const array& indices,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  if (axis + a.ndim() < 0 || axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << "[take_along_axis] Received invalid axis for array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (indices.ndim() != a.ndim()) {
+    std::ostringstream msg;
+    msg << "[take_along_axis] Indices of dimension " << indices.ndim()
+        << " does not match array of dimension " << a.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Allow negative axis
+  axis = axis < 0 ? a.ndim() + axis : axis;
+
+  // Broadcast indices and input ignoring the take axis
+  auto inputs =
+      broadcast_arrays({a, indices}, std::vector<int>{axis - int(a.ndim())}, s);
+
+  auto out_shape = inputs[1].shape();
+  return array(
+      std::move(out_shape),
+      a.dtype(),
+      std::make_shared<GatherAxis>(to_stream(s), axis),
+      std::move(inputs));
+}
+
+array scatter_axis(
+    const array& a,
+    const array& indices,
+    const array& values,
+    int axis,
+    ScatterAxis::ReduceType mode,
+    StreamOrDevice s) {
+  std::string prefix =
+      (mode == ScatterAxis::None) ? "[put_along_axis]" : "[scatter_add_axis]";
+  if (axis + a.ndim() < 0 || axis >= static_cast<int>(a.ndim())) {
+    std::ostringstream msg;
+    msg << prefix << " Received invalid axis for array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (indices.ndim() != a.ndim()) {
+    std::ostringstream msg;
+    msg << prefix << " Indices of dimension " << indices.ndim()
+        << " does not match array of dimension " << a.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.size() == 0) {
+    return a;
+  }
+
+  auto upd = astype(values, a.dtype(), s);
+
+  // Squeeze leading singletons out of update
+  if (upd.ndim() > indices.ndim()) {
+    std::vector<int> sq_ax(upd.ndim() - indices.ndim());
+    std::iota(sq_ax.begin(), sq_ax.end(), 0);
+    upd = squeeze(upd, sq_ax, s);
+  }
+
+  auto inputs = broadcast_arrays({indices, upd}, s);
+  inputs.insert(inputs.begin(), a);
+
+  // Allow negative axis
+  axis = axis < 0 ? a.ndim() + axis : axis;
+
+  // Broadcast src, indices, values while ignoring the take axis
+  inputs = broadcast_arrays(inputs, {axis - int(a.ndim())}, s);
+
+  auto out_shape = inputs[0].shape();
+  return array(
+      std::move(out_shape),
+      a.dtype(),
+      std::make_shared<ScatterAxis>(to_stream(s), mode, axis),
+      std::move(inputs));
+}
+
+array put_along_axis(
+    const array& a,
+    const array& indices,
+    const array& values,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  return scatter_axis(a, indices, values, axis, ScatterAxis::None, s);
+}
+
+array scatter_add_axis(
+    const array& a,
+    const array& indices,
+    const array& values,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  return scatter_axis(a, indices, values, axis, ScatterAxis::Sum, s);
+}
+
+/** Scatter updates to given indices */
+array scatter(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    Scatter::ReduceType mode,
+    StreamOrDevice s) {
+  // Checks that indices, dimensions, and slice_sizes are all valid
+  if (indices.size() > a.ndim()) {
+    std::ostringstream msg;
+    msg << "[scatter] Too many index arrays. Got " << indices.size()
+        << " index arrays for input with " << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  for (auto& x : indices) {
+    if (x.dtype() == bool_) {
+      throw("[scatter] Boolean indices not supported.");
+    }
+  }
+
+  std::set dims(axes.begin(), axes.end());
+  if (dims.size() != axes.size()) {
+    throw std::invalid_argument(
+        "[scatter] Repeat axes not allowed in scatter.");
+  }
+  if (!dims.empty() && (*dims.begin() < 0 || *dims.rbegin() >= a.ndim())) {
+    throw std::invalid_argument("[scatter] Axes don't match array dimensions.");
+  }
+  if (indices.size() != axes.size()) {
+    throw std::invalid_argument(
+        "[scatter] Number of index arrays does not match number of axes.");
+  }
+
+  // Broadcast and cast indices if necessary
+  auto inputs = broadcast_arrays(indices);
+
+  Shape idx_shape;
+  if (!inputs.empty()) {
+    idx_shape = inputs[0].shape();
+  }
+
+  if (updates.ndim() != (a.ndim() + idx_shape.size())) {
+    std::ostringstream msg;
+    msg << "[scatter] Updates with " << updates.ndim()
+        << " dimensions does not match the sum of the array (" << a.ndim()
+        << ") and indices (" << idx_shape.size() << ") dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  for (int i = 0; i < idx_shape.size(); ++i) {
+    if (updates.shape(i) != idx_shape[i]) {
+      std::ostringstream msg;
+      msg << "[scatter] Update shape " << updates.shape()
+          << " is not valid for broadcasted index shape " << idx_shape << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  for (int i = 0; i < a.ndim(); ++i) {
+    auto up_shape = updates.shape(i + idx_shape.size());
+    if (up_shape > a.shape(i)) {
+      std::ostringstream msg;
+      msg << "[scatter] Updates with shape " << updates.shape()
+          << " are too large for array with shape " << a.shape() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  // Promote indices to the same type
+  auto dtype = result_type(indices);
+  if (issubdtype(dtype, inexact)) {
+    throw std::invalid_argument(
+        "[scatter] Got indices with invalid dtype. Indices must be integral.");
+  }
+  for (auto& idx : inputs) {
+    idx = astype(idx, dtype, s);
+  }
+
+  // TODO, remove when scatter supports 64-bit outputs
+  if (to_stream(s).device == Device::gpu && size_of(a.dtype()) == 8) {
+    std::ostringstream msg;
+    msg << "[scatter] GPU scatter does not yet support " << a.dtype()
+        << " for the input or updates.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  inputs.insert(inputs.begin(), a);
+  inputs.push_back(astype(updates, a.dtype(), s));
+
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Scatter>(to_stream(s), mode, axes),
+      std::move(inputs));
+}
+
+array scatter(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s /*= {}*/) {
+  return scatter(a, indices, updates, axes, Scatter::None, s);
+}
+
+array scatter_add(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s /*= {}*/) {
+  return scatter(a, indices, updates, axes, Scatter::Sum, s);
+}
+
+array scatter_prod(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s /*= {}*/) {
+  return scatter(a, indices, updates, axes, Scatter::Prod, s);
+}
+
+array scatter_max(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s /*= {}*/) {
+  return scatter(a, indices, updates, axes, Scatter::Max, s);
+}
+
+array scatter_min(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s /*= {}*/) {
+  return scatter(a, indices, updates, axes, Scatter::Min, s);
+}
+
+array masked_scatter(
+    const array& a,
+    const array& mask,
+    const array& value,
+    StreamOrDevice s /* =  {} */) {
+  if (mask.dtype() != bool_) {
+    throw std::invalid_argument("[masked_scatter] The mask has to be boolean.");
+  }
+
+  if (mask.ndim() > a.ndim()) {
+    throw std::invalid_argument(
+        "[masked_scatter] The mask cannot have more dimensions than the target.");
+  }
+
+  int unmasked_dims = a.ndim() - mask.ndim();
+
+  if (value.ndim() > unmasked_dims + 1) {
+    std::ostringstream msg;
+    msg << "[masked_scatter] Value array shape must be broadcastable with the last "
+        << unmasked_dims << " dimensions of the input.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Check if the start of the mask is compatible
+  if (!std::equal(
+          mask.shape().begin(), mask.shape().end(), a.shape().begin())) {
+    std::ostringstream msg;
+    msg << "[masked_scatter] The boolean mask should have the same shape as the "
+        << "beginning of the indexed array but the mask has shape "
+        << mask.shape() << " and the array has shape " << a.shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  array expanded_mask = mask;
+  array expanded_value = astype(value, a.dtype(), s);
+
+  // Broadcast both the mask with the last unmasked_dims of a
+  if (unmasked_dims > 0) {
+    auto mask_shape = mask.shape();
+    while (mask_shape.size() < a.ndim()) {
+      mask_shape.push_back(1);
+    }
+    expanded_mask = broadcast_to(reshape(mask, mask_shape, s), a.shape(), s);
+  }
+
+  // Broadcast the value with the unmasked dims plus one extra dimension of
+  // size mask.size(). If that dim is already provided leave it as is.
+  if (value.ndim() < unmasked_dims + 1) {
+    Shape value_shape(unmasked_dims + 1 - value.ndim(), 1);
+    value_shape.insert(
+        value_shape.end(), value.shape().begin(), value.shape().end());
+    expanded_value = reshape(expanded_value, value_shape, s);
+
+    value_shape[0] = mask.size();
+    for (int i = 1; i < unmasked_dims + 1; i++) {
+      value_shape[i] = a.shape(i - unmasked_dims - 1);
+    }
+    expanded_value = broadcast_to(expanded_value, value_shape, s);
+  } else if (!std::equal(
+                 value.shape().begin() + 1,
+                 value.shape().end(),
+                 a.shape().end() - unmasked_dims)) {
+    auto value_shape = value.shape();
+    for (int i = 1; i < unmasked_dims + 1; i++) {
+      value_shape[i] = a.shape(i - unmasked_dims - 1);
+    }
+    expanded_value = broadcast_to(expanded_value, value_shape, s);
+  }
+
+  array expanded_a = expand_dims(a, 0, s);
+  expanded_mask = expand_dims(expanded_mask, 0, s);
+  expanded_value = expand_dims(expanded_value, 0, s);
+
+  return squeeze(
+      array(
+          expanded_a.shape(),
+          expanded_a.dtype(),
+          std::make_shared<MaskedScatter>(to_stream(s)),
+          {expanded_a, expanded_mask, expanded_value}),
+      0,
+      s);
+}
+
+array sqrt(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<Sqrt>(to_stream(s)),
+      {astype(a, dtype, s)});
+}
+
+array rsqrt(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = at_least_float(a.dtype());
+  return array(
+      a.shape(),
+      dtype,
+      std::make_shared<Sqrt>(to_stream(s), true),
+      {astype(a, dtype, s)});
+}
+
+array softmax(
+    const array& a,
+    const std::vector<int>& axes,
+    bool precise /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  if (a.size() == 0) {
+    return a;
+  }
+  if (a.ndim() == 0 && !axes.empty()) {
+    throw std::invalid_argument(
+        "[softmax] Received non-empty axes for array with 0 dimensions.");
+  }
+  bool reduce_last_dim =
+      !axes.empty() && (axes.back() == a.ndim() - 1 || axes.back() == -1);
+  if (reduce_last_dim) {
+    // For more than 2 axes check if axes is [0, 1, ..., NDIM - 1] and shape
+    // is [1, 1, ..., N].
+    for (int i = axes.size() - 2; i >= 0; --i) {
+      if ((axes[i] + 1 != axes[i + 1]) || (a.shape(axes[i]) != 1)) {
+        reduce_last_dim = false;
+        break;
+      }
+    }
+  }
+  bool is_complex = issubdtype(a.dtype(), complexfloating);
+  if (!is_complex && reduce_last_dim) {
+    auto dtype = at_least_float(a.dtype());
+    return array(
+        a.shape(),
+        dtype,
+        std::make_shared<Softmax>(to_stream(s), precise),
+        {astype(a, dtype, s)});
+  } else {
+    auto in = a;
+    if (precise && !is_complex) {
+      in = astype(a, float32, s);
+    }
+    auto a_max = stop_gradient(max(in, axes, /*keepdims = */ true, s), s);
+    auto ex = exp(subtract(in, a_max, s), s);
+    return astype(
+        divide(ex, sum(ex, axes, /*keepdims = */ true, s), s), a.dtype(), s);
+  }
+}
+
+array softmax(
+    const array& a,
+    bool precise /* = false */,
+    StreamOrDevice s /* = {}*/) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return softmax(a, axes, precise, s);
+}
+
+array power(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto dtype = promote_types(a.dtype(), b.dtype());
+  std::vector<array> inputs = {astype(a, dtype, s), astype(b, dtype, s)};
+  if (a.shape() != b.shape()) {
+    inputs = broadcast_arrays(inputs, s);
+  }
+  return array(
+      inputs[0].shape(), dtype, std::make_shared<Power>(to_stream(s)), inputs);
+}
+
+array cumsum(
+    const array& a,
+    int axis,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  int ndim = a.ndim();
+  if (axis >= ndim || axis < -ndim) {
+    std::ostringstream msg;
+    msg << "[cumsum] Axis " << axis << " is out of bounds for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  axis = (axis + a.ndim()) % a.ndim();
+  auto out_type = a.dtype() == bool_ ? int32 : a.dtype();
+  return array(
+      a.shape(),
+      out_type,
+      std::make_shared<Scan>(
+          to_stream(s), Scan::ReduceType::Sum, axis, reverse, inclusive),
+      {a});
+}
+
+array cumsum(
+    const array& a,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  return cumsum(flatten(a, to_stream(s)), 0, reverse, inclusive, to_stream(s));
+}
+
+array cumprod(
+    const array& a,
+    int axis,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  int ndim = a.ndim();
+  if (axis >= ndim || axis < -ndim) {
+    std::ostringstream msg;
+    msg << "[cumprod] Axis " << axis << " is out of bounds for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  axis = (axis + a.ndim()) % a.ndim();
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Scan>(
+          to_stream(s), Scan::ReduceType::Prod, axis, reverse, inclusive),
+      {a});
+}
+
+array cumprod(
+    const array& a,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  return cumprod(flatten(a, s), 0, reverse, inclusive, s);
+}
+
+array cummax(
+    const array& a,
+    int axis,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  int ndim = a.ndim();
+  if (axis >= ndim || axis < -ndim) {
+    std::ostringstream msg;
+    msg << "[cummax] Axis " << axis << " is out of bounds for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  axis = (axis + a.ndim()) % a.ndim();
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Scan>(
+          to_stream(s), Scan::ReduceType::Max, axis, reverse, inclusive),
+      {a});
+}
+
+array cummax(
+    const array& a,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  return cummax(flatten(a, s), 0, reverse, inclusive, s);
+}
+
+array cummin(
+    const array& a,
+    int axis,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  int ndim = a.ndim();
+  if (axis >= ndim || axis < -ndim) {
+    std::ostringstream msg;
+    msg << "[cummin] Axis " << axis << " is out of bounds for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  axis = (axis + a.ndim()) % a.ndim();
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Scan>(
+          to_stream(s), Scan::ReduceType::Min, axis, reverse, inclusive),
+      {a});
+}
+
+array cummin(
+    const array& a,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  return cummin(flatten(a, s), 0, reverse, inclusive, s);
+}
+
+array logcumsumexp(
+    const array& a,
+    int axis,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  int ndim = a.ndim();
+  if (axis >= ndim || axis < -ndim) {
+    std::ostringstream msg;
+    msg << "[logcumsumexp] Axis " << axis << " is out of bounds for array with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  axis = (axis + a.ndim()) % a.ndim();
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Scan>(
+          to_stream(s), Scan::ReduceType::LogAddExp, axis, reverse, inclusive),
+      {a});
+}
+
+array logcumsumexp(
+    const array& a,
+    bool reverse /* = false*/,
+    bool inclusive /* = true*/,
+    StreamOrDevice s /* = {}*/) {
+  return logcumsumexp(
+      flatten(a, to_stream(s)), 0, reverse, inclusive, to_stream(s));
+}
+
+/** Convolution operations */
+
+namespace {
+
+inline void
+run_conv_checks(const array& in, const array& wt, int n_dim, int groups) {
+  if (!issubdtype(in.dtype(), floating)) {
+    std::ostringstream msg;
+    msg << "[conv] Invalid input array with type " << in.dtype() << "."
+        << " Convolution currently only supports floating point types";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (in.ndim() != n_dim + 2) {
+    std::ostringstream msg;
+    msg << "[conv] Invalid input array with " << in.ndim() << " dimensions for "
+        << n_dim << "D convolution. Expected an array with " << n_dim + 2
+        << " dimensions following the format [N, ..., C_in].";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (wt.ndim() != n_dim + 2) {
+    std::ostringstream msg;
+    msg << "[conv] Invalid weight array with " << wt.ndim()
+        << " dimensions for " << n_dim << "D convolution."
+        << " Expected an array with " << n_dim + 2
+        << " dimensions following the format [C_out, ..., C_in].";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (in.shape(n_dim + 1) % groups != 0) {
+    std::ostringstream msg;
+    msg << "[conv] The input channels must be divisible by the number"
+        << " of groups. Got input with shape " << in.shape() << " and "
+        << groups << " groups.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (groups > 1 && wt.shape(0) % groups != 0) {
+    std::ostringstream msg;
+    msg << "[conv] If groups > 1, the output channels must be divisible by the number"
+        << " of groups. Got " << wt.shape(0) << " output channels and "
+        << groups << " groups.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (in.shape(n_dim + 1) != (groups * wt.shape(n_dim + 1))) {
+    std::ostringstream msg;
+    if (groups == 1) {
+      msg << "[conv] Expect the input channels in the input"
+          << " and weight array to match but got shapes -"
+          << " input: " << in.shape() << " and weight: " << wt.shape();
+
+    } else {
+      msg << "Given groups=" << groups << " and weights of shape " << wt.shape()
+          << ", expected to have " << (groups * wt.shape(n_dim + 1))
+          << " input channels but got " << in.shape(n_dim + 1)
+          << " input channels instead.";
+    }
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+} // namespace
+
+/** 1D convolution with a filter */
+array conv1d(
+    const array& in_,
+    const array& wt_,
+    int stride /* = 1 */,
+    int padding /* = 0 */,
+    int dilation /* = 1 */,
+    int groups /* = 1 */,
+    StreamOrDevice s /* = {} */) {
+  return conv_general(
+      /* const array& input = */ in_,
+      /* const array& weight = */ wt_,
+      /* std::vector<int> stride = */ {stride},
+      /* std::vector<int> padding = */ {padding},
+      /* std::vector<int> kernel_dilation = */ {dilation},
+      /* std::vector<int> input_dilation = */ {1},
+      /* int groups = */ groups,
+      /* bool flip = */ false,
+      s);
+}
+
+/** 2D convolution with a filter */
+array conv2d(
+    const array& in_,
+    const array& wt_,
+    const std::pair<int, int>& stride /* = {1, 1} */,
+    const std::pair<int, int>& padding /* = {0, 0} */,
+    const std::pair<int, int>& dilation /* = {1, 1} */,
+    int groups /* = 1 */,
+    StreamOrDevice s /* = {} */) {
+  return conv_general(
+      /* const array& input = */ in_,
+      /* const array& weight = */ wt_,
+      /* std::vector<int> stride = */ {stride.first, stride.second},
+      /* std::vector<int> padding = */ {padding.first, padding.second},
+      /* std::vector<int> kernel_dilation = */
+      {dilation.first, dilation.second},
+      /* std::vector<int> input_dilation = */ {1, 1},
+      /* int groups = */ groups,
+      /* bool flip = */ false,
+      s);
+}
+
+/** 3D convolution with a filter */
+array conv3d(
+    const array& in_,
+    const array& wt_,
+    const std::tuple<int, int, int>& stride /* = {1, 1, 1} */,
+    const std::tuple<int, int, int>& padding /* = {0, 0, 0} */,
+    const std::tuple<int, int, int>& dilation /* = {1, 1, 1} */,
+    int groups /* = 1 */,
+    StreamOrDevice s /* = {} */) {
+  return conv_general(
+      /* const array& input = */ in_,
+      /* const array& weight = */ wt_,
+      /* std::vector<int> stride = */
+      {std::get<0>(stride), std::get<1>(stride), std::get<2>(stride)},
+      /* std::vector<int> padding = */
+      {std::get<0>(padding), std::get<1>(padding), std::get<2>(padding)},
+      /* std::vector<int> kernel_dilation = */
+      {std::get<0>(dilation), std::get<1>(dilation), std::get<2>(dilation)},
+      /* std::vector<int> input_dilation = */ {1, 1, 1},
+      /* int groups = */ groups,
+      /* bool flip = */ false,
+      s);
+}
+
+// Helper function for transposed convolutions
+array conv_transpose_general(
+    const array& input,
+    const array& weight,
+    std::vector<int> stride,
+    std::vector<int> padding,
+    std::vector<int> dilation,
+    std::vector<int> output_padding,
+    int groups,
+    StreamOrDevice s) {
+  std::vector<int> padding_lo(padding.size());
+  std::vector<int> padding_hi(padding.size());
+  for (int i = 0; i < padding.size(); ++i) {
+    int wt_size = 1 + dilation[i] * (weight.shape(1 + i) - 1);
+    padding_lo[i] = wt_size - padding[i] - 1;
+
+    int conv_output_shape = (input.shape(i + 1) - 1) * stride[i] -
+        2 * padding[i] + dilation[i] * (weight.shape(i + 1) - 1) + 1;
+
+    int in_size = 1 + (conv_output_shape - 1);
+    int out_size = 1 + stride[i] * (input.shape(1 + i) - 1);
+    padding_hi[i] = in_size - out_size + padding[i] +
+        output_padding[i]; // Adjust with output_padding
+  }
+
+  return conv_general(
+      /* const array& input = */ input,
+      /* const array& weight = */ weight,
+      /* std::vector<int> stride = */ std::vector(stride.size(), 1),
+      /* std::vector<int> padding_lo = */ std::move(padding_lo),
+      /* std::vector<int> padding_hi = */ std::move(padding_hi),
+      /* std::vector<int> kernel_dilation = */ std::move(dilation),
+      /* std::vector<int> input_dilation = */ std::move(stride),
+      /* int groups = */ groups,
+      /* bool flip = */ true,
+      s);
+}
+
+/** 1D transposed convolution with a filter */
+array conv_transpose1d(
+    const array& in_,
+    const array& wt_,
+    int stride /* = 1 */,
+    int padding /* = 0 */,
+    int dilation /* = 1 */,
+    int output_padding /* = 0 */,
+    int groups /* = 1 */,
+    StreamOrDevice s /* = {} */) {
+  return conv_transpose_general(
+      in_, wt_, {stride}, {padding}, {dilation}, {output_padding}, groups, s);
+}
+
+/** 2D transposed convolution with a filter */
+array conv_transpose2d(
+    const array& in_,
+    const array& wt_,
+    const std::pair<int, int>& stride /* = {1, 1} */,
+    const std::pair<int, int>& padding /* = {0, 0} */,
+    const std::pair<int, int>& dilation /* = {1, 1} */,
+    const std::pair<int, int>& output_padding /* = {0, 0} */,
+    int groups /* = 1 */,
+    StreamOrDevice s /* = {} */) {
+  return conv_transpose_general(
+      in_,
+      wt_,
+      {stride.first, stride.second},
+      {padding.first, padding.second},
+      {dilation.first, dilation.second},
+      {output_padding.first, output_padding.second},
+      groups,
+      s);
+}
+
+/** 3D transposed convolution with a filter */
+array conv_transpose3d(
+    const array& in_,
+    const array& wt_,
+    const std::tuple<int, int, int>& stride /* = {1, 1, 1} */,
+    const std::tuple<int, int, int>& padding /* = {0, 0, 0} */,
+    const std::tuple<int, int, int>& dilation /* = {1, 1, 1} */,
+    const std::tuple<int, int, int>& output_padding /* = {0, 0, 0} */,
+    int groups /* = 1 */,
+    StreamOrDevice s /* = {} */) {
+  return conv_transpose_general(
+      in_,
+      wt_,
+      {std::get<0>(stride), std::get<1>(stride), std::get<2>(stride)},
+      {std::get<0>(padding), std::get<1>(padding), std::get<2>(padding)},
+      {std::get<0>(dilation), std::get<1>(dilation), std::get<2>(dilation)},
+      {std::get<0>(output_padding),
+       std::get<1>(output_padding),
+       std::get<2>(output_padding)},
+      groups,
+      s);
+}
+
+/** General convolution with a filter */
+array conv_general(
+    array in,
+    array wt,
+    std::vector<int> stride /* = {} */,
+    std::vector<int> padding_lo /* = {} */,
+    std::vector<int> padding_hi /* = {} */,
+    std::vector<int> kernel_dilation /* = {} */,
+    std::vector<int> input_dilation /* = {} */,
+    int groups /* = 1 */,
+    bool flip /* = false */,
+    StreamOrDevice s /* = {} */) {
+  // Run checks
+  if (groups != 1 && in.ndim() != 3 && in.ndim() != 4) {
+    throw std::invalid_argument(
+        "[conv] Can only handle groups != 1 in 1D or 2D convolutions.");
+  }
+
+  int spatial_dims = in.ndim() - 2;
+
+  if (spatial_dims < 1 || spatial_dims > 3) {
+    throw std::invalid_argument(
+        "[conv] Only works for inputs with 1-3 spatial dimensions."
+        " The inputs must be in the format [N, ..., C_in]");
+  }
+
+  // Run checks
+  run_conv_checks(in, wt, spatial_dims, groups);
+
+  // Type promotion
+  auto out_type = promote_types(in.dtype(), wt.dtype());
+  in = astype(in, out_type, s);
+  wt = astype(wt, out_type, s);
+
+  if (stride.size() <= 1) {
+    int stride_int = stride.size() ? stride[0] : 1;
+    stride = std::vector<int>(spatial_dims, stride_int);
+  }
+
+  if (padding_lo.size() <= 1) {
+    int padding_int = padding_lo.size() ? padding_lo[0] : 0;
+    padding_lo = std::vector<int>(spatial_dims, padding_int);
+  }
+
+  if (padding_hi.size() <= 1) {
+    int padding_int = padding_hi.size() ? padding_hi[0] : 0;
+    padding_hi = std::vector<int>(spatial_dims, padding_int);
+  }
+
+  if (kernel_dilation.size() <= 1) {
+    int kernel_dilation_int = kernel_dilation.size() ? kernel_dilation[0] : 1;
+    kernel_dilation = std::vector<int>(spatial_dims, kernel_dilation_int);
+  }
+
+  if (input_dilation.size() <= 1) {
+    int input_dilation_int = input_dilation.size() ? input_dilation[0] : 1;
+    input_dilation = std::vector<int>(spatial_dims, input_dilation_int);
+  }
+
+  // Check for negative padding
+  bool has_neg_padding = false;
+  for (auto& pd : padding_lo) {
+    has_neg_padding |= (pd < 0);
+  }
+  for (auto& pd : padding_hi) {
+    has_neg_padding |= (pd < 0);
+  }
+
+  // Handle negative padding
+  if (has_neg_padding) {
+    Shape starts(in.ndim(), 0);
+    auto stops = in.shape();
+
+    for (int i = 0; i < spatial_dims; i++) {
+      if (padding_lo[i] < 0) {
+        starts[i + 1] -= padding_lo[i];
+        padding_lo[i] = 0;
+      }
+
+      if (padding_hi[i] < 0) {
+        stops[i + 1] += padding_hi[i];
+        padding_hi[i] = 0;
+      }
+    }
+
+    in = slice(in, std::move(starts), std::move(stops), s);
+  }
+
+  // Get output shapes
+  auto out_shape = Convolution::conv_out_shape(
+      in.shape(),
+      wt.shape(),
+      stride,
+      padding_lo,
+      padding_hi,
+      kernel_dilation,
+      input_dilation);
+
+  return array(
+      std::move(out_shape),
+      in.dtype(),
+      std::make_shared<Convolution>(
+          to_stream(s),
+          stride,
+          padding_lo,
+          padding_hi,
+          kernel_dilation,
+          input_dilation,
+          groups,
+          flip),
+      {in, wt});
+}
+
+std::pair<int, int> quantization_params_from_mode(
+    QuantizationMode mode,
+    std::optional<int> group_size_,
+    std::optional<int> bits_) {
+  int default_group_size;
+  int default_bits;
+  switch (mode) {
+    case QuantizationMode::Affine:
+      default_group_size = 64;
+      default_bits = 4;
+      break;
+    case QuantizationMode::Nvfp4:
+      default_group_size = 16;
+      default_bits = 4;
+      break;
+    case QuantizationMode::Mxfp4:
+      default_group_size = 32;
+      default_bits = 4;
+      break;
+    case QuantizationMode::Mxfp8:
+      default_group_size = 32;
+      default_bits = 8;
+      break;
+  }
+  return {
+      group_size_.has_value() ? *group_size_ : default_group_size,
+      bits_.has_value() ? *bits_ : default_bits};
+}
+
+std::pair<Dtype, QuantizationMode> validate_mode_with_type(
+    std::string_view tag,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<Dtype> out_type,
+    const std::string& mode) {
+  auto qmode = string_to_quantization_mode(mode, tag);
+  if (out_type.has_value() && !issubdtype(*out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] Only real floating types are supported but "
+        << "output dtype == " << *out_type << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (qmode == QuantizationMode::Affine) {
+    if (!biases) {
+      std::ostringstream msg;
+      msg << "[" << tag << "] Biases must be provided for affine quantization.";
+      throw std::invalid_argument(msg.str());
+    }
+    auto dtype = result_type(scales, *biases);
+    if (!issubdtype(dtype, floating)) {
+      std::ostringstream msg;
+      msg << "[" << tag << "] Only real floating types are supported but "
+          << "scales.dtype() == " << scales.dtype()
+          << " and biases.dtype() == " << biases->dtype() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+    if (out_type.has_value()) {
+      return {*out_type, qmode};
+    } else {
+      return {dtype, qmode};
+    }
+  } else if (scales.dtype() != uint8) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] Scale type must be uint8 but received type "
+        << scales.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (biases) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] Biases must be null for quantization mode '" << mode
+        << "'.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (out_type.has_value()) {
+    return {*out_type, qmode};
+  } else {
+    return {bfloat16, qmode};
+  }
+}
+
+array quantized_matmul(
+    array x,
+    array w,
+    array scales,
+    std::optional<array> biases /* = std::nullopt */,
+    bool transpose /* = true */,
+    std::optional<int> group_size_ /* = std::nullopt */,
+    std::optional<int> bits_ /* = std::nullopt */,
+    const std::string& mode /* = "affine" */,
+    StreamOrDevice s /* = {} */) {
+  auto [dtype, qmode] = validate_mode_with_type(
+      "quantized_matmul", scales, biases, std::nullopt, mode);
+
+  auto [group_size, bits] =
+      quantization_params_from_mode(qmode, group_size_, bits_);
+  // Check and extract the quantized matrix shape against x
+  auto [w_inner_dims, w_outer_dims] = extract_quantized_matmul_dims(
+      "quantized_matmul", x, w, scales, biases, transpose, group_size, bits);
+
+  if (qmode == QuantizationMode::Affine) {
+    dtype = promote_types(x.dtype(), dtype);
+  } else {
+    dtype = x.dtype();
+  }
+
+  if (!issubdtype(dtype, floating)) {
+    std::ostringstream msg;
+    msg << "[quantized_matmul] Only real floating types are supported but "
+        << "x.dtype() == " << x.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  std::vector<array> inputs;
+  if (qmode == QuantizationMode::Affine) {
+    inputs = {
+        astype(x, dtype), w, astype(scales, dtype), astype(*biases, dtype)};
+  } else {
+    inputs = {x, w, scales};
+  }
+
+  if (x.ndim() > 2 && w.ndim() > 2) {
+    inputs = broadcast_arrays(inputs, {-2, -1}, s);
+  }
+
+  auto out_shape = inputs[0].shape();
+  out_shape.back() = w_outer_dims;
+  return array(
+      std::move(out_shape),
+      dtype,
+      std::make_shared<QuantizedMatmul>(
+          to_stream(s), group_size, bits, qmode, transpose),
+      std::move(inputs));
+}
+
+void validate_qqmm_inputs(
+    array x,
+    array w,
+    std::optional<array> scales_w,
+    int group_size,
+    int bits) {
+  // check 2D (for now)
+  if (x.ndim() > 2 || w.ndim() > 2) {
+    std::ostringstream msg;
+    msg << "[qqmm] Only 2D inputs are supported but "
+        << "x.ndim() == " << x.ndim() << " and "
+        << "w.ndim() == " << w.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (w.dtype() == uint32) {
+    // if w is quantized, scales are provided
+    if (!scales_w.has_value()) {
+      std::ostringstream msg;
+      throw std::invalid_argument(
+          "[qqmm] Scales must be provided if second argument is quantized.");
+    }
+    // if scales are provided, check compatibility with quantized w
+    else {
+      validate_quantized_input("qqmm", w, *scales_w, group_size, bits);
+    }
+  }
+  // if w is not quantized, dtype must be in {f16, bf16, fp32}
+  else {
+    if (!issubdtype(w.dtype(), floating) || w.dtype() == float64) {
+      std::ostringstream msg;
+      msg << "[qqmm] Only real floating types except float64 are supported but "
+          << "second argument dtype == " << w.dtype() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  // x dtype must be in {f16, bf16, fp32}
+  if (!issubdtype(x.dtype(), floating) || x.dtype() == float64) {
+    std::ostringstream msg;
+    msg << "[qqmm] Only real floating types except float64 are supported but "
+        << "first argument dtype == " << x.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+std::pair<int, int> extract_qqmm_dims(
+    array x,
+    array w,
+    std::optional<array> scales_w,
+    int group_size,
+    int bits) {
+  if (w.dtype() != uint32) {
+    // if w is not quantized, check that last dims match
+    if (x.shape(-1) != w.shape(-1)) {
+      std::ostringstream msg;
+      msg << "[qqmm] Last dimension of first input with shape " << x.shape()
+          << " must match last dimension of"
+          << " second input with shape " << w.shape() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+    return std::make_pair(w.shape(-1), w.shape(-2));
+  } else {
+    // if w is quantized, extract dims from quantized w
+    return extract_quantized_matmul_dims(
+        "qqmm",
+        x,
+        w,
+        *scales_w,
+        std::nullopt,
+        /* transpose = */ true,
+        group_size,
+        bits);
+  }
+}
+
+array qqmm(
+    array in_x,
+    array w,
+    std::optional<array> scales_w,
+    std::optional<int> group_size_ /* = std::nullopt */,
+    std::optional<int> bits_ /* = std::nullopt */,
+    const std::string& mode /* = "nvfp4" */,
+    StreamOrDevice s /* = {} */) {
+  auto stream = to_stream(s);
+  auto qmode = string_to_quantization_mode(mode, "qqmm");
+  // cuBLAS block scaled matmul only supports nvfp4 and mxfp8
+  if (qmode != QuantizationMode::Nvfp4 && qmode != QuantizationMode::Mxfp8) {
+    std::ostringstream msg;
+    msg << "[qqmm] Only 'nvfp4' and 'mxfp8' quantization modes are supported but '"
+        << mode << "' was provided.";
+    throw std::invalid_argument(msg.str());
+  }
+  // we need to check 2 cases:
+  // 1. w is quantized, scales is provided
+  // 2. w is not quantized, scales is not provided
+  auto [group_size, bits] =
+      quantization_params_from_mode(qmode, group_size_, bits_);
+
+  // Allow gemv
+  auto x = in_x;
+  if (x.ndim() == 1) {
+    // Insert a singleton dim in the beginning
+    x = expand_dims(x, 0, s);
+  } else if (w.ndim() == 2 && x.ndim() > 2) {
+    x = flatten(x, 0, -2, s);
+  }
+
+  // validate inputs
+  validate_qqmm_inputs(x, w, scales_w, group_size, bits);
+  // validate and extract shapes
+  auto [w_inner_dims, w_outer_dims] =
+      extract_qqmm_dims(x, w, scales_w, group_size, bits);
+  std::vector<array> inputs = {
+      x,
+      w,
+  };
+  if (scales_w.has_value()) {
+    inputs.push_back(*scales_w);
+  }
+  auto out_shape = inputs[0].shape();
+  out_shape.back() = w_outer_dims;
+  auto out = array(
+      std::move(out_shape),
+      x.dtype(), // output dtype is the same as x dtype
+      std::make_shared<QQMatmul>(stream, group_size, bits, qmode),
+      std::move(inputs));
+  if (in_x.ndim() > 2) {
+    auto orig_shape = in_x.shape();
+    orig_shape.pop_back();
+    out = unflatten(out, 0, std::move(orig_shape), s);
+  } else if (in_x.ndim() == 1) {
+    out = squeeze(out, 0, s);
+  }
+  return out;
+}
+
+array pack_and_quantize(
+    array& packed_w,
+    const array& scales,
+    const array& biases,
+    int bits,
+    const Stream& s) {
+  int el_per_int = 32 / bits;
+  array zero(0, packed_w.dtype());
+  array n_bins((1 << bits) - 1, packed_w.dtype()); // 2**bits - 1
+  packed_w = astype(
+      clip(
+          round(divide(subtract(packed_w, biases, s), scales, s), s),
+          zero,
+          n_bins,
+          s),
+      uint32,
+      s);
+  if (is_power_of_2(bits)) {
+    array shifts = power(array(2, uint32), arange(0, 32, bits, uint32, s), s);
+    packed_w = reshape(packed_w, {packed_w.shape(0), -1, el_per_int}, s);
+    packed_w =
+        sum(multiply(packed_w, shifts, s),
+            /* axis= */ 2,
+            /* keepdims= */ false,
+            s);
+  } else {
+    // This is slow but we have fast GPU/CPU versions of this function so we
+    // shouldn't be here often.
+    packed_w = expand_dims(packed_w, /* axis= */ -1, s);
+    packed_w = bitwise_and(
+        right_shift(packed_w, arange(bits, uint32, s), s),
+        array({1}, uint32),
+        s);
+    auto new_shape = packed_w.shape();
+    new_shape[new_shape.size() - 2] = -1;
+    new_shape.back() = 32;
+    packed_w = reshape(packed_w, new_shape, s);
+    array shifts = arange(32, uint32, s);
+    packed_w =
+        sum(left_shift(packed_w, shifts, s),
+            /* axis= */ -1,
+            /* keepdims= */ false,
+            s);
+  }
+  return packed_w;
+}
+
+std::vector<array>
+affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) {
+  auto s = to_stream(s_);
+  if (group_size != 32 && group_size != 64 && group_size != 128) {
+    std::ostringstream msg;
+    msg << "[quantize] The requested group size " << group_size
+        << " is not supported. The supported group sizes are 32, 64, and 128.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (bits < 2 || bits > 8 || bits == 7) {
+    std::ostringstream msg;
+    msg << "[quantize] The requested number of bits " << bits
+        << " is not supported. The supported bits are 2, 3, 4, 5, 6 and 8.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto fallback = [group_size, bits, s](
+                      const std::vector<array>& inputs) -> std::vector<array> {
+    auto& w = inputs[0];
+    auto wshape = w.shape();
+    wshape.back() = -1;
+
+    array zero(0, float32);
+    array n_bins((1 << bits) - 1, float32); // 2**bits - 1
+    array eps(1e-7, float32);
+
+    array packed_w = reshape(w, {-1, w.shape(-1) / group_size, group_size}, s);
+
+    array w_max = max(packed_w, /* axis= */ -1, /* keepdims= */ true, s);
+    array w_min = min(packed_w, /* axis= */ -1, /* keepdims= */ true, s);
+    w_max = astype(w_max, float32, s);
+    w_min = astype(w_min, float32, s);
+
+    array mask = greater(abs(w_min, s), abs(w_max, s), s);
+    array scales =
+        maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s);
+    scales = where(mask, scales, negative(scales, s), s);
+    array edge = where(mask, w_min, w_max, s);
+    array q0 = round(divide(edge, scales, s), s);
+    scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales);
+    array biases = where(equal(q0, zero, s), zero, edge, s);
+
+    packed_w = pack_and_quantize(packed_w, scales, biases, bits, s);
+
+    scales = astype(scales, w.dtype(), s);
+    biases = astype(biases, w.dtype(), s);
+    return {
+        reshape(packed_w, wshape, s),
+        reshape(scales, wshape, s),
+        reshape(biases, wshape, s),
+    };
+  };
+
+  auto wq_shape = w.shape();
+  wq_shape.back() = w.shape(-1) * bits / 32;
+  auto sshape = w.shape();
+  sshape.back() = w.shape(-1) / group_size;
+  return array::make_arrays(
+      {std::move(wq_shape), sshape, sshape},
+      {uint32, w.dtype(), w.dtype()},
+      std::make_shared<fast::Quantize>(
+          s, fallback, group_size, bits, QuantizationMode::Affine, false),
+      {w});
+}
+
+std::vector<array> fp_quantize(
+    const array& w,
+    int group_size,
+    int bits,
+    QuantizationMode mode,
+    Stream s) {
+  int expected_gs = mode == QuantizationMode::Nvfp4 ? 16 : 32;
+  int expected_bits = mode == QuantizationMode::Mxfp8 ? 8 : 4;
+  if (group_size != expected_gs) {
+    std::ostringstream msg;
+    msg << "[quantize] " << quantization_mode_to_string(mode)
+        << " quantization requires group size " << expected_gs << " but got "
+        << group_size << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (bits != expected_bits) {
+    std::ostringstream msg;
+    msg << "[quantize] " << quantization_mode_to_string(mode)
+        << " quantization requires bits to be " << expected_bits << " but got "
+        << bits << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  auto fallback = [bits = bits, group_size = group_size, s](
+                      const std::vector<array>& inputs) -> std::vector<array> {
+    auto& w = inputs[0];
+    float maxval = (bits == 4) ? 6.0f : 448.0f;
+    auto new_shape = w.shape();
+    new_shape.back() = -1;
+    auto wq = reshape(w, {-1, group_size}, s);
+    auto scales =
+        divide(max(abs(wq, s), -1, true, s), array(maxval, w.dtype()), s);
+    if (group_size == 16) {
+      // convert to e4m3
+      scales = to_fp8(scales, s);
+      wq = divide(wq, from_fp8(scales, w.dtype(), s), s);
+    } else {
+      // convert to e8m0
+      auto z = array(0, scales.dtype());
+      scales = where(
+          equal(scales, z, s),
+          z,
+          astype(round(log2(scales, s), s), int32, s),
+          s);
+
+      wq = divide(wq, power(array(2.0f, w.dtype()), scales, s), s);
+      scales = astype(add(scales, array(127, int32), s), uint8, s);
+    }
+    if (bits == 4) {
+      auto lut = array({
+          +0.0f,
+          +0.5f,
+          +1.0f,
+          +1.5f,
+          +2.0f,
+          +3.0f,
+          +4.0f,
+          +6.0f,
+          -0.0f,
+          -0.5f,
+          -1.0f,
+          -1.5f,
+          -2.0f,
+          -3.0f,
+          -4.0f,
+          -6.0f,
+      });
+      lut = astype(lut, w.dtype(), s);
+      wq = argmin(
+          abs(subtract(expand_dims(wq, -1, s), lut, s), s), -1, false, s);
+      auto shifts = power(array(2, uint32), arange(0, 32, 4, uint32, s), s);
+      wq = reshape(wq, {-1, 4, 8}, s);
+      wq = sum(multiply(wq, shifts, s), -1, false, s);
+    } else {
+      wq = view(to_fp8(wq, s), uint32, s);
+    }
+    wq = reshape(wq, new_shape, s);
+    scales = reshape(scales, new_shape, s);
+    return {std::move(wq), std::move(scales)};
+  };
+
+  if (s.device == Device::gpu) {
+    auto wq_shape = w.shape();
+    wq_shape.back() = w.shape(-1) * bits / 32;
+    auto sshape = w.shape();
+    sshape.back() = w.shape(-1) / group_size;
+    return array::make_arrays(
+        {std::move(wq_shape), std::move(sshape)},
+        {uint32, uint8},
+        std::make_shared<fast::Quantize>(
+            s, fallback, group_size, bits, mode, false),
+        {w});
+  }
+  return fallback({w});
+}
+
+std::vector<array> quantize(
+    const array& w,
+    std::optional<int> group_size_ /* = std::nullopt */,
+    std::optional<int> bits_ /* = std::nullopt */,
+    const std::string& mode /* = "affine" */,
+    StreamOrDevice s /* = {} */) {
+  auto qmode = string_to_quantization_mode(mode, "quantize");
+  auto [group_size, bits] =
+      quantization_params_from_mode(qmode, group_size_, bits_);
+  if (!issubdtype(w.dtype(), floating)) {
+    std::ostringstream msg;
+    msg << "[quantize] Only real floating types can be quantized "
+        << "but w has type " << w.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (w.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[quantize] The matrix to be quantized must have at least 2 dimension "
+        << "but it has only " << w.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if ((w.shape(-1) % group_size) != 0) {
+    std::ostringstream msg;
+    msg << "[quantize] The last dimension of the matrix needs to be divisible by "
+        << "the quantization group size " << group_size
+        << ". However the provided "
+        << " matrix has shape " << w.shape();
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (qmode == QuantizationMode::Affine) {
+    return affine_quantize(w, group_size, bits, s);
+  } else {
+    return fp_quantize(w, group_size, bits, qmode, to_stream(s));
+  }
+}
+
+array affine_dequantize(
+    const array& w,
+    const array& scales,
+    const array& biases,
+    int group_size,
+    int bits,
+    StreamOrDevice s_) {
+  auto wshape = w.shape();
+  auto sshape = scales.shape();
+  auto bshape = biases.shape();
+  if (wshape.size() != sshape.size() || wshape.size() != bshape.size()) {
+    throw std::invalid_argument(
+        "[dequantize] Shape of scales and biases does not match the matrix");
+  }
+  wshape.back() = -1;
+  sshape.back() = -1;
+  bshape.back() = -1;
+
+  if (wshape != sshape || wshape != bshape) {
+    throw std::invalid_argument(
+        "[dequantize] Shape of scales and biases does not match the matrix");
+  }
+
+  // Packing into uint32
+  int out_size = w.shape(-1) * 32 / bits;
+
+  if (out_size != scales.shape(-1) * group_size) {
+    std::ostringstream msg;
+    msg << "[dequantize] Shape of scales and biases does not match the matrix "
+        << "given the quantization parameters. Provided matrix of shape "
+        << w.shape() << " and scales/biases of shape " << scales.shape()
+        << " with group_size=" << group_size << " and bits=" << bits << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto s = to_stream(s_);
+
+  auto fallback =
+      [wshape = std::move(wshape),
+       sshape = std::move(sshape),
+       group_size,
+       bits,
+       s](const std::vector<array>& inputs) mutable -> std::vector<array> {
+    auto w = inputs[0];
+    auto& scales = inputs[1];
+    auto& biases = inputs[2];
+    if (is_power_of_2(bits)) {
+      std::vector<array> parts;
+      for (int start = 0; start < 32; start += bits) {
+        parts.push_back(expand_dims(
+            right_shift(
+                left_shift(w, array(32 - (start + bits), uint32), s),
+                array(32 - bits, uint32),
+                s),
+            -1,
+            s));
+      }
+      w = concatenate(parts, -1, s);
+    } else {
+      w = expand_dims(w, /* axis= */ -1, s);
+      w = bitwise_and(
+          right_shift(w, arange(32, uint32, s), s), array({1}, uint32), s);
+      auto new_shape = w.shape();
+      new_shape[new_shape.size() - 2] = -1;
+      new_shape.back() = bits;
+      w = reshape(w, new_shape, s);
+      array shifts = arange(bits, uint32, s);
+      w = sum(
+          left_shift(w, shifts, s), /* axis= */ -1, /* keepdims= */ false, s);
+    }
+
+    // Dequantize
+    wshape.push_back(group_size);
+    w = reshape(w, wshape, s);
+    w = multiply(w, expand_dims(scales, -1, s), s);
+    w = add(w, expand_dims(biases, -1, s), s);
+    w = reshape(w, sshape, s);
+
+    return {w};
+  };
+
+  if (s.device == Device::gpu) {
+    auto out_shape = w.shape();
+    out_shape.back() = out_size;
+    return array(
+        std::move(out_shape),
+        scales.dtype(),
+        std::make_shared<fast::Quantize>(
+            s, fallback, group_size, bits, QuantizationMode::Affine, true),
+        {w, scales, biases});
+  }
+  return fallback({w, scales, biases})[0];
+}
+
+array fp_dequantize(
+    const array& w,
+    const array& scales,
+    int group_size,
+    int bits,
+    Dtype out_type,
+    QuantizationMode mode,
+    Stream s) {
+  int expected_gs = mode == QuantizationMode::Nvfp4 ? 16 : 32;
+  int expected_bits = mode == QuantizationMode::Mxfp8 ? 8 : 4;
+  if (group_size != expected_gs) {
+    std::ostringstream msg;
+    msg << "[dequantize] " << quantization_mode_to_string(mode)
+        << " quantization requires group size " << expected_gs << " but got "
+        << group_size << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (bits != expected_bits) {
+    std::ostringstream msg;
+    msg << "[dequantize] " << quantization_mode_to_string(mode)
+        << " quantization requires bits to be " << expected_bits << " but got "
+        << bits << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto wshape = w.shape();
+  auto sshape = scales.shape();
+  if (wshape.size() != sshape.size()) {
+    throw std::invalid_argument(
+        "[dequantize] Shape of scales does not match the matrix");
+  }
+
+  wshape.back() = -1;
+  sshape.back() = -1;
+
+  if (wshape != sshape) {
+    throw std::invalid_argument(
+        "[dequantize] Shape of scales does not match the matrix");
+  }
+
+  // Packing into uint32
+  int out_size = w.shape(-1) * 32 / bits;
+
+  if (out_size != scales.shape(-1) * group_size) {
+    std::ostringstream msg;
+    msg << "[dequantize] Shape of scales does not match the matrix "
+        << "given the quantization parameters. Provided matrix of shape "
+        << w.shape() << " and scales of shape " << scales.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto fallback =
+      [wshape = std::move(wshape),
+       sshape = std::move(sshape),
+       group_size,
+       bits,
+       out_type,
+       s](const std::vector<array>& inputs) mutable -> std::vector<array> {
+    auto out = inputs[0];
+    auto scales = inputs[1];
+    if (bits == 4) {
+      auto lut = array(
+          {
+              +0.0f,
+              +0.5f,
+              +1.0f,
+              +1.5f,
+              +2.0f,
+              +3.0f,
+              +4.0f,
+              +6.0f,
+              -0.0f,
+              -0.5f,
+              -1.0f,
+              -1.5f,
+              -2.0f,
+              -3.0f,
+              -4.0f,
+              -6.0f,
+          },
+          out_type);
+      out = view(reshape(out, {-1, 4}, s), int8, s);
+      auto idx_lo = bitwise_and(out, array(0x0F, int8), s);
+      auto idx_hi = right_shift(out, array(4, int8), s);
+      auto lo = gather(lut, idx_lo, 0, {1}, s);
+      auto hi = gather(lut, idx_hi, 0, {1}, s);
+      out = concatenate({lo, hi}, -1, s);
+    } else {
+      out = from_fp8(view(out, uint8, s), out_type, s);
+    }
+    out = reshape(out, {-1, group_size}, s);
+    scales = reshape(scales, {-1, 1}, s);
+    if (group_size == 16) {
+      scales = from_fp8(scales, out_type, s);
+    } else {
+      scales = subtract(astype(scales, out_type, s), array(127, out_type), s);
+      scales = power(array(2.0f, out_type), scales, s);
+    }
+    return {reshape(multiply(out, scales, s), wshape, s)};
+  };
+  if (s.device == Device::gpu) {
+    auto out_shape = w.shape();
+    out_shape.back() = out_size;
+    return array(
+        std::move(out_shape),
+        out_type,
+        std::make_shared<fast::Quantize>(
+            s, fallback, group_size, bits, mode, true),
+        {w, scales});
+  }
+  return fallback({w, scales})[0];
+}
+
+array dequantize(
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases /* = std::nullopt */,
+    std::optional<int> group_size_ /* = std::nullopt */,
+    std::optional<int> bits_ /* = std::nullopt */,
+    const std::string& mode /* = "affine" */,
+    std::optional<Dtype> dtype /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto [out_type, qmode] =
+      validate_mode_with_type("dequantize", scales, biases, dtype, mode);
+  auto [group_size, bits] =
+      quantization_params_from_mode(qmode, group_size_, bits_);
+  if (bits <= 0) {
+    std::ostringstream msg;
+    msg << "[dequantize] Invalid value for bits: " << bits;
+    throw std::invalid_argument(msg.str());
+  }
+  if (group_size <= 0) {
+    std::ostringstream msg;
+    msg << "[dequantize] Invalid value for group_size: " << group_size;
+    throw std::invalid_argument(msg.str());
+  }
+  if (w.dtype() != uint32) {
+    throw std::invalid_argument(
+        "[dequantize] The matrix should be given as a uint32");
+  }
+  if (w.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[dequantize] The matrix to be dequantized must have at least 2 dimension "
+        << "but it has only " << w.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (qmode == QuantizationMode::Affine) {
+    return astype(
+        affine_dequantize(w, scales, *biases, group_size, bits, s),
+        out_type,
+        s);
+  } else {
+    return fp_dequantize(
+        w, scales, group_size, bits, out_type, qmode, to_stream(s));
+  }
+}
+
+array from_fp8(array x, Dtype dtype, StreamOrDevice s) {
+  if (x.dtype() != uint8) {
+    std::ostringstream msg;
+    msg << "[from_fp8] Input must have type uint8 but "
+        << "x.dtype() == " << x.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (!issubdtype(dtype, floating)) {
+    std::ostringstream msg;
+    msg << "[from_fp8] Only real floating types are supported but "
+        << "dtype == " << dtype << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  return array(
+      x.shape(),
+      dtype,
+      std::make_shared<fast::ConvertFP8>(to_stream(s), false),
+      {x});
+}
+
+array to_fp8(array x, StreamOrDevice s) {
+  if (!issubdtype(x.dtype(), floating)) {
+    std::ostringstream msg;
+    msg << "[to_fp8] Only real floating types are supported but "
+        << "x.dtype() == " << x.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  return array(
+      x.shape(),
+      uint8,
+      std::make_shared<fast::ConvertFP8>(to_stream(s), true),
+      {x});
+}
+
+array gather_qmm(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases /* = std::nullopt */,
+    std::optional<array> lhs_indices_ /* = std::nullopt */,
+    std::optional<array> rhs_indices_ /* = std::nullopt */,
+    bool transpose /* = true */,
+    std::optional<int> group_size_ /* = std::nullopt */,
+    std::optional<int> bits_ /* = std::nullopt */,
+    const std::string& mode /* = "affine" */,
+    bool sorted_indices /* = false */,
+    StreamOrDevice s /* = {} */) {
+  if (!lhs_indices_ && !rhs_indices_) {
+    return quantized_matmul(
+        x, w, scales, biases, transpose, group_size_, bits_, mode, s);
+  }
+
+  auto [out_type, qmode] =
+      validate_mode_with_type("gather_qmm", scales, biases, std::nullopt, mode);
+  auto [group_size, bits] =
+      quantization_params_from_mode(qmode, group_size_, bits_);
+  auto [w_inner_dims, w_outer_dims] = extract_quantized_matmul_dims(
+      "gather_qmm", x, w, scales, biases, transpose, group_size, bits);
+  if (qmode == QuantizationMode::Affine) {
+    out_type = promote_types(x.dtype(), out_type);
+  } else {
+    out_type = x.dtype();
+  }
+
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[gather_qmm] Only real floating types are supported but "
+        << "x.dtype() == " << x.dtype() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Extract indices and broadcast them
+  array lhs_indices = indices_or_default(lhs_indices_, x, s);
+  array rhs_indices = indices_or_default(rhs_indices_, w, s);
+  std::tie(lhs_indices, rhs_indices) =
+      broadcast_arrays(lhs_indices, rhs_indices, s);
+
+  if (!issubdtype(lhs_indices.dtype(), integer)) {
+    throw std::invalid_argument(
+        "[gather_qmm] Got lhs_indices with invalid dtype. Indices must be integral.");
+  }
+
+  if (!issubdtype(rhs_indices.dtype(), integer)) {
+    throw std::invalid_argument(
+        "[gather_qmm] Got rhs_indices with invalid dtype. Indices must be integral.");
+  }
+  if (x.ndim() < 2) {
+    std::ostringstream msg;
+    msg << "[gather_qmm] Non-quantized input must have at least two"
+        << " dimensions but got input with shape " << x.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  lhs_indices = astype(lhs_indices, uint32, s);
+  rhs_indices = astype(rhs_indices, uint32, s);
+
+  // Compute the full output shape
+  auto out_shape = lhs_indices.shape();
+  out_shape.push_back(x.shape(-2));
+  out_shape.push_back(w_outer_dims);
+  std::vector<array> inputs;
+  if (qmode == QuantizationMode::Affine) {
+    inputs = {
+        astype(x, out_type, s),
+        std::move(w),
+        astype(scales, out_type, s),
+        astype(*biases, out_type, s),
+        std::move(lhs_indices),
+        std::move(rhs_indices)};
+  } else {
+    inputs = {
+        astype(x, out_type, s),
+        std::move(w),
+        std::move(scales),
+        std::move(lhs_indices),
+        std::move(rhs_indices)};
+  }
+  return array(
+      std::move(out_shape),
+      out_type,
+      std::make_shared<GatherQMM>(
+          to_stream(s),
+          group_size,
+          bits,
+          qmode,
+          transpose,
+          sorted_indices && !rhs_indices_,
+          sorted_indices && !lhs_indices_),
+      std::move(inputs));
+}
+
+array tensordot(
+    const array& a,
+    const array& b,
+    const int axis /* = 2 */,
+    StreamOrDevice s /* = {} */
+) {
+  if (axis < 0) {
+    throw std::invalid_argument(
+        "[tensordot] axis must be greater or equal to 0.");
+  }
+  if (axis > std::min(a.ndim(), b.ndim())) {
+    throw std::invalid_argument(
+        "[tensordot] axis must be less than the number of dimensions of a and b.");
+  }
+  std::vector<int> adims;
+  std::vector<int> bdims;
+  for (int i = 0; i < axis; i++) {
+    bdims.emplace_back(i);
+    adims.emplace_back(i - axis);
+  }
+  return tensordot(a, b, {adims}, {bdims}, s);
+}
+
+array tensordot(
+    const array& a,
+    const array& b,
+    const std::vector<int>& axes_a,
+    const std::vector<int>& axes_b,
+    StreamOrDevice s /* = {} */) {
+  if (axes_a.size() != axes_b.size()) {
+    throw std::invalid_argument("[tensordot] axes must have the same size.");
+  }
+  int csize = 1;
+  auto x = a;
+  auto y = b;
+  for (int i = 0; i < axes_a.size(); i++) {
+    if (x.shape(axes_a.at(i)) == y.shape(axes_b.at(i))) {
+      csize *= x.shape(axes_a.at(i));
+    } else {
+      throw std::invalid_argument(
+          "[tensordot] a and b must have the same shape on the contracted axes.");
+    }
+  }
+
+  std::vector<bool> cdims1(x.ndim(), false);
+  std::vector<bool> cdims2(y.ndim(), false);
+  for (const auto n : axes_a) {
+    int n_ = (n < 0) ? n + x.ndim() : n;
+    cdims1[n_] = true;
+  }
+  for (const auto n : axes_b) {
+    int n_ = (n < 0) ? n + y.ndim() : n;
+    cdims2[n_] = true;
+  }
+
+  std::vector<int> t1;
+  std::vector<int> t2;
+  Shape rshape;
+  int size1 = 1;
+  int size2 = 1;
+  for (int i = 0; i < a.ndim(); i++) {
+    if (!cdims1[i]) {
+      t1.emplace_back(i);
+      size1 *= a.shape(i);
+      rshape.emplace_back(a.shape(i));
+    }
+  }
+  for (const auto x : axes_a) {
+    t1.emplace_back(x);
+  }
+  for (const auto x : axes_b) {
+    t2.emplace_back(x);
+  }
+  for (int i = 0; i < b.ndim(); i++) {
+    if (!cdims2[i]) {
+      t2.emplace_back(i);
+      size2 *= b.shape(i);
+      rshape.emplace_back(b.shape(i));
+    }
+  }
+  x = reshape(transpose(x, t1, s), {size1, csize}, s);
+  y = reshape(transpose(y, t2, s), {csize, size2}, s);
+  return reshape(matmul(x, y, s), rshape, s);
+}
+
+array outer(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  return multiply(
+      reshape(a, {static_cast<int>(a.size()), 1}, s), flatten(b, s), s);
+}
+
+array inner(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  if (a.ndim() == 0 || b.ndim() == 0) {
+    return multiply(a, b, s);
+  }
+  if (a.shape(-1) != b.shape(-1)) {
+    throw std::invalid_argument(
+        "[inner] a and b must have the same last dimension.");
+  }
+
+  return tensordot(a, b, {-1}, {-1}, s);
+}
+
+/** Compute D = beta * C + alpha * (A @ B) */
+array addmm(
+    array c,
+    array a,
+    array b,
+    const float& alpha /* = 1.f */,
+    const float& beta /* = 1.f */,
+    StreamOrDevice s /* = {} */) {
+  int in_a_ndim = a.ndim();
+  int in_b_ndim = b.ndim();
+
+  if (a.ndim() == 0 || b.ndim() == 0) {
+    throw std::invalid_argument(
+        "[addmm] Got 0 dimension input. Inputs must "
+        "have at least one dimension.");
+  }
+
+  // Type promotion
+  auto out_type = result_type(a, b, c);
+
+  if (out_type == complex64) {
+    return add(
+        multiply(matmul(a, b, s), array(alpha), s),
+        multiply(array(beta), c, s),
+        s);
+  }
+
+  if (a.ndim() == 1) {
+    // Insert a singleton dim in the beginning
+    a = expand_dims(a, 0, s);
+  }
+  if (b.ndim() == 1) {
+    // Insert a singleton dim at the end
+    b = expand_dims(b, 1, s);
+  }
+
+  if (a.shape(-1) != b.shape(-2)) {
+    std::ostringstream msg;
+    msg << "[addmm] Last dimension of first input with shape " << a.shape()
+        << " must match second to last dimension of"
+        << " second input with shape " << b.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[addmm] Only real floating point types are supported but "
+        << c.dtype() << ", " << a.dtype() << " and " << b.dtype()
+        << " were provided which results in " << out_type
+        << ", which is not a real floating point type.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  a = astype(a, out_type, s);
+  b = astype(b, out_type, s);
+  c = astype(c, out_type, s);
+
+  // We can batch the multiplication by reshaping a
+  if (a.ndim() > 2 && b.ndim() == 2 && c.ndim() <= 1) {
+    auto out_shape = a.shape();
+    a = reshape(a, {-1, out_shape.back()}, s);
+    out_shape.back() = b.shape(-1);
+
+    if (in_b_ndim == 1) {
+      out_shape.pop_back();
+    }
+
+    c = broadcast_to(c, {a.shape(0), b.shape(1)}, s);
+
+    auto out = array(
+        {a.shape(0), b.shape(1)},
+        out_type,
+        std::make_shared<AddMM>(to_stream(s), alpha, beta),
+        {a, b, c});
+    return reshape(out, out_shape, s);
+  }
+
+  if (a.ndim() > 2 || b.ndim() > 2) {
+    Shape bsx_a(a.shape().begin(), a.shape().end() - 2);
+    Shape bsx_b(b.shape().begin(), b.shape().end() - 2);
+    auto inner_shape = broadcast_shapes(bsx_a, bsx_b);
+
+    // Broadcast a
+    inner_shape.push_back(a.shape(-2));
+    inner_shape.push_back(a.shape(-1));
+    a = broadcast_to(a, inner_shape, s);
+
+    // Broadcast b
+    *(inner_shape.end() - 2) = b.shape(-2);
+    *(inner_shape.end() - 1) = b.shape(-1);
+    b = broadcast_to(b, inner_shape, s);
+  }
+
+  auto out_shape = a.shape();
+  out_shape.back() = b.shape(-1);
+
+  auto out_shape_adjusted = out_shape;
+
+  if (in_a_ndim == 1 || in_b_ndim == 1) {
+    out_shape_adjusted.erase(
+        out_shape_adjusted.end() - ((in_a_ndim == 1) ? 2 : 1),
+        out_shape_adjusted.end() - ((in_b_ndim == 1) ? 0 : 1));
+  }
+
+  auto c_broadcast_shape = broadcast_shapes(c.shape(), out_shape_adjusted);
+  c = broadcast_to(c, c_broadcast_shape, s);
+
+  if (in_a_ndim == 1 || in_b_ndim == 1) {
+    auto c_reshape = c.shape();
+    if (in_b_ndim == 1) {
+      c_reshape.push_back(1);
+    }
+
+    if (in_a_ndim == 1) {
+      c_reshape.push_back(c_reshape.back());
+      c_reshape[c_reshape.size() - 2] = 1;
+    }
+
+    c = reshape(c, c_reshape, s);
+  }
+  if (c.shape() != out_shape) {
+    throw std::invalid_argument(
+        "[addmm] input c must broadcast to the output shape");
+  }
+
+  auto out = array(
+      std::move(out_shape),
+      out_type,
+      std::make_shared<AddMM>(to_stream(s), alpha, beta),
+      {a, b, c});
+
+  // Remove the possibly inserted singleton dimensions
+  std::vector<int> axes;
+  if (in_a_ndim == 1) {
+    axes.push_back(out.ndim() - 2);
+  }
+  if (in_b_ndim == 1) {
+    axes.push_back(out.ndim() - 1);
+  }
+  return axes.empty() ? out : squeeze(out, axes, s);
+}
+
+/** Compute matrix product with tile-level masking */
+array block_masked_mm(
+    array a,
+    array b,
+    int block_size,
+    std::optional<array> mask_out /* = std::nullopt */,
+    std::optional<array> mask_lhs /* = std::nullopt */,
+    std::optional<array> mask_rhs /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  // If no masks, just perform regular matmul
+  if (!mask_out && !mask_lhs && !mask_rhs) {
+    return matmul(a, b, s);
+  }
+
+  bool has_out_mask = mask_out.has_value();
+  bool has_operand_mask = mask_lhs.has_value() || mask_rhs.has_value();
+
+  // Check valid tile sizes
+  // TODO: Add support for 16x16 tile
+  if (block_size != 32 && block_size != 64) {
+    std::ostringstream msg;
+    msg << "[block_masked_mm] Only block_sizes 32, 64 are supported."
+        << "Got block size " << block_size << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Do shape checks for operands
+  int in_a_ndim = a.ndim();
+  int in_b_ndim = b.ndim();
+
+  if (a.ndim() == 0 || b.ndim() == 0) {
+    throw std::invalid_argument(
+        "[block_masked_mm] Got 0 dimension input. Inputs must "
+        "have at least one dimension.");
+  }
+
+  if (a.ndim() == 1) {
+    // Insert a singleton dim in the beginning
+    a = expand_dims(a, 0, s);
+  }
+  if (b.ndim() == 1) {
+    // Insert a singleton dim at the end
+    b = expand_dims(b, 1, s);
+  }
+
+  if (a.shape(-1) != b.shape(-2)) {
+    std::ostringstream msg;
+    msg << "[block_masked_mm] Last dimension of first input with shape "
+        << a.shape() << " must match second to last dimension of"
+        << " second input with shape " << b.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Type promotion
+  auto out_type = result_type(a, b);
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[block_masked_mm] Only real floating point types are supported but "
+        << a.dtype() << " and " << b.dtype()
+        << " were provided which results in " << out_type
+        << ", which is not a real floating point type.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  a = astype(a, out_type, s);
+  b = astype(b, out_type, s);
+
+  // Handle broadcasting
+  Shape bsx_a(a.shape().begin(), a.shape().end() - 2);
+  Shape bsx_b(b.shape().begin(), b.shape().end() - 2);
+
+  auto bsx_shape = broadcast_shapes(bsx_a, bsx_b);
+
+  bsx_shape.push_back(1);
+  bsx_shape.push_back(1);
+  int nd = bsx_shape.size();
+
+  int M = a.shape(-2);
+  int N = b.shape(-1);
+  int K = a.shape(-1);
+
+  // Prepare A
+  bsx_shape[nd - 2] = M;
+  bsx_shape[nd - 1] = K;
+  a = broadcast_to(a, bsx_shape, s);
+
+  // Prepare B
+  bsx_shape[nd - 2] = K;
+  bsx_shape[nd - 1] = N;
+  b = broadcast_to(b, bsx_shape, s);
+
+  // Get output shape
+  auto out_shape = bsx_shape;
+  out_shape[nd - 2] = M;
+  out_shape[nd - 1] = N;
+
+  // Determine mask shape requirments
+  int tm = (M + block_size - 1) / block_size;
+  int tn = (N + block_size - 1) / block_size;
+  int tk = (K + block_size - 1) / block_size;
+
+  std::vector<array> inputs = {a, b};
+
+  // Broadcast and astype mask
+  auto broadcast_mask = [](array mask,
+                           Shape& bs_shape,
+                           int y,
+                           int x,
+                           Dtype mask_dtype,
+                           StreamOrDevice s) {
+    int nd_bsx = bs_shape.size();
+    bs_shape[nd_bsx - 2] = y;
+    bs_shape[nd_bsx - 1] = x;
+    mask = astype(mask, mask_dtype, s);
+    return broadcast_to(mask, bs_shape, s);
+  };
+
+  // Out mask
+  if (has_out_mask) {
+    array mask_out_p = mask_out.value_or(array({true}));
+    if (in_a_ndim == 1 || in_b_ndim == 1) {
+      std::vector<int> ex_dims;
+      if (in_a_ndim == 1)
+        ex_dims.push_back(-2);
+      if (in_b_ndim == 1)
+        ex_dims.push_back(-1);
+      mask_out_p = expand_dims(mask_out_p, ex_dims, s);
+    }
+    auto maskout_dtype = mask_out_p.dtype() == bool_ ? bool_ : out_type;
+    mask_out_p =
+        broadcast_mask(mask_out_p, bsx_shape, tm, tn, maskout_dtype, s);
+
+    inputs.push_back(mask_out_p);
+  }
+
+  // Operand masks
+  if (has_operand_mask) {
+    // Pull masks
+    array mask_lhs_p = mask_lhs.value_or(array({true}));
+    array mask_rhs_p = mask_rhs.value_or(array({true}));
+    auto mask_dtype =
+        (mask_lhs_p.dtype() == bool_ && mask_rhs_p.dtype() == bool_) ? bool_
+                                                                     : out_type;
+
+    // LHS mask
+    if (in_a_ndim == 1) {
+      mask_lhs_p = expand_dims(mask_lhs_p, -2, s);
+    }
+    mask_lhs_p = broadcast_mask(mask_lhs_p, bsx_shape, tm, tk, mask_dtype, s);
+
+    // RHS mask
+    if (in_b_ndim == 1) {
+      mask_rhs_p = expand_dims(mask_rhs_p, -1, s);
+    }
+    mask_rhs_p = broadcast_mask(mask_rhs_p, bsx_shape, tk, tn, mask_dtype, s);
+
+    inputs.push_back(mask_lhs_p);
+    inputs.push_back(mask_rhs_p);
+  }
+
+  // Caculate array
+  auto out = array(
+      std::move(out_shape),
+      out_type,
+      std::make_shared<BlockMaskedMM>(to_stream(s), block_size),
+      std::move(inputs));
+  // Remove the possibly inserted singleton dimensions
+  std::vector<int> axes;
+  if (in_a_ndim == 1) {
+    axes.push_back(out.ndim() - 2);
+  }
+  if (in_b_ndim == 1) {
+    axes.push_back(out.ndim() - 1);
+  }
+  return axes.empty() ? out : squeeze(out, axes, s);
+}
+
+/** Compute matrix product with matrix-level gather */
+array gather_mm(
+    array a,
+    array b,
+    std::optional<array> lhs_indices_ /* = std::nullopt */,
+    std::optional<array> rhs_indices_ /* = std::nullopt */,
+    bool sorted_indices /* = false */,
+    StreamOrDevice s /* = {} */) {
+  // If no indices, fall back to full matmul
+  if (!lhs_indices_ && !rhs_indices_) {
+    return matmul(a, b, s);
+  }
+
+  // Do shape checks for operands
+  int in_a_ndim = a.ndim();
+  int in_b_ndim = b.ndim();
+
+  if (a.ndim() == 0 || b.ndim() == 0) {
+    throw std::invalid_argument(
+        "[gather_mm] Got 0 dimension input. Inputs must "
+        "have at least one dimension.");
+  }
+
+  if (a.ndim() == 1) {
+    // Insert a singleton dim in the beginning
+    a = expand_dims(a, 0, s);
+  }
+  if (b.ndim() == 1) {
+    // Insert a singleton dim at the end
+    b = expand_dims(b, 1, s);
+  }
+
+  if (a.shape(-1) != b.shape(-2)) {
+    std::ostringstream msg;
+    msg << "[gather_mm] Last dimension of first input with shape " << a.shape()
+        << " must match second to last dimension of"
+        << " second input with shape " << b.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Type promotion
+  auto out_type = result_type(a, b);
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[gather_mm] Only real floating point types are supported but "
+        << a.dtype() << " and " << b.dtype()
+        << " were provided which results in " << out_type
+        << ", which is not a real floating point type.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  a = astype(a, out_type, s);
+  b = astype(b, out_type, s);
+
+  // Handle broadcasting
+  array lhs_indices = indices_or_default(lhs_indices_, a, s);
+  array rhs_indices = indices_or_default(rhs_indices_, b, s);
+
+  if (!issubdtype(lhs_indices.dtype(), integer)) {
+    throw std::invalid_argument(
+        "[gather_mm] Got lhs_indices with invalid dtype. Indices must be integral.");
+  }
+
+  if (!issubdtype(rhs_indices.dtype(), integer)) {
+    throw std::invalid_argument(
+        "[gather_mm] Got rhs_indices with invalid dtype. Indices must be integral.");
+  }
+
+  lhs_indices = astype(lhs_indices, uint32, s);
+  rhs_indices = astype(rhs_indices, uint32, s);
+
+  int M = a.shape(-2);
+  int N = b.shape(-1);
+
+  std::tie(lhs_indices, rhs_indices) =
+      broadcast_arrays(lhs_indices, rhs_indices, s);
+
+  auto out_shape = lhs_indices.shape();
+  out_shape.push_back(M);
+  out_shape.push_back(N);
+
+  // Make the output array
+  auto out = array(
+      std::move(out_shape),
+      out_type,
+      std::make_shared<GatherMM>(
+          to_stream(s),
+          sorted_indices && !rhs_indices_,
+          sorted_indices && !lhs_indices_),
+      {std::move(a),
+       std::move(b),
+       std::move(lhs_indices),
+       std::move(rhs_indices)});
+
+  // Remove the possibly inserted singleton dimensions
+  std::vector<int> axes;
+  if (in_a_ndim == 1) {
+    axes.push_back(out.ndim() - 2);
+  }
+  if (in_b_ndim == 1) {
+    axes.push_back(out.ndim() - 1);
+  }
+  return axes.empty() ? out : squeeze(out, axes, s);
+}
+
+array segmented_mm(
+    array a,
+    array b,
+    array segments,
+    StreamOrDevice s /* = {} */) {
+  if (a.ndim() != 2 || b.ndim() != 2) {
+    throw std::invalid_argument("[segmented_mm] Batched matmul not supported");
+  }
+
+  if (segments.ndim() < 1 || segments.shape().back() != 2) {
+    std::ostringstream msg;
+    msg << "[segmented_mm] The segments should have shape (..., 2) but "
+        << segments.shape() << " was provided.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Type promotion
+  auto out_type = result_type(a, b);
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[segmented_mm] Only real floating point types are supported but "
+        << a.dtype() << " and " << b.dtype()
+        << " were provided which results in " << out_type
+        << ", which is not a real floating point type.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (!issubdtype(segments.dtype(), integer)) {
+    throw std::invalid_argument(
+        "[segmented_mm] Got segments with invalid dtype. Segments must be integral.");
+  }
+
+  a = astype(a, out_type, s);
+  b = astype(b, out_type, s);
+  segments = astype(segments, uint32, s);
+
+  Shape out_shape = segments.shape();
+  out_shape.pop_back();
+  out_shape.push_back(a.shape(0));
+  out_shape.push_back(b.shape(1));
+
+  return array(
+      std::move(out_shape),
+      out_type,
+      std::make_shared<SegmentedMM>(to_stream(s)),
+      {std::move(a), std::move(b), std::move(segments)});
+}
+
+array diagonal(
+    const array& a,
+    int offset /* = 0 */,
+    int axis1 /* = 0 */,
+    int axis2 /* = 1 */,
+    StreamOrDevice s /* = {} */
+) {
+  int ndim = a.ndim();
+  if (ndim < 2) {
+    std::ostringstream msg;
+    msg << "[diagonal] Array must have at least two dimensions, but got "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto ax1 = (axis1 < 0) ? axis1 + ndim : axis1;
+  if (ax1 < 0 || ax1 >= ndim) {
+    std::ostringstream msg;
+    msg << "[diagonal] Invalid axis1 " << axis1 << " for array with " << ndim
+        << " dimensions.";
+    throw std::out_of_range(msg.str());
+  }
+
+  auto ax2 = (axis2 < 0) ? axis2 + ndim : axis2;
+  if (ax2 < 0 || ax2 >= ndim) {
+    std::ostringstream msg;
+    msg << "[diagonal] Invalid axis2 " << axis2 << " for array with " << ndim
+        << " dimensions.";
+    throw std::out_of_range(msg.str());
+  }
+
+  if (ax1 == ax2) {
+    throw std::invalid_argument(
+        "[diagonal] axis1 and axis2 cannot be the same axis");
+  }
+
+  ShapeElem off1 = std::max(-offset, 0);
+  ShapeElem off2 = std::max(offset, 0);
+
+  auto diag_size = std::min(a.shape(ax1) - off1, a.shape(ax2) - off2);
+  diag_size = diag_size < 0 ? 0 : diag_size;
+
+  std::vector<array> indices = {
+      arange(off1, off1 + diag_size, s), arange(off2, off2 + diag_size, s)};
+
+  Shape slice_sizes = a.shape();
+  slice_sizes[ax1] = 1;
+  slice_sizes[ax2] = 1;
+
+  auto out = gather(a, indices, {ax1, ax2}, slice_sizes, s);
+  return moveaxis(squeeze(out, {ax1 + 1, ax2 + 1}, s), 0, -1, s);
+}
+
+array diag(const array& a, int k /* = 0 */, StreamOrDevice s /* = {} */) {
+  if (a.ndim() == 1) {
+    int a_size = a.size();
+    int n = a_size + std::abs(k);
+    auto res = zeros({n, n}, a.dtype(), s);
+
+    std::vector<array> indices;
+    auto s1 = std::max(0, -k);
+    auto s2 = std::max(0, k);
+    indices.push_back(arange(s1, a_size + s1, uint32, s));
+    indices.push_back(arange(s2, a_size + s2, uint32, s));
+
+    return scatter(res, indices, reshape(a, {a_size, 1, 1}, s), {0, 1}, s);
+  } else if (a.ndim() == 2) {
+    return diagonal(a, k, 0, 1, s);
+  } else {
+    std::ostringstream msg;
+    msg << "[diag] array must be 1-D or 2-D, got array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+array trace(
+    const array& a,
+    int offset,
+    int axis1,
+    int axis2,
+    Dtype dtype,
+    StreamOrDevice s /* = {} */) {
+  int ndim = a.ndim();
+  if (ndim < 2) {
+    std::ostringstream msg;
+    msg << "[trace] Array must have at least two dimensions, but got " << ndim
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto ax1 = (axis1 < 0) ? axis1 + ndim : axis1;
+  if (ax1 < 0 || ax1 >= ndim) {
+    std::ostringstream msg;
+    msg << "[trace] Invalid axis1 " << axis1 << " for array with " << ndim
+        << " dimensions.";
+    throw std::out_of_range(msg.str());
+  }
+
+  auto ax2 = (axis2 < 0) ? axis2 + ndim : axis2;
+  if (ax2 < 0 || ax2 >= ndim) {
+    std::ostringstream msg;
+    msg << "[trace] Invalid axis2 " << axis2 << " for array with " << ndim
+        << " dimensions.";
+    throw std::out_of_range(msg.str());
+  }
+
+  if (ax1 == ax2) {
+    throw std::invalid_argument(
+        "[trace] axis1 and axis2 cannot be the same axis");
+  }
+
+  return sum(
+      astype(diagonal(a, offset, axis1, axis2, s), dtype, s),
+      /* axis = */ -1,
+      /* keepdims = */ false,
+      s);
+}
+array trace(
+    const array& a,
+    int offset,
+    int axis1,
+    int axis2,
+    StreamOrDevice s /* = {} */) {
+  auto dtype = a.dtype();
+  return trace(a, offset, axis1, axis2, dtype, s);
+}
+array trace(const array& a, StreamOrDevice s /* = {} */) {
+  auto dtype = a.dtype();
+  return trace(a, 0, 0, 1, dtype, s);
+}
+
+std::vector<array> depends(
+    const std::vector<array>& inputs,
+    const std::vector<array>& dependencies) {
+  std::vector<array> all_inputs = inputs;
+  all_inputs.insert(all_inputs.end(), dependencies.begin(), dependencies.end());
+
+  // Compute the stream. Maybe do it in a smarter way at some point in the
+  // future.
+  Stream s = (inputs[0].has_primitive()) ? inputs[0].primitive().stream()
+                                         : to_stream({});
+  // Make the output info
+  std::vector<Shape> shapes;
+  std::vector<Dtype> dtypes;
+  for (const auto& in : inputs) {
+    shapes.emplace_back(in.shape());
+    dtypes.emplace_back(in.dtype());
+  }
+
+  return array::make_arrays(
+      std::move(shapes),
+      dtypes,
+      std::make_shared<Depends>(to_stream(s)),
+      all_inputs);
+}
+
+array atleast_1d(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.ndim() == 0) {
+    return reshape(a, {1}, s);
+  }
+  return a;
+}
+
+std::vector<array> atleast_1d(
+    const std::vector<array>& arrays,
+    StreamOrDevice s /* = {} */) {
+  std::vector<array> out;
+  out.reserve(arrays.size());
+  for (const auto& a : arrays) {
+    out.push_back(atleast_1d(a, s));
+  }
+  return out;
+}
+
+array atleast_2d(const array& a, StreamOrDevice s /* = {} */) {
+  switch (a.ndim()) {
+    case 0:
+      return reshape(a, {1, 1}, s);
+    case 1:
+      return reshape(a, {1, a.shape(0)}, s);
+    default:
+      return a;
+  }
+}
+
+std::vector<array> atleast_2d(
+    const std::vector<array>& arrays,
+    StreamOrDevice s /* = {} */) {
+  std::vector<array> out;
+  out.reserve(arrays.size());
+  for (const auto& a : arrays) {
+    out.push_back(atleast_2d(a, s));
+  }
+  return out;
+}
+
+array atleast_3d(const array& a, StreamOrDevice s /* = {} */) {
+  switch (a.ndim()) {
+    case 0:
+      return reshape(a, {1, 1, 1}, s);
+    case 1:
+      return reshape(a, {1, a.shape(0), 1}, s);
+    case 2:
+      return reshape(a, {a.shape(0), a.shape(1), 1}, s);
+    default:
+      return a;
+  }
+}
+
+std::vector<array> atleast_3d(
+    const std::vector<array>& arrays,
+    StreamOrDevice s /* = {} */) {
+  std::vector<array> out;
+  out.reserve(arrays.size());
+  for (const auto& a : arrays) {
+    out.push_back(atleast_3d(a, s));
+  }
+  return out;
+}
+
+array number_of_elements(
+    const array& a,
+    std::vector<int> axes,
+    bool inverted,
+    Dtype dtype /* = int32 */,
+    StreamOrDevice s /* = {} */) {
+  for (auto& ax : axes) {
+    int normal_axis = (ax + a.ndim()) % a.ndim();
+    if (normal_axis >= a.ndim() || normal_axis < 0) {
+      std::ostringstream msg;
+      msg << "[number_of_elements] Can't get the shape for axis " << ax
+          << " from an array with " << a.ndim() << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    ax = normal_axis;
+  }
+
+  if (!detail::in_dynamic_tracing()) {
+    double numel = 1;
+    for (auto ax : axes) {
+      numel *= a.shape(ax);
+    }
+    return array(inverted ? 1.0 / numel : numel, dtype);
+  }
+  return stop_gradient(array(
+      Shape{},
+      dtype,
+      std::make_shared<NumberOfElements>(
+          to_stream(s), std::move(axes), inverted, dtype),
+      {a}));
+}
+
+array conjugate(const array& a, StreamOrDevice s /* = {} */) {
+  // Mirror NumPy's behaviour for real input
+  if (a.dtype() != complex64) {
+    return a;
+  }
+  return array(
+      a.shape(), a.dtype(), std::make_shared<Conjugate>(to_stream(s)), {a});
+}
+
+array bitwise_impl(
+    const array& a,
+    const array& b,
+    BitwiseBinary::Op op,
+    const std::string& op_name,
+    const StreamOrDevice& s,
+    std::optional<Dtype> out_type_ = std::nullopt) {
+  auto out_type = out_type_ ? *out_type_ : promote_types(a.dtype(), b.dtype());
+  if (!(issubdtype(out_type, integer) || out_type == bool_)) {
+    std::ostringstream msg;
+    msg << "[" << op_name
+        << "] Only allowed on integer or boolean types "
+           "but got types "
+        << a.dtype() << " and " << b.dtype() << ".";
+    throw std::runtime_error(msg.str());
+  }
+  auto inputs =
+      broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s);
+  auto& out_shape = inputs[0].shape();
+  return array(
+      out_shape,
+      out_type,
+      std::make_shared<BitwiseBinary>(to_stream(s), op),
+      std::move(inputs));
+}
+
+array bitwise_and(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  return bitwise_impl(a, b, BitwiseBinary::Op::And, "bitwise_and", s);
+}
+array operator&(const array& a, const array& b) {
+  return bitwise_and(a, b);
+}
+
+array bitwise_or(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  return bitwise_impl(a, b, BitwiseBinary::Op::Or, "bitwise_or", s);
+}
+array operator|(const array& a, const array& b) {
+  return bitwise_or(a, b);
+}
+
+array bitwise_xor(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  return bitwise_impl(a, b, BitwiseBinary::Op::Xor, "bitwise_xor", s);
+}
+array operator^(const array& a, const array& b) {
+  return bitwise_xor(a, b);
+}
+
+array left_shift(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto t = result_type(a, b);
+  if (t == bool_) {
+    t = uint8;
+  }
+  return bitwise_impl(a, b, BitwiseBinary::Op::LeftShift, "left_shift", s, t);
+}
+array operator<<(const array& a, const array& b) {
+  return left_shift(a, b);
+}
+
+array right_shift(const array& a, const array& b, StreamOrDevice s /* = {} */) {
+  auto t = result_type(a, b);
+  if (t == bool_) {
+    t = uint8;
+  }
+  return bitwise_impl(
+      astype(a, t, s),
+      astype(b, t, s),
+      BitwiseBinary::Op::RightShift,
+      "right_shift",
+      s,
+      t);
+}
+array operator>>(const array& a, const array& b) {
+  return right_shift(a, b);
+}
+
+array bitwise_invert(const array& a, StreamOrDevice s /* = {} */) {
+  if (issubdtype(a.dtype(), inexact)) {
+    throw std::invalid_argument(
+        "[bitwise_invert] Bitwise inverse only allowed on integer types.");
+  } else if (a.dtype() == bool_) {
+    return logical_not(a, s);
+  }
+  return array(
+      a.shape(), a.dtype(), std::make_shared<BitwiseInvert>(to_stream(s)), {a});
+}
+
+array operator~(const array& a) {
+  return bitwise_invert(a);
+}
+
+array view(const array& a, const Dtype& dtype, StreamOrDevice s /* = {} */) {
+  if (a.dtype() == dtype) {
+    return a;
+  }
+  auto out_shape = a.shape();
+  auto ibytes = size_of(a.dtype());
+  auto obytes = size_of(dtype);
+  if (a.ndim() == 0 && ibytes != obytes) {
+    throw std::invalid_argument(
+        "[view] Changing the type of a scalar is only allowed"
+        " for types with the same size.");
+  } else {
+    if (ibytes < obytes) {
+      if (out_shape.back() % (obytes / ibytes) != 0) {
+        throw std::invalid_argument(
+            "[view] When viewing as a larger dtype, the size in bytes of the last"
+            " axis must be a multiple of the requested type size.");
+      }
+      out_shape.back() /= (obytes / ibytes);
+    } else if (ibytes > obytes) {
+      // Type size ratios are always integers
+      out_shape.back() *= (ibytes / obytes);
+    }
+  }
+  return array(
+      out_shape, dtype, std::make_shared<View>(to_stream(s), dtype), {a});
+}
+
+array roll(
+    const array& a,
+    const Shape& shift,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+
+  if (shift.size() < axes.size()) {
+    std::ostringstream msg;
+    msg << "[roll] At least one shift value per axis is required, "
+        << shift.size() << " provided for " << axes.size() << " axes.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  array result = a;
+  for (int i = 0; i < axes.size(); i++) {
+    int ax = axes[i];
+    if (ax < 0) {
+      ax += a.ndim();
+    }
+    if (ax < 0 || ax >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[roll] Invalid axis " << axes[i] << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+
+    auto sh = shift[i];
+    auto size = a.shape(ax);
+    if (size == 0) {
+      continue; // skip rolling this axis if it has size 0
+    }
+    auto split_index = (sh < 0) ? (-sh) % size : size - sh % size;
+
+    auto parts = split(result, Shape{split_index}, ax, s);
+    std::swap(parts[0], parts[1]);
+    result = concatenate(parts, ax, s);
+  }
+
+  return result;
+}
+
+array roll(const array& a, int shift, StreamOrDevice s /* = {} */) {
+  auto shape = a.shape();
+  return reshape(
+      roll(flatten(a, s), Shape{shift}, std::vector<int>{0}, s),
+      std::move(shape),
+      s);
+}
+
+array roll(const array& a, const Shape& shift, StreamOrDevice s /* = {} */) {
+  int total_shift = 0;
+  for (auto& s : shift) {
+    total_shift += s;
+  }
+  return roll(a, total_shift, s);
+}
+
+array roll(const array& a, int shift, int axis, StreamOrDevice s /* = {} */) {
+  return roll(a, Shape{shift}, std::vector<int>{axis}, s);
+}
+
+array roll(
+    const array& a,
+    int shift,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  Shape shifts(axes.size(), shift);
+  return roll(a, shifts, axes, s);
+}
+
+array roll(
+    const array& a,
+    const Shape& shift,
+    int axis,
+    StreamOrDevice s /* = {} */) {
+  int total_shift = 0;
+  for (auto& s : shift) {
+    total_shift += s;
+  }
+  return roll(a, Shape{total_shift}, std::vector<int>{axis}, s);
+}
+
+array real(const array& a, StreamOrDevice s /* = {} */) {
+  if (!issubdtype(a.dtype(), complexfloating)) {
+    return a;
+  }
+  return array(a.shape(), float32, std::make_shared<Real>(to_stream(s)), {a});
+}
+
+array imag(const array& a, StreamOrDevice s /* = {} */) {
+  if (!issubdtype(a.dtype(), complexfloating)) {
+    return zeros_like(a);
+  }
+  return array(a.shape(), float32, std::make_shared<Imag>(to_stream(s)), {a});
+}
+
+array contiguous(
+    const array& a,
+    bool allow_col_major /* = false */,
+    StreamOrDevice s /* = {} */) {
+  return array(
+      a.shape(),
+      a.dtype(),
+      std::make_shared<Contiguous>(to_stream(s), allow_col_major),
+      {a});
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/ops.h b/macos/third_party/mlx/mlx/ops.h
new file mode 100644
index 0000000..1ff3bbf
--- /dev/null
+++ b/macos/third_party/mlx/mlx/ops.h
@@ -0,0 +1,1610 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <optional>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/device.h"
+#include "mlx/stream.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+/**
+ * \defgroup ops Core array operations
+ * @{
+ */
+
+/**
+ * A 1D array of numbers starting at `start` (optional),
+ * stopping at stop, stepping by `step` (optional). */
+MLX_API array arange(
+    double start,
+    double stop,
+    double step,
+    Dtype dtype,
+    StreamOrDevice s = {});
+MLX_API array
+arange(double start, double stop, double step, StreamOrDevice s = {});
+MLX_API array
+arange(double start, double stop, Dtype dtype, StreamOrDevice s = {});
+MLX_API array arange(double start, double stop, StreamOrDevice s = {});
+MLX_API array arange(double stop, Dtype dtype, StreamOrDevice s = {});
+MLX_API array arange(double stop, StreamOrDevice s = {});
+
+MLX_API array arange(int start, int stop, int step, StreamOrDevice s = {});
+MLX_API array arange(int start, int stop, StreamOrDevice s = {});
+MLX_API array arange(int stop, StreamOrDevice s = {});
+
+/** A 1D array of `num` evenly spaced numbers in the range `[start, stop]` */
+MLX_API array linspace(
+    double start,
+    double stop,
+    int num = 50,
+    Dtype dtype = float32,
+    StreamOrDevice s = {});
+
+/** Convert an array to the given data type. */
+MLX_API array astype(array a, Dtype dtype, StreamOrDevice s = {});
+
+/** Create a view of an array with the given shape and strides. */
+MLX_API array as_strided(
+    array a,
+    Shape shape,
+    Strides strides,
+    size_t offset,
+    StreamOrDevice s = {});
+
+/** Copy another array. */
+MLX_API array copy(array a, StreamOrDevice s = {});
+
+/** Fill an array of the given shape with the given value(s). */
+MLX_API array full(Shape shape, array vals, Dtype dtype, StreamOrDevice s = {});
+MLX_API array full(Shape shape, array vals, StreamOrDevice s = {});
+template <typename T>
+array full(Shape shape, T val, Dtype dtype, StreamOrDevice s = {}) {
+  return full(std::move(shape), array(val, dtype), to_stream(s));
+}
+template <typename T>
+array full(Shape shape, T val, StreamOrDevice s = {}) {
+  return full(std::move(shape), array(val), to_stream(s));
+}
+
+MLX_API array
+full_like(const array& a, array vals, Dtype dtype, StreamOrDevice s = {});
+MLX_API array full_like(const array& a, array vals, StreamOrDevice s = {});
+template <typename T>
+array full_like(const array& a, T val, Dtype dtype, StreamOrDevice s = {}) {
+  return full_like(a, array(val, dtype), dtype, to_stream(s));
+}
+template <typename T>
+array full_like(const array& a, T val, StreamOrDevice s = {}) {
+  return full_like(a, array(val, a.dtype()), to_stream(s));
+}
+
+/** Fill an array of the given shape with zeros. */
+MLX_API array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s = {});
+inline array zeros(const Shape& shape, StreamOrDevice s = {}) {
+  return zeros(shape, float32, s);
+}
+MLX_API array zeros_like(const array& a, StreamOrDevice s = {});
+
+/** Fill an array of the given shape with ones. */
+MLX_API array ones(const Shape& shape, Dtype dtype, StreamOrDevice s = {});
+inline array ones(const Shape& shape, StreamOrDevice s = {}) {
+  return ones(shape, float32, s);
+}
+MLX_API array ones_like(const array& a, StreamOrDevice s = {});
+
+/** Fill an array of the given shape (n,m) with ones in the specified diagonal
+ * k, and zeros everywhere else. */
+MLX_API array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s = {});
+inline array eye(int n, Dtype dtype, StreamOrDevice s = {}) {
+  return eye(n, n, 0, dtype, s);
+}
+inline array eye(int n, int m, StreamOrDevice s = {}) {
+  return eye(n, m, 0, float32, s);
+}
+inline array eye(int n, int m, int k, StreamOrDevice s = {}) {
+  return eye(n, m, k, float32, s);
+}
+inline array eye(int n, StreamOrDevice s = {}) {
+  return eye(n, n, 0, float32, s);
+}
+
+/** Create a square matrix of shape (n,n) of zeros, and ones in the major
+ * diagonal. */
+MLX_API array identity(int n, Dtype dtype, StreamOrDevice s = {});
+inline array identity(int n, StreamOrDevice s = {}) {
+  return identity(n, float32, s);
+}
+
+MLX_API array tri(int n, int m, int k, Dtype type, StreamOrDevice s = {});
+inline array tri(int n, Dtype type, StreamOrDevice s = {}) {
+  return tri(n, n, 0, type, s);
+}
+
+MLX_API array tril(array x, int k = 0, StreamOrDevice s = {});
+MLX_API array triu(array x, int k = 0, StreamOrDevice s = {});
+
+/** Reshape an array to the given shape. */
+MLX_API array reshape(const array& a, Shape shape, StreamOrDevice s = {});
+
+/** Unflatten the axis to the given shape. */
+MLX_API array
+unflatten(const array& a, int axis, Shape shape, StreamOrDevice s = {});
+
+/** Flatten the dimensions in the range `[start_axis, end_axis]` . */
+MLX_API array flatten(
+    const array& a,
+    int start_axis,
+    int end_axis = -1,
+    StreamOrDevice s = {});
+
+/** Flatten the array to 1D. */
+MLX_API array flatten(const array& a, StreamOrDevice s = {});
+
+/** Multiply the array by the Hadamard matrix of corresponding size. */
+MLX_API array hadamard_transform(
+    const array& a,
+    std::optional<float> scale = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Remove singleton dimensions at the given axes. */
+MLX_API array
+squeeze(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+
+/** Remove singleton dimensions at the given axis. */
+MLX_API array squeeze(const array& a, int axis, StreamOrDevice s = {});
+
+/** Remove all singleton dimensions. */
+MLX_API array squeeze(const array& a, StreamOrDevice s = {});
+
+/** Add a singleton dimension at the given axes. */
+MLX_API array expand_dims(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+
+/** Add a singleton dimension at the given axis. */
+MLX_API array expand_dims(const array& a, int axis, StreamOrDevice s = {});
+
+/** Slice an array. */
+MLX_API array slice(
+    const array& a,
+    Shape start,
+    Shape stop,
+    Shape strides,
+    StreamOrDevice s = {});
+inline array slice(
+    const array& a,
+    std::initializer_list<int> start,
+    Shape stop,
+    Shape strides,
+    StreamOrDevice s = {}) {
+  return slice(a, Shape(start), std::move(stop), std::move(strides), s);
+}
+
+/** Slice an array with a stride of 1 in each dimension. */
+MLX_API array
+slice(const array& a, Shape start, Shape stop, StreamOrDevice s = {});
+
+/** Slice an array with dynamic starting indices. */
+MLX_API array slice(
+    const array& a,
+    const array& start,
+    std::vector<int> axes,
+    Shape slice_size,
+    StreamOrDevice s = {});
+
+/** Update a slice from the source array. */
+MLX_API array slice_update(
+    const array& src,
+    const array& update,
+    Shape start,
+    Shape stop,
+    Shape strides,
+    StreamOrDevice s = {});
+
+/** Update a slice from the source array with stride 1 in each dimension. */
+MLX_API array slice_update(
+    const array& src,
+    const array& update,
+    Shape start,
+    Shape stop,
+    StreamOrDevice s = {});
+
+/** Update a slice from the source array with dynamic starting indices. */
+MLX_API array slice_update(
+    const array& src,
+    const array& update,
+    const array& start,
+    std::vector<int> axes,
+    StreamOrDevice s = {});
+
+/** Split an array into sub-arrays along a given axis. */
+MLX_API std::vector<array>
+split(const array& a, int num_splits, int axis, StreamOrDevice s = {});
+MLX_API std::vector<array>
+split(const array& a, int num_splits, StreamOrDevice s = {});
+MLX_API std::vector<array>
+split(const array& a, const Shape& indices, int axis, StreamOrDevice s = {});
+MLX_API std::vector<array>
+split(const array& a, const Shape& indices, StreamOrDevice s = {});
+
+/** A vector of coordinate arrays from coordinate vectors. */
+MLX_API std::vector<array> meshgrid(
+    const std::vector<array>& arrays,
+    bool sparse = false,
+    const std::string& indexing = "xy",
+    StreamOrDevice s = {});
+
+/**
+ * Clip (limit) the values in an array.
+ */
+MLX_API array clip(
+    const array& a,
+    const std::optional<array>& a_min = std::nullopt,
+    const std::optional<array>& a_max = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Concatenate arrays along a given axis. */
+MLX_API array
+concatenate(std::vector<array> arrays, int axis, StreamOrDevice s = {});
+MLX_API array concatenate(std::vector<array> arrays, StreamOrDevice s = {});
+
+/** Stack arrays along a new axis. */
+MLX_API array
+stack(const std::vector<array>& arrays, int axis, StreamOrDevice s = {});
+MLX_API array stack(const std::vector<array>& arrays, StreamOrDevice s = {});
+
+/** Repeat an array along an axis. */
+MLX_API array
+repeat(const array& arr, int repeats, int axis, StreamOrDevice s = {});
+MLX_API array repeat(const array& arr, int repeats, StreamOrDevice s = {});
+
+MLX_API array
+tile(const array& arr, std::vector<int> reps, StreamOrDevice s = {});
+
+/** Permutes the dimensions according to the given axes. */
+MLX_API array
+transpose(const array& a, std::vector<int> axes, StreamOrDevice s = {});
+inline array transpose(
+    const array& a,
+    std::initializer_list<int> axes,
+    StreamOrDevice s = {}) {
+  return transpose(a, std::vector<int>(axes), s);
+}
+
+/** Swap two axes of an array. */
+MLX_API array
+swapaxes(const array& a, int axis1, int axis2, StreamOrDevice s = {});
+
+/** Move an axis of an array. */
+MLX_API array
+moveaxis(const array& a, int source, int destination, StreamOrDevice s = {});
+
+/** Pad an array with a constant value */
+MLX_API array
+pad(const array& a,
+    const std::vector<int>& axes,
+    const Shape& low_pad_size,
+    const Shape& high_pad_size,
+    const array& pad_value = array(0),
+    const std::string& mode = "constant",
+    StreamOrDevice s = {});
+
+/** Pad an array with a constant value along all axes */
+MLX_API array
+pad(const array& a,
+    const std::vector<std::pair<int, int>>& pad_width,
+    const array& pad_value = array(0),
+    const std::string& mode = "constant",
+    StreamOrDevice s = {});
+MLX_API array
+pad(const array& a,
+    const std::pair<int, int>& pad_width,
+    const array& pad_value = array(0),
+    const std::string& mode = "constant",
+    StreamOrDevice s = {});
+MLX_API array
+pad(const array& a,
+    int pad_width,
+    const array& pad_value = array(0),
+    const std::string& mode = "constant",
+    StreamOrDevice s = {});
+
+/** Permutes the dimensions in reverse order. */
+MLX_API array transpose(const array& a, StreamOrDevice s = {});
+
+/** Broadcast an array to a given shape. */
+MLX_API array
+broadcast_to(const array& a, const Shape& shape, StreamOrDevice s = {});
+
+/** Broadcast a vector of arrays against one another. */
+MLX_API std::vector<array> broadcast_arrays(
+    const std::vector<array>& inputs,
+    StreamOrDevice s = {});
+
+/** Returns the bool array with (a == b) element-wise. */
+MLX_API array equal(const array& a, const array& b, StreamOrDevice s = {});
+inline array operator==(const array& a, const array& b) {
+  return equal(a, b);
+}
+template <typename T>
+array operator==(T a, const array& b) {
+  return equal(array(a), b);
+}
+template <typename T>
+array operator==(const array& a, T b) {
+  return equal(a, array(b));
+}
+
+/** Returns the bool array with (a != b) element-wise. */
+MLX_API array not_equal(const array& a, const array& b, StreamOrDevice s = {});
+inline array operator!=(const array& a, const array& b) {
+  return not_equal(a, b);
+}
+template <typename T>
+array operator!=(T a, const array& b) {
+  return not_equal(array(a), b);
+}
+template <typename T>
+array operator!=(const array& a, T b) {
+  return not_equal(a, array(b));
+}
+
+/** Returns bool array with (a > b) element-wise. */
+MLX_API array greater(const array& a, const array& b, StreamOrDevice s = {});
+inline array operator>(const array& a, const array& b) {
+  return greater(a, b);
+}
+template <typename T>
+array operator>(T a, const array& b) {
+  return greater(array(a), b);
+}
+template <typename T>
+array operator>(const array& a, T b) {
+  return greater(a, array(b));
+}
+
+/** Returns bool array with (a >= b) element-wise. */
+MLX_API array
+greater_equal(const array& a, const array& b, StreamOrDevice s = {});
+inline array operator>=(const array& a, const array& b) {
+  return greater_equal(a, b);
+}
+template <typename T>
+array operator>=(T a, const array& b) {
+  return greater_equal(array(a), b);
+}
+template <typename T>
+array operator>=(const array& a, T b) {
+  return greater_equal(a, array(b));
+}
+
+/** Returns bool array with (a < b) element-wise. */
+MLX_API array less(const array& a, const array& b, StreamOrDevice s = {});
+inline array operator<(const array& a, const array& b) {
+  return less(a, b);
+}
+template <typename T>
+array operator<(T a, const array& b) {
+  return less(array(a), b);
+}
+template <typename T>
+array operator<(const array& a, T b) {
+  return less(a, array(b));
+}
+
+/** Returns bool array with (a <= b) element-wise. */
+MLX_API array less_equal(const array& a, const array& b, StreamOrDevice s = {});
+inline array operator<=(const array& a, const array& b) {
+  return less_equal(a, b);
+}
+template <typename T>
+array operator<=(T a, const array& b) {
+  return less_equal(array(a), b);
+}
+template <typename T>
+array operator<=(const array& a, T b) {
+  return less_equal(a, array(b));
+}
+
+/** True if two arrays have the same shape and elements. */
+MLX_API array array_equal(
+    const array& a,
+    const array& b,
+    bool equal_nan,
+    StreamOrDevice s = {});
+inline array
+array_equal(const array& a, const array& b, StreamOrDevice s = {}) {
+  return array_equal(a, b, false, s);
+}
+
+MLX_API array isnan(const array& a, StreamOrDevice s = {});
+
+MLX_API array isinf(const array& a, StreamOrDevice s = {});
+
+MLX_API array isfinite(const array& a, StreamOrDevice s = {});
+
+MLX_API array isposinf(const array& a, StreamOrDevice s = {});
+
+MLX_API array isneginf(const array& a, StreamOrDevice s = {});
+
+/** Select from x or y depending on condition. */
+MLX_API array where(
+    const array& condition,
+    const array& x,
+    const array& y,
+    StreamOrDevice s = {});
+
+/** Replace NaN and infinities with finite numbers. */
+MLX_API array nan_to_num(
+    const array& a,
+    float nan = 0.0f,
+    const std::optional<float> posinf = std::nullopt,
+    const std::optional<float> neginf = std::nullopt,
+    StreamOrDevice s = {});
+
+/** True if all elements in the array are true (or non-zero). **/
+MLX_API array all(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array all(const array& a, StreamOrDevice s = {}) {
+  return all(a, false, to_stream(s));
+}
+
+/** True if the two arrays are equal within the specified tolerance. */
+MLX_API array allclose(
+    const array& a,
+    const array& b,
+    double rtol = 1e-5,
+    double atol = 1e-8,
+    bool equal_nan = false,
+    StreamOrDevice s = {});
+
+/** Returns a boolean array where two arrays are element-wise equal within the
+ * specified tolerance. */
+MLX_API array isclose(
+    const array& a,
+    const array& b,
+    double rtol = 1e-5,
+    double atol = 1e-8,
+    bool equal_nan = false,
+    StreamOrDevice s = {});
+
+/**
+ *  Reduces the input along the given axes. An output value is true
+ *  if all the corresponding inputs are true.
+ **/
+MLX_API array
+all(const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/**
+ *  Reduces the input along the given axis. An output value is true
+ *  if all the corresponding inputs are true.
+ **/
+MLX_API array
+all(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** True if any elements in the array are true (or non-zero). **/
+MLX_API array any(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array any(const array& a, StreamOrDevice s = {}) {
+  return any(a, false, to_stream(s));
+}
+
+/**
+ *  Reduces the input along the given axes. An output value is true
+ *  if any of the corresponding inputs are true.
+ **/
+MLX_API array
+any(const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/**
+ *  Reduces the input along the given axis. An output value is true
+ *  if any of the corresponding inputs are true.
+ **/
+MLX_API array
+any(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** Sums the elements of an array. */
+MLX_API array sum(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array sum(const array& a, StreamOrDevice s = {}) {
+  return sum(a, false, to_stream(s));
+}
+
+/** Sums the elements of an array along the given axes. */
+MLX_API array
+sum(const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** Sums the elements of an array along the given axis. */
+MLX_API array
+sum(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** Computes the mean of the elements of an array. */
+MLX_API array mean(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array mean(const array& a, StreamOrDevice s = {}) {
+  return mean(a, false, to_stream(s));
+}
+
+/** Computes the mean of the elements of an array along the given axes */
+MLX_API array mean(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** Computes the mean of the elements of an array along the given axis */
+MLX_API array
+mean(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** Computes the median of the elements of an array. */
+MLX_API array median(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array median(const array& a, StreamOrDevice s = {}) {
+  return median(a, false, to_stream(s));
+}
+
+/** Computes the median of the elements of an array along the given axes */
+MLX_API array median(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** Computes the median of the elements of an array along the given axis */
+MLX_API array
+median(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** Computes the variance of the elements of an array. */
+MLX_API array
+var(const array& a, bool keepdims, int ddof = 0, StreamOrDevice s = {});
+inline array var(const array& a, StreamOrDevice s = {}) {
+  return var(a, false, 0, to_stream(s));
+}
+
+/** Computes the variance of the elements of an array along the given
+ * axes */
+MLX_API array
+var(const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    int ddof = 0,
+    StreamOrDevice s = {});
+
+/** Computes the variance of the elements of an array along the given
+ * axis */
+MLX_API array
+var(const array& a,
+    int axis,
+    bool keepdims = false,
+    int ddof = 0,
+    StreamOrDevice s = {});
+
+/** Computes the standard deviation of the elements of an array. */
+MLX_API array
+std(const array& a, bool keepdims, int ddof = 0, StreamOrDevice s = {});
+inline array std(const array& a, StreamOrDevice s = {}) {
+  return std(a, false, 0, to_stream(s));
+}
+
+/** Computes the standard deviation of the elements of an array along the given
+ * axes */
+MLX_API array
+std(const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    int ddof = 0,
+    StreamOrDevice s = {});
+
+/** Computes the standard deviation of the elements of an array along the given
+ * axis */
+MLX_API array
+std(const array& a,
+    int axis,
+    bool keepdims = false,
+    int ddof = 0,
+    StreamOrDevice s = {});
+
+/** The product of all elements of the array. */
+MLX_API array prod(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array prod(const array& a, StreamOrDevice s = {}) {
+  return prod(a, false, to_stream(s));
+}
+
+/** The product of the elements of an array along the given axes. */
+MLX_API array prod(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** The product of the elements of an array along the given axis. */
+MLX_API array
+prod(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** The maximum of all elements of the array. */
+MLX_API array max(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array max(const array& a, StreamOrDevice s = {}) {
+  return max(a, false, to_stream(s));
+}
+
+/** The maximum of the elements of an array along the given axes. */
+MLX_API array
+max(const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** The maximum of the elements of an array along the given axis. */
+MLX_API array
+max(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** The minimum of all elements of the array. */
+MLX_API array min(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array min(const array& a, StreamOrDevice s = {}) {
+  return min(a, false, to_stream(s));
+}
+
+/** The minimum of the elements of an array along the given axes. */
+MLX_API array
+min(const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** The minimum of the elements of an array along the given axis. */
+MLX_API array
+min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** Returns the index of the minimum value in the array. */
+MLX_API array argmin(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array argmin(const array& a, StreamOrDevice s = {}) {
+  return argmin(a, false, s);
+}
+
+/** Returns the indices of the minimum values along a given axis. */
+MLX_API array
+argmin(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** Returns the index of the maximum value in the array. */
+MLX_API array argmax(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array argmax(const array& a, StreamOrDevice s = {}) {
+  return argmax(a, false, s);
+}
+
+/** Returns the indices of the maximum values along a given axis. */
+MLX_API array
+argmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {});
+
+/** Returns a sorted copy of the flattened array. */
+MLX_API array sort(const array& a, StreamOrDevice s = {});
+
+/** Returns a sorted copy of the array along a given axis. */
+MLX_API array sort(const array& a, int axis, StreamOrDevice s = {});
+
+/** Returns indices that sort the flattened array. */
+MLX_API array argsort(const array& a, StreamOrDevice s = {});
+
+/** Returns indices that sort the array along a given axis. */
+MLX_API array argsort(const array& a, int axis, StreamOrDevice s = {});
+
+/**
+ * Returns a partitioned copy of the flattened array
+ * such that the smaller kth elements are first.
+ **/
+MLX_API array partition(const array& a, int kth, StreamOrDevice s = {});
+
+/**
+ * Returns a partitioned copy of the array along a given axis
+ * such that the smaller kth elements are first.
+ **/
+MLX_API array
+partition(const array& a, int kth, int axis, StreamOrDevice s = {});
+
+/**
+ * Returns indices that partition the flattened array
+ * such that the smaller kth elements are first.
+ **/
+MLX_API array argpartition(const array& a, int kth, StreamOrDevice s = {});
+
+/**
+ * Returns indices that partition the array along a given axis
+ * such that the smaller kth elements are first.
+ **/
+MLX_API array
+argpartition(const array& a, int kth, int axis, StreamOrDevice s = {});
+
+/** Returns topk elements of the flattened array. */
+MLX_API array topk(const array& a, int k, StreamOrDevice s = {});
+
+/** Returns topk elements of the array along a given axis. */
+MLX_API array topk(const array& a, int k, int axis, StreamOrDevice s = {});
+
+/** Cumulative logsumexp of an array. */
+MLX_API array logcumsumexp(
+    const array& a,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative logsumexp of an array along the given axis. */
+MLX_API array logcumsumexp(
+    const array& a,
+    int axis,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** The logsumexp of all elements of the array. */
+MLX_API array logsumexp(const array& a, bool keepdims, StreamOrDevice s = {});
+inline array logsumexp(const array& a, StreamOrDevice s = {}) {
+  return logsumexp(a, false, to_stream(s));
+}
+
+/** The logsumexp of the elements of an array along the given axes. */
+MLX_API array logsumexp(
+    const array& a,
+    const std::vector<int>& axes,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** The logsumexp of the elements of an array along the given axis. */
+MLX_API array logsumexp(
+    const array& a,
+    int axis,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+
+/** Absolute value of elements in an array. */
+MLX_API array abs(const array& a, StreamOrDevice s = {});
+
+/** Negate an array. */
+MLX_API array negative(const array& a, StreamOrDevice s = {});
+MLX_API array operator-(const array& a);
+
+/** The sign of the elements in an array. */
+MLX_API array sign(const array& a, StreamOrDevice s = {});
+
+/** Logical not of an array */
+MLX_API array logical_not(const array& a, StreamOrDevice s = {});
+
+/** Logical and of two arrays */
+MLX_API array
+logical_and(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator&&(const array& a, const array& b);
+
+/** Logical or of two arrays */
+MLX_API array logical_or(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator||(const array& a, const array& b);
+
+/** The reciprocal (1/x) of the elements in an array. */
+MLX_API array reciprocal(const array& a, StreamOrDevice s = {});
+
+/** Add two arrays. */
+MLX_API array add(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator+(const array& a, const array& b);
+template <typename T>
+array operator+(T a, const array& b) {
+  return add(array(a), b);
+}
+template <typename T>
+array operator+(const array& a, T b) {
+  return add(a, array(b));
+}
+
+/** Subtract two arrays. */
+MLX_API array subtract(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator-(const array& a, const array& b);
+template <typename T>
+array operator-(T a, const array& b) {
+  return subtract(array(a), b);
+}
+template <typename T>
+array operator-(const array& a, T b) {
+  return subtract(a, array(b));
+}
+
+/** Multiply two arrays. */
+MLX_API array multiply(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator*(const array& a, const array& b);
+template <typename T>
+array operator*(T a, const array& b) {
+  return multiply(array(a), b);
+}
+template <typename T>
+array operator*(const array& a, T b) {
+  return multiply(a, array(b));
+}
+
+/** Divide two arrays. */
+MLX_API array divide(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator/(const array& a, const array& b);
+MLX_API array operator/(double a, const array& b);
+MLX_API array operator/(const array& a, double b);
+
+/** Compute the element-wise quotient and remainder. */
+MLX_API std::vector<array>
+divmod(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Compute integer division. Equivalent to doing floor(a / x). */
+MLX_API array
+floor_divide(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Compute the element-wise remainder of division */
+MLX_API array remainder(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator%(const array& a, const array& b);
+template <typename T>
+array operator%(T a, const array& b) {
+  return remainder(array(a), b);
+}
+template <typename T>
+array operator%(const array& a, T b) {
+  return remainder(a, array(b));
+}
+
+/** Element-wise maximum between two arrays. */
+MLX_API array maximum(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Element-wise minimum between two arrays. */
+MLX_API array minimum(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Floor the element of an array. **/
+MLX_API array floor(const array& a, StreamOrDevice s = {});
+
+/** Ceil the element of an array. **/
+MLX_API array ceil(const array& a, StreamOrDevice s = {});
+
+/** Square the elements of an array. */
+MLX_API array square(const array& a, StreamOrDevice s = {});
+
+/** Exponential of the elements of an array. */
+MLX_API array exp(const array& a, StreamOrDevice s = {});
+
+/** Sine of the elements of an array */
+MLX_API array sin(const array& a, StreamOrDevice s = {});
+
+/** Cosine of the elements of an array */
+MLX_API array cos(const array& a, StreamOrDevice s = {});
+
+/** Tangent of the elements of an array */
+MLX_API array tan(const array& a, StreamOrDevice s = {});
+
+/** Arc Sine of the elements of an array */
+MLX_API array arcsin(const array& a, StreamOrDevice s = {});
+
+/** Arc Cosine of the elements of an array */
+MLX_API array arccos(const array& a, StreamOrDevice s = {});
+
+/** Arc Tangent of the elements of an array */
+MLX_API array arctan(const array& a, StreamOrDevice s = {});
+
+/** Inverse tangent of the ratio of two arrays */
+MLX_API array arctan2(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Hyperbolic Sine of the elements of an array */
+MLX_API array sinh(const array& a, StreamOrDevice s = {});
+
+/** Hyperbolic Cosine of the elements of an array */
+MLX_API array cosh(const array& a, StreamOrDevice s = {});
+
+/** Hyperbolic Tangent of the elements of an array */
+MLX_API array tanh(const array& a, StreamOrDevice s = {});
+
+/** Inverse Hyperbolic Sine of the elements of an array */
+MLX_API array arcsinh(const array& a, StreamOrDevice s = {});
+
+/** Inverse Hyperbolic Cosine of the elements of an array */
+MLX_API array arccosh(const array& a, StreamOrDevice s = {});
+
+/** Inverse Hyperbolic Tangent of the elements of an array */
+MLX_API array arctanh(const array& a, StreamOrDevice s = {});
+
+/** Convert the elements of an array from Radians to Degrees **/
+MLX_API array degrees(const array& a, StreamOrDevice s = {});
+
+/** Convert the elements of an array from Degrees to Radians **/
+MLX_API array radians(const array& a, StreamOrDevice s = {});
+
+/** Natural logarithm of the elements of an array. */
+MLX_API array log(const array& a, StreamOrDevice s = {});
+
+/** Log base 2 of the elements of an array. */
+MLX_API array log2(const array& a, StreamOrDevice s = {});
+
+/** Log base 10 of the elements of an array. */
+MLX_API array log10(const array& a, StreamOrDevice s = {});
+
+/** Natural logarithm of one plus elements in the array: `log(1 + a)`. */
+MLX_API array log1p(const array& a, StreamOrDevice s = {});
+
+/** Log-add-exp of one elements in the array: `log(exp(a) + exp(b))`. */
+MLX_API array logaddexp(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Element-wise logistic sigmoid of the array: `1 / (1 + exp(-x)`. */
+MLX_API array sigmoid(const array& a, StreamOrDevice s = {});
+
+/** Computes the error function of the elements of an array. */
+MLX_API array erf(const array& a, StreamOrDevice s = {});
+
+/** Computes the inverse error function of the elements of an array. */
+MLX_API array erfinv(const array& a, StreamOrDevice s = {});
+
+/** Computes the expm1 function of the elements of an array. */
+MLX_API array expm1(const array& a, StreamOrDevice s = {});
+
+/** Stop the flow of gradients. */
+MLX_API array stop_gradient(const array& a, StreamOrDevice s = {});
+
+/** Round a floating point number */
+MLX_API array round(const array& a, int decimals, StreamOrDevice s = {});
+inline array round(const array& a, StreamOrDevice s = {}) {
+  return round(a, 0, s);
+}
+
+/** Matrix-matrix multiplication. */
+MLX_API array matmul(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Gather array entries given indices and slices */
+MLX_API array gather(
+    const array& a,
+    const std::vector<array>& indices,
+    const std::vector<int>& axes,
+    const Shape& slice_sizes,
+    StreamOrDevice s = {});
+inline array gather(
+    const array& a,
+    const array& indices,
+    int axis,
+    const Shape& slice_sizes,
+    StreamOrDevice s = {}) {
+  return gather(a, {indices}, std::vector<int>{axis}, slice_sizes, s);
+}
+
+/**  Compute the Kronecker product of two arrays. */
+MLX_API array kron(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Take array slices at the given indices of the specified axis. */
+MLX_API array
+take(const array& a, const array& indices, int axis, StreamOrDevice s = {});
+MLX_API array take(const array& a, int index, int axis, StreamOrDevice s = {});
+
+/** Take array entries at the given indices treating the array as flattened. */
+MLX_API array take(const array& a, const array& indices, StreamOrDevice s = {});
+MLX_API array take(const array& a, int index, StreamOrDevice s = {});
+
+/** Take array entries given indices along the axis */
+MLX_API array take_along_axis(
+    const array& a,
+    const array& indices,
+    int axis,
+    StreamOrDevice s = {});
+
+/** Put the values into the array at the given indices along the axis */
+MLX_API array put_along_axis(
+    const array& a,
+    const array& indices,
+    const array& values,
+    int axis,
+    StreamOrDevice s = {});
+
+/** Add the values into the array at the given indices along the axis */
+MLX_API array scatter_add_axis(
+    const array& a,
+    const array& indices,
+    const array& values,
+    int axis,
+    StreamOrDevice s = {});
+
+/** Scatter updates to the given indices.
+ *
+ * The parameters ``indices`` and ``axes`` determine the locations of ``a``
+ * that are updated with the values in ``updates``. Assuming 1-d ``indices``
+ * for simplicity, ``indices[i]`` are the indices on axis ``axes[i]`` to which
+ * the values in ``updates`` will be applied. Note each array in
+ * ``indices`` is assigned to a corresponding axis and hence ``indices.size() ==
+ * axes.size()``. If an index/axis pair is not provided then indices along that
+ * axis are assumed to be zero.
+ *
+ * Note the rank of ``updates`` must be equal to the sum of the rank of the
+ * broadcasted ``indices`` and the rank of ``a``. In other words, assuming the
+ * arrays in ``indices`` have the same shape, ``updates.ndim() ==
+ * indices[0].ndim() + a.ndim()``. The leading dimensions of ``updates``
+ * correspond to the indices, and the remaining ``a.ndim()`` dimensions are the
+ * values that will be applied to the given location in ``a``.
+ *
+ * For example:
+ *
+ * @code
+ * auto in = zeros({4, 4}, float32);
+ * auto indices = array({2});
+ * auto updates = reshape(arange(1, 3, float32), {1, 1, 2});
+ * std::vector<int> axes{0};
+ *
+ * auto out = scatter(in, {indices}, updates, axes);
+ * @endcode
+ *
+ * will produce:
+ *
+ * @code
+ * array([[0, 0, 0, 0],
+ *        [0, 0, 0, 0],
+ *        [1, 2, 0, 0],
+ *        [0, 0, 0, 0]], dtype=float32)
+ * @endcode
+ *
+ * This scatters the two-element row vector ``[1, 2]`` starting at the ``(2,
+ * 0)`` position of ``a``.
+ *
+ * Adding another element to ``indices`` will scatter into another location of
+ * ``a``. We also have to add an another update for the new index:
+ *
+ * @code
+ * auto in = zeros({4, 4}, float32);
+ * auto indices = array({2, 0});
+ * auto updates = reshape(arange(1, 5, float32), {2, 1, 2});
+ * std::vector<int> axes{0};
+ *
+ * auto out = scatter(in, {indices}, updates, axes):
+ * @endcode
+ *
+ * will produce:
+ *
+ * @code
+ * array([[3, 4, 0, 0],
+ *        [0, 0, 0, 0],
+ *        [1, 2, 0, 0],
+ *        [0, 0, 0, 0]], dtype=float32)
+ * @endcode
+ *
+ * To control the scatter location on an additional axis, add another index
+ * array to ``indices`` and another axis to ``axes``:
+ *
+ * @code
+ * auto in = zeros({4, 4}, float32);
+ * auto indices = std::vector{array({2, 0}), array({1, 2})};
+ * auto updates = reshape(arange(1, 5, float32), {2, 1, 2});
+ * std::vector<int> axes{0, 1};
+ *
+ * auto out = scatter(in, indices, updates, axes);
+ * @endcode
+ *
+ * will produce:
+ *
+ * @code
+ * array([[0, 0, 3, 4],
+ *       [0, 0, 0, 0],
+ *       [0, 1, 2, 0],
+ *       [0, 0, 0, 0]], dtype=float32)
+ * @endcode
+ *
+ * Items in indices are broadcasted together. This means:
+ *
+ * @code
+ * auto indices = std::vector{array({2, 0}), array({1})};
+ * @endcode
+ *
+ * is equivalent to:
+ *
+ * @code
+ * auto indices = std::vector{array({2, 0}), array({1, 1})};
+ * @endcode
+ *
+ * Note, ``scatter`` does not perform bounds checking on the indices and
+ * updates.  Out-of-bounds accesses on ``a`` are undefined and typically result
+ * in unintended or invalid memory writes.
+ */
+MLX_API array scatter(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+inline array scatter(
+    const array& a,
+    const array& indices,
+    const array& updates,
+    int axis,
+    StreamOrDevice s = {}) {
+  return scatter(a, {indices}, updates, std::vector<int>{axis}, s);
+}
+
+/** Scatter and add updates to given indices */
+MLX_API array scatter_add(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+inline array scatter_add(
+    const array& a,
+    const array& indices,
+    const array& updates,
+    int axis,
+    StreamOrDevice s = {}) {
+  return scatter_add(a, {indices}, updates, std::vector<int>{axis}, s);
+}
+
+/** Scatter and prod updates to given indices */
+MLX_API array scatter_prod(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+inline array scatter_prod(
+    const array& a,
+    const array& indices,
+    const array& updates,
+    int axis,
+    StreamOrDevice s = {}) {
+  return scatter_prod(a, {indices}, updates, std::vector<int>{axis}, s);
+}
+
+/** Scatter and max updates to given linear indices */
+MLX_API array scatter_max(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+inline array scatter_max(
+    const array& a,
+    const array& indices,
+    const array& updates,
+    int axis,
+    StreamOrDevice s = {}) {
+  return scatter_max(a, {indices}, updates, std::vector<int>{axis}, s);
+}
+/** Scatter and min updates to given linear indices */
+MLX_API array scatter_min(
+    const array& a,
+    const std::vector<array>& indices,
+    const array& updates,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+inline array scatter_min(
+    const array& a,
+    const array& indices,
+    const array& updates,
+    int axis,
+    StreamOrDevice s = {}) {
+  return scatter_min(a, {indices}, updates, std::vector<int>{axis}, s);
+}
+
+MLX_API array masked_scatter(
+    const array& a,
+    const array& mask,
+    const array& src,
+    StreamOrDevice s = {});
+
+/** Square root the elements of an array. */
+MLX_API array sqrt(const array& a, StreamOrDevice s = {});
+
+/** Square root and reciprocal the elements of an array. */
+MLX_API array rsqrt(const array& a, StreamOrDevice s = {});
+
+/** Softmax of an array. */
+MLX_API array softmax(
+    const array& a,
+    const std::vector<int>& axes,
+    bool precise = false,
+    StreamOrDevice s = {});
+
+/** Softmax of an array. */
+MLX_API array
+softmax(const array& a, bool precise = false, StreamOrDevice s = {});
+
+/** Softmax of an array. */
+inline array
+softmax(const array& a, int axis, bool precise = false, StreamOrDevice s = {}) {
+  return softmax(a, std::vector<int>{axis}, precise, s);
+}
+
+/** Raise elements of a to the power of b element-wise */
+MLX_API array power(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Cumulative sum of an array. */
+MLX_API array cumsum(
+    const array& a,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative sum of an array along the given axis. */
+MLX_API array cumsum(
+    const array& a,
+    int axis,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative product of an array. */
+MLX_API array cumprod(
+    const array& a,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative product of an array along the given axis. */
+MLX_API array cumprod(
+    const array& a,
+    int axis,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative max of an array. */
+MLX_API array cummax(
+    const array& a,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative max of an array along the given axis. */
+MLX_API array cummax(
+    const array& a,
+    int axis,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative min of an array. */
+MLX_API array cummin(
+    const array& a,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** Cumulative min of an array along the given axis. */
+MLX_API array cummin(
+    const array& a,
+    int axis,
+    bool reverse = false,
+    bool inclusive = true,
+    StreamOrDevice s = {});
+
+/** General convolution with a filter */
+MLX_API array conv_general(
+    array input,
+    array weight,
+    std::vector<int> stride = {},
+    std::vector<int> padding_lo = {},
+    std::vector<int> padding_hi = {},
+    std::vector<int> kernel_dilation = {},
+    std::vector<int> input_dilation = {},
+    int groups = 1,
+    bool flip = false,
+    StreamOrDevice s = {});
+
+/** General convolution with a filter */
+inline array conv_general(
+    const array& input,
+    const array& weight,
+    std::vector<int> stride = {},
+    std::vector<int> padding = {},
+    std::vector<int> kernel_dilation = {},
+    std::vector<int> input_dilation = {},
+    int groups = 1,
+    bool flip = false,
+    StreamOrDevice s = {}) {
+  return conv_general(
+      /* const array& input = */ input,
+      /* const array& weight = */ weight,
+      /* std::vector<int> stride = */ stride,
+      /* std::vector<int> padding_lo = */ padding,
+      /* std::vector<int> padding_hi = */ padding,
+      /* std::vector<int> kernel_dilation = */ kernel_dilation,
+      /* std::vector<int> input_dilation = */ input_dilation,
+      /* int groups = */ groups,
+      /* bool flip = */ flip,
+      /* StreamOrDevice s = */ s);
+}
+
+/** 1D convolution with a filter */
+MLX_API array conv1d(
+    const array& input,
+    const array& weight,
+    int stride = 1,
+    int padding = 0,
+    int dilation = 1,
+    int groups = 1,
+    StreamOrDevice s = {});
+
+/** 2D convolution with a filter */
+MLX_API array conv2d(
+    const array& input,
+    const array& weight,
+    const std::pair<int, int>& stride = {1, 1},
+    const std::pair<int, int>& padding = {0, 0},
+    const std::pair<int, int>& dilation = {1, 1},
+    int groups = 1,
+    StreamOrDevice s = {});
+
+/** 3D convolution with a filter */
+MLX_API array conv3d(
+    const array& input,
+    const array& weight,
+    const std::tuple<int, int, int>& stride = {1, 1, 1},
+    const std::tuple<int, int, int>& padding = {0, 0, 0},
+    const std::tuple<int, int, int>& dilation = {1, 1, 1},
+    int groups = 1,
+    StreamOrDevice s = {});
+
+/** 1D transposed convolution with a filter */
+MLX_API array conv_transpose1d(
+    const array& input,
+    const array& weight,
+    int stride = 1,
+    int padding = 0,
+    int dilation = 1,
+    int output_padding = 0,
+    int groups = 1,
+    StreamOrDevice s = {});
+
+/** 2D transposed convolution with a filter */
+MLX_API array conv_transpose2d(
+    const array& input,
+    const array& weight,
+    const std::pair<int, int>& stride = {1, 1},
+    const std::pair<int, int>& padding = {0, 0},
+    const std::pair<int, int>& dilation = {1, 1},
+    const std::pair<int, int>& output_padding = {0, 0},
+    int groups = 1,
+    StreamOrDevice s = {});
+
+/** 3D transposed convolution with a filter */
+MLX_API array conv_transpose3d(
+    const array& input,
+    const array& weight,
+    const std::tuple<int, int, int>& stride = {1, 1, 1},
+    const std::tuple<int, int, int>& padding = {0, 0, 0},
+    const std::tuple<int, int, int>& dilation = {1, 1, 1},
+    const std::tuple<int, int, int>& output_padding = {0, 0, 0},
+    int groups = 1,
+    StreamOrDevice s = {});
+
+/** Quantized matmul multiplies x with a quantized matrix w*/
+MLX_API array quantized_matmul(
+    array x,
+    array w,
+    array scales,
+    std::optional<array> biases = std::nullopt,
+    bool transpose = true,
+    std::optional<int> group_size = std::nullopt,
+    std::optional<int> bits = std::nullopt,
+    const std::string& mode = "affine",
+    StreamOrDevice s = {});
+
+/** Quantize a matrix along its last axis */
+MLX_API std::vector<array> quantize(
+    const array& w,
+    std::optional<int> group_size = std::nullopt,
+    std::optional<int> bits = std::nullopt,
+    const std::string& mode = "affine",
+    StreamOrDevice s = {});
+
+/** Dequantize a matrix produced by quantize() */
+MLX_API array dequantize(
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases = std::nullopt,
+    std::optional<int> group_size = std::nullopt,
+    std::optional<int> bits = std::nullopt,
+    const std::string& mode = "affine",
+    std::optional<Dtype> dtype = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array qqmm(
+    array x, // input activations
+    array w, // maybe quantized weights
+    std::optional<array> w_scales = std::nullopt, // optional scales if w is
+                                                  // quantized
+    std::optional<int> group_size = std::nullopt,
+    std::optional<int> bits = std::nullopt,
+    const std::string& mode = "nvfp4",
+    StreamOrDevice s = {});
+
+/** Convert an E4M3 float8 to the given floating point dtype. */
+MLX_API array from_fp8(array x, Dtype dtype, StreamOrDevice s = {});
+
+/** Convert a floating point matrix to E4M3 float8. */
+MLX_API array to_fp8(array x, StreamOrDevice s = {});
+
+/** Compute matrix products with matrix-level gather. */
+MLX_API array gather_qmm(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases = std::nullopt,
+    std::optional<array> lhs_indices = std::nullopt,
+    std::optional<array> rhs_indices = std::nullopt,
+    bool transpose = true,
+    std::optional<int> group_size = std::nullopt,
+    std::optional<int> bits = std::nullopt,
+    const std::string& mode = "affine",
+    bool sorted_indices = false,
+    StreamOrDevice s = {});
+
+/** Returns a contraction of a and b over multiple dimensions. */
+MLX_API array tensordot(
+    const array& a,
+    const array& b,
+    const int axis = 2,
+    StreamOrDevice s = {});
+
+MLX_API array tensordot(
+    const array& a,
+    const array& b,
+    const std::vector<int>& axes_a,
+    const std::vector<int>& axes_b,
+    StreamOrDevice s = {});
+
+/** Compute the outer product of two vectors. */
+MLX_API array outer(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Compute the inner product of two vectors. */
+MLX_API array inner(const array& a, const array& b, StreamOrDevice s = {});
+
+/** Compute D = beta * C + alpha * (A @ B) */
+MLX_API array addmm(
+    array c,
+    array a,
+    array b,
+    const float& alpha = 1.f,
+    const float& beta = 1.f,
+    StreamOrDevice s = {});
+
+/** Compute matrix product with block masking */
+MLX_API array block_masked_mm(
+    array a,
+    array b,
+    int block_size,
+    std::optional<array> mask_out = std::nullopt,
+    std::optional<array> mask_lhs = std::nullopt,
+    std::optional<array> mask_rhs = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Compute matrix product with matrix-level gather */
+MLX_API array gather_mm(
+    array a,
+    array b,
+    std::optional<array> lhs_indices = std::nullopt,
+    std::optional<array> rhs_indices = std::nullopt,
+    bool sorted_indices = false,
+    StreamOrDevice s = {});
+
+/**
+ * Compute a matrix product but segment the inner dimension and write the
+ * result separately for each segment.
+ */
+MLX_API array
+segmented_mm(array a, array b, array segments, StreamOrDevice s = {});
+
+/** Extract a diagonal or construct a diagonal array */
+MLX_API array diagonal(
+    const array& a,
+    int offset = 0,
+    int axis1 = 0,
+    int axis2 = 1,
+    StreamOrDevice s = {});
+
+/** Extract diagonal from a 2d array or create a diagonal matrix. */
+MLX_API array diag(const array& a, int k = 0, StreamOrDevice s = {});
+
+/** Return the sum along a specified diagonal in the given array. */
+MLX_API array trace(
+    const array& a,
+    int offset,
+    int axis1,
+    int axis2,
+    Dtype dtype,
+    StreamOrDevice s = {});
+MLX_API array
+trace(const array& a, int offset, int axis1, int axis2, StreamOrDevice s = {});
+MLX_API array trace(const array& a, StreamOrDevice s = {});
+
+/**
+ * Implements the identity function but allows injecting dependencies to other
+ * arrays. This ensures that these other arrays will have been computed
+ * when the outputs of this function are computed.
+ */
+MLX_API std::vector<array> depends(
+    const std::vector<array>& inputs,
+    const std::vector<array>& dependencies);
+
+/** convert an array to an atleast ndim array */
+MLX_API array atleast_1d(const array& a, StreamOrDevice s = {});
+MLX_API std::vector<array> atleast_1d(
+    const std::vector<array>& a,
+    StreamOrDevice s = {});
+MLX_API array atleast_2d(const array& a, StreamOrDevice s = {});
+MLX_API std::vector<array> atleast_2d(
+    const std::vector<array>& a,
+    StreamOrDevice s = {});
+MLX_API array atleast_3d(const array& a, StreamOrDevice s = {});
+MLX_API std::vector<array> atleast_3d(
+    const std::vector<array>& a,
+    StreamOrDevice s = {});
+
+/**
+ * Extract the number of elements along some axes as a scalar array. Used to
+ * allow shape dependent shapeless compilation (pun intended).
+ */
+MLX_API array number_of_elements(
+    const array& a,
+    std::vector<int> axes,
+    bool inverted,
+    Dtype dtype = int32,
+    StreamOrDevice s = {});
+
+MLX_API array conjugate(const array& a, StreamOrDevice s = {});
+
+/** Bitwise and. */
+MLX_API array
+bitwise_and(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator&(const array& a, const array& b);
+
+/** Bitwise inclusive or. */
+MLX_API array bitwise_or(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator|(const array& a, const array& b);
+
+/** Bitwise exclusive or. */
+MLX_API array
+bitwise_xor(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator^(const array& a, const array& b);
+
+/** Shift bits to the left. */
+MLX_API array left_shift(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator<<(const array& a, const array& b);
+
+/** Shift bits to the right. */
+MLX_API array
+right_shift(const array& a, const array& b, StreamOrDevice s = {});
+MLX_API array operator>>(const array& a, const array& b);
+
+/** Invert the bits. */
+MLX_API array bitwise_invert(const array& a, StreamOrDevice s = {});
+MLX_API array operator~(const array& a);
+
+MLX_API array view(const array& a, const Dtype& dtype, StreamOrDevice s = {});
+
+/** Roll elements along an axis and introduce them on the other side */
+MLX_API array roll(const array& a, int shift, StreamOrDevice s = {});
+MLX_API array roll(const array& a, const Shape& shift, StreamOrDevice s = {});
+MLX_API array roll(const array& a, int shift, int axis, StreamOrDevice s = {});
+MLX_API array roll(
+    const array& a,
+    int shift,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+roll(const array& a, const Shape& shift, int axis, StreamOrDevice s = {});
+MLX_API array roll(
+    const array& a,
+    const Shape& shift,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+
+/* The real part of a complex array. */
+MLX_API array real(const array& a, StreamOrDevice s = {});
+
+/* The imaginary part of a complex array. */
+MLX_API array imag(const array& a, StreamOrDevice s = {});
+
+/* Ensure the array's underlying memory is contiguous. */
+MLX_API array
+contiguous(const array& a, bool allow_col_major = false, StreamOrDevice s = {});
+
+/** @} */
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/primitives.cpp b/macos/third_party/mlx/mlx/primitives.cpp
new file mode 100644
index 0000000..b078625
--- /dev/null
+++ b/macos/third_party/mlx/mlx/primitives.cpp
@@ -0,0 +1,5850 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// Required for using M_2_SQRTPI in MSVC.
+#define _USE_MATH_DEFINES
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <numeric>
+#include <sstream>
+#include <stdexcept>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/fft.h"
+#include "mlx/linalg.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+std::tuple<array, array, int> vmap_binary_op(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes,
+    const Stream& stream) {
+  assert(inputs.size() == 2);
+  assert(axes.size() == 2);
+
+  if (axes[0] == -1 && axes[1] == -1) {
+    return {inputs[0], inputs[1], -1};
+  }
+
+  auto a = inputs[0];
+  auto b = inputs[1];
+  int ndim = std::max(a.ndim() + (axes[0] == -1), b.ndim() + (axes[1] == -1));
+
+  auto expand_dims = [stream, ndim](auto in) {
+    auto shape = in.shape();
+    shape.insert(shape.begin(), ndim - shape.size(), 1);
+    return reshape(in, shape, stream);
+  };
+
+  int to_ax = (ndim - a.ndim()) + axes[0];
+  int from_ax = (ndim - b.ndim()) + axes[1];
+  a = expand_dims(a);
+  b = expand_dims(b);
+
+  if (from_ax != to_ax) {
+    std::vector<int> tdims(b.ndim());
+    std::iota(tdims.begin(), tdims.end(), 0);
+    tdims.erase(tdims.begin() + from_ax);
+    tdims.insert(tdims.begin() + to_ax, from_ax);
+    b = transpose(b, tdims, stream);
+  }
+  return {a, b, to_ax};
+}
+
+std::tuple<array, array, array, int> vmap_ternary_op(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes,
+    const Stream& stream) {
+  assert(inputs.size() == 3);
+  assert(axes.size() == 3);
+
+  if (axes[0] == -1 && axes[1] == -1 && axes[2] == -1) {
+    return {inputs[0], inputs[1], inputs[2], -1};
+  }
+
+  auto a = inputs[0];
+  auto b = inputs[1];
+  auto c = inputs[2];
+  int ndim = std::max(
+      {a.ndim() + (axes[0] == -1),
+       b.ndim() + (axes[1] == -1),
+       c.ndim() + (axes[2] == -1)});
+
+  auto expand_dims = [stream, ndim](auto in) {
+    auto shape = in.shape();
+    shape.insert(shape.begin(), ndim - shape.size(), 1);
+    return reshape(in, shape, stream);
+  };
+
+  int to_ax = (ndim - a.ndim()) + axes[0];
+  int from_ax1 = (ndim - b.ndim()) + axes[1];
+  int from_ax2 = (ndim - c.ndim()) + axes[2];
+  a = expand_dims(a);
+  b = expand_dims(b);
+  c = expand_dims(c);
+
+  auto find_tdims = [](auto x, int to_ax, int from_ax) {
+    std::vector<int> tdims(x.ndim());
+    std::iota(tdims.begin(), tdims.end(), 0);
+    tdims.erase(tdims.begin() + from_ax);
+    tdims.insert(tdims.begin() + to_ax, from_ax);
+    return tdims;
+  };
+
+  if (to_ax != from_ax1) {
+    std::vector<int> tdims = find_tdims(b, to_ax, from_ax1);
+    b = transpose(b, tdims, stream);
+  }
+
+  if (to_ax != from_ax2) {
+    std::vector<int> tdims = find_tdims(c, to_ax, from_ax2);
+    c = transpose(c, tdims, stream);
+  }
+  return {a, b, c, to_ax};
+}
+
+// Calculate the gradient wrt to the weights of the following calculation
+//
+// y = gather_mm(x, w.T, lhs_indices, rhs_indices, sorted)
+//
+// Note the transpose above. This function returns the gradient for w.T so if w
+// was used instead then one needs to transpose the returned gradient.
+//
+// We define it as a separate function to reuse it for gather_mm and
+// gather_qmm.
+array gather_mm_grad(
+    const array& x,
+    const array& dy,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    bool sorted,
+    Shape batch_shape,
+    const Stream& s) {
+  int M = x.shape(-2);
+  int K = x.shape(-1);
+  int N = dy.shape(-1);
+  int num_segments = std::accumulate(
+      batch_shape.begin(), batch_shape.end(), 1, std::multiplies<int>());
+  batch_shape.push_back(N);
+  batch_shape.push_back(K);
+
+  // If the indices are sorted then it means that we can do the whole gradient
+  // computation via a segmented matmul. We just need to calculate the segments
+  // using the indices.
+  if (sorted) {
+    auto segments = zeros({num_segments}, uint32, s);
+    segments = scatter_add_axis(segments, rhs_indices, array(M, uint32), 0, s);
+    segments = cumsum(segments, 0, false, true, s);
+    segments = concatenate({array({0}, {1}, uint32), segments}, 0, s);
+    segments = as_strided(segments, {num_segments, 2}, {1, 1}, 0, s);
+
+    return reshape(
+        segmented_mm(
+            swapaxes(flatten(dy, 0, -2, s), 0, 1, s),
+            flatten(x, 0, -2, s),
+            segments,
+            s),
+        std::move(batch_shape),
+        s);
+  }
+
+  // Otherwise we need to gather matmul the dy and then scatter add it to the
+  // correct locations.
+  else {
+    // TODO: If the lhs indices wasn't provided, this is always a sorted matmul
+    //       so we should add that check.
+    auto dw = gather_mm(
+        swapaxes(dy, -1, -2, s), x, std::nullopt, lhs_indices, false, s);
+    return reshape(
+        scatter_add(
+            zeros({num_segments, N, K}, dw.dtype(), s),
+            rhs_indices,
+            expand_dims(dw, -3, s),
+            0,
+            s),
+        std::move(batch_shape),
+        s);
+  }
+}
+
+} // namespace
+
+std::vector<array> Primitive::jvp(
+    const std::vector<array>&,
+    const std::vector<array>&,
+    const std::vector<int>&) {
+  std::ostringstream msg;
+  msg << "[Primitive::jvp] Not implemented for ";
+  msg << name();
+  msg << ".";
+  throw std::invalid_argument(msg.str());
+}
+
+std::vector<array> Primitive::vjp(
+    const std::vector<array>&,
+    const std::vector<array>&,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  std::ostringstream msg;
+  msg << "[Primitive::vjp] Not implemented for ";
+  msg << name();
+  msg << ".";
+  throw std::invalid_argument(msg.str());
+}
+
+std::pair<std::vector<array>, std::vector<int>> Primitive::vmap(
+    const std::vector<array>&,
+    const std::vector<int>&) {
+  std::ostringstream msg;
+  msg << "[Primitive::vmap] Not implemented for ";
+  msg << name();
+  msg << ".";
+  throw std::invalid_argument(msg.str());
+}
+
+std::vector<Shape> Primitive::output_shapes(const std::vector<array>&) {
+  std::ostringstream msg;
+  msg << "[Primitive::output_shapes] ";
+  msg << name();
+  msg << " cannot infer output shapes.";
+  throw std::invalid_argument(msg.str());
+}
+
+std::vector<array> Abs::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Abs::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(tangents[0], sign(primals[0], stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Abs::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{abs(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Add::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  return {
+      tangents.size() > 1 ? add(tangents[0], tangents[1], stream())
+                          : tangents[0]};
+}
+
+std::vector<array> Add::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  if (argnums.size() == 1) {
+    return cotangents;
+  } else {
+    return {cotangents[0], cotangents[0]};
+  }
+}
+
+std::pair<std::vector<array>, std::vector<int>> Add::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{add(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> AddMM::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  auto& cotan = cotangents[0];
+  std::vector<int> reorder(cotan.ndim());
+  std::iota(reorder.begin(), reorder.end(), 0);
+  std::iter_swap(reorder.end() - 1, reorder.end() - 2);
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      // M X N * (K X N).T -> M X K
+      auto cotan_scaled = cotan;
+      if (alpha_ != 1.) {
+        auto alpha_arr = array(alpha_, cotan.dtype());
+        cotan_scaled = (multiply(alpha_arr, cotan_scaled, stream()));
+      }
+      vjps.push_back(matmul(
+          cotan_scaled, transpose(primals[1], reorder, stream()), stream()));
+    } else if (arg == 1) {
+      // (M X K).T * M X N -> K X N
+      auto cotan_scaled = cotan;
+      if (alpha_ != 1.) {
+        auto alpha_arr = array(alpha_, cotan.dtype());
+        cotan_scaled = (multiply(alpha_arr, cotan_scaled, stream()));
+      }
+      vjps.push_back(matmul(
+          transpose(primals[0], reorder, stream()), cotan_scaled, stream()));
+    } else {
+      auto cotan_scaled = cotan;
+      if (beta_ != 1.) {
+        auto beta_arr = array(beta_, cotan.dtype());
+        cotan_scaled = (multiply(beta_arr, cotan_scaled, stream()));
+      }
+      vjps.push_back(cotan_scaled);
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> AddMM::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  std::vector<array> jvp;
+  for (int i = 0; i < argnums.size(); ++i) {
+    auto arg = argnums[i];
+    if (arg == 0) {
+      if (jvp.empty()) {
+        jvp.push_back(matmul(tangents[i], primals[1], stream()));
+      } else {
+        jvp[0] = addmm(jvp[0], tangents[i], primals[1], 1.0f, 1.0f, stream());
+      }
+    } else if (arg == 1) {
+      if (jvp.empty()) {
+        jvp.push_back(matmul(primals[0], tangents[i], stream()));
+      } else {
+        jvp[0] = addmm(jvp[0], primals[0], tangents[i], 1.0f, 1.0f, stream());
+      }
+    } else {
+      if (jvp.empty()) {
+        jvp.push_back(tangents[i]);
+      } else {
+        jvp[0] = add(jvp[0], tangents[i], stream());
+      }
+    }
+  }
+  return jvp;
+}
+
+bool AddMM::is_equivalent(const Primitive& other) const {
+  const AddMM& a_other = static_cast<const AddMM&>(other);
+  return (alpha_ == a_other.alpha_ && beta_ == a_other.beta_);
+}
+
+std::pair<std::vector<array>, std::vector<int>> AddMM::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto maybe_move_ax = [this](auto& arr, auto ax) {
+    return ax > 0 ? moveaxis(arr, ax, 0, stream()) : arr;
+  };
+  auto a = maybe_move_ax(inputs[0], axes[0]);
+  auto b = maybe_move_ax(inputs[1], axes[1]);
+  auto c = maybe_move_ax(inputs[2], axes[2]);
+  return {{addmm(c, a, b, alpha_, beta_, stream())}, {0}};
+}
+
+bool Arange::is_equivalent(const Primitive& other) const {
+  const Arange& a_other = static_cast<const Arange&>(other);
+  return (
+      start_ == a_other.start_ && stop_ == a_other.stop_ &&
+      step_ == a_other.step_);
+}
+
+std::vector<Shape> Arange::output_shapes(const std::vector<array>&) {
+  auto real_size = std::ceil((stop_ - start_) / step_);
+  return {{std::max(static_cast<int>(real_size), 0)}};
+}
+
+std::vector<array> ArcCos::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> ArcCos::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array one = array(1., primals[0].dtype());
+  array t = subtract(one, square(primals[0], stream()), stream());
+  array denom = negative(rsqrt(t, stream()), stream());
+  return {multiply(tangents[0], denom, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArcCos::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{arccos(inputs[0], stream())}, axes};
+}
+
+std::vector<array> ArcCosh::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> ArcCosh::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array one = array(1., primals[0].dtype());
+  array t = subtract(square(primals[0], stream()), one, stream());
+  return {multiply(tangents[0], rsqrt(t, stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArcCosh::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{arccosh(inputs[0], stream())}, axes};
+}
+
+std::vector<array> ArcSin::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> ArcSin::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array one = array(1., primals[0].dtype());
+  array t = subtract(one, square(primals[0], stream()), stream());
+  return {multiply(tangents[0], rsqrt(t, stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArcSin::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{arcsin(inputs[0], stream())}, axes};
+}
+
+std::vector<array> ArcSinh::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> ArcSinh::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array one = array(1., primals[0].dtype());
+  array t = add(square(primals[0], stream()), one, stream());
+  return {multiply(tangents[0], rsqrt(t, stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArcSinh::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{arcsinh(inputs[0], stream())}, axes};
+}
+
+std::vector<array> ArcTan::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> ArcTan::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array one = array(1., primals[0].dtype());
+  array t = add(one, square(primals[0], stream()), stream());
+  return {divide(tangents[0], t, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArcTan::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{arctan(inputs[0], stream())}, axes};
+}
+
+std::vector<array> ArcTan2::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 2);
+  assert(argnums.size() == 2);
+
+  const auto& s = stream();
+  const array& x1 = primals[0];
+  const array& x2 = primals[1];
+  const array& dy = cotangents[0];
+
+  std::vector<array> grads;
+  array dy_over_x1_x2_squared =
+      divide(dy, add(square(x1, s), square(x2, s)), s);
+
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      grads.emplace_back(multiply(x2, dy_over_x1_x2_squared, s));
+    } else {
+      grads.emplace_back(multiply(negative(x1, s), dy_over_x1_x2_squared, s));
+    }
+  }
+
+  return grads;
+}
+
+std::vector<array> ArcTan2::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 2);
+  assert(argnums.size() == 2);
+
+  const auto& s = stream();
+  const array& x1 = primals[0];
+  const array& x2 = primals[1];
+  const array& dx1 = tangents[0];
+  const array& dx2 = tangents[1];
+
+  return {divide(
+      subtract(multiply(x2, dx1, s), multiply(x1, dx2, s), s),
+      add(square(x1, s), square(x2, s), s),
+      s)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArcTan2::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 2);
+  assert(axes.size() == 2);
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{arctan2(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> ArcTanh::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> ArcTanh::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array one = array(1., primals[0].dtype());
+  array t = subtract(one, square(primals[0], stream()), stream());
+  return {divide(tangents[0], t, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArcTanh::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{arctanh(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArgPartition::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  int axis_left = axes[0] >= 0 && axes[0] <= axis_;
+  return {{argpartition(inputs[0], axis_ + axis_left, stream())}, axes};
+}
+
+std::vector<array> ArgPartition::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>&,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {zeros_like(primals[0], stream())};
+}
+
+std::vector<array> ArgPartition::jvp(
+    const std::vector<array>&,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {zeros_like(tangents[0], stream())};
+}
+
+bool ArgPartition::is_equivalent(const Primitive& other) const {
+  const ArgPartition& r_other = static_cast<const ArgPartition&>(other);
+  return axis_ == r_other.axis_ && kth_ == r_other.kth_;
+}
+
+bool ArgReduce::is_equivalent(const Primitive& other) const {
+  const ArgReduce& r_other = static_cast<const ArgReduce&>(other);
+  return reduce_type_ == r_other.reduce_type_ && axis_ == r_other.axis_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArgReduce::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  int reduce_ax = axis_ + (axes[0] >= 0 && axis_ >= axes[0]);
+  auto& in = inputs[0];
+  std::vector<array> out;
+  if (reduce_type_ == ArgReduce::ArgMin) {
+    out.push_back(argmin(in, reduce_ax, true, stream()));
+  } else {
+    out.push_back(argmax(in, reduce_ax, true, stream()));
+  }
+  return {out, axes};
+}
+
+std::vector<array> ArgReduce::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>&,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {zeros_like(primals[0], stream())};
+}
+
+std::vector<array> ArgReduce::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>&,
+    const std::vector<int>&) {
+  auto shape = output_shapes(primals)[0];
+  return {zeros(shape, uint32, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ArgSort::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  int axis_left = axes[0] >= 0 && axes[0] <= axis_;
+  return {{argsort(inputs[0], axis_ + axis_left, stream())}, axes};
+}
+
+std::vector<Shape> ArgReduce::output_shapes(const std::vector<array>& inputs) {
+  auto out_shape = inputs[0].shape();
+  out_shape[axis_] = 1;
+  return {std::move(out_shape)};
+}
+
+bool ArgSort::is_equivalent(const Primitive& other) const {
+  const ArgSort& r_other = static_cast<const ArgSort&>(other);
+  return axis_ == r_other.axis_;
+}
+
+std::vector<array> ArgSort::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>&,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {zeros_like(primals[0], stream())};
+}
+
+std::vector<array> ArgSort::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>&,
+    const std::vector<int>&) {
+  return {zeros(primals[0].shape(), uint32, stream())};
+}
+
+std::vector<array> AsType::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  if (cotangents[0].dtype() != dtype_) {
+    throw std::invalid_argument(
+        "[astype] Type of cotangents does not match primal output type.");
+  }
+  return {astype(cotangents[0], primals[0].dtype(), stream())};
+}
+
+std::vector<array> AsType::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  return {astype(tangents[0], dtype_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> AsType::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  return {{astype(inputs[0], dtype_, stream())}, axes};
+}
+
+bool AsType::is_equivalent(const Primitive& other) const {
+  const AsType& a_other = static_cast<const AsType&>(other);
+  return dtype_ == a_other.dtype_;
+}
+
+std::vector<array> AsStrided::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(argnums.size() == 1);
+
+  // Extract the sizes and cast them to ints
+  int grad_size = primals[0].size();
+  int cotangents_size = cotangents[0].size();
+
+  // Make a flat container to hold the gradients
+  auto grad = zeros_like(primals[0], stream());
+  grad = reshape(grad, {grad_size}, stream());
+
+  // Create the indices that map output to input
+  auto idx = arange(grad_size, stream());
+  idx = as_strided(idx, shape_, strides_, offset_, stream());
+  idx = reshape(idx, {cotangents_size}, stream());
+
+  // Reshape the cotangentsgent for use with scatter
+  auto flat_cotangents = reshape(cotangents[0], {cotangents_size, 1}, stream());
+
+  // Finally accumulate the gradients and reshape them to look like the input
+  grad = scatter_add(grad, idx, flat_cotangents, 0, stream());
+  grad = reshape(grad, primals[0].shape(), stream());
+
+  return {grad};
+}
+
+std::vector<array> AsStrided::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+
+  return {as_strided(tangents[0], shape_, strides_, offset_, stream())};
+}
+
+bool AsStrided::is_equivalent(const Primitive& other) const {
+  const AsStrided& a_other = static_cast<const AsStrided&>(other);
+  return shape_ == a_other.shape_ && strides_ == a_other.strides_ &&
+      offset_ == a_other.offset_;
+}
+
+bool BitwiseBinary::is_equivalent(const Primitive& other) const {
+  const BitwiseBinary& a_other = static_cast<const BitwiseBinary&>(other);
+  return op_ == a_other.op_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> BitwiseBinary::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {
+      {array(
+          a.shape(),
+          a.dtype(),
+          std::make_shared<BitwiseBinary>(stream(), op_),
+          {a, b})},
+      {to_ax}};
+}
+
+std::vector<array> BitwiseBinary::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 2);
+  std::vector<array> vjps = {zeros_like(tangents[0], stream())};
+  if (argnums.size() > 1) {
+    vjps.push_back(vjps.back());
+  }
+  return vjps;
+}
+
+std::vector<array> BitwiseBinary::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array>
+broadcast_vjp(const array& primal, const array& cotan, const Stream& s) {
+  // Reduce cotangents to the shape of the primal
+  auto& shape = primal.shape();
+  int diff = cotan.ndim() - shape.size();
+  std::vector<int> squeeze_axes(diff);
+  std::iota(squeeze_axes.begin(), squeeze_axes.end(), 0);
+  auto reduce_axes = squeeze_axes;
+  for (int i = diff; i < cotan.ndim(); ++i) {
+    if (shape[i - diff] != cotan.shape(i)) {
+      reduce_axes.push_back(i);
+    }
+  }
+  return {squeeze(sum(cotan, reduce_axes, true, s), squeeze_axes, s)};
+}
+
+std::vector<array> Broadcast::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return broadcast_vjp(primals[0], cotangents[0], stream());
+}
+
+std::vector<array> Broadcast::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  return {array(
+      shape_,
+      tangents[0].dtype(),
+      std::make_shared<Broadcast>(stream(), shape_),
+      tangents)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Broadcast::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0];
+  auto& in = inputs[0];
+  if (ax >= 0) {
+    int diff = shape_.size() - in.ndim() + 1;
+    assert(diff >= 0);
+    shape_.insert(shape_.begin() + ax + diff, in.shape(ax));
+    ax += diff;
+  }
+  return {{broadcast_to(in, shape_, stream())}, {ax}};
+}
+
+bool Broadcast::is_equivalent(const Primitive& other) const {
+  const Broadcast& b_other = static_cast<const Broadcast&>(other);
+  return shape_ == b_other.shape_;
+}
+
+Shape Broadcast::output_shape(const std::vector<array>& inputs) {
+  auto shape = inputs[0].shape();
+  for (int i = 1; i < inputs.size(); ++i) {
+    shape = broadcast_shapes(shape, inputs[i].shape());
+  }
+  return shape;
+}
+
+std::vector<Shape> Broadcast::output_shapes(const std::vector<array>& inputs) {
+  if (inputs.size() < 2) {
+    if (broadcast_shapes(inputs[0].shape(), shape_) != shape_) {
+      throw std::invalid_argument(
+          "[Broadcast] Unable to infer broadcast shape");
+    }
+    return {shape_};
+  }
+  return {output_shape(inputs)};
+};
+
+std::vector<array> BroadcastAxes::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return broadcast_vjp(primals[0], cotangents[0], stream());
+}
+
+std::vector<array> BroadcastAxes::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  return {array(
+      output_shape(primals, ignore_axes_),
+      tangents[0].dtype(),
+      std::make_shared<BroadcastAxes>(stream(), ignore_axes_),
+      tangents)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> BroadcastAxes::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  throw std::invalid_argument("[BroadcastAxes] VMAP NYI");
+}
+
+bool BroadcastAxes::is_equivalent(const Primitive& other) const {
+  const auto& b_other = static_cast<const BroadcastAxes&>(other);
+  return ignore_axes_ == b_other.ignore_axes_;
+}
+
+Shape BroadcastAxes::output_shape(
+    const std::vector<array>& inputs,
+    const std::vector<int>& ignore_axes) {
+  auto shape = Shape{};
+  for (auto& in : inputs) {
+    auto in_shape = in.shape();
+    for (auto it = ignore_axes.rbegin(); it != ignore_axes.rend(); ++it) {
+      in_shape.erase(in_shape.begin() + in.ndim() + *it);
+    }
+    shape = broadcast_shapes(shape, in_shape);
+  }
+  int dims = ignore_axes.size() + shape.size();
+  for (auto ax : ignore_axes) {
+    shape.insert(shape.begin() + dims + ax, inputs[0].shape(ax));
+  }
+  return shape;
+}
+
+std::vector<Shape> BroadcastAxes::output_shapes(
+    const std::vector<array>& inputs) {
+  return {output_shape(inputs, ignore_axes_)};
+}
+
+std::vector<array> Ceil::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Ceil::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {zeros_like(primals[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Ceil::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{ceil(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Cholesky::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0] >= 0 ? 0 : -1;
+  auto a = axes[0] > 0 ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  return {{linalg::cholesky(a, upper_, stream())}, {ax}};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Eig::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  bool needs_move = axes[0] >= (inputs[0].ndim() - 2);
+  auto a = needs_move ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  auto ax = needs_move ? 0 : axes[0];
+
+  std::vector<array> outputs;
+  if (compute_eigenvectors_) {
+    auto [values, vectors] = linalg::eig(a, stream());
+    outputs = {values, vectors};
+  } else {
+    outputs = {linalg::eigvals(a, stream())};
+  }
+
+  return {outputs, std::vector<int>(outputs.size(), ax)};
+}
+
+std::vector<Shape> Eig::output_shapes(const std::vector<array>& inputs) {
+  auto shape = inputs[0].shape();
+  shape.pop_back(); // Remove last dimension for eigenvalues
+  if (compute_eigenvectors_) {
+    return {
+        std::move(shape), inputs[0].shape()}; // Eigenvalues and eigenvectors
+  } else {
+    return {std::move(shape)}; // Only eigenvalues
+  }
+}
+
+bool Eig::is_equivalent(const Primitive& other) const {
+  auto& e_other = static_cast<const Eig&>(other);
+  return compute_eigenvectors_ == e_other.compute_eigenvectors_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Eigh::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  bool needs_move = axes[0] >= (inputs[0].ndim() - 2);
+  auto a = needs_move ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  auto ax = needs_move ? 0 : axes[0];
+
+  std::vector<array> outputs;
+  if (compute_eigenvectors_) {
+    auto [values, vectors] = linalg::eigh(a, uplo_, stream());
+    outputs = {values, vectors};
+  } else {
+    outputs = {linalg::eigvalsh(a, uplo_, stream())};
+  }
+
+  return {outputs, std::vector<int>(outputs.size(), ax)};
+}
+
+std::vector<Shape> Eigh::output_shapes(const std::vector<array>& inputs) {
+  auto shape = inputs[0].shape();
+  shape.pop_back(); // Remove last dimension for eigenvalues
+  if (compute_eigenvectors_) {
+    return {
+        std::move(shape), inputs[0].shape()}; // Eigenvalues and eigenvectors
+  } else {
+    return {std::move(shape)}; // Only eigenvalues
+  }
+}
+
+bool Eigh::is_equivalent(const Primitive& other) const {
+  auto& e_other = static_cast<const Eigh&>(other);
+  return uplo_ == e_other.uplo_ &&
+      compute_eigenvectors_ == e_other.compute_eigenvectors_;
+}
+
+std::vector<array> Concatenate::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto& cotan = cotangents[0];
+  Shape start(cotan.ndim(), 0);
+  Shape stop = cotan.shape();
+
+  Shape sizes;
+  sizes.push_back(0);
+  for (auto& p : primals) {
+    sizes.push_back(p.shape(axis_));
+  }
+  std::partial_sum(sizes.cbegin(), sizes.cend(), sizes.begin());
+
+  std::vector<array> grads;
+  for (auto i : argnums) {
+    start[axis_] = sizes[i];
+    stop[axis_] = sizes[i + 1];
+    grads.push_back(slice(cotan, start, stop, stream()));
+  }
+  return grads;
+}
+
+std::vector<array> Concatenate::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  std::vector<int> argidx(argnums.size());
+  std::iota(argidx.begin(), argidx.end(), 0);
+  std::sort(argidx.begin(), argidx.end(), [&argnums](int a, int b) {
+    return argnums[a] < argnums[b];
+  });
+
+  std::vector<array> vals;
+  for (int i = 0, j = 0; i < primals.size(); ++i) {
+    if (j < argnums.size() && argnums[argidx[j]] == i) {
+      vals.push_back(tangents[argidx[j++]]);
+    } else {
+      vals.push_back(zeros_like(primals[i], stream()));
+    }
+  }
+  return {concatenate(vals, axis_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Concatenate::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  int out_ax = -1;
+  int first_vmap = -1;
+
+  // Find the first vmapped input
+  for (int i = 0; i < axes.size(); i++) {
+    if (axes[i] >= 0) {
+      out_ax = axes[i];
+      first_vmap = i;
+      break;
+    }
+  }
+
+  // No vmap, should we even be in here?
+  if (out_ax < 0) {
+    return {{concatenate(inputs, axis_, stream())}, {out_ax}};
+  }
+
+  // Make sure vmapped arrays have all vmapped axes in the same location and
+  // expand non-vmapped arrays to be compatible with the vmapped ones.
+  std::vector<array> t_inputs;
+  int axis = axis_ + (axis_ >= out_ax);
+  auto cat_shape = inputs[first_vmap].shape();
+  for (int i = 0; i < axes.size(); i++) {
+    if (axes[i] >= 0) {
+      if (out_ax != axes[i]) {
+        t_inputs.push_back(moveaxis(inputs[i], axes[i], out_ax, stream()));
+      } else {
+        t_inputs.push_back(inputs[i]);
+      }
+    } else {
+      cat_shape[axis] = inputs[i].shape(axis_);
+      t_inputs.push_back(broadcast_to(
+          expand_dims(inputs[i], out_ax, stream()), cat_shape, stream()));
+    }
+  }
+
+  return {{concatenate(t_inputs, axis, stream())}, {out_ax}};
+}
+
+bool Concatenate::is_equivalent(const Primitive& other) const {
+  const Concatenate& c_other = static_cast<const Concatenate&>(other);
+  return axis_ == c_other.axis_;
+}
+
+std::vector<Shape> Concatenate::output_shapes(
+    const std::vector<array>& inputs) {
+  auto shape = inputs[0].shape();
+  for (int i = 1; i < inputs.size(); ++i) {
+    shape[axis_] += inputs[i].shape(axis_);
+  }
+  return {std::move(shape)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Conjugate::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{conjugate(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Contiguous::vjp(
+    const std::vector<array>&,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {cotangents};
+}
+
+std::vector<array> Contiguous::jvp(
+    const std::vector<array>&,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {tangents};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Contiguous::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  return {{contiguous(inputs[0], allow_col_major_, stream())}, axes};
+}
+
+bool Contiguous::is_equivalent(const Primitive& other) const {
+  const Contiguous& c_other = static_cast<const Contiguous&>(other);
+  return allow_col_major_ == c_other.allow_col_major_;
+}
+
+array conv_weight_backward_patches(
+    const array& in,
+    const array& wt,
+    const array& cotan,
+    const std::vector<int>& kernel_strides,
+    const std::vector<int>& padding_lo,
+    const std::vector<int>& padding_hi,
+    StreamOrDevice s) {
+  // Resolve Padded input shapes and strides
+  Shape padding_starts(in.ndim(), 0);
+  auto padding_ends = in.shape();
+  auto in_padded_shape = in.shape();
+
+  // padded shape
+  for (int i = 1; i < in.ndim() - 1; i++) {
+    in_padded_shape[i] += padding_lo[i - 1] + padding_hi[i - 1];
+    padding_ends[i] += padding_lo[i - 1];
+    padding_starts[i] += padding_lo[i - 1];
+  }
+
+  // padded strides (contiguous)
+  Strides in_padded_strides(in.ndim(), 1);
+  for (int i = in.ndim() - 2; i >= 0; --i) {
+    in_padded_strides[i] = in_padded_strides[i + 1] * in_padded_shape[i + 1];
+  }
+
+  // Pad input
+  std::vector<int> padded_axes(in.ndim() - 2, 0);
+  std::iota(padded_axes.begin(), padded_axes.end(), 1);
+  auto in_padded =
+      pad(in,
+          padded_axes,
+          Shape(padding_lo.begin(), padding_lo.end()),
+          Shape(padding_hi.begin(), padding_hi.end()),
+          array(0, in.dtype()),
+          "constant",
+          s);
+
+  // Resolve strided patches
+
+  // patches are shaped as
+  // (batch_dim, out_spatial_dims, weight_spatial_dims, in_channels)
+  Shape patches_shape{cotan.shape().begin(), cotan.shape().end() - 1};
+  patches_shape.insert(
+      patches_shape.end(), wt.shape().begin() + 1, wt.shape().end());
+
+  // Resolve patch strides
+  int n_spatial_dim = in.ndim() - 2;
+  Strides patches_strides(patches_shape.size(), 1);
+  patches_strides[0] = in_padded_strides[0];
+  for (int i = 1; i < n_spatial_dim + 1; i++) {
+    patches_strides[i] = in_padded_strides[i] * kernel_strides[i - 1];
+  }
+  for (int i = 1; i < in.ndim(); i++) {
+    patches_strides[n_spatial_dim + i] = in_padded_strides[i];
+  }
+
+  // Make patches from in
+  auto in_patches = as_strided(in_padded, patches_shape, patches_strides, 0, s);
+
+  // Prepare for matmul
+  int O = wt.shape(0);
+  auto cotan_mat = reshape(cotan, {-1, O}, s);
+  in_patches = reshape(in_patches, {cotan_mat.shape(0), -1}, s);
+
+  auto grad = matmul(transpose(cotan_mat, {1, 0}, s), in_patches, s);
+  grad = reshape(grad, wt.shape(), s);
+  return grad;
+}
+
+namespace {
+
+// Conv helpers
+inline int conv_out_axis_size(int in_dim, int wt_dim, int stride, int padding) {
+  return ((in_dim + padding - wt_dim) / stride) + 1;
+}
+
+// Conv helpers
+inline int dilate_size(int dim, int dil) {
+  return 1 + dil * (dim - 1);
+}
+
+} // namespace
+
+Shape Convolution::conv_out_shape(
+    const Shape& in_shape,
+    const Shape& wt_shape,
+    const std::vector<int>& strides,
+    const std::vector<int>& pads_lo,
+    const std::vector<int>& pads_hi,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation) {
+  int N = in_shape[0];
+  int O = wt_shape[0];
+  Shape out_shape(in_shape.size());
+  int i = 0;
+  out_shape[i++] = N;
+
+  int spatial_dims = in_shape.size() - 2;
+
+  if (strides.size() != spatial_dims) {
+    std::ostringstream msg;
+    msg << "[conv] Invalid strides " << strides << " for " << spatial_dims
+        << "D convolution.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (pads_lo.size() != spatial_dims || pads_hi.size() != spatial_dims) {
+    std::ostringstream msg;
+    msg << "[conv] Invalid padding " << pads_lo << " | " << pads_hi << " for "
+        << spatial_dims << "D convolution.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (kernel_dilation.size() != spatial_dims) {
+    std::ostringstream msg;
+    msg << "[conv] Invalid kernel dilation " << kernel_dilation << " for "
+        << spatial_dims << "D convolution.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (input_dilation.size() != spatial_dims) {
+    std::ostringstream msg;
+    msg << "[conv] Invalid input dilation " << input_dilation << " for "
+        << spatial_dims << "D convolution.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  for (; i < in_shape.size() - 1; i++) {
+    if (kernel_dilation[i - 1] <= 0) {
+      std::ostringstream msg;
+      msg << "[conv] Kernel dilation sizes must be positive."
+          << " Got kernel dilation " << kernel_dilation << ".";
+      throw std::invalid_argument(msg.str());
+    }
+
+    if (input_dilation[i - 1] <= 0) {
+      std::ostringstream msg;
+      msg << "[conv] Input dilation sizes must be positive."
+          << " Got input dilation " << input_dilation << ".";
+      throw std::invalid_argument(msg.str());
+    }
+
+    if (pads_lo[i - 1] < 0 || pads_hi[i - 1] < 0) {
+      std::ostringstream msg;
+      msg << "[conv] Padding sizes must be non-negative. Got padding "
+          << pads_lo << " | " << pads_hi << ".";
+      throw std::invalid_argument(msg.str());
+    }
+
+    if (strides[i - 1] <= 0) {
+      std::ostringstream msg;
+      msg << "[conv] Stride sizes must be positive."
+          << " Got strides " << strides << ".";
+      throw std::invalid_argument(msg.str());
+    }
+
+    int kd = dilate_size(wt_shape[i], kernel_dilation[i - 1]);
+    int id = dilate_size(in_shape[i], input_dilation[i - 1]);
+
+    out_shape[i] = conv_out_axis_size(
+        id, kd, strides[i - 1], pads_lo[i - 1] + pads_hi[i - 1]);
+
+    if (out_shape[i] <= 0) {
+      std::ostringstream msg;
+      msg << "[conv] Spatial dimensions of input after padding"
+          << " cannot be smaller than weight spatial dimensions."
+          << " Got error at axis " << i << " for input with shape " << in_shape
+          << ", padding low " << pads_lo << ", padding high " << pads_hi
+          << ", and weight of shape " << wt_shape << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  out_shape[i] = O;
+
+  return out_shape;
+}
+
+std::vector<array> Convolution::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 2);
+  std::vector<array> grads;
+
+  // Collect info
+  auto& in = primals[0];
+  auto& wt = primals[1];
+  auto& cotan = cotangents[0];
+
+  auto group_transpose =
+      [this](const array& x, int group_dim, int ax_a, int ax_b) {
+        if (groups_ > 1) {
+          auto shape = x.shape();
+          if (group_dim < 0) {
+            group_dim += shape.size();
+          }
+          shape.insert(shape.begin() + group_dim, groups_);
+          shape[group_dim + 1] = shape[group_dim + 1] / groups_;
+          auto x_trans = swapaxes(
+              reshape(x, std::move(shape), stream()), ax_a, ax_b, stream());
+          return flatten(x_trans, group_dim, group_dim + 1, stream());
+        } else {
+          return swapaxes(x, 0, -1, stream());
+        }
+      };
+
+  for (int a : argnums) {
+    // Grads for input
+    if (a == 0) {
+      std::vector<int> padding_lo = padding_lo_;
+      std::vector<int> padding_hi = padding_hi_;
+
+      for (int i = 0; i < padding_lo.size(); ++i) {
+        int wt_size = 1 + kernel_dilation_[i] * (wt.shape(1 + i) - 1);
+        padding_lo[i] = wt_size - padding_lo_[i] - 1;
+
+        int in_size = 1 + input_dilation_[i] * (in.shape(1 + i) - 1);
+        int out_size = 1 + kernel_strides_[i] * (cotan.shape(1 + i) - 1);
+        padding_hi[i] = in_size - out_size + padding_hi_[i];
+      }
+
+      // Check for negative padding
+      bool has_neg_padding = false;
+      for (auto& pd : padding_lo) {
+        has_neg_padding |= (pd < 0);
+      }
+      for (auto& pd : padding_hi) {
+        has_neg_padding |= (pd < 0);
+      }
+
+      auto wt_trans = group_transpose(wt, 0, 1, -1);
+      auto grad = conv_general(
+          /* const array& input = */ cotan,
+          /* const array& weight = */ wt_trans,
+          /* std::vector<int> stride = */ input_dilation_,
+          /* std::vector<int> padding_lo = */ padding_lo,
+          /* std::vector<int> padding_hi = */ padding_hi,
+          /* std::vector<int> kernel_dilation = */ kernel_dilation_,
+          /* std::vector<int> input_dilation = */ kernel_strides_,
+          /* int groups = */ groups_,
+          /* bool flip = */ !flip_,
+          stream());
+
+      // Handle negative padding
+      if (has_neg_padding) {
+        Shape starts(grad.ndim(), 0);
+        auto stops = grad.shape();
+
+        for (int i = 0; i < grad.ndim() - 2; i++) {
+          if (padding_lo[i] < 0) {
+            starts[i + 1] -= padding_lo[i];
+          }
+          if (padding_hi[i] < 0) {
+            stops[i + 1] += padding_hi[i];
+          }
+        }
+
+        grad = slice(grad, std::move(starts), std::move(stops), stream());
+      }
+
+      grads.push_back(grad);
+    }
+    // Grads for weight
+    else if (a == 1) {
+      bool no_dilation = true;
+
+      for (int i = 0; i < input_dilation_.size(); i++) {
+        no_dilation &= (input_dilation_[i] == 1) && (kernel_dilation_[i] == 1);
+      }
+
+      if (no_dilation && !flip_ && groups_ == 1) {
+        auto grad = conv_weight_backward_patches(
+            in, wt, cotan, kernel_strides_, padding_lo_, padding_hi_, stream());
+        grads.push_back(grad);
+      } else {
+        auto padding_hi = padding_lo_;
+
+        for (int i = 0; i < padding_hi.size(); ++i) {
+          int in_size = 1 + input_dilation_[i] * (in.shape(1 + i) - 1);
+          int out_size = 1 + kernel_strides_[i] * (cotan.shape(1 + i) - 1);
+          int wt_size = 1 + kernel_dilation_[i] * (wt.shape(1 + i) - 1);
+          padding_hi[i] = out_size - in_size + wt_size - padding_hi[i] - 1;
+        }
+
+        auto cotan_trans = swapaxes(cotan, 0, -1, stream());
+        auto in_trans = group_transpose(in, -1, 0, -1);
+
+        auto grad_trans = conv_general(
+            /* const array& input = */ in_trans,
+            /* const array& weight = */ cotan_trans,
+            /* std::vector<int> stride = */ kernel_dilation_,
+            /* std::vector<int> padding_lo = */ padding_lo_,
+            /* std::vector<int> padding_hi = */ padding_hi,
+            /* std::vector<int> kernel_dilation = */ kernel_strides_,
+            /* std::vector<int> input_dilation = */ input_dilation_,
+            /* int groups = */ groups_,
+            /* bool flip = */ false,
+            stream());
+        if (flip_) {
+          auto start = Shape(grad_trans.ndim(), 0);
+          auto stop = Shape(grad_trans.ndim(), 0);
+          auto strides = Shape(grad_trans.ndim(), 1);
+          for (int i = 0; i < stop.size(); ++i) {
+            if (i >= 1 && i < stop.size() - 1) {
+              start[i] = grad_trans.shape(i);
+              stop[i] = -start[i] - 1;
+              strides[i] = -1;
+            } else {
+              stop[i] = grad_trans.shape(i);
+            }
+          }
+          grad_trans = slice(grad_trans, start, stop, strides, stream());
+        }
+        grads.push_back(swapaxes(grad_trans, 0, -1, stream()));
+      }
+    }
+  }
+
+  return grads;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Convolution::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto do_conv = [&](const array& in, const array& w, int groups) {
+    return conv_general(
+        in,
+        w,
+        kernel_strides_,
+        padding_lo_,
+        padding_hi_,
+        kernel_dilation_,
+        input_dilation_,
+        groups,
+        flip_,
+        stream());
+  };
+  bool in_vmap = axes[0] >= 0;
+  bool w_vmap = axes[1] >= 0;
+  auto in = inputs[0];
+  auto w = inputs[1];
+  if (in_vmap && !w_vmap) {
+    // flatten / unflatten the batch dimension
+    // of the input / output
+    if (axes[0] > 0) {
+      in = moveaxis(in, axes[0], 0, stream());
+    }
+    auto out = do_conv(flatten(in, 0, 1, stream()), w, groups_);
+    out = unflatten(out, 0, {in.shape(0), in.shape(1)}, stream());
+    return {{out}, {0}};
+  } else if (!in_vmap && w_vmap) {
+    // flatten into the output channels of w
+    // unflatten the channels of the output
+    if (axes[1] > 0) {
+      w = moveaxis(w, axes[1], 0, stream());
+    }
+    auto out = do_conv(in, flatten(w, 0, 1, stream()), groups_);
+    out = unflatten(out, -1, {w.shape(0), w.shape(1)}, stream());
+    return {{out}, {static_cast<int>(out.ndim() - 2)}};
+  } else if (in_vmap && w_vmap) {
+    // use a group convolution when both inputs are vmapped
+    auto b = in.shape(axes[0]);
+    in = moveaxis(in, axes[0], -2, stream());
+    in = flatten(in, -2, -1, stream());
+    if (axes[1] > 0) {
+      w = moveaxis(w, axes[1], 0, stream());
+    }
+    auto c_out = w.shape(1);
+    w = flatten(w, 0, 1, stream());
+    auto out = do_conv(in, w, groups_ * b);
+    out = unflatten(out, -1, {b, c_out}, stream());
+    return {{out}, {static_cast<int>(out.ndim() - 2)}};
+  } else {
+    return {{do_conv(in, w, groups_)}, {-1}};
+  }
+}
+
+bool Convolution::is_equivalent(const Primitive& other) const {
+  const Convolution& c_other = static_cast<const Convolution&>(other);
+  return padding_lo_ == c_other.padding_lo_ &&
+      padding_hi_ == c_other.padding_hi_ &&
+      kernel_strides_ == c_other.kernel_strides_ &&
+      kernel_dilation_ == c_other.kernel_dilation_ &&
+      input_dilation_ == c_other.input_dilation_ &&
+      groups_ == c_other.groups_ && flip_ == c_other.flip_;
+}
+
+std::vector<Shape> Convolution::output_shapes(
+    const std::vector<array>& inputs) {
+  return {conv_out_shape(
+      inputs[0].shape(), // in_shape
+      inputs[1].shape(), // wt_shape
+      kernel_strides_,
+      padding_lo_,
+      padding_hi_,
+      kernel_dilation_,
+      input_dilation_)};
+}
+
+std::vector<array> Copy::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return cotangents;
+}
+
+std::vector<array> Copy::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return tangents;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Copy::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{copy(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Cos::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return {jvp(primals, cotangents, argnums)};
+}
+
+std::vector<array> Cos::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(
+      tangents[0], negative(sin(primals[0], stream()), stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Cos::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{cos(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Cosh::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Cosh::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(tangents[0], sinh(primals[0], stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Cosh::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{cosh(inputs[0], stream())}, axes};
+}
+
+std::vector<array> CustomTransforms::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  // Extract the inputs to the VJP function
+  std::vector<array> inputs(primals.begin(), primals.end() - num_outputs_);
+
+  // Compute all the vjps
+  auto all_vjps = vjp_fun_(inputs, cotangents, outputs);
+  for (const auto& cot : cotangents) {
+    all_vjps.emplace_back(cot);
+  }
+
+  // Select the vjps requested
+  std::vector<array> vjps;
+  vjps.reserve(argnums.size());
+  for (auto arg : argnums) {
+    vjps.push_back(all_vjps[arg]);
+  }
+
+  return vjps;
+}
+
+std::vector<array> CustomTransforms::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  // Extract the inputs to the JVP function
+  std::vector<array> inputs(primals.begin(), primals.end() - num_outputs_);
+
+  // Compute the jvps
+  return jvp_fun_(inputs, tangents, argnums);
+}
+
+std::pair<std::vector<array>, std::vector<int>> CustomTransforms::vmap(
+    const std::vector<array>& inputs_,
+    const std::vector<int>& axes_) {
+  // Extract the inputs to the vmap function
+  std::vector<array> inputs(inputs_.begin(), inputs_.end() - num_outputs_);
+  std::vector<int> axes(axes_.begin(), axes_.end() - num_outputs_);
+  return vmap_fun_(inputs, axes);
+}
+
+std::vector<array> Depends::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  std::vector<array> vjps;
+
+  for (auto arg : argnums) {
+    if (arg < cotangents.size()) {
+      vjps.push_back(cotangents[arg]);
+    } else {
+      vjps.push_back(zeros_like(primals[arg]));
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> Divide::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  array denominator_bar = conjugate(primals[1], stream());
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      vjps.push_back(divide(cotangents[0], denominator_bar, stream()));
+    } else {
+      vjps.push_back(negative(
+          divide(
+              multiply(
+                  cotangents[0], conjugate(primals[0], stream()), stream()),
+              square(denominator_bar, stream()),
+              stream()),
+          stream()));
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> DivMod::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(zeros_like(primals[arg], stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> DivMod::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  return {zeros_like(primals[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> DivMod::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {divmod(a, b, stream()), {to_ax}};
+}
+
+std::vector<array> Divide::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+    if (arg == 0) {
+      return divide(tangents[i], primals[1], stream());
+    } else {
+      return negative(
+          divide(
+              multiply(tangents[i], primals[0], stream()),
+              square(primals[1], stream()),
+              stream()),
+          stream());
+    }
+  };
+  auto out = jvp_fun(0);
+  if (argnums.size() > 1) {
+    out = add(out, jvp_fun(1), stream());
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Divide::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{divide(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Remainder::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      vjps.push_back(cotangents[0]);
+    } else {
+      auto x_over_y = divide(primals[0], primals[1], stream());
+      x_over_y = floor(x_over_y, stream());
+      vjps.push_back(
+          negative(multiply(x_over_y, cotangents[0], stream()), stream()));
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> Remainder::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+    if (arg == 0) {
+      return tangents[i];
+    } else {
+      auto x_over_y = divide(primals[0], primals[1], stream());
+      x_over_y = floor(x_over_y, stream());
+      return negative(multiply(x_over_y, tangents[i], stream()), stream());
+    }
+  };
+  auto out = jvp_fun(0);
+  if (argnums.size() > 1) {
+    out = add(out, jvp_fun(1), stream());
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Remainder::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{remainder(a, b, stream())}, {to_ax}};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Equal::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{equal(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Equal::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(zeros_like(primals[arg], stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> Equal::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape());
+  return {zeros(shape, bool_, stream())};
+}
+
+std::vector<array> Erf::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Erf::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  auto dtype = primals[0].dtype();
+  auto scale = multiply(array(M_2_SQRTPI, dtype), tangents[0], stream());
+  return {multiply(
+      scale,
+      exp(negative(square(primals[0], stream()), stream()), stream()),
+      stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Erf::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{erf(inputs[0], stream())}, axes};
+}
+
+std::vector<array> ErfInv::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  auto dtype = primals[0].dtype();
+  auto scale =
+      multiply(array(1.0 / M_2_SQRTPI, dtype), cotangents[0], stream());
+  return {
+      multiply(scale, exp(square(outputs[0], stream()), stream()), stream())};
+}
+
+std::vector<array> ErfInv::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  auto dtype = primals[0].dtype();
+  auto scale = multiply(array(1.0 / M_2_SQRTPI, dtype), tangents[0], stream());
+  return {multiply(
+      scale,
+      exp(square(erfinv(primals[0], stream()), stream()), stream()),
+      stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ErfInv::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{erfinv(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Exp::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  return {multiply(cotangents[0], outputs[0], stream())};
+}
+
+std::vector<array> Exp::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(tangents[0], exp(primals[0], stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Exp::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{exp(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Expm1::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  return {multiply(
+      cotangents[0],
+      add(outputs[0], array(1.0f, outputs[0].dtype()), stream()),
+      stream())};
+}
+
+std::vector<array> Expm1::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(tangents[0], exp(primals[0], stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Expm1::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{expm1(inputs[0], stream())}, axes};
+}
+
+std::vector<array> ExpandDims::vjp(
+    const std::vector<array>&,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {squeeze(cotangents[0], axes_, stream())};
+}
+
+std::vector<array> ExpandDims::jvp(
+    const std::vector<array>&,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {expand_dims(tangents[0], axes_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> ExpandDims::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0];
+  auto expand_axes = axes_;
+  for (auto& s : expand_axes) {
+    if (s >= axes[0]) {
+      s++;
+    } else {
+      ax++;
+    }
+  }
+  return {{expand_dims(inputs[0], std::move(expand_axes), stream())}, {ax}};
+}
+
+bool ExpandDims::is_equivalent(const Primitive& other) const {
+  const ExpandDims& a_other = static_cast<const ExpandDims&>(other);
+  return (axes_ == a_other.axes_);
+}
+
+Shape ExpandDims::output_shape(
+    const array& input,
+    const std::vector<int>& axes) {
+  auto shape = input.shape();
+  for (auto ax : axes) {
+    shape.insert(shape.begin() + ax, 1);
+  }
+  return shape;
+}
+
+std::vector<Shape> ExpandDims::output_shapes(const std::vector<array>& inputs) {
+  return {ExpandDims::output_shape(inputs[0], axes_)};
+}
+
+std::vector<array> Flatten::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  auto& in = primals[0];
+  Shape unflatten_shape(
+      in.shape().begin() + start_axis_, in.shape().begin() + end_axis_ + 1);
+  return {unflatten(
+      cotangents[0], start_axis_, std::move(unflatten_shape), stream())};
+}
+
+std::vector<array> Flatten::jvp(
+    const std::vector<array>&,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {flatten(tangents[0], start_axis_, end_axis_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Flatten::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0];
+  auto start_axis = start_axis_;
+  auto end_axis = end_axis_;
+  auto in = inputs[0];
+  if (ax < start_axis) {
+    start_axis++;
+    end_axis++;
+  } else if (ax <= end_axis_) {
+    start_axis++;
+    end_axis++;
+    in = moveaxis(in, ax, 0, stream());
+    ax = 0;
+  } else {
+    ax -= (end_axis - start_axis);
+  }
+  return {{flatten(in, start_axis, end_axis, stream())}, {ax}};
+}
+
+bool Flatten::is_equivalent(const Primitive& other) const {
+  const Flatten& a_other = static_cast<const Flatten&>(other);
+  return start_axis_ == a_other.start_axis_ && end_axis_ == a_other.end_axis_;
+}
+
+Shape Flatten::output_shape(const array& input, int start_axis, int end_axis) {
+  Shape shape = input.shape();
+  auto flat_size = input.shape(start_axis);
+  for (int ax = start_axis + 1; ax <= end_axis; ++ax) {
+    flat_size *= input.shape(ax);
+  }
+  shape.erase(shape.begin() + start_axis + 1, shape.begin() + end_axis + 1);
+  shape[start_axis] = flat_size;
+  return shape;
+}
+
+std::vector<Shape> Flatten::output_shapes(const std::vector<array>& inputs) {
+  return {Flatten::output_shape(inputs[0], start_axis_, end_axis_)};
+}
+
+bool FFT::is_equivalent(const Primitive& other) const {
+  const FFT& r_other = static_cast<const FFT&>(other);
+  return axes_ == r_other.axes_ && inverse_ == r_other.inverse_ &&
+      real_ == r_other.real_;
+}
+
+std::vector<array> Unflatten::vjp(
+    const std::vector<array>&,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {flatten(cotangents[0], axis_, axis_ + shape_.size() - 1, stream())};
+}
+
+std::vector<array> Unflatten::jvp(
+    const std::vector<array>&,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {unflatten(tangents[0], axis_, shape_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Unflatten::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0];
+  auto axis = axis_;
+  if (ax <= axis_) {
+    axis++;
+  } else {
+    ax += (shape_.size() - 1);
+  }
+  return {{unflatten(inputs[0], axis, shape_, stream())}, {ax}};
+}
+
+bool Unflatten::is_equivalent(const Primitive& other) const {
+  const auto& a_other = static_cast<const Unflatten&>(other);
+  return axis_ == a_other.axis_ && shape_ == a_other.shape_;
+}
+
+Shape Unflatten::output_shape(
+    const array& input,
+    int axis,
+    const Shape& shape) {
+  Shape out_shape = input.shape();
+  out_shape[axis] = shape[0];
+  out_shape.insert(
+      out_shape.begin() + axis + 1, shape.begin() + 1, shape.end());
+  return out_shape;
+}
+
+std::vector<Shape> Unflatten::output_shapes(const std::vector<array>& inputs) {
+  return {Unflatten::output_shape(inputs[0], axis_, shape_)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> FFT::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto& in = inputs[0];
+  int ax = axes[0];
+  auto fft_axes = axes_;
+  auto out_shape = in.shape();
+  if (ax >= 0) {
+    for (auto& fft_ax : fft_axes) {
+      if (fft_ax >= ax) {
+        fft_ax++;
+      }
+      if (real_) {
+        auto n = out_shape[fft_ax];
+        out_shape[fft_ax] = inverse_ ? 2 * (n - 1) : n / 2 + 1;
+      }
+    }
+  }
+  return {
+      {array(
+          out_shape,
+          real_ && inverse_ ? float32 : complex64,
+          std::make_shared<FFT>(stream(), fft_axes, inverse_, real_),
+          {in})},
+      {ax}};
+}
+
+std::vector<array> FFT::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  auto& in = primals[0];
+  std::vector<int> axes(axes_.begin(), axes_.end());
+
+  // TODO: Add it as an option to do an unnormalized or scaled fft so that this
+  //       isn't part of the graph.
+  double n_elements = 1;
+  for (auto ax : axes) {
+    n_elements *= inverse_ ? cotangents[0].shape(ax) : primals[0].shape(ax);
+  }
+
+  if (real_ && inverse_) {
+    // Make a mask to account for the double use in the forward pass.
+    // Everything except the DC and nyquist frequencies gets doubled.
+    int N = in.shape(axes_.back());
+    bool odd = cotangents[0].shape(axes_.back()) % 2;
+    Shape c(in.ndim(), 1);
+    c[axes_.back()] = N;
+    array indices = reshape(arange(N, stream()), std::move(c), stream());
+    array first(0, indices.dtype());
+    array last(N - 1 + odd, indices.dtype());
+    array one(1 / n_elements, in.dtype());
+    array two(2 / n_elements, in.dtype());
+    array mask = where(
+        logical_and(
+            greater(indices, first, stream()),
+            less(indices, last, stream()),
+            stream()),
+        two,
+        one,
+        stream());
+    return {
+        multiply(fft::rfftn(cotangents[0], axes, stream()), mask, stream())};
+  } else if (real_) {
+    Shape n;
+    for (auto ax : axes_) {
+      n.push_back(in.shape(ax));
+    }
+    // Make a mask to account for the double use in the forward pass.
+    // Everything except the DC and nyquist frequencies gets halved.
+    int N = cotangents[0].shape(axes_.back());
+    bool odd = in.shape(axes_.back()) % 2;
+    Shape c(in.ndim(), 1);
+    c[axes_.back()] = N;
+    array indices = reshape(arange(N, stream()), std::move(c), stream());
+    array first(0, indices.dtype());
+    array last(N - 1 + odd, indices.dtype());
+    array one(1, complex64);
+    array half(0.5, complex64);
+    array mask = where(
+        logical_and(
+            greater(indices, first, stream()),
+            less(indices, last, stream()),
+            stream()),
+        half,
+        one,
+        stream());
+    return {multiply(
+        fft::irfftn(multiply(cotangents[0], mask, stream()), n, axes, stream()),
+        array(n_elements, in.dtype()),
+        stream())};
+  } else if (inverse_) {
+    return {multiply(
+        fft::fftn(cotangents[0], axes, stream()),
+        array(1 / n_elements, complex64),
+        stream())};
+  } else {
+    return {multiply(
+        fft::ifftn(cotangents[0], axes, stream()),
+        array(n_elements, complex64),
+        stream())};
+  }
+}
+
+std::vector<array> FFT::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  auto& tan = tangents[0];
+  if (real_ & inverse_) {
+    return {fft::irfftn(tan, stream())};
+  } else if (real_) {
+    return {fft::rfftn(tan, stream())};
+  } else if (inverse_) {
+    return {fft::ifftn(tan, stream())};
+  } else {
+    return {fft::fftn(tan, stream())};
+  }
+}
+
+std::vector<array> Floor::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Floor::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {zeros_like(primals[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Floor::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{floor(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Full::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(cotangents[0], primals[0], stream())};
+}
+
+std::vector<array> Full::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return tangents;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Full::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  auto& in = inputs[0];
+  auto out =
+      array(in.shape(), in.dtype(), std::make_shared<Full>(stream()), {in});
+  return {{out}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Gather::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto& src = inputs[0];
+  std::vector<array> indices(inputs.begin() + 1, inputs.end());
+  auto gather_axes = axes_;
+  auto slice_sizes = slice_sizes_;
+  auto src_vmapped = axes[0] >= 0;
+  auto ind_vmap_ax_ptr =
+      std::find_if(axes.begin() + 1, axes.end(), [](int a) { return a >= 0; });
+  int out_ax = -1;
+  bool indices_vmapped = (ind_vmap_ax_ptr != axes.end());
+  if (indices_vmapped) {
+    out_ax = *ind_vmap_ax_ptr;
+  } else if (src_vmapped) {
+    out_ax = axes[0];
+  }
+
+  // Reorder all the index arrays so the vmap axis is in the same spot.
+  if (indices_vmapped) {
+    for (int i = 1; i < axes.size(); ++i) {
+      if (out_ax != axes[i] && axes[i] >= 0) {
+        indices[i - 1] = moveaxis(indices[i - 1], axes[i], out_ax, stream());
+      } else if (axes[i] < 0) {
+        indices[i - 1] = expand_dims(indices[i - 1], out_ax, stream());
+      }
+    }
+  }
+
+  int idx_dims = indices.empty() ? 0 : indices[0].ndim();
+
+  if (src_vmapped) {
+    for (auto& ax : gather_axes) {
+      if (ax >= axes[0]) {
+        ax++;
+      }
+    }
+    if (indices_vmapped) {
+      // Make a new index array for the vmapped dimension
+      auto vmap_inds =
+          arange(static_cast<ShapeElem>(0), src.shape(axes[0]), stream());
+      // Reshape it so it broadcasts with other index arrays
+      {
+        auto shape = Shape(idx_dims, 1);
+        shape[out_ax] = vmap_inds.size();
+        vmap_inds = reshape(vmap_inds, std::move(shape), stream());
+      }
+      // Update gather axes and slice sizes accordingly
+      slice_sizes.insert(slice_sizes.begin() + axes[0], 1);
+      gather_axes.push_back(axes[0]);
+      indices.push_back(vmap_inds);
+    } else {
+      slice_sizes.insert(slice_sizes.begin() + out_ax, src.shape(out_ax));
+      out_ax += idx_dims;
+    }
+  }
+  auto out = gather(src, indices, gather_axes, slice_sizes, stream());
+  if (src_vmapped && indices_vmapped) {
+    out = squeeze(out, idx_dims + axes[0], stream());
+  }
+  return {{out}, {out_ax}};
+}
+
+std::vector<array> Gather::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (int argnum : argnums) {
+    if (argnum > 0) {
+      // Grads w.r.t. indices are zero
+      vjps.push_back(
+          zeros(primals[argnum].shape(), primals[argnum].dtype(), stream()));
+    } else {
+      auto src = zeros_like(primals[0], stream());
+      std::vector<array> inds(primals.begin() + 1, primals.end());
+      vjps.push_back(scatter_add(src, inds, cotangents[0], axes_, stream()));
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> Gather::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  if (argnums.size() > 1 || argnums[0] != 0) {
+    throw std::invalid_argument(
+        "[gather] Cannot calculate JVP with respect to indices.");
+  }
+  std::vector<array> inds(primals.begin() + 1, primals.end());
+  return {gather(tangents[0], inds, axes_, slice_sizes_, stream())};
+}
+
+bool Gather::is_equivalent(const Primitive& other) const {
+  const Gather& g_other = static_cast<const Gather&>(other);
+  return axes_ == g_other.axes_ && slice_sizes_ == g_other.slice_sizes_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> GatherAxis::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  bool vmap_in = axes[0] >= 0;
+  bool vmap_idx = axes[1] >= 0;
+
+  auto in = inputs[0];
+  auto idx = inputs[1];
+  int out_ax;
+  if (vmap_in && vmap_idx) {
+    // reorder the vmap axes to the same location
+    idx = moveaxis(idx, axes[1], axes[0], stream());
+    out_ax = axes[0];
+  } else if (vmap_in) {
+    // expand just the indices dimension
+    idx = expand_dims(idx, axes[0], stream());
+    out_ax = axes[0];
+  } else if (vmap_idx) {
+    // expand just the input dimension
+    in = expand_dims(in, axes[1], stream());
+    out_ax = axes[1];
+  } else {
+    out_ax = -1;
+  }
+  int axis = (out_ax >= 0 && axis_ >= out_ax) ? axis_ + 1 : axis_;
+  return {{take_along_axis(in, idx, axis, stream())}, {out_ax}};
+}
+
+std::vector<array> GatherAxis::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (int argnum : argnums) {
+    if (argnum > 0) {
+      // Grads w.r.t. indices are zero
+      vjps.push_back(
+          zeros(primals[argnum].shape(), primals[argnum].dtype(), stream()));
+    } else {
+      auto src = zeros_like(primals[0], stream());
+      vjps.push_back(array(
+          src.shape(),
+          src.dtype(),
+          std::make_shared<ScatterAxis>(stream(), ScatterAxis::Sum, axis_),
+          {src, primals[1], cotangents[0]}));
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> GatherAxis::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  if (argnums.size() > 1 || argnums[0] != 0) {
+    throw std::invalid_argument(
+        "[gather_axis] Cannot calculate JVP with respect to indices.");
+  }
+  return {take_along_axis(tangents[0], primals[1], axis_, stream())};
+}
+
+std::vector<Shape> GatherAxis::output_shapes(const std::vector<array>& inputs) {
+  return {inputs[1].shape()};
+}
+
+bool GatherAxis::is_equivalent(const Primitive& other) const {
+  auto& g_other = static_cast<const GatherAxis&>(other);
+  return axis_ == g_other.axis_;
+}
+
+std::vector<Shape> Gather::output_shapes(const std::vector<array>& inputs) {
+  Shape out_shape;
+  if (inputs.size() > 1) {
+    out_shape = inputs[1].shape();
+  }
+  out_shape.insert(out_shape.end(), slice_sizes_.begin(), slice_sizes_.end());
+  return {std::move(out_shape)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Greater::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{greater(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Greater::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(zeros_like(primals[arg], stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> Greater::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape());
+  return {zeros(shape, bool_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> GreaterEqual::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{greater_equal(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> GreaterEqual::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(zeros_like(primals[arg], stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> GreaterEqual::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape());
+  return {zeros(shape, bool_, stream())};
+}
+
+std::vector<array> Imag::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(
+      array(complex64_t{0.0f, 1.0f}, primals[0].dtype()),
+      cotangents[0],
+      stream())};
+}
+
+std::vector<array> Imag::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {imag(tangents[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Imag::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{imag(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Less::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{less(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Less::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(zeros_like(primals[arg], stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> Less::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape());
+  return {zeros(shape, bool_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> LessEqual::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{less_equal(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> LessEqual::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(zeros_like(primals[arg], stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> LessEqual::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape());
+  return {zeros(shape, bool_, stream())};
+}
+
+std::vector<array> Log::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Log::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  auto out = divide(tangents[0], primals[0], stream());
+  if (base_ != Base::e) {
+    auto scale = 1 / std::log(base_ == Base::ten ? 10.0f : 2.0f);
+    out = multiply(array(scale, out.dtype()), out, stream());
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Log::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  auto& in = inputs[0];
+  return {
+      {array(
+          in.shape(),
+          in.dtype(),
+          std::make_shared<Log>(stream(), base_),
+          {in})},
+      axes};
+}
+
+std::vector<array> Log1p::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Log1p::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  auto dtype = primals[0].dtype();
+  return {divide(
+      tangents[0], add(array(1.0f, dtype), primals[0], stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Log1p::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{log1p(inputs[0], stream())}, axes};
+}
+
+std::vector<array> LogicalNot::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> LogicalNot::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {zeros_like(tangents[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> LogicalNot::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{logical_not(inputs[0], stream())}, axes};
+}
+
+std::vector<array> LogicalAnd::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 2);
+  std::vector<array> vjps = {zeros_like(cotangents[0], stream())};
+  if (argnums.size() > 1) {
+    vjps.push_back(vjps.back());
+  }
+  return vjps;
+}
+
+std::vector<array> LogicalAnd::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 2);
+  assert(argnums.size() <= 2);
+  return {zeros_like(primals[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> LogicalAnd::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 2);
+  assert(axes.size() == 2);
+
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{logical_and(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> LogicalOr::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 2);
+  std::vector<array> vjps = {zeros_like(cotangents[0], stream())};
+  if (argnums.size() > 1) {
+    vjps.push_back(vjps.back());
+  }
+  return vjps;
+}
+
+std::vector<array> LogicalOr::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 2);
+  assert(argnums.size() <= 2);
+
+  return {zeros_like(primals[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> LogicalOr::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 2);
+  assert(axes.size() == 2);
+
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{logical_or(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> LogAddExp::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto a = primals[0];
+  auto b = primals[1];
+  auto s = sigmoid(subtract(a, b, stream()), stream());
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(multiply(
+        cotangents[0],
+        arg == 0 ? s : subtract(array(1.0f, s.dtype()), s, stream()),
+        stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> LogAddExp::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto a = primals[0];
+  auto b = primals[1];
+  auto s = sigmoid(subtract(a, b, stream()), stream());
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+    return multiply(
+        tangents[i],
+        arg == 0 ? s : subtract(array(1.0f, s.dtype()), s, stream()),
+        stream());
+  };
+  auto out = jvp_fun(0);
+  if (argnums.size() > 1) {
+    out = add(out, jvp_fun(1), stream());
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> LogAddExp::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{logaddexp(a, b, stream())}, {to_ax}};
+}
+
+std::pair<std::vector<array>, std::vector<int>> LogSumExp::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0];
+  auto in = inputs[0];
+  if (ax == (in.ndim() - 1)) {
+    in = swapaxes(in, -1, -2, stream());
+    ax = in.ndim() - 2;
+  }
+  return {{logsumexp(in, -1, true, stream())}, {ax}};
+}
+
+std::vector<array> LogSumExp::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(cotangents.size() == 1);
+  return {multiply(
+      cotangents[0],
+      softmax(primals[0], std::vector<int>{-1}, true, stream()),
+      stream())};
+}
+
+std::vector<array> LogSumExp::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(tangents.size() == 1);
+  return {multiply(
+      tangents[0],
+      softmax(primals[0], std::vector<int>{-1}, true, stream()),
+      stream())};
+}
+
+std::vector<Shape> LogSumExp::output_shapes(const std::vector<array>& inputs) {
+  auto s = inputs[0].shape();
+  s.back() = 1;
+  return {s};
+}
+
+std::vector<array> Matmul::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  auto& cotan = cotangents[0];
+  std::vector<int> reorder(cotan.ndim());
+  std::iota(reorder.begin(), reorder.end(), 0);
+  std::iter_swap(reorder.end() - 1, reorder.end() - 2);
+  auto& s = stream();
+
+  auto complex_transpose = [&](const array& x) {
+    return transpose(conjugate(x, s), reorder, s);
+  };
+
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      // M X N * (K X N).T -> M X K
+      vjps.push_back(matmul(cotan, complex_transpose(primals[1]), s));
+    } else {
+      // (M X K).T * M X N -> K X N
+      vjps.push_back(matmul(complex_transpose(primals[0]), cotan, s));
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> Matmul::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  std::vector<array> jvp;
+  for (int i = 0; i < argnums.size(); ++i) {
+    auto arg = argnums[i];
+    if (arg == 0 && i == 0) {
+      jvp.push_back(matmul(tangents[0], primals[1], stream()));
+    } else if (arg == 0 && i == 1) {
+      jvp[0] = addmm(jvp[0], tangents[1], primals[1], 1.0f, 1.0f, stream());
+    } else if (i == 0) {
+      jvp.push_back(matmul(primals[0], tangents[0], stream()));
+    } else if (i == 1) {
+      jvp[0] = addmm(jvp[0], primals[0], tangents[1], 1.0f, 1.0f, stream());
+    }
+  }
+  return jvp;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Matmul::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto maybe_move_ax = [this](auto& arr, auto ax) {
+    return ax > 0 ? moveaxis(arr, ax, 0, stream()) : arr;
+  };
+  auto a = maybe_move_ax(inputs[0], axes[0]);
+  auto b = maybe_move_ax(inputs[1], axes[1]);
+  return {{matmul(a, b, stream())}, {0}};
+}
+
+std::vector<Shape> Matmul::output_shapes(const std::vector<array>& inputs) {
+  auto out_shape = inputs[0].shape();
+  out_shape.back() = inputs[1].shape(-1);
+  return {std::move(out_shape)};
+}
+
+std::vector<array> Maximum::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto& a = primals[0];
+  auto& b = primals[1];
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    auto mask =
+        (arg == 0) ? greater(a, b, stream()) : less_equal(a, b, stream());
+    vjps.push_back(multiply(cotangents[0], mask, stream()));
+  }
+  return {vjps};
+}
+
+std::vector<array> Maximum::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto& a = primals[0];
+  auto& b = primals[1];
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+    auto mask =
+        (arg == 0) ? greater(a, b, stream()) : less_equal(a, b, stream());
+    return multiply(tangents[i], mask, stream());
+  };
+  auto out = jvp_fun(0);
+  if (argnums.size() > 1) {
+    out = add(out, jvp_fun(1), stream());
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Maximum::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{maximum(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Minimum::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto& a = primals[0];
+  auto& b = primals[1];
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    auto mask =
+        (arg == 0) ? less(a, b, stream()) : greater_equal(a, b, stream());
+    vjps.push_back(multiply(cotangents[0], mask, stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> Minimum::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto& a = primals[0];
+  auto& b = primals[1];
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+    auto mask =
+        (arg == 0) ? less(a, b, stream()) : greater_equal(a, b, stream());
+    return multiply(tangents[i], mask, stream());
+  };
+  auto out = jvp_fun(0);
+  if (argnums.size() > 1) {
+    out = add(out, jvp_fun(1), stream());
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Minimum::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{minimum(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Multiply::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto arg = argnums[0];
+  auto jvp = multiply(tangents[0], primals[1 - arg], stream());
+  if (argnums.size() > 1) {
+    arg = argnums[1];
+    jvp = add(jvp, multiply(tangents[1], primals[1 - arg], stream()), stream());
+  }
+  return {jvp};
+}
+
+std::vector<array> Multiply::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(multiply(
+        conjugate(primals[1 - arg], stream()), cotangents[0], stream()));
+  }
+  return vjps;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Multiply::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{multiply(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Select::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 3);
+  assert(tangents.size() == 3);
+
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+
+    if (arg == 0) {
+      return zeros_like(primals[0], stream());
+    } else if (arg == 1) {
+      return multiply(
+          astype(primals[0], tangents[1].dtype(), stream()),
+          tangents[1],
+          stream());
+    } else {
+      return multiply(
+          astype(
+              logical_not(primals[0], stream()), tangents[2].dtype(), stream()),
+          tangents[2],
+          stream());
+    }
+  };
+
+  array jvp = jvp_fun(argnums[0]);
+  for (int i = 1; i < argnums.size(); i++) {
+    jvp = add(jvp, jvp_fun(argnums[i]));
+  }
+  return {jvp};
+}
+
+std::vector<array> Select::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 3);
+  assert(cotangents.size() == 1);
+
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      vjps.push_back(zeros_like(primals[0], stream()));
+    } else if (arg == 1) {
+      vjps.push_back(multiply(
+          astype(primals[0], cotangents[0].dtype(), stream()),
+          cotangents[0],
+          stream()));
+    } else if (arg == 2) {
+      vjps.push_back(multiply(
+          astype(
+              logical_not(primals[0], stream()),
+              cotangents[0].dtype(),
+              stream()),
+          cotangents[0],
+          stream()));
+    }
+  }
+  return vjps;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Select::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, c, to_ax] = vmap_ternary_op(inputs, axes, stream());
+  return {{where(a, b, c, stream())}, {to_ax}};
+}
+
+std::vector<array> Negative::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Negative::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {negative(tangents[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Negative::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{negative(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> NotEqual::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{not_equal(a, b, stream())}, axes};
+}
+
+std::vector<array> NotEqual::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    vjps.push_back(zeros_like(primals[arg], stream()));
+  }
+  return vjps;
+}
+
+std::vector<array> NotEqual::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape());
+  return {zeros(shape, bool_, stream())};
+}
+
+std::vector<array> Pad::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(argnums.size() == 1 && argnums[0] == 0);
+
+  auto& cotan = cotangents[0];
+  Shape start(cotan.ndim(), 0);
+  auto stop = cotan.shape();
+
+  for (auto i : axes_) {
+    start[i] = low_pad_size_[i];
+    stop[i] -= high_pad_size_[i];
+  }
+
+  auto out = slice(cotan, start, stop, stream());
+
+  return {out};
+}
+
+std::vector<array> Pad::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(argnums.size() == 1 && argnums[0] == 0);
+
+  return {
+      pad(tangents[0],
+          axes_,
+          low_pad_size_,
+          high_pad_size_,
+          array(0, tangents[0].dtype()),
+          "constant",
+          stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Pad::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  throw std::runtime_error("Pad vmap is NYI.");
+}
+
+bool Pad::is_equivalent(const Primitive& other) const {
+  const Pad& p_other = static_cast<const Pad&>(other);
+  return (
+      p_other.axes_ == axes_ && p_other.low_pad_size_ == low_pad_size_ &&
+      p_other.high_pad_size_ == high_pad_size_);
+}
+
+std::vector<array> Partition::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto sort_idx = argpartition(primals[0], kth_, axis_, stream());
+  return {put_along_axis(
+      zeros_like(primals[0], stream()),
+      sort_idx,
+      cotangents[0],
+      axis_,
+      stream())};
+}
+
+std::vector<array> Partition::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(tangents.size() == 1);
+  auto sort_idx = argpartition(primals[0], kth_, axis_, stream());
+  auto out = take_along_axis(tangents[0], sort_idx, axis_, stream());
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Partition::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  int axis_left = axes[0] >= 0 && axes[0] <= axis_;
+  return {{partition(inputs[0], axis_ + axis_left, stream())}, axes};
+}
+
+bool Partition::is_equivalent(const Primitive& other) const {
+  const Partition& r_other = static_cast<const Partition&>(other);
+  return axis_ == r_other.axis_ && kth_ == r_other.kth_;
+}
+
+std::vector<array> Power::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      vjps.push_back(multiply(
+          power(
+              primals[0],
+              subtract(primals[1], array(1, primals[0].dtype()), stream()),
+              stream()),
+          primals[1],
+          stream()));
+    } else {
+      auto& exp = outputs[0];
+      auto exp_vjp = multiply(log(primals[0], stream()), outputs[0], stream());
+      // 0 * log 0 -> 0
+      vjps.push_back(where(exp, exp_vjp, array(0.0f, exp.dtype()), stream()));
+    }
+    vjps.back() = multiply(cotangents[0], vjps.back(), stream());
+  }
+  return vjps;
+}
+
+std::vector<array> Power::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto output = power(primals[0], primals[1], stream());
+  auto grads = vjp(primals, tangents, argnums, {output});
+  if (argnums.size() > 1) {
+    return {add(grads[0], grads[1], stream())};
+  } else {
+    return grads;
+  }
+}
+
+std::pair<std::vector<array>, std::vector<int>> Power::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{power(a, b, stream())}, {to_ax}};
+}
+
+std::string quantization_mode_to_string(QuantizationMode mode) {
+  switch (mode) {
+    case QuantizationMode::Affine:
+      return "affine";
+    case QuantizationMode::Mxfp4:
+      return "mxfp4";
+    case QuantizationMode::Mxfp8:
+      return "mxfp8";
+    case QuantizationMode::Nvfp4:
+    default:
+      return "nvfp4";
+  }
+}
+
+QuantizationMode string_to_quantization_mode(
+    const std::string& mode,
+    std::string_view tag /* = "" */) {
+  if (mode == "affine") {
+    return QuantizationMode::Affine;
+  } else if (mode == "mxfp4") {
+    return QuantizationMode::Mxfp4;
+  } else if (mode == "mxfp8") {
+    return QuantizationMode::Mxfp8;
+  } else if (mode == "nvfp4") {
+    return QuantizationMode::Nvfp4;
+  }
+  std::string msg;
+  if (!tag.empty()) {
+    msg += "[" + std::string(tag) + "]";
+  }
+  msg += " Invalid quantization mode '" + mode + "'.";
+  throw std::invalid_argument(msg);
+}
+
+std::pair<std::vector<array>, std::vector<int>> QuantizedMatmul::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  throw std::runtime_error("[QuantizedMatmul::vmap] NYI");
+}
+
+std::vector<array> QuantizedMatmul::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+
+  // We rely on the fact that w is always 2D so transpose is simple
+  std::optional<array> dsb = std::nullopt;
+  for (auto arg : argnums) {
+    // gradient wrt to x
+    if (arg == 0) {
+      vjps.push_back(quantized_matmul(
+          cotangents[0],
+          primals[1],
+          primals[2],
+          mode_ == QuantizationMode::Affine ? std::optional<array>(primals[3])
+                                            : std::nullopt,
+          !transpose_,
+          group_size_,
+          bits_,
+          quantization_mode_to_string(mode_),
+          stream()));
+    }
+
+    // gradient wrt to w_q, scales or biases
+    else if (arg == 1) {
+      throw std::runtime_error(
+          "[QuantizedMatmul::vjp] no gradient wrt the quantized weights.");
+    } else {
+      if (mode_ != QuantizationMode::Affine) {
+        std::ostringstream msg;
+        msg << "[QuantizedMatmul::vjp] no gradient wrt scales in "
+            << quantization_mode_to_string(mode_) << " quantization.";
+        throw std::invalid_argument(msg.str());
+      }
+      if (!dsb) {
+        int ndim = primals[1].ndim();
+        auto fc = flatten(cotangents[0], 0, -ndim, stream());
+        auto fx = flatten(primals[0], 0, -ndim, stream());
+        auto dw = transpose_
+            ? matmul(swapaxes(fc, -1, -2, stream()), fx, stream())
+            : matmul(swapaxes(fx, -1, -2, stream()), fc, stream());
+        dsb = unflatten(dw, -1, {-1, group_size_}, stream());
+      }
+      if (arg == 3) {
+        // biases
+        vjps.push_back(sum(*dsb, -1, false, stream()));
+      } else {
+        // scales
+        auto wq = dequantize(
+            primals[1],
+            ones_like(primals[2], stream()),
+            zeros_like(primals[3], stream()),
+            group_size_,
+            bits_,
+            quantization_mode_to_string(mode_),
+            std::nullopt,
+            stream());
+        wq = unflatten(wq, -1, {-1, group_size_}, stream());
+        vjps.push_back(sum(multiply(*dsb, wq, stream()), -1, false, stream()));
+      }
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> QuantizedMatmul::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  if (argnums.size() > 1 || argnums[0] != 0) {
+    throw std::runtime_error(
+        "[QuantizedMatmul::jvp] No JVP wrt the quantized matrix yet.");
+  }
+  return {quantized_matmul(
+      tangents[0],
+      primals[1],
+      primals[2],
+      mode_ == QuantizationMode::Affine ? std::optional<array>(primals[3])
+                                        : std::nullopt,
+      transpose_,
+      group_size_,
+      bits_,
+      quantization_mode_to_string(mode_),
+      stream())};
+}
+
+bool QuantizedMatmul::is_equivalent(const Primitive& other) const {
+  const QuantizedMatmul& qm_other = static_cast<const QuantizedMatmul&>(other);
+  return group_size_ == qm_other.group_size_ && bits_ == qm_other.bits_ &&
+      mode_ == qm_other.mode_ && transpose_ == qm_other.transpose_;
+}
+
+std::vector<Shape> QuantizedMatmul::output_shapes(
+    const std::vector<array>& inputs) {
+  auto& w = inputs[1];
+  int w_outer_dims = (transpose_) ? w.shape(-2) : w.shape(-1) * 32 / bits_;
+  auto out_shape = inputs[0].shape();
+  out_shape.back() = w_outer_dims;
+  return {std::move(out_shape)};
+}
+
+bool QQMatmul::is_equivalent(const Primitive& other) const {
+  const QQMatmul& qm_other = static_cast<const QQMatmul&>(other);
+  return group_size_ == qm_other.group_size_ && bits_ == qm_other.bits_ &&
+      mode_ == qm_other.mode_;
+}
+
+std::vector<Shape> QQMatmul::output_shapes(const std::vector<array>& inputs) {
+  auto out_shape = inputs[0].shape();
+  int w_outer_dims = inputs[1].shape(-2);
+  out_shape.back() = w_outer_dims;
+  return {std::move(out_shape)};
+}
+
+std::vector<array> QQMatmul::vjp(
+    const std::vector<array>& primals, // non quantized x, non quantized w
+    const std::vector<array>& cotangents, // non quantized upstream grads
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  if (primals.size() != 2) {
+    throw std::runtime_error(
+        "[QQMatmul::vjp] Expected exactly 2 non-quantized primal inputs (x, w).");
+  }
+  std::vector<array> vjps;
+  auto& cotan = cotangents[0];
+  auto& s = stream();
+  // primal[1] -- non quantized w (N, K)
+  // primal[0] -- non quantized activations (M, K)
+  // cotan -- non quantized grads (M, N)
+  auto qmode = quantization_mode_to_string(mode_);
+  for (auto arg : argnums) {
+    if (arg == 0) { // gradient wrt to x
+      // We transpose weights -> quantize along N
+      vjps.push_back(qqmm(
+          cotan, //  M X N
+          swapaxes(primals[1], -1, -2, s), // assuming that w is 2D
+          {},
+          group_size_,
+          bits_,
+          qmode,
+          s));
+    } else if (arg == 1) { // gradient wrt to weights
+      vjps.push_back(qqmm(
+          swapaxes(cotan, -1, -2, s), // (N, M)
+          swapaxes(primals[0], -1, -2, s), // (K, M)
+          {},
+          group_size_,
+          bits_,
+          qmode,
+          s));
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> QQMatmul::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  throw std::runtime_error("QQMM::jvp NYI");
+}
+
+std::pair<std::vector<array>, std::vector<int>> GatherQMM::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  throw std::runtime_error("GatherQMM::vmap NYI");
+}
+
+std::vector<array> GatherQMM::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+
+  auto& cotan = cotangents[0];
+
+  auto& x = primals[0];
+  auto& w = primals[1];
+  auto& scales = primals[2];
+  auto& lhs_indices = primals[primals.size() - 2];
+  auto& rhs_indices = primals[primals.size() - 1];
+  auto biases = (mode_ == QuantizationMode::Affine)
+      ? std::optional<array>(primals[3])
+      : std::nullopt;
+
+  int M = cotan.shape(-2);
+  int K = x.shape(-1);
+
+  bool sorted = left_sorted_ || right_sorted_;
+  bool no_broadcast = rhs_indices.size() * M * K == x.size();
+  std::optional<array> dsb = std::nullopt;
+
+  for (auto arg : argnums) {
+    // gradient wrt to x
+    if (arg == 0) {
+      auto g = gather_qmm(
+          cotan,
+          w,
+          scales,
+          biases,
+          std::nullopt,
+          rhs_indices,
+          !transpose_,
+          group_size_,
+          bits_,
+          quantization_mode_to_string(mode_),
+          sorted,
+          stream());
+      if (sorted && no_broadcast) {
+        vjps.push_back(g);
+      } else {
+        vjps.push_back(reshape(
+            scatter_add(
+                flatten(zeros_like(x, stream()), 0, -3, stream()),
+                lhs_indices,
+                expand_dims(g, -3, stream()),
+                0,
+                stream()),
+            x.shape(),
+            stream()));
+      }
+    }
+
+    // gradient wrt to the indices is undefined
+    else if (arg > 3) {
+      throw std::runtime_error(
+          "[GatherQMM::vjp] cannot compute the gradient wrt the indices.");
+    }
+
+    // gradient wrt to w_q, scales or biases
+    else if (arg == 1) {
+      throw std::runtime_error(
+          "[GatherQMM::vjp] no gradient wrt the quantized weights.");
+    } else {
+      if (mode_ != QuantizationMode::Affine) {
+        std::ostringstream msg;
+        msg << "[GatherQMM::vjp] no gradient wrt scales in "
+            << quantization_mode_to_string(mode_) << " quantization.";
+        throw std::invalid_argument(msg.str());
+      }
+
+      if (!dsb) {
+        auto shape = w.shape();
+        shape.pop_back();
+        shape.pop_back();
+        dsb = unflatten(
+            gather_mm_grad(
+                x,
+                cotan,
+                lhs_indices,
+                rhs_indices,
+                sorted,
+                std::move(shape),
+                stream()),
+            -1,
+            {-1, group_size_},
+            stream());
+      }
+      if (arg == 3) {
+        vjps.push_back(sum(*dsb, -1, false, stream()));
+      } else {
+        vjps.push_back(
+            sum(multiply(
+                    *dsb,
+                    unflatten(
+                        dequantize(
+                            w,
+                            ones_like(scales, stream()),
+                            zeros_like(*biases, stream()),
+                            group_size_,
+                            bits_,
+                            quantization_mode_to_string(mode_),
+                            std::nullopt,
+                            stream()),
+                        -1,
+                        {-1, group_size_},
+                        stream()),
+                    stream()),
+                -1,
+                false,
+                stream()));
+      }
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> GatherQMM::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  throw std::runtime_error("GatherQMM::jvp NYI");
+}
+
+bool GatherQMM::is_equivalent(const Primitive& other) const {
+  const GatherQMM& qm_other = static_cast<const GatherQMM&>(other);
+  return group_size_ == qm_other.group_size_ && bits_ == qm_other.bits_ &&
+      mode_ == qm_other.mode_ && transpose_ == qm_other.transpose_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> RandomBits::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  // The last dimension of the key is always a key pair
+  auto key = inputs[0];
+  auto kax = axes[0];
+  if (kax == key.ndim() - 1) {
+    std::vector<int> reorder(key.ndim());
+    std::iota(reorder.begin(), reorder.end(), 0);
+    std::swap(reorder[kax], reorder[kax - 1]);
+    key = transpose(key, reorder, stream());
+    kax--;
+  }
+
+  auto shape = shape_;
+  if (kax >= 0) {
+    shape.insert(shape.begin() + kax, key.shape()[kax]);
+  }
+
+  auto get_dtype = [width = width_]() {
+    switch (width) {
+      case 1:
+        return uint8;
+      case 2:
+        return uint16;
+      default:
+        return uint32;
+    }
+  };
+
+  auto out = array(
+      shape,
+      get_dtype(),
+      std::make_shared<RandomBits>(stream(), shape, width_),
+      {key});
+  return {{out}, {kax}};
+}
+
+bool RandomBits::is_equivalent(const Primitive& other) const {
+  const RandomBits& r_other = static_cast<const RandomBits&>(other);
+  return shape_ == r_other.shape_ && width_ == r_other.width_;
+}
+
+std::vector<array> Real::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {astype(cotangents[0], primals[0].dtype(), stream())};
+}
+
+std::vector<array> Real::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {real(tangents[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Real::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{real(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Reshape::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  // Transpose the input so that the vmap dim is first.
+  auto& in = inputs[0];
+  auto ax = axes[0];
+  if (ax >= 0) {
+    std::vector<int> reorder(in.ndim());
+    std::iota(reorder.begin(), reorder.end(), 0);
+    reorder.erase(reorder.begin() + ax);
+    reorder.insert(reorder.begin(), ax);
+    // Insert the vmap dim into the shape at the beginning.
+    auto out = transpose(in, reorder, stream());
+    shape_.insert(shape_.begin(), in.shape()[ax]);
+    // Reshape the transposed input to the new shape.
+    return {{reshape(out, shape_, stream())}, {0}};
+  } else {
+    return {{reshape(in, shape_, stream())}, {ax}};
+  }
+}
+
+std::vector<array> Reshape::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  assert(argnums[0] == 0);
+  return {reshape(cotangents[0], primals[0].shape(), stream())};
+}
+
+std::vector<array> Reshape::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  assert(argnums[0] == 0);
+  return {reshape(tangents[0], shape_, stream())};
+}
+
+bool Reshape::is_equivalent(const Primitive& other) const {
+  const Reshape& r_other = static_cast<const Reshape&>(other);
+  return shape_ == r_other.shape_;
+}
+
+Shape Reshape::output_shape(const array& input, Shape shape) {
+  size_t size = 1;
+  int infer_idx = -1;
+  for (int i = 0; i < shape.size(); ++i) {
+    if (shape[i] == -1) {
+      if (infer_idx >= 0) {
+        throw std::invalid_argument(
+            "[reshape] Reshape can only infer one dimension.");
+      }
+      infer_idx = i;
+    } else {
+      size *= shape[i];
+    }
+  }
+
+  // Infer the shape
+  if (size > 0 && infer_idx >= 0) {
+    shape[infer_idx] = input.size() / size;
+    size *= shape[infer_idx];
+  } else if (infer_idx >= 0) {
+    throw std::invalid_argument(
+        "[reshape] Cannot infer the shape of an empty array");
+  }
+
+  // Check that the reshaping is valid
+  if (input.size() != size) {
+    std::ostringstream msg;
+    msg << "[reshape] Cannot reshape array of size " << input.size()
+        << " into shape " << shape << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  return shape;
+}
+
+std::vector<Shape> Reshape::output_shapes(const std::vector<array>& inputs) {
+  return {output_shape(inputs[0], shape_)};
+}
+
+std::vector<array> Reduce::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  auto in = primals[0];
+
+  auto& cotan = cotangents[0];
+  if (reduce_type_ == Reduce::Sum) {
+    return {broadcast_arrays({cotan, in}, stream())[0]};
+  } else if (reduce_type_ == Reduce::Prod) {
+    auto s = stream();
+    auto prod_grad_single_axis =
+        [&s](const array& x, const array& cotan, int axis) {
+          auto p1 = cumprod(x, axis, /*reverse=*/false, /*inclusive=*/false, s);
+          auto p2 = cumprod(x, axis, /*reverse=*/true, /*inclusive=*/false, s);
+          auto exclusive_prod = multiply(p1, p2, s);
+          return multiply(exclusive_prod, cotan, s);
+        };
+
+    // To compute a numerically stable gradient for prod we need an exclusive
+    // product of all elements in axes_ . To achieve that we move axes_ to the
+    // last dim and perform two exclusive cumprods. Afterwards we move
+    // everything back to the original axes.
+    if (axes_.size() > 1) {
+      std::vector<int> transpose_to;
+      std::vector<int> transpose_back;
+      Shape shape_flat;
+      {
+        // Find the transpose needed to move axes_ to the back and the shape
+        // except the reduced over axes.
+        int j = 0;
+        for (int i = 0; i < in.ndim(); i++) {
+          if (j < axes_.size() && axes_[j] == i) {
+            j++;
+          } else {
+            transpose_to.push_back(i);
+            shape_flat.push_back(in.shape(i));
+          }
+        }
+        for (auto ax : axes_) {
+          transpose_to.push_back(ax);
+        }
+        shape_flat.push_back(-1);
+        transpose_back.resize(transpose_to.size());
+        for (int i = 0; i < transpose_to.size(); i++) {
+          transpose_back[transpose_to[i]] = i;
+        }
+      }
+
+      // Move axes to the back
+      auto x = transpose(in, transpose_to, s);
+      // Keep the shape in order to reshape back to the original
+      auto shape_to = x.shape();
+
+      // Flatten and compute the gradient
+      x = reshape(x, shape_flat, stream());
+      auto grad = prod_grad_single_axis(x, reshape(cotan, shape_flat, s), -1);
+
+      // Reshape and transpose to the original shape
+      grad = reshape(grad, shape_to, s);
+      grad = transpose(grad, transpose_back, s);
+
+      return {grad};
+    } else {
+      return {prod_grad_single_axis(in, cotan, axes_[0])};
+    }
+
+  } else if (reduce_type_ == Reduce::Min || reduce_type_ == Reduce::Max) {
+    auto out = outputs[0];
+    if (out.ndim() != in.ndim()) {
+      out = expand_dims(out, axes_, stream());
+    }
+    auto mask = equal(in, out, stream());
+    auto normalizer = sum(mask, axes_, true, stream());
+    return {multiply(divide(cotan, normalizer, stream()), mask, stream())};
+  }
+
+  else {
+    return {zeros_like(in, stream())};
+  }
+}
+
+std::vector<array> Reduce::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto in = primals[0];
+  auto s = stream();
+
+  auto grad_op = [&s, reduce_type = reduce_type_](
+                     const array& x, const array& tan, int axis) {
+    if (reduce_type == Reduce::Min) {
+      auto idx = argmin(x, axis, true, s);
+      return take_along_axis(tan, idx, axis, s);
+    } else if (reduce_type == Reduce::Max) {
+      auto idx = argmax(x, axis, true, s);
+      return take_along_axis(tan, idx, axis, s);
+    } else {
+      auto p1 = cumprod(x, axis, /*reverse=*/false, /*inclusive=*/false, s);
+      auto p2 = cumprod(x, axis, /*reverse=*/true, /*inclusive=*/false, s);
+      auto out = multiply(multiply(p1, p2, s), tan, s);
+      return sum(out, axis, true, s);
+    }
+  };
+
+  auto tan = tangents[0];
+  if (reduce_type_ == Reduce::Sum) {
+    return {sum(tan, axes_, true, s)};
+  } else {
+    if (axes_.size() > 1) {
+      std::vector<int> transpose_to;
+      {
+        // Find the transpose needed to move axes_ to the back.
+        int j = 0;
+        for (int i = 0; i < in.ndim(); i++) {
+          if (j < axes_.size() && axes_[j] == i) {
+            j++;
+          } else {
+            transpose_to.push_back(i);
+          }
+        }
+        for (auto ax : axes_) {
+          transpose_to.push_back(ax);
+        }
+      }
+
+      int start_ax = in.ndim() - axes_.size();
+      in = flatten(transpose(in, transpose_to, s), start_ax, -1, s);
+      tan = flatten(transpose(tan, transpose_to, s), start_ax, -1, s);
+
+      auto grad = squeeze(grad_op(in, tan, -1), -1, s);
+      return {expand_dims(grad, axes_, s)};
+    } else {
+      return {grad_op(in, tan, axes_[0])};
+    }
+  }
+}
+
+std::pair<std::vector<array>, std::vector<int>> Reduce::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0];
+  auto reduce_axes = axes_;
+  if (ax >= 0) {
+    for (auto& rax : reduce_axes) {
+      if (rax >= ax) {
+        rax++;
+      }
+    }
+  }
+  auto& in = inputs[0];
+  std::vector<array> out;
+  switch (reduce_type_) {
+    case Reduce::And:
+      out.push_back(all(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Or:
+      out.push_back(any(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Sum:
+      out.push_back(sum(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Prod:
+      out.push_back(prod(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Min:
+      out.push_back(min(in, reduce_axes, true, stream()));
+      break;
+    case Reduce::Max:
+      out.push_back(max(in, reduce_axes, true, stream()));
+      break;
+  }
+  return {out, axes};
+}
+
+bool Reduce::is_equivalent(const Primitive& other) const {
+  const Reduce& r_other = static_cast<const Reduce&>(other);
+  return reduce_type_ == r_other.reduce_type_ && axes_ == r_other.axes_;
+}
+
+std::vector<Shape> Reduce::output_shapes(const std::vector<array>& inputs) {
+  auto out_shape = inputs[0].shape();
+  for (auto i : axes_) {
+    out_shape[i] = 1;
+  }
+  return {std::move(out_shape)};
+}
+
+std::vector<array> Round::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Round::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {zeros_like(primals[0], stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Round::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{round(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Scan::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto& in = inputs[0];
+  auto out_dtype =
+      (in.dtype() == bool_ && reduce_type_ == Scan::Sum) ? int32 : in.dtype();
+  int axis_left = axes[0] >= 0 && axes[0] <= axis_;
+  return {
+      {array(
+          in.shape(),
+          out_dtype,
+          std::make_shared<Scan>(
+              stream(), reduce_type_, axis_ + axis_left, reverse_, inclusive_),
+          {in})},
+      axes};
+}
+
+std::vector<array> Scan::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  assert(primals.size() == 1);
+  assert(argnums[0] == 0);
+
+  if (reduce_type_ == Scan::Sum) {
+    return {cumsum(cotangents[0], axis_, !reverse_, inclusive_, stream())};
+  } else if (reduce_type_ == Scan::LogAddExp) {
+    // Ref:
+    // https://github.com/tensorflow/tensorflow/blob/2a5910906a0e0f3dbc186ff9db6386d81a63448c/tensorflow/python/ops/math_grad.py#L1832-L1863
+
+    auto x = primals[0];
+    auto grad = cotangents[0];
+    auto results = outputs[0];
+
+    auto zero = zeros({1}, grad.dtype(), stream());
+    auto grad_min = array(finfo(grad.dtype()).min, grad.dtype());
+
+    // Split the incoming gradient into positive and negative part
+    // in order to take logs. This is required for stable results.
+    auto log_abs_grad = log(abs(grad, stream()), stream());
+    auto log_grad_positive =
+        where(greater(grad, zero, stream()), log_abs_grad, grad_min, stream());
+    auto log_grad_negative =
+        where(less(grad, zero, stream()), log_abs_grad, grad_min, stream());
+
+    auto output_pos = exp(
+        add(logcumsumexp(
+                subtract(log_grad_positive, results, stream()),
+                axis_,
+                !reverse_,
+                inclusive_,
+                stream()),
+            x,
+            stream()));
+    auto output_neg = exp(
+        add(logcumsumexp(
+                subtract(log_grad_negative, results, stream()),
+                axis_,
+                !reverse_,
+                inclusive_,
+                stream()),
+            x,
+            stream()));
+
+    return {subtract(output_pos, output_neg, stream())};
+  } else if (reduce_type_ == Scan::Prod) {
+    auto in = primals[0];
+    // Find the location of the first 0 and set it to 1:
+    // - A: Exclusive cumprod
+    // - B: Inclusive cumprod
+    // - Find the location that is 0 in A and not zero B
+    // Compute the gradient by:
+    // - Compute the regular gradient for everything before the first zero
+    // - Set the first zero to 1 and redo the computation, use this for the
+    //   gradient of the first zero
+    // - Everything after the first zero has a gradient of 0
+
+    // Get inclusive and exclusive cum prods
+    auto cprod_exclusive = cumprod(in, axis_, reverse_, !inclusive_, stream());
+    auto cprod_inclusive = outputs[0];
+    if (!inclusive_) {
+      std::swap(cprod_exclusive, cprod_inclusive);
+    }
+
+    // Make the mask for the first zero
+    auto z = array(0, in.dtype());
+    auto eq_zero = equal(cprod_inclusive, z, stream());
+    auto first_zero =
+        logical_and(eq_zero, not_equal(cprod_exclusive, z, stream()), stream());
+
+    auto to_partial_grad = [this, &cotangents](const array& arr) {
+      return cumsum(
+          multiply(arr, cotangents[0], stream()),
+          axis_,
+          !reverse_,
+          inclusive_,
+          stream());
+    };
+
+    auto cprod_with_one = cumprod(
+        where(first_zero, array(1, in.dtype()), in, stream()),
+        axis_,
+        reverse_,
+        inclusive_,
+        stream());
+    auto grad_with_one = to_partial_grad(cprod_with_one);
+    auto grad = divide(to_partial_grad(outputs[0]), in, stream());
+    return {where(
+        first_zero,
+        grad_with_one,
+        where(eq_zero, z, grad, stream()),
+        stream())};
+  } else {
+    // Can probably be implemented by equals and then cummax to make the mask
+    throw std::runtime_error("VJP is not implemented for cumulative min/max");
+  }
+}
+
+std::vector<array> Scan::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(tangents.size() == 1);
+  assert(argnums[0] == 0);
+
+  if (reduce_type_ == Scan::Sum) {
+    return {cumsum(tangents[0], axis_, reverse_, inclusive_, stream())};
+  } else {
+    throw std::runtime_error(
+        "JVP is not implemented for cumulative prod/min/max");
+  }
+}
+
+bool Scan::is_equivalent(const Primitive& other) const {
+  const Scan& s_other = static_cast<const Scan&>(other);
+  return (
+      reduce_type_ == s_other.reduce_type_ && axis_ == s_other.axis_ &&
+      reverse_ == s_other.reverse_ && inclusive_ == s_other.inclusive_);
+}
+
+bool Scatter::is_equivalent(const Primitive& other) const {
+  const Scatter& s_other = static_cast<const Scatter&>(other);
+  return reduce_type_ == s_other.reduce_type_ && axes_ == s_other.axes_;
+}
+
+std::vector<array> Scatter::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  switch (reduce_type_) {
+    case Scatter::None:
+    case Scatter::Sum:
+    case Scatter::Max:
+    case Scatter::Min:
+      break;
+    default:
+      throw std::runtime_error(
+          "[scatter] VJP not implemented for scatter_prod");
+  }
+
+  const array& result = outputs[0];
+  const array& values = primals[0];
+  const array& updates = primals.back();
+  const std::vector<array> indices(primals.begin() + 1, primals.end() - 1);
+
+  std::vector<array> vjps;
+  for (auto num : argnums) {
+    // Gradient wrt to the input array
+    if (num == 0) {
+      switch (reduce_type_) {
+        case Scatter::None:
+          // Scatter 0s to the locations that were updated with the updates
+          vjps.push_back(scatter(
+              cotangents[0],
+              indices,
+              zeros_like(updates, stream()),
+              axes_,
+              stream()));
+          break;
+        case Scatter::Sum:
+          // The input array values are kept so they all get gradients
+          vjps.push_back(cotangents[0]);
+          break;
+        case Scatter::Max:
+        case Scatter::Min: {
+          vjps.push_back(where(
+              equal(result, values, stream()),
+              cotangents[0],
+              array(0, cotangents[0].dtype()),
+              stream()));
+          break;
+        }
+        default:
+          // Should never reach here
+          throw std::invalid_argument("");
+      }
+    } else if (num == primals.size() - 1) {
+      switch (reduce_type_) {
+        case Scatter::None:
+        case Scatter::Sum: {
+          // Gather the values from the cotangent
+          auto slice_sizes = cotangents[0].shape();
+          for (auto ax : axes_) {
+            slice_sizes[ax] = 1;
+          }
+          vjps.push_back(
+              gather(cotangents[0], indices, axes_, slice_sizes, stream()));
+          break;
+        }
+        case Scatter::Max:
+        case Scatter::Min: {
+          auto slice_sizes = cotangents[0].shape();
+          for (auto ax : axes_) {
+            slice_sizes[ax] = 1;
+          }
+          auto gathered_cotan =
+              gather(cotangents[0], indices, axes_, slice_sizes, stream());
+          auto gathered_result =
+              gather(result, indices, axes_, slice_sizes, stream());
+          vjps.push_back(
+              multiply(gathered_cotan, gathered_result == updates, stream()));
+          break;
+        }
+        default: {
+          // Should never reach here
+          throw std::invalid_argument("");
+        }
+      }
+    } else {
+      throw std::invalid_argument(
+          "[scatter] Cannot calculate VJP with respect to indices.");
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> Scatter::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  throw std::runtime_error("[scatter] JVP not yet implemented");
+}
+
+std::pair<std::vector<array>, std::vector<int>> Scatter::vmap(
+    const std::vector<array>& inputs_,
+    const std::vector<int>& vmap_axes) {
+  assert(inputs_.size() >= 2);
+  assert(inputs_.size() == vmap_axes.size());
+
+  auto inputs = inputs_;
+
+  auto scatter_axes = axes_;
+  int src_ax = vmap_axes[0];
+
+  auto vmap_ax_it = std::find_if(
+      vmap_axes.begin(), vmap_axes.end(), [](int a) { return a >= 0; });
+  auto vmap_ax = *vmap_ax_it;
+  if (vmap_ax >= 0) {
+    auto vmap_size = inputs[vmap_ax_it - vmap_axes.begin()].shape(vmap_ax);
+    if (src_ax < 0) {
+      src_ax = 0;
+      inputs[0] =
+          repeat(expand_dims(inputs[0], 0, stream()), vmap_size, 0, stream());
+    }
+    for (int i = 1; i < vmap_axes.size() - 1; ++i) {
+      // vmap axis for indices goes to 0
+      if (vmap_axes[i] >= 0) {
+        inputs[i] = moveaxis(inputs[i], vmap_axes[i], 0, stream());
+      }
+      // insert a vmap axis and repeat
+      if (vmap_axes[i] < 0) {
+        auto idx_shape = inputs[i].shape();
+        inputs[i] =
+            repeat(expand_dims(inputs[i], 0, stream()), vmap_size, 0, stream());
+      }
+      // Adjust non-vmapped index axes to account for the extra vmap dimension.
+      if (scatter_axes[i - 1] >= src_ax) {
+        scatter_axes[i - 1]++;
+      }
+    }
+
+    auto vmap_inds = arange(vmap_size, inputs[1].dtype(), stream());
+    auto vmap_inds_shape = Shape(inputs[1].ndim(), 1);
+    vmap_inds_shape[0] = vmap_inds.size();
+    vmap_inds = reshape(vmap_inds, std::move(vmap_inds_shape), stream());
+    inputs.insert(
+        inputs.end() - 1, broadcast_to(vmap_inds, inputs[1].shape(), stream()));
+    scatter_axes.push_back(src_ax);
+
+    // Clone updates along the vmap dimension so they can be applied to each
+    // source tensor in the vmap.
+    auto& updates = inputs.back();
+    if (vmap_axes.back() < 0) {
+      updates = expand_dims(
+          updates, {0, static_cast<int>(inputs[1].ndim())}, stream());
+      updates = repeat(updates, vmap_size, 0, stream());
+    } else {
+      updates =
+          expand_dims(updates, static_cast<int>(inputs[1].ndim()), stream());
+      updates = moveaxis(updates, vmap_axes.back(), 0, stream());
+    }
+  }
+
+  auto& shape = inputs[0].shape();
+  auto dtype = inputs[0].dtype();
+  auto out = array(
+      shape,
+      dtype,
+      std::make_shared<Scatter>(stream(), reduce_type_, scatter_axes),
+      std::move(inputs));
+
+  return {{out}, {src_ax}};
+}
+
+std::vector<array> ScatterAxis::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  const auto& indices = primals[1];
+  const auto& updates = primals[2];
+
+  std::vector<array> vjps;
+  for (auto num : argnums) {
+    // Gradient wrt to the input array
+    if (num == 0) {
+      if (reduce_type_ == ScatterAxis::None) {
+        // Scatter 0s to the locations that were updated with the updates
+        vjps.push_back(put_along_axis(
+            cotangents[0],
+            indices,
+            zeros_like(updates, stream()),
+            axis_,
+            stream()));
+      } else {
+        // The input array values are kept so they all get gradients
+        vjps.push_back(cotangents[0]);
+      }
+    } else if (num == 2) {
+      vjps.push_back(take_along_axis(cotangents[0], indices, axis_, stream()));
+    } else {
+      throw std::invalid_argument(
+          "[scatter_axis] Cannot calculate VJP with respect to indices.");
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> ScatterAxis::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  for (auto arg : argnums) {
+    if (arg == 1) {
+      throw std::invalid_argument(
+          "[scatter_axis] Cannot calculate JVP with respect to indices.");
+    }
+  }
+  if (argnums.size() == 2) {
+    return {array(
+        primals[0].shape(),
+        primals[0].dtype(),
+        std::make_shared<ScatterAxis>(stream(), reduce_type_, axis_),
+        {tangents[0], primals[1], tangents[1]})};
+  } else {
+    auto tan_a =
+        argnums[0] == 0 ? tangents[0] : zeros_like(primals[0], stream());
+    auto tan_b =
+        argnums[0] == 2 ? tangents[0] : zeros_like(primals[2], stream());
+    return {array(
+        primals[0].shape(),
+        primals[0].dtype(),
+        std::make_shared<ScatterAxis>(stream(), reduce_type_, axis_),
+        {tan_a, primals[1], tan_b})};
+  }
+}
+
+std::pair<std::vector<array>, std::vector<int>> ScatterAxis::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  // Find the first vmap axis
+  int out_ax = -1;
+  for (auto ax : axes) {
+    if (ax >= 0) {
+      out_ax = ax;
+      break;
+    }
+  }
+
+  if (out_ax < 0) {
+    return {
+        {array(
+            inputs[0].shape(),
+            inputs[0].dtype(),
+            std::make_shared<ScatterAxis>(stream(), reduce_type_, axis_),
+            inputs)},
+        {-1}};
+  }
+
+  auto v_in = inputs;
+  for (int i = 0; i < axes.size(); ++i) {
+    if (axes[i] >= 0) {
+      // if out_ax >= 0 move axis o/w set out_ax
+      if (out_ax != axes[i]) {
+        v_in[i] = moveaxis(v_in[i], axes[i], out_ax, stream());
+      }
+    } else {
+      v_in[i] = expand_dims(v_in[i], out_ax, stream());
+    }
+  }
+  int axis = axis_ >= out_ax ? axis_ + 1 : axis_;
+  auto fn = reduce_type_ == Sum ? scatter_add_axis : put_along_axis;
+  return {{fn(v_in[0], v_in[1], v_in[2], axis, stream())}, {out_ax}};
+}
+
+std::vector<Shape> ScatterAxis::output_shapes(
+    const std::vector<array>& inputs) {
+  return {inputs[0].shape()};
+}
+
+bool ScatterAxis::is_equivalent(const Primitive& other) const {
+  auto& s_other = static_cast<const ScatterAxis&>(other);
+  return reduce_type_ == s_other.reduce_type_ && axis_ == s_other.axis_;
+}
+
+std::vector<array> MaskedScatter::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto& s = stream();
+  const array& dst = primals[0];
+  const array& mask = primals[1];
+  const array& src = primals[2];
+  const array mask_b = broadcast_to(mask, dst.shape(), s);
+  const array& cotan = cotangents[0];
+
+  std::vector<array> vjps;
+  vjps.reserve(argnums.size());
+
+  for (int arg : argnums) {
+    if (arg == 0) {
+      vjps.push_back(where(mask_b, zeros_like(cotan, s), cotan, s));
+    } else if (arg == 2) {
+      const array mask_flat = flatten(mask_b, s);
+      const array cotan_flat = flatten(cotan, s);
+
+      const array idx_src =
+          cumsum(astype(mask_flat, int32, s), 0, false, false, s);
+      const array cotan_src =
+          where(mask_flat, cotan_flat, array(0, cotan_flat.dtype()), s);
+
+      array gsrc_flat =
+          zeros({static_cast<int>(src.size())}, cotan_src.dtype(), s);
+      if (src.size() > 0) {
+        const array cotan_updates =
+            reshape(cotan_src, {static_cast<int>(idx_src.size()), 1}, s);
+        gsrc_flat = scatter_add(gsrc_flat, idx_src, cotan_updates, 0, s);
+      }
+
+      vjps.push_back(reshape(gsrc_flat, src.shape(), s));
+    } else {
+      throw std::invalid_argument(
+          "[masked_scatter] Cannot calculate VJP with respect to mask.");
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> MaskedScatter::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto& s = stream();
+  const array& dst = primals[0];
+  const array& mask = primals[1];
+  array mask_b = mask;
+  if (mask_b.ndim() < dst.ndim()) {
+    std::vector<int> axes(dst.ndim() - mask_b.ndim(), 0);
+    std::iota(axes.begin(), axes.end(), mask_b.ndim());
+    mask_b = expand_dims(mask_b, axes, s);
+  }
+
+  array out = zeros_like(dst, s);
+  for (int arg : argnums) {
+    if (arg == 0) {
+      out = where(mask_b, out, tangents[0], s);
+    } else if (arg == 2) {
+      out = array(
+          out.shape(),
+          out.dtype(),
+          std::make_shared<MaskedScatter>(s),
+          {out, mask, tangents[1]});
+    } else {
+      throw std::invalid_argument("[masked_scatter] invalid arg index in JVP");
+    }
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> MaskedScatter::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto& s = stream();
+
+  // The inputs all had batching in the 0-th dim. So vectorization amounts to
+  //  - Move the vectorized axis first
+  //  - Expand and broadcast the unvectorized inputs
+  //  - Flatten the first two dims (the new and old batch axes)
+  //  - Masked scatter
+  //  - Unflatten the vectorized axis again
+
+  // Find the batch dim if any
+  int batch_dim = -1;
+  for (int i = 0; i < axes.size(); i++) {
+    if (axes[i] >= 0) {
+      batch_dim = inputs[i].shape(axes[i]);
+    }
+  }
+
+  // Early exit if it's not vmapped
+  if (batch_dim < 0) {
+    return {
+        {array(
+            inputs[0].shape(),
+            inputs[0].dtype(),
+            std::make_shared<MaskedScatter>(to_stream(s)),
+            inputs)},
+        {-1}};
+  }
+
+  // Move vmapped axis to 0-th dim and broadcast the non-vectorized ones
+  auto v_in = inputs;
+  for (int i = 0; i < axes.size(); i++) {
+    if (axes[i] > 0) {
+      v_in[i] = moveaxis(v_in[i], axes[i], 0, s);
+    } else if (axes[i] < 0) {
+      v_in[i] = expand_dims(v_in[i], 0, s);
+      auto in_shape = v_in[i].shape();
+      in_shape[0] = batch_dim;
+      v_in[i] = broadcast_to(v_in[i], in_shape, s);
+    }
+  }
+
+  // Flatten the first 2 dims
+  for (int i = 0; i < 3; i++) {
+    v_in[i] = flatten(v_in[i], 0, 1, s);
+  }
+
+  // Masked scatter
+  const auto result_shape = v_in[0].shape();
+  const auto result_dtype = v_in[0].dtype();
+  array result(
+      result_shape,
+      result_dtype,
+      std::make_shared<MaskedScatter>(to_stream(s)),
+      std::move(v_in));
+
+  // Now unflatten so the vectorized axis is nice and separate
+  result = unflatten(result, 0, {batch_dim, -1}, s);
+
+  return {{result}, {0}};
+}
+
+std::vector<array> Sigmoid::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  auto& s = outputs[0];
+  auto sprime =
+      multiply(s, subtract(array(1.0f, s.dtype()), s, stream()), stream());
+  return {multiply(cotangents[0], sprime, stream())};
+}
+
+std::vector<array> Sigmoid::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  auto s = sigmoid(primals[0], stream());
+  auto sprime =
+      multiply(s, subtract(array(1.0f, s.dtype()), s, stream()), stream());
+  return {multiply(tangents[0], sprime, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Sigmoid::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{sigmoid(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Sign::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Sign::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {zeros(primals[0].shape(), primals[0].dtype(), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Sign::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{sign(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Sin::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Sin::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(tangents[0], cos(primals[0], stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Sin::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{sin(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Sinh::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Sinh::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(tangents[0], cosh(primals[0], stream()), stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Sinh::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{sinh(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Slice::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto start = start_indices_;
+  auto stop = end_indices_;
+  auto strides = strides_;
+  auto ax = axes[0];
+  auto& input = inputs[0];
+  if (ax >= 0) {
+    start.insert(start.begin() + ax, 0);
+    stop.insert(stop.begin() + ax, input.shape(ax));
+    strides.insert(strides.begin() + ax, 1);
+  }
+  return {{slice(input, start, stop, strides, stream())}, {ax}};
+}
+
+std::vector<array> Slice::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  // Check inputs
+  assert(primals.size() == 1);
+  auto out = zeros_like(primals[0], stream());
+  return {slice_update(
+      out, cotangents[0], start_indices_, end_indices_, strides_, stream())};
+}
+
+std::vector<array> Slice::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  // Check inputs
+  assert(primals.size() == 1);
+  return {slice(tangents[0], start_indices_, end_indices_, strides_, stream())};
+}
+
+bool Slice::is_equivalent(const Primitive& other) const {
+  const Slice& s_other = static_cast<const Slice&>(other);
+  return (
+      start_indices_ == s_other.start_indices_ &&
+      end_indices_ == s_other.end_indices_ && strides_ == s_other.strides_);
+}
+
+std::pair<std::vector<array>, std::vector<int>> SliceUpdate::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 2);
+  assert(axes.size() == 2);
+
+  auto start = start_indices_;
+  auto stop = end_indices_;
+  auto strides = strides_;
+
+  auto src = inputs[0];
+  auto upd = inputs[1];
+
+  auto src_ax = axes[0];
+  auto upd_ax = axes[1];
+
+  // No vmapping needed
+  if (src_ax == -1 && upd_ax == -1) {
+    return {{slice_update(src, upd, start, stop, strides, stream())}, {-1}};
+  }
+
+  // Broadcast src
+  if (src_ax == -1) {
+    src = expand_dims(src, upd_ax, stream());
+    auto shape = src.shape();
+    shape[upd_ax] = upd.shape(upd_ax);
+    src = broadcast_to(src, shape, stream());
+    src_ax = upd_ax;
+  }
+
+  // Broadcast upd
+  if (upd_ax == -1) {
+    upd = expand_dims(upd, src_ax, stream());
+    upd_ax = src_ax;
+  }
+
+  if (src_ax != upd_ax) {
+    upd = moveaxis(upd, upd_ax, src_ax, stream());
+  }
+
+  start.insert(start.begin() + src_ax, 0);
+  stop.insert(stop.begin() + src_ax, src.shape(src_ax));
+  strides.insert(strides.begin() + src_ax, 1);
+
+  return {{slice_update(src, upd, start, stop, strides, stream())}, {src_ax}};
+}
+
+std::vector<array> SliceUpdate::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  // Check inputs
+  assert(primals.size() == 2);
+
+  auto& cotan = cotangents[0];
+  auto& upd = primals[1];
+
+  std::vector<array> vjps;
+
+  for (int num : argnums) {
+    // Vjp for source
+    if (num == 0) {
+      vjps.push_back(slice_update(
+          cotan,
+          zeros_like(upd, stream()),
+          start_indices_,
+          end_indices_,
+          strides_,
+          stream()));
+    }
+    // Vjp fpr updates
+    else {
+      vjps.push_back(
+          slice(cotan, start_indices_, end_indices_, strides_, stream()));
+    }
+  }
+
+  return vjps;
+}
+
+std::vector<array> SliceUpdate::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  // Check inputs
+  assert(primals.size() == 2);
+  return {slice_update(
+      tangents[0],
+      tangents[1],
+      start_indices_,
+      end_indices_,
+      strides_,
+      stream())};
+}
+
+bool SliceUpdate::is_equivalent(const Primitive& other) const {
+  const auto& s_other = static_cast<const SliceUpdate&>(other);
+  return (
+      start_indices_ == s_other.start_indices_ &&
+      end_indices_ == s_other.end_indices_ && strides_ == s_other.strides_);
+}
+
+std::pair<std::vector<array>, std::vector<int>> DynamicSlice::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto& in = inputs[0];
+  auto& start = inputs[1];
+  auto vax = axes[0];
+  if (axes[1] >= 0) {
+    throw std::invalid_argument(
+        "[DynamicSlice::vmap] vmap over start indices not yet supported.");
+  }
+  auto slice_size = slice_size_;
+  auto slice_axes = axes_;
+  if (vax >= 0) {
+    for (auto& ax : slice_axes) {
+      if (ax >= vax) {
+        ax++;
+      }
+    }
+    slice_size.insert(slice_size.begin() + vax, in.shape(vax));
+  }
+  return {
+      {slice(
+          in, start, std::move(slice_axes), std::move(slice_size), stream())},
+      {vax}};
+}
+
+std::vector<array> DynamicSlice::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  if (argnums[0] == 1 || argnums.size() > 1) {
+    throw std::invalid_argument(
+        "[DynamicSlice::vjp] Not supported for start indices.");
+  }
+  auto out = zeros_like(primals[0], stream());
+  return {slice_update(out, cotangents[0], primals[1], axes_, stream())};
+}
+
+std::vector<array> DynamicSlice::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {slice(tangents[0], primals[1], axes_, slice_size_, stream())};
+}
+
+bool DynamicSlice::is_equivalent(const Primitive& other) const {
+  const auto& s_other = static_cast<const DynamicSlice&>(other);
+  return (axes_ == s_other.axes_ && slice_size_ == s_other.slice_size_);
+}
+
+std::vector<Shape> DynamicSlice::output_shapes(const std::vector<array>&) {
+  return {slice_size_};
+}
+
+std::pair<std::vector<array>, std::vector<int>> DynamicSliceUpdate::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto src = inputs[0];
+  auto upd = inputs[1];
+  auto& start = inputs[2];
+  auto src_ax = axes[0];
+  auto upd_ax = axes[1];
+  if (axes[2] >= 0) {
+    throw std::runtime_error(
+        "[DynamicSliceUpdate::vmap] vmap over start indices not yet supported.");
+  }
+  // No vmapping needed
+  if (src_ax == -1 && upd_ax == -1) {
+    return {{slice_update(src, upd, start, axes_, stream())}, {-1}};
+  }
+
+  // Broadcast src
+  if (src_ax == -1) {
+    src = expand_dims(src, upd_ax, stream());
+    auto shape = src.shape();
+    shape[upd_ax] = upd.shape(upd_ax);
+    src = broadcast_to(src, shape, stream());
+    src_ax = upd_ax;
+  }
+
+  // Broadcast upd
+  if (upd_ax == -1) {
+    upd = expand_dims(upd, src_ax, stream());
+    upd_ax = src_ax;
+  }
+
+  if (src_ax != upd_ax) {
+    upd = moveaxis(upd, upd_ax, src_ax, stream());
+  }
+
+  auto slice_axes = axes_;
+  for (auto& ax : slice_axes) {
+    if (ax >= src_ax) {
+      ax++;
+    }
+  }
+  return {
+      {slice_update(src, upd, start, std::move(slice_axes), stream())},
+      {src_ax}};
+}
+
+std::vector<array> DynamicSliceUpdate::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  auto& cotan = cotangents[0];
+  auto& upd = primals[1];
+  auto& start = primals[2];
+
+  std::vector<array> vjps;
+
+  for (int num : argnums) {
+    if (num == 0) {
+      // Vjp for source
+      vjps.push_back(slice_update(
+          cotan, zeros_like(upd, stream()), start, axes_, stream()));
+    } else if (num == 1) {
+      // Vjp fpr updates
+      vjps.push_back(slice(cotan, start, axes_, upd.shape(), stream()));
+    } else {
+      throw std::invalid_argument(
+          "[DynamicSliceUpdate::vjp] Not supported for start indices");
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> DynamicSliceUpdate::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {slice_update(tangents[0], tangents[1], primals[2], axes_, stream())};
+}
+
+bool DynamicSliceUpdate::is_equivalent(const Primitive& other) const {
+  const auto& s_other = static_cast<const DynamicSliceUpdate&>(other);
+  return axes_ == s_other.axes_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Softmax::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  std::vector<int> softmax_axes;
+
+  // We are vectorizing over an axis other than the last one so keep the
+  // softmax axis unchanged
+  if (axes[0] >= 0 && axes[0] < inputs[0].ndim() - 1) {
+    softmax_axes.push_back(-1);
+  } else {
+    softmax_axes.push_back(-2);
+  }
+  return {{softmax(inputs[0], softmax_axes, precise_, stream())}, axes};
+}
+
+std::vector<array> Softmax::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  assert(primals.size() == 1);
+  assert(cotangents.size() == 1);
+  auto& s = outputs[0];
+  auto sv = multiply(s, cotangents[0], stream());
+  return {subtract(
+      sv,
+      multiply(s, sum(sv, std::vector<int>{-1}, true, stream()), stream()))};
+}
+
+std::vector<array> Softmax::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(tangents.size() == 1);
+  auto s = softmax(primals[0], std::vector<int>{-1}, precise_, stream());
+  auto sv = multiply(s, tangents[0], stream());
+  return {subtract(
+      sv,
+      multiply(s, sum(sv, std::vector<int>{-1}, true, stream()), stream()))};
+}
+
+bool Softmax::is_equivalent(const Primitive& other) const {
+  const Softmax& s_other = static_cast<const Softmax&>(other);
+  return precise_ == s_other.precise_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Sort::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  int axis_left = axes[0] >= 0 && axes[0] <= axis_;
+  return {{sort(inputs[0], axis_ + axis_left, stream())}, axes};
+}
+
+std::vector<array> Sort::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Sort::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(tangents.size() == 1);
+  auto sort_idx = argsort(primals[0], axis_, stream());
+  auto out = take_along_axis(tangents[0], sort_idx, axis_, stream());
+  return {out};
+}
+
+bool Sort::is_equivalent(const Primitive& other) const {
+  const Sort& r_other = static_cast<const Sort&>(other);
+  return axis_ == r_other.axis_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> Split::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  int axis_left = axes[0] >= 0 && axes[0] <= axis_;
+  auto output = split(inputs[0], indices_, axis_ + axis_left, stream());
+  std::vector<int> output_axes(output.size(), axes[0]);
+  return {std::move(output), std::move(output_axes)};
+}
+
+std::vector<array> Split::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return {concatenate(cotangents, axis_, stream())};
+}
+
+std::vector<array> Split::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  return split(tangents[0], indices_, axis_, stream());
+}
+
+bool Split::is_equivalent(const Primitive& other) const {
+  const Split& s_other = static_cast<const Split&>(other);
+  return axis_ == s_other.axis_ && indices_ == s_other.indices_;
+}
+
+std::vector<array> Square::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Square::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(tangents.size() == 1);
+  return {multiply(
+      primals[0],
+      multiply(array(2, primals[0].dtype()), tangents[0], stream()),
+      stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Square::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{square(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Sqrt::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>& outputs) {
+  assert(primals.size() == 1);
+  assert(cotangents.size() == 1);
+  auto dtype = primals[0].dtype();
+  if (recip_) {
+    auto one_over_x_root_x = divide(outputs[0], primals[0], stream());
+    return {multiply(
+        multiply(array(-0.5, dtype), cotangents[0], stream()),
+        one_over_x_root_x,
+        stream())};
+  } else {
+    return {divide(
+        multiply(array(0.5, dtype), cotangents[0], stream()),
+        outputs[0],
+        stream())};
+  }
+}
+
+std::vector<array> Sqrt::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  if (recip_) {
+    return vjp(primals, tangents, argnums, {rsqrt(primals[0], stream())});
+  } else {
+    return vjp(primals, tangents, argnums, {sqrt(primals[0], stream())});
+  }
+}
+
+std::pair<std::vector<array>, std::vector<int>> Sqrt::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  if (recip_) {
+    return {{rsqrt(inputs[0], stream())}, axes};
+  }
+  return {{sqrt(inputs[0], stream())}, axes};
+}
+
+bool Sqrt::is_equivalent(const Primitive& other) const {
+  const Sqrt& s_other = static_cast<const Sqrt&>(other);
+  return recip_ == s_other.recip_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> StopGradient::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  return {{stop_gradient(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Subtract::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    auto vjp = cotangents[0];
+    if (arg == 1) {
+      vjp = negative(vjp, stream());
+    }
+    vjps.push_back(vjp);
+  }
+  return vjps;
+}
+
+std::vector<array> Subtract::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  auto jvp_fun = [&](int i) {
+    int arg = argnums[i];
+    return arg == 1 ? negative(tangents[i], stream()) : tangents[i];
+  };
+  auto out = jvp_fun(0);
+  if (argnums.size() > 1) {
+    out = add(out, jvp_fun(1), stream());
+  }
+  return {out};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Subtract::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream());
+  return {{subtract(a, b, stream())}, {to_ax}};
+}
+
+std::vector<array> Squeeze::vjp(
+    const std::vector<array>&,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {expand_dims(cotangents[0], axes_, stream())};
+}
+
+std::vector<array> Squeeze::jvp(
+    const std::vector<array>&,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {squeeze(tangents[0], axes_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Squeeze::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0];
+  auto squeeze_axes = axes_;
+  for (auto& s : squeeze_axes) {
+    if (s >= axes[0]) {
+      s++;
+    } else {
+      ax--;
+    }
+  }
+  return {{squeeze(inputs[0], std::move(squeeze_axes), stream())}, {ax}};
+}
+
+bool Squeeze::is_equivalent(const Primitive& other) const {
+  const Squeeze& a_other = static_cast<const Squeeze&>(other);
+  return (axes_ == a_other.axes_);
+}
+
+Shape Squeeze::output_shape(const array& input, const std::vector<int>& axes) {
+  Shape shape;
+  for (int i = 0, j = 0; i < input.ndim(); ++i) {
+    if (j < axes.size() && i == axes[j]) {
+      j++;
+    } else {
+      shape.push_back(input.shape(i));
+    }
+  }
+  return shape;
+}
+
+std::vector<Shape> Squeeze::output_shapes(const std::vector<array>& inputs) {
+  return {Squeeze::output_shape(inputs[0], axes_)};
+}
+
+std::vector<array> Tan::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Tan::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array cos_sq = square(cos(primals[0], stream()), stream());
+  return {divide(tangents[0], cos_sq, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Tan::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{tan(inputs[0], stream())}, axes};
+}
+
+std::vector<array> Tanh::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Tanh::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  array cosh_sq = square(cosh(primals[0], stream()), stream());
+  return {divide(tangents[0], cosh_sq, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Tanh::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{tanh(inputs[0], stream())}, axes};
+}
+
+std::pair<std::vector<array>, std::vector<int>> BitwiseInvert::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{bitwise_invert(inputs[0], stream())}, axes};
+}
+
+std::vector<array> BlockMaskedMM::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  /////////////////////////////////////////////////////////////////////////////
+  // The operation that is done w/o intermediates by the primitive is
+  //    - tm = (M + block_size - 1) // block_size; MP = tm * block_size;
+  //    - tn = (N + block_size - 1) // block_size; NP = tn * block_size;
+  //    - tm = (K + block_size - 1) // block_size; KP = tk * block_size;
+  //    - mask_b <- mask broadcasted to block sizes
+  //    - A_m = A [..., M, K] * mask_b_lhs [..., MP, KP]
+  //    - B_m = B [..., K, N] * mask_b_rhs [..., KP, MP]
+  //    - C = A_m [..., M, K]  @ B_m [..., K, N]
+  //    - C_m = C [..., M, N] * mask_b_out [..., MP, NP]
+  //
+  // The grads are therefore
+  //    - dC_m = cotan [..., M, N]
+  //    - dmask_b_out = cotan [..., M, N] * C [..., M, N]
+  //    - dC = cotan [..., M, N] * mask_b_out [..., MP, NP]
+  //    - dA_m = dC [..., M, N] @ B_m.T [..., N, K]
+  //    - dB_m = A_m.T [..., K, M] @ dC [..., M, N]
+  //    - dA = dA_m * mask_b_lhs [..., MP, KP]
+  //    - dB = dB_m * mask_b_rhs [..., KP, MP]
+  //    - dmask_b_lhs = dA_m [..., M, K] * A [..., M, K] // need [..., MP,
+  //    KP]
+  //    - dmask_b_rhs = dB_m [..., K, N] * B [..., K, N] // need [..., KP,
+  //    NP]
+  //
+  // Observations:
+  //  * If dmask_b_lhs is not needed, then dA can be calulated in one go as
+  //  a
+  //    as a block_masked_mm with mask_b_lhs as the out_mask without needing
+  //    to materialize the intermediate dA_m. Similar for dB.
+  //  * If dmask_b_lhs is needed, we need to materialize dA_m directly and
+  //  then
+  //    point-wise multiply with A. But the output needs to be padded
+
+  std::vector<array> vjps;
+  auto& cotan = cotangents[0];
+  std::vector<int> reorder(cotan.ndim());
+  std::iota(reorder.begin(), reorder.end(), 0);
+  std::iter_swap(reorder.end() - 1, reorder.end() - 2);
+
+  bool has_op_mask = primals.size() > 3;
+  bool has_out_mask = primals.size() == 3 || primals.size() == 5;
+
+  const int op_mask_idx = has_out_mask ? 3 : 2;
+  bool needs_lhs_mask_vjp = has_op_mask;
+  bool needs_rhs_mask_vjp = has_op_mask;
+
+  for (auto arg : argnums) {
+    needs_lhs_mask_vjp = arg == op_mask_idx;
+    needs_rhs_mask_vjp = arg == op_mask_idx + 1;
+  }
+
+  if ((needs_lhs_mask_vjp && primals[op_mask_idx].dtype() == bool_) ||
+      (needs_rhs_mask_vjp && primals[op_mask_idx + 1].dtype() == bool_)) {
+    throw std::invalid_argument(
+        "[BlockMaskedMM] Cannot calculate VJP with respect to boolean masks.");
+  }
+
+  auto expand_mask = [&](array mask, int Y, int X) {
+    // Exapnd mask
+    auto mask_reshape = mask.shape();
+    mask = expand_dims(mask, {-3, -1}, stream());
+    auto mask_shape = mask.shape();
+    int mask_ndim = mask_shape.size();
+
+    // Broadcast mask
+    mask_shape[mask_ndim - 1] = block_size_;
+    mask_shape[mask_ndim - 3] = block_size_;
+    mask = broadcast_to(mask, mask_shape, stream());
+
+    // Reshape mask to squeeze in braodcasted dims
+    mask_ndim = mask_reshape.size();
+    mask_reshape[mask_ndim - 2] *= block_size_;
+    mask_reshape[mask_ndim - 1] *= block_size_;
+    mask = reshape(mask, mask_reshape, stream());
+
+    // Slice mask
+    mask_reshape[mask_ndim - 2] = Y;
+    mask_reshape[mask_ndim - 1] = X;
+    mask = slice(mask, Shape(mask_ndim, 0), mask_reshape, stream());
+
+    return mask;
+  };
+
+  array zero = array(0, cotan.dtype());
+
+  auto multiply_pad_reduce = [&](array p, array q, int align_Y, int align_X) {
+    // Multiply with cotan
+    auto r = multiply(p, q, stream());
+
+    // Pad if needed
+    if ((align_Y != 0) || (align_X != 0)) {
+      r = pad(
+          r, {-2, -1}, {0, 0}, {align_Y, align_X}, zero, "constant", stream());
+    }
+
+    // Reshape
+    Shape r_reshape(r.shape().begin(), r.shape().end() - 2);
+    r_reshape.push_back(r.shape(-2) / block_size_);
+    r_reshape.push_back(block_size_);
+    r_reshape.push_back(r.shape(-1) / block_size_);
+    r_reshape.push_back(block_size_);
+    r = reshape(r, r_reshape, stream());
+
+    // Reduce
+    return sum(r, {-3, -1}, false, stream());
+  };
+
+  // Prepare for padding if needed
+  const int M = cotan.shape(-2);
+  const int N = cotan.shape(-1);
+  const int K = primals[0].shape(-1);
+  const int tm = (M + block_size_ - 1) / block_size_;
+  const int tn = (N + block_size_ - 1) / block_size_;
+  const int tk = (K + block_size_ - 1) / block_size_;
+  const int align_M = tm * block_size_ - M;
+  const int align_N = tn * block_size_ - N;
+  const int align_K = tk * block_size_ - K;
+
+  // Potential intermediates
+  array unmasked_lhs_grad = primals[0];
+  array unmasked_rhs_grad = primals[1];
+
+  bool unmasked_lhs_grad_calculated = false;
+  bool unmasked_rhs_grad_calculated = false;
+
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      // M X N * (K X N).T -> M X K
+      auto b_t = transpose(primals[1], reorder, stream());
+      auto out_mask =
+          has_out_mask ? std::make_optional<array>(primals[2]) : std::nullopt;
+      auto lhs_mask = has_op_mask && !needs_lhs_mask_vjp
+          ? std::make_optional<array>(primals[op_mask_idx])
+          : std::nullopt;
+      auto rhs_mask_t = has_op_mask
+          ? std::make_optional<array>(
+                transpose(primals[op_mask_idx + 1], reorder, stream()))
+          : std::nullopt;
+
+      auto grad = block_masked_mm(
+          cotan, b_t, block_size_, lhs_mask, out_mask, rhs_mask_t, stream());
+
+      if (needs_lhs_mask_vjp) {
+        unmasked_lhs_grad = grad;
+        unmasked_lhs_grad_calculated = true;
+        auto exp_mask = expand_mask(primals[op_mask_idx], M, K);
+        grad = multiply(grad, exp_mask, stream());
+      }
+
+      vjps.push_back(grad);
+
+    } else if (arg == 1) {
+      // (M X K).T * M X N -> K X N
+      auto a_t = transpose(primals[0], reorder, stream());
+      auto out_mask =
+          has_out_mask ? std::make_optional<array>(primals[2]) : std::nullopt;
+      auto lhs_mask_t = has_op_mask
+          ? std::make_optional<array>(
+                transpose(primals[op_mask_idx], reorder, stream()))
+          : std::nullopt;
+      auto rhs_mask = has_op_mask && !needs_rhs_mask_vjp
+          ? std::make_optional<array>(primals[op_mask_idx + 1])
+          : std::nullopt;
+
+      auto grad = block_masked_mm(
+          a_t, cotan, block_size_, rhs_mask, lhs_mask_t, out_mask, stream());
+
+      if (needs_rhs_mask_vjp) {
+        unmasked_rhs_grad = grad;
+        unmasked_rhs_grad_calculated = true;
+        auto exp_mask = expand_mask(primals[op_mask_idx + 1], K, N);
+        grad = multiply(grad, exp_mask, stream());
+      }
+
+      vjps.push_back(grad);
+
+    } else if (arg == 2 && has_out_mask) {
+      // Produce the forward result
+      auto lhs_mask = has_op_mask
+          ? std::make_optional<array>(primals[op_mask_idx])
+          : std::nullopt;
+      auto rhs_mask = has_op_mask
+          ? std::make_optional<array>(primals[op_mask_idx + 1])
+          : std::nullopt;
+
+      auto C = block_masked_mm(
+          primals[0],
+          primals[1],
+          block_size_,
+          primals[2],
+          lhs_mask,
+          rhs_mask,
+          stream());
+
+      // Multiply, Pad and Reduce if needed
+      auto grad = multiply_pad_reduce(cotan, C, align_M, align_N);
+      vjps.push_back(grad);
+
+    } else if (arg == op_mask_idx && has_op_mask) {
+      if (!unmasked_lhs_grad_calculated) {
+        // (M X K).T * M X N -> K X N
+        auto b_t = transpose(primals[1], reorder, stream());
+        auto out_mask =
+            has_out_mask ? std::make_optional<array>(primals[2]) : std::nullopt;
+        auto rhs_mask_t =
+            transpose(primals[op_mask_idx + 1], reorder, stream());
+
+        unmasked_lhs_grad = block_masked_mm(
+            cotan,
+            b_t,
+            block_size_,
+            std::nullopt,
+            out_mask,
+            rhs_mask_t,
+            stream());
+
+        unmasked_lhs_grad_calculated = true;
+      }
+
+      // Multiply, Pad and Reduce if needed
+      auto grad =
+          multiply_pad_reduce(primals[0], unmasked_lhs_grad, align_M, align_K);
+      vjps.push_back(grad);
+
+    } else if (arg == op_mask_idx + 1 && has_op_mask) {
+      if (!unmasked_rhs_grad_calculated) {
+        // (M X K).T * M X N -> K X N
+        auto a_t = transpose(primals[0], reorder, stream());
+        auto out_mask =
+            has_out_mask ? std::make_optional<array>(primals[2]) : std::nullopt;
+        auto lhs_mask_t = transpose(primals[op_mask_idx], reorder, stream());
+
+        unmasked_rhs_grad = block_masked_mm(
+            a_t,
+            cotan,
+            block_size_,
+            std::nullopt,
+            lhs_mask_t,
+            out_mask,
+            stream());
+
+        unmasked_rhs_grad_calculated = true;
+      }
+
+      // Multiply, Pad and Reduce if needed
+      auto grad =
+          multiply_pad_reduce(primals[1], unmasked_rhs_grad, align_K, align_N);
+      vjps.push_back(grad);
+
+    } else {
+      throw std::invalid_argument(
+          "[BlockMaskedMM] Cannot calculate VJP with respect to masks.");
+    }
+  }
+  return vjps;
+}
+
+std::vector<array> GatherMM::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  std::vector<array> vjps;
+  auto& cotan = cotangents[0];
+
+  auto& a = primals[0];
+  auto& b = primals[1];
+  auto& lhs_indices = primals[2];
+  auto& rhs_indices = primals[3];
+
+  int M = cotan.shape(-2);
+  int K = primals[0].shape(-1);
+
+  bool sorted = left_sorted_ || right_sorted_;
+  bool no_broadcast = rhs_indices.size() * M * K == primals[0].size();
+
+  for (auto arg : argnums) {
+    if (arg == 0) {
+      auto g = gather_mm(
+          cotan,
+          swapaxes(b, -1, -2, stream()),
+          std::nullopt,
+          rhs_indices,
+          sorted,
+          stream());
+      if (sorted && no_broadcast) {
+        vjps.push_back(g);
+      } else {
+        vjps.push_back(reshape(
+            scatter_add(
+                flatten(zeros_like(a, stream()), 0, -3, stream()),
+                lhs_indices,
+                expand_dims(g, -3, stream()),
+                0,
+                stream()),
+            a.shape(),
+            stream()));
+      }
+    } else if (arg == 1) {
+      auto shape = b.shape();
+      shape.pop_back();
+      shape.pop_back();
+      vjps.push_back(swapaxes(
+          gather_mm_grad(
+              a,
+              cotan,
+              lhs_indices,
+              rhs_indices,
+              sorted,
+              std::move(shape),
+              stream()),
+          -1,
+          -2,
+          stream()));
+    } else {
+      throw std::invalid_argument(
+          "[GatherMM] Cannot calculate VJP with respect to indices.");
+    }
+  }
+  return vjps;
+}
+
+bool GatherMM::is_equivalent(const Primitive& other) const {
+  const GatherMM& g_other = static_cast<const GatherMM&>(other);
+  return left_sorted_ == g_other.left_sorted_ &&
+      right_sorted_ == g_other.right_sorted_;
+}
+
+bool BlockMaskedMM::is_equivalent(const Primitive& other) const {
+  const BlockMaskedMM& a_other = static_cast<const BlockMaskedMM&>(other);
+  return (block_size_ == a_other.block_size_);
+}
+
+std::vector<array> Transpose::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  std::vector<int> iaxes(axes_.size());
+  for (int i = 0; i < axes_.size(); ++i) {
+    iaxes[axes_[i]] = i;
+  }
+  return {transpose(cotangents[0], iaxes, stream())};
+}
+
+std::vector<array> Transpose::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(tangents.size() == 1);
+  return {transpose(tangents[0], axes_, stream())};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Transpose::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  auto vdim = axes[0];
+  if (vdim >= 0) {
+    for (auto& dim : axes_) {
+      if (dim >= vdim) {
+        dim++;
+      }
+    }
+    axes_.insert(axes_.begin() + vdim, vdim);
+  }
+  return {{transpose(inputs[0], axes_, stream())}, {vdim}};
+}
+
+bool Transpose::is_equivalent(const Primitive& other) const {
+  const Transpose& t_other = static_cast<const Transpose&>(other);
+  return axes_ == t_other.axes_;
+}
+
+std::vector<Shape> Transpose::output_shapes(const std::vector<array>& inputs) {
+  auto& in = inputs[0];
+  Shape shape(in.ndim(), 0);
+  for (int i = 0; i < axes_.size(); ++i) {
+    shape[i] = in.shape()[axes_[i]];
+  }
+  return {std::move(shape)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> NumberOfElements::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+
+  std::vector<int> new_axes = axes_;
+  auto vdim = axes[0];
+  if (vdim >= 0) {
+    for (auto& dim : new_axes) {
+      if (dim >= vdim) {
+        dim++;
+      }
+    }
+  }
+
+  array out = array(
+      {},
+      dtype_,
+      std::make_shared<NumberOfElements>(stream(), new_axes, inverted_, dtype_),
+      inputs);
+
+  return {{out}, {-1}};
+}
+
+bool NumberOfElements::is_equivalent(const Primitive& other) const {
+  const NumberOfElements& n_other = static_cast<const NumberOfElements&>(other);
+  return axes_ == n_other.axes_ && inverted_ == n_other.inverted_ &&
+      dtype_ == n_other.dtype_;
+}
+
+std::pair<std::vector<array>, std::vector<int>> SVD::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0] >= 0 ? 0 : -1;
+  auto a = axes[0] > 0 ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  std::vector<int> new_axes(compute_uv_ ? 3 : 1, ax);
+  return {linalg::svd(a, compute_uv_, stream()), std::move(new_axes)};
+}
+
+std::pair<std::vector<array>, std::vector<int>> Inverse::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  auto ax = axes[0] >= 0 ? 0 : -1;
+  auto a = axes[0] > 0 ? moveaxis(inputs[0], axes[0], 0, stream()) : inputs[0];
+  return {{linalg::inv(a, stream())}, {ax}};
+}
+
+std::pair<std::vector<array>, std::vector<int>> View::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  return {{view(inputs[0], dtype_, stream())}, axes};
+}
+
+const char* View::name() const {
+  if (name_.empty()) {
+    std::ostringstream os;
+    os << "View " << dtype_;
+    name_ = os.str();
+  }
+  return name_.c_str();
+}
+
+bool View::is_equivalent(const Primitive& other) const {
+  const View& a_other = static_cast<const View&>(other);
+  return (dtype_ == a_other.dtype_);
+}
+
+std::pair<std::vector<array>, std::vector<int>> Hadamard::vmap(
+    const std::vector<array>& inputs,
+    const std::vector<int>& axes) {
+  assert(inputs.size() == 1);
+  assert(axes.size() == 1);
+  auto& s = stream();
+  if (axes[0] == inputs[0].ndim() - 1) {
+    auto a = moveaxis(inputs[0], axes[0], 0, s);
+    auto b = hadamard_transform(a, scale_, s);
+    return {{b}, {0}};
+  }
+  return {{hadamard_transform(inputs[0], scale_, s)}, axes};
+}
+
+std::vector<array> Hadamard::vjp(
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents,
+    const std::vector<int>& argnums,
+    const std::vector<array>&) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return jvp(primals, cotangents, argnums);
+}
+
+std::vector<array> Hadamard::jvp(
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents,
+    const std::vector<int>& argnums) {
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {hadamard_transform(tangents[0], scale_, stream())};
+}
+
+bool Hadamard::is_equivalent(const Primitive& other) const {
+  const Hadamard& h_other = static_cast<const Hadamard&>(other);
+  return scale_ == h_other.scale_;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/primitives.h b/macos/third_party/mlx/mlx/primitives.h
new file mode 100644
index 0000000..4091aaf
--- /dev/null
+++ b/macos/third_party/mlx/mlx/primitives.h
@@ -0,0 +1,2525 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <unordered_set>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/device.h"
+#include "mlx/io/load.h"
+#include "mlx/stream.h"
+
+#define DEFINE_VMAP()                                                 \
+  virtual std::pair<std::vector<array>, std::vector<int>> vmap(       \
+      const std::vector<array>& inputs, const std::vector<int>& axes) \
+      override;
+
+#define DEFINE_GRADS()                           \
+  std::vector<array> jvp(                        \
+      const std::vector<array>& primals,         \
+      const std::vector<array>& tangents,        \
+      const std::vector<int>& argnums) override; \
+                                                 \
+  std::vector<array> vjp(                        \
+      const std::vector<array>& primals,         \
+      const std::vector<array>& cotangents,      \
+      const std::vector<int>& argnums,           \
+      const std::vector<array>& outputs) override;
+
+#define DEFINE_NAME(PRIMITIVE)        \
+  const char* name() const override { \
+    return #PRIMITIVE;                \
+  }
+
+#define DEFINE_DEFAULT_IS_EQUIVALENT()                        \
+  bool is_equivalent(const Primitive& other) const override { \
+    return true;                                              \
+  }
+
+#define DEFINE_INPUT_OUTPUT_SHAPE()                                  \
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) \
+      override {                                                     \
+    return {inputs[0].shape()};                                      \
+  }
+
+namespace mlx::core {
+
+// Abstract base class
+class MLX_API Primitive {
+ public:
+  explicit Primitive(Stream stream) : stream_(stream) {}
+
+  /** The device the primitive will run on. */
+  const Device& device() {
+    return stream().device;
+  }
+
+  /** The stream the primitive will run on. */
+  const Stream& stream() {
+    return stream_;
+  }
+
+  /**
+   * A primitive must know how to evaluate itself on
+   * the CPU/GPU for the given inputs and populate the output arrays.
+   *
+   * To avoid unnecessary allocations, the evaluation function
+   * is responsible for allocating space for the array.
+   */
+  virtual void eval_cpu(
+      const std::vector<array>& inputs,
+      std::vector<array>& outputs) = 0;
+  virtual void eval_gpu(
+      const std::vector<array>& inputs,
+      std::vector<array>& outputs) = 0;
+
+  /**
+   * The Jacobian-vector product.
+   */
+  virtual std::vector<array> jvp(
+      const std::vector<array>& primals,
+      const std::vector<array>& tangents,
+      const std::vector<int>& argnums);
+
+  /**
+   * The vector-Jacobian product.
+   */
+  virtual std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs);
+
+  /**
+   * The primitive must know how to vectorize itself across
+   * the given axes. The output is a pair containing the output arrays
+   * representing the vectorized computation and the axes which
+   * corresponds to the vectorized dimensions of each output.
+   */
+  virtual std::pair<std::vector<array>, std::vector<int>> vmap(
+      const std::vector<array>& inputs,
+      const std::vector<int>& axes);
+
+  /** Get the name of primitive. */
+  virtual const char* name() const = 0;
+
+  /** Equivalence check defaults to false unless overridden by the primitive */
+  virtual bool is_equivalent(const Primitive& other) const {
+    return false;
+  }
+
+  /** Get the output shapes of the primitive. This is not required to be
+   * implemented by derived classes, in which case it will throw. */
+  virtual std::vector<Shape> output_shapes(const std::vector<array>& inputs);
+
+  virtual ~Primitive() = default;
+  Primitive(const Primitive& other) = delete;
+  Primitive(Primitive&& other) = delete;
+  Primitive& operator=(const Primitive& other) = delete;
+  Primitive& operator=(Primitive&& other) = delete;
+
+ private:
+  // Every primitive stores the stream it should run in
+  Stream stream_;
+};
+
+class MLX_API UnaryPrimitive : public Primitive {
+  /**
+   * An abstract base class for a primitive with a single output.
+   */
+ public:
+  explicit UnaryPrimitive(Stream stream) : Primitive(stream) {}
+
+  virtual void eval_cpu(const std::vector<array>& inputs, array& output) = 0;
+  virtual void eval_gpu(const std::vector<array>& inputs, array& output) = 0;
+
+  inline void eval_cpu(
+      const std::vector<array>& inputs,
+      std::vector<array>& outputs) override {
+    eval_cpu(inputs, outputs[0]);
+  }
+  inline void eval_gpu(
+      const std::vector<array>& inputs,
+      std::vector<array>& outputs) override {
+    eval_gpu(inputs, outputs[0]);
+  }
+
+  virtual ~UnaryPrimitive() = default;
+  UnaryPrimitive(const UnaryPrimitive& other) = delete;
+  UnaryPrimitive(UnaryPrimitive&& other) = delete;
+  UnaryPrimitive& operator=(const UnaryPrimitive& other) = delete;
+  UnaryPrimitive& operator=(UnaryPrimitive&& other) = delete;
+};
+
+enum class QuantizationMode { Affine, Mxfp4, Mxfp8, Nvfp4 };
+
+std::string quantization_mode_to_string(QuantizationMode mode);
+QuantizationMode string_to_quantization_mode(
+    const std::string& mode,
+    std::string_view error_tag = "");
+
+class Abs : public UnaryPrimitive {
+ public:
+  explicit Abs(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Abs)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class MLX_API Add : public UnaryPrimitive {
+ public:
+  explicit Add(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Add)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class AddMM : public UnaryPrimitive {
+ public:
+  explicit AddMM(Stream stream, float alpha, float beta)
+      : UnaryPrimitive(stream), alpha_(alpha), beta_(beta) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_GRADS()
+  DEFINE_VMAP()
+  DEFINE_NAME(AddMM)
+
+  bool is_equivalent(const Primitive& other) const override;
+  std::pair<float, float> state() const {
+    return {alpha_, beta_};
+  };
+
+ private:
+  const float alpha_;
+  const float beta_;
+};
+
+class Arange : public UnaryPrimitive {
+ public:
+  explicit Arange(Stream stream, double start, double stop, double step)
+      : UnaryPrimitive(stream), start_(start), stop_(stop), step_(step) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_NAME(Arange)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  std::tuple<double, double, double> state() const {
+    return {start_, stop_, step_};
+  };
+
+ private:
+  double start_;
+  double stop_;
+  double step_;
+};
+
+class ArcCos : public UnaryPrimitive {
+ public:
+  explicit ArcCos(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArcCos)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ArcCosh : public UnaryPrimitive {
+ public:
+  explicit ArcCosh(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArcCosh)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ArcSin : public UnaryPrimitive {
+ public:
+  explicit ArcSin(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArcSin)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ArcSinh : public UnaryPrimitive {
+ public:
+  explicit ArcSinh(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArcSinh)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ArcTan : public UnaryPrimitive {
+ public:
+  explicit ArcTan(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArcTan)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ArcTan2 : public UnaryPrimitive {
+ public:
+  explicit ArcTan2(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArcTan2)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ArcTanh : public UnaryPrimitive {
+ public:
+  explicit ArcTanh(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArcTanh)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ArgPartition : public UnaryPrimitive {
+ public:
+  explicit ArgPartition(Stream stream, int kth, int axis)
+      : UnaryPrimitive(stream), kth_(kth), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArgPartition)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
+  std::pair<int, int> state() const {
+    return {kth_, axis_};
+  };
+
+ private:
+  int kth_;
+  int axis_;
+};
+
+class MLX_API ArgReduce : public UnaryPrimitive {
+ public:
+  enum ReduceType {
+    ArgMin,
+    ArgMax,
+  };
+
+  explicit ArgReduce(Stream stream, ReduceType reduce_type, int axis)
+      : UnaryPrimitive(stream), reduce_type_(reduce_type), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArgReduce)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  std::pair<ReduceType, int> state() const {
+    return {reduce_type_, axis_};
+  };
+
+ private:
+  ReduceType reduce_type_;
+  int axis_;
+};
+
+class ArgSort : public UnaryPrimitive {
+ public:
+  explicit ArgSort(Stream stream, int axis)
+      : UnaryPrimitive(stream), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ArgSort)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
+  int state() const {
+    return axis_;
+  };
+
+ private:
+  int axis_;
+};
+
+class AsType : public UnaryPrimitive {
+ public:
+  explicit AsType(Stream stream, Dtype dtype)
+      : UnaryPrimitive(stream), dtype_(dtype) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(AsType)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
+  Dtype state() const {
+    return dtype_;
+  };
+
+ private:
+  Dtype dtype_;
+};
+
+class AsStrided : public UnaryPrimitive {
+ public:
+  explicit AsStrided(Stream stream, Shape shape, Strides strides, size_t offset)
+      : UnaryPrimitive(stream),
+        shape_(std::move(shape)),
+        strides_(std::move(strides)),
+        offset_(offset) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_GRADS()
+  DEFINE_NAME(AsStrided)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(shape_, strides_, offset_);
+  }
+
+ private:
+  Shape shape_;
+  Strides strides_;
+  size_t offset_;
+
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+class BitwiseBinary : public UnaryPrimitive {
+ public:
+  enum Op { And, Or, Xor, LeftShift, RightShift };
+
+  explicit BitwiseBinary(Stream stream, Op op)
+      : UnaryPrimitive(stream), op_(op) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+
+  const char* name() const override {
+    switch (op_) {
+      case BitwiseBinary::And:
+        return "BitwiseAnd";
+      case BitwiseBinary::Or:
+        return "BitwiseOr";
+      case BitwiseBinary::Xor:
+        return "BitwiseXor";
+      case BitwiseBinary::LeftShift:
+        return "LeftShift";
+      case BitwiseBinary::RightShift:
+        return "RightShift";
+    }
+    return "<unknwon BitwiseBinary>";
+  }
+
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return op_;
+  }
+
+ private:
+  Op op_;
+};
+
+class BitwiseInvert : public UnaryPrimitive {
+ public:
+  explicit BitwiseInvert(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(BitwiseInvert)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class BlockMaskedMM : public UnaryPrimitive {
+ public:
+  explicit BlockMaskedMM(Stream stream, int block_size)
+      : UnaryPrimitive(stream), block_size_(block_size) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(BlockMaskedMM)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return block_size_;
+  }
+
+ private:
+  int block_size_;
+};
+
+class GatherMM : public UnaryPrimitive {
+ public:
+  explicit GatherMM(
+      Stream stream,
+      bool left_sorted = false,
+      bool right_sorted = false)
+      : UnaryPrimitive(stream),
+        left_sorted_(left_sorted),
+        right_sorted_(right_sorted) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(GatherMM)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_pair(left_sorted_, right_sorted_);
+  }
+
+ private:
+  bool left_sorted_;
+  bool right_sorted_;
+};
+
+class SegmentedMM : public UnaryPrimitive {
+ public:
+  explicit SegmentedMM(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_NAME(SegmentedMM)
+};
+
+class BroadcastAxes : public UnaryPrimitive {
+ public:
+  explicit BroadcastAxes(Stream stream, std::vector<int> ignore_axes = {})
+      : UnaryPrimitive(stream), ignore_axes_(std::move(ignore_axes)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(BroadcastAxes)
+  bool is_equivalent(const Primitive& other) const override;
+  static Shape output_shape(
+      const std::vector<array>& inputs,
+      const std::vector<int>& ignore_axes);
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return ignore_axes_;
+  }
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+  std::vector<int> ignore_axes_;
+};
+
+class Broadcast : public UnaryPrimitive {
+ public:
+  explicit Broadcast(Stream stream, const Shape& shape)
+      : UnaryPrimitive(stream), shape_(shape) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Broadcast)
+  static Shape output_shape(const std::vector<array>& inputs);
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  bool is_equivalent(const Primitive& other) const override;
+  Shape state() const {
+    return shape_;
+  };
+
+ private:
+  Shape shape_;
+
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+class Ceil : public UnaryPrimitive {
+ public:
+  explicit Ceil(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Ceil)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class MLX_API Compiled : public Primitive {
+ public:
+  /*
+   * The inputs, outputs and tape are either tracers or constants.
+   * - The tape should not contain the inputs, but it should contain the
+   *   outputs.
+   * - The tape should also have only one array per primitive for multi-output
+   *   primitives.
+   * - The constant_ids contains ids of arrays in the input list that are safe
+   *   to treat as scalar constants.
+   */
+  explicit Compiled(
+      Stream stream,
+      std::vector<array> inputs,
+      std::vector<array> outputs,
+      std::vector<array> tape,
+      std::unordered_set<uintptr_t> constant_ids);
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  const char* name() const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  bool is_equivalent(const Primitive& other) const override;
+
+  std::string lib_name() const {
+    return kernel_lib_;
+  }
+
+ private:
+  const std::vector<array> inputs_;
+  const std::vector<array> outputs_;
+  const std::vector<array> tape_;
+  const std::unordered_set<uintptr_t> constant_ids_;
+  const std::function<bool(size_t)> is_constant_;
+
+  mutable std::string name_;
+  std::string kernel_lib_;
+};
+
+class Concatenate : public UnaryPrimitive {
+ public:
+  explicit Concatenate(Stream stream, int axis)
+      : UnaryPrimitive(stream), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Concatenate)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return axis_;
+  }
+
+ private:
+  int axis_;
+};
+
+class Conjugate : public UnaryPrimitive {
+ public:
+  explicit Conjugate(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(Conjugate)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Contiguous : public UnaryPrimitive {
+ public:
+  explicit Contiguous(Stream stream, bool allow_col_major)
+      : UnaryPrimitive(stream), allow_col_major_(allow_col_major) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Contiguous)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  bool is_equivalent(const Primitive& other) const override;
+
+ private:
+  bool allow_col_major_;
+};
+
+class Convolution : public UnaryPrimitive {
+ public:
+  explicit Convolution(
+      Stream stream,
+      const std::vector<int>& kernel_strides,
+      const std::vector<int>& padding_lo,
+      const std::vector<int>& padding_hi,
+      const std::vector<int>& kernel_dilation,
+      const std::vector<int>& input_dilation,
+      const int groups = 1,
+      const bool flip = false)
+      : UnaryPrimitive(stream),
+        padding_lo_(padding_lo),
+        padding_hi_(padding_hi),
+        kernel_strides_(kernel_strides),
+        kernel_dilation_(kernel_dilation),
+        input_dilation_(input_dilation),
+        groups_(groups),
+        flip_(flip) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotangents,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(Convolution)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return std::make_tuple(
+        kernel_strides_,
+        padding_lo_,
+        padding_hi_,
+        kernel_dilation_,
+        input_dilation_,
+        groups_,
+        flip_);
+  }
+
+  static Shape conv_out_shape(
+      const Shape& in_shape,
+      const Shape& wt_shape,
+      const std::vector<int>& strides,
+      const std::vector<int>& pads_lo,
+      const std::vector<int>& pads_hi,
+      const std::vector<int>& kernel_dilation,
+      const std::vector<int>& input_dilation);
+
+ private:
+  std::vector<int> padding_lo_;
+  std::vector<int> padding_hi_;
+  std::vector<int> kernel_strides_;
+  std::vector<int> kernel_dilation_;
+  std::vector<int> input_dilation_;
+  int groups_;
+  bool flip_;
+};
+
+class Copy : public UnaryPrimitive {
+ public:
+  explicit Copy(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Copy)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+class Cos : public UnaryPrimitive {
+ public:
+  explicit Cos(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Cos)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Cosh : public UnaryPrimitive {
+ public:
+  explicit Cosh(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Cosh)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class CustomTransforms : public Primitive {
+ public:
+  explicit CustomTransforms(
+      Stream stream,
+      int num_outputs,
+      std::function<std::vector<array>(
+          const std::vector<array>&,
+          const std::vector<array>&,
+          const std::vector<array>&)> vjp,
+      std::function<std::vector<array>(
+          const std::vector<array>&,
+          const std::vector<array>&,
+          const std::vector<int>&)> jvp,
+      std::function<std::pair<std::vector<array>, std::vector<int>>(
+          const std::vector<array>&,
+          const std::vector<int>&)> vmap)
+      : Primitive(stream),
+        num_outputs_(num_outputs),
+        vjp_fun_(std::move(vjp)),
+        jvp_fun_(std::move(jvp)),
+        vmap_fun_(std::move(vmap)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_GRADS();
+  DEFINE_VMAP();
+  DEFINE_NAME(CustomTransforms);
+
+ private:
+  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
+
+  int num_outputs_;
+
+  std::function<std::vector<array>(
+      const std::vector<array>&,
+      const std::vector<array>&,
+      const std::vector<array>&)>
+      vjp_fun_;
+  std::function<std::vector<array>(
+      const std::vector<array>&,
+      const std::vector<array>&,
+      const std::vector<int>&)>
+      jvp_fun_;
+  std::function<std::pair<std::vector<array>, std::vector<int>>(
+      const std::vector<array>&,
+      const std::vector<int>&)>
+      vmap_fun_;
+};
+
+class Depends : public Primitive {
+ public:
+  explicit Depends(Stream stream) : Primitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  std::vector<array> vjp(
+      const std::vector<array>& primals,
+      const std::vector<array>& cotan,
+      const std::vector<int>& argnums,
+      const std::vector<array>& outputs) override;
+
+  DEFINE_NAME(Depends);
+
+ private:
+  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
+};
+
+class Divide : public UnaryPrimitive {
+ public:
+  explicit Divide(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Divide)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class DivMod : public Primitive {
+ public:
+  explicit DivMod(Stream stream) : Primitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(DivMod)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override {
+    return std::vector{inputs[0].shape(), inputs[0].shape()};
+  }
+};
+
+class Select : public UnaryPrimitive {
+ public:
+  explicit Select(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Select)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Remainder : public UnaryPrimitive {
+ public:
+  explicit Remainder(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Remainder)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Equal : public UnaryPrimitive {
+ public:
+  explicit Equal(Stream stream, bool equal_nan = false)
+      : UnaryPrimitive(stream), equal_nan_(equal_nan) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  const char* name() const override {
+    if (equal_nan_) {
+      return "NaNEqual";
+    } else {
+      return "Equal";
+    }
+  }
+  auto state() const {
+    return equal_nan_;
+  };
+
+ private:
+  bool equal_nan_;
+};
+
+class Erf : public UnaryPrimitive {
+ public:
+  explicit Erf(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Erf)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ErfInv : public UnaryPrimitive {
+ public:
+  explicit ErfInv(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ErfInv)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class MLX_API Exp : public UnaryPrimitive {
+ public:
+  explicit Exp(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Exp)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Expm1 : public UnaryPrimitive {
+ public:
+  explicit Expm1(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Expm1)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class ExpandDims : public UnaryPrimitive {
+ public:
+  explicit ExpandDims(Stream stream, std::vector<int> axes)
+      : UnaryPrimitive(stream), axes_(std::move(axes)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(ExpandDims)
+
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  bool is_equivalent(const Primitive& other) const override;
+
+  static Shape output_shape(const array& input, const std::vector<int>& axes);
+  auto state() const {
+    return axes_;
+  }
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+  std::vector<int> axes_;
+};
+
+class FFT : public UnaryPrimitive {
+ public:
+  explicit FFT(
+      Stream stream,
+      const std::vector<size_t>& axes,
+      bool inverse,
+      bool real)
+      : UnaryPrimitive(stream), axes_(axes), inverse_(inverse), real_(real) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(FFT)
+
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(axes_, inverse_, real_);
+  }
+
+ private:
+  std::vector<size_t> axes_;
+  bool inverse_;
+  bool real_;
+};
+
+class Flatten : public UnaryPrimitive {
+ public:
+  explicit Flatten(Stream stream, int start_axis, int end_axis)
+      : UnaryPrimitive(stream), start_axis_(start_axis), end_axis_(end_axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Flatten)
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  bool is_equivalent(const Primitive& other) const override;
+
+  static Shape output_shape(const array& input, int start_axis, int end_axis);
+  auto state() const {
+    return std::make_pair(start_axis_, end_axis_);
+  }
+
+ private:
+  int start_axis_;
+  int end_axis_;
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+class Floor : public UnaryPrimitive {
+ public:
+  explicit Floor(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Floor)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Full : public UnaryPrimitive {
+ public:
+  explicit Full(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Full)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Gather : public UnaryPrimitive {
+ public:
+  explicit Gather(Stream stream, std::vector<int> axes, Shape slice_sizes)
+      : UnaryPrimitive(stream),
+        axes_(std::move(axes)),
+        slice_sizes_(std::move(slice_sizes)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Gather)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  std::pair<std::vector<int>, Shape> state() const {
+    return {axes_, slice_sizes_};
+  }
+
+ private:
+  std::vector<int> axes_;
+  Shape slice_sizes_;
+};
+
+class GatherAxis : public UnaryPrimitive {
+ public:
+  explicit GatherAxis(Stream stream, int axis)
+      : UnaryPrimitive(stream), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(GatherAxis)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return axis_;
+  }
+
+ private:
+  int axis_;
+};
+
+class Greater : public UnaryPrimitive {
+ public:
+  explicit Greater(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Greater)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class GreaterEqual : public UnaryPrimitive {
+ public:
+  explicit GreaterEqual(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(GreaterEqual)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Hadamard : public UnaryPrimitive {
+ public:
+  explicit Hadamard(Stream stream, float scale)
+      : UnaryPrimitive(stream), scale_(scale) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Hadamard)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return scale_;
+  }
+
+ private:
+  float scale_;
+};
+
+class Imag : public UnaryPrimitive {
+ public:
+  explicit Imag(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Imag)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Less : public UnaryPrimitive {
+ public:
+  explicit Less(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Less)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class LessEqual : public UnaryPrimitive {
+ public:
+  explicit LessEqual(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(LessEqual)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Load : public UnaryPrimitive {
+ public:
+  explicit Load(
+      Stream stream,
+      std::shared_ptr<io::Reader> reader,
+      size_t offset,
+      bool swap_endianness = false)
+      : UnaryPrimitive(stream),
+        reader_(std::move(reader)),
+        offset_(offset),
+        swap_endianness_(swap_endianness) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_NAME(Load)
+
+ private:
+  std::shared_ptr<io::Reader> reader_;
+  size_t offset_;
+  bool swap_endianness_;
+};
+
+class Log : public UnaryPrimitive {
+ public:
+  enum Base { two, ten, e };
+
+  explicit Log(Stream stream, Base base)
+      : UnaryPrimitive(stream), base_(base) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  Base state() const {
+    return base_;
+  };
+
+  const char* name() const override {
+    switch (base_) {
+      case e:
+        return "Log";
+      case two:
+        return "Log2";
+      case ten:
+        return "Log10";
+    }
+    return "<unknwon Log>";
+  }
+
+ private:
+  Base base_;
+};
+
+class Log1p : public UnaryPrimitive {
+ public:
+  explicit Log1p(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Log1p)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class LogicalNot : public UnaryPrimitive {
+ public:
+  explicit LogicalNot(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(LogicalNot)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class LogicalAnd : public UnaryPrimitive {
+ public:
+  explicit LogicalAnd(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(LogicalAnd)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class LogicalOr : public UnaryPrimitive {
+ public:
+  explicit LogicalOr(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(LogicalOr)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class LogAddExp : public UnaryPrimitive {
+ public:
+  explicit LogAddExp(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(LogAddExp)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class LogSumExp : public UnaryPrimitive {
+ public:
+  explicit LogSumExp(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(LogSumExp)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+};
+
+class Matmul : public UnaryPrimitive {
+ public:
+  explicit Matmul(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_GRADS()
+  DEFINE_VMAP()
+  DEFINE_NAME(Matmul)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+};
+
+class Maximum : public UnaryPrimitive {
+ public:
+  explicit Maximum(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Maximum)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Minimum : public UnaryPrimitive {
+ public:
+  explicit Minimum(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Minimum)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Multiply : public UnaryPrimitive {
+ public:
+  explicit Multiply(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Multiply)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Negative : public UnaryPrimitive {
+ public:
+  explicit Negative(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Negative)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class NotEqual : public UnaryPrimitive {
+ public:
+  explicit NotEqual(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(NotEqual)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class NumberOfElements : public UnaryPrimitive {
+ public:
+  explicit NumberOfElements(
+      Stream stream,
+      std::vector<int> axes,
+      bool inverted,
+      Dtype dtype)
+      : UnaryPrimitive(stream),
+        axes_(std::move(axes)),
+        inverted_(inverted),
+        dtype_(dtype) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(NumberOfElements)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override {
+    return {{}};
+  }
+  std::tuple<std::vector<int>, bool, Dtype> state() const {
+    return {axes_, inverted_, dtype_};
+  }
+
+ private:
+  std::vector<int> axes_;
+  bool inverted_;
+  Dtype dtype_;
+
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+class Pad : public UnaryPrimitive {
+ public:
+  explicit Pad(
+      Stream stream,
+      const std::vector<int>& axes,
+      const Shape& low_pad_size,
+      const Shape& high_pad_size)
+      : UnaryPrimitive(stream),
+        axes_(axes),
+        low_pad_size_(low_pad_size),
+        high_pad_size_(high_pad_size) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Pad)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(axes_, low_pad_size_, high_pad_size_);
+  }
+
+ private:
+  std::vector<int> axes_;
+  Shape low_pad_size_;
+  Shape high_pad_size_;
+};
+
+class Partition : public UnaryPrimitive {
+ public:
+  explicit Partition(Stream stream, int kth, int axis)
+      : UnaryPrimitive(stream), kth_(kth), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Partition)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_pair(kth_, axis_);
+  };
+
+ private:
+  int kth_;
+  int axis_;
+};
+
+class Power : public UnaryPrimitive {
+ public:
+  explicit Power(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Power)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class QuantizedMatmul : public UnaryPrimitive {
+ public:
+  explicit QuantizedMatmul(
+      Stream stream,
+      int group_size,
+      int bits,
+      QuantizationMode mode,
+      bool transpose)
+      : UnaryPrimitive(stream),
+        group_size_(group_size),
+        bits_(bits),
+        mode_(mode),
+        transpose_(transpose) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(QuantizedMatmul)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return std::make_tuple(group_size_, bits_, mode_, transpose_);
+  }
+
+ private:
+  int group_size_;
+  int bits_;
+  QuantizationMode mode_;
+  bool transpose_;
+};
+
+class QQMatmul : public UnaryPrimitive {
+ public:
+  explicit QQMatmul(
+      Stream stream,
+      int group_size,
+      int bits,
+      QuantizationMode mode)
+      : UnaryPrimitive(stream),
+        group_size_(group_size),
+        bits_(bits),
+        mode_(mode) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  // DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(QQMatmul)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return std::make_tuple(group_size_, bits_, mode_);
+  }
+
+ private:
+  int group_size_;
+  int bits_;
+  QuantizationMode mode_;
+};
+
+class GatherQMM : public UnaryPrimitive {
+ public:
+  explicit GatherQMM(
+      Stream stream,
+      int group_size,
+      int bits,
+      QuantizationMode mode,
+      bool transpose,
+      bool left_sorted = false,
+      bool right_sorted = false)
+      : UnaryPrimitive(stream),
+        group_size_(group_size),
+        bits_(bits),
+        mode_(mode),
+        transpose_(transpose),
+        left_sorted_(left_sorted),
+        right_sorted_(right_sorted) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(GatherQMM)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(
+        group_size_, bits_, mode_, transpose_, left_sorted_, right_sorted_);
+  }
+
+ private:
+  int group_size_;
+  int bits_;
+  QuantizationMode mode_;
+  bool transpose_;
+  bool left_sorted_;
+  bool right_sorted_;
+};
+
+class RandomBits : public UnaryPrimitive {
+ public:
+  explicit RandomBits(Stream stream, const Shape& shape, int width)
+      : UnaryPrimitive(stream), shape_(shape), width_(width) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(RandomBits)
+  bool is_equivalent(const Primitive& other) const override;
+  std::pair<Shape, int> state() const {
+    return {shape_, width_};
+  };
+
+ private:
+  Shape shape_;
+  int width_;
+};
+
+class Real : public UnaryPrimitive {
+ public:
+  explicit Real(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Real)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Reshape : public UnaryPrimitive {
+ public:
+  explicit Reshape(Stream stream, const Shape& shape)
+      : UnaryPrimitive(stream), shape_(shape) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Reshape)
+  bool is_equivalent(const Primitive& other) const override;
+  Shape state() const {
+    return shape_;
+  };
+  static Shape output_shape(const array& input, Shape shape);
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+
+ private:
+  Shape shape_;
+};
+
+class MLX_API Reduce : public UnaryPrimitive {
+ public:
+  enum ReduceType { And, Or, Sum, Prod, Min, Max };
+
+  explicit Reduce(
+      Stream stream,
+      ReduceType reduce_type,
+      const std::vector<int>& axes)
+      : UnaryPrimitive(stream), reduce_type_(reduce_type), axes_(axes) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS();
+
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+
+  const char* name() const override {
+    switch (reduce_type_) {
+      case And:
+        return "And";
+      case Or:
+        return "Or";
+      case Sum:
+        return "Sum";
+      case Prod:
+        return "Prod";
+      case Min:
+        return "Min";
+      case Max:
+        return "Max";
+    }
+    return "<unknwon Reduce>";
+  }
+
+  bool is_equivalent(const Primitive& other) const override;
+  std::pair<ReduceType, std::vector<int>> state() const {
+    return {reduce_type_, axes_};
+  };
+
+ private:
+  ReduceType reduce_type_;
+  std::vector<int> axes_;
+};
+
+class Round : public UnaryPrimitive {
+ public:
+  explicit Round(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Round)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Scan : public UnaryPrimitive {
+ public:
+  enum ReduceType { Max, Min, Sum, Prod, LogAddExp };
+
+  explicit Scan(
+      Stream stream,
+      ReduceType reduce_type,
+      int axis,
+      bool reverse,
+      bool inclusive)
+      : UnaryPrimitive(stream),
+        reduce_type_(reduce_type),
+        axis_(axis),
+        reverse_(reverse),
+        inclusive_(inclusive) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS();
+
+  const char* name() const override {
+    switch (reduce_type_) {
+      case Sum:
+        return "CumSum";
+      case Prod:
+        return "CumProd";
+      case Min:
+        return "CumMin";
+      case Max:
+        return "CumMax";
+      case LogAddExp:
+        return "CumLogAddExp";
+    }
+    return "<unknwon Scan>";
+  }
+
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(reduce_type_, axis_, reverse_, inclusive_);
+  }
+
+ private:
+  ReduceType reduce_type_;
+  int axis_;
+  bool reverse_;
+  bool inclusive_;
+};
+
+class Scatter : public UnaryPrimitive {
+ public:
+  enum ReduceType { Max, Min, Sum, Prod, None };
+
+  explicit Scatter(
+      Stream stream,
+      ReduceType reduce_type,
+      const std::vector<int>& axes)
+      : UnaryPrimitive(stream), reduce_type_(reduce_type), axes_(axes) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP();
+  DEFINE_GRADS();
+
+  const char* name() const override {
+    switch (reduce_type_) {
+      case Sum:
+        return "Scatter Sum";
+      case Prod:
+        return "Scatter Prod";
+      case Min:
+        return "Scatter Min";
+      case Max:
+        return "Scatter Max";
+      case None:
+        return "Scatter";
+    }
+    return "<unknwon Scatter>";
+  }
+
+  bool is_equivalent(const Primitive& other) const override;
+  std::pair<ReduceType, std::vector<int>> state() const {
+    return {reduce_type_, axes_};
+  };
+
+ private:
+  ReduceType reduce_type_;
+  std::vector<int> axes_;
+};
+
+class ScatterAxis : public UnaryPrimitive {
+ public:
+  enum ReduceType { Sum, None };
+
+  explicit ScatterAxis(Stream stream, ReduceType reduce_type, int axis)
+      : UnaryPrimitive(stream), reduce_type_(reduce_type), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+
+  const char* name() const override {
+    switch (reduce_type_) {
+      case Sum:
+        return "ScatterAxis Sum";
+      case None:
+        return "ScatterAxis";
+    }
+    return "<unknwon ScatterAxis>";
+  }
+
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  std::pair<ReduceType, int> state() const {
+    return {reduce_type_, axis_};
+  }
+
+ private:
+  ReduceType reduce_type_;
+  int axis_;
+};
+
+class MaskedScatter : public UnaryPrimitive {
+ public:
+  explicit MaskedScatter(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP();
+  DEFINE_GRADS();
+  DEFINE_NAME(MaskedScatter);
+  DEFINE_DEFAULT_IS_EQUIVALENT();
+  DEFINE_INPUT_OUTPUT_SHAPE();
+};
+
+class Sigmoid : public UnaryPrimitive {
+ public:
+  explicit Sigmoid(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Sigmoid)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Sign : public UnaryPrimitive {
+ public:
+  explicit Sign(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Sign)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Sin : public UnaryPrimitive {
+ public:
+  explicit Sin(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Sin)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Sinh : public UnaryPrimitive {
+ public:
+  explicit Sinh(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Sinh)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Slice : public UnaryPrimitive {
+ public:
+  explicit Slice(
+      Stream stream,
+      const Shape& start_indices,
+      const Shape& end_indices,
+      const Shape& strides)
+      : UnaryPrimitive(stream),
+        start_indices_(start_indices),
+        end_indices_(end_indices),
+        strides_(strides) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Slice)
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_tuple(start_indices_, end_indices_, strides_);
+  }
+
+ private:
+  Shape start_indices_;
+  Shape end_indices_;
+  Shape strides_;
+};
+
+class SliceUpdate : public UnaryPrimitive {
+ public:
+  explicit SliceUpdate(
+      Stream stream,
+      const Shape& start_indices,
+      const Shape& end_indices,
+      const Shape& strides)
+      : UnaryPrimitive(stream),
+        start_indices_(start_indices),
+        end_indices_(end_indices),
+        strides_(strides) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(SliceUpdate)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return std::make_tuple(start_indices_, end_indices_, strides_);
+  }
+
+ private:
+  Shape start_indices_;
+  Shape end_indices_;
+  Shape strides_;
+};
+
+class DynamicSlice : public UnaryPrimitive {
+ public:
+  explicit DynamicSlice(Stream stream, std::vector<int> axes, Shape slice_size)
+      : UnaryPrimitive(stream),
+        axes_(std::move(axes)),
+        slice_size_(std::move(slice_size)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(DynamicSlice)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  auto state() const {
+    return std::make_pair(axes_, slice_size_);
+  }
+
+ private:
+  std::vector<int> axes_;
+  Shape slice_size_;
+};
+
+class DynamicSliceUpdate : public UnaryPrimitive {
+ public:
+  explicit DynamicSliceUpdate(Stream stream, std::vector<int> axes)
+      : UnaryPrimitive(stream), axes_(std::move(axes)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(DynamicSliceUpdate)
+  bool is_equivalent(const Primitive& other) const override;
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  auto state() const {
+    return axes_;
+  }
+
+ private:
+  std::vector<int> axes_;
+};
+
+class Softmax : public UnaryPrimitive {
+ public:
+  explicit Softmax(Stream stream, bool precise)
+      : UnaryPrimitive(stream), precise_(precise) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Softmax)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return precise_;
+  };
+
+ private:
+  bool precise_;
+};
+
+class Sort : public UnaryPrimitive {
+ public:
+  explicit Sort(Stream stream, int axis)
+      : UnaryPrimitive(stream), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Sort)
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return axis_;
+  }
+
+ private:
+  int axis_;
+};
+
+class Split : public Primitive {
+ public:
+  explicit Split(Stream stream, const Shape& indices, int axis)
+      : Primitive(stream), indices_(indices), axis_(axis) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Split)
+  bool is_equivalent(const Primitive& other) const override;
+  std::pair<Shape, int> state() const {
+    return {indices_, axis_};
+  };
+
+ private:
+  void eval(const std::vector<array>& inputs, std::vector<array>& outputs);
+
+  Shape indices_;
+  int axis_;
+};
+
+class Square : public UnaryPrimitive {
+ public:
+  explicit Square(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Square)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Sqrt : public UnaryPrimitive {
+ public:
+  explicit Sqrt(Stream stream, bool recip = false)
+      : UnaryPrimitive(stream), recip_(recip) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return recip_;
+  }
+
+  const char* name() const override {
+    if (recip_) {
+      return "Rsqrt";
+    } else {
+      return "Sqrt";
+    }
+  }
+
+ private:
+  bool recip_;
+};
+
+class StopGradient : public UnaryPrimitive {
+ public:
+  explicit StopGradient(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(StopGradient)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+class Subtract : public UnaryPrimitive {
+ public:
+  explicit Subtract(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Subtract)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Squeeze : public UnaryPrimitive {
+ public:
+  explicit Squeeze(Stream stream, std::vector<int> axes)
+      : UnaryPrimitive(stream), axes_(std::move(axes)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Squeeze)
+
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  bool is_equivalent(const Primitive& other) const override;
+
+  static Shape output_shape(const array& input, const std::vector<int>& axes);
+  auto state() const {
+    return axes_;
+  };
+
+ private:
+  void eval(const std::vector<array>& inputs, array& out);
+  std::vector<int> axes_;
+};
+
+class Tan : public UnaryPrimitive {
+ public:
+  explicit Tan(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Tan)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Tanh : public UnaryPrimitive {
+ public:
+  explicit Tanh(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Tanh)
+  DEFINE_DEFAULT_IS_EQUIVALENT()
+  DEFINE_INPUT_OUTPUT_SHAPE()
+};
+
+class Unflatten : public UnaryPrimitive {
+ public:
+  explicit Unflatten(Stream stream, int axis, Shape shape)
+      : UnaryPrimitive(stream), axis_(axis), shape_(std::move(shape)) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Unflatten)
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  bool is_equivalent(const Primitive& other) const override;
+
+  static Shape output_shape(const array& input, int axis, const Shape& shape);
+  auto state() const {
+    return std::make_pair(axis_, shape_);
+  }
+
+ private:
+  int axis_;
+  Shape shape_;
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+class View : public UnaryPrimitive {
+ public:
+  explicit View(Stream stream, Dtype dtype)
+      : UnaryPrimitive(stream), dtype_(dtype) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  const char* name() const override;
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return dtype_;
+  }
+
+ private:
+  Dtype dtype_;
+  mutable std::string name_;
+};
+
+class Transpose : public UnaryPrimitive {
+ public:
+  explicit Transpose(Stream stream, const std::vector<int>& axes)
+      : UnaryPrimitive(stream), axes_(axes) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_VMAP()
+  DEFINE_GRADS()
+  DEFINE_NAME(Transpose)
+  bool is_equivalent(const Primitive& other) const override;
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+  std::vector<int> state() const {
+    return axes_;
+  };
+
+ private:
+  std::vector<int> axes_;
+
+  void eval(const std::vector<array>& inputs, array& out);
+};
+
+/* QR Factorization primitive. */
+class QRF : public Primitive {
+ public:
+  explicit QRF(Stream stream) : Primitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(QRF)
+};
+
+/* SVD primitive. */
+class SVD : public Primitive {
+ public:
+  explicit SVD(Stream stream, bool compute_uv)
+      : Primitive(stream), compute_uv_(compute_uv) {}
+
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(SVD)
+  auto state() const {
+    return compute_uv_;
+  }
+
+ private:
+  bool compute_uv_;
+};
+
+/* Matrix inversion primitive. */
+class Inverse : public UnaryPrimitive {
+ public:
+  explicit Inverse(Stream stream, bool tri, bool upper)
+      : UnaryPrimitive(stream), tri_(tri), upper_(upper) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& output) override;
+  void eval_gpu(const std::vector<array>& inputs, array& output) override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(Inverse)
+  auto state() const {
+    return std::make_pair(tri_, upper_);
+  }
+
+ private:
+  bool tri_;
+  bool upper_;
+};
+
+class Cholesky : public UnaryPrimitive {
+ public:
+  explicit Cholesky(Stream stream, bool upper)
+      : UnaryPrimitive(stream), upper_(upper) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+  auto state() const {
+    return upper_;
+  }
+
+  DEFINE_VMAP()
+  DEFINE_NAME(Cholesky)
+
+ private:
+  bool upper_;
+};
+
+class Eig : public Primitive {
+ public:
+  explicit Eig(Stream stream, bool compute_eigenvectors)
+      : Primitive(stream), compute_eigenvectors_(compute_eigenvectors) {}
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(Eig)
+
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return compute_eigenvectors_;
+  }
+
+ private:
+  bool compute_eigenvectors_;
+};
+
+class Eigh : public Primitive {
+ public:
+  explicit Eigh(Stream stream, std::string uplo, bool compute_eigenvectors)
+      : Primitive(stream),
+        uplo_(std::move(uplo)),
+        compute_eigenvectors_(compute_eigenvectors) {}
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_VMAP()
+  DEFINE_NAME(Eigh)
+
+  std::vector<Shape> output_shapes(const std::vector<array>& inputs) override;
+
+  bool is_equivalent(const Primitive& other) const override;
+  auto state() const {
+    return std::make_pair(uplo_, compute_eigenvectors_);
+  }
+
+ private:
+  std::string uplo_;
+  bool compute_eigenvectors_;
+};
+
+/* LU Factorization primitive. */
+class LUF : public Primitive {
+ public:
+  explicit LUF(Stream stream) : Primitive(stream) {}
+  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
+      override;
+
+  DEFINE_NAME(LUF)
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/random.cpp b/macos/third_party/mlx/mlx/random.cpp
new file mode 100644
index 0000000..def3169
--- /dev/null
+++ b/macos/third_party/mlx/mlx/random.cpp
@@ -0,0 +1,492 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <cmath>
+#include <sstream>
+
+#include "mlx/linalg.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/random.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::random {
+
+KeySequence::KeySequence(uint64_t seed) : key_(key(seed)) {}
+
+void KeySequence::seed(uint64_t seed) {
+  key_ = key((seed));
+}
+
+array KeySequence::next() {
+  auto out = split(key_);
+  key_ = out.first;
+  return out.second;
+}
+
+void seed(uint64_t seed) {
+  KeySequence::default_().seed(seed);
+}
+
+array key(uint64_t seed) {
+  uint32_t k1 = static_cast<uint32_t>(seed >> 32);
+  uint32_t k2 = static_cast<uint32_t>(seed);
+  return array({k1, k2});
+}
+
+array bits(
+    const Shape& shape,
+    int width /* 4 */,
+    const std::optional<array>& key_ /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto key = key_ ? *key_ : KeySequence::default_().next();
+  if (key.dtype() != uint32) {
+    std::ostringstream msg;
+    msg << "[bits] Expected key type uint32 but received " << key.dtype()
+        << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (key.shape() != Shape{2}) {
+    std::ostringstream msg;
+    msg << "[bits] Expected key shape (2) but received " << key.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto get_dtype = [width]() {
+    switch (width) {
+      case 4:
+        return uint32;
+      case 2:
+        return uint16;
+      case 1:
+        return uint8;
+      default:
+        std::ostringstream msg;
+        msg << "[bits] Bit width must be in {1, 2, 4} but got " << width << ".";
+        throw std::invalid_argument(msg.str());
+    }
+  };
+  return array(
+      shape,
+      get_dtype(),
+      std::make_shared<RandomBits>(to_stream(s), shape, width),
+      {key});
+}
+
+std::pair<array, array> split(const array& key, StreamOrDevice s /* = {} */) {
+  auto stream = to_stream(s);
+  auto out = mlx::core::split(random::split(key, 2, stream), 2, stream);
+  return {reshape(out[0], {2}, stream), reshape(out[1], {2}, stream)};
+}
+
+array split(const array& key, int num, StreamOrDevice s /* = {} */) {
+  return bits({num, 2}, 4, key, s);
+}
+
+// Get the next representable value below 1.0 for half precision
+// floating point types (fp16, bf16)
+template <typename T>
+T below_one() {
+  T f = T(1.0);
+  uint16_t* m = (uint16_t*)&f;
+  *m -= 1;
+  return f;
+}
+
+array uniform(
+    const array& low,
+    const array& high,
+    const Shape& shape,
+    Dtype dtype /* = float32 */,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  if (!issubdtype(dtype, floating)) {
+    throw std::invalid_argument(
+        "[uniform] Can only generate uniform numbers with real "
+        "floating point type.");
+  }
+
+  auto stream = to_stream(s);
+  auto lo = astype(low, dtype, stream);
+  auto hi = astype(high, dtype, stream);
+  auto range = subtract(hi, lo, stream);
+  auto out_shape = broadcast_shapes(shape, range.shape());
+  if (out_shape != shape) {
+    std::ostringstream msg;
+    msg << "[uniform] Cannot generate random values of shape " << shape
+        << " from broadcasted shape " << out_shape << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // Get random values between [0, nextafter(1.0, 0.0)] since samples must
+  // be in [low, high)
+  auto get_upper = [&dtype]() {
+    switch (dtype) {
+      case float32:
+        return array(std::nextafter(1.0f, 0.0f), float32);
+      case float16:
+        return array(below_one<float16_t>(), float32);
+      case bfloat16:
+        return array(below_one<bfloat16_t>(), float32);
+      default:
+        throw std::runtime_error("[uniform] Unsupported type.");
+    }
+  };
+
+  auto upper = get_upper();
+  auto maxval = array(std::numeric_limits<uint32_t>::max(), float32);
+  auto out = bits(shape, size_of(float32), key, stream);
+  out = divide(out, maxval, stream);
+  out = astype(minimum(out, upper, stream), dtype, stream);
+  return add(multiply(range, out, stream), lo, stream);
+}
+
+array uniform(
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return uniform(
+      array(0.0, dtype), array(1.0, dtype), shape, dtype, key, to_stream(s));
+}
+
+inline array complex_normal(
+    Shape shape,
+    const std::optional<array>& loc,
+    const std::optional<array>& scale,
+    const std::optional<array>& key,
+    StreamOrDevice s) {
+  auto stream = to_stream(s);
+  auto low = array(std::nextafter(-1.0f, 0.0f), float32);
+  auto high = array(1.0f, float32);
+  shape.push_back(2);
+  auto samples =
+      erfinv(uniform(low, high, shape, float32, key, stream), stream);
+  samples = squeeze(view(samples, complex64, stream), -1, stream);
+  if (scale.has_value()) {
+    samples = multiply(*scale, samples, stream);
+  }
+  if (loc.has_value()) {
+    samples = add(*loc, samples, stream);
+  }
+  return samples;
+}
+
+array normal(
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& loc,
+    const std::optional<array>& scale,
+    const std::optional<array>& key,
+    StreamOrDevice s /* = {} */) {
+  if (dtype == complex64) {
+    return complex_normal(shape, loc, scale, key, s);
+  } else if (!issubdtype(dtype, floating)) {
+    throw std::invalid_argument(
+        "[normal] Can only generate uniform numbers with "
+        "floating point type.");
+  }
+
+  auto stream = to_stream(s);
+  auto low = array(std::nextafter(-1.0f, 0.0f), float32);
+  auto high = array(1.0f, float32);
+  auto samples = uniform(low, high, shape, float32, key, stream);
+  auto applied_scale = array(std::sqrt(2.0), dtype);
+  if (scale.has_value()) {
+    applied_scale =
+        multiply(applied_scale, astype(*scale, dtype, stream), stream);
+  }
+  samples = astype(erfinv(samples, stream), dtype, stream);
+  samples = multiply(applied_scale, samples, stream);
+  if (loc.has_value()) {
+    samples = add(astype(*loc, dtype, stream), samples, stream);
+  }
+  return samples;
+}
+
+array multivariate_normal(
+    const array& mean,
+    const array& cov,
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& key /* = nullopt */,
+    StreamOrDevice s) {
+  auto stream = to_stream(s);
+
+  if (dtype != float32) {
+    throw std::invalid_argument("[multivariate_normal] dtype must be float32.");
+  }
+
+  if (mean.ndim() < 1) {
+    throw std::invalid_argument(
+        "[multivariate_normal] mean must have at least one dimension.");
+  }
+
+  if (cov.ndim() < 2) {
+    throw std::invalid_argument(
+        "[multivariate_normal] cov must have at least two dimensions.");
+  }
+
+  auto n = mean.shape(-1);
+
+  // Check shapes compatibility of mean and cov
+  if (cov.shape(-1) != cov.shape(-2)) {
+    throw std::invalid_argument(
+        "[multivariate_normal] last two dimensions of cov must be equal.");
+  }
+  if (n != cov.shape(-1)) {
+    throw std::invalid_argument(
+        "[multivariate_normal] mean and cov must have compatible shapes.");
+  }
+
+  // Compute output shape
+  auto truncated_mean_shape =
+      Shape(mean.shape().begin(), mean.shape().end() - 1);
+  auto truncated_cov_shape = Shape(cov.shape().begin(), cov.shape().end() - 2);
+  auto output_shape =
+      broadcast_shapes(truncated_cov_shape, truncated_mean_shape);
+  output_shape = broadcast_shapes(output_shape, shape);
+  output_shape.push_back(n);
+
+  // Compute the square-root of the covariance matrix, using the SVD
+  auto covariance = astype(cov, float32, stream);
+  auto SVD = linalg::svd(covariance, true, stream);
+  auto std = astype(
+      matmul(
+          multiply(
+              SVD[0], expand_dims(sqrt(SVD[1], stream), -2, stream), stream),
+          SVD[2],
+          stream),
+      dtype,
+      stream);
+
+  // Generate standard the samples
+  auto standard_normal = normal(output_shape, dtype, 0.0, 1.0, key, stream);
+  auto scaled_out = squeeze(
+      matmul(expand_dims(standard_normal, -2, stream), std, stream),
+      -2,
+      stream);
+  return add(mean, scaled_out, stream);
+}
+
+array randint(
+    const array& low,
+    const array& high,
+    const Shape& shape,
+    Dtype dtype /* = int32 */,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  if (issubdtype(dtype, inexact)) {
+    throw std::invalid_argument(
+        "[randint] randint only accepts integer dtypes and bool.");
+  }
+  auto u = uniform(low, high, shape, float32, key, s);
+  return astype(maximum(u, low, s), dtype, s);
+}
+
+array bernoulli(
+    const array& p,
+    const Shape& shape,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  if (!issubdtype(p.dtype(), floating)) {
+    throw std::invalid_argument(
+        "[bernoulli] bernoulli probability `p` must be a float type.");
+  }
+
+  // Place p on the scale [0, nexthigher(UINT32_MAX)] so that if p >= 1.0 we
+  // get all true and if p <= 0.0 we get all false
+  auto upper = array(
+      std::nextafter(
+          static_cast<float>(std::numeric_limits<uint32_t>::max()),
+          std::numeric_limits<float>::max()),
+      float32);
+  auto res = less(bits(shape, key, s), multiply(p, upper, s), s);
+  if (res.shape() != shape) {
+    throw std::invalid_argument(
+        "[bernoulli] shape of `p` is incompatible with argument `shape`.");
+  }
+  return res;
+}
+
+array bernoulli(
+    const array& p,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return bernoulli(p, p.shape(), key, s);
+}
+
+array bernoulli(
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return bernoulli(array(0.5f), key, s);
+}
+
+array truncated_normal(
+    const array& lower,
+    const array& upper,
+    const Shape& shape,
+    Dtype dtype /* = float32 */,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  // Same as
+  // https://jax.readthedocs.io/en/latest/_modules/jax/_src/random.html#truncated_normal
+
+  if (!issubdtype(dtype, floating)) {
+    throw std::invalid_argument(
+        "[trunc_normal] trunc_normal only accepts floating point dtypes.");
+  }
+
+  auto sqrt2 = array(std::sqrt(2.0), dtype);
+  auto lower_t = astype(lower, dtype, s);
+  auto upper_t = astype(upper, dtype, s);
+  auto a = erf(divide(lower_t, sqrt2, s), s);
+  auto b = erf(divide(upper_t, sqrt2, s), s);
+  auto u = uniform(a, b, shape, dtype, key, s);
+  auto out = multiply(sqrt2, erfinv(u, s), s);
+
+  // Clip in bounds
+  return maximum(minimum(upper_t, out, s), lower_t, s);
+}
+
+array truncated_normal(
+    const array& lower,
+    const array& upper,
+    Dtype dtype /* = float32 */,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto shape = broadcast_shapes(lower.shape(), upper.shape());
+  return truncated_normal(lower, upper, shape, dtype, key, s);
+}
+
+array gumbel(
+    const Shape& shape,
+    Dtype dtype /* = float32 */,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  // -log(-log(uniform(shape)))
+  return negative(
+      log(negative(log(uniform(shape, dtype, key, s), s), s), s), s);
+}
+
+int get_valid_axis(int axis, int ndim) {
+  int ax = axis < 0 ? axis + ndim : axis;
+  if (ax < 0 || ax >= ndim) {
+    std::ostringstream msg;
+    msg << "[categorical] Invalid axis " << axis << " for logits with " << ndim
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  return ax;
+}
+
+array categorical_impl(
+    const array& logits,
+    int axis,
+    const Shape& shape,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s) {
+  auto gumbel_shape = shape;
+  auto offset = axis + shape.size() - logits.ndim() + 1;
+  gumbel_shape.insert(gumbel_shape.begin() + offset, logits.shape(axis));
+  auto g = gumbel(gumbel_shape, float32, key, s);
+  return argmax(add(g, logits, s), offset, false, s);
+}
+
+array categorical(
+    const array& logits,
+    int axis,
+    const Shape& shape,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  // Validate and normalize axis
+  axis = get_valid_axis(axis, logits.ndim());
+
+  // Check that shape broadcasts with reduce(logits, axis)
+  auto reduced_shape = logits.shape();
+  reduced_shape.erase(reduced_shape.begin() + axis);
+  if (broadcast_shapes(shape, reduced_shape) != shape) {
+    std::ostringstream msg;
+    msg << "[categorical] Requested shape " << shape
+        << " is not broadcast compatible with reduced logits shape"
+        << reduced_shape << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  return categorical_impl(logits, axis, shape, key, s);
+}
+
+array categorical(
+    const array& logits_,
+    int axis,
+    int num_samples,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  axis = get_valid_axis(axis, logits_.ndim());
+  auto logits = expand_dims(logits_, -1);
+  auto shape = logits.shape();
+  shape.erase(shape.begin() + axis);
+  shape.back() = num_samples;
+  return categorical_impl(logits, axis, shape, key, s);
+}
+
+array categorical(
+    const array& logits,
+    int axis /* = -1 */,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  axis = get_valid_axis(axis, logits.ndim());
+  auto shape = logits.shape();
+  shape.erase(shape.begin() + axis);
+  return categorical_impl(logits, axis, shape, key, s);
+}
+
+array laplace(
+    const Shape& shape,
+    Dtype dtype,
+    const float loc /* = 0.0 */,
+    const float scale /* = 1.0 */,
+    const std::optional<array>& key /*= nullopt */,
+    StreamOrDevice s /* = {} */) {
+  if (!issubdtype(dtype, floating)) {
+    throw std::invalid_argument(
+        "[laplace] Can only generate uniform numbers with real"
+        "floating point type.");
+  }
+
+  auto stream = to_stream(s);
+  auto low = array(std::nextafter(-1.0f, 0.0f), float32);
+  auto high = array(1.0f, float32);
+  auto samples = uniform(low, high, shape, float32, key, stream);
+  // Use inverse CDF to generate Laplacian noise
+  samples = multiply(
+      sign(samples, stream),
+      log1p(
+          multiply(array(-1.0f, dtype), abs(samples, stream), stream), stream),
+      stream);
+  samples = astype(samples, dtype, stream);
+
+  if (scale != 1.0) {
+    samples = multiply(array(scale, dtype), samples, stream);
+  }
+  if (loc != 0.0) {
+    samples = add(array(loc, dtype), samples, stream);
+  }
+  return samples;
+}
+
+array permutation(
+    const array& x,
+    int axis /* = 0 */,
+    const std::optional<array>& key /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return take(x, permutation(x.shape(axis), key, s), axis, s);
+}
+
+array permutation(
+    int x,
+    const std::optional<array>& key /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return argsort(bits({x}, key, s), s);
+}
+
+} // namespace mlx::core::random
diff --git a/macos/third_party/mlx/mlx/random.h b/macos/third_party/mlx/mlx/random.h
new file mode 100644
index 0000000..a23c255
--- /dev/null
+++ b/macos/third_party/mlx/mlx/random.h
@@ -0,0 +1,283 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <chrono>
+#include <optional>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/stream.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::random {
+
+class MLX_API KeySequence {
+ public:
+  explicit KeySequence(uint64_t seed);
+
+  void seed(uint64_t seed);
+  array next();
+
+  // static default
+  static KeySequence& default_() {
+    static KeySequence ks(get_current_time_seed());
+    return ks;
+  }
+
+ private:
+  array key_;
+  static uint64_t get_current_time_seed() {
+    auto now = std::chrono::system_clock::now();
+    return std::chrono::duration_cast<std::chrono::milliseconds>(
+               now.time_since_epoch())
+        .count();
+  }
+};
+
+/** Get a PRNG key from a seed. */
+MLX_API array key(uint64_t seed);
+
+/** Seed the default PRNG key. */
+MLX_API void seed(uint64_t seed);
+
+/** Generate an array with type uint32 filled with random bits. */
+MLX_API array bits(
+    const Shape& shape,
+    int width,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+inline array bits(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return bits(shape, 4, key, s);
+}
+
+/** Split the rng key into a pair of keys. */
+MLX_API std::pair<array, array> split(const array& key, StreamOrDevice s = {});
+
+/** Split the rng key into `num` keys. */
+MLX_API array split(const array& key, int num, StreamOrDevice s = {});
+
+/** Generate uniform random numbers between low and high. */
+MLX_API array uniform(
+    const array& low,
+    const array& high,
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+template <typename T, typename U>
+array uniform(
+    T low,
+    U high,
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return uniform(array(low), array(high), shape, dtype, key, to_stream(s));
+}
+
+/** Generate uniform random numbers between 0 and 1. */
+MLX_API array uniform(
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+inline array uniform(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return uniform(shape, float32, key, s);
+}
+
+/** Generate samples from the standard normal distribution. */
+MLX_API array normal(
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& loc,
+    const std::optional<array>& scale,
+    const std::optional<array>& key,
+    StreamOrDevice s = {});
+inline array normal(
+    const Shape& shape,
+    Dtype dtype,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  auto loc_ = loc == 0 ? std::nullopt : std::make_optional(array(loc, dtype));
+  auto scale_ =
+      scale == 1 ? std::nullopt : std::make_optional(array(scale, dtype));
+  return normal(shape, dtype, loc_, scale_, key, s);
+}
+inline array normal(
+    const Shape& shape,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return normal(shape, float32, loc, scale, key, s);
+}
+inline array normal(
+    const Shape& shape,
+    const Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return normal(shape, dtype, std::nullopt, std::nullopt, key, s);
+}
+inline array normal(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return normal(shape, float32, std::nullopt, std::nullopt, key, s);
+}
+
+/** Generate samples from a multivariate normal distribution. **/
+MLX_API array multivariate_normal(
+    const array& mean,
+    const array& cov,
+    const Shape& shape,
+    Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Generate integer samples uniformly at random */
+MLX_API array randint(
+    const array& low,
+    const array& high,
+    const Shape& shape,
+    Dtype dtype = int32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+template <typename T, typename U>
+array randint(
+    T low,
+    U high,
+    const Shape& shape,
+    Dtype dtype = int32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return randint(array(low), array(high), shape, dtype, key, to_stream(s));
+}
+
+/** Generate binary variables with probability to be true equal to p */
+MLX_API array bernoulli(
+    const array& p,
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+MLX_API array bernoulli(
+    const array& p,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+template <typename T>
+array bernoulli(
+    T p,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return bernoulli(array(p), key, s);
+}
+
+template <typename T>
+array bernoulli(
+    T p,
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return bernoulli(array(p), shape, key, s);
+}
+
+MLX_API array bernoulli(
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array truncated_normal(
+    const array& lower,
+    const array& upper,
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array truncated_normal(
+    const array& lower,
+    const array& upper,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array gumbel(
+    const Shape& shape,
+    Dtype dtype = float32,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array categorical(
+    const array& logits,
+    int axis,
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array categorical(
+    const array& logits_,
+    int axis,
+    int num_samples,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+MLX_API array categorical(
+    const array& logits,
+    int axis = -1,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+/** Generate samples from the laplace distribution. */
+MLX_API array laplace(
+    const Shape& shape,
+    Dtype dtype,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+inline array laplace(
+    const Shape& shape,
+    const float loc,
+    const float scale,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return laplace(shape, float32, loc, scale, key, s);
+}
+inline array laplace(
+    const Shape& shape,
+    const Dtype dtype,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return laplace(shape, dtype, 0.0, 1.0, key, s);
+}
+inline array laplace(
+    const Shape& shape,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {}) {
+  return laplace(shape, float32, 0.0, 1.0, key, s);
+}
+
+/* Randomly permute the elements of x along the given axis. */
+MLX_API array permutation(
+    const array& x,
+    int axis = 0,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+/* A random permutation of `arange(x)` */
+MLX_API array permutation(
+    int x,
+    const std::optional<array>& key = std::nullopt,
+    StreamOrDevice s = {});
+
+} // namespace mlx::core::random
diff --git a/macos/third_party/mlx/mlx/scheduler.cpp b/macos/third_party/mlx/mlx/scheduler.cpp
new file mode 100644
index 0000000..87ff98e
--- /dev/null
+++ b/macos/third_party/mlx/mlx/scheduler.cpp
@@ -0,0 +1,73 @@
+// Copyright © 2023 Apple Inc.
+
+#include "mlx/scheduler.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/gpu/eval.h"
+
+namespace mlx::core {
+
+Stream default_stream(Device d) {
+  if (!gpu::is_available() && d == Device::gpu) {
+    throw std::invalid_argument(
+        "[default_stream] Cannot get gpu stream without gpu backend.");
+  }
+  return scheduler::scheduler().get_default_stream(d);
+}
+
+void set_default_stream(Stream s) {
+  if (!gpu::is_available() && s.device == Device::gpu) {
+    throw std::invalid_argument(
+        "[set_default_stream] Cannot set gpu stream without gpu backend.");
+  }
+  return scheduler::scheduler().set_default_stream(s);
+}
+
+Stream get_stream(int index) {
+  return scheduler::scheduler().get_stream(index);
+}
+
+Stream new_stream(Device d) {
+  if (!gpu::is_available() && d == Device::gpu) {
+    throw std::invalid_argument(
+        "[new_stream] Cannot make gpu stream without gpu backend.");
+  }
+  return scheduler::scheduler().new_stream(d);
+}
+
+Stream new_stream() {
+  return scheduler::scheduler().new_stream(default_device());
+}
+
+void synchronize(Stream s) {
+  if (s.device == mlx::core::Device::cpu) {
+    auto p = std::make_shared<std::promise<void>>();
+    std::future<void> f = p->get_future();
+    scheduler::enqueue(s, [p = std::move(p)]() { p->set_value(); });
+    f.wait();
+  } else {
+    gpu::synchronize(s);
+  }
+}
+
+void synchronize() {
+  synchronize(default_stream(default_device()));
+}
+
+namespace scheduler {
+
+/** A singleton scheduler to manage devices, streams, and task execution. */
+Scheduler& scheduler() {
+  // Leak the scheduler on Windows to avoid joining threads on exit, can be
+  // removed after Visual Studio fixes bug:
+  // https://developercommunity.visualstudio.com/t/1654756
+#ifdef _WIN32
+  static Scheduler* scheduler = new Scheduler;
+  return *scheduler;
+#else
+  static Scheduler scheduler;
+  return scheduler;
+#endif
+}
+
+} // namespace scheduler
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/scheduler.h b/macos/third_party/mlx/mlx/scheduler.h
new file mode 100644
index 0000000..4b94945
--- /dev/null
+++ b/macos/third_party/mlx/mlx/scheduler.h
@@ -0,0 +1,189 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <atomic>
+#include <future>
+#include <queue>
+#include <thread>
+#include <unordered_map>
+
+#include "mlx/api.h"
+#include "mlx/backend/gpu/eval.h"
+#include "mlx/device.h"
+#include "mlx/stream.h"
+
+namespace mlx::core::scheduler {
+
+struct StreamThread {
+  std::mutex mtx;
+  std::queue<std::function<void()>> q;
+  std::condition_variable cond;
+  bool stop;
+  std::thread thread;
+
+  StreamThread() : stop(false), thread(&StreamThread::thread_fn, this) {}
+
+  ~StreamThread() {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      stop = true;
+    }
+    cond.notify_one();
+    thread.join();
+  }
+
+  void thread_fn() {
+    while (true) {
+      std::function<void()> task;
+      {
+        std::unique_lock<std::mutex> lk(mtx);
+        cond.wait(lk, [this] { return !this->q.empty() || this->stop; });
+        if (q.empty() && stop) {
+          return;
+        }
+        task = std::move(q.front());
+        q.pop();
+      }
+
+      task();
+    }
+  }
+
+  template <typename F>
+  void enqueue(F&& f) {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      if (stop) {
+        throw std::runtime_error(
+            "Cannot enqueue work after stream is stopped.");
+      }
+      q.emplace(std::forward<F>(f));
+    }
+    cond.notify_one();
+  }
+};
+
+class Scheduler {
+ public:
+  Scheduler() : n_active_tasks_(0) {
+    if (is_available(Device::gpu)) {
+      default_streams_.insert({Device::gpu, new_stream(Device::gpu)});
+    }
+    default_streams_.insert({Device::cpu, new_stream(Device::cpu)});
+  }
+
+  // Not copyable or moveable
+  Scheduler(const Scheduler&) = delete;
+  Scheduler(Scheduler&&) = delete;
+  Scheduler& operator=(const Scheduler&) = delete;
+  Scheduler& operator=(Scheduler&&) = delete;
+
+  Stream new_stream(const Device& d) {
+    streams_.emplace_back(streams_.size(), d);
+    if (d == Device::gpu) {
+      threads_.push_back(nullptr);
+      gpu::new_stream(streams_.back());
+    } else {
+      threads_.push_back(new StreamThread{});
+    }
+    return streams_.back();
+  }
+
+  template <typename F>
+  void enqueue(const Stream& stream, F&& f);
+
+  Stream get_default_stream(const Device& d) const {
+    return default_streams_.at(d.type);
+  }
+  Stream get_stream(int index) const {
+    return streams_.at(index);
+  }
+
+  void set_default_stream(const Stream& s) {
+    default_streams_.at(s.device.type) = s;
+  }
+
+  void notify_new_task(const Stream& stream) {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      n_active_tasks_++;
+    }
+    completion_cv.notify_all();
+  }
+
+  void notify_task_completion(const Stream& stream) {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      n_active_tasks_--;
+    }
+    completion_cv.notify_all();
+  }
+
+  int n_active_tasks() const {
+    return n_active_tasks_;
+  }
+
+  void wait_for_one() {
+    std::unique_lock<std::mutex> lk(mtx);
+    int n_tasks_old = n_active_tasks();
+    if (n_tasks_old > 1) {
+      completion_cv.wait(lk, [this, n_tasks_old] {
+        return this->n_active_tasks() < n_tasks_old;
+      });
+    }
+  }
+
+  ~Scheduler() {
+    for (auto s : streams_) {
+      try {
+        synchronize(s);
+      } catch (const std::runtime_error&) {
+        // ignore errors if synch fails
+      }
+    }
+    for (auto t : threads_) {
+      if (t != nullptr) {
+        delete t;
+      }
+    }
+  }
+
+ private:
+  int n_active_tasks_;
+  std::vector<StreamThread*> threads_;
+  std::vector<Stream> streams_;
+  std::unordered_map<Device::DeviceType, Stream> default_streams_;
+  std::condition_variable completion_cv;
+  std::mutex mtx;
+};
+
+template <typename F>
+void Scheduler::enqueue(const Stream& stream, F&& f) {
+  threads_[stream.index]->enqueue(std::forward<F>(f));
+}
+
+MLX_API Scheduler& scheduler();
+
+template <typename F>
+void enqueue(const Stream& stream, F&& f) {
+  scheduler().enqueue(stream, std::forward<F>(f));
+}
+
+inline int n_active_tasks() {
+  return scheduler().n_active_tasks();
+}
+
+inline void notify_new_task(const Stream& stream) {
+  scheduler().notify_new_task(stream);
+}
+
+inline void notify_task_completion(const Stream& stream) {
+  scheduler().notify_task_completion(stream);
+}
+
+inline void wait_for_one() {
+  scheduler().wait_for_one();
+}
+
+} // namespace mlx::core::scheduler
diff --git a/macos/third_party/mlx/mlx/small_vector.h b/macos/third_party/mlx/mlx/small_vector.h
new file mode 100644
index 0000000..143101c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/small_vector.h
@@ -0,0 +1,540 @@
+// Copyright © 2025 Apple Inc.
+// Copyright © 2018 the V8 project authors.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//     * Redistributions of source code must retain the above copyright
+//       notice, this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above
+//       copyright notice, this list of conditions and the following
+//       disclaimer in the documentation and/or other materials provided
+//       with the distribution.
+//     * Neither the name of Google Inc. nor the names of its
+//       contributors may be used to endorse or promote products derived
+//       from this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <type_traits>
+#include <utility>
+
+namespace mlx::core {
+
+#if defined(__has_builtin)
+#define MLX_HAS_BUILTIN(x) __has_builtin(x)
+#else
+#define MLX_HAS_BUILTIN(x) 0
+#endif
+
+#if defined(__has_attribute)
+#define MLX_HAS_ATTRIBUTE(x) __has_attribute(x)
+#else
+#define MLX_HAS_ATTRIBUTE(x) 0
+#endif
+
+#if MLX_HAS_BUILTIN(__builtin_expect)
+#define MLX_LIKELY(condition) (__builtin_expect(!!(condition), 1))
+#define MLX_UNLIKELY(condition) (__builtin_expect(!!(condition), 0))
+#else
+#define MLX_LIKELY(condition) (condition)
+#define MLX_UNLIKELY(condition) (condition)
+#endif
+
+#if MLX_HAS_ATTRIBUTE(noinline)
+#define MLX_NOINLINE __attribute__((noinline))
+#else
+#define MLX_NOINLINE
+#endif
+
+template <typename T, typename = void>
+struct is_iterator : std::false_type {};
+
+template <typename T>
+struct is_iterator<
+    T,
+    std::void_t<
+        typename std::iterator_traits<T>::difference_type,
+        typename std::iterator_traits<T>::iterator_category,
+        typename std::iterator_traits<T>::pointer,
+        typename std::iterator_traits<T>::reference,
+        typename std::iterator_traits<T>::value_type>> : std::true_type {};
+
+template <typename T>
+constexpr bool is_iterator_v = is_iterator<T>::value;
+
+// Minimal SmallVector implementation. Uses inline storage first, switches to
+// dynamic storage when it overflows.
+//
+// Notes:
+// * The default inline storage size is MAX_NDIM, as it is mainly used for
+//   shapes and strides, users should choose a better size for other cases.
+// * The data() returns real address even for empty vector.
+// * The pointer returned by data() will change after moving the vector as it
+//   points to the inline storage.
+// * For trivial elements the storage will not be default constructed,
+//   i.e. SmallVector<int>(10) will not be filled with 0 by default.
+template <typename T, size_t kSize = 10, typename Allocator = std::allocator<T>>
+class SmallVector {
+ public:
+  using value_type = T;
+  using reference = T&;
+  using const_reference = const T&;
+  using iterator = T*;
+  using const_iterator = const T*;
+  using difference_type = std::ptrdiff_t;
+  using size_type = std::size_t;
+
+  SmallVector() = default;
+
+  explicit SmallVector(const Allocator& allocator) : allocator_(allocator) {}
+
+  explicit SmallVector(size_t size, const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    resize(size);
+  }
+
+  SmallVector(
+      size_t size,
+      const T& initial_value,
+      const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    resize(size, initial_value);
+  }
+
+  SmallVector(
+      std::initializer_list<T> init,
+      const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    if (init.size() > capacity()) {
+      grow(init.size());
+    }
+    assert(capacity() >= init.size()); // sanity check
+    std::uninitialized_move(init.begin(), init.end(), begin_);
+    end_ = begin_ + init.size();
+  }
+
+  template <typename Iter, typename = std::enable_if_t<is_iterator_v<Iter>>>
+  SmallVector(Iter begin, Iter end, const Allocator& allocator = Allocator())
+      : allocator_(allocator) {
+    size_t size = std::distance(begin, end);
+    if (size > capacity()) {
+      grow(size);
+    }
+    assert(capacity() >= size); // sanity check
+    std::uninitialized_copy(begin, end, begin_);
+    end_ = begin_ + size;
+  }
+
+  SmallVector(const SmallVector& other) : allocator_(other.allocator_) {
+    *this = other;
+  }
+  SmallVector(const SmallVector& other, const Allocator& allocator)
+      : allocator_(allocator) {
+    *this = other;
+  }
+  SmallVector(SmallVector&& other) : allocator_(std::move(other.allocator_)) {
+    *this = std::move(other);
+  }
+  SmallVector(SmallVector&& other, const Allocator& allocator)
+      : allocator_(allocator) {
+    *this = std::move(other);
+  }
+
+  ~SmallVector() {
+    free_storage();
+  }
+
+  SmallVector& operator=(const SmallVector& other) {
+    if (this == &other) {
+      return *this;
+    }
+    size_t other_size = other.size();
+    if (capacity() < other_size) {
+      // Create large-enough heap-allocated storage.
+      free_storage();
+      begin_ = allocator_.allocate(other_size);
+      end_of_storage_ = begin_ + other_size;
+      std::uninitialized_copy(other.begin_, other.end_, begin_);
+    } else if constexpr (kHasTrivialElement) {
+      std::copy(other.begin_, other.end_, begin_);
+    } else {
+      ptrdiff_t to_copy =
+          std::min(static_cast<ptrdiff_t>(other_size), end_ - begin_);
+      std::copy(other.begin_, other.begin_ + to_copy, begin_);
+      if (other.begin_ + to_copy < other.end_) {
+        std::uninitialized_copy(
+            other.begin_ + to_copy, other.end_, begin_ + to_copy);
+      } else {
+        std::destroy_n(begin_ + to_copy, size() - to_copy);
+      }
+    }
+    end_ = begin_ + other_size;
+    return *this;
+  }
+
+  SmallVector& operator=(SmallVector&& other) {
+    if (this == &other) {
+      return *this;
+    }
+    if (other.is_big()) {
+      free_storage();
+      begin_ = other.begin_;
+      end_ = other.end_;
+      end_of_storage_ = other.end_of_storage_;
+    } else {
+      assert(capacity() >= other.size()); // sanity check
+      size_t other_size = other.size();
+      if constexpr (kHasTrivialElement) {
+        std::move(other.begin_, other.end_, begin_);
+      } else {
+        ptrdiff_t to_move =
+            std::min(static_cast<ptrdiff_t>(other_size), end_ - begin_);
+        std::move(other.begin_, other.begin_ + to_move, begin_);
+        if (other.begin_ + to_move < other.end_) {
+          std::uninitialized_move(
+              other.begin_ + to_move, other.end_, begin_ + to_move);
+        } else {
+          std::destroy_n(begin_ + to_move, size() - to_move);
+        }
+      }
+      end_ = begin_ + other_size;
+    }
+    other.reset_to_inline_storage();
+    return *this;
+  }
+
+  bool operator==(const SmallVector& other) const {
+    if (size() != other.size()) {
+      return false;
+    }
+    return std::equal(begin_, end_, other.begin_);
+  }
+
+  bool operator!=(const SmallVector& other) const {
+    return !(*this == other);
+  }
+
+  T* data() {
+    return begin_;
+  }
+  const T* data() const {
+    return begin_;
+  }
+
+  iterator begin() {
+    return begin_;
+  }
+  const_iterator begin() const {
+    return begin_;
+  }
+
+  iterator end() {
+    return end_;
+  }
+  const_iterator end() const {
+    return end_;
+  }
+
+  const_iterator cbegin() const {
+    return begin_;
+  }
+
+  const_iterator cend() const {
+    return end_;
+  }
+
+  auto rbegin() {
+    return std::make_reverse_iterator(end_);
+  }
+  auto rbegin() const {
+    return std::make_reverse_iterator(end_);
+  }
+
+  auto rend() {
+    return std::make_reverse_iterator(begin_);
+  }
+  auto rend() const {
+    return std::make_reverse_iterator(begin_);
+  }
+
+  size_t size() const {
+    return end_ - begin_;
+  }
+  bool empty() const {
+    return end_ == begin_;
+  }
+  size_t capacity() const {
+    return end_of_storage_ - begin_;
+  }
+
+  T& front() {
+    assert(size() != 0);
+    return begin_[0];
+  }
+  const T& front() const {
+    assert(size() != 0);
+    return begin_[0];
+  }
+
+  T& back() {
+    assert(size() != 0);
+    return end_[-1];
+  }
+  const T& back() const {
+    assert(size() != 0);
+    return end_[-1];
+  }
+
+  T& at(size_t index) {
+    if (index >= size()) {
+      throw std::out_of_range("SmallVector out of range.");
+    }
+    return begin_[index];
+  }
+  const T& at(size_t index) const {
+    return const_cast<SmallVector*>(this)->at(index);
+  }
+
+  T& operator[](size_t index) {
+    assert(size() > index);
+    return begin_[index];
+  }
+  const T& operator[](size_t index) const {
+    return const_cast<SmallVector*>(this)->operator[](index);
+  }
+
+  template <typename... Args>
+  void emplace_back(Args&&... args) {
+    if (MLX_UNLIKELY(end_ == end_of_storage_)) {
+      grow();
+    }
+    void* storage = end_;
+    end_ += 1;
+    new (storage) T(std::forward<Args>(args)...);
+  }
+
+  void push_back(T x) {
+    emplace_back(std::move(x));
+  }
+
+  void pop_back(size_t count = 1) {
+    assert(size() >= count);
+    end_ -= count;
+    std::destroy_n(end_, count);
+  }
+
+  iterator insert(iterator pos, T value) {
+    return insert(pos, static_cast<size_t>(1), std::move(value));
+  }
+
+  iterator insert(iterator pos, size_t count, T value) {
+    assert(pos <= end_);
+    size_t offset = pos - begin_;
+    size_t old_size = size();
+    resize(old_size + count);
+    pos = begin_ + offset;
+    iterator old_end = begin_ + old_size;
+    assert(old_end <= end_);
+    std::move_backward(pos, old_end, end_);
+    if constexpr (kHasTrivialElement) {
+      std::fill_n(pos, count, value);
+    } else {
+      std::fill_n(pos + 1, count - 1, value);
+      *pos = std::move(value);
+    }
+    return pos;
+  }
+
+  template <typename Iter, typename = std::enable_if_t<is_iterator_v<Iter>>>
+  iterator insert(iterator pos, Iter begin, Iter end) {
+    if constexpr (std::is_same_v<std::decay_t<Iter>, iterator>) {
+      // The implementation can not take overlapping range.
+      assert(!(begin >= pos && begin < pos + std::distance(begin, end)));
+      assert(!(end > pos && end <= pos + std::distance(begin, end)));
+    }
+
+    assert(pos <= end_);
+    size_t offset = pos - begin_;
+    size_t count = std::distance(begin, end);
+    size_t old_size = size();
+    resize(old_size + count);
+    pos = begin_ + offset;
+    iterator old_end = begin_ + old_size;
+    assert(old_end <= end_);
+    std::move_backward(pos, old_end, end_);
+    std::copy(begin, end, pos);
+    return pos;
+  }
+
+  iterator insert(iterator pos, std::initializer_list<const T> values) {
+    return insert(pos, values.begin(), values.end());
+  }
+
+  iterator erase(iterator erase_start, iterator erase_end) {
+    assert(erase_start >= begin_);
+    assert(erase_start <= erase_end);
+    assert(erase_end <= end_);
+    iterator new_end = std::move(erase_end, end_, erase_start);
+    std::destroy_n(new_end, std::distance(new_end, end_));
+    end_ = new_end;
+    return erase_start;
+  }
+
+  iterator erase(iterator pos) {
+    return erase(pos, pos + 1);
+  }
+
+  void resize(size_t new_size) {
+    if (new_size > capacity()) {
+      grow(new_size);
+    }
+    T* new_end = begin_ + new_size;
+    if constexpr (!kHasTrivialElement) {
+      if (new_end > end_) {
+        std::uninitialized_default_construct(end_, new_end);
+      } else {
+        std::destroy_n(new_end, end_ - new_end);
+      }
+    }
+    end_ = new_end;
+  }
+
+  void resize(size_t new_size, const T& initial_value) {
+    if (new_size > capacity()) {
+      grow(new_size);
+    }
+    T* new_end = begin_ + new_size;
+    if (new_end > end_) {
+      std::uninitialized_fill(end_, new_end, initial_value);
+    } else {
+      std::destroy_n(new_end, end_ - new_end);
+    }
+    end_ = new_end;
+  }
+
+  void reserve(size_t new_capacity) {
+    if (new_capacity > capacity()) {
+      grow(new_capacity);
+    }
+  }
+
+  // Clear without reverting back to inline storage.
+  void clear() {
+    std::destroy_n(begin_, end_ - begin_);
+    end_ = begin_;
+  }
+
+ private:
+  // Grows the backing store by a factor of two, and at least to {min_capacity}.
+  // TODO: Move to private after removing external code using this method.
+  MLX_NOINLINE void grow(size_t min_capacity = 0) {
+    size_t new_capacity = std::max(min_capacity, 2 * capacity());
+    // Round up to power of 2.
+    new_capacity--;
+    new_capacity |= new_capacity >> 1;
+    new_capacity |= new_capacity >> 2;
+    new_capacity |= new_capacity >> 4;
+    new_capacity |= new_capacity >> 8;
+    new_capacity |= new_capacity >> 16;
+    if constexpr (sizeof(size_t) == sizeof(uint64_t)) {
+      new_capacity |= new_capacity >> 32;
+    }
+    new_capacity++;
+
+    T* new_storage = allocator_.allocate(new_capacity);
+    if (new_storage == nullptr) {
+      throw std::bad_alloc();
+    }
+
+    size_t in_use = end_ - begin_;
+    std::uninitialized_move(begin_, end_, new_storage);
+    free_storage();
+    begin_ = new_storage;
+    end_ = new_storage + in_use;
+    end_of_storage_ = new_storage + new_capacity;
+  }
+
+  MLX_NOINLINE void free_storage() {
+    std::destroy_n(begin_, end_ - begin_);
+    if (is_big()) {
+      allocator_.deallocate(begin_, end_of_storage_ - begin_);
+    }
+  }
+
+  // Clear and go back to inline storage. Dynamic storage is *not* freed. For
+  // internal use only.
+  void reset_to_inline_storage() {
+    if constexpr (!kHasTrivialElement) {
+      if (!is_big())
+        std::destroy_n(begin_, end_ - begin_);
+    }
+    begin_ = inline_storage_begin();
+    end_ = begin_;
+    end_of_storage_ = begin_ + kSize;
+  }
+
+  bool is_big() const {
+    return begin_ != inline_storage_begin();
+  }
+
+  T* inline_storage_begin() {
+    return reinterpret_cast<T*>(inline_storage_);
+  }
+  const T* inline_storage_begin() const {
+    return reinterpret_cast<const T*>(inline_storage_);
+  }
+
+  Allocator allocator_;
+
+  // Invariants:
+  // 1. The elements in the range between `begin_` (included) and `end_` (not
+  //    included) will be initialized at all times.
+  // 2. All other elements outside the range, both in the inline storage and in
+  //    the dynamic storage (if it exists), will be uninitialized at all times.
+
+  T* begin_ = inline_storage_begin();
+  T* end_ = begin_;
+  T* end_of_storage_ = begin_ + kSize;
+
+  alignas(T) char inline_storage_[sizeof(T) * kSize];
+
+  static constexpr bool kHasTrivialElement =
+      std::is_trivially_copyable<T>::value &&
+      std::is_trivially_destructible<T>::value;
+};
+
+template <typename>
+struct is_vector : std::false_type {};
+
+template <typename T, size_t Size, typename Allocator>
+struct is_vector<SmallVector<T, Size, Allocator>> : std::true_type {};
+
+template <typename T, typename Allocator>
+struct is_vector<std::vector<T, Allocator>> : std::true_type {};
+
+template <typename Vec>
+inline constexpr bool is_vector_v = is_vector<Vec>::value;
+
+#undef MLX_HAS_BUILTIN
+#undef MLX_HAS_ATTRIBUTE
+#undef MLX_LIKELY
+#undef MLX_UNLIKELY
+#undef MLX_NOINLINE
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/stream.h b/macos/third_party/mlx/mlx/stream.h
new file mode 100644
index 0000000..9bd14db
--- /dev/null
+++ b/macos/third_party/mlx/mlx/stream.h
@@ -0,0 +1,42 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/api.h"
+#include "mlx/device.h"
+
+namespace mlx::core {
+
+struct MLX_API Stream {
+  int index;
+  Device device;
+  explicit Stream(int index, Device device) : index(index), device(device) {}
+};
+
+/** Get the default stream for the given device. */
+MLX_API Stream default_stream(Device d);
+
+/** Make the stream the default for its device. */
+MLX_API void set_default_stream(Stream s);
+
+/** Make a new stream on the given device. */
+MLX_API Stream new_stream(Device d);
+
+/** Get the stream with the given index. */
+MLX_API Stream get_stream(int index);
+
+inline bool operator==(const Stream& lhs, const Stream& rhs) {
+  return lhs.index == rhs.index;
+}
+
+inline bool operator!=(const Stream& lhs, const Stream& rhs) {
+  return !(lhs == rhs);
+}
+
+/* Synchronize with the default stream. */
+MLX_API void synchronize();
+
+/* Synchronize with the provided stream. */
+MLX_API void synchronize(Stream);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/threadpool.h b/macos/third_party/mlx/mlx/threadpool.h
new file mode 100644
index 0000000..b0e56d0
--- /dev/null
+++ b/macos/third_party/mlx/mlx/threadpool.h
@@ -0,0 +1,133 @@
+// This code was modified from https://github.com/progschj/ThreadPool
+// The original License is copied below:
+//
+// Copyright (c) 2012 Jakob Progsch, Václav Zeman
+// This software is provided 'as-is', without any express or implied
+// warranty. In no event will the authors be held liable for any damages
+// arising from the use of this software.
+//
+// Permission is granted to anyone to use this software for any purpose,
+// including commercial applications, and to alter it and redistribute it
+// freely, subject to the following restrictions:
+//
+//    1. The origin of this software must not be misrepresented; you must not
+//    claim that you wrote the original software. If you use this software
+//    in a product, an acknowledgment in the product documentation would be
+//    appreciated but is not required.
+//
+//    2. Altered source versions must be plainly marked as such, and must not be
+//    misrepresented as being the original software.
+//
+//    3. This notice may not be removed or altered from any source
+//    distribution.
+#pragma once
+
+#include <condition_variable>
+#include <functional>
+#include <future>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <stdexcept>
+#include <thread>
+#include <vector>
+
+class ThreadPool {
+ public:
+  ThreadPool(size_t);
+  template <class F, class... Args>
+  auto enqueue(F&& f, Args&&... args)
+      -> std::future<typename std::invoke_result_t<F, Args...>>;
+  void resize(size_t);
+  ~ThreadPool();
+
+ private:
+  void stop_and_wait();
+  void start_threads(size_t);
+
+  std::vector<std::thread> workers;
+  std::queue<std::function<void()>> tasks;
+  std::mutex queue_mutex;
+  std::condition_variable condition;
+  bool stop;
+};
+
+inline ThreadPool::ThreadPool(size_t threads) : stop(false) {
+  start_threads(threads);
+}
+
+template <class F, class... Args>
+auto ThreadPool::enqueue(F&& f, Args&&... args)
+    -> std::future<typename std::invoke_result_t<F, Args...>> {
+  using return_type = typename std::invoke_result_t<F, Args...>;
+
+  auto task = std::make_shared<std::packaged_task<return_type()>>(
+      std::bind(std::forward<F>(f), std::forward<Args>(args)...));
+
+  std::future<return_type> res = task->get_future();
+  {
+    std::unique_lock<std::mutex> lock(queue_mutex);
+
+    if (stop) {
+      throw std::runtime_error(
+          "[ThreadPool::enqueue] Not allowed on stopped ThreadPool");
+    }
+
+    tasks.emplace([task]() { (*task)(); });
+  }
+  condition.notify_one();
+  return res;
+}
+
+inline void ThreadPool::resize(size_t threads) {
+  if (workers.size() == threads) {
+    return;
+  }
+
+  if (workers.size() > threads) {
+    stop_and_wait();
+  }
+  start_threads(threads - workers.size());
+}
+
+inline ThreadPool::~ThreadPool() {
+  stop_and_wait();
+}
+
+inline void ThreadPool::stop_and_wait() {
+  // Stop the current threads and wait until they finish
+  {
+    std::unique_lock<std::mutex> lock(queue_mutex);
+    stop = true;
+  }
+  condition.notify_all();
+  for (std::thread& worker : workers) {
+    worker.join();
+  }
+
+  // Reset the member variables so that the threadpool is reusable
+  stop = false;
+  workers.clear();
+}
+
+inline void ThreadPool::start_threads(size_t threads) {
+  for (size_t i = 0; i < threads; ++i) {
+    workers.emplace_back([this] {
+      for (;;) {
+        std::function<void()> task;
+
+        {
+          std::unique_lock<std::mutex> lock(this->queue_mutex);
+          this->condition.wait(
+              lock, [this] { return this->stop || !this->tasks.empty(); });
+          if (this->stop && this->tasks.empty())
+            return;
+          task = std::move(this->tasks.front());
+          this->tasks.pop();
+        }
+
+        task();
+      }
+    });
+  }
+}
diff --git a/macos/third_party/mlx/mlx/transforms.cpp b/macos/third_party/mlx/mlx/transforms.cpp
new file mode 100644
index 0000000..4967c50
--- /dev/null
+++ b/macos/third_party/mlx/mlx/transforms.cpp
@@ -0,0 +1,1065 @@
+// Copyright © 2023-2024 Apple Inc.
+#include <algorithm>
+#include <deque>
+#include <future>
+#include <numeric>
+#include <set>
+#include <sstream>
+#include <stack>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/backend/cpu/eval.h"
+#include "mlx/backend/gpu/eval.h"
+#include "mlx/fence.h"
+#include "mlx/memory.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/scheduler.h"
+#include "mlx/transforms.h"
+#include "mlx/transforms_impl.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+static constexpr int MAX_ACTIVE_TASKS = 10;
+
+/* This class is only meant to be used in eval
+ * for synchronizing with the main thread. */
+class Synchronizer : public Primitive {
+ public:
+  explicit Synchronizer(Stream stream) : Primitive(stream) {}
+
+  void eval_cpu(const std::vector<array>&, std::vector<array>&) override {}
+  void eval_gpu(const std::vector<array>&, std::vector<array>&) override {}
+
+  DEFINE_NAME(Synchronize);
+};
+
+// Initialize the static tracing members from transforms_impl.h
+//
+// These are used to implement the in_tracing() function the returns true if we
+// are currently under a function transformation and the retain_graph()
+// function which returns true if we are forced to retain the graph during
+// evaluation.
+std::vector<std::pair<char, char>>& detail::InTracing::trace_stack() {
+  static std::vector<std::pair<char, char>> trace_stack_;
+  return trace_stack_;
+}
+int detail::InTracing::grad_counter{0};
+int detail::RetainGraph::tracing_counter{0};
+
+array eval_impl(std::vector<array> outputs, bool async) {
+  std::deque<array> tape;
+
+  // Make an effort to choose a good output stream
+  Stream stream = default_stream(default_device());
+  for (auto& o : outputs) {
+    if (o.status() == array::Status::unscheduled && o.has_primitive()) {
+      stream = o.primitive().stream();
+      break;
+    }
+  }
+
+  // Map of array id that needs fence and stream it's computed on
+  std::unordered_map<uintptr_t, std::pair<uint32_t, bool>> needs_fence;
+
+  auto synchronizer = array(
+      {}, bool_, std::make_shared<Synchronizer>(stream), std::move(outputs));
+
+  // Stream fences for inter-stream synchronization
+  std::unordered_map<uint32_t, Fence> fences;
+
+  // Stream events for synchronization after eval
+  std::unordered_map<uint32_t, Event> events;
+  {
+    auto e = Event{stream};
+    e.set_value(1);
+    synchronizer.attach_event(e);
+    events.emplace(stream.index, std::move(e));
+  }
+
+  {
+    // Record the degree of each input
+    std::unordered_map<std::uintptr_t, int> cache;
+
+    std::stack<std::pair<std::reference_wrapper<array>, int>> dfs;
+    dfs.emplace(synchronizer, 0);
+    while (!dfs.empty()) {
+      auto& [a_ref, idx] = dfs.top();
+      auto& a = a_ref.get();
+
+      if (idx < a.inputs().size()) {
+        // Add an input, and continue
+        auto& in = a.inputs()[idx++];
+
+        if (in.status() == array::Status::unscheduled) {
+          if (async && in.is_tracer()) {
+            throw std::invalid_argument(
+                "[async_eval] Not allowed inside a graph transformation.");
+          }
+          if (!in.has_primitive()) {
+            if (in.is_tracer()) {
+              throw std::invalid_argument(
+                  "[eval] Attempting to eval an array during function"
+                  " transformations like compile or vmap is not allowed.");
+            }
+            throw std::runtime_error(
+                "[eval] Attempting to eval an array without a primitive.\n"
+                "If you are compiling a function, make sure all the inputs "
+                "and outputs are captured:\n"
+                "https://ml-explore.github.io/mlx/build/html/usage/compile.html#pure-functions.\n"
+                "If you are not using compile, this may be a bug. "
+                "Please file an issue here:\n"
+                "https://github.com/ml-explore/mlx/issues.");
+          }
+          if (a.primitive().stream() != in.primitive().stream()) {
+            bool device_switch =
+                a.primitive().stream().device != in.primitive().stream().device;
+            auto [it, inserted] = needs_fence.emplace(
+                in.id(),
+                std::make_pair(in.primitive().stream().index, device_switch));
+            if (!inserted) {
+              it->second.second |= device_switch;
+            }
+          }
+        }
+
+        // All siblings have the same degree
+        auto cache_it = cache.find(in.id());
+        if (cache_it == cache.end()) {
+          dfs.emplace(in, 0);
+          cache.insert({in.id(), 1});
+          for (auto& s : in.siblings()) {
+            cache.insert({s.id(), 1});
+          }
+        } else {
+          cache_it->second++;
+          for (auto& s : in.siblings()) {
+            cache[s.id()]++;
+          }
+        }
+        continue;
+      }
+      if ((a.status() != array::Status::unscheduled) && !a.is_tracer() &&
+          a.has_primitive()) {
+        // If the array is evaluated and is no longer a tracer, detach it
+        a.detach();
+      }
+      dfs.pop();
+    }
+
+    // Build the tape in BFS order with a width limit
+    int max_width = env::bfs_max_width();
+    dfs = std::stack<std::pair<std::reference_wrapper<array>, int>>();
+    tape.push_back(synchronizer);
+    for (int i = 0; !cache.empty() && (i < tape.size() || !dfs.empty());) {
+      auto& a = (i >= tape.size()) ? dfs.top().first.get() : tape[i];
+      int j = 0;
+      if (i >= tape.size()) {
+        j = dfs.top().second;
+        dfs.pop();
+      } else {
+        i++;
+      }
+      for (; j < a.inputs().size(); ++j) {
+        auto& in = a.inputs()[j];
+        if (in.status() != array::Status::unscheduled) {
+          continue;
+        }
+
+        // If the width limit is exceeded, push the array on the stack
+        // and go down a level
+        if ((tape.size() - i) >= max_width) {
+          dfs.emplace(a, j);
+          break;
+        }
+
+        auto it = cache.find(in.id());
+        it->second -= 1;
+
+        if (it->second != 0) {
+          for (auto& s : in.siblings()) {
+            cache[s.id()] -= 1;
+          }
+          continue;
+        }
+
+        // Remove input and siblings from cache
+        cache.erase(it);
+        for (auto& s : in.siblings()) {
+          cache.erase(s.id());
+        }
+
+        tape.push_back(in);
+      }
+    }
+  }
+
+  std::unordered_set<int> open_streams;
+  while (!tape.empty()) {
+    auto arr = std::move(tape.back());
+    tape.pop_back();
+
+    auto stream = arr.primitive().stream();
+    open_streams.insert(stream.index);
+
+    if (async) {
+      // Lookup corresponding event
+      auto e = events.find(stream.index);
+      if (e == events.end()) {
+        e = events.emplace(stream.index, Event{stream}).first;
+      }
+      e->second.set_value(1);
+      arr.attach_event(e->second);
+      for (auto& s : arr.siblings()) {
+        s.attach_event(e->second);
+      }
+    }
+
+    for (auto& in : arr.inputs()) {
+      if (auto it = needs_fence.find(in.id()); it != needs_fence.end()) {
+        // Use fence to wait within a single eval
+        // Get the input array's stream fence and wait on the
+        // output arrays stream
+        fences[it->second.first].wait(stream, in);
+      } else if (in.event().valid()) {
+        if (in.event().is_signaled()) {
+          in.detach_event();
+        } else if (in.event().stream() != stream) {
+          // Use event to wait across async eval
+          in.event().wait(stream);
+        }
+      }
+    }
+
+    if (arr.primitive().device() == Device::gpu) {
+      gpu::eval(arr);
+    } else {
+      cpu::eval(arr);
+    }
+
+    if (scheduler::n_active_tasks() > MAX_ACTIVE_TASKS ||
+        (get_active_memory() > get_memory_limit() &&
+         scheduler::n_active_tasks() > 0)) {
+      // Commit any open streams
+      for (auto i : open_streams) {
+        auto s = get_stream(i);
+        if (s.device == Device::gpu) {
+          gpu::finalize(s);
+        }
+      }
+      scheduler::wait_for_one();
+      while (get_active_memory() > get_memory_limit() &&
+             scheduler::n_active_tasks() > 0) {
+        scheduler::wait_for_one();
+      }
+    }
+
+    auto maybe_update_fence = [&fences, &needs_fence, stream](const array& a) {
+      if (auto nf = needs_fence.find(a.id()); nf != needs_fence.end()) {
+        auto it = fences.find(stream.index);
+        if (it == fences.end()) {
+          it = fences.emplace(stream.index, Fence{stream}).first;
+        }
+        it->second.update(stream, a, nf->second.second);
+      }
+    };
+
+    arr.set_status(array::Status::evaluated);
+    // TODO Maybe always want the fence coherent kernel in the same cbuf
+    // as the other kernels?
+    maybe_update_fence(arr);
+    for (auto& sib : arr.siblings()) {
+      sib.set_status(array::Status::evaluated);
+      maybe_update_fence(sib);
+    }
+    if (!arr.is_tracer()) {
+      arr.detach();
+    }
+  }
+
+  // Signal the event in its stream
+  for (auto i : open_streams) {
+    auto s = get_stream(i);
+    if (auto e = events.find(i); e != events.end()) {
+      e->second.signal(s);
+    }
+    if (s.device == Device::gpu) {
+      gpu::finalize(s);
+    }
+  }
+
+  return synchronizer;
+}
+
+void async_eval(std::vector<array> outputs) {
+  if (outputs.empty()) {
+    return;
+  }
+
+  if (std::none_of(outputs.begin(), outputs.end(), [](array& x) {
+        return x.status() == array::Status::unscheduled;
+      })) {
+    return;
+  }
+
+  eval_impl(std::move(outputs), true);
+}
+
+void eval(std::vector<array> outputs) {
+  if (outputs.empty()) {
+    return;
+  }
+
+  if (std::none_of(outputs.begin(), outputs.end(), [](array& x) {
+        return x.status() == array::Status::unscheduled;
+      })) {
+    for (auto& x : outputs) {
+      x.wait();
+    }
+    return;
+  }
+
+  eval_impl(std::move(outputs), false).wait();
+}
+
+std::pair<std::vector<array>, std::vector<array>> vjp(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& primals,
+    const std::vector<array>& cotans,
+    const std::vector<int>& argnums) {
+  // Set the global tracing flag.
+  detail::InTracing in_tracing{false, true};
+
+  // Make tracers from given primals
+  std::vector<array> primals_;
+  for (auto& p : primals) {
+    auto s = p.has_primitive() ? p.primitive().stream()
+                               : default_stream(default_device());
+    primals_.push_back(copy(p, s)); // Does not do a deep copy
+    primals_.back().set_tracer(true);
+  }
+
+  // Pass tracer primals through the function
+  // Any variables that depend on the primals are marked as tracers
+  auto outputs = fun(primals_);
+
+  // Map outputs to passed cotans while ignoring the outputs
+  // that have stop_gradient called on them
+  int cotan_index = 0;
+  std::vector<std::pair<int, int>> output_cotan_pairs;
+  for (int i = 0; i < outputs.size(); ++i) {
+    auto& out = outputs[i];
+    if (out.has_primitive()) {
+      if (auto& p = out.primitive(); typeid(p) == typeid(StopGradient)) {
+        continue;
+      }
+    }
+    if (cotan_index >= cotans.size()) {
+      std::ostringstream msg;
+      msg << "[vjp] Number of outputs to compute gradients for ("
+          << outputs.size() << ") does not match number of cotangents ("
+          << cotans.size() << ").";
+      throw std::invalid_argument(msg.str());
+    }
+    if (out.shape() != cotans[cotan_index].shape()) {
+      std::ostringstream msg;
+      msg << "[vjp] Output shape " << out.shape()
+          << " does not match cotangent shape " << cotans[cotan_index].shape()
+          << ".";
+      if (outputs.size() == 1 && out.size() == 1) {
+        msg << " If you are using grad your function must return a scalar.";
+      }
+      throw std::invalid_argument(msg.str());
+    }
+    output_cotan_pairs.emplace_back(i, cotan_index++);
+  }
+
+  // Topologically sort the compute graph, add graph nodes
+  // to the tape which need a gradient.
+  std::unordered_set<std::uintptr_t> cache;
+  std::unordered_set<std::uintptr_t> calc_grad;
+  for (int i = 0, j = 0; i < primals_.size(); ++i) {
+    auto& primal = primals_[i];
+    primal.set_tracer(false);
+    cache.insert(primal.id());
+    if (j < argnums.size() && argnums[j] == i) {
+      j++;
+      calc_grad.insert(primal.id());
+    }
+  }
+
+  std::vector<array> tape;
+
+  std::function<void(array&)> recurse;
+  recurse = [&](auto& a) {
+    // Check if visited and add to cache if not
+    if (auto inserted = cache.insert(a.id()); !inserted.second) {
+      return;
+    }
+    a.set_tracer(false);
+    for (auto& s : a.siblings()) {
+      s.set_tracer(false);
+      cache.insert(s.id());
+    }
+
+    for (auto& input : a.inputs()) {
+      recurse(input);
+    }
+
+    // Stop grad
+    if (a.has_primitive()) {
+      if (auto& p = a.primitive(); typeid(p) == typeid(StopGradient)) {
+        return;
+      }
+    }
+
+    // Calculate gradient if any inputs require gradient
+    for (auto& input : a.inputs()) {
+      if (calc_grad.find(input.id()) != calc_grad.end()) {
+        tape.push_back(a);
+        calc_grad.insert(a.id());
+        for (auto& s : a.siblings()) {
+          calc_grad.insert(s.id());
+        }
+        break;
+      }
+    }
+  };
+
+  for (auto out : outputs) {
+    recurse(out);
+  }
+
+  // Run the tape backwards, computing vector-jacobian
+  // products for each primitive
+  std::unordered_map<std::uintptr_t, array> cotan_map;
+  for (auto [out_idx, cotan_idx] : output_cotan_pairs) {
+    auto& o = outputs[out_idx];
+    auto s = o.has_primitive() ? o.primitive().stream()
+                               : default_stream(default_device());
+    cotan_map.insert({o.id(), astype(cotans[cotan_idx], o.dtype(), s)});
+  }
+  for (auto it = tape.rbegin(); it != tape.rend(); ++it) {
+    auto& a = *it;
+
+    // Get the arguments whose gradients are needed
+    std::vector<int> argnums;
+    for (int i = 0; i < a.inputs().size(); ++i) {
+      if (calc_grad.find(a.inputs()[i].id()) != calc_grad.end()) {
+        argnums.push_back(i);
+      }
+    }
+
+    // Check if any of the array or its siblings have cotangents,
+    // if not, we can skip this primitive
+    auto outputs = a.outputs();
+    bool has_cotans =
+        std::any_of(outputs.cbegin(), outputs.cend(), [&cotan_map](auto& s) {
+          return cotan_map.find(s.id()) != cotan_map.end();
+        });
+    if (!has_cotans) {
+      continue;
+    }
+
+    auto s = a.primitive().stream();
+    std::vector<array> cotangents{};
+    for (auto& o : outputs) {
+      if (auto cotan_it = cotan_map.find(o.id()); cotan_it != cotan_map.end()) {
+        cotangents.push_back(cotan_map.extract(cotan_it).mapped());
+      } else {
+        cotangents.push_back(zeros_like(o, s));
+      }
+    }
+
+    std::vector<array> vjps;
+    {
+      detail::RetainGraph retain;
+      vjps = a.primitive().vjp(a.inputs(), cotangents, argnums, outputs);
+    }
+    // Accumulate the vector-jacobian products for each input
+    for (int i = 0; i < argnums.size(); ++i) {
+      auto in_id = a.inputs()[argnums[i]].id();
+      if (auto cotan_it = cotan_map.find(in_id); cotan_it != cotan_map.end()) {
+        cotan_it->second = add(cotan_it->second, vjps[i], s);
+      } else {
+        cotan_map.insert({in_id, vjps[i]});
+      }
+    }
+  }
+  std::vector<array> vjps;
+  for (auto arg : argnums) {
+    auto& primal = primals_[arg];
+    if (auto cotan_it = cotan_map.find(primal.id());
+        cotan_it != cotan_map.end()) {
+      vjps.push_back(cotan_it->second);
+    } else {
+      auto s = primal.has_primitive() ? primal.primitive().stream()
+                                      : default_stream(default_device());
+      vjps.push_back(zeros_like(primal, s));
+    }
+  }
+  return {outputs, vjps};
+}
+
+std::pair<std::vector<array>, std::vector<array>> vjp(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& primals,
+    const std::vector<array>& cotans) {
+  std::vector<int> argnums(primals.size());
+  std::iota(argnums.begin(), argnums.end(), 0);
+  return vjp(fun, primals, cotans, argnums);
+}
+
+std::pair<array, array> vjp(
+    const std::function<array(const array&)>& fun,
+    const array& primal,
+    const array& cotan) {
+  auto vec_fun = [fun](const std::vector<array>& inputs) {
+    return std::vector<array>{fun(inputs[0])};
+  };
+  auto [outputs, vjps] = vjp(vec_fun, {primal}, {cotan});
+  return {outputs[0], vjps[0]};
+}
+
+std::pair<std::vector<array>, std::vector<array>> jvp(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents) {
+  // Set the global tracing flag.
+  detail::InTracing in_tracing{false, true};
+
+  if (primals.size() != tangents.size()) {
+    throw std::invalid_argument(
+        "[jvp] Number of inputs does not match number of tangents.");
+  }
+  for (int i = 0; i < primals.size(); ++i) {
+    if (primals[i].shape() != tangents[i].shape()) {
+      throw std::invalid_argument(
+          "[jvp] Input shape does not match shape of tangent.");
+    }
+  }
+
+  std::vector<array> primals_;
+  for (auto& p : primals) {
+    auto s = p.has_primitive() ? p.primitive().stream()
+                               : default_stream(default_device());
+    primals_.push_back(copy(p, s)); // Does not do a deep copy
+    primals_.back().set_tracer(true);
+  }
+  auto outputs = fun(primals_);
+
+  // Topologically sort the compute graph, record outputs
+  // in the tape if a gradient is needed.
+  std::unordered_set<std::uintptr_t> cache;
+  std::unordered_set<std::uintptr_t> calc_grad;
+  for (auto& primal : primals_) {
+    primal.set_tracer(false);
+    calc_grad.insert(primal.id());
+    cache.insert(primal.id());
+  }
+
+  std::vector<array> tape;
+
+  std::function<void(array&)> recurse;
+  recurse = [&](auto& a) {
+    // Check if visited and add to cache if not
+    if (auto inserted = cache.insert(a.id()); !inserted.second) {
+      return;
+    }
+    a.set_tracer(false);
+    for (auto& s : a.siblings()) {
+      s.set_tracer(false);
+      cache.insert(s.id());
+    }
+
+    for (auto input : a.inputs()) {
+      recurse(input);
+    }
+
+    // Stop grad
+    if (a.has_primitive()) {
+      if (auto& p = a.primitive(); typeid(p) == typeid(StopGradient)) {
+        return;
+      }
+    }
+
+    // Calculate gradient if any inputs require gradient
+    for (auto& input : a.inputs()) {
+      if (calc_grad.find(input.id()) != calc_grad.end()) {
+        tape.push_back(a);
+        calc_grad.insert(a.id());
+        for (auto& s : a.siblings()) {
+          calc_grad.insert(s.id());
+        }
+        break;
+      }
+    }
+  };
+
+  for (auto out : outputs) {
+    recurse(out);
+  }
+
+  std::unordered_map<std::uintptr_t, array> tan_map;
+  for (int i = 0; i < primals_.size(); ++i) {
+    tan_map.insert({primals_[i].id(), tangents[i]});
+  }
+
+  for (auto& a : tape) {
+    // Get the arguments used in the jvp
+    std::vector<int> argnums;
+    std::vector<array> tangents;
+    for (int i = 0; i < a.inputs().size(); ++i) {
+      if (auto it = tan_map.find(a.inputs()[i].id()); it != tan_map.end()) {
+        argnums.push_back(i);
+        tangents.push_back(it->second);
+      }
+    }
+
+    auto jvps = a.primitive().jvp(a.inputs(), tangents, argnums);
+    auto outputs = a.outputs();
+    for (int i = 0; i < jvps.size(); ++i) {
+      tan_map.insert({outputs[i].id(), jvps[i]});
+    }
+  }
+
+  std::vector<array> jvps;
+  for (auto& out : outputs) {
+    if (auto it = tan_map.find(out.id()); it != tan_map.end()) {
+      jvps.push_back(it->second);
+    } else {
+      auto s = out.has_primitive() ? out.primitive().stream()
+                                   : default_stream(default_device());
+      jvps.push_back(zeros_like(out, s));
+    }
+  }
+  return {outputs, jvps};
+}
+
+std::pair<array, array> jvp(
+    const std::function<array(const array&)>& fun,
+    const array& primal,
+    const array& tangent) {
+  auto vec_fun = [fun](const std::vector<array>& inputs) {
+    return std::vector<array>{fun(inputs[0])};
+  };
+  auto [outputs, jvps] = jvp(vec_fun, {primal}, {tangent});
+  return {outputs[0], jvps[0]};
+}
+
+ValueAndGradFn value_and_grad(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<int>& argnums) {
+  if (argnums.empty()) {
+    throw std::invalid_argument("[grad] Must specify at least one argument.");
+  }
+  return [fun, argnums](const std::vector<array>& inputs) {
+    std::set<int> args;
+    for (auto& arg : argnums) {
+      args.insert(arg < 0 ? arg + inputs.size() : arg);
+    }
+    if (args.size() != argnums.size()) {
+      throw std::invalid_argument(
+          "[grad] Repeat argument number not allowed in grad.");
+    }
+    if (*args.begin() < 0 || *args.rbegin() >= inputs.size()) {
+      std::ostringstream msg;
+      msg << "[grad] Invalid argument number for function with "
+          << inputs.size() << " inputs.";
+      throw std::invalid_argument(msg.str());
+    }
+    std::vector<int> sorted_argnums(args.begin(), args.end());
+
+    auto gfun = [&fun](const std::vector<array>& inputs) {
+      auto outputs = fun(inputs);
+      for (int i = 1; i < outputs.size(); i++) {
+        auto& out = outputs[i];
+        auto s = out.has_primitive() ? out.primitive().stream()
+                                     : default_stream(default_device());
+        outputs[i] = stop_gradient(out, s);
+      }
+      return outputs;
+    };
+
+    // Set the incoming gradient to float32, vjp will cast it to the output type
+    auto [outputs, grads] = vjp(gfun, inputs, {array(1.0f)}, sorted_argnums);
+    return std::make_pair(outputs, grads);
+  };
+}
+
+namespace detail {
+
+std::pair<std::vector<array>, std::vector<array>> vmap_trace(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& inputs,
+    const std::vector<int>& in_axes) {
+  // Set the global tracing flag.
+  detail::InTracing in_tracing;
+
+  if (in_axes.size() != inputs.size()) {
+    std::stringstream ss;
+    ss << "[vmap] The number of in axes (" << in_axes.size()
+       << ") must match the number of inputs (" << inputs.size() << ").";
+    throw std::invalid_argument(ss.str());
+  }
+
+  // Some error checking and get the vmap axis size
+  size_t vmap_ax_size;
+  for (int i = 0; i < inputs.size(); ++i) {
+    if (in_axes[i] != -1) {
+      if (inputs[i].ndim() == 0) {
+        throw std::invalid_argument(
+            "[vmap] Cannot vmap an input with zero dimensions.");
+      }
+      if (in_axes[i] > inputs[i].ndim()) {
+        std::ostringstream msg;
+        msg << "[vmap] Axis " << in_axes[i] << " invalid for input with "
+            << inputs[i].ndim() << " dimensions.";
+        throw std::invalid_argument(msg.str());
+      }
+      vmap_ax_size = inputs[i].shape(in_axes[i]);
+    }
+  }
+  // Check that all vmapped axes have the same size
+  for (int i = 0; i < inputs.size(); ++i) {
+    if (in_axes[i] != -1) {
+      if (size_t in_ax = inputs[i].shape(in_axes[i]); vmap_ax_size != in_ax) {
+        std::ostringstream msg;
+        msg << "[vmap] Inconsistent axis sizes: " << in_ax << " and "
+            << vmap_ax_size << ".";
+        throw std::invalid_argument(msg.str());
+      }
+    }
+  }
+
+  // Run the function on placeholder inputs
+  // to get the original graph
+  std::vector<array> s_inputs;
+  for (int i = 0; i < inputs.size(); ++i) {
+    if (in_axes[i] != -1) {
+      auto shape = inputs[i].shape();
+      shape.erase(shape.begin() + in_axes[i]);
+      array in(shape, inputs[i].dtype(), nullptr, {});
+      s_inputs.push_back(in);
+      s_inputs.back().set_tracer(true);
+    } else {
+      s_inputs.push_back(inputs[i]);
+    }
+  }
+  return {s_inputs, fun(s_inputs)};
+}
+
+std::vector<array> vmap_replace(
+    const std::vector<array>& inputs,
+    const std::vector<array>& s_inputs,
+    const std::vector<array>& s_outputs,
+    const std::vector<int>& in_axes,
+    const std::vector<int>& out_axes) {
+  if (out_axes.size() != s_outputs.size()) {
+    std::stringstream msg;
+    msg << "[vmap] The number of out axes (" << out_axes.size()
+        << ") must match the number of outputs (" << s_outputs.size() << ").";
+    throw std::invalid_argument(msg.str());
+  }
+
+  int vmap_size = -1;
+  for (int i = 0; i < inputs.size(); ++i) {
+    if (in_axes[i] >= 0) {
+      vmap_size = inputs[i].shape(in_axes[i]);
+      break;
+    }
+  }
+  if (vmap_size == -1) {
+    throw std::invalid_argument("At least one of in_axes must be non-None.");
+  }
+
+  std::unordered_map<std::uintptr_t, std::pair<array, int>> tmap;
+  std::unordered_set<std::uintptr_t> needs_vmap;
+  std::unordered_set<std::uintptr_t> cache;
+  for (int i = 0; i < s_inputs.size(); ++i) {
+    auto in = s_inputs[i];
+    if (in_axes[i] != -1) {
+      tmap.insert({in.id(), {inputs[i], in_axes[i]}});
+      needs_vmap.insert(in.id());
+      in.set_tracer(false);
+    }
+    cache.insert(in.id());
+  }
+
+  // Topologically sort the graph
+  std::vector<array> tape;
+
+  std::function<void(const array&)> recurse;
+
+  recurse = [&](const array& a) {
+    auto id = a.id();
+    if (cache.find(id) != cache.end()) {
+      return;
+    }
+    cache.insert(id);
+    for (auto& s : a.siblings()) {
+      cache.insert(s.id());
+    }
+
+    // Recurse on inputs
+    for (auto& input : a.inputs()) {
+      recurse(input);
+    }
+    // If any input needs a vmap, then the outputs also need
+    // a vmap
+    for (auto& input : a.inputs()) {
+      if (needs_vmap.find(input.id()) != needs_vmap.end()) {
+        tape.push_back(a);
+        tape.back().set_tracer(false);
+        needs_vmap.insert(a.id());
+        for (auto s : a.siblings()) {
+          needs_vmap.insert(s.id());
+          s.set_tracer(false);
+        }
+        break;
+      }
+    }
+  };
+
+  for (auto& out : s_outputs) {
+    if (out.has_primitive()) {
+      recurse(out);
+    }
+  }
+
+  // Transform each primitive in the graph with
+  // its vmap implementation
+  for (auto& a : tape) {
+    std::vector<array> v_inputs;
+    std::vector<int> v_axes;
+    for (auto& in : a.inputs()) {
+      auto map_it = tmap.find(in.id());
+      if (map_it != tmap.end()) {
+        v_inputs.push_back(map_it->second.first);
+        v_axes.push_back(map_it->second.second);
+      } else {
+        v_inputs.push_back(in);
+        v_axes.push_back(-1);
+      }
+    }
+
+    auto [v_outputs, v_out_axes] = a.primitive().vmap(v_inputs, v_axes);
+
+    // For each primitive's outputs add its id, the vout id and the vax
+    auto outputs = a.outputs();
+    for (int i = 0; i < v_outputs.size(); ++i) {
+      tmap.insert({outputs[i].id(), {v_outputs[i], v_out_axes[i]}});
+    }
+  }
+
+  // Populate the outputs and make sure all the output axes are
+  // in the right place
+  std::vector<array> outputs;
+  for (int i = 0; i < s_outputs.size(); ++i) {
+    if (auto map_it = tmap.find(s_outputs[i].id()); map_it != tmap.end()) {
+      auto& [out, vdim] = map_it->second;
+      if (vdim != out_axes[i]) {
+        if (out_axes[i] >= out.ndim()) {
+          std::ostringstream msg;
+          msg << "[vmap] Axis " << out_axes[i] << " invalid for output with "
+              << out.ndim() << " dimensions.";
+          throw std::invalid_argument(msg.str());
+        }
+        out = moveaxis(out, vdim, out_axes[i]);
+      }
+      outputs.push_back(out);
+    } else {
+      // When the output has no input dependencies
+      // use the size of the vmapped axis in the inputs to expand the output
+      array output = expand_dims(s_outputs[i], out_axes[i]);
+      output = repeat(output, vmap_size, out_axes[i]);
+      outputs.push_back(output);
+    }
+  }
+  return outputs;
+}
+
+} // namespace detail
+
+std::function<std::vector<array>(const std::vector<array>&)> vmap(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<int>& in_axes /* = {} */,
+    const std::vector<int>& out_axes /* = {} */) {
+  auto infer_axes = [](auto axes) {
+    return !axes.empty() &&
+        std::all_of(axes.begin(), axes.end(), [](int ax) { return ax < 0; });
+  };
+  if (infer_axes(in_axes) != infer_axes(out_axes)) {
+    throw std::invalid_argument(
+        "[vmap] Input (or output) axes must be "
+        "specified if output (or input) axes are.");
+  }
+  auto vfun = [fun, in_axes = in_axes, out_axes = out_axes](
+                  const std::vector<array>& inputs) mutable {
+    if (in_axes.size() == 0) {
+      in_axes.resize(inputs.size(), 0);
+    }
+
+    auto [trace_inputs, trace_outputs] =
+        detail::vmap_trace(fun, inputs, in_axes);
+
+    if (out_axes.size() == 0) {
+      out_axes.resize(trace_outputs.size(), 0);
+    }
+
+    return detail::vmap_replace(
+        inputs, trace_inputs, trace_outputs, in_axes, out_axes);
+  };
+
+  return vfun;
+}
+
+std::function<array(const array&, const array&)> vmap(
+    const std::function<array(const array&, const array&)>& fun,
+    int in_axis_a /* = 0 */,
+    int in_axis_b /* = 0 */,
+    int out_axis /* = 0 */) {
+  auto vfun = vmap(
+      [fun](const std::vector<array>& inputs) {
+        return std::vector<array>{fun(inputs[0], inputs[1])};
+      },
+      {in_axis_a, in_axis_b},
+      {out_axis});
+  return [vfun](const array& a, const array& b) { return vfun({a, b})[0]; };
+}
+
+std::function<array(const array&)> vmap(
+    const std::function<array(const array&)>& fun,
+    int in_axis /* = 0 */,
+    int out_axis /* = 0 */) {
+  auto vfun = vmap(
+      [fun](const std::vector<array>& inputs) {
+        return std::vector<array>{fun(inputs[0])};
+      },
+      {in_axis},
+      {out_axis});
+  return [vfun](const array& a) { return vfun({a})[0]; };
+}
+
+std::function<std::vector<array>(const std::vector<array>&)> custom_function(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::optional<std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<array>&)>> fun_vjp /* = std::nullopt */,
+    std::optional<std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<int>&)>> fun_jvp /* = std::nullopt */,
+    std::optional<std::function<std::pair<std::vector<array>, std::vector<int>>(
+        const std::vector<array>&,
+        const std::vector<int>&)>> fun_vmap /* = std::nullopt */) {
+  if (!fun_vjp.has_value() && !fun_jvp.has_value() && !fun_vmap.has_value()) {
+    return fun;
+  }
+
+  return [fun = std::move(fun),
+          fun_vjp = std::move(fun_vjp),
+          fun_jvp = std::move(fun_jvp),
+          fun_vmap = std::move(fun_vmap)](const std::vector<array>& args) {
+    // Compute the outputs
+    auto outputs = fun(args);
+    for (auto& out : outputs) {
+      out = stop_gradient(out);
+    }
+
+    // Prepare the inputs to the primitive
+    // We also add the outputs to the primitive so that it can "run" the forward
+    // pass.
+    std::vector<array> inputs = args;
+    inputs.insert(inputs.end(), outputs.begin(), outputs.end());
+
+    // Compute the stream. Maybe do it in a smarter way at some point in the
+    // future.
+    Stream s = (outputs[0].has_primitive()) ? outputs[0].primitive().stream()
+                                            : default_stream(default_device());
+
+    // Make the output info
+    std::vector<Shape> shapes;
+    std::vector<Dtype> dtypes;
+    for (const auto& out : outputs) {
+      shapes.emplace_back(out.shape());
+      dtypes.emplace_back(out.dtype());
+    }
+
+    return array::make_arrays(
+        std::move(shapes),
+        dtypes,
+        std::make_shared<CustomTransforms>(
+            to_stream(s),
+            outputs.size(),
+
+            // We use the passed vjp function or compute it from the inputs and
+            // passed cotangents. Note that this may be less efficient than
+            // using `fun` directly because we may not be able to fully reuse
+            // the outputs of the forward pass.
+            fun_vjp.value_or(
+                [fun](auto primals, auto cotangents, auto outputs) {
+                  auto [__, vjps] = vjp(fun, primals, cotangents);
+                  return vjps;
+                }),
+
+            // We use the passed jvp function or compute it from the primals
+            // and tangents. Similarly we can't take full advantage of the
+            // argnums so it is best to use `fun` directly if we don't need a
+            // custom transform.
+            //
+            // TODO: Use stop_gradient to make full use of argnums and not
+            //       waste computation.
+            fun_jvp.value_or([fun](auto primals, auto tangents, auto argnums) {
+              std::vector<array> all_tangents;
+              for (int i = 0, j = 0; i < primals.size(); i++) {
+                if (j < argnums.size() && i == argnums[j]) {
+                  all_tangents.emplace_back(tangents[j++]);
+                } else {
+                  all_tangents.emplace_back(zeros_like(primals[i]));
+                }
+              }
+              auto [__, jvps] = jvp(fun, primals, all_tangents);
+              return jvps;
+            }),
+
+            // Same as above, we use the passed vmap function or we compute it
+            // from `fun`. The output axes is selected to be all 0s which again
+            // may be suboptimal but the only thing we can do without any
+            // information for `fun`.
+            fun_vmap.value_or(
+                [fun, out_size = outputs.size()](auto inputs, auto in_axes)
+                    -> std::pair<std::vector<array>, std::vector<int>> {
+                  std::vector<int> out_axes(out_size, 0);
+                  return {vmap(fun, in_axes, out_axes)(inputs), out_axes};
+                })),
+        inputs);
+  };
+}
+
+std::function<std::vector<array>(const std::vector<array>&)> custom_vjp(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<array>&)> fun_vjp) {
+  return custom_function(fun, fun_vjp, std::nullopt, std::nullopt);
+}
+
+std::function<std::vector<array>(const std::vector<array>&)> checkpoint(
+    std::function<std::vector<array>(const std::vector<array>&)> fun) {
+  auto vjp_fun = [fun](
+                     const std::vector<array>& primals,
+                     const std::vector<array>& cotangents,
+                     const std::vector<array>& outputs) -> std::vector<array> {
+    auto [__, vjps] = vjp(fun, depends(primals, outputs), cotangents);
+    return vjps;
+  };
+
+  return custom_vjp(fun, vjp_fun);
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/transforms.h b/macos/third_party/mlx/mlx/transforms.h
new file mode 100644
index 0000000..1848be7
--- /dev/null
+++ b/macos/third_party/mlx/mlx/transforms.h
@@ -0,0 +1,231 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <optional>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+MLX_API void async_eval(std::vector<array> outputs);
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+void async_eval(Arrays&&... outputs) {
+  async_eval(std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+MLX_API void eval(std::vector<array> outputs);
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+void eval(Arrays&&... outputs) {
+  eval(std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+/**
+ *  Computes the output and vector-Jacobian product (VJP) of a function.
+ *
+ *  Computes the vector-Jacobian product of the vector of cotangents with the
+ *  Jacobian of the function evaluated at the primals. Returns a pair of
+ *  vectors of output arrays and VJP arrays.
+ **/
+MLX_API std::pair<std::vector<array>, std::vector<array>> vjp(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& primals,
+    const std::vector<array>& cotangents);
+
+/**
+ *  Computes the output and vector-Jacobian product (VJP) of a unary function.
+ */
+MLX_API std::pair<array, array> vjp(
+    const std::function<array(const array&)>& fun,
+    const array& primal,
+    const array& cotangent);
+
+/**
+ *  Computes the output and Jacobian-vector product (JVP) of a function.
+ *
+ *  Computes the Jacobian-vector product of the Jacobian of the function
+ *  evaluated at the primals with the vector of tangents. Returns a pair of
+ *  vectors of output arrays and JVP arrays.
+ **/
+MLX_API std::pair<std::vector<array>, std::vector<array>> jvp(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& primals,
+    const std::vector<array>& tangents);
+
+/**
+ *  Computes the output and Jacobian-vector product (JVP) of a unary function.
+ */
+MLX_API std::pair<array, array> jvp(
+    const std::function<array(const array&)>& fun,
+    const array& primal,
+    const array& tangent);
+
+// Return type of general value_and_grad: a function which takes an input
+// vector of arrays and returns a pair of vectors of arrays one for the
+// values and one for the gradients wrt the first value.
+using ValueAndGradFn =
+    std::function<std::pair<std::vector<array>, std::vector<array>>(
+        const std::vector<array>&)>;
+using SimpleValueAndGradFn = std::function<std::pair<array, std::vector<array>>(
+    const std::vector<array>&)>;
+
+/**
+ *  Returns a function which computes the value and gradient of the input
+ *  function with respect to a vector of input arrays.
+ **/
+MLX_API ValueAndGradFn value_and_grad(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<int>& argnums);
+
+/**
+ *  Returns a function which computes the value and gradient of the input
+ *  function with respect to a single input array.
+ **/
+ValueAndGradFn inline value_and_grad(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    int argnum = 0) {
+  return value_and_grad(fun, std::vector<int>{argnum});
+}
+
+/**
+ *  Returns a function which computes the value and gradient of the unary
+ *  input function.
+ **/
+std::function<std::pair<array, array>(const array&)> inline value_and_grad(
+    const std::function<array(const array&)>& fun) {
+  return [fun](auto inputs) { return vjp(fun, inputs, array(1.0f)); };
+}
+
+SimpleValueAndGradFn inline value_and_grad(
+    const std::function<array(const std::vector<array>&)>& fun,
+    const std::vector<int>& argnums) {
+  return [fun, argnums](auto inputs) {
+    auto result = value_and_grad(
+        [fun](auto inputs) { return std::vector<array>{fun(inputs)}; },
+        argnums)(inputs);
+
+    return std::make_pair(result.first[0], result.second);
+  };
+}
+
+SimpleValueAndGradFn inline value_and_grad(
+    const std::function<array(const std::vector<array>&)>& fun,
+    int argnum = 0) {
+  return value_and_grad(fun, std::vector<int>{argnum});
+}
+
+/**
+ *  Returns a function which computes the gradient of the input function with
+ *  respect to a vector of input arrays.
+ *
+ *  The function being differentiated takes a vector of arrays and returns an
+ *  array. The vector of `argnums` specifies which the arguments to compute
+ *  the gradient with respect to. At least one argument must be specified.
+ **/
+std::function<std::vector<array>(const std::vector<array>&)> inline grad(
+    const std::function<array(const std::vector<array>&)>& fun,
+    const std::vector<int>& argnums) {
+  auto fn = value_and_grad(fun, argnums);
+  return [fn](const std::vector<array>& inputs) { return fn(inputs).second; };
+}
+
+/**
+ *  Returns a function which computes the gradient of the input function with
+ *  respect to a single input array.
+ *
+ *  The function being differentiated takes a vector of arrays and returns an
+ *  array. The optional `argnum` index specifies which the argument to compute
+ *  the gradient with respect to and defaults to 0.
+ **/
+std::function<std::vector<array>(const std::vector<array>&)> inline grad(
+    const std::function<array(const std::vector<array>&)>& fun,
+    int argnum = 0) {
+  return grad(fun, std::vector<int>{argnum});
+}
+
+/**
+ *  Returns a function which computes the gradient of the unary input function.
+ **/
+std::function<array(const array&)> inline grad(
+    const std::function<array(const array&)>& fun) {
+  auto fn = value_and_grad(fun);
+  return [fn](const array& input) { return fn(input).second; };
+}
+
+/**
+ * Automatically vectorize a unary function over the requested axes.
+ */
+MLX_API std::function<array(const array&)> vmap(
+    const std::function<array(const array&)>& fun,
+    int in_axis = 0,
+    int out_axis = 0);
+
+/**
+ * Automatically vectorize a binary function over the requested axes.
+ */
+MLX_API std::function<array(const array&, const array&)> vmap(
+    const std::function<array(const array&, const array&)>& fun,
+    int in_axis_a = 0,
+    int in_axis_b = 0,
+    int out_axis = 0);
+
+/**
+ * Automatically vectorize a function over the requested axes.
+ *
+ * The input function to `vmap` takes as an argument a vector of arrays and
+ * returns a vector of arrays. Optionally specify the axes to vectorize over
+ * with `in_axes` and `out_axes`, otherwise a default of 0 is used.
+ * Returns a vectorized function with the same signature as the input
+ * function.
+ */
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> vmap(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<int>& in_axes = {},
+    const std::vector<int>& out_axes = {});
+
+/**
+ * Redefine the transformations of `fun` according to the provided functions.
+ *
+ * Namely when calling the vjp of `fun` then `fun_vjp` will be called,
+ * `fun_jvp` for the jvp and `fun_vmap` for vmap.
+ *
+ * If any transformation is not provided, then a default one is created by
+ * calling `vjp`, `jvp` and `vmap` on the function directly.
+ */
+MLX_API std::function<std::vector<array>(const std::vector<array>&)>
+custom_function(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::optional<std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<array>&)>> fun_vjp = std::nullopt,
+    std::optional<std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<int>&)>> fun_jvp = std::nullopt,
+    std::optional<std::function<std::pair<std::vector<array>, std::vector<int>>(
+        const std::vector<array>&,
+        const std::vector<int>&)>> fun_vmap = std::nullopt);
+
+/**
+ * Return a function that behaves exactly like `fun` but if the vjp of the
+ * results is computed `fun_vjp` will be used instead of `vjp(fun, ...)` .
+ */
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> custom_vjp(
+    std::function<std::vector<array>(const std::vector<array>&)> fun,
+    std::function<std::vector<array>(
+        const std::vector<array>&,
+        const std::vector<array>&,
+        const std::vector<array>&)> fun_vjp);
+
+/**
+ * Checkpoint the gradient of a function. Namely, discard all intermediate
+ * state and recalculate it when we need to compute the gradient.
+ */
+MLX_API std::function<std::vector<array>(const std::vector<array>&)> checkpoint(
+    std::function<std::vector<array>(const std::vector<array>&)> fun);
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/transforms_impl.h b/macos/third_party/mlx/mlx/transforms_impl.h
new file mode 100644
index 0000000..eff458c
--- /dev/null
+++ b/macos/third_party/mlx/mlx/transforms_impl.h
@@ -0,0 +1,88 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/api.h"
+
+namespace mlx::core::detail {
+
+MLX_API std::pair<std::vector<array>, std::vector<array>> vmap_trace(
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    const std::vector<array>& inputs,
+    const std::vector<int>& in_axes);
+
+MLX_API std::vector<array> vmap_replace(
+    const std::vector<array>& inputs,
+    const std::vector<array>& s_inputs,
+    const std::vector<array>& s_outputs,
+    const std::vector<int>& in_axes,
+    const std::vector<int>& out_axes);
+
+// Create an InTracing object during tracing operations to signify to the rest
+// of the codebase that we are during tracing so evals should not throw away
+// the graph.
+struct InTracing {
+  explicit InTracing(bool dynamic = false, bool grad = false) {
+    grad_counter += grad;
+    trace_stack().push_back({dynamic, grad});
+  }
+  ~InTracing() {
+    grad_counter -= trace_stack().back().second;
+    trace_stack().pop_back();
+  }
+
+  static bool in_tracing() {
+    return !trace_stack().empty();
+  }
+  static bool in_dynamic_tracing() {
+    // compile is always and only the outer-most transform
+    return in_tracing() && trace_stack().front().first;
+  }
+
+  static bool in_grad_tracing() {
+    return grad_counter > 0;
+  }
+
+ private:
+  static int grad_counter;
+  static std::vector<std::pair<char, char>>& trace_stack();
+};
+
+struct RetainGraph {
+  RetainGraph() {
+    tracing_counter++;
+  }
+  ~RetainGraph() {
+    tracing_counter--;
+  }
+
+  static bool retain_graph() {
+    return tracing_counter > 0;
+  }
+
+ private:
+  static int tracing_counter;
+};
+
+/** Return true if we are currently performing a function transformation in
+ * order to keep the graph when evaluating tracer arrays. */
+inline bool in_tracing() {
+  return detail::InTracing::in_tracing();
+}
+
+/** Return true if we are in a dynamic (shapeless) trace used for compiling or
+ * exporting graphs with dynamic shapes.  */
+inline bool in_dynamic_tracing() {
+  return detail::InTracing::in_dynamic_tracing();
+}
+
+/** Return true if we are in a gradient trace (vjp, jvp, etc).  */
+inline bool in_grad_tracing() {
+  return detail::InTracing::in_grad_tracing();
+}
+
+inline bool retain_graph() {
+  return detail::RetainGraph::retain_graph();
+}
+
+} // namespace mlx::core::detail
diff --git a/macos/third_party/mlx/mlx/types/bf16.h b/macos/third_party/mlx/mlx/types/bf16.h
new file mode 100644
index 0000000..5951941
--- /dev/null
+++ b/macos/third_party/mlx/mlx/types/bf16.h
@@ -0,0 +1,187 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <vector>
+
+#define __MLX_BFLOAT_NAN__ 0x7FC0
+
+namespace mlx::core {
+
+namespace {
+union float_bits_bf16 {
+  float f;
+  uint32_t u;
+};
+} // namespace
+
+struct _MLX_BFloat16 {
+  uint16_t bits_;
+
+  // Default constructor
+  _MLX_BFloat16() = default;
+
+  // Default copy constructor
+  _MLX_BFloat16(_MLX_BFloat16 const&) = default;
+
+  // Appease std::vector<bool> for being special
+  _MLX_BFloat16& operator=(std::vector<bool>::reference x) {
+    bits_ = x;
+    return *this;
+  }
+
+  _MLX_BFloat16& operator=(const float& x) {
+    return (*this = _MLX_BFloat16(x));
+  }
+
+  // From float32
+  _MLX_BFloat16(const float& x) {
+    if (std::isnan(x)) {
+      bits_ = __MLX_BFLOAT_NAN__;
+    } else {
+      // Union
+      float_bits_bf16 in;
+
+      // Take bits
+      in.f = x;
+
+      // Round to nearest even
+      in.u += (in.u >> 16 & 1) + uint32_t(0x7FFF);
+
+      // Take upper 16 bits
+      bits_ = in.u >> 16;
+    }
+  }
+
+  // To float32
+  operator float() const {
+    // Union
+    float_bits_bf16 out;
+
+    // Upper 16 bits are the data and lower 16 bits are 0s
+    out.u = ((uint32_t)bits_) << 16;
+
+    return out.f;
+  }
+};
+
+#define bfloat_binop_base(__op__, __operator__, otype, atype, btype, ctype) \
+  inline otype __operator__(atype lhs, btype rhs) {                         \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);          \
+  }
+
+#define bfloat_binop_helper(__op__, __operator__, otype, itype, ctype) \
+  inline otype __operator__(_MLX_BFloat16 lhs, itype rhs) {            \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);     \
+  }                                                                    \
+  inline otype __operator__(itype lhs, _MLX_BFloat16 rhs) {            \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);     \
+  }
+
+// Operators
+#define bfloat_binop(_op_, _operator_)                                       \
+  bfloat_binop_base(                                                         \
+      _op_, _operator_, _MLX_BFloat16, _MLX_BFloat16, _MLX_BFloat16, float); \
+  bfloat_binop_helper(_op_, _operator_, float, float, float);                \
+  bfloat_binop_helper(_op_, _operator_, double, double, double);             \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, bool, float);         \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, int32_t, float);      \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, uint32_t, float);     \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, int64_t, float);      \
+  bfloat_binop_helper(_op_, _operator_, _MLX_BFloat16, uint64_t, float);
+
+bfloat_binop(+, operator+);
+bfloat_binop(-, operator-);
+bfloat_binop(*, operator*);
+bfloat_binop(/, operator/);
+
+#undef bfloat_binop
+
+// Comparison ops
+#define bfloat_compop(__op__, __operator__)                             \
+  bfloat_binop_base(                                                    \
+      __op__, __operator__, bool, _MLX_BFloat16, _MLX_BFloat16, float); \
+  bfloat_binop_helper(__op__, __operator__, bool, float, float);        \
+  bfloat_binop_helper(__op__, __operator__, bool, double, double);      \
+  bfloat_binop_helper(__op__, __operator__, bool, int32_t, float);      \
+  bfloat_binop_helper(__op__, __operator__, bool, uint32_t, float);     \
+  bfloat_binop_helper(__op__, __operator__, bool, int64_t, float);      \
+  bfloat_binop_helper(__op__, __operator__, bool, uint64_t, float);
+
+bfloat_compop(>, operator>);
+bfloat_compop(<, operator<);
+bfloat_compop(>=, operator>=);
+bfloat_compop(<=, operator<=);
+bfloat_compop(==, operator==);
+bfloat_compop(!=, operator!=);
+
+#undef bfloat_compop
+
+// Negative
+inline _MLX_BFloat16 operator-(_MLX_BFloat16 lhs) {
+  return -static_cast<float>(lhs);
+}
+
+// Inplace ops
+#define bfloat_inplace_op(__op__, __operator__)                              \
+  inline _MLX_BFloat16& __operator__(_MLX_BFloat16& lhs, const float& rhs) { \
+    lhs = lhs __op__ rhs;                                                    \
+    return lhs;                                                              \
+  }                                                                          \
+  inline float& __operator__(float& lhs, _MLX_BFloat16 rhs) {                \
+    lhs = lhs __op__ rhs;                                                    \
+    return lhs;                                                              \
+  }
+
+bfloat_inplace_op(+, operator+=);
+bfloat_inplace_op(-, operator-=);
+bfloat_inplace_op(*, operator*=);
+bfloat_inplace_op(/, operator/=);
+
+#undef bfloat_inplace_op
+
+// Bitwise ops
+
+#define bfloat_bitop(__op__, __operator__)                                  \
+  inline _MLX_BFloat16 __operator__(_MLX_BFloat16 lhs, _MLX_BFloat16 rhs) { \
+    _MLX_BFloat16 out;                                                      \
+    out.bits_ = lhs.bits_ __op__ rhs.bits_;                                 \
+    return out;                                                             \
+  }                                                                         \
+  inline _MLX_BFloat16 __operator__(_MLX_BFloat16 lhs, uint16_t rhs) {      \
+    _MLX_BFloat16 out;                                                      \
+    out.bits_ = lhs.bits_ __op__ rhs;                                       \
+    return out;                                                             \
+  }                                                                         \
+  inline _MLX_BFloat16 __operator__(uint16_t lhs, _MLX_BFloat16 rhs) {      \
+    _MLX_BFloat16 out;                                                      \
+    out.bits_ = lhs __op__ rhs.bits_;                                       \
+    return out;                                                             \
+  }
+
+bfloat_bitop(|, operator|);
+bfloat_bitop(&, operator&);
+bfloat_bitop(^, operator^);
+
+#undef bfloat_bitop
+
+#define bfloat_inplace_bitop(__op__, __operator__)                            \
+  inline _MLX_BFloat16& __operator__(_MLX_BFloat16& lhs, _MLX_BFloat16 rhs) { \
+    lhs.bits_ = lhs.bits_ __op__ rhs.bits_;                                   \
+    return lhs;                                                               \
+  }                                                                           \
+  inline _MLX_BFloat16& __operator__(_MLX_BFloat16& lhs, uint16_t rhs) {      \
+    lhs.bits_ = lhs.bits_ __op__ rhs;                                         \
+    return lhs;                                                               \
+  }
+
+bfloat_inplace_bitop(|, operator|=);
+bfloat_inplace_bitop(&, operator&=);
+bfloat_inplace_bitop(^, operator^=);
+
+#undef bfloat_inplace_bitop
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/types/complex.h b/macos/third_party/mlx/mlx/types/complex.h
new file mode 100644
index 0000000..51101cc
--- /dev/null
+++ b/macos/third_party/mlx/mlx/types/complex.h
@@ -0,0 +1,113 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+#include <complex>
+#include "mlx/types/half_types.h"
+
+namespace mlx::core {
+
+struct complex64_t;
+struct complex128_t;
+
+template <typename T>
+inline constexpr bool can_convert_to_complex128 =
+    !std::is_same_v<T, complex128_t> && std::is_convertible_v<T, double>;
+
+struct complex128_t : public std::complex<double> {
+  complex128_t() : std::complex<double>() {};
+  complex128_t(double v, double u) : std::complex<double>(v, u) {};
+  complex128_t(std::complex<double> v) : std::complex<double>(v) {};
+
+  template <
+      typename T,
+      typename = typename std::enable_if<can_convert_to_complex128<T>>::type>
+  complex128_t(T x) : std::complex<double>(x){};
+
+  operator float() const {
+    return real();
+  };
+};
+
+template <typename T>
+inline constexpr bool can_convert_to_complex64 =
+    !std::is_same_v<T, complex64_t> && std::is_convertible_v<T, float>;
+
+struct complex64_t : public std::complex<float> {
+  complex64_t() : std::complex<float>() {};
+  complex64_t(float v, float u) : std::complex<float>(v, u) {};
+  complex64_t(std::complex<float> v) : std::complex<float>(v) {};
+
+  template <
+      typename T,
+      typename = typename std::enable_if<can_convert_to_complex64<T>>::type>
+  complex64_t(T x) : std::complex<float>(x){};
+
+  operator float() const {
+    return real();
+  };
+};
+
+inline bool operator>=(const complex64_t& a, const complex64_t& b) {
+  return (a.real() > b.real()) ||
+      (a.real() == b.real() && a.imag() >= b.imag());
+}
+
+inline bool operator>(const complex64_t& a, const complex64_t& b) {
+  return (a.real() > b.real()) || (a.real() == b.real() && a.imag() > b.imag());
+}
+
+inline complex64_t operator%(complex64_t a, complex64_t b) {
+  auto real = a.real() - (b.real() * static_cast<int64_t>(a.real() / b.real()));
+  auto imag = a.imag() - (b.imag() * static_cast<int64_t>(a.imag() / b.imag()));
+  if (real != 0 && ((real < 0) != (b.real() < 0)))
+    real += b.real();
+  if (imag != 0 && ((imag < 0) != (b.imag() < 0)))
+    imag += b.imag();
+  return {real, imag};
+}
+
+inline bool operator<=(const complex64_t& a, const complex64_t& b) {
+  return operator>=(b, a);
+}
+
+inline bool operator<(const complex64_t& a, const complex64_t& b) {
+  return operator>(b, a);
+}
+
+inline complex64_t operator-(const complex64_t& v) {
+  return -static_cast<std::complex<float>>(v);
+}
+
+// clang-format off
+#define complex_binop_helper(_op_, _operator_, itype)            \
+  inline complex64_t _operator_(itype x, const complex64_t& y) { \
+    return static_cast<complex64_t>(x) _op_ y;           \
+  }                                                              \
+  inline complex64_t _operator_(const complex64_t& x, itype y) { \
+    return x _op_ static_cast<complex64_t>(y);           \
+  }
+
+#define complex_binop(_op_, _operator_)                                               \
+  inline complex64_t _operator_(const std::complex<float>& x, const complex64_t& y) { \
+    return x _op_ static_cast<std::complex<float>>(y);                                \
+  }                                                                                   \
+  inline complex64_t _operator_(const complex64_t& x, const std::complex<float>& y) { \
+    return static_cast<std::complex<float>>(x) _op_ y;                                \
+  }                                                                                   \
+  inline complex64_t _operator_(const complex64_t& x, const complex64_t& y) {         \
+    return static_cast<std::complex<float>>(x)                                        \
+        _op_ static_cast<std::complex<float>>(y);                                     \
+  }                                                                                   \
+  complex_binop_helper(_op_, _operator_, bool)                                        \
+  complex_binop_helper(_op_, _operator_, uint32_t)                                    \
+  complex_binop_helper(_op_, _operator_, uint64_t)                                    \
+  complex_binop_helper(_op_, _operator_, int32_t)                                     \
+  complex_binop_helper(_op_, _operator_, int64_t)                                     \
+  complex_binop_helper(_op_, _operator_, float16_t)                                   \
+  complex_binop_helper(_op_, _operator_, bfloat16_t)                                  \
+  complex_binop_helper(_op_, _operator_, float)
+// clang-format on
+
+complex_binop(+, operator+)
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/types/fp16.h b/macos/third_party/mlx/mlx/types/fp16.h
new file mode 100644
index 0000000..c174afe
--- /dev/null
+++ b/macos/third_party/mlx/mlx/types/fp16.h
@@ -0,0 +1,234 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <vector>
+
+#define __MLX_HALF_NAN__ 0x7D00
+
+namespace mlx::core {
+
+namespace {
+union float_bits_fp16 {
+  float f;
+  uint32_t u;
+};
+} // namespace
+
+struct _MLX_Float16 {
+  uint16_t bits_;
+
+  // Default constructor
+  _MLX_Float16() = default;
+
+  // Default copy constructor
+  _MLX_Float16(_MLX_Float16 const&) = default;
+
+  // Appease std::vector<bool> for being special
+  _MLX_Float16& operator=(std::vector<bool>::reference x) {
+    bits_ = x;
+    return *this;
+  }
+
+  _MLX_Float16& operator=(const float& x) {
+    return (*this = _MLX_Float16(x));
+  }
+
+  // From float32
+  _MLX_Float16(const float& x) : bits_(0) {
+    // Conversion following
+    // https://github.com/Maratyszcza/FP16/blob/master/include/fp16/fp16.h
+
+    // Union
+    float_bits_fp16 in;
+
+    // Take fp32 bits
+    in.f = x;
+
+    // Find and take sign bit
+    uint32_t x_sign_32 = in.u & uint32_t(0x80000000);
+    uint16_t x_sign_16 = (x_sign_32 >> 16);
+
+    if (std::isnan(x)) {
+      bits_ = x_sign_16 | uint16_t(__MLX_HALF_NAN__);
+    } else {
+      // Union
+      float_bits_fp16 inf_scale, zero_scale, magic_bits;
+
+      // Find exponent bits and take the max supported by half
+      uint32_t x_expo_32 = in.u & uint32_t(0x7f800000);
+      uint32_t max_expo_32 = uint32_t(0x38800000);
+      x_expo_32 = x_expo_32 < max_expo_32 ? max_expo_32 : x_expo_32;
+      x_expo_32 += uint32_t(15) << 23;
+
+      // Handle scaling to inf as needed
+      inf_scale.u = uint32_t(0x77800000);
+      zero_scale.u = uint32_t(0x08800000);
+
+      // Combine with magic and let addition do rounding
+      magic_bits.u = x_expo_32;
+      magic_bits.f += (std::abs(x) * inf_scale.f) * zero_scale.f;
+
+      // Take the lower 5 bits of the exponent
+      uint32_t x_expo_16 = ((magic_bits.u >> 13) & uint32_t(0x7c00));
+
+      // Collect the lower 12 bits which have the mantissa
+      uint32_t x_mant_16 = magic_bits.u & uint32_t(0x0fff);
+
+      // Combine sign, exp and mantissa
+      bits_ = (x_sign_16 | uint16_t(x_expo_16 + x_mant_16));
+    }
+  }
+
+  // To float32
+  operator float() const {
+    // Conversion following
+    // https://github.com/Maratyszcza/FP16/blob/master/include/fp16/fp16.h
+
+    // Union
+    float_bits_fp16 out;
+
+    uint32_t x_sign_32 = (bits_ << 16) & uint32_t(0x80000000);
+    uint32_t base = (bits_ << 16);
+    uint32_t two_base = base + base;
+
+    uint32_t denorm_max = 1u << 27;
+    if (two_base < denorm_max) {
+      out.u = uint32_t(126) << 23; // magic mask
+      out.u |= (two_base >> 17); // Bits from fp16
+      out.f -= 0.5f; // magic bias
+    } else {
+      out.u = uint32_t(0xE0) << 23; // exponent offset
+      out.u += (two_base >> 4); // Bits from fp16
+      float out_unscaled = out.f; // Store value
+      out.u = uint32_t(0x7800000); // exponent scale
+      out.f *= out_unscaled;
+    }
+
+    // Add sign
+    out.u |= x_sign_32;
+
+    return out.f;
+  }
+};
+
+#define half_binop_base(__op__, __operator__, otype, atype, btype, ctype) \
+  inline otype __operator__(atype lhs, btype rhs) {                       \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);        \
+  }
+
+#define half_binop_helper(__op__, __operator__, otype, itype, ctype) \
+  inline otype __operator__(_MLX_Float16 lhs, itype rhs) {           \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);   \
+  }                                                                  \
+  inline otype __operator__(itype lhs, _MLX_Float16 rhs) {           \
+    return static_cast<ctype>(lhs) __op__ static_cast<ctype>(rhs);   \
+  }
+
+// Operators
+#define half_binop(__op__, __operator__)                                      \
+  half_binop_base(                                                            \
+      __op__, __operator__, _MLX_Float16, _MLX_Float16, _MLX_Float16, float); \
+  half_binop_helper(__op__, __operator__, float, float, float);               \
+  half_binop_helper(__op__, __operator__, double, double, double);            \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, bool, float);         \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, int32_t, float);      \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, uint32_t, float);     \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, int64_t, float);      \
+  half_binop_helper(__op__, __operator__, _MLX_Float16, uint64_t, float);
+
+half_binop(+, operator+);
+half_binop(-, operator-);
+half_binop(*, operator*);
+half_binop(/, operator/);
+
+#undef half_binop
+
+// Comparison ops
+#define half_compop(__op__, __operator__)                             \
+  half_binop_base(                                                    \
+      __op__, __operator__, bool, _MLX_Float16, _MLX_Float16, float); \
+  half_binop_helper(__op__, __operator__, bool, float, float);        \
+  half_binop_helper(__op__, __operator__, bool, double, double);      \
+  half_binop_helper(__op__, __operator__, bool, int32_t, float);      \
+  half_binop_helper(__op__, __operator__, bool, uint32_t, float);     \
+  half_binop_helper(__op__, __operator__, bool, int64_t, float);      \
+  half_binop_helper(__op__, __operator__, bool, uint64_t, float);
+
+half_compop(>, operator>);
+half_compop(<, operator<);
+half_compop(>=, operator>=);
+half_compop(<=, operator<=);
+half_compop(==, operator==);
+half_compop(!=, operator!=);
+
+#undef half_compop
+
+// Negative
+inline _MLX_Float16 operator-(_MLX_Float16 lhs) {
+  return -static_cast<float>(lhs);
+}
+
+// Inplace ops
+#define half_inplace_op(__op__, __operator__)                              \
+  inline _MLX_Float16& __operator__(_MLX_Float16& lhs, const float& rhs) { \
+    lhs = lhs __op__ rhs;                                                  \
+    return lhs;                                                            \
+  }                                                                        \
+  inline float& __operator__(float& lhs, _MLX_Float16 rhs) {               \
+    lhs = lhs __op__ rhs;                                                  \
+    return lhs;                                                            \
+  }
+
+half_inplace_op(+, operator+=);
+half_inplace_op(-, operator-=);
+half_inplace_op(*, operator*=);
+half_inplace_op(/, operator/=);
+
+#undef half_inplace_op
+
+// Bitwise ops
+
+#define half_bitop(__op__, __operator__)                                 \
+  inline _MLX_Float16 __operator__(_MLX_Float16 lhs, _MLX_Float16 rhs) { \
+    _MLX_Float16 out;                                                    \
+    out.bits_ = lhs.bits_ __op__ rhs.bits_;                              \
+    return out;                                                          \
+  }                                                                      \
+  inline _MLX_Float16 __operator__(_MLX_Float16 lhs, uint16_t rhs) {     \
+    _MLX_Float16 out;                                                    \
+    out.bits_ = lhs.bits_ __op__ rhs;                                    \
+    return out;                                                          \
+  }                                                                      \
+  inline _MLX_Float16 __operator__(uint16_t lhs, _MLX_Float16 rhs) {     \
+    _MLX_Float16 out;                                                    \
+    out.bits_ = lhs __op__ rhs.bits_;                                    \
+    return out;                                                          \
+  }
+
+half_bitop(|, operator|);
+half_bitop(&, operator&);
+half_bitop(^, operator^);
+
+#undef half_bitop
+
+#define half_inplace_bitop(__op__, __operator__)                           \
+  inline _MLX_Float16& __operator__(_MLX_Float16& lhs, _MLX_Float16 rhs) { \
+    lhs.bits_ = lhs.bits_ __op__ rhs.bits_;                                \
+    return lhs;                                                            \
+  }                                                                        \
+  inline _MLX_Float16& __operator__(_MLX_Float16& lhs, uint16_t rhs) {     \
+    lhs.bits_ = lhs.bits_ __op__ rhs;                                      \
+    return lhs;                                                            \
+  }
+
+half_inplace_bitop(|, operator|=);
+half_inplace_bitop(&, operator&=);
+half_inplace_bitop(^, operator^=);
+
+#undef half_inplace_bitop
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/types/half_types.h b/macos/third_party/mlx/mlx/types/half_types.h
new file mode 100644
index 0000000..d9d6b9b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/types/half_types.h
@@ -0,0 +1,58 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+
+#include <arm_fp16.h>
+namespace mlx::core {
+using ::float16_t;
+} // namespace mlx::core
+
+#else
+
+#define ADD_HALF_BINOPS
+#include "mlx/types/fp16.h"
+namespace mlx::core {
+typedef struct _MLX_Float16 float16_t;
+} // namespace mlx::core
+
+#endif // __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
+
+#ifdef __ARM_FEATURE_BF16
+
+#include <arm_bf16.h>
+namespace mlx::core {
+using ::bfloat16_t;
+} // namespace mlx::core
+
+#else
+
+#define ADD_HALF_BINOPS
+#include "mlx/types/bf16.h"
+namespace mlx::core {
+typedef struct _MLX_BFloat16 bfloat16_t;
+} // namespace mlx::core
+
+#endif // __ARM_FEATURE_BF16
+
+#ifdef ADD_HALF_BINOPS
+namespace mlx::core {
+
+// clang-format off
+#define fp16_bf16_binop_helper(__op__, __operator__)               \
+  inline float __operator__(float16_t lhs, bfloat16_t rhs) {       \
+    return static_cast<float>(lhs) __op__ static_cast<float>(rhs); \
+  }                                                                \
+  inline float __operator__(bfloat16_t lhs, float16_t rhs) {       \
+    return static_cast<float>(lhs) __op__ static_cast<float>(rhs); \
+  }
+
+fp16_bf16_binop_helper(+, operator+)
+fp16_bf16_binop_helper(-, operator-)
+fp16_bf16_binop_helper(*, operator*)
+fp16_bf16_binop_helper(/, operator/)
+// clang-format on
+
+} // namespace mlx::core
+#endif
diff --git a/macos/third_party/mlx/mlx/types/limits.h b/macos/third_party/mlx/mlx/types/limits.h
new file mode 100644
index 0000000..5f2b1e9
--- /dev/null
+++ b/macos/third_party/mlx/mlx/types/limits.h
@@ -0,0 +1,70 @@
+// Copyright © 2024 Apple Inc.
+#pragma once
+
+#include <limits>
+#include "mlx/types/half_types.h"
+
+namespace mlx::core {
+
+template <typename T>
+struct numeric_limits;
+
+template <>
+struct numeric_limits<float> : public std::numeric_limits<float> {};
+
+template <>
+struct numeric_limits<double> : public std::numeric_limits<double> {};
+
+template <>
+struct numeric_limits<float16_t> {
+ private:
+  union half_or_bits {
+    uint16_t bits;
+    float16_t value;
+  };
+  constexpr static float16_t bits_to_half(uint16_t v) {
+    return half_or_bits{v}.value;
+  }
+
+ public:
+  constexpr static float16_t lowest() {
+    return bits_to_half(0xFBFF);
+  }
+  static constexpr float16_t max() {
+    return bits_to_half(0x7BFF);
+  }
+  static constexpr float16_t epsilon() {
+    return bits_to_half(0x1400);
+  }
+  static constexpr float16_t infinity() {
+    return bits_to_half(0x7C00);
+  }
+};
+
+template <>
+struct numeric_limits<bfloat16_t> {
+ private:
+  union bfloat_or_bits {
+    uint16_t bits;
+    bfloat16_t value;
+  };
+  constexpr static bfloat16_t bits_to_bfloat(uint16_t v) {
+    return bfloat_or_bits{v}.value;
+  }
+
+ public:
+  constexpr static bfloat16_t lowest() {
+    return bits_to_bfloat(0xFF7F);
+  }
+  static constexpr bfloat16_t max() {
+    return bits_to_bfloat(0x7F7F);
+  }
+  static constexpr bfloat16_t epsilon() {
+    return bits_to_bfloat(0x3C00);
+  }
+  static constexpr bfloat16_t infinity() {
+    return bits_to_bfloat(0x7F80);
+  }
+};
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/utils.cpp b/macos/third_party/mlx/mlx/utils.cpp
new file mode 100644
index 0000000..1ae259b
--- /dev/null
+++ b/macos/third_party/mlx/mlx/utils.cpp
@@ -0,0 +1,302 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cstdlib>
+#include <iostream>
+#include <sstream>
+#include <vector>
+
+#include "mlx/dtype_utils.h"
+#include "mlx/types/limits.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+Stream to_stream(StreamOrDevice s) {
+  if (std::holds_alternative<std::monostate>(s)) {
+    return default_stream(default_device());
+  } else if (std::holds_alternative<Device>(s)) {
+    return default_stream(std::get<Device>(s));
+  } else {
+    return std::get<Stream>(s);
+  }
+}
+
+Stream to_stream(StreamOrDevice s, Device default_) {
+  if (std::holds_alternative<std::monostate>(s)) {
+    return default_stream(default_);
+  } else if (std::holds_alternative<Device>(s)) {
+    return default_stream(std::get<Device>(s));
+  } else {
+    return std::get<Stream>(s);
+  }
+}
+
+void PrintFormatter::print(std::ostream& os, bool val) {
+  if (capitalize_bool) {
+    os << (val ? "True" : "False");
+  } else {
+    os << val;
+  }
+}
+inline void PrintFormatter::print(std::ostream& os, int16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, uint16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, int32_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, uint32_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, int64_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, uint64_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, float16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, bfloat16_t val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, float val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, double val) {
+  os << val;
+}
+inline void PrintFormatter::print(std::ostream& os, complex64_t val) {
+  os << val.real();
+  if (val.imag() >= 0 || std::isnan(val.imag())) {
+    os << "+" << val.imag() << "j";
+  } else {
+    os << "-" << -val.imag() << "j";
+  }
+}
+
+PrintFormatter& get_global_formatter() {
+  static PrintFormatter formatter;
+  return formatter;
+}
+
+void abort_with_exception(const std::exception& error) {
+  std::ostringstream msg;
+  msg << "Terminating due to uncaught exception: " << error.what();
+  std::cerr << msg.str() << std::endl;
+  std::abort();
+}
+
+Dtype result_type(const std::vector<array>& arrays) {
+  Dtype t = bool_;
+  for (auto& arr : arrays) {
+    t = promote_types(t, arr.dtype());
+  }
+  return t;
+}
+
+Shape broadcast_shapes(const Shape& s1, const Shape& s2) {
+  // Use the same broadcasting rules as numpy
+  // https://numpy.org/doc/1.20/user/theory.broadcasting.html
+  // "The size of the trailing axes for both arrays in an operation must
+  // either be the same size or one of them must be one."
+  int ndim1 = s1.size();
+  int ndim2 = s2.size();
+  int ndim = std::max(ndim1, ndim2);
+  int diff = std::abs(ndim1 - ndim2);
+  const auto& big = ndim1 > ndim2 ? s1 : s2;
+  const auto& small = ndim1 > ndim2 ? s2 : s1;
+  Shape out_shape(ndim);
+  for (int i = ndim - 1; i >= diff; --i) {
+    auto a = big[i];
+    auto b = small[i - diff];
+    if (b == a) {
+      out_shape[i] = a;
+    } else if (a == 1 || b == 1) {
+      // 0 if a or b is 0 otherwise max(a, b)
+      out_shape[i] = a * b;
+    } else {
+      std::ostringstream msg;
+      msg << "[broadcast_shapes] Shapes " << s1 << " and " << s2
+          << " cannot be broadcast.";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  for (int i = diff - 1; i >= 0; --i) {
+    out_shape[i] = big[i];
+  }
+  return out_shape;
+}
+
+int normalize_axis_index(
+    int axis,
+    int ndim,
+    const std::string& msg_prefix /* = "" */) {
+  if (axis < -ndim || axis >= ndim) {
+    std::ostringstream msg;
+    msg << msg_prefix << "Axis " << axis << " is out of bounds for array with "
+        << ndim << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+  return axis < 0 ? axis + ndim : axis;
+}
+
+std::ostream& operator<<(std::ostream& os, const Device& d) {
+  os << "Device(";
+  switch (d.type) {
+    case Device::cpu:
+      os << "cpu";
+      break;
+    case Device::gpu:
+      os << "gpu";
+      break;
+  }
+  os << ", " << d.index << ")";
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, const Stream& s) {
+  os << "Stream(";
+  os << s.device;
+  os << ", " << s.index << ")";
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, int8_t x) {
+  os << static_cast<int>(x);
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, uint8_t x) {
+  os << static_cast<unsigned int>(x);
+  return os;
+}
+
+namespace {
+
+template <typename T>
+void print_subarray(std::ostream& os, const array& a, size_t index, int dim) {
+  int num_print = 3;
+  int n = a.shape(dim);
+  size_t s = a.strides()[dim];
+  bool is_last = dim == a.ndim() - 1;
+  auto prefix = is_last ? "" : std::string(7 + dim, ' ');
+  auto postfix = is_last ? ", " : ",\n";
+  os << "[";
+  for (int i = 0; i < n; ++i) {
+    os << (i == 0 ? "" : prefix);
+    if (i == num_print && n > 2 * num_print) {
+      os << "...";
+      i = n - num_print - 1;
+      index += s * (n - 2 * num_print - 1);
+    } else if (is_last) {
+      get_global_formatter().print(os, a.data<T>()[index]);
+    } else {
+      print_subarray<T>(os, a, index, dim + 1);
+    }
+    os << (i == n - 1 ? "" : postfix);
+    index += s;
+  }
+  os << "]";
+}
+
+template <typename T>
+void print_array(std::ostream& os, const array& a) {
+  os << std::boolalpha;
+  os << "array(";
+  if (a.ndim() == 0) {
+    auto data = a.data<T>();
+    get_global_formatter().print(os, data[0]);
+  } else {
+    print_subarray<T>(os, a, 0, 0);
+  }
+  os << ", dtype=" << a.dtype() << ")";
+  os << std::noboolalpha;
+}
+
+} // namespace
+
+std::ostream& operator<<(std::ostream& os, const Dtype& dtype) {
+  return os << dtype_to_string(dtype);
+}
+
+std::ostream& operator<<(std::ostream& os, const Dtype::Kind& k) {
+  switch (k) {
+    case Dtype::Kind::b:
+      return os << "b";
+    case Dtype::Kind::i:
+      return os << "i";
+    case Dtype::Kind::u:
+      return os << "u";
+    case Dtype::Kind::f:
+      return os << "f";
+    case Dtype::Kind::c:
+      return os << "c";
+    case Dtype::Kind::V:
+      return os << "V";
+  }
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, array a) {
+  a.eval();
+  dispatch_all_types(a.dtype(), [&](auto type_tag) {
+    print_array<MLX_GET_TYPE(type_tag)>(os, a);
+  });
+  return os;
+}
+
+namespace env {
+
+int get_var(const char* name, int default_value) {
+  if (const char* buff_str = std::getenv(name)) {
+    return atoi(buff_str);
+  } else {
+    return default_value;
+  }
+}
+
+} // namespace env
+
+template <typename T>
+void set_finfo_limits(double& min, double& max, double& eps) {
+  min = numeric_limits<T>::lowest();
+  max = numeric_limits<T>::max();
+  eps = numeric_limits<T>::epsilon();
+}
+
+finfo::finfo(Dtype dtype) : dtype(dtype) {
+  if (!issubdtype(dtype, inexact)) {
+    std::ostringstream msg;
+    msg << "[finfo] dtype " << dtype << " is not inexact.";
+    throw std::invalid_argument(msg.str());
+  }
+  if (dtype == float32) {
+    set_finfo_limits<float>(min, max, eps);
+  } else if (dtype == float16) {
+    set_finfo_limits<float16_t>(min, max, eps);
+  } else if (dtype == bfloat16) {
+    set_finfo_limits<bfloat16_t>(min, max, eps);
+  } else if (dtype == float64) {
+    set_finfo_limits<double>(min, max, eps);
+  } else if (dtype == complex64) {
+    this->dtype = float32;
+    set_finfo_limits<float>(min, max, eps);
+  }
+}
+
+template <typename T>
+void set_iinfo_limits(int64_t& min, uint64_t& max) {
+  min = std::numeric_limits<T>::min();
+  max = std::numeric_limits<T>::max();
+}
+
+iinfo::iinfo(Dtype dtype) : dtype(dtype) {
+  dispatch_int_types(dtype, "[iinfo]", [&](auto type_tag) {
+    set_iinfo_limits<MLX_GET_TYPE(type_tag)>(min, max);
+  });
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/utils.h b/macos/third_party/mlx/mlx/utils.h
new file mode 100644
index 0000000..bb2de46
--- /dev/null
+++ b/macos/third_party/mlx/mlx/utils.h
@@ -0,0 +1,174 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <exception>
+#include <variant>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/device.h"
+#include "mlx/dtype.h"
+#include "mlx/stream.h"
+
+namespace mlx::core {
+
+using StreamOrDevice = std::variant<std::monostate, Stream, Device>;
+MLX_API Stream to_stream(StreamOrDevice s);
+MLX_API Stream to_stream(StreamOrDevice s, Device default_);
+
+struct StreamContext {
+ public:
+  StreamContext(StreamOrDevice s) : _stream(default_stream(default_device())) {
+    if (std::holds_alternative<std::monostate>(s)) {
+      throw std::runtime_error(
+          "[StreamContext] Invalid argument, please specify a stream or device.");
+    }
+    auto _s = to_stream(s);
+    set_default_device(_s.device);
+    set_default_stream(_s);
+  }
+
+  ~StreamContext() {
+    set_default_device(_stream.device);
+    set_default_stream(_stream);
+  }
+
+ private:
+  Stream _stream;
+};
+
+struct PrintFormatter {
+  inline void print(std::ostream& os, bool val);
+  inline void print(std::ostream& os, int16_t val);
+  inline void print(std::ostream& os, uint16_t val);
+  inline void print(std::ostream& os, int32_t val);
+  inline void print(std::ostream& os, uint32_t val);
+  inline void print(std::ostream& os, int64_t val);
+  inline void print(std::ostream& os, uint64_t val);
+  inline void print(std::ostream& os, float16_t val);
+  inline void print(std::ostream& os, bfloat16_t val);
+  inline void print(std::ostream& os, float val);
+  inline void print(std::ostream& os, double val);
+  inline void print(std::ostream& os, complex64_t val);
+
+  bool capitalize_bool{false};
+};
+
+MLX_API PrintFormatter& get_global_formatter();
+
+/** Print the exception and then abort. */
+MLX_API void abort_with_exception(const std::exception& error);
+
+/** Holds information about floating-point types. */
+struct MLX_API finfo {
+  explicit finfo(Dtype dtype);
+  Dtype dtype;
+  double min;
+  double max;
+  double eps;
+};
+
+/** Holds information about integral types. */
+struct MLX_API iinfo {
+  explicit iinfo(Dtype dtype);
+  Dtype dtype;
+  int64_t min;
+  uint64_t max;
+};
+
+/** The type from promoting the arrays' types with one another. */
+inline Dtype result_type(const array& a, const array& b) {
+  return promote_types(a.dtype(), b.dtype());
+}
+inline Dtype result_type(const array& a, const array& b, const array& c) {
+  return promote_types(result_type(a, b), c.dtype());
+}
+MLX_API Dtype result_type(const std::vector<array>& arrays);
+
+MLX_API Shape broadcast_shapes(const Shape& s1, const Shape& s2);
+
+/**
+ * Returns the axis normalized to be in the range [0, ndim).
+ */
+MLX_API int
+normalize_axis_index(int axis, int ndim, const std::string& msg_prefix = "");
+
+MLX_API std::ostream& operator<<(std::ostream& os, const Device& d);
+MLX_API std::ostream& operator<<(std::ostream& os, const Stream& s);
+MLX_API std::ostream& operator<<(std::ostream& os, const Dtype& d);
+MLX_API std::ostream& operator<<(std::ostream& os, const Dtype::Kind& k);
+MLX_API std::ostream& operator<<(std::ostream& os, array a);
+inline std::ostream& operator<<(std::ostream& os, const complex64_t& v) {
+  return os << v.real() << (v.imag() >= 0 ? "+" : "") << v.imag() << "j";
+}
+inline std::ostream& operator<<(std::ostream& os, const float16_t& v) {
+  return os << static_cast<float>(v);
+}
+inline std::ostream& operator<<(std::ostream& os, const bfloat16_t& v) {
+  return os << static_cast<float>(v);
+}
+
+template <typename Vec, typename = std::enable_if_t<is_vector_v<Vec>>>
+inline std::ostream& operator<<(std::ostream& os, const Vec& v) {
+  os << "(";
+  for (auto it = v.begin(); it != v.end(); ++it) {
+    os << *it;
+    if (it != std::prev(v.end())) {
+      os << ",";
+    }
+  }
+  os << ")";
+  return os;
+}
+
+inline bool is_power_of_2(int n) {
+  return ((n & (n - 1)) == 0) && n != 0;
+}
+
+inline int next_power_of_2(int n) {
+  if (is_power_of_2(n)) {
+    return n;
+  }
+  return pow(2, std::ceil(std::log2(n)));
+}
+
+namespace env {
+
+int get_var(const char* name, int default_value);
+
+inline int bfs_max_width() {
+  static int bfs_max_width_ = get_var("MLX_BFS_MAX_WIDTH", 20);
+  return bfs_max_width_;
+}
+
+inline int max_ops_per_buffer(int default_value) {
+  static int max_ops_per_buffer_ =
+      get_var("MLX_MAX_OPS_PER_BUFFER", default_value);
+  return max_ops_per_buffer_;
+}
+
+inline int max_mb_per_buffer(int default_value) {
+  static int max_mb_per_buffer_ =
+      get_var("MLX_MAX_MB_PER_BUFFER", default_value);
+  return max_mb_per_buffer_;
+}
+
+inline bool metal_fast_synch() {
+  static bool metal_fast_synch = get_var("MLX_METAL_FAST_SYNCH", 0);
+  return metal_fast_synch;
+}
+
+inline bool enable_tf32() {
+  static bool enable_tf32_ = get_var("MLX_ENABLE_TF32", 1);
+  return enable_tf32_;
+}
+
+inline int nccl_timeout(int default_value) {
+  static int nccl_timeout = get_var("MLX_NCCL_TIMEOUT", default_value);
+  return nccl_timeout;
+}
+
+} // namespace env
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/version.cpp b/macos/third_party/mlx/mlx/version.cpp
new file mode 100644
index 0000000..79f5948
--- /dev/null
+++ b/macos/third_party/mlx/mlx/version.cpp
@@ -0,0 +1,11 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/version.h"
+
+namespace mlx::core {
+
+const char* version() {
+  return MLX_VERSION;
+}
+
+} // namespace mlx::core
diff --git a/macos/third_party/mlx/mlx/version.h b/macos/third_party/mlx/mlx/version.h
new file mode 100644
index 0000000..0667a98
--- /dev/null
+++ b/macos/third_party/mlx/mlx/version.h
@@ -0,0 +1,22 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/api.h"
+
+#define MLX_VERSION_MAJOR 0
+#define MLX_VERSION_MINOR 30
+#define MLX_VERSION_PATCH 6
+#define MLX_VERSION_NUMERIC \
+  (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH)
+
+namespace mlx::core {
+
+/* A string representation of the MLX version in the format
+ * "major.minor.patch".
+ *
+ * For dev builds, the version will include the suffix ".devYYYYMMDD+hash"
+ */
+MLX_API const char* version();
+
+} // namespace mlx::core
diff --git a/patches/llama.cpp/README.md b/patches/llama.cpp/README.md
new file mode 100644
index 0000000..b2acc27
--- /dev/null
+++ b/patches/llama.cpp/README.md
@@ -0,0 +1,85 @@
+# patches/llama.cpp — the llama engine supply ledger
+
+**Topology discipline**: this directory + `vendor/llama.cpp` (the single pristine source of truth)
+are the entire formal surface of the llama engine supply. The product trees (`windows/`,
+`android/`) may only depend on the repo's `vendor/` + `patches/` + `wt/`. The shared vendor tree
+is a detached checkout of upstream and is **never patched in place** (the build scripts refuse a
+dirty or moved tree).
+
+- **Pin = `vendor.llama.pin` one-liner at the repo root** (currently `7ab4ee7` = tag `b11100`
+  of [ggml-org/llama.cpp](https://github.com/ggml-org/llama.cpp)). The tree is **not** a submodule:
+  the platform scripts clone it on first run (set `EDGE0_LLAMA_URL` for a mirror) and detach exactly
+  at the pin. Shared tree stays detached, and is **never patched in place** (the build scripts refuse
+  a dirty or moved tree).
+- **Consumption model = worktree isolation**: each platform replays its own bands with `git am`
+  inside its own worktree (`wt/win` / `wt/and`) and builds its own artifacts there; the shared
+  tree and the bands never contaminate each other, and builds can run in parallel.
+- **Surface rule (unchanged)**: patches only touch hook points, and every hook is annotated
+  `edge0 private surface #N`. Platform-owned logic lives in glob-injected files
+  (Windows: `windows/serve/*`, Android: `android/engine` pieces wired via `build_vendor_libs.sh`),
+  not in patch caps.
+
+## Band structure & assembly formula
+
+| Platform | Assembly sequence | Count | Golden result-tree hash |
+|---|---|---|---|
+| Windows | `common/*` → `windows/*` | 6+2=8 | `df9d12dcb59aa3314f3356c2cdb777ee4ca4919c` |
+| Android | `common/*` → `android/*` | 6+14=20 | `e974be50ba5c184bf0ba9a4a26adbbae4235c2ab` |
+
+Verify with `git rev-parse HEAD^{tree}` on the worktree **after replay and before copying in the
+own-owned files**. The bands share no files (windows band = `common/arg.cpp` + server params;
+android band = `ggml-cpu/moe_pool` + `src/edge0` hooks + bench debug gates), so switching bands
+is safe: `git am --abort; git reset --hard 7ab4ee7`, then replay.
+
+## common band (6 patches, hook surface shared by both platforms)
+
+| File | Surface | Hook |
+|---|---|---|
+| 0001 | #1 | bailingmoe3 main-graph LoRA wire-up (9 raw `mul_mat` → `build_lora_mm`) |
+| 0002 | #2 | qwen3next/bailingmoe3: capture router input `m_in` into `t_layer_inp` |
+| 0003 | #3 | CMake glob `src/edge0/*.cc` into the llama target (injection channel for own files) |
+| 0004 | #4 | server: wire prerouter advisory `init`/`on_step` |
+| 0005 | #5 | server: wire `edge0_apply_mem_budget` (before model load, env `E0_MEM_BUDGET_MB`) |
+| 0006 | #7 | `ggml_edge0_mmid_base` expert-base resolver: three read points (ggml-cpu.c / iqp.cpp / repack.cpp) + declaration in `ggml-cpu-impl.h`; **NULL = upstream semantics** |
+
+> Why #7 sits in common: the android band registers its pool's row resolver into the same #7
+> surface at first registration, so "swappable expert row base" is one hook point on both
+> platforms instead of two hand-edited rows.
+
+## windows band (2 patches)
+
+| File | Surface | Hook |
+|---|---|---|
+| 0001 | #6 | `--mem-budget-mb` first-class CLI (common_params + arg + server) |
+| 0002 | #8 | `--pool-mb` first-class CLI (env `LLAMA_ARG_POOL_MB` on the same line) |
+
+## android band (14 patches)
+
+0001–0013 = the full demand-pool sequence (phase-1 → v4.7.1, item-by-item parity with the
+engine's streaming/expert_pool contract): copy-in private frames + equal-size slot state machine
++ IO-thread queues + pin/blob/trim/metrics + NEON predictor head + watermark throttle + keepwarm
+refill + reset_all + race hardening. **Hook surface**: `ensure`/`note` dispatch at `MUL_MAT_ID`,
+IQP pooled gate (arm64 never enables IQP anyway = double insurance), qwen3next layer-input bind,
+`moe_wire` registration in `llama-model.cpp` and its call site. 0001/0005 also carry the
+`E0_MMAP_NOPREFETCH` / `E0_NO_REPACK` debug gates.
+
+0014 = consolidation: registers the moe_pool row resolver into the #7 surface (`ggml_edge0_set_mmid_resolver`
+at `g_ready=1`), replacing the direct row rewrite at `src0_cur`; zero disturbance while the pool
+is not ready.
+
+## Refreshing the upstream pin (rebase recipe)
+
+```powershell
+# 1) move the pin: update vendor.llama.pin first (it is the single source of the pin), then
+#    refresh the materialized tree — the scripts detach at the new pin on their own
+git -C vendor/llama.cpp fetch origin
+git -C vendor/llama.cpp checkout --detach <new-commit>
+# 2) re-replay each band three-way on a scratch worktree and re-export it (bands are independent; do them one at a time)
+git -C wt/win reset --hard <new-commit>
+git -C wt/win am --3way (Get-ChildItem patches/llama.cpp/common/*.patch).FullName
+git -C wt/win am --3way (Get-ChildItem patches/llama.cpp/windows/*.patch).FullName
+#    conflicts: resolve by the hook-point semantics above, then `git am --3way --continue`
+# 3) after the gate + tree parity + smoke pass: re-export with format-patch into this directory,
+#    update the golden hashes in the table above, and bump vendor.llama.pin (all scripts read it)
+pwsh -File scripts/patch-green.ps1 -Bands common,windows
+```
diff --git a/patches/llama.cpp/android/0001-moe_pool-android-demand-pool-phase-1-copy-in-private.patch b/patches/llama.cpp/android/0001-moe_pool-android-demand-pool-phase-1-copy-in-private.patch
new file mode 100644
index 0000000..876323d
--- /dev/null
+++ b/patches/llama.cpp/android/0001-moe_pool-android-demand-pool-phase-1-copy-in-private.patch
@@ -0,0 +1,414 @@
+From 6f0158a08cd763eac4bbc2a830ea52fe27996a87 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 20:54:33 +0800
+Subject: [PATCH 01/14] moe_pool: android demand-pool phase-1 (copy-in private
+ frames + lazy row-ptr + expert LRU)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+- ggml-cpu/moe_pool.{c,h}: E0_POOL_MB gate (off by default = zero disturbance); copy-in private
+  frames (Android mlock RLIMIT=64KB rules out the pin-file-page design -- PROBE B-track closed);
+  demand miss = thread0 bulk pread before the barrier ("wholesale"); eviction = self-managed LRU
+  at whole-expert granularity + current-generation protection (rows still in compute are not
+  recycled); non-resident row pointers resolve lazily back to the original mmap (kernel backstop).
+- ggml-cpu.c MUL_MAT_ID, four hooks: include / IQP bypass disabled (self-computed row pointers
+  would skip the pool) / thread0 dispatches ensure after scanning ids / src0_cur row-pointer resolve
+- llama-model.cpp load-time registration: walk ml.contexts + name contains 'exps' + weight.offs
+  taken directly (stride%4096 conservative gate; pread does not move the file cursor; pointer idempotent)
+- Android path gates: llama-model.cpp AUTO mmap_support downgrade skippable (E0_MMAP_NOPREFETCH) +
+  bench E0_NO_REPACK + main.cpp setvbuf (root cause of the 217B truncation misattribution)
+- concurrency protocol = phase-1 single ctx (node barriers keep readers quiescent while ensure
+  writes); server parallel ctx = phase 2
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/CMakeLists.txt  |   1 +
+ ggml/src/ggml-cpu/ggml-cpu.c      |   6 +-
+ ggml/src/ggml-cpu/moe_pool.c      | 188 ++++++++++++++++++++++++++++++
+ ggml/src/ggml-cpu/moe_pool.h      |  40 +++++++
+ src/llama-model.cpp               |  37 +++++-
+ tools/llama-bench/llama-bench.cpp |   2 +
+ tools/llama-bench/main.cpp        |   3 +
+ 7 files changed, 274 insertions(+), 3 deletions(-)
+ create mode 100644 ggml/src/ggml-cpu/moe_pool.c
+ create mode 100644 ggml/src/ggml-cpu/moe_pool.h
+
+diff --git a/ggml/src/ggml-cpu/CMakeLists.txt b/ggml/src/ggml-cpu/CMakeLists.txt
+index 1c7338e..f3c945a 100644
+--- a/ggml/src/ggml-cpu/CMakeLists.txt
++++ b/ggml/src/ggml-cpu/CMakeLists.txt
+@@ -30,6 +30,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
+         ggml-cpu/ggml-cpu.c
+         ggml-cpu/ggml-cpu.cpp
+         ggml-cpu/repack.cpp
++        ggml-cpu/moe_pool.c
+         ggml-cpu/repack.h
+         ggml-cpu/iqp.cpp
+         ggml-cpu/iqp.h
+diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c
+index be7c2db..54f0d9a 100644
+--- a/ggml/src/ggml-cpu/ggml-cpu.c
++++ b/ggml/src/ggml-cpu/ggml-cpu.c
+@@ -15,6 +15,7 @@
+ #include "ops.h"
+ #include "ggml.h"
+ #include "common.h"
++#include "moe_pool.h"
+ 
+ #if defined(_MSC_VER) || defined(__MINGW32__)
+ #include <malloc.h> // using malloc.h with MSC/MINGW
+@@ -1610,7 +1611,7 @@ static void ggml_compute_forward_mul_mat_id(
+ 
+     // IQ panel gemm (see iqp.h); per expert eligibility is decided below, but the work buffer is
+     // reserved for the whole node (ggml_graph_plan sizes it without params, use_ref only skips the dispatch)
+-    const bool iqp = ggml_cpu_iqp_supports_mul_mat_id(dst) && !params->use_ref;
++    const bool iqp = ggml_cpu_iqp_supports_mul_mat_id(dst) && !params->use_ref && !moe_pool_pooled(src0);
+ 
+     char * iqp_panels = NULL;
+ 
+@@ -1672,6 +1673,9 @@ static void ggml_compute_forward_mul_mat_id(
+                 matrix_row_counts[i02] += 1;
+             }
+         }
++
++        // moe_pool: demand-pool dispatch -- bulk-read this ubatch's experts (unregistered tensors return immediately)
++        moe_pool_ensure(src0, (const int32_t *) ids->data, (int) (ids->ne[0] * ids->ne[1]));
+     }
+ 
+     // reset current_chunk
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+new file mode 100644
+index 0000000..2f75ab1
+--- /dev/null
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -0,0 +1,188 @@
++// moe_pool.c -- expert demand pool (Android phase 1, edge0/android-pool; no prediction, synchronous bulk reads).
++// Design = copy-in private frames + lazy row-pointer resolution: a hit = zero-fault direct read of the private frame; a miss =
++// thread0 preads the whole block before the ggml_barrier ("wholesale"); eviction = pool-managed LRU (whole-expert granularity).
++// Concurrency protocol (the premise phase 1 holds on; server parallel ctx is a phase-2 topic):
++//   the ggml-cpu thread pool barriers at EVERY graph-node boundary; ensure runs only on thread0 at the mul_mat_id entry,
++//   before the global barrier -> all readers are quiescent when writing; cross-ctx parallel decode is unguarded (documented red line).
++// Experts outside the pool / that don't fit: resident_slot=-1 -> the caller gets the default pointer and falls back to the original mmap (kernel backstop).
++// env: E0_POOL_MB (0/unset = off, zero disturbance).
++#include "moe_pool.h"
++#include "ggml.h"
++
++#include <pthread.h>
++#include <stdio.h>
++#include <stdlib.h>
++#include <string.h>
++#include <sys/mman.h>
++#include <unistd.h>
++
++#define MOE_POOL_MAX_TENSORS 192   // 40 layers x3 (gate/up/down) = 120 @35b; headroom kept
++
++struct slot_ent {
++    char *   base;           // arena private frame
++    int32_t  expert;         // -1 = free
++    int32_t  tix;
++    uint64_t last_gen;
++    size_t   bytes;
++    struct slot_ent * prev, * next;
++};
++
++struct tensor_pool {
++    const void * key;                // src0->data is stable after load; used as identity
++    int      fd;
++    uint64_t file_off;
++    size_t   slice;                  // private-frame size (page-aligned, >= stride)
++    size_t   stride;                 // per-expert step within the file
++    int64_t  n_experts;
++    int64_t * slot_of;               // [n_experts] slot index or -1 (no tearing concern under the quiesce protocol)
++    long long hits, misses;
++};
++
++static pthread_mutex_t g_lock = PTHREAD_MUTEX_INITIALIZER;
++static size_t   g_cap, g_used;
++static char *   g_arena;
++static size_t   g_bump;
++static uint64_t g_gen;
++static struct tensor_pool g_tp[MOE_POOL_MAX_TENSORS];
++static int      g_ntp;
++static struct slot_ent ** g_slots; static int g_nslots, g_slotcap;
++static struct slot_ent * g_lru_head, * g_lru_tail;
++static int      g_ready;
++static long long g_loads, g_evicts, g_bytes_loaded;
++
++static size_t align_up(size_t x) { return (x + 4095) & ~(size_t)4095; }
++
++void moe_pool_init(void) {
++    if (g_ready) return;
++    const char * e = getenv("E0_POOL_MB");
++    size_t mb = e ? (size_t) atol(e) : 0;
++    if (!mb) return;
++    g_cap = mb * 1024 * 1024;
++    void * m = mmap(NULL, g_cap, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
++    if (m == MAP_FAILED) return;
++    g_arena = m; g_ready = 1;
++}
++
++int moe_pool_enabled(void) { return g_ready; }
++static struct tensor_pool * find_tp(const struct ggml_tensor * t);
++
++void moe_pool_register(struct ggml_tensor * t, int fd, uint64_t file_off) {
++    if (!g_ready) return;
++    if (t->ne[2] < 8) return;                       // expert-stack shape
++    if (find_tp(t)) return;                                 // idempotent: same pointer never re-registered
++    const size_t total = (size_t) ggml_nbytes(t);
++    const size_t stride = total / (size_t) t->ne[2];
++    // conservative gate: pool only when the expert slice is naturally page-aligned (holds for this artifact); otherwise leave it to the kernel
++    if (stride % 4096) return;
++    const size_t slice = stride;
++    pthread_mutex_lock(&g_lock);
++    if (g_ntp >= MOE_POOL_MAX_TENSORS) { pthread_mutex_unlock(&g_lock); return; }
++    struct tensor_pool * p = &g_tp[g_ntp++];
++    p->key = t->data; p->fd = fd; p->file_off = file_off; p->slice = slice; p->stride = stride;
++    p->n_experts = t->ne[2];
++    p->slot_of = calloc(p->n_experts, sizeof(int64_t));
++    for (int64_t i = 0; i < p->n_experts; i++) p->slot_of[i] = -1;
++    pthread_mutex_unlock(&g_lock);
++}
++
++static struct tensor_pool * find_tp(const struct ggml_tensor * t) {
++    for (int i = 0; i < g_ntp; i++) if (g_tp[i].key == t->data) return &g_tp[i];
++    return NULL;
++}
++
++int moe_pool_pooled(const struct ggml_tensor * t) {
++    if (!g_ready || !t || !t->data) return 0;
++    return find_tp(t) != NULL;
++}
++
++// ---- LRU ----
++static void lru_unlink(struct slot_ent * s) {
++    if (s->prev) s->prev->next = s->next; else g_lru_head = s->next;
++    if (s->next) s->next->prev = s->prev; else g_lru_tail = s->prev;
++    s->prev = s->next = NULL;
++}
++static void lru_front(struct slot_ent * s) {
++    s->prev = NULL; s->next = g_lru_head;
++    if (g_lru_head) g_lru_head->prev = s; else g_lru_tail = s;
++    g_lru_head = s;
++}
++
++static struct slot_ent * slot_alloc(size_t bytes, int tix) {
++    struct slot_ent * s = NULL;
++    for (int i = 0; i < g_nslots; i++) {                       // first pick a free slot of the same size
++        if (g_slots[i]->expert == -1 && g_slots[i]->bytes == bytes) { s = g_slots[i]; break; }
++    }
++    if (!s && g_bump + bytes <= g_cap) {                       // then bump-allocate
++        s = calloc(1, sizeof(*s));
++        s->base = g_arena + g_bump; g_bump += bytes; s->bytes = bytes;
++        if (g_nslots == g_slotcap) {
++            g_slotcap = g_slotcap ? g_slotcap * 2 : 512;
++            g_slots = realloc(g_slots, (size_t) g_slotcap * sizeof(*g_slots));
++        }
++        g_slots[g_nslots++] = s;
++    }
++    if (!s) {                                                   // finally evict an older-generation LRU entry
++        for (struct slot_ent * c = g_lru_tail; c; c = c->prev) {
++            if (c->expert < 0 || c->last_gen == g_gen) continue;
++            g_tp[c->tix].slot_of[c->expert] = -1;
++            g_used -= c->bytes;
++            s = c; break;
++        }
++        if (s) { g_evicts++; }
++    }
++    if (!s) return NULL;                                        // pool full this generation / huge block: kernel backstop
++    if (s->expert >= 0) lru_unlink(s);
++    s->tix = tix; s->expert = -1;
++    g_used += s->bytes;
++    return s;
++}
++
++void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_ids) {
++    if (!g_ready) return;
++    pthread_mutex_lock(&g_lock);
++    g_gen++;
++    struct tensor_pool * p = find_tp(t);
++    if (!p) { pthread_mutex_unlock(&g_lock); return; }
++    const int tix = (int) (p - g_tp);
++    for (int i = 0; i < n_ids; i++) {
++        const int32_t e = ids[i];
++        if (e < 0 || e >= p->n_experts) continue;
++        if (p->slot_of[e] >= 0) {                               // hit (slot_of = slot_ent address)
++            p->hits++;
++            struct slot_ent * s = (struct slot_ent *) p->slot_of[e];
++            if (s->last_gen != g_gen) { lru_unlink(s); lru_front(s); s->last_gen = g_gen; }
++            continue;
++        }
++        p->misses++;
++        struct slot_ent * s = slot_alloc(p->slice, tix);
++        if (!s) break;
++        const off_t off = (off_t) (p->file_off + (uint64_t) e * p->stride);
++        const ssize_t got = pread(p->fd, s->base, p->slice, off);
++        if (got != (ssize_t) p->slice) { s->expert = -1; break; }
++        s->expert = e; s->last_gen = g_gen;
++        lru_front(s);
++        g_loads++; g_bytes_loaded += got;
++        p->slot_of[e] = (int64_t) (intptr_t) s;
++    }
++    pthread_mutex_unlock(&g_lock);
++}
++
++const void * moe_pool_row_ptr(const struct ggml_tensor * t, int64_t expert, const void * def) {
++    if (!g_ready || expert < 0) return def;
++    const struct tensor_pool * p = find_tp(t);
++    if (!p || expert >= p->n_experts) return def;
++    const int64_t si = p->slot_of[expert];
++    if (si < 0) return def;
++    const struct slot_ent * s = (const struct slot_ent *) si;
++    return (s->expert == expert) ? s->base : def;               // always true under the quiesce protocol; kept as a defensive check
++}
++
++void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
++                    long long * bytes_loaded, size_t * cap_bytes, size_t * used_bytes) {
++    long long h = 0, m = 0;
++    for (int i = 0; i < g_ntp; i++) { h += g_tp[i].hits; m += g_tp[i].misses; }
++    if (hits) *hits = h; if (misses) *misses = m;
++    if (loads) *loads = g_loads; if (evicts) *evicts = g_evicts;
++    if (bytes_loaded) *bytes_loaded = g_bytes_loaded;
++    if (cap_bytes) *cap_bytes = g_cap; if (used_bytes) *used_bytes = g_used;
++}
+diff --git a/ggml/src/ggml-cpu/moe_pool.h b/ggml/src/ggml-cpu/moe_pool.h
+new file mode 100644
+index 0000000..ab9f7a1
+--- /dev/null
++++ b/ggml/src/ggml-cpu/moe_pool.h
+@@ -0,0 +1,40 @@
++// moe_pool.h -- expert demand pool (Android phase 1, no prediction; fork edge0/android-pool).
++// Foundation = copy-in private frames (Android mlock RLIMIT=64KB is hard-wired; the pin-file-page design is unusable -- see PROBE B-track closure).
++// hit = zero-fault direct read of private memory; miss = ensure() bulk preads the whole block ("wholesale" ~1GB/s vs page-fault "retail" 16MB/s);
++// eviction = this pool's own LRU (whole-expert granularity, never half an expert); long-tail experts outside the pool fall back to their original mmap pages, kernel-backed (lazy pointer resolution).
++#ifndef MOE_POOL_H
++#define MOE_POOL_H
++
++#include <stddef.h>
++#include <stdint.h>
++
++#ifdef __cplusplus
++extern "C" {
++#endif
++
++struct ggml_tensor;
++
++// called from the llama side before the first registration; E0_POOL_MB unset/0 -> never ready (zero disturbance to current behavior)
++void moe_pool_init(void);
++int moe_pool_enabled(void);
++
++// load-time registration: t = expert 3D tensor (src0), file_off = t's data offset in the model file, fd = file descriptor
++void moe_pool_register(struct ggml_tensor * t, int fd, uint64_t file_off);
++
++// query: is this tensor pool-managed (used to turn off the IQP/repack bypass)
++int moe_pool_pooled(const struct ggml_tensor * t);
++
++// thread0 only (before the barrier): bulk-read the experts named by the ids tensor into the pool (LRU makes room)
++void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_ids);
++
++// row-base resolution: resident -> private-frame pointer; otherwise -> default_ptr (original mmap address, kernel backstop)
++const void * moe_pool_row_ptr(const struct ggml_tensor * t, int64_t expert, const void * default_ptr);
++
++// stats snapshot (for debugging / accounting)
++void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
++                    long long * bytes_loaded, size_t * cap_bytes, size_t * used_bytes);
++
++#ifdef __cplusplus
++}
++#endif
++#endif // MOE_POOL_H
+diff --git a/src/llama-model.cpp b/src/llama-model.cpp
+index e195f50..ec1089d 100644
+--- a/src/llama-model.cpp
++++ b/src/llama-model.cpp
+@@ -1,3 +1,13 @@
++#include <cstdlib>
++#include <cstring>
++#include <unistd.h>
++
++// moe_pool in-place declarations (same convention as fork edge0/prefetch.cc; implementation in ggml-cpu/moe_pool.c, linked into the same libllama)
++extern "C" {
++void        moe_pool_init(void);
++int         moe_pool_enabled(void);
++void        moe_pool_register(struct ggml_tensor * t, int fd, uint64_t file_off);
++}
+ #include "llama-model.h"
+ 
+ #include "llama-arch.h"
+@@ -1441,7 +1451,7 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
+ 
+     this->ml = &ml; // to be used by create_tensor() and load_arch_tensors()
+ 
+-    if (ml.use_mmap && params.load_mode == LLAMA_LOAD_MODE_AUTO) {
++    if (ml.use_mmap && params.load_mode == LLAMA_LOAD_MODE_AUTO && !getenv("E0_MMAP_NOPREFETCH")) {
+         for (const auto & dev : devices) {
+             ggml_backend_dev_props props;
+             ggml_backend_dev_get_props(dev.dev, &props);
+@@ -1731,7 +1741,9 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
+         }
+     }
+ 
+-    ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr);
++    // android gate (exp/llama-gguf-perf): E0_MMAP_NOPREFETCH=1 skips whole-file WILLNEED
++    const bool mmap_prefetch = !use_mlock && getenv("E0_MMAP_NOPREFETCH") == nullptr;
++    ml.init_mappings(mmap_prefetch, use_mlock ? &pimpl->mlock_mmaps : nullptr);
+     pimpl->mappings.reserve(ml.mappings.size());
+ 
+     // create the backend buffers
+@@ -1879,6 +1891,27 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
+         }
+     }
+ 
++    // moe_pool registration (Android phase 1): expert 3D tensors -> pool. With E0_POOL_MB unset this block has no side effects.
++    moe_pool_init();
++    if (moe_pool_enabled() && ml.use_mmap && !ml.files.empty()) {
++        auto reg = [&](ggml_tensor * t) {
++            if (!t || !strstr(t->name, "exps")) {
++                return;
++            }
++            const auto * w = ml.get_weight(t->name);
++            if (!w || w->idx >= ml.files.size()) {
++                return;
++            }
++            // pread does not depend on the file cursor, so the loader's fd is safe to use directly (no dup, avoids leaks)
++            moe_pool_register(t, ml.files[w->idx]->file_id(), (uint64_t) w->offs);
++        };
++        for (auto & ctx : ml.contexts) {
++            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t; t = ggml_get_next_tensor(ctx.get(), t)) {
++                reg(t);
++            }
++        }
++    }
++
+     return true;
+ }
+ 
+diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp
+index 6541bb5..eff23ca 100644
+--- a/tools/llama-bench/llama-bench.cpp
++++ b/tools/llama-bench/llama-bench.cpp
+@@ -1229,6 +1229,8 @@ struct cmd_params_instance {
+         llama_model_params mparams = llama_model_default_params();
+ 
+         mparams.n_gpu_layers = n_gpu_layers;
++        // android gate: E0_NO_REPACK=1 -> no extra bufts (plain mmap zero-copy, kernel LRU)
++        if (getenv("E0_NO_REPACK")) { mparams.use_extra_bufts = false; }
+         if (!devices.empty()) {
+             mparams.devices = const_cast<ggml_backend_dev_t *>(devices.data());
+         }
+diff --git a/tools/llama-bench/main.cpp b/tools/llama-bench/main.cpp
+index 0c18bb0..339602e 100644
+--- a/tools/llama-bench/main.cpp
++++ b/tools/llama-bench/main.cpp
+@@ -1,5 +1,8 @@
++#include <cstdio>
++
+ int llama_bench(int argc, char ** argv);
+ 
+ int main(int argc, char ** argv) {
++    setvbuf(stderr, nullptr, _IONBF, 0);  // edge0 diag: unbuffered load log (217B truncation = stderr block-buffer)
+     return llama_bench(argc, argv);
+ }
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0002-moe_pool-demand-exclusive-mode-flash-pool.patch b/patches/llama.cpp/android/0002-moe_pool-demand-exclusive-mode-flash-pool.patch
new file mode 100644
index 0000000..c296183
--- /dev/null
+++ b/patches/llama.cpp/android/0002-moe_pool-demand-exclusive-mode-flash-pool.patch
@@ -0,0 +1,84 @@
+From b8b7b832fad28d5926aea27e0667e98e24e45306 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 21:34:34 +0800
+Subject: [PATCH 02/14] moe_pool: demand-exclusive mode (flash-pool
+ requirements 1 and 3 complete)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+- the 4096 pool-size arm read negative (0.173 < no-pool 0.188), convicting the mixed dual-track:
+  pool hits were bytes the kernel could have served anyway, and the pool's anonymous pages squeezed
+  the page cache -> net negative. Fix = decode small batches (<=16 ids) force evict-old/admit-new
+  and never fall back to mmap -- the pool becomes the only path to experts, expert file pages stay
+  entirely out of the page cache, and the full cache budget goes to dense (1.5GB kept hot); the LRU
+  head/tail order guarantees this batch's just-loaded rows are not evicted (safe).
+- prefill large batches stay best-effort (prevents ping-pong; overflow goes to the kernel).
+- vs the engine flash-pool: (1) sole path and (3) whole-granularity self-eviction are met; (2) the
+  async window needs phase-2 predictive lookahead (the fork #4 hook is already in place); (4) the
+  explicit budget accounting is on the record.
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 20 +++++++++++++++++---
+ 1 file changed, 17 insertions(+), 3 deletions(-)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 2f75ab1..06f4999 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -107,7 +107,7 @@ static void lru_front(struct slot_ent * s) {
+     g_lru_head = s;
+ }
+ 
+-static struct slot_ent * slot_alloc(size_t bytes, int tix) {
++static struct slot_ent * slot_alloc(size_t bytes, int tix, int force) {
+     struct slot_ent * s = NULL;
+     for (int i = 0; i < g_nslots; i++) {                       // first pick a free slot of the same size
+         if (g_slots[i]->expert == -1 && g_slots[i]->bytes == bytes) { s = g_slots[i]; break; }
+@@ -123,7 +123,8 @@ static struct slot_ent * slot_alloc(size_t bytes, int tix) {
+     }
+     if (!s) {                                                   // finally evict an older-generation LRU entry
+         for (struct slot_ent * c = g_lru_tail; c; c = c->prev) {
+-            if (c->expert < 0 || c->last_gen == g_gen) continue;
++            if (c->expert < 0) continue;
++            if (c->last_gen == g_gen && !force) continue;   // demand mode: allow evicting stale entries of THIS generation (just-used rows were already bumped to the LRU head)
+             g_tp[c->tix].slot_of[c->expert] = -1;
+             g_used -= c->bytes;
+             s = c; break;
+@@ -144,6 +145,9 @@ void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_id
+     struct tensor_pool * p = find_tp(t);
+     if (!p) { pthread_mutex_unlock(&g_lock); return; }
+     const int tix = (int) (p - g_tp);
++    // demand-exclusive mode: small batches (decode 1~4 tok) evict-old/admit-new with no mmap fallback -- pool = the only path to experts (flash-pool requirements 1 and 3)
++    // large batches (prefill) load best-effort: whatever doesn't fit is left to the kernel, preventing cross-layer ping-pong
++    const int demand = n_ids <= 16;
+     for (int i = 0; i < n_ids; i++) {
+         const int32_t e = ids[i];
+         if (e < 0 || e >= p->n_experts) continue;
+@@ -154,7 +158,7 @@ void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_id
+             continue;
+         }
+         p->misses++;
+-        struct slot_ent * s = slot_alloc(p->slice, tix);
++        struct slot_ent * s = slot_alloc(p->slice, tix, demand);
+         if (!s) break;
+         const off_t off = (off_t) (p->file_off + (uint64_t) e * p->stride);
+         const ssize_t got = pread(p->fd, s->base, p->slice, off);
+@@ -186,3 +190,13 @@ void moe_pool_stats(long long * hits, long long * misses, long long * loads, lon
+     if (bytes_loaded) *bytes_loaded = g_bytes_loaded;
+     if (cap_bytes) *cap_bytes = g_cap; if (used_bytes) *used_bytes = g_used;
+ }
++
++__attribute__((destructor)) static void moe_pool_dump(void) {
++    if (!g_ready || !getenv("E0_POOL_STATS")) {
++        return;
++    }
++    long long h = 0, m = 0, l = 0, e = 0, bl = 0; size_t cap = 0, used = 0;
++    moe_pool_stats(&h, &m, &l, &e, &bl, &cap, &used);
++    fprintf(stderr, "moe_pool stats: hits=%lld misses=%lld loads=%lld evicts=%lld loaded_mb=%lld cap_mb=%zu used_mb=%zu\n",
++            h, m, l, e, bl / (1024 * 1024), cap / (1024 * 1024), used / (1024 * 1024));
++}
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0003-moe_pool-periodic-stats-dump-destructor-path-bypasse.patch b/patches/llama.cpp/android/0003-moe_pool-periodic-stats-dump-destructor-path-bypasse.patch
new file mode 100644
index 0000000..a37669f
--- /dev/null
+++ b/patches/llama.cpp/android/0003-moe_pool-periodic-stats-dump-destructor-path-bypasse.patch
@@ -0,0 +1,32 @@
+From 4e039a16242eb10852bff812fa0e8acd2400a20e Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 21:35:19 +0800
+Subject: [PATCH 03/14] moe_pool: periodic stats dump (destructor path bypassed
+ by bench exit)
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 7 +++++++
+ 1 file changed, 7 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 06f4999..d27c2fb 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -148,6 +148,13 @@ void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_id
+     // demand-exclusive mode: small batches (decode 1~4 tok) evict-old/admit-new with no mmap fallback -- pool = the only path to experts (flash-pool requirements 1 and 3)
+     // large batches (prefill) load best-effort: whatever doesn't fit is left to the kernel, preventing cross-layer ping-pong
+     const int demand = n_ids <= 16;
++    static long s_calls;
++    if (getenv("E0_POOL_STATS") && (++s_calls % 2000) == 0) {
++        long long h = 0, m = 0, l = 0, e = 0, bl = 0; size_t cap = 0, used = 0;
++        moe_pool_stats(&h, &m, &l, &e, &bl, &cap, &used);
++        fprintf(stderr, "moe_pool@%ld: hits=%lld misses=%lld loads=%lld evicts=%lld used_mb=%zu\n",
++                s_calls, h, m, l, e, used / (1024 * 1024));
++    }
+     for (int i = 0; i < n_ids; i++) {
+         const int32_t e = ids[i];
+         if (e < 0 || e >= p->n_experts) continue;
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0004-moe_pool-v4-full-engine-parity-pool-staging-queues-p.patch b/patches/llama.cpp/android/0004-moe_pool-v4-full-engine-parity-pool-staging-queues-p.patch
new file mode 100644
index 0000000..0b7bb7e
--- /dev/null
+++ b/patches/llama.cpp/android/0004-moe_pool-v4-full-engine-parity-pool-staging-queues-p.patch
@@ -0,0 +1,1013 @@
+From b2ce036a670a9b3657d2342eadf4a6949275601c Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 00:36:44 +0800
+Subject: [PATCH 04/14] moe_pool v4: full engine-parity pool (staging queues,
+ pin, blob, trim, metrics)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Feature parity with the mainline engine's streaming/expert_pool, landed item by item:
+- equal-size slot state machine EMPTY/LOADING/READY/FAILED + IO-thread queue (E0_IO_THREADS, default 4)
+- dual queues: the demand-dedicated queue pops first (mixed FIFO was confirmed to suffer prefetch
+  head-of-line blocking -- the 14s/token pseudo-hang incident)
+- pin protocol: slots in compute are unevictable; released wholesale at the next ensure entry
+  (barrier quiescence contract)
+- acquire hitting LOADING waits on a cv + FAILED falls back synchronously (stall_ms metered, ~ the
+  engine's acq metric)
+- resident blob: one copy-in per non-expert mmap-view tensor (E0_BLOB_MB, isomorphic to the engine's
+  img mirror)
+- watermark trim E0_MEM_FLOOR_MB (when MemAvailable drops below the line, evict non-pinned READY
+  down to 60%)
+- metrics aligned with the engine's PoolMetrics: hits/misses/stall/pf_issued-used-waste-drop-miss/flash/evicts/trims
+- prefetch K=4->8 (teacher n_ids=8 evidences 35b top-8) + throttled drop when the main queue backs up >64
+- registration point moved: tensors at the tail of load_tensors have data=NULL (ml.contexts copies),
+  so llama_model_load now calls llama_model::moe_wire() after load_tensors (tensors_by_name primary
+  set, virtual-interface hook)
+- registered fds are always dup()'d (hang incident: pread got EBADF after the loader closed the original fd)
+- find_tp matches by name + self-corrects the pointer (graph-copy pointers != load-set pointers)
+
+Known differences from the engine (on record): the worker head math is the scalar version (NEON
+queued); trim does no physical tail-block release (the bump arena's middle is not contiguous);
+prep-slot tail pre-conversion is N/A (llama dequantizes ad hoc).
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 714 +++++++++++++++++++++++++++++------
+ ggml/src/ggml-cpu/moe_pool.h |  18 +-
+ src/llama-model.cpp          |  70 +++-
+ src/llama-model.h            |   2 +
+ src/llama.cpp                |   1 +
+ 5 files changed, 667 insertions(+), 138 deletions(-)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index d27c2fb..6f76c90 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -1,92 +1,535 @@
+-// moe_pool.c -- expert demand pool (Android phase 1, edge0/android-pool; no prediction, synchronous bulk reads).
+-// Design = copy-in private frames + lazy row-pointer resolution: a hit = zero-fault direct read of the private frame; a miss =
+-// thread0 preads the whole block before the ggml_barrier ("wholesale"); eviction = pool-managed LRU (whole-expert granularity).
+-// Concurrency protocol (the premise phase 1 holds on; server parallel ctx is a phase-2 topic):
+-//   the ggml-cpu thread pool barriers at EVERY graph-node boundary; ensure runs only on thread0 at the mul_mat_id entry,
+-//   before the global barrier -> all readers are quiescent when writing; cross-ctx parallel decode is unguarded (documented red line).
+-// Experts outside the pool / that don't fit: resident_slot=-1 -> the caller gets the default pointer and falls back to the original mmap (kernel backstop).
+-// env: E0_POOL_MB (0/unset = off, zero disturbance).
++// moe_pool.c -- expert demand pool v4 (implements the engine streaming/expert_pool contract item by item, edge0/android-pool).
++// Design = copy-in private frames (Android mlock RLIMIT=64KB hard-wired; pin-file-page unusable) + equal-size slot state machine
++//   (EMPTY/LOADING/READY/FAILED) + IO-thread queue reads (prefetch enqueues and returns; acquire hitting LOADING waits on a cv;
++//   demand misses are dispatched in batches then awaited together -- mirroring the engine's "async serves prefetch only, cold reads metered as stall" structure, plus batch overlap) +
++//   pin protocol (slots in compute are unevictable; released wholesale at the next ensure entry -- barrier quiescence premise) +
++//   resident blob (the non-expert region pread'd once into a private frame, killing dense mmap "retail" page faults).
++// Determinism contract (same as the engine): slot-state transitions occur only on the main thread under lock; IO threads only write bytes + set terminal states; completion timing affects only
++//   stall (speed), never the generation result. Cross-ctx parallel decode is unguarded (documented red line; phase-2 topic).
++// env: E0_POOL_MB (pool budget, 0/unset = off, zero disturbance) E0_BLOB_MB (non-expert copy-in budget, 0 = off)
++//      E0_IO_THREADS (default 4) E0_MEM_FLOOR_MB (watermark: when MemAvailable falls below this, trim down to 60%)
++//      E0_HEADS / E0_PREROUTER / E0_POOL_STATS (same meanings as v3).
+ #include "moe_pool.h"
+ #include "ggml.h"
+ 
++#include <errno.h>
++#include <fcntl.h>
+ #include <pthread.h>
++#include <stdatomic.h>
+ #include <stdio.h>
+ #include <stdlib.h>
+ #include <string.h>
+ #include <sys/mman.h>
++#include <time.h>
+ #include <unistd.h>
+ 
+ #define MOE_POOL_MAX_TENSORS 192   // 40 layers x3 (gate/up/down) = 120 @35b; headroom kept
++#define MOE_Q_MAX            4096  // IO queue depth (over the limit: prefetch dropped / demand goes synchronous)
++#define MOE_PIN_MAX          512   // per-ensure pin cap on slots (unique experts <= n_experts=256)
++
++enum { ST_EMPTY = 0, ST_LOADING = 1, ST_READY = 2, ST_FAILED = 3 };
++#define SLOT_NONE (-1)
+ 
+ struct slot_ent {
+-    char *   base;           // arena private frame
+-    int32_t  expert;         // -1 = free
++    void *   mem;
++    int32_t  expert;                 // -1 = EMPTY
+     int32_t  tix;
+-    uint64_t last_gen;
+-    size_t   bytes;
+-    struct slot_ent * prev, * next;
++    uint8_t  state;
++    uint8_t  pinned;
++    uint8_t  prefetched;             // slot claimed by prefetch (source of used/wasted accounting)
++    uint8_t  consumed;               // prefetch consumed by acquire (counted once per slot)
++    uint64_t clock;
++    int32_t  lru_prev, lru_next;
++    uint8_t  on_list;
+ };
+ 
+ struct tensor_pool {
+-    const void * key;                // src0->data is stable after load; used as identity
++    const void * key;                // identity pointer (self-corrected by name matching at runtime)
++    char     name[96];               // v4 primary key: the compute tensor may be a graph copy -- pointers differ, names match
++    int      layer;                  // parsed from blk.N.
++    int      role;                   // 0 gate / 1 up / 2 down / -1 non-expert
+     int      fd;
+     uint64_t file_off;
+-    size_t   slice;                  // private-frame size (page-aligned, >= stride)
++    size_t   slice;                  // slot bytes (= stride, the equal-size premise)
+     size_t   stride;                 // per-expert step within the file
+     int64_t  n_experts;
+-    int64_t * slot_of;               // [n_experts] slot index or -1 (no tearing concern under the quiesce protocol)
++    int32_t *slot;                   // [n_experts] -> slot number or SLOT_NONE
+     long long hits, misses;
+ };
+ 
+ static pthread_mutex_t g_lock = PTHREAD_MUTEX_INITIALIZER;
+-static size_t   g_cap, g_used;
++static pthread_cond_t  g_scv  = PTHREAD_COND_INITIALIZER;   // slot-terminal-state / watermark wakeup
++static pthread_cond_t  g_qcv  = PTHREAD_COND_INITIALIZER;   // IO queue wakeup
++
++static size_t   g_cap_bytes;                                 // converted from E0_POOL_MB
++static size_t   g_slice;                                     // fixed at first registration (equal-size)
+ static char *   g_arena;
+-static size_t   g_bump;
+-static uint64_t g_gen;
++static struct slot_ent * g_slots;
++static int      g_nslots, g_used_slots;
++static int32_t  g_empty_cursor;
++static int32_t  g_lru_head = SLOT_NONE, g_lru_tail = SLOT_NONE;  // contains READY/FAILED only
++static uint64_t g_clock;
++static int32_t  g_inflight;
++static int      g_ready, g_io_on, g_stop;
++static long long g_stuckq;   // cumulative count of inflight>0 with both queues empty (tasks vanished)
++
++static int      g_qd[64];    static int g_qdh, g_qdt, g_qdn;              // demand-dedicated queue (higher priority)
++static int      g_qp[MOE_Q_MAX]; static int g_qph, g_qpt, g_qpn;          // prefetch main queue
++
+ static struct tensor_pool g_tp[MOE_POOL_MAX_TENSORS];
+ static int      g_ntp;
+-static struct slot_ent ** g_slots; static int g_nslots, g_slotcap;
+-static struct slot_ent * g_lru_head, * g_lru_tail;
+-static int      g_ready;
+-static long long g_loads, g_evicts, g_bytes_loaded;
++static int32_t  g_pins[MOE_PIN_MAX]; static int g_npins;
+ 
+-static size_t align_up(size_t x) { return (x + 4095) & ~(size_t)4095; }
++// accounting (subset aligned with the engine's PoolMetrics)
++static long long g_loads, g_evicts, g_trims;
++static long long g_flash;                                    // bytes admitted to the pool (worker + synchronous)
++static long long g_iofail;
++static long long g_stall_ns;                                 // cumulative main-thread read waits (~ the engine's stall/acq metric)
++static long long g_pf_issued, g_pf_used, g_pf_waste, g_pf_drop, g_pf_miss, g_rescue;
++
++// resident blob (copy-in of the non-expert region)
++static char *   g_blob; static size_t g_blob_cap, g_blob_bump; static int g_blob_on;
++static volatile int g_heads_on, g_pred_enabled;   // referenced early by init (defined in the pred section)
++
++static uint64_t now_ns(void) {
++    struct timespec ts; clock_gettime(CLOCK_MONOTONIC, &ts);
++    return (uint64_t) ts.tv_sec * 1000000000ull + (uint64_t) ts.tv_nsec;
++}
+ 
+ void moe_pool_init(void) {
+     if (g_ready) return;
+     const char * e = getenv("E0_POOL_MB");
+     size_t mb = e ? (size_t) atol(e) : 0;
++    const char * be = getenv("E0_BLOB_MB");
++    size_t bmb = be ? (size_t) atol(be) : 0;
++    if (bmb) {
++        g_blob_cap = bmb * 1024 * 1024;
++        void * m = mmap(NULL, g_blob_cap, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
++        if (m != MAP_FAILED) { g_blob = m; g_blob_on = 1; }
++    }
+     if (!mb) return;
+-    g_cap = mb * 1024 * 1024;
+-    void * m = mmap(NULL, g_cap, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+-    if (m == MAP_FAILED) return;
+-    g_arena = m; g_ready = 1;
++    g_cap_bytes = mb * 1024 * 1024;                          // slots are created at first registration (once slice is known)
++    if (getenv("E0_POOL_STATS")) fprintf(stderr, "moe_pool init: pool_mb=%zu blob_on=%d heads_on=%d pred=%d\n", mb, g_blob_on, g_heads_on, g_pred_enabled);
++    const char * hp = getenv("E0_HEADS");
++    if (hp && hp[0]) {
++        moe_pool_load_heads(hp);
++    }
+ }
+ 
+-int moe_pool_enabled(void) { return g_ready; }
+ static struct tensor_pool * find_tp(const struct ggml_tensor * t);
+ 
++int moe_pool_enabled(void) { return g_ready; }
++int moe_pool_blob_enabled(void) { return g_blob_on; }
++
++void * moe_pool_blob_add(size_t nbytes) {
++    if (!g_blob_on || !nbytes) return NULL;
++    const size_t need = (nbytes + 4095) & ~(size_t)4095;
++    if (g_blob_bump + need > g_blob_cap) return NULL;        // budget exhausted: this tensor stays in mmap (see dump for accounting)
++    void * p = g_blob + g_blob_bump;
++    g_blob_bump += need;
++    return p;
++}
++
++// ===================== v3 -> v4: predictive prefetch engine (advisory head; the pool is read-only and never touches logits) =====================
++// Math = PREROUTER_SPEC sec.2: logits = linear_init(x) + fc2(gelu_erf(fc1(x))),
++//   x = [m_in(2048) | cur_onehot(256) | prev_onehot(256)]; the top-K prediction set is used only for prefetch; routing remains teacher.
++#include <math.h>
++#define PRED_RINGSZ 128
++struct pred_job {
++    int      layer;
++    int      ntok;
++    float    x[2560];                    // feats: 0..2047 = m_in last-token row, 2048..2303 = cur_oh, 2304..2559 = prev_oh
++    int32_t  cur_ids[8];  int n_cur;
++    int32_t  prev_ids[8]; int n_prev;
++};
++static struct {
++    int      bound;
++    const struct ggml_tensor * min;
++    int64_t  tix[3];                     // g_tp indices of gate/up/down, filled back at registration
++    struct { const uint16_t * fc1, * fc2, * lin; } hd;
++    int32_t  prev_exec[8]; int n_prev;
++} g_lr[64];
++static struct pred_job g_ring[PRED_RINGSZ];
++static volatile int g_ring_head, g_ring_tail, g_worker_on;
++static char * g_headmap; static size_t g_headmap_sz;
++
++static void * pred_worker(void * arg);
++struct tensor_pool;
++void moe_pool_prefetch_impl(struct tensor_pool * p, int64_t e);
++
++int moe_pool_load_heads(const char * path) {
++    int fd = open(path, O_RDONLY);
++    if (fd < 0) return 0;
++    size_t sz = (size_t) lseek(fd, 0, SEEK_END); lseek(fd, 0, SEEK_SET);
++    void * m = mmap(NULL, sz, PROT_READ, MAP_PRIVATE, fd, 0);
++    close(fd);
++    if (m == MAP_FAILED) return 0;
++    g_headmap = m; g_headmap_sz = sz;
++    uint32_t magic; memcpy(&magic, m, 4);
++    if (magic != 0x44483045u) return 0;
++    int n; memcpy(&n, (char *) m + 4, 4);
++    size_t off = 8;
++    for (int i = 0; i < n; i++) {
++        int layer; memcpy(&layer, (char *) m + off, 4); off += 4;
++        uint32_t l1, l2, l3;
++        memcpy(&l1, (char *) m + off, 4); memcpy(&l2, (char *) m + off + 4, 4); memcpy(&l3, (char *) m + off + 8, 4);
++        off += 12;
++        if (layer >= 0 && layer < 64) {
++            g_lr[layer].hd.fc1 = (const uint16_t *) ((char *) m + off);
++            g_lr[layer].hd.fc2 = (const uint16_t *) ((char *) m + off + l1);
++            g_lr[layer].hd.lin = (const uint16_t *) ((char *) m + off + l1 + l2);
++        }
++        off += (size_t) l1 + l2 + l3;
++    }
++    g_heads_on = 1;
++    g_pred_enabled = getenv("E0_PREROUTER") != NULL;
++    return 1;
++}
++
++void moe_pool_bind_min(int layer, const struct ggml_tensor * t) {
++    if (layer >= 0 && layer < 64) { g_lr[layer].min = t; g_lr[layer].bound = 1; }
++}
++
++void moe_pool_note_routing(const struct ggml_tensor * src0, const int32_t * ids, int n_ids) {
++    static int s_nr1;
++    if (getenv("E0_POOL_STATS") && !s_nr1) { s_nr1 = 1; fprintf(stderr, "moe_pool note#1: pred=%d heads=%d n_ids=%d name=%s\n", g_pred_enabled, g_heads_on, n_ids, src0->name); }
++    if (!g_pred_enabled || !g_heads_on || n_ids <= 0 || n_ids > 8) return;
++    int layer = -1; const char * q = strstr(src0->name, "blk.");
++    if (q) layer = (int) strtol(q + 4, NULL, 10);
++    if (layer < 0 || layer >= 64) return;
++    if (!g_lr[layer].bound || !g_lr[layer].min || !g_lr[layer].hd.fc1) return;
++    const struct ggml_tensor * min = g_lr[layer].min;
++    if (!min->data) return;
++    const int64_t n_embd = min->ne[0];
++    if (n_embd > 2048 || min->ne[1] < 1) return;
++    pthread_mutex_lock(&g_lock);
++    if (g_ring_head - g_ring_tail >= PRED_RINGSZ) { pthread_mutex_unlock(&g_lock); return; }  // backpressure: drop the prefetch
++    struct pred_job * j = &g_ring[g_ring_head % PRED_RINGSZ];
++    g_ring_head++;
++    j->layer = layer; j->n_cur = n_ids;
++    const float * row = (const float *) ((const char *) min->data + (min->ne[1] - 1) * min->nb[1]);
++    memcpy(j->x, row, (size_t) n_embd * sizeof(float));
++    memset(j->x + n_embd, 0, (2048 - (size_t) n_embd) * sizeof(float));
++    memset(&j->x[2048], 0, (2560 - 2048) * sizeof(float));
++    for (int i = 0; i < n_ids; i++) if (ids[i] >= 0 && ids[i] < 256) j->x[2048 + ids[i]] = 1.0f;
++    j->n_prev = g_lr[layer].n_prev;
++    for (int i = 0; i < j->n_prev; i++) {
++        const int32_t e = g_lr[layer].prev_exec[i];
++        if (e >= 0 && e < 256) j->x[2304 + e] = 1.0f;
++    }
++    memcpy(j->cur_ids, ids, (size_t) n_ids * sizeof(int32_t));
++    memcpy(g_lr[layer].prev_exec, ids, (size_t) n_ids * sizeof(int32_t));   // teacher facts go into prev
++    g_lr[layer].n_prev = n_ids;
++    if (!g_worker_on) {
++        g_worker_on = 1;
++        pthread_t th;
++        pthread_create(&th, NULL, pred_worker, NULL);
++        pthread_detach(th);
++    }
++    pthread_mutex_unlock(&g_lock);
++}
++
++static inline float f16tof(uint16_t h) {
++    const uint32_t w = (uint32_t) h << 16;
++    float f; memcpy(&f, &w, 4); return f;
++}
++static inline float gelu_erf_f(float x) { return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f)); }
++
++static void pred_do_job(struct pred_job * j) {
++    float h[512], out[256];
++    const uint16_t * fc1 = g_lr[j->layer].hd.fc1;
++    const uint16_t * fc2 = g_lr[j->layer].hd.fc2;
++    const uint16_t * lin = g_lr[j->layer].hd.lin;
++    for (int o = 0; o < 512; o++) {
++        const uint16_t * wrow = fc1 + (size_t) o * 2560;
++        float acc = 0;
++        for (int i = 0; i < 2560; i++) acc += f16tof(wrow[i]) * j->x[i];
++        h[o] = gelu_erf_f(acc);
++    }
++    for (int o = 0; o < 256; o++) {
++        const uint16_t * wr = lin + (size_t) o * 2560;
++        const uint16_t * wc = fc2 + (size_t) o * 512;
++        float acc = 0;
++        for (int i = 0; i < 2560; i++) acc += f16tof(wr[i]) * j->x[i];
++        for (int i = 0; i < 512; i++) acc += f16tof(wc[i]) * h[i];
++        out[o] = acc;
++    }
++    int top[8]; float tv[8]; for (int k = 0; k < 8; k++) { tv[k] = -1e30f; top[k] = -1; }
++    for (int e = 0; e < 256; e++) {
++        for (int k = 0; k < 8; k++) {
++            if (out[e] > tv[k]) {
++                for (int m = 7; m > k; m--) { tv[m] = tv[m-1]; top[m] = top[m-1]; }
++                tv[k] = out[e]; top[k] = e; break;
++            }
++        }
++    }
++    for (int k = 0; k < 8; k++) {
++        if (top[k] < 0) break;
++        for (int role = 0; role < 3; role++) {
++            const int64_t tix = g_lr[j->layer].tix[role];
++            if (tix < 0) continue;
++            moe_pool_prefetch_impl(&g_tp[tix], top[k]);
++        }
++    }
++}
++
++static void * pred_worker(void * arg) {
++    (void) arg;
++    while (1) {
++        int head_i = -1;
++        pthread_mutex_lock(&g_lock);
++        if (g_ring_head > g_ring_tail) head_i = g_ring_tail++;
++        pthread_mutex_unlock(&g_lock);
++        if (head_i < 0) { usleep(200); continue; }
++        pred_do_job(&g_ring[head_i % PRED_RINGSZ]);
++    }
++    return NULL;
++}
++
++// ===================== slot / LRU / queue core (v4) =====================
++
++static void list_unlink(int32_t i) {
++    struct slot_ent * s = &g_slots[i];
++    if (!s->on_list) return;
++    if (s->lru_prev != SLOT_NONE) g_slots[s->lru_prev].lru_next = s->lru_next; else g_lru_head = s->lru_next;
++    if (s->lru_next != SLOT_NONE) g_slots[s->lru_next].lru_prev = s->lru_prev; else g_lru_tail = s->lru_prev;
++    s->lru_prev = s->lru_next = SLOT_NONE; s->on_list = 0;
++}
++static void list_push_back(int32_t i) {
++    struct slot_ent * s = &g_slots[i];
++    if (s->on_list) list_unlink(i);
++    s->lru_prev = g_lru_tail; s->lru_next = SLOT_NONE;
++    if (g_lru_tail != SLOT_NONE) g_slots[g_lru_tail].lru_next = i; else g_lru_head = i;
++    g_lru_tail = i; s->on_list = 1;
++}
++static void account_evict(struct slot_ent * s) {              // same as the engine's account_prefetch_eviction
++    if (s->prefetched && !s->consumed) g_pf_miss++;
++    s->prefetched = 0; s->consumed = 0;
++}
++static int32_t evict_victim_locked(void) {                    // first non-pinned entry at the LRU head (chain contains READY/FAILED only)
++    for (int32_t i = g_lru_head; i != SLOT_NONE; ) {
++        int32_t nxt = g_slots[i].lru_next;
++        if (!g_slots[i].pinned) return i;
++        i = nxt;
++    }
++    return SLOT_NONE;
++}
++static void unclaim_locked(int32_t si) {                      // claim a slot then back out (e.g. queue full)
++    struct slot_ent * s = &g_slots[si];
++    g_tp[s->tix].slot[s->expert] = SLOT_NONE;
++    account_evict(s);
++    s->state = ST_EMPTY; s->expert = -1; s->tix = -1;
++    g_used_slots--;
++}
++// Claim a slot for (tix,e): EMPTY cursor first, then the non-pinned LRU head (READY/FAILED). LOADING is never evictable.
++static int32_t claim_slot_locked(int tix, int32_t e) {
++    int32_t si = SLOT_NONE;
++    for (int t = 0; t < g_nslots; t++) {
++        int32_t i = (g_empty_cursor + t) % g_nslots;
++        if (g_slots[i].state == ST_EMPTY) { si = i; g_empty_cursor = (i + 1) % g_nslots; g_used_slots++; break; }
++    }
++    if (si < 0) {
++        si = evict_victim_locked();
++        if (si < 0) return SLOT_NONE;
++        struct slot_ent * vs = &g_slots[si];
++        g_tp[vs->tix].slot[vs->expert] = SLOT_NONE;
++        if (vs->state == ST_READY) g_evicts++;
++        account_evict(vs);
++        list_unlink(si);
++    }
++    struct slot_ent * s = &g_slots[si];
++    s->tix = tix; s->expert = e; s->state = ST_LOADING;
++    s->pinned = 0; s->prefetched = 0; s->consumed = 0; s->clock = ++g_clock;
++    g_tp[tix].slot[e] = si;
++    return si;
++}
++static int enqueue_io_locked(int32_t si, int demand) {   // demand=1 goes to the dedicated queue (popped first)
++    if (demand) {
++        if (g_qdn >= 64) return 0;
++        g_qd[g_qdt] = si; g_qdt = (g_qdt + 1) % 64; g_qdn++;
++    } else {
++        if (g_qpn >= MOE_Q_MAX) return 0;
++        g_qp[g_qpt] = si; g_qpt = (g_qpt + 1) % MOE_Q_MAX; g_qpn++;
++    }
++    g_inflight++;
++    pthread_cond_broadcast(&g_qcv);
++    return 1;
++}
++static ssize_t pread_full(int fd, void * dst, size_t n, off_t off) {
++    size_t got = 0;
++    int tries = 0;
++    while (got < n) {
++        ssize_t r = pread(fd, (char *) dst + got, n - got, off + (off_t) got);
++        if (r > 0) { got += (size_t) r; continue; }
++        if (r == 0) break;                                  // EOF: real failure
++        if (errno == EINTR) continue;
++        if (++tries < 200) { struct timespec ts = { 0, 1000000 }; nanosleep(&ts, NULL); continue; }
++        static int s_once;
++        if (!s_once) { s_once = 1; fprintf(stderr, "moe_pool pread ERRNO=%d fd=%d off=%lld\n", errno, fd, (long long) off); }
++        break;                                              // permanent error: give up (FAILED -> kernel fallback, never hangs)
++    }
++    return (ssize_t) got == n ? (ssize_t) n : -1;
++}
++static void slot_finalize_locked(int32_t si, int ok) {
++    static int s_fin8;
++    if (getenv("E0_POOL_STATS") && s_fin8 < 8) { s_fin8++; fprintf(stderr, "moe_pool io-fin#%d: slot=%d ok=%d qn=%d inf=%d\n", s_fin8, si, ok, g_qdn + g_qpn, g_inflight); }
++    struct slot_ent * s = &g_slots[si];
++    if (ok) { s->state = ST_READY; g_flash += g_slice; }
++    else    { s->state = ST_FAILED; g_iofail++; }
++    list_push_back(si);
++    if (g_inflight > 0) g_inflight--;
++    pthread_cond_broadcast(&g_scv);
++}
++// synchronous read on the main thread outside the lock (same as the engine's sync_load_locked; no one can move a LOADING slot, so it's safe)
++static int load_sync_locked(int32_t si) {
++    struct slot_ent * s = &g_slots[si];
++    struct tensor_pool * p = &g_tp[s->tix];
++    const uint64_t off = p->file_off + (uint64_t) s->expert * p->stride;
++    pthread_mutex_unlock(&g_lock);
++    const ssize_t got = pread_full(p->fd, s->mem, p->slice, (off_t) off);
++    pthread_mutex_lock(&g_lock);
++    if (got == (ssize_t) p->slice) { g_loads++; slot_finalize_locked(si, 1); return 1; }
++    slot_finalize_locked(si, 0);
++    return 0;
++}
++static void wait_ready_locked(int32_t si) {
++    const uint64_t t0 = now_ns();
++    while (g_slots[si].state == ST_LOADING && !g_stop) { if (g_inflight > 0 && g_qdn == 0 && g_qpn == 0) { g_stuckq++; if (g_stuckq < 6) fprintf(stderr, "moe_pool STUCKQ#%lld: si=%d infl=%d used=%d\n", g_stuckq, si, g_inflight, g_used_slots); } pthread_cond_wait(&g_scv, &g_lock); }
++    g_stall_ns += (long long) (now_ns() - t0);
++}
++static void pin_locked(int32_t si) {
++    struct slot_ent * s = &g_slots[si];
++    if (s->pinned) return;
++    s->pinned = 1;
++    if (g_npins < MOE_PIN_MAX) g_pins[g_npins++] = si;
++}
++
++static void * io_thread(void * arg) {
++    (void) arg;
++    for (;;) {
++        int32_t si;
++        pthread_mutex_lock(&g_lock);
++        while (g_qdn + g_qpn == 0 && !g_stop) pthread_cond_wait(&g_qcv, &g_lock);
++        if (g_stop) { pthread_mutex_unlock(&g_lock); return NULL; }
++        if (g_qdn > 0)      { si = g_qd[g_qdh]; g_qdh = (g_qdh + 1) % 64; g_qdn--; }
++        else if (g_qpn > 0) { si = g_qp[g_qph]; g_qph = (g_qph + 1) % MOE_Q_MAX; g_qpn--; }
++        else                { pthread_mutex_unlock(&g_lock); continue; }
++        struct slot_ent * s = &g_slots[si];
++        const int fd = g_tp[s->tix].fd;
++        const uint64_t off = g_tp[s->tix].file_off + (uint64_t) s->expert * g_tp[s->tix].stride;
++        void * mem = s->mem;
++        const size_t slice = g_slice;
++        pthread_mutex_unlock(&g_lock);
++        const ssize_t got = pread_full(fd, mem, slice, (off_t) off);
++        pthread_mutex_lock(&g_lock);
++        if (g_slots[si].state == ST_LOADING && g_slots[si].mem == mem) {
++            if (got == (ssize_t) slice) { g_loads++; slot_finalize_locked(si, 1); }
++            else slot_finalize_locked(si, 0);
++        } else {
++            if (g_inflight > 0) g_inflight--;                 // main thread already voided it: this read's result is discarded
++            pthread_cond_broadcast(&g_scv);
++        }
++        pthread_mutex_unlock(&g_lock);
++    }
++}
++
++static long long mem_avail_mb(void) {
++    FILE * f = fopen("/proc/meminfo", "r");
++    if (!f) return -1;
++    char ln[256]; long long v = -1;
++    while (fgets(ln, sizeof(ln), f)) {
++        if (!strncmp(ln, "MemAvailable:", 13)) { v = atoll(ln + 13) / 1024; break; }
++    }
++    fclose(f);
++    return v;
++}
++// watermark trim (simplified form of the engine's trim): evict only non-pinned READY; no physical frame reclaim (the arena middle is not contiguous -- see ALIGNMENT.md).
++static void water_trim_locked(void) {
++    const char * fe = getenv("E0_MEM_FLOOR_MB");
++    if (!fe) return;
++    const long long floor_mb = atoll(fe);
++    if (floor_mb <= 0) return;
++    const long long avail = mem_avail_mb();
++    if (avail < 0 || avail >= floor_mb) return;
++    const int target = (int) (g_nslots * 6 / 10);
++    int n = 0;
++    while (g_used_slots > target) {
++        int32_t v = SLOT_NONE;
++        for (int32_t i = g_lru_head; i != SLOT_NONE; ) {
++            int32_t nxt = g_slots[i].lru_next;
++            if (!g_slots[i].pinned) { v = i; break; }
++            i = nxt;
++        }
++        if (v < 0) break;
++        struct slot_ent * s = &g_slots[v];
++        g_tp[s->tix].slot[s->expert] = SLOT_NONE;
++        account_evict(s);
++        list_unlink(v);
++        s->state = ST_EMPTY; s->expert = -1;
++        g_used_slots--; g_trims++; n++;
++    }
++    if (n) fprintf(stderr, "moe_pool trim: evicted=%d used=%d MemAvailable=%lldMB\n", n, g_used_slots, avail);
++}
++
++// ===================== public API =====================
++
+ void moe_pool_register(struct ggml_tensor * t, int fd, uint64_t file_off) {
+-    if (!g_ready) return;
+-    if (t->ne[2] < 8) return;                       // expert-stack shape
+-    if (find_tp(t)) return;                                 // idempotent: same pointer never re-registered
++    if (!g_cap_bytes) return;
++    if (getenv("E0_POOL_STATS") && g_ntp == 0) {
++        fprintf(stderr, "moe_pool reg#1: %s ne2=%lld nbytes=%zu\n", t->name, (long long) t->ne[2], (size_t) ggml_nbytes(t));
++    }
++    if (t->ne[2] < 8) return;                                // expert-stack shape
++    if (find_tp(t)) return;                                  // idempotent: same pointer never re-registered
+     const size_t total = (size_t) ggml_nbytes(t);
+     const size_t stride = total / (size_t) t->ne[2];
+-    // conservative gate: pool only when the expert slice is naturally page-aligned (holds for this artifact); otherwise leave it to the kernel
+-    if (stride % 4096) return;
+-    const size_t slice = stride;
++    if (stride % 4096) return;                               // conservative gate: only page-aligned slices get pooled (holds for this artifact)
++    if (!g_ready) {                                          // fixed at first registration: equal-size slot array
++        g_slice = stride;
++        g_nslots = (int) (g_cap_bytes / g_slice);
++        g_arena = mmap(NULL, g_slice * (size_t) g_nslots, PROT_READ | PROT_WRITE,
++                       MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
++        if (g_arena == MAP_FAILED) { g_arena = NULL; g_nslots = 0; return; }
++        g_slots = calloc((size_t) g_nslots, sizeof(struct slot_ent));
++        for (int i = 0; i < g_nslots; i++) { g_slots[i].mem = g_arena + (size_t) i * g_slice; g_slots[i].expert = -1; g_slots[i].tix = -1; }
++        g_ready = 1;
++        const char * ie = getenv("E0_IO_THREADS");
++        int nth = ie ? (int) atol(ie) : 4;
++        if (nth < 1) nth = 1; if (nth > 8) nth = 8;
++        if (!g_io_on) {
++            g_io_on = 1;
++            for (int i = 0; i < nth; i++) {
++                pthread_t th; pthread_create(&th, NULL, io_thread, NULL); pthread_detach(th);
++            }
++        }
++    }
++    if (stride != g_slice) return;                           // the v4 premise is equal-size slots (same as the engine); odd-shaped tensors go to the kernel
+     pthread_mutex_lock(&g_lock);
+     if (g_ntp >= MOE_POOL_MAX_TENSORS) { pthread_mutex_unlock(&g_lock); return; }
+     struct tensor_pool * p = &g_tp[g_ntp++];
+-    p->key = t->data; p->fd = fd; p->file_off = file_off; p->slice = slice; p->stride = stride;
++    p->layer = -1; p->role = -1;
++    { const char * q = strstr(t->name, "blk."); if (q) p->layer = (int) strtol(q + 4, NULL, 10); }
++    if (strstr(t->name, "gate_exps")) p->role = 0;
++    else if (strstr(t->name, "up_exps")) p->role = 1;
++    else if (strstr(t->name, "down_exps")) p->role = 2;
++    if (p->role >= 0 && p->layer >= 0 && p->layer < 64) {
++        g_lr[p->layer].tix[p->role] = g_ntp - 1;
++    }
++    p->key = t->data; snprintf(p->name, sizeof(p->name), "%s", t->name); p->fd = dup(fd);   // independent ownership: pread still works after the loader closed its fd (v4 hang incident) p->file_off = file_off; p->slice = stride; p->stride = stride;
+     p->n_experts = t->ne[2];
+-    p->slot_of = calloc(p->n_experts, sizeof(int64_t));
+-    for (int64_t i = 0; i < p->n_experts; i++) p->slot_of[i] = -1;
++    p->slot = calloc((size_t) p->n_experts, sizeof(int32_t));
++    for (int64_t i = 0; i < p->n_experts; i++) p->slot[i] = SLOT_NONE;
++    if (getenv("E0_POOL_STATS") && g_ntp <= 2) {
++        fprintf(stderr, "moe_pool reg-ok#%d: %s nslots=%d\n", g_ntp, t->name, g_nslots);
++    }
+     pthread_mutex_unlock(&g_lock);
+ }
+ 
+ static struct tensor_pool * find_tp(const struct ggml_tensor * t) {
+     for (int i = 0; i < g_ntp; i++) if (g_tp[i].key == t->data) return &g_tp[i];
++    // v4: the tensor used by compute may be a graph copy (pointer != load-time set) -- match by name and re-bind the pointer
++    for (int i = 0; i < g_ntp; i++) {
++        if (!strcmp(g_tp[i].name, t->name)) { g_tp[i].key = t->data; return &g_tp[i]; }
++    }
+     return NULL;
+ }
+ 
+@@ -95,115 +538,152 @@ int moe_pool_pooled(const struct ggml_tensor * t) {
+     return find_tp(t) != NULL;
+ }
+ 
+-// ---- LRU ----
+-static void lru_unlink(struct slot_ent * s) {
+-    if (s->prev) s->prev->next = s->next; else g_lru_head = s->next;
+-    if (s->next) s->next->prev = s->prev; else g_lru_tail = s->prev;
+-    s->prev = s->next = NULL;
+-}
+-static void lru_front(struct slot_ent * s) {
+-    s->prev = NULL; s->next = g_lru_head;
+-    if (g_lru_head) g_lru_head->prev = s; else g_lru_tail = s;
+-    g_lru_head = s;
+-}
+-
+-static struct slot_ent * slot_alloc(size_t bytes, int tix, int force) {
+-    struct slot_ent * s = NULL;
+-    for (int i = 0; i < g_nslots; i++) {                       // first pick a free slot of the same size
+-        if (g_slots[i]->expert == -1 && g_slots[i]->bytes == bytes) { s = g_slots[i]; break; }
+-    }
+-    if (!s && g_bump + bytes <= g_cap) {                       // then bump-allocate
+-        s = calloc(1, sizeof(*s));
+-        s->base = g_arena + g_bump; g_bump += bytes; s->bytes = bytes;
+-        if (g_nslots == g_slotcap) {
+-            g_slotcap = g_slotcap ? g_slotcap * 2 : 512;
+-            g_slots = realloc(g_slots, (size_t) g_slotcap * sizeof(*g_slots));
+-        }
+-        g_slots[g_nslots++] = s;
+-    }
+-    if (!s) {                                                   // finally evict an older-generation LRU entry
+-        for (struct slot_ent * c = g_lru_tail; c; c = c->prev) {
+-            if (c->expert < 0) continue;
+-            if (c->last_gen == g_gen && !force) continue;   // demand mode: allow evicting stale entries of THIS generation (just-used rows were already bumped to the LRU head)
+-            g_tp[c->tix].slot_of[c->expert] = -1;
+-            g_used -= c->bytes;
+-            s = c; break;
+-        }
+-        if (s) { g_evicts++; }
+-    }
+-    if (!s) return NULL;                                        // pool full this generation / huge block: kernel backstop
+-    if (s->expert >= 0) lru_unlink(s);
+-    s->tix = tix; s->expert = -1;
+-    g_used += s->bytes;
+-    return s;
+-}
+-
+ void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_ids) {
+-    if (!g_ready) return;
++    if (!g_ready || n_ids <= 0) return;
+     pthread_mutex_lock(&g_lock);
+-    g_gen++;
++    // release all pins from the previous node wholesale (barrier quiescence premise: all compute threads of the previous tensor finished)
++    for (int i = 0; i < g_npins; i++) g_slots[g_pins[i]].pinned = 0;
++    g_npins = 0;
+     struct tensor_pool * p = find_tp(t);
++    static int s_np;
++    if (getenv("E0_POOL_STATS") && s_np < 3) { s_np++; fprintf(stderr, "moe_pool ensure#%d: %s ptr=%p buf=%p ntp=%d key0=%p\n", s_np, t->name, t->data, (void*)t->buffer, g_ntp, g_ntp ? g_tp[0].key : NULL); }
+     if (!p) { pthread_mutex_unlock(&g_lock); return; }
+     const int tix = (int) (p - g_tp);
+-    // demand-exclusive mode: small batches (decode 1~4 tok) evict-old/admit-new with no mmap fallback -- pool = the only path to experts (flash-pool requirements 1 and 3)
+-    // large batches (prefill) load best-effort: whatever doesn't fit is left to the kernel, preventing cross-layer ping-pong
+-    const int demand = n_ids <= 16;
++    const int exclusive = n_ids <= 16;                       // decode: demand-exclusive; prefill: best-effort
+     static long s_calls;
+     if (getenv("E0_POOL_STATS") && (++s_calls % 2000) == 0) {
+-        long long h = 0, m = 0, l = 0, e = 0, bl = 0; size_t cap = 0, used = 0;
+-        moe_pool_stats(&h, &m, &l, &e, &bl, &cap, &used);
+-        fprintf(stderr, "moe_pool@%ld: hits=%lld misses=%lld loads=%lld evicts=%lld used_mb=%zu\n",
+-                s_calls, h, m, l, e, used / (1024 * 1024));
++        long long h = 0, m = 0;
++        for (int i = 0; i < g_ntp; i++) { h += g_tp[i].hits; m += g_tp[i].misses; }
++        fprintf(stderr, "moe_pool@%ld: hits=%lld misses=%lld loads=%lld evicts=%lld used_mb=%zu"
++                        " stall_ms=%lld pf=%lld/%lld/%lld/%lld flash_mb=%lld trims=%lld\n",
++                s_calls, h, m, g_loads, g_evicts, (size_t) g_used_slots * g_slice / (1024 * 1024),
++                g_stall_ns / 1000000, g_pf_issued, g_pf_used, g_pf_waste, g_pf_miss,
++                g_flash / (1024 * 1024), g_trims);
+     }
++    water_trim_locked();
++
++    int32_t waiters[64]; int nwait = 0;                      // slots enqueued in this batch; awaited together after dispatch (queue parallelism)
+     for (int i = 0; i < n_ids; i++) {
+         const int32_t e = ids[i];
+         if (e < 0 || e >= p->n_experts) continue;
+-        if (p->slot_of[e] >= 0) {                               // hit (slot_of = slot_ent address)
+-            p->hits++;
+-            struct slot_ent * s = (struct slot_ent *) p->slot_of[e];
+-            if (s->last_gen != g_gen) { lru_unlink(s); lru_front(s); s->last_gen = g_gen; }
++        int32_t si = p->slot[e];
++        if (si >= 0) {                                       // already in the pool
++            struct slot_ent * s = &g_slots[si];
++            if (s->state == ST_READY) {
++                p->hits++;
++                if (s->prefetched && !s->consumed) { g_pf_used++; s->consumed = 1; }
++                list_push_back(si);
++                pin_locked(si);
++                continue;
++            }
++            if (s->state == ST_LOADING) {                    // wait for the async prefetch (same as the engine)
++                pin_locked(si);
++                g_rescue++;
++                wait_ready_locked(si);
++                if (s->state == ST_FAILED && !load_sync_locked(si)) break;
++                p->hits++;
++                if (s->prefetched && !s->consumed) { g_pf_used++; s->consumed = 1; }
++                list_push_back(si);
++                continue;
++            }
++            // stale FAILED slot: reload synchronously
++            pin_locked(si);
++            g_rescue++;
++            if (!load_sync_locked(si)) break;
++            p->hits++; list_push_back(si);
+             continue;
+         }
++        // cold path: claim a slot + enqueue (or synchronous when the queue is full)
+         p->misses++;
+-        struct slot_ent * s = slot_alloc(p->slice, tix, demand);
+-        if (!s) break;
+-        const off_t off = (off_t) (p->file_off + (uint64_t) e * p->stride);
+-        const ssize_t got = pread(p->fd, s->base, p->slice, off);
+-        if (got != (ssize_t) p->slice) { s->expert = -1; break; }
+-        s->expert = e; s->last_gen = g_gen;
+-        lru_front(s);
+-        g_loads++; g_bytes_loaded += got;
+-        p->slot_of[e] = (int64_t) (intptr_t) s;
++        int32_t slot;
++        for (;;) {
++            slot = claim_slot_locked(tix, e);
++            if (slot >= 0) break;
++            if (!exclusive || g_inflight <= 0) break;
++            static int s_c40;
++            if (getenv("E0_POOL_STATS") && s_c40 < 40) { s_c40++; fprintf(stderr, "moe_pool corner: infl=%d used=%d qd=%d qp=%d\n", g_inflight, g_used_slots, g_qdn, g_qpn); }
++            pthread_cond_wait(&g_scv, &g_lock);
++        }
++        if (slot < 0) continue;
++        pin_locked(slot);
++        if (enqueue_io_locked(slot, 1)) {
++            if (nwait < 64) waiters[nwait++] = slot;
++            else wait_ready_locked(slot);                    // rare window overrun: wait in place
++        } else {
++            if (!load_sync_locked(slot)) break;
++        }
++    }
++    static int s_nw8;
++    if (getenv("E0_POOL_STATS") && s_nw8 < 4) { s_nw8++; fprintf(stderr, "moe_pool enq-done#%d: nwait=%d miss=%d qn=%d inf=%d\n", s_nw8, nwait, p->misses, g_qdn + g_qpn, g_inflight); }
++    for (int i = 0; i < nwait; i++) {
++        static int s_w400;
++        if (getenv("E0_POOL_STATS") && s_w400 < 400) { s_w400++; fprintf(stderr, "moe_pool w: slot=%d st=%d infl=%d qd=%d qp=%d\n", waiters[i], (int)g_slots[waiters[i]].state, g_inflight, g_qdn, g_qpn); }
++        wait_ready_locked(waiters[i]);
++        struct slot_ent * s = &g_slots[waiters[i]];
++        if (s->state == ST_FAILED && !load_sync_locked(waiters[i])) break;
+     }
+     pthread_mutex_unlock(&g_lock);
+ }
+ 
++// prefetch entry (called by pred_worker): enqueue and return immediately, never waits; dropped if there is no room (acquire backstops).
++void moe_pool_prefetch_impl(struct tensor_pool * p, int64_t e) {
++    if (!g_ready || e < 0 || e >= p->n_experts) return;
++    const int tix = (int) (p - g_tp);
++    pthread_mutex_lock(&g_lock);
++    const int32_t si = p->slot[e];
++    if (si >= 0) {
++        if (g_slots[si].state == ST_FAILED) {                // failed slot: retry once (enqueue)
++            g_slots[si].state = ST_LOADING;
++            if (!enqueue_io_locked(si, 0)) {
++                load_sync_locked(si);
++            }
++        }
++        g_pf_waste++;
++        pthread_mutex_unlock(&g_lock);
++        return;
++    }
++    if (g_qpn > 64) { g_pf_drop++; pthread_mutex_unlock(&g_lock); return; }   // throttle: drop prefetches when the main queue backs up
++    g_pf_issued++;
++    const int32_t slot = claim_slot_locked(tix, (int32_t) e);
++    if (slot < 0) { g_pf_drop++; pthread_mutex_unlock(&g_lock); return; }
++    g_slots[slot].prefetched = 1;
++    if (!enqueue_io_locked(slot, 0)) { unclaim_locked(slot); g_pf_drop++; }
++    pthread_mutex_unlock(&g_lock);
++}
++
+ const void * moe_pool_row_ptr(const struct ggml_tensor * t, int64_t expert, const void * def) {
+     if (!g_ready || expert < 0) return def;
+     const struct tensor_pool * p = find_tp(t);
+     if (!p || expert >= p->n_experts) return def;
+-    const int64_t si = p->slot_of[expert];
++    const int32_t si = p->slot[expert];
+     if (si < 0) return def;
+-    const struct slot_ent * s = (const struct slot_ent *) si;
+-    return (s->expert == expert) ? s->base : def;               // always true under the quiesce protocol; kept as a defensive check
++    const struct slot_ent * s = &g_slots[si];
++    return (s->state == ST_READY && s->expert == expert) ? s->mem : def;
+ }
+ 
+ void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
+                     long long * bytes_loaded, size_t * cap_bytes, size_t * used_bytes) {
+     long long h = 0, m = 0;
+     for (int i = 0; i < g_ntp; i++) { h += g_tp[i].hits; m += g_tp[i].misses; }
+-    if (hits) *hits = h; if (misses) *misses = m;
+-    if (loads) *loads = g_loads; if (evicts) *evicts = g_evicts;
+-    if (bytes_loaded) *bytes_loaded = g_bytes_loaded;
+-    if (cap_bytes) *cap_bytes = g_cap; if (used_bytes) *used_bytes = g_used;
++    if (hits) *hits = h;
++    if (misses) *misses = m;
++    if (loads) *loads = g_loads;
++    if (evicts) *evicts = g_evicts;
++    if (bytes_loaded) *bytes_loaded = g_flash;
++    if (cap_bytes) *cap_bytes = g_cap_bytes;
++    if (used_bytes) *used_bytes = (size_t) g_used_slots * g_slice;
+ }
+ 
+ __attribute__((destructor)) static void moe_pool_dump(void) {
+-    if (!g_ready || !getenv("E0_POOL_STATS")) {
+-        return;
+-    }
+-    long long h = 0, m = 0, l = 0, e = 0, bl = 0; size_t cap = 0, used = 0;
+-    moe_pool_stats(&h, &m, &l, &e, &bl, &cap, &used);
+-    fprintf(stderr, "moe_pool stats: hits=%lld misses=%lld loads=%lld evicts=%lld loaded_mb=%lld cap_mb=%zu used_mb=%zu\n",
+-            h, m, l, e, bl / (1024 * 1024), cap / (1024 * 1024), used / (1024 * 1024));
++    if (!g_ready && !g_blob_on) return;
++    long long h = 0, m = 0;
++    for (int i = 0; i < g_ntp; i++) { h += g_tp[i].hits; m += g_tp[i].misses; }
++    fprintf(stderr, "moe_pool stats: hits=%lld misses=%lld loads=%lld evicts=%lld trims=%lld stuck=%lld "
++                    "stall_ms=%lld pf=%lld/%lld/%lld/%lld/%lld flash_mb=%lld iofail=%lld "
++                    "cap_mb=%zu used_mb=%d blob_mb=%zu\n",
++            h, m, g_loads, g_evicts, g_trims, g_stuckq, g_stall_ns / 1000000,
++            g_pf_issued, g_pf_used, g_pf_waste, g_pf_drop, g_pf_miss,
++            g_flash / (1024 * 1024), g_iofail,
++            g_cap_bytes / (1024 * 1024), (int) ((size_t) g_used_slots * g_slice / (1024 * 1024)),
++            g_blob_bump / (1024 * 1024));
+ }
+diff --git a/ggml/src/ggml-cpu/moe_pool.h b/ggml/src/ggml-cpu/moe_pool.h
+index ab9f7a1..e4c48ca 100644
+--- a/ggml/src/ggml-cpu/moe_pool.h
++++ b/ggml/src/ggml-cpu/moe_pool.h
+@@ -1,7 +1,9 @@
+-// moe_pool.h -- expert demand pool (Android phase 1, no prediction; fork edge0/android-pool).
++// moe_pool.h -- expert demand pool v4 (aligned with the engine streaming/expert_pool contract; fork edge0/android-pool).
+ // Foundation = copy-in private frames (Android mlock RLIMIT=64KB is hard-wired; the pin-file-page design is unusable -- see PROBE B-track closure).
+-// hit = zero-fault direct read of private memory; miss = ensure() bulk preads the whole block ("wholesale" ~1GB/s vs page-fault "retail" 16MB/s);
+-// eviction = this pool's own LRU (whole-expert granularity, never half an expert); long-tail experts outside the pool fall back to their original mmap pages, kernel-backed (lazy pointer resolution).
++// v4 design = engine-isomorphic: equal-size slot state machine (EMPTY/LOADING/READY/FAILED) + IO-thread queue (prefetch enqueues and returns;
++//   acquire hitting LOADING waits on a cv) + pin protocol (unevictable while in compute) + resident blob (non-expert region copy-in) +
++//   watermark trim + stall/prefetch metrics. Hit = zero page faults; cold read = wholesale whole blocks (batched dispatch, awaited together).
++// env: E0_POOL_MB / E0_BLOB_MB / E0_IO_THREADS / E0_MEM_FLOOR_MB / E0_HEADS / E0_PREROUTER.
+ #ifndef MOE_POOL_H
+ #define MOE_POOL_H
+ 
+@@ -30,6 +32,16 @@ void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_id
+ // row-base resolution: resident -> private-frame pointer; otherwise -> default_ptr (original mmap address, kernel backstop)
+ const void * moe_pool_row_ptr(const struct ggml_tensor * t, int64_t expert, const void * default_ptr);
+ 
++// v4 resident blob (isomorphic to the engine's network.cpp img mirror): one copy-in per non-expert mmap-view tensor.
++// E0_BLOB_MB unset / budget exhausted -> blob_add returns NULL (caller stays as-is, zero disturbance).
++int    moe_pool_blob_enabled(void);
++void * moe_pool_blob_add(size_t nbytes);
++
++// phase 2: predictive prefetch (E0_HEADS weight files + E0_PREROUTER switch)
++int  moe_pool_load_heads(const char * path);
++void moe_pool_bind_min(int layer, const struct ggml_tensor * t);
++void moe_pool_note_routing(const struct ggml_tensor * src0, const int32_t * ids, int n_ids);
++
+ // stats snapshot (for debugging / accounting)
+ void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
+                     long long * bytes_loaded, size_t * cap_bytes, size_t * used_bytes);
+diff --git a/src/llama-model.cpp b/src/llama-model.cpp
+index ec1089d..b9e0907 100644
+--- a/src/llama-model.cpp
++++ b/src/llama-model.cpp
+@@ -7,6 +7,8 @@ extern "C" {
+ void        moe_pool_init(void);
+ int         moe_pool_enabled(void);
+ void        moe_pool_register(struct ggml_tensor * t, int fd, uint64_t file_off);
++void      * moe_pool_blob_add(size_t nbytes);
++int         moe_pool_blob_enabled(void);
+ }
+ #include "llama-model.h"
+ 
+@@ -1891,28 +1893,60 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
+         }
+     }
+ 
+-    // moe_pool registration (Android phase 1): expert 3D tensors -> pool. With E0_POOL_MB unset this block has no side effects.
++    // moe_pool: init stays here (heads/env parsing has zero dependencies); registration and blob live in moe_wire() --
++    // inside load_tensors the ml.contexts tensor set has data unbound (measured NULL); the primary set lives in tensors_by_name.
+     moe_pool_init();
+-    if (moe_pool_enabled() && ml.use_mmap && !ml.files.empty()) {
+-        auto reg = [&](ggml_tensor * t) {
+-            if (!t || !strstr(t->name, "exps")) {
+-                return;
+-            }
+-            const auto * w = ml.get_weight(t->name);
+-            if (!w || w->idx >= ml.files.size()) {
+-                return;
+-            }
+-            // pread does not depend on the file cursor, so the loader's fd is safe to use directly (no dup, avoids leaks)
+-            moe_pool_register(t, ml.files[w->idx]->file_id(), (uint64_t) w->offs);
+-        };
+-        for (auto & ctx : ml.contexts) {
+-            for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t; t = ggml_get_next_tensor(ctx.get(), t)) {
+-                reg(t);
++
++    return true;
++}
++
++// edge0 v4 (llama_model_load calls this after load_tensors): the tensors_by_name primary set now has its data pointers
++// bound. Expert tensors are registered into the pool; with E0_BLOB_MB on, each non-expert mmap-view tensor gets one copy-in (isomorphic to the engine's img mirror).
++void llama_model_base::moe_wire(llama_model_loader & ml) {
++    if (!ml.use_mmap || ml.files.empty()) {
++        return;
++    }
++    for (auto & [name, t] : tensors_by_name) {
++        if (!t || !t->data || !strstr(t->name, "exps")) {
++            continue;
++        }
++        const auto * w = ml.get_weight(name.c_str());
++        if (!w || w->idx >= ml.files.size()) {
++            continue;
++        }
++        // pread does not depend on the file cursor, so the loader's fd is safe to use directly (no dup, avoids leaks)
++        moe_pool_register(t, ml.files[w->idx]->file_id(), (uint64_t) w->offs);
++    }
++    if (!moe_pool_blob_enabled()) {
++        return;
++    }
++    size_t blob_copied = 0;
++    for (auto & [name, t] : tensors_by_name) {
++        if (!t || !t->data || strstr(t->name, "exps") || ggml_nbytes(t) == 0) {
++            continue;
++        }
++        bool in_mapping = false;
++        for (const auto & mapping : pimpl->mappings) {
++            const uint8_t * lo = (const uint8_t *) mapping->addr();
++            const uint8_t * hi = lo + mapping->size();
++            const uint8_t * p  = (const uint8_t *) t->data;
++            if (p >= lo && p < hi) {
++                in_mapping = true;
++                break;
+             }
+         }
++        if (!in_mapping) {
++            continue;   // heap tensors (already converted/copied) are resident anyway -- leave them alone
++        }
++        void * dst = moe_pool_blob_add(ggml_nbytes(t));
++        if (!dst) {
++            continue;   // budget exhausted: stay in mmap (visible as blob_mb in the dump)
++        }
++        memcpy(dst, t->data, ggml_nbytes(t));
++        t->data = dst;
++        blob_copied += ggml_nbytes(t);
+     }
+-
+-    return true;
++    fprintf(stderr, "moe_pool blob copy-in: %zu MiB\n", blob_copied / (1024 * 1024));
+ }
+ 
+ ggml_tensor * llama_model_base::create_tensor(llama_model_loader & ml, const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags) {
+diff --git a/src/llama-model.h b/src/llama-model.h
+index 984b2cf..47aaf83 100644
+--- a/src/llama-model.h
++++ b/src/llama-model.h
+@@ -751,6 +751,7 @@ struct llama_model {
+     bool has_tensor_overrides() const;
+ 
+     const struct ggml_tensor * get_tensor(const char * name) const;
++    virtual void moe_wire(llama_model_loader & ml) { (void) ml; }  // edge0 v4: pool+blob registration hook
+ 
+     float get_rope_freq_base (const llama_cparams & cparams, int il) const;
+     float get_rope_freq_scale(const llama_cparams & cparams, int il) const;
+@@ -821,6 +822,7 @@ struct llama_model_base : public llama_model {
+     void load_hparams(llama_model_loader & ml) override;
+     void load_vocab  (llama_model_loader & ml) override;
+     bool load_tensors(llama_model_loader & ml) override;
++    void moe_wire(llama_model_loader & ml) override;  // edge0 v4
+ 
+     // model must define these
+     void load_arch_hparams(llama_model_loader & ml) override = 0;
+diff --git a/src/llama.cpp b/src/llama.cpp
+index ad8e443..b09ffb7 100644
+--- a/src/llama.cpp
++++ b/src/llama.cpp
+@@ -369,6 +369,7 @@ static std::pair<int, llama_model *> llama_model_load(struct gguf_context * meta
+         if (!model->load_tensors(ml)) {
+             return {-2, nullptr};
+         }
++        model->moe_wire(ml);   // edge0 v4: expert pool + resident blob registration (see llama-model.cpp)
+ 
+         return {0, model_ptr.release()};
+     } catch (const std::exception & err) {
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0005-ggml-cpu-qwen3next-moe_pool-hooks-IQP-bypass-gate-en.patch b/patches/llama.cpp/android/0005-ggml-cpu-qwen3next-moe_pool-hooks-IQP-bypass-gate-en.patch
new file mode 100644
index 0000000..9d516fd
--- /dev/null
+++ b/patches/llama.cpp/android/0005-ggml-cpu-qwen3next-moe_pool-hooks-IQP-bypass-gate-en.patch
@@ -0,0 +1,59 @@
+From cda991a728b79384b05f1cb58c1ebca34eb31a49 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 00:36:57 +0800
+Subject: [PATCH 05/14] ggml-cpu/qwen3next: moe_pool hooks (IQP bypass gate,
+ ensure/note at MUL_MAT_ID, row-ptr lazy resolve, layer-input bind)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+- ggml-cpu.c: #include moe_pool.h; at the MUL_MAT_ID entry thread0 runs moe_pool_ensure +
+  moe_pool_note_routing before the barrier (prefetch signal; the pool reads only and never touches
+  logits); row bases resolve lazily via moe_pool_row_ptr; iqf bypass gate && !moe_pool_pooled(src0)
+  (pooled tensors are barred from the repack dual-track)
+- qwen3next.cpp: moe_pool_bind_min after both t_layer_inp sites (feature source m_in for the
+  predictor head)
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/ggml-cpu.c | 1 +
+ src/models/qwen3next.cpp     | 6 ++++++
+ 2 files changed, 7 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c
+index 54f0d9a..656f780 100644
+--- a/ggml/src/ggml-cpu/ggml-cpu.c
++++ b/ggml/src/ggml-cpu/ggml-cpu.c
+@@ -1676,6 +1676,7 @@ static void ggml_compute_forward_mul_mat_id(
+ 
+         // moe_pool: demand-pool dispatch -- bulk-read this ubatch's experts (unregistered tensors return immediately)
+         moe_pool_ensure(src0, (const int32_t *) ids->data, (int) (ids->ne[0] * ids->ne[1]));
++        moe_pool_note_routing(src0, (const int32_t *) ids->data, (int) (ids->ne[0] * ids->ne[1]));
+     }
+ 
+     // reset current_chunk
+diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp
+index 9673e3b..c081fd5 100644
+--- a/src/models/qwen3next.cpp
++++ b/src/models/qwen3next.cpp
+@@ -1,4 +1,9 @@
+ #include "models.h"
++
++// moe_pool in-place declarations (phase-2 Android: the builder binds m_in; implementation in ggml-cpu/moe_pool.c)
++extern "C" {
++void moe_pool_bind_min(int layer, const struct ggml_tensor * t);
++}
+ #include "llama-memory-recurrent.h"
+ 
+ void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) {
+@@ -187,6 +192,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
+         // only consumed when embeddings_layer_inp[il] is active (prerouter mode);
+         // with the feature off, set_outputs() skips this slot entirely.
+         res->t_layer_inp[il] = attn_post_norm;
++        moe_pool_bind_min(il, res->t_layer_inp[il]);
+ 
+         // FFN layer (MoE or dense) - without residual connection
+         cur = build_layer_ffn(attn_post_norm, il);
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0006-moe_pool-v4.2.1-FAILED-4-LOADING-evict-file_off-SYNC.patch b/patches/llama.cpp/android/0006-moe_pool-v4.2.1-FAILED-4-LOADING-evict-file_off-SYNC.patch
new file mode 100644
index 0000000..6c1144e
--- /dev/null
+++ b/patches/llama.cpp/android/0006-moe_pool-v4.2.1-FAILED-4-LOADING-evict-file_off-SYNC.patch
@@ -0,0 +1,91 @@
+From a07f9c8ad77f4ec1b169a540626cc172932071d3 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 01:18:49 +0800
+Subject: [PATCH 06/14] =?UTF-8?q?moe=5Fpool=20v4.2.1:=20FAILED=20=E9=87=8D?=
+ =?UTF-8?q?=E8=AF=95=E5=85=A5=E9=98=9F=E8=B7=AF=E5=BE=84=E5=85=88=E7=A6=BB?=
+ =?UTF-8?q?=E9=93=BE=EF=BC=88=E4=BA=8C=E5=B7=A1#4=EF=BC=9ALOADING=20?=
+ =?UTF-8?q?=E5=9C=A8=E9=93=BE=E7=A0=B4=E5=9D=8F=E4=B8=8D=E5=8F=98=E5=BC=8F?=
+ =?UTF-8?q?=E2=86=92evict=20=E5=8F=AF=E9=80=90=E5=9C=A8=E9=80=94=E6=A7=BD?=
+ =?UTF-8?q?=E5=8F=8C=E5=86=99=EF=BC=89+=20file=5Foff=20=E5=90=9E=E7=A0=81?=
+ =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20+=20SYNC=20=E4=BA=8C=E5=88=86=E5=BC=80?=
+ =?UTF-8?q?=E5=85=B3=20+=20=E6=8C=87=E7=BA=B9=E8=B0=83=E8=AF=95=E4=BB=B6?=
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 22 +++++++++++++++++++---
+ 1 file changed, 19 insertions(+), 3 deletions(-)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 6f76c90..fa89ffd 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -383,6 +383,7 @@ static void slot_finalize_locked(int32_t si, int ok) {
+ // synchronous read on the main thread outside the lock (same as the engine's sync_load_locked; no one can move a LOADING slot, so it's safe)
+ static int load_sync_locked(int32_t si) {
+     struct slot_ent * s = &g_slots[si];
++    list_unlink(si);        // unlink: during the retry window this slot cannot be taken over by a claim (out-of-lock read vs double write)
+     struct tensor_pool * p = &g_tp[s->tix];
+     const uint64_t off = p->file_off + (uint64_t) s->expert * p->stride;
+     pthread_mutex_unlock(&g_lock);
+@@ -443,7 +444,7 @@ static long long mem_avail_mb(void) {
+     fclose(f);
+     return v;
+ }
+-// watermark trim (simplified form of the engine's trim): evict only non-pinned READY; no physical frame reclaim (the arena middle is not contiguous -- see ALIGNMENT.md).
++// watermark trim (isomorphic to the engine's trim): evict non-pinned READY + madvise(DONTNEED) on evicted slots to truly drop RSS (the engine frees tail blocks via os_anon_free; we use a single arena, so per-slot DONTNEED is the equivalent).
+ static void water_trim_locked(void) {
+     const char * fe = getenv("E0_MEM_FLOOR_MB");
+     if (!fe) return;
+@@ -467,6 +468,7 @@ static void water_trim_locked(void) {
+         list_unlink(v);
+         s->state = ST_EMPTY; s->expert = -1;
+         g_used_slots--; g_trims++; n++;
++        madvise(s->mem, g_slice, MADV_DONTNEED);   // the slot region is a contiguous 640KB: physical pages released at once and read back as zeros (rewritten by pread on reuse)
+     }
+     if (n) fprintf(stderr, "moe_pool trim: evicted=%d used=%d MemAvailable=%lldMB\n", n, g_used_slots, avail);
+ }
+@@ -514,7 +516,9 @@ void moe_pool_register(struct ggml_tensor * t, int fd, uint64_t file_off) {
+     if (p->role >= 0 && p->layer >= 0 && p->layer < 64) {
+         g_lr[p->layer].tix[p->role] = g_ntp - 1;
+     }
+-    p->key = t->data; snprintf(p->name, sizeof(p->name), "%s", t->name); p->fd = dup(fd);   // independent ownership: pread still works after the loader closed its fd (v4 hang incident) p->file_off = file_off; p->slice = stride; p->stride = stride;
++    p->key = t->data; snprintf(p->name, sizeof(p->name), "%s", t->name);
++    p->fd = dup(fd);   // independent ownership: pread still works after the loader closed its fd (v4 hang incident)
++    p->file_off = file_off; p->slice = stride; p->stride = stride;
+     p->n_experts = t->ne[2];
+     p->slot = calloc((size_t) p->n_experts, sizeof(int32_t));
+     for (int64_t i = 0; i < p->n_experts; i++) p->slot[i] = SLOT_NONE;
+@@ -633,6 +637,7 @@ void moe_pool_prefetch_impl(struct tensor_pool * p, int64_t e) {
+     const int32_t si = p->slot[e];
+     if (si >= 0) {
+         if (g_slots[si].state == ST_FAILED) {                // failed slot: retry once (enqueue)
++            list_unlink(si);                                 // second-audit case #4: LOADING must be unlinked (otherwise eviction could take an in-flight slot -> double write)
+             g_slots[si].state = ST_LOADING;
+             if (!enqueue_io_locked(si, 0)) {
+                 load_sync_locked(si);
+@@ -658,7 +663,18 @@ const void * moe_pool_row_ptr(const struct ggml_tensor * t, int64_t expert, cons
+     const int32_t si = p->slot[expert];
+     if (si < 0) return def;
+     const struct slot_ent * s = &g_slots[si];
+-    return (s->state == ST_READY && s->expert == expert) ? s->mem : def;
++    if (s->state == ST_READY && s->expert == expert) {
++        static int s_fp6;
++        if (getenv("E0_POOL_STATS") && s_fp6 < 6) {
++            s_fp6++;
++            const unsigned char * q = (const unsigned char *) s->mem;
++            const unsigned char * d2 = (const unsigned char *) def;
++            fprintf(stderr, "moe_pool fp: %s e=%d slot=%02x%02x%02x%02x def=%02x%02x%02x%02x\n",
++                    t->name, (int) expert, q[0], q[1], q[2], q[3], d2[0], d2[1], d2[2], d2[3]); fprintf(stderr, "moe_pool fp2: tix=%d file_off=%llu stride=%llu fd=%d\n", (int)(p - g_tp), (unsigned long long) p->file_off, (unsigned long long) p->stride, p->fd);
++        }
++        return s->mem;
++    }
++    return def;
+ }
+ 
+ void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0007-moe_pool-v4.3-NEON-vmovl-vcvt-vfma-8-lane.patch b/patches/llama.cpp/android/0007-moe_pool-v4.3-NEON-vmovl-vcvt-vfma-8-lane.patch
new file mode 100644
index 0000000..c700ebe
--- /dev/null
+++ b/patches/llama.cpp/android/0007-moe_pool-v4.3-NEON-vmovl-vcvt-vfma-8-lane.patch
@@ -0,0 +1,81 @@
+From 16a0020f3f5a2430d2156a7462e75887d7ef7bfc Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 01:25:31 +0800
+Subject: [PATCH 07/14] moe_pool v4.3: NEON-ify the predictor-head math
+ (vmovl+vcvt+vfma 8-lane, scalar tail fallback)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Head math fc1(512x2560)+fc2/lin(256x2560/512) ~= 2.1MMAC per layer: the scalar version costs the
+worker 40-80ms per token, directly squeezing the prefetch lead time. The NEON version gives ~4-8x
+speedup = the prefetch window recovers.
+(Verified on a real arm64 device, on the record: the dualpath discipline does not apply to this
+piece -- it is pure added math, and routing/numerics are already covered by the text-gate three-way
+consistency guardrail; the NEON change only affects how fast prefetch targets are selected, never
+logits.)
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 36 ++++++++++++++++++++++++++----------
+ 1 file changed, 26 insertions(+), 10 deletions(-)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index fa89ffd..903f0f7 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -233,6 +233,29 @@ static inline float f16tof(uint16_t h) {
+     const uint32_t w = (uint32_t) h << 16;
+     float f; memcpy(&f, &w, 4); return f;
+ }
++#if defined(__aarch64__)
++#include <arm_neon.h>
++static float dot_f16_f32(const uint16_t * w, const float * x, int n) {
++    float32x4_t acc0 = vdupq_n_f32(0.f), acc1 = vdupq_n_f32(0.f);
++    int i = 0;
++    for (; i + 8 <= n; i += 8) {
++        uint16x8_t wv = vld1q_u16(w + i);
++        float32x4_t w0 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(wv)));
++        float32x4_t w1 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(wv)));
++        acc0 = vfmaq_f32(acc0, w0, vld1q_f32(x + i));
++        acc1 = vfmaq_f32(acc1, w1, vld1q_f32(x + i + 4));
++    }
++    float r = vaddvq_f32(vaddq_f32(acc0, acc1));
++    for (; i < n; i++) r += f16tof(w[i]) * x[i];
++    return r;
++}
++#else
++static float dot_f16_f32(const uint16_t * w, const float * x, int n) {
++    float r = 0;
++    for (int i = 0; i < n; i++) r += f16tof(w[i]) * x[i];
++    return r;
++}
++#endif
+ static inline float gelu_erf_f(float x) { return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f)); }
+ 
+ static void pred_do_job(struct pred_job * j) {
+@@ -241,18 +264,11 @@ static void pred_do_job(struct pred_job * j) {
+     const uint16_t * fc2 = g_lr[j->layer].hd.fc2;
+     const uint16_t * lin = g_lr[j->layer].hd.lin;
+     for (int o = 0; o < 512; o++) {
+-        const uint16_t * wrow = fc1 + (size_t) o * 2560;
+-        float acc = 0;
+-        for (int i = 0; i < 2560; i++) acc += f16tof(wrow[i]) * j->x[i];
+-        h[o] = gelu_erf_f(acc);
++        h[o] = gelu_erf_f(dot_f16_f32(fc1 + (size_t) o * 2560, j->x, 2560));
+     }
+     for (int o = 0; o < 256; o++) {
+-        const uint16_t * wr = lin + (size_t) o * 2560;
+-        const uint16_t * wc = fc2 + (size_t) o * 512;
+-        float acc = 0;
+-        for (int i = 0; i < 2560; i++) acc += f16tof(wr[i]) * j->x[i];
+-        for (int i = 0; i < 512; i++) acc += f16tof(wc[i]) * h[i];
+-        out[o] = acc;
++        out[o] = dot_f16_f32(lin + (size_t) o * 2560, j->x, 2560)
++               + dot_f16_f32(fc2 + (size_t) o * 512, h, 512);
+     }
+     int top[8]; float tv[8]; for (int k = 0; k < 8; k++) { tv[k] = -1e30f; top[k] = -1; }
+     for (int e = 0; e < 256; e++) {
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0008-moe_pool-v4.4-2s-ensure-proc-12ms-token.patch b/patches/llama.cpp/android/0008-moe_pool-v4.4-2s-ensure-proc-12ms-token.patch
new file mode 100644
index 0000000..5698896
--- /dev/null
+++ b/patches/llama.cpp/android/0008-moe_pool-v4.4-2s-ensure-proc-12ms-token.patch
@@ -0,0 +1,34 @@
+From 03c5ca68489cede393308a1cd12e89e68aa476f3 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 01:44:29 +0800
+Subject: [PATCH 08/14] =?UTF-8?q?moe=5Fpool=20v4.4:=20=E6=B0=B4=E4=BD=8D?=
+ =?UTF-8?q?=E6=8E=A2=E6=B5=8B=202s=20=E8=8A=82=E6=B5=81=EF=BC=88=E6=AF=8F?=
+ =?UTF-8?q?=20ensure=20=E5=BC=80=20/proc=E2=89=8812ms/token=20=E6=B5=8B?=
+ =?UTF-8?q?=E9=87=8F=E7=A8=8E=EF=BC=89?=
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 4 ++++
+ 1 file changed, 4 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 903f0f7..4d62904 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -466,6 +466,10 @@ static void water_trim_locked(void) {
+     if (!fe) return;
+     const long long floor_mb = atoll(fe);
+     if (floor_mb <= 0) return;
++    static uint64_t s_next = 0;                        // 2s throttle (third-audit case: opening /proc on every ensure cost a ~3% measurement tax)
++    const uint64_t now = now_ns();
++    if (now < s_next) return;
++    s_next = now + 2000000000ull;
+     const long long avail = mem_avail_mb();
+     if (avail < 0 || avail >= floor_mb) return;
+     const int target = (int) (g_nslots * 6 / 10);
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0009-moe_pool-v4.5-K8-K4-64-32-p6144pr-3.10-UFS-K8-1.23GB.patch b/patches/llama.cpp/android/0009-moe_pool-v4.5-K8-K4-64-32-p6144pr-3.10-UFS-K8-1.23GB.patch
new file mode 100644
index 0000000..375be39
--- /dev/null
+++ b/patches/llama.cpp/android/0009-moe_pool-v4.5-K8-K4-64-32-p6144pr-3.10-UFS-K8-1.23GB.patch
@@ -0,0 +1,55 @@
+From 437ebfe1a299b86f7e299dd0db93f79961c0a68a Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 02:23:09 +0800
+Subject: [PATCH 09/14] =?UTF-8?q?moe=5Fpool=20v4.5:=20=E9=A2=84=E5=8F=96?=
+ =?UTF-8?q?=20K8=E2=86=92K4=20+=20=E7=A7=AF=E5=8E=8B=E8=8A=82=E6=B5=81=206?=
+ =?UTF-8?q?4=E2=86=9232=EF=BC=88p6144pr=3D3.10=20=E5=8F=8D=E5=99=AC?=
+ =?UTF-8?q?=E6=A1=88=E2=80=94=E2=80=94=E9=A2=84=E5=8F=96=E9=A2=84=E7=AE=97?=
+ =?UTF-8?q?=3D=E6=9C=AC=E7=AB=AF=20UFS=20=E5=B8=A6=E5=AE=BD=EF=BC=8CK8=3D1?=
+ =?UTF-8?q?.23GB/tok=20=E8=B6=85=E8=BD=BD=202=20=E5=80=8D=EF=BC=89?=
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 10 +++++-----
+ 1 file changed, 5 insertions(+), 5 deletions(-)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 4d62904..7584680 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -270,16 +270,16 @@ static void pred_do_job(struct pred_job * j) {
+         out[o] = dot_f16_f32(lin + (size_t) o * 2560, j->x, 2560)
+                + dot_f16_f32(fc2 + (size_t) o * 512, h, 512);
+     }
+-    int top[8]; float tv[8]; for (int k = 0; k < 8; k++) { tv[k] = -1e30f; top[k] = -1; }
++    int top[4]; float tv[4]; for (int k = 0; k < 4; k++) { tv[k] = -1e30f; top[k] = -1; }
+     for (int e = 0; e < 256; e++) {
+-        for (int k = 0; k < 8; k++) {
++        for (int k = 0; k < 4; k++) {
+             if (out[e] > tv[k]) {
+-                for (int m = 7; m > k; m--) { tv[m] = tv[m-1]; top[m] = top[m-1]; }
++                for (int m = 3; m > k; m--) { tv[m] = tv[m-1]; top[m] = top[m-1]; }
+                 tv[k] = out[e]; top[k] = e; break;
+             }
+         }
+     }
+-    for (int k = 0; k < 8; k++) {
++    for (int k = 0; k < 4; k++) {
+         if (top[k] < 0) break;
+         for (int role = 0; role < 3; role++) {
+             const int64_t tix = g_lr[j->layer].tix[role];
+@@ -667,7 +667,7 @@ void moe_pool_prefetch_impl(struct tensor_pool * p, int64_t e) {
+         pthread_mutex_unlock(&g_lock);
+         return;
+     }
+-    if (g_qpn > 64) { g_pf_drop++; pthread_mutex_unlock(&g_lock); return; }   // throttle: drop prefetches when the main queue backs up
++    if (g_qpn > 32) { g_pf_drop++; pthread_mutex_unlock(&g_lock); return; }   // throttle v4.5 (K8 backlash case: 64 still exceeded the device bandwidth)
+     g_pf_issued++;
+     const int32_t slot = claim_slot_locked(tix, (int32_t) e);
+     if (slot < 0) { g_pf_drop++; pthread_mutex_unlock(&g_lock); return; }
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0010-moe_pool-v4.6-moe_pool_extra-stall-App-metrics-pool-.patch b/patches/llama.cpp/android/0010-moe_pool-v4.6-moe_pool_extra-stall-App-metrics-pool-.patch
new file mode 100644
index 0000000..f00c737
--- /dev/null
+++ b/patches/llama.cpp/android/0010-moe_pool-v4.6-moe_pool_extra-stall-App-metrics-pool-.patch
@@ -0,0 +1,56 @@
+From 961da91d656cecdeede567d05545e69352aee3f2 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 09:05:06 +0800
+Subject: [PATCH 10/14] =?UTF-8?q?moe=5Fpool=20v4.6:=20moe=5Fpool=5Fextra?=
+ =?UTF-8?q?=20=E5=AF=BC=E5=87=BA=EF=BC=88stall/=E6=A7=BD=E4=BD=8D/?=
+ =?UTF-8?q?=E9=A2=84=E5=8F=96=E6=B6=88=E8=B4=B9=E8=AE=A1=E6=95=B0=E2=80=94?=
+ =?UTF-8?q?=E2=80=94App=20metrics=20=E5=8D=A1=20pool/rss=20=E4=BE=9B?=
+ =?UTF-8?q?=E6=95=B0=E9=80=9A=E9=81=93=EF=BC=89?=
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 9 +++++++++
+ ggml/src/ggml-cpu/moe_pool.h | 4 ++++
+ 2 files changed, 13 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 7584680..d7be8ca 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -710,6 +710,15 @@ void moe_pool_stats(long long * hits, long long * misses, long long * loads, lon
+     if (used_bytes) *used_bytes = (size_t) g_used_slots * g_slice;
+ }
+ 
++void moe_pool_extra(long long * stall_ns, long long * slots_used, long long * slots_total,
++                    long long * pf_used, long long * pf_mispredict) {
++    if (stall_ns) *stall_ns = g_stall_ns;
++    if (slots_used) *slots_used = g_used_slots;
++    if (slots_total) *slots_total = g_nslots;
++    if (pf_used) *pf_used = g_pf_used;
++    if (pf_mispredict) *pf_mispredict = g_pf_miss;
++}
++
+ __attribute__((destructor)) static void moe_pool_dump(void) {
+     if (!g_ready && !g_blob_on) return;
+     long long h = 0, m = 0;
+diff --git a/ggml/src/ggml-cpu/moe_pool.h b/ggml/src/ggml-cpu/moe_pool.h
+index e4c48ca..533d6ac 100644
+--- a/ggml/src/ggml-cpu/moe_pool.h
++++ b/ggml/src/ggml-cpu/moe_pool.h
+@@ -42,6 +42,10 @@ int  moe_pool_load_heads(const char * path);
+ void moe_pool_bind_min(int layer, const struct ggml_tensor * t);
+ void moe_pool_note_routing(const struct ggml_tensor * src0, const int32_t * ids, int n_ids);
+ 
++// v4 extended counters (App metrics channel: stall / slots / prefetch consumption)
++void moe_pool_extra(long long * stall_ns, long long * slots_used, long long * slots_total,
++                    long long * pf_used, long long * pf_mispredict);
++
+ // stats snapshot (for debugging / accounting)
+ void moe_pool_stats(long long * hits, long long * misses, long long * loads, long long * evicts,
+                     long long * bytes_loaded, size_t * cap_bytes, size_t * used_bytes);
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0011-moe_pool-v4.7-moe_pool_reset_all-inflight-blob-IO-9-.patch b/patches/llama.cpp/android/0011-moe_pool-v4.7-moe_pool_reset_all-inflight-blob-IO-9-.patch
new file mode 100644
index 0000000..d0804dd
--- /dev/null
+++ b/patches/llama.cpp/android/0011-moe_pool-v4.7-moe_pool_reset_all-inflight-blob-IO-9-.patch
@@ -0,0 +1,76 @@
+From b946e2ad525538a04c784bbd730e890004be9997 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 09:22:43 +0800
+Subject: [PATCH 11/14] =?UTF-8?q?moe=5Fpool=20v4.7:=20moe=5Fpool=5Freset?=
+ =?UTF-8?q?=5Fall=EF=BC=88=E8=B7=A8=E6=A8=A1=E5=9E=8B=E5=88=87=E6=8D=A2?=
+ =?UTF-8?q?=E5=85=A8=E5=A4=8D=E4=BD=8D=E2=80=94=E2=80=94=E7=AD=89=20inflig?=
+ =?UTF-8?q?ht=20=E6=8E=92=E7=A9=BA=E3=80=81=E6=A7=BD=E8=A1=A8/=E6=B3=A8?=
+ =?UTF-8?q?=E5=86=8C/blob=20=E5=BD=92=E9=9B=B6=E3=80=81IO=20=E7=BA=BF?=
+ =?UTF-8?q?=E7=A8=8B=E7=95=99=E5=9C=BA=E5=A4=8D=E7=94=A8=EF=BC=9B9/29=208b?=
+ =?UTF-8?q?->35b=20=E9=97=AA=E9=80=80=E6=A1=88=E7=9A=84=E5=90=8C=E6=97=8F?=
+ =?UTF-8?q?=E5=9C=B0=E9=9B=B7=E6=8E=92=E9=9B=B7=EF=BC=89?=
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 27 +++++++++++++++++++++++++++
+ ggml/src/ggml-cpu/moe_pool.h |  3 +++
+ 2 files changed, 30 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index d7be8ca..e9d3415 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -719,6 +719,33 @@ void moe_pool_extra(long long * stall_ns, long long * slots_used, long long * sl
+     if (pf_mispredict) *pf_mispredict = g_pf_miss;
+ }
+ 
++// cross-model-switch reset (insurance against the 9/29 crash-family incident): called by nativeFree. After waiting for the IO
++// drain, all state is zeroed; the next moe_pool_init/register rebuilds under the new env. Head mappings are kept (those files don't change with the model).
++void moe_pool_reset_all(void) {
++    pthread_mutex_lock(&g_lock);
++    while (g_inflight > 0) pthread_cond_wait(&g_scv, &g_lock);       // wait for in-flight reads to terminate
++    for (int i = 0; i < g_ntp; i++) { free(g_tp[i].slot); g_tp[i].slot = NULL; }
++    g_ntp = 0;
++    for (int i = 0; i < 64; i++) { g_lr[i].tix[0] = g_lr[i].tix[1] = g_lr[i].tix[2] = -1; }
++    if (g_slots) {
++        for (int i = 0; i < g_nslots; i++) {
++            struct slot_ent * sl = &g_slots[i];
++            if (sl->on_list) list_unlink(i);
++            sl->state = ST_EMPTY; sl->expert = -1; sl->tix = -1; sl->pinned = 0;
++            sl->prefetched = 0; sl->consumed = 0;
++        }
++        free(g_slots); g_slots = NULL;
++    }
++    if (g_arena) { munmap(g_arena, g_slice * (size_t) g_nslots); g_arena = NULL; }
++    g_nslots = 0; g_used_slots = 0; g_empty_cursor = 0; g_slice = 0;
++    g_lru_head = g_lru_tail = SLOT_NONE;
++    g_qdn = g_qpn = 0; g_qph = g_qpt = g_qdh = g_qdt = 0;
++    if (g_blob) { madvise(g_blob, g_blob_cap, MADV_DONTNEED); g_blob_bump = 0; }
++    g_ready = 0; g_cap_bytes = 0;
++    g_npins = 0;
++    pthread_mutex_unlock(&g_lock);
++}
++
+ __attribute__((destructor)) static void moe_pool_dump(void) {
+     if (!g_ready && !g_blob_on) return;
+     long long h = 0, m = 0;
+diff --git a/ggml/src/ggml-cpu/moe_pool.h b/ggml/src/ggml-cpu/moe_pool.h
+index 533d6ac..82dfc74 100644
+--- a/ggml/src/ggml-cpu/moe_pool.h
++++ b/ggml/src/ggml-cpu/moe_pool.h
+@@ -42,6 +42,9 @@ int  moe_pool_load_heads(const char * path);
+ void moe_pool_bind_min(int layer, const struct ggml_tensor * t);
+ void moe_pool_note_routing(const struct ggml_tensor * src0, const int32_t * ids, int n_ids);
+ 
++// cross-model-switch reset (called by the llama shell's nativeFree)
++void moe_pool_reset_all(void);
++
+ // v4 extended counters (App metrics channel: stall / slots / prefetch consumption)
+ void moe_pool_extra(long long * stall_ns, long long * slots_used, long long * slots_total,
+                     long long * pf_used, long long * pf_mispredict);
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0012-moe_pool-v4.7-keepwarm-demand-ensure-ring-8192-turn-.patch b/patches/llama.cpp/android/0012-moe_pool-v4.7-keepwarm-demand-ensure-ring-8192-turn-.patch
new file mode 100644
index 0000000..1dcd9c2
--- /dev/null
+++ b/patches/llama.cpp/android/0012-moe_pool-v4.7-keepwarm-demand-ensure-ring-8192-turn-.patch
@@ -0,0 +1,129 @@
+From 24f7d4c33b0fc3e176da129a986c4d2d7a7e3596 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 10:37:19 +0800
+Subject: [PATCH 12/14] =?UTF-8?q?moe=5Fpool=20v4.7:=20keepwarm=20=E8=BD=AE?=
+ =?UTF-8?q?=E6=9C=AB=E5=9B=9E=E7=81=8C=EF=BC=88demand=20ensure=20=E5=8F=8C?=
+ =?UTF-8?q?=E6=94=AF=E8=AE=B0=20ring=208192=EF=BC=8Cturn=20=E6=9C=AB=20det?=
+ =?UTF-8?q?ached=20=E7=BA=BF=E7=A8=8B=E6=8C=89=20E0=5FKEEPWARM=5FMB=20?=
+ =?UTF-8?q?=E9=A2=84=E7=AE=97=E6=8A=8A=E8=A2=AB=E9=80=90=E8=80=85=E4=BB=A5?=
+ =?UTF-8?q?=E9=A2=84=E5=8F=96=E7=BA=A7=E6=90=AC=E5=9B=9E=EF=BC=9BApp=2035b?=
+ =?UTF-8?q?=20=E6=A1=A3=3D512=EF=BC=89+=20pred=20feats=20f16=20=E9=87=8F?=
+ =?UTF-8?q?=E5=8C=96=EF=BC=88SPEC=20=C2=A71=20=E8=AE=AD=E7=BB=83=E5=88=86?=
+ =?UTF-8?q?=E5=B8=83=E5=AF=B9=E9=BD=90=EF=BC=8Cpf=20=E5=91=BD=E4=B8=AD?=
+ =?UTF-8?q?=E7=8E=87=E5=B7=AE=E5=80=BC=E5=80=99=E9=80=89=E2=91=A0=EF=BC=9B?=
+ =?UTF-8?q?pred=20OFF=20=E6=97=B6=E9=9B=B6=E8=A7=A6=E8=BE=BE=EF=BC=89?=
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 47 ++++++++++++++++++++++++++++++++++++
+ ggml/src/ggml-cpu/moe_pool.h |  3 +++
+ 2 files changed, 50 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index e9d3415..2098376 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -91,6 +91,12 @@ static long long g_pf_issued, g_pf_used, g_pf_waste, g_pf_drop, g_pf_miss, g_res
+ static char *   g_blob; static size_t g_blob_cap, g_blob_bump; static int g_blob_on;
+ static volatile int g_heads_on, g_pred_enabled;   // referenced early by init (defined in the pred section)
+ 
++// keepwarm (9/29 20s-between-turns case): accounts the (tix,e) touched by the current decode round; at turn end, evicted entries are refilled asynchronously when idle.
++#define KW_DEPTH 8192
++static struct { int32_t tix; int32_t e; } g_kw[KW_DEPTH];
++static unsigned g_kwh, g_kwt;
++static volatile int g_kw_busy;
++
+ static uint64_t now_ns(void) {
+     struct timespec ts; clock_gettime(CLOCK_MONOTONIC, &ts);
+     return (uint64_t) ts.tv_sec * 1000000000ull + (uint64_t) ts.tv_nsec;
+@@ -258,8 +264,15 @@ static float dot_f16_f32(const uint16_t * w, const float * x, int n) {
+ #endif
+ static inline float gelu_erf_f(float x) { return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f)); }
+ 
++static void moe_pool_f16_quant_feats(float * x, int n) {
++    for (int i = 0; i < n; i++) x[i] = (float) (__fp16) x[i];
++}
++
+ static void pred_do_job(struct pred_job * j) {
+     float h[512], out[256];
++    // SPEC sec.1: feats are cast to f16 before entering the network (training distribution = f16; feeding f32 directly would
++    // drift the top-K ranking on near-ties -- prime suspect #1 for the gap between the 10% pf hit rate and the 55% recorded profile).
++    { extern void moe_pool_f16_quant_feats(float * x, int n); moe_pool_f16_quant_feats(j->x, 2560); }
+     const uint16_t * fc1 = g_lr[j->layer].hd.fc1;
+     const uint16_t * fc2 = g_lr[j->layer].hd.fc2;
+     const uint16_t * lin = g_lr[j->layer].hd.lin;
+@@ -595,6 +608,7 @@ void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_id
+             struct slot_ent * s = &g_slots[si];
+             if (s->state == ST_READY) {
+                 p->hits++;
++                { unsigned h = g_kwh % KW_DEPTH; g_kw[h].tix = tix; g_kw[h].e = e; g_kwh++; }
+                 if (s->prefetched && !s->consumed) { g_pf_used++; s->consumed = 1; }
+                 list_push_back(si);
+                 pin_locked(si);
+@@ -629,6 +643,7 @@ void moe_pool_ensure(const struct ggml_tensor * t, const int32_t * ids, int n_id
+             pthread_cond_wait(&g_scv, &g_lock);
+         }
+         if (slot < 0) continue;
++        { unsigned h = g_kwh % KW_DEPTH; g_kw[h].tix = tix; g_kw[h].e = e; g_kwh++; }
+         pin_locked(slot);
+         if (enqueue_io_locked(slot, 1)) {
+             if (nwait < 64) waiters[nwait++] = slot;
+@@ -676,6 +691,38 @@ void moe_pool_prefetch_impl(struct tensor_pool * p, int64_t e) {
+     pthread_mutex_unlock(&g_lock);
+ }
+ 
++// turn-end refill: move (tix,e) pairs in the ring that are currently not in the pool back at prefetch level
++// (demand=0 queue, evictable). budget_mb=0 -> read E0_KEEPWARM_MB (default 256MB, low intensity, never competes with the foreground).
++void moe_pool_keepwarm(int budget_mb) {
++    if (!g_ready) return;
++    if (budget_mb <= 0) {
++        const char * e = getenv("E0_KEEPWARM_MB");
++        budget_mb = e ? (int) atol(e) : 0;
++        if (budget_mb <= 0) budget_mb = 256;
++    }
++    if (g_kw_busy) return;
++    g_kw_busy = 1;
++    long long bytes = (long long) budget_mb * 1024 * 1024;
++    pthread_mutex_lock(&g_lock);
++    unsigned seen_end = g_kwh;
++    for (unsigned i = g_kwt; i < seen_end && bytes > 0; i++) {
++        int k = i % KW_DEPTH;
++        const int tix = g_kw[k].tix; const int32_t e2 = g_kw[k].e;
++        if (tix < 0 || tix >= g_ntp) continue;
++        struct tensor_pool * p = &g_tp[tix];
++        if (e2 < 0 || e2 >= p->n_experts) continue;
++        if (p->slot[e2] >= 0) continue;                     // already resident
++        const int32_t slot = claim_slot_locked(tix, e2);
++        if (slot < 0) break;
++        g_slots[slot].prefetched = 1;
++        if (!enqueue_io_locked(slot, 0)) { unclaim_locked(slot); break; }
++        bytes -= (long long) p->slice;
++    }
++    g_kwt = seen_end;                                       // consumed this round (whether or not fully refilled)
++    pthread_mutex_unlock(&g_lock);
++    g_kw_busy = 0;
++}
++
+ const void * moe_pool_row_ptr(const struct ggml_tensor * t, int64_t expert, const void * def) {
+     if (!g_ready || expert < 0) return def;
+     const struct tensor_pool * p = find_tp(t);
+diff --git a/ggml/src/ggml-cpu/moe_pool.h b/ggml/src/ggml-cpu/moe_pool.h
+index 82dfc74..3341e1d 100644
+--- a/ggml/src/ggml-cpu/moe_pool.h
++++ b/ggml/src/ggml-cpu/moe_pool.h
+@@ -42,6 +42,9 @@ int  moe_pool_load_heads(const char * path);
+ void moe_pool_bind_min(int layer, const struct ggml_tensor * t);
+ void moe_pool_note_routing(const struct ggml_tensor * src0, const int32_t * ids, int n_ids);
+ 
++// turn-end keepwarm refill (E0_KEEPWARM_MB budget, 0 = default 256MB)
++void moe_pool_keepwarm(int budget_mb);
++
+ // cross-model-switch reset (called by the llama shell's nativeFree)
+ void moe_pool_reset_all(void);
+ 
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0013-moe_pool-v4.7.1-reset_all-keepwarm-g_kw_busy-2s-g_re.patch b/patches/llama.cpp/android/0013-moe_pool-v4.7.1-reset_all-keepwarm-g_kw_busy-2s-g_re.patch
new file mode 100644
index 0000000..536d02a
--- /dev/null
+++ b/patches/llama.cpp/android/0013-moe_pool-v4.7.1-reset_all-keepwarm-g_kw_busy-2s-g_re.patch
@@ -0,0 +1,42 @@
+From ee1237ee6ce144d24097406654a64d14cad8ce82 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 17:45:30 +0800
+Subject: [PATCH 13/14] =?UTF-8?q?moe=5Fpool=20v4.7.1:=20reset=5Fall=C3=97k?=
+ =?UTF-8?q?eepwarm=20=E7=AB=9E=E4=BA=89=E5=8A=A0=E5=9B=BA=E2=80=94?=
+ =?UTF-8?q?=E2=80=94=E5=85=A5=E5=8F=A3=E7=AD=89=20g=5Fkw=5Fbusy=20?=
+ =?UTF-8?q?=E6=B8=85=EF=BC=88=E2=89=A42s=EF=BC=89+=20=E5=9B=9E=E7=81=8C?=
+ =?UTF-8?q?=E5=BE=AA=E7=8E=AF=E9=80=90=E6=A7=BD=E9=87=8D=E6=A3=80=20g=5Fre?=
+ =?UTF-8?q?ady=EF=BC=88unload=20=E4=B8=AD=E9=80=94=E7=AB=8B=E5=81=9C?=
+ =?UTF-8?q?=EF=BC=8C=E9=98=B2=20freed=20arena=20UAF=EF=BC=89?=
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 2 ++
+ 1 file changed, 2 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 2098376..9bd55dc 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -706,6 +706,7 @@ void moe_pool_keepwarm(int budget_mb) {
+     pthread_mutex_lock(&g_lock);
+     unsigned seen_end = g_kwh;
+     for (unsigned i = g_kwt; i < seen_end && bytes > 0; i++) {
++        if (!g_ready) break;                                   // stop immediately if reset came in midway
+         int k = i % KW_DEPTH;
+         const int tix = g_kw[k].tix; const int32_t e2 = g_kw[k].e;
+         if (tix < 0 || tix >= g_ntp) continue;
+@@ -769,6 +770,7 @@ void moe_pool_extra(long long * stall_ns, long long * slots_used, long long * sl
+ // cross-model-switch reset (insurance against the 9/29 crash-family incident): called by nativeFree. After waiting for the IO
+ // drain, all state is zeroed; the next moe_pool_init/register rebuilds under the new env. Head mappings are kept (those files don't change with the model).
+ void moe_pool_reset_all(void) {
++    for (int w = 0; w < 200 && g_kw_busy; w++) usleep(10000);   // keepwarm in-flight hard gate (prevents the UAF race)
+     pthread_mutex_lock(&g_lock);
+     while (g_inflight > 0) pthread_cond_wait(&g_scv, &g_lock);       // wait for in-flight reads to terminate
+     for (int i = 0; i < g_ntp; i++) { free(g_tp[i].slot); g_tp[i].slot = NULL; }
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/android/0014-edge0-moe_pool-7-mmid-base-resolver.patch b/patches/llama.cpp/android/0014-edge0-moe_pool-7-mmid-base-resolver.patch
new file mode 100644
index 0000000..5dd0708
--- /dev/null
+++ b/patches/llama.cpp/android/0014-edge0-moe_pool-7-mmid-base-resolver.patch
@@ -0,0 +1,56 @@
+From 0456216585cebdf82191547240e9ddc66c7284a4 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Tue, 29 Sep 2026 20:25:27 +0800
+Subject: [PATCH 14/14] edge0 consolidation: register the moe_pool row
+ resolver into the #7 mmid-base resolver surface
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Moves the Android demand-pool's lazy row-pointer resolution out of the direct ggml-cpu.c row
+rewrite and into a registration, at the pool's first registration (g_ready=1), into common surface
+#7's ggml_edge0_mmid_base resolver. NULL keeps the original semantics; zero disturbance while the
+pool is not ready. Both platforms now share one resolver entry point, removing the duplicate
+src0_cur rewrites in #7 and moe_pool.
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/moe_pool.c | 8 ++++++++
+ 1 file changed, 8 insertions(+)
+
+diff --git a/ggml/src/ggml-cpu/moe_pool.c b/ggml/src/ggml-cpu/moe_pool.c
+index 9bd55dc..388c3b5 100644
+--- a/ggml/src/ggml-cpu/moe_pool.c
++++ b/ggml/src/ggml-cpu/moe_pool.c
+@@ -11,6 +11,7 @@
+ //      E0_HEADS / E0_PREROUTER / E0_POOL_STATS (same meanings as v3).
+ #include "moe_pool.h"
+ #include "ggml.h"
++#include "ggml-cpu-impl.h"   // edge0 #7 resolver surface (ggml_edge0_mmid_base_fn / _set_mmid_resolver)
+ 
+ #include <errno.h>
+ #include <fcntl.h>
+@@ -102,6 +103,12 @@ static uint64_t now_ns(void) {
+     return (uint64_t) ts.tv_sec * 1000000000ull + (uint64_t) ts.tv_nsec;
+ }
+ 
++// edge0 consolidation: hang the pool's row resolution on the #7 mmid-base resolver surface (resident -> private frame; otherwise fall back to upstream arithmetic).
++// registered at first registration where g_ready=1; zero disturbance when the pool never becomes ready (E0_POOL_MB unset); NULL keeps the original semantics.
++static const char * moe_pool_mmid_resolver(const struct ggml_tensor * src0, int64_t id, size_t nb02) {
++    return (const char *) moe_pool_row_ptr(src0, id, (const char *) src0->data + id * nb02);
++}
++
+ void moe_pool_init(void) {
+     if (g_ready) return;
+     const char * e = getenv("E0_POOL_MB");
+@@ -527,6 +534,7 @@ void moe_pool_register(struct ggml_tensor * t, int fd, uint64_t file_off) {
+         g_slots = calloc((size_t) g_nslots, sizeof(struct slot_ent));
+         for (int i = 0; i < g_nslots; i++) { g_slots[i].mem = g_arena + (size_t) i * g_slice; g_slots[i].expert = -1; g_slots[i].tix = -1; }
+         g_ready = 1;
++        ggml_edge0_set_mmid_resolver(moe_pool_mmid_resolver);   // edge0 #7 consolidation: row bases go through the resolver surface
+         const char * ie = getenv("E0_IO_THREADS");
+         int nth = ie ? (int) atol(ie) : 4;
+         if (nth < 1) nth = 1; if (nth > 8) nth = 8;
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/common/0001-bailingmoe3-wire-LoRA-into-main-graph-edge0-private-.patch b/patches/llama.cpp/common/0001-bailingmoe3-wire-LoRA-into-main-graph-edge0-private-.patch
new file mode 100644
index 0000000..fbf31a4
--- /dev/null
+++ b/patches/llama.cpp/common/0001-bailingmoe3-wire-LoRA-into-main-graph-edge0-private-.patch
@@ -0,0 +1,125 @@
+From a49736d435951f641bd4bd4ba22fd6792ddb7839 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 15:57:45 +0800
+Subject: [PATCH 1/6] bailingmoe3: wire LoRA into main graph (edge0 private
+ surface #1)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Reroutes the 9 raw ggml_mul_mat call sites in the main graph through build_lora_mm (same shape as
+qwen3next): GDN branch wq/wk/wv (the causal_conv1d helper now takes the already-projected tensor)
++ ssm_f_a/ssm_beta/ssm_g_a/wo; MLA branch wq_b/wq/wqkv_gate/wo. wq_a/wkv_a_mqa/output have no
+adapter target and stay untouched; wk_b/wv_b are 3D batched with no upstream LoRA precedent,
+deferred to phase 2. The MTP graph has no adapter coverage and stays untouched.
+
+For edge0-8b the 147 LoRA pairs previously loaded but silently never entered the graph (the graph
+did not reference them); after this patch the 12-step golden probes went 9/12 -> 11/12 and R3
+free-gen accuracy 35.5% -> 43.6%.
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ src/models/bailingmoe3.cpp | 30 ++++++++++++++----------------
+ 1 file changed, 14 insertions(+), 16 deletions(-)
+
+diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp
+index e208c7d..2643859 100644
+--- a/src/models/bailingmoe3.cpp
++++ b/src/models/bailingmoe3.cpp
+@@ -171,8 +171,7 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
+         ggml_tensor * conv_states_all,
+         ggml_tensor * conv_state_all,
+         int64_t qkv,
+-        ggml_tensor * x,
+-        ggml_tensor * proj_w,
++        ggml_tensor * x_proj,
+         ggml_tensor * conv_w,
+         int64_t d_conv,
+         int64_t head_dim,
+@@ -192,9 +191,8 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
+             total_state_size * ggml_element_size(conv_state_all),
+             qkv * conv_state_size * ggml_element_size(conv_state_all));
+ 
+-    ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
+-    x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
+-    ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);
++    ggml_tensor * x_proj_r = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
++    ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj_r), 0);
+ 
+     const int64_t K = (int64_t) n_rs_seq + 1;
+     const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
+@@ -260,16 +258,16 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
+             ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
+ 
+             ggml_tensor * q = bailingmoe3_causal_conv1d(
+-                    gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
++                    gf, ctx0, conv_states_all, conv_state_all, 0, build_lora_mm(layer.wq, cur), layer.ssm_q_conv,
+                     d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+             ggml_tensor * k = bailingmoe3_causal_conv1d(
+-                    gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
++                    gf, ctx0, conv_states_all, conv_state_all, 1, build_lora_mm(layer.wk, cur), layer.ssm_k_conv,
+                     d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+             ggml_tensor * v = bailingmoe3_causal_conv1d(
+-                    gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
++                    gf, ctx0, conv_states_all, conv_state_all, 2, build_lora_mm(layer.wv, cur), layer.ssm_v_conv,
+                     d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+ 
+-            ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
++            ggml_tensor * gate = build_lora_mm(layer.ssm_f_a, cur);
+             gate = ggml_add(ctx0, gate, layer.ssm_dt_b);
+             gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens);
+             ggml_tensor * a = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1);
+@@ -277,7 +275,7 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
+             gate = ggml_reshape_4d(ctx0, gate, head_dim, n_head, n_seq_tokens, n_seqs);
+             cb(gate, "kda_gate", il);
+ 
+-            ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
++            ggml_tensor * beta = build_lora_mm(layer.ssm_beta, cur);
+             beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs));
+ 
+             q = build_gdn_l2_norm(ctx0, q, hparams.f_norm_rms_eps);
+@@ -290,12 +288,12 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
+             ggml_tensor * out = ggml_cont(ctx0, build_recurrent_attn(
+                     inp_rs, states_all, q, k, v, gate, beta, state, il));
+ 
+-            ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur);
++            ggml_tensor * out_gate = build_lora_mm(layer.ssm_g_a, cur);
+             out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens);
+             out = ggml_reshape_3d(ctx0, out, head_dim, n_head, n_tokens);
+             out = build_norm(out, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
+             out = ggml_mul(ctx0, out, ggml_sigmoid(ctx0, out_gate));
+-            cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens));
++            cur = build_lora_mm(layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens));
+             cb(cur, "kda_out", il);
+         } else {
+             ggml_tensor * attn_input = cur;
+@@ -305,10 +303,10 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
+                 cb(q_all, "q_a", il);
+                 q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+                 cb(q_all, "q_a_norm", il);
+-                q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all);
++                q_all = build_lora_mm(layer.wq_b, q_all);
+                 cb(q_all, "q_b", il);
+             } else {
+-                q_all = ggml_mul_mat(ctx0, layer.wq, cur);
++                q_all = build_lora_mm(layer.wq, cur);
+             }
+             ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens,
+                     ggml_row_size(q_all->type, qk_head_dim),
+@@ -343,11 +341,11 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
+             cur = build_attn(inp_attn, nullptr, nullptr, nullptr,
+                     q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il);
+ 
+-            ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input);
++            ggml_tensor * attn_gate = build_lora_mm(layer.wqkv_gate, attn_input);
+             attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens));
+             cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens);
+             cur = ggml_mul(ctx0, cur, attn_gate);
+-            cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens));
++            cur = build_lora_mm(layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens));
+             cb(cur, "mla_out", il);
+         }
+ 
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/common/0002-models-qwen3next-bailingmoe3-capture-router-input-m_.patch b/patches/llama.cpp/common/0002-models-qwen3next-bailingmoe3-capture-router-input-m_.patch
new file mode 100644
index 0000000..ef7ceb1
--- /dev/null
+++ b/patches/llama.cpp/common/0002-models-qwen3next-bailingmoe3-capture-router-input-m_.patch
@@ -0,0 +1,53 @@
+From c009a31c93bfb4d03d53a0b82ff4c756356583e6 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 16:58:30 +0800
+Subject: [PATCH 2/6] models/qwen3next,bailingmoe3: capture router input (m_in)
+ into t_layer_inp (edge0 private surface #2)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+On both edge0 archs the layer_inp channel semantics is redirected from "layer input" to "router
+input (post-ffn/attn-post-norm)" so the prerouter advisory heads can read features (W4 hook 1).
+Takes effect only while cparams.embeddings_layer_inp[il] is active (the set_outputs gate); other
+archs and the off state see zero difference. The existing D2H readback loop = extract_layer_inputs
+(the production-proven eagle3/dflash channel).
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ src/models/bailingmoe3.cpp | 3 +++
+ src/models/qwen3next.cpp   | 4 ++++
+ 2 files changed, 7 insertions(+)
+
+diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp
+index 2643859..f736435 100644
+--- a/src/models/bailingmoe3.cpp
++++ b/src/models/bailingmoe3.cpp
+@@ -356,6 +356,9 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
+ 
+         ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+         cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
++        // edge0 W4 hook#1: capture router input (m_in) via the layer_inp channel
++        // (consumed only in prerouter mode; same annotation as qwen3next.cpp)
++        res->t_layer_inp[il] = cur;
+ 
+         if ((uint32_t) il < hparams.n_layer_dense_lead) {
+             cur = build_ffn(cur,
+diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp
+index b63fc9c..9673e3b 100644
+--- a/src/models/qwen3next.cpp
++++ b/src/models/qwen3next.cpp
+@@ -183,6 +183,10 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
+         // Post-attention norm
+         ggml_tensor * attn_post_norm = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il);
+         cb(attn_post_norm, "attn_post_norm", il);
++        // edge0 W4 hook#1: capture router input (m_in) via the layer_inp channel.
++        // only consumed when embeddings_layer_inp[il] is active (prerouter mode);
++        // with the feature off, set_outputs() skips this slot entirely.
++        res->t_layer_inp[il] = attn_post_norm;
+ 
+         // FFN layer (MoE or dense) - without residual connection
+         cur = build_layer_ffn(attn_post_norm, il);
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/common/0003-CMake-glob-src-edge0-.cc-into-llama-target-edge0-pri.patch b/patches/llama.cpp/common/0003-CMake-glob-src-edge0-.cc-into-llama-target-edge0-pri.patch
new file mode 100644
index 0000000..422806d
--- /dev/null
+++ b/patches/llama.cpp/common/0003-CMake-glob-src-edge0-.cc-into-llama-target-edge0-pri.patch
@@ -0,0 +1,43 @@
+From d1c5a393fa294dbac96f54c4002d68bf7b3a5147 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 17:04:59 +0800
+Subject: [PATCH 3/6] CMake: glob src/edge0/*.cc into llama target (edge0
+ private surface #3)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Release assembly (scripts/vendor-build.ps1) copies this repo's serve/ pieces into src/edge0/ and
+this glob compiles them into llama.lib; a missing/empty directory = pure no-op. The three W4 hooks
+exhaust the surface caps (#2 capture / #3 glob / #4 context wire-up), so all prerouter logic lives
+in serve/ and consumes no caps.
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ src/CMakeLists.txt | 3 +++
+ 1 file changed, 3 insertions(+)
+
+diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
+index afdaddc..961fc1c 100644
+--- a/src/CMakeLists.txt
++++ b/src/CMakeLists.txt
+@@ -7,6 +7,8 @@ llama_add_compile_flags()
+ # llama
+ 
+ file(GLOB LLAMA_MODELS_SOURCES "models/*.cpp")
++# edge0 private surface #3: directory for our own sources (the release path copies them in from serve/ via the vendor-build script; empty directory = no-op)
++file(GLOB LLAMA_EDGE0_SOURCES "edge0/*.cc" "edge0/*.cpp")
+ 
+ set(LLAMA_CORE_SOURCES
+     llama.cpp
+@@ -48,6 +50,7 @@ add_library(llama
+             ${LLAMA_CORE_SOURCES}
+             unicode.h
+             ${LLAMA_MODELS_SOURCES}
++            ${LLAMA_EDGE0_SOURCES}
+             )
+ 
+ set_target_properties(llama PROPERTIES
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/common/0004-server-wire-edge0-prerouter-advisory-init-on_step-ed.patch b/patches/llama.cpp/common/0004-server-wire-edge0-prerouter-advisory-init-on_step-ed.patch
new file mode 100644
index 0000000..faeea9c
--- /dev/null
+++ b/patches/llama.cpp/common/0004-server-wire-edge0-prerouter-advisory-init-on_step-ed.patch
@@ -0,0 +1,80 @@
+From 8688715e945011c1c192a611a2334e968ddd4283 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 18:25:12 +0800
+Subject: [PATCH 4/6] server: wire edge0 prerouter advisory init/on_step (edge0
+ private surface #4, app-layer)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+edge0_pref_init once load_model has ctx ready; edge0_pref_on_step(ctx, batch_view.n_tokens-1)
+after every llama_decode with ret==0. The implementation lives in src/edge0/prefetch.cc
+(= this repo's serve/ pieces, gitignored there, assembled by the vendor-build script on the
+release path). With E0_PREROUTER unset both call sites return immediately -- zero footprint.
+
+Smoke (35b r3+lora, 48tok, accounting mode): INIT heads=33 gates=33 exps=33 map=ok;
+head 54.0% ~ the reference machine's 55.2% => the ported formula closes the loop.
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ .gitignore                      |  1 +
+ tools/server/server-context.cpp | 22 ++++++++++++++++++++++
+ 2 files changed, 23 insertions(+)
+
+diff --git a/.gitignore b/.gitignore
+index 9b58961..eae8d8d 100644
+--- a/.gitignore
++++ b/.gitignore
+@@ -153,3 +153,4 @@ a.out.*
+ 
+ AGENTS.local.md
+ .pi/SYSTEM.md
++src/edge0/
+diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
+index b6835e4..35f5804 100644
+--- a/tools/server/server-context.cpp
++++ b/tools/server/server-context.cpp
+@@ -1,4 +1,18 @@
+ #include "server-context.h"
++
++// edge0 W4 prerouter advisory (impl: serve/prefetch.cc; DECISIONS 28-B)
++#include <cstdint>
++struct llama_model;
++struct llama_context;
++bool edge0_pref_init(struct llama_model * model, struct llama_context * ctx);
++void edge0_pref_on_step(struct llama_context * ctx, int64_t last_row);
++
++
++// edge0 W4 prerouter advisory (DECISIONS 28-B; implementation in serve/prefetch.cc, compiled into the llama target)
++struct llama_model;
++struct llama_context;
++bool edge0_pref_init(llama_model * model, llama_context * ctx);
++void edge0_pref_on_step(llama_context * ctx, int64_t last_row);
+ #include "server-chat.h"
+ #include "server-common.h"
+ #include "server-http.h"
+@@ -1111,6 +1125,9 @@ private:
+             return false;
+         }
+ 
++        // edge0 W4: prerouter advisory init (E0_PREROUTER unset => sleeps immediately)
++        edge0_pref_init(model_tgt, ctx_tgt);
++
+         vocab = llama_model_get_vocab(model_tgt);
+ 
+         n_ctx = llama_n_ctx(ctx_tgt);
+@@ -3685,6 +3702,11 @@ private:
+             }
+         });
+ 
++        // edge0 W4: prerouter advisory accounting/prefetch (returns immediately when dormant)
++        if (ret == 0) {
++            edge0_pref_on_step(ctx_tgt, (int64_t) batch_view.n_tokens - 1);
++        }
++
+         if (ret != 0) {
+             {
+                 std::string err;
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/common/0005-server-wire-edge0_apply_mem_budget-P-tier-memory-bud.patch b/patches/llama.cpp/common/0005-server-wire-edge0_apply_mem_budget-P-tier-memory-bud.patch
new file mode 100644
index 0000000..624ad2a
--- /dev/null
+++ b/patches/llama.cpp/common/0005-server-wire-edge0_apply_mem_budget-P-tier-memory-bud.patch
@@ -0,0 +1,44 @@
+From 63f3f978913f5307bb3a5b10b0b521c749e9dea7 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 18:51:32 +0800
+Subject: [PATCH 5/6] server: wire edge0_apply_mem_budget (P-tier memory
+ budget, env E0_MEM_BUDGET_MB)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+Apply the process working-set hard cap before load_model (implementation in serve/mem_budget.cc).
+Verified: under a 12GB budget the server loads and generates normally, ws stays bounded and pages
+keep getting trimmed; no slowdown on this machine -- the standby cache backstops it (a large-RAM
+machine cannot measure true cold reads; interpretation logged).
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ tools/server/server-context.cpp | 4 ++++
+ 1 file changed, 4 insertions(+)
+
+diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
+index 35f5804..bc4c7a3 100644
+--- a/tools/server/server-context.cpp
++++ b/tools/server/server-context.cpp
+@@ -6,6 +6,7 @@ struct llama_model;
+ struct llama_context;
+ bool edge0_pref_init(struct llama_model * model, struct llama_context * ctx);
+ void edge0_pref_on_step(struct llama_context * ctx, int64_t last_row);
++bool edge0_apply_mem_budget(void);   // serve/mem_budget.cc: E0_MEM_BUDGET_MB working-set hard cap
+ 
+ 
+ // edge0 W4 prerouter advisory (DECISIONS 28-B; implementation in serve/prefetch.cc, compiled into the llama target)
+@@ -1110,6 +1111,9 @@ private:
+             params_base.load_progress_callback_user_data = &load_progress_text;
+         }
+ 
++        // edge0: P-tier memory budget (E0_MEM_BUDGET_MB unset = no-op; must take effect before the model loads)
++        edge0_apply_mem_budget();
++
+         llama_init = common_init_from_params(params_base);
+ 
+         model_tgt = llama_init->model();
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/common/0006-edge0-private-surface-7-mul_mat_id-expert-base-resol.patch b/patches/llama.cpp/common/0006-edge0-private-surface-7-mul_mat_id-expert-base-resol.patch
new file mode 100644
index 0000000..d1c99b3
--- /dev/null
+++ b/patches/llama.cpp/common/0006-edge0-private-surface-7-mul_mat_id-expert-base-resol.patch
@@ -0,0 +1,114 @@
+From 5fa766c64759650fc259a7d45eed2c3b31034917 Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 21:59:15 +0800
+Subject: [PATCH 6/6] edge0 private surface #7: mul_mat_id expert-base resolver
+ (L1 pool hook)
+
+Three expert-base computations (ggml-cpu.c forward_mul_mat_id, iqp.cpp panel
+path, repack.cpp template path) now route through ggml_edge0_mmid_base(src0,
+id, nb02). Default NULL resolver = exact upstream arithmetic (zero behavior
+change); the llama-side edge0 pool installs a resolver at init to serve expert
+bytes from committed private pages (byte-verified vs the file view by
+E0_POOL_SELFCHECK sampling). Export macro keyed on ggml_cpu_EXPORTS (this fork
+builds ggml-cpu as its own DLL); decls wrapped in extern "C".
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ ggml/src/ggml-cpu/ggml-cpu-impl.h | 20 ++++++++++++++++++++
+ ggml/src/ggml-cpu/ggml-cpu.c      | 14 +++++++++++++-
+ ggml/src/ggml-cpu/iqp.cpp         |  2 +-
+ ggml/src/ggml-cpu/repack.cpp      |  2 +-
+ 4 files changed, 35 insertions(+), 3 deletions(-)
+
+diff --git a/ggml/src/ggml-cpu/ggml-cpu-impl.h b/ggml/src/ggml-cpu/ggml-cpu-impl.h
+index 5dd9ec8..85cfc1c 100644
+--- a/ggml/src/ggml-cpu/ggml-cpu-impl.h
++++ b/ggml/src/ggml-cpu/ggml-cpu-impl.h
+@@ -5,6 +5,26 @@
+ #include "ggml.h"
+ #include "ggml-impl.h"
+ 
++// edge0 private surface #7 — MoE expert-base indirection for the edge0 L1 pool
++// (installed at runtime by serve/prefetch.cc; NULL = upstream pass-through semantics).
++// this fork builds ggml-cpu as a standalone DLL: decoration is bidirectional via the target's EXPORTS macro (GGML_API resolves to import inside this target's translation units, so it cannot be used).
++#if defined(_WIN32) && defined(ggml_cpu_EXPORTS)
++#  define EDGE0_MMID_API __declspec(dllexport)
++#elif defined(_WIN32)
++#  define EDGE0_MMID_API __declspec(dllimport)
++#else
++#  define EDGE0_MMID_API
++#endif
++#ifdef __cplusplus
++extern "C" {
++#endif
++typedef const char * (*ggml_edge0_mmid_base_fn)(const struct ggml_tensor * src0, int64_t id, size_t nb02);
++EDGE0_MMID_API void         ggml_edge0_set_mmid_resolver(ggml_edge0_mmid_base_fn fn);
++EDGE0_MMID_API const char * ggml_edge0_mmid_base(const struct ggml_tensor * src0, int64_t id, size_t nb02);
++#ifdef __cplusplus
++}
++#endif
++
+ #include <stdlib.h> // load `stdlib.h` before other headers to work around MinGW bug: https://sourceforge.net/p/mingw-w64/bugs/192/
+ //#include <stddef.h>
+ #include <stdbool.h>
+diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c
+index 8bb0ff7..be7c2db 100644
+--- a/ggml/src/ggml-cpu/ggml-cpu.c
++++ b/ggml/src/ggml-cpu/ggml-cpu.c
+@@ -1469,6 +1469,18 @@ UseGgmlGemm2:;
+ 
+ // ggml_compute_forward_mul_mat_id
+ 
++// edge0 private surface #7 — expert-base indirection (pool hook; default = upstream math)
++static ggml_edge0_mmid_base_fn g_edge0_mmid_base_fn = NULL;
++
++void ggml_edge0_set_mmid_resolver(ggml_edge0_mmid_base_fn fn) {
++    g_edge0_mmid_base_fn = fn;
++}
++
++const char * ggml_edge0_mmid_base(const struct ggml_tensor * src0, int64_t id, size_t nb02) {
++    return g_edge0_mmid_base_fn ? g_edge0_mmid_base_fn(src0, id, nb02)
++                                : (const char *) src0->data + id * nb02;
++}
++
+ #define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id)*ids->ne[0]*ids->ne[1] + (i1)]
+ 
+ struct mmid_row_mapping {
+@@ -1684,7 +1696,7 @@ static void ggml_compute_forward_mul_mat_id(
+             continue;
+         }
+ 
+-        const char * src0_cur = (const char *) src0->data + cur_a * nb02;
++        const char * src0_cur = ggml_edge0_mmid_base(src0, cur_a, nb02);  // edge0 #7
+         const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
+         const size_t row_size = ggml_row_size(vec_dot_type, ne10);
+ 
+diff --git a/ggml/src/ggml-cpu/iqp.cpp b/ggml/src/ggml-cpu/iqp.cpp
+index b9201db..031d7b3 100644
+--- a/ggml/src/ggml-cpu/iqp.cpp
++++ b/ggml/src/ggml-cpu/iqp.cpp
+@@ -1147,7 +1147,7 @@ void ggml_compute_forward_mul_mat_id_iqp(const struct ggml_compute_params * para
+ 
+     block_iqp_x8 * panel = (block_iqp_x8 *) ((char *) panels + (size_t) ith * ggml_cpu_iqp_scratch_size(dst));
+ 
+-    const char * src0_cur = (const char *) src0->data + cur_a * nb02;
++    const char * src0_cur = ggml_edge0_mmid_base(src0, cur_a, nb02);  // edge0 #7
+ 
+     const int64_t ngroups = ne01 / IQP_NB_ROWS;
+ 
+diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp
+index d56db98..f402385 100644
+--- a/ggml/src/ggml-cpu/repack.cpp
++++ b/ggml/src/ggml-cpu/repack.cpp
+@@ -4873,7 +4873,7 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
+                 continue;
+             }
+ 
+-            const auto * src0_cur = (const char *) src0->data + cur_a*nb02;
++            const auto * src0_cur = ggml_edge0_mmid_base(src0, cur_a, nb02);  // edge0 #7
+ 
+             //const int64_t nr0 = ne01; // src0 rows
+             const int64_t nr1 = cne1; // src1 rows
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/windows/0001-common-server-mem-budget-mb-CLI-edge0-private-surfac.patch b/patches/llama.cpp/windows/0001-common-server-mem-budget-mb-CLI-edge0-private-surfac.patch
new file mode 100644
index 0000000..5aef1b6
--- /dev/null
+++ b/patches/llama.cpp/windows/0001-common-server-mem-budget-mb-CLI-edge0-private-surfac.patch
@@ -0,0 +1,80 @@
+From edef0ec1f8d5e43239b80fe3d91e01aa47a27e5a Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 19:54:14 +0800
+Subject: [PATCH 1/2] common+server: --mem-budget-mb CLI (edge0 private surface
+ #6)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+common_params.mem_budget_mb + arg registration (env LLAMA_ARG_MEM_BUDGET_MB on the same line);
+consumed by server before load_model => edge0_apply_mem_budget(mb). The env channel stays as a
+debug fallback. Smoke: --mem-budget-mb 16384 -> [mem-budget] applied; visible in --help.
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ common/arg.cpp                  | 12 ++++++++++++
+ common/common.h                 |  1 +
+ tools/server/server-context.cpp |  4 ++--
+ 3 files changed, 15 insertions(+), 2 deletions(-)
+
+diff --git a/common/arg.cpp b/common/arg.cpp
+index 996ea75..111e4ed 100644
+--- a/common/arg.cpp
++++ b/common/arg.cpp
+@@ -2780,6 +2780,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
+             llm_add_n_cpu_ffn_overrides(value, LLM_FFN_DENSE_REGEX, params.tensor_buft_overrides);
+         }
+     ).set_env("LLAMA_ARG_N_CPU_FFN"));
++    add_opt(common_arg(
++        {"--mem-budget-mb"}, "N",
++        "edge0: cap this process' max working set at N MB (hard trim on overflow; 0=off)\n"
++
++        "(P-tier memory budget; predicted-hot pages can later be pinned within it)",
++        [](common_params & params, int value) {
++            if (value < 0) {
++                throw std::invalid_argument("invalid value");
++            }
++            params.mem_budget_mb = value;
++        }
++    ).set_env("LLAMA_ARG_MEM_BUDGET_MB"));
+     GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0
+     add_opt(common_arg(
+         {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N",
+diff --git a/common/common.h b/common/common.h
+index 63d0bad..bfa67d9 100644
+--- a/common/common.h
++++ b/common/common.h
+@@ -471,6 +471,7 @@ struct common_params {
+     std::vector<ggml_backend_dev_t> devices; // devices to use for offloading
+ 
+     int32_t n_gpu_layers       = -1;    // number of layers to store in VRAM, -1 is auto, <= -2 is all
++    int32_t mem_budget_mb      = 0;     // edge0: cap process max working set in MB (0=off; P-tier memory budget)
+     int32_t main_gpu           = 0;     // the GPU that is used for scratch and small tensors
+     float   tensor_split[128]  = {0};   // how split tensors should be distributed across GPUs
+     bool    fit_params         = true;  // whether to fit unset model/context parameters to free device memory
+diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
+index bc4c7a3..1859a24 100644
+--- a/tools/server/server-context.cpp
++++ b/tools/server/server-context.cpp
+@@ -6,7 +6,7 @@ struct llama_model;
+ struct llama_context;
+ bool edge0_pref_init(struct llama_model * model, struct llama_context * ctx);
+ void edge0_pref_on_step(struct llama_context * ctx, int64_t last_row);
+-bool edge0_apply_mem_budget(void);   // serve/mem_budget.cc: E0_MEM_BUDGET_MB working-set hard cap
++bool edge0_apply_mem_budget(int mb);   // serve/mem_budget.cc: --mem-budget-mb working-set hard cap
+ 
+ 
+ // edge0 W4 prerouter advisory (DECISIONS 28-B; implementation in serve/prefetch.cc, compiled into the llama target)
+@@ -1112,7 +1112,7 @@ private:
+         }
+ 
+         // edge0: P-tier memory budget (E0_MEM_BUDGET_MB unset = no-op; must take effect before the model loads)
+-        edge0_apply_mem_budget();
++        edge0_apply_mem_budget(params_base.mem_budget_mb);
+ 
+         llama_init = common_init_from_params(params_base);
+ 
+-- 
+2.52.0.windows.1
+
diff --git a/patches/llama.cpp/windows/0002-common-server-pool-mb-CLI-edge0-private-surface-8.patch b/patches/llama.cpp/windows/0002-common-server-pool-mb-CLI-edge0-private-surface-8.patch
new file mode 100644
index 0000000..d8a17f8
--- /dev/null
+++ b/patches/llama.cpp/windows/0002-common-server-pool-mb-CLI-edge0-private-surface-8.patch
@@ -0,0 +1,95 @@
+From 7f9c2513d789728292e72c6708ed4ee497dadbab Mon Sep 17 00:00:00 2001
+From: Alex Zhou <alex.zhou@edge0.ai>
+Date: Mon, 28 Sep 2026 23:28:28 +0800
+Subject: [PATCH 2/2] common+server: --pool-mb CLI (edge0 private surface #8)
+MIME-Version: 1.0
+Content-Type: text/plain; charset=UTF-8
+Content-Transfer-Encoding: 8bit
+
+common_params.pool_mb + arg registration (env LLAMA_ARG_POOL_MB on the same line); server passes
+params_base.pool_mb + model.path into edge0_pref_init (signature extended; env E0_GGUF/E0_POOL_MB
+stay as harness/debug fallbacks, CLI takes precedence). Also removes the duplicate declaration
+block left over from #4. Smoke: 35b --pool-mb 1024 on a clean env -> POOL2 resolver=on tt=120,
+hit/fillD incrementing normally.
+
+Co-Authored-By: Claude Code <noreply@anthropic.com>
+---
+ common/arg.cpp                  | 11 +++++++++++
+ common/common.h                 |  1 +
+ tools/server/server-context.cpp | 14 ++++----------
+ 3 files changed, 16 insertions(+), 10 deletions(-)
+
+diff --git a/common/arg.cpp b/common/arg.cpp
+index 111e4ed..73ded46 100644
+--- a/common/arg.cpp
++++ b/common/arg.cpp
+@@ -2792,6 +2792,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
+             params.mem_budget_mb = value;
+         }
+     ).set_env("LLAMA_ARG_MEM_BUDGET_MB"));
++    add_opt(common_arg(
++        {"--pool-mb"}, "N",
++        "edge0: expert L1 cache pool of N MB (private pages, mul_mat_id direct read; 0=off)\n"
++        "(P-tier contract: pool bytes are the hard commitment, remaining RAM stays opportunistic OS cache)",
++        [](common_params & params, int value) {
++            if (value < 0) {
++                throw std::invalid_argument("invalid value");
++            }
++            params.pool_mb = value;
++        }
++    ).set_env("LLAMA_ARG_POOL_MB"));
+     GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0
+     add_opt(common_arg(
+         {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N",
+diff --git a/common/common.h b/common/common.h
+index bfa67d9..683b4e5 100644
+--- a/common/common.h
++++ b/common/common.h
+@@ -472,6 +472,7 @@ struct common_params {
+ 
+     int32_t n_gpu_layers       = -1;    // number of layers to store in VRAM, -1 is auto, <= -2 is all
+     int32_t mem_budget_mb      = 0;     // edge0: cap process max working set in MB (0=off; P-tier memory budget)
++    int32_t pool_mb            = 0;     // edge0: expert L1 cache pool in MB, private pages direct-read (0=off; DECISIONS 28/16)
+     int32_t main_gpu           = 0;     // the GPU that is used for scratch and small tensors
+     float   tensor_split[128]  = {0};   // how split tensors should be distributed across GPUs
+     bool    fit_params         = true;  // whether to fit unset model/context parameters to free device memory
+diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
+index 1859a24..f607072 100644
+--- a/tools/server/server-context.cpp
++++ b/tools/server/server-context.cpp
+@@ -1,19 +1,13 @@
+ #include "server-context.h"
+ 
+-// edge0 W4 prerouter advisory (impl: serve/prefetch.cc; DECISIONS 28-B)
++// edge0 W4 prerouter advisory + 28/16 real pool (impl: src/edge0/prefetch.cc; DECISIONS 28-B/16)
+ #include <cstdint>
+ struct llama_model;
+ struct llama_context;
+-bool edge0_pref_init(struct llama_model * model, struct llama_context * ctx);
++bool edge0_pref_init(struct llama_model * model, struct llama_context * ctx, int pool_mb, const char * gguf_path);
+ void edge0_pref_on_step(struct llama_context * ctx, int64_t last_row);
+ bool edge0_apply_mem_budget(int mb);   // serve/mem_budget.cc: --mem-budget-mb working-set hard cap
+ 
+-
+-// edge0 W4 prerouter advisory (DECISIONS 28-B; implementation in serve/prefetch.cc, compiled into the llama target)
+-struct llama_model;
+-struct llama_context;
+-bool edge0_pref_init(llama_model * model, llama_context * ctx);
+-void edge0_pref_on_step(llama_context * ctx, int64_t last_row);
+ #include "server-chat.h"
+ #include "server-common.h"
+ #include "server-http.h"
+@@ -1129,8 +1123,8 @@ private:
+             return false;
+         }
+ 
+-        // edge0 W4: prerouter advisory init (E0_PREROUTER unset => sleeps immediately)
+-        edge0_pref_init(model_tgt, ctx_tgt);
++        // edge0 W4+16: prerouter/pool init (both E0_PREROUTER and --pool-mb unset => sleeps immediately)
++        edge0_pref_init(model_tgt, ctx_tgt, params_base.pool_mb, params_base.model.path.c_str());
+ 
+         vocab = llama_model_get_vocab(model_tgt);
+ 
+-- 
+2.52.0.windows.1
+
diff --git a/python/README.md b/python/README.md
new file mode 100644
index 0000000..3a323f4
--- /dev/null
+++ b/python/README.md
@@ -0,0 +1,204 @@
+# edge0 — Python framework
+
+English | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+This document is for **developers installing and running the edge0 Python
+framework from source**. It covers: quick start (install → models → run),
+measured performance and quality, and the design behind the stack.
+
+edge0 is published as a **monorepo** —
+[`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — whose top level
+holds the shared docs (`docs/`) and the platform subprojects (`python/` =
+this framework, plus `macos/`, `ios/`, `android/`, `windows/`). The Python
+framework is the reference implementation of the edge0 recipe — **SSD
+expert offload + Recover-LoRA + prerouter routing prediction** — running
+large sparse-MoE models on Apple Silicon via MLX, with peak memory bounded
+by the *active* expert set rather than the parameter count.
+
+Third-party attribution: see [`NOTICE`](../NOTICE) at the repo root. Model
+cards: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+· [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+All commands below run from this directory (`python/`).
+
+---
+
+## Quick Start
+
+### Prerequisites
+
+| component | requirement |
+|---|---|
+| OS / hardware | macOS on Apple Silicon (M1/M2/M3/M4) — the MLX backend is Apple-Silicon-only; a CUDA backend is on the roadmap |
+| Python | 3.10+ (3.12 recommended) |
+| MLX | `mlx==0.30.6` / `mlx-metal==0.30.6` with `mlx-lm==0.31.0` (see `pyproject.toml`) |
+| Memory | ~2.9 GB peak active (edge0-35b), ~1.0 GB (edge0-8b), short contexts |
+| Disk | ~23 GB (edge0-35b) / ~4.2 GB (edge0-8b); expert weights are mmapped and read on demand, not loaded into RAM up front |
+
+> Garbled, mixed-language output on Apple A18 / A18 Pro means an older
+> `mlx`: `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
+> ([#8](https://github.com/Edge0-AI/Edge0/issues/8)).
+
+### Install
+
+```bash
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+### Models
+
+Both tiers are published on Hugging Face and ModelScope; each repo bundles
+the base checkpoint and the trained LoRA + prerouter adapters in **one
+directory**, so a single download is a ready-to-run model:
+
+```bash
+# the repo's helper (defaults to the two published tiers):
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b  --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview  --local-dir models/edge0-8b
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+```
+
+Point edge0 at a tier through an environment variable, or pass the
+directory directly (the tier is auto-detected from `config.json`):
+
+```bash
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+### Run
+
+```bash
+edge0 demo edge0-8b                 # one-shot generation demo
+edge0 chat edge0-8b --prompt "Explain streaming inference in one sentence."
+edge0 serve edge0-8b                # OpenAI-compatible server on http://127.0.0.1:8000
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-8b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+```
+
+`python -m edge0 ...` is equivalent to `edge0 ...`.
+
+### Python API
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`examples/demo.py` is the same minimal walkthrough (`edge0 demo` runs this
+exact path).
+
+---
+
+## Performance
+
+Measured with `examples/bench.py` (3.3k-token prompt prefill → 10 sampled
+warmup steps → 200 timed sampled decode tokens, 2 runs per tier):
+
+| Tier | Decode speed | Prefill throughput (cold / warm)* | Peak active memory | Test machine |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*Cold = first request after process start (expert weights fault in from
+SSD); warm = subsequent requests (page cache resident).*
+
+```bash
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+```
+
+## Quality
+
+Run with [OpenCompass](https://github.com/open-compass/opencompass) under
+identical settings for the edge0 models (int4 + trained adapters +
+prerouter) and the original fp16 bases. The edge0 pipeline loses little:
+**3.9 points on average for edge0-35b, 2.8 for edge0-8b** (MMLU-Pro is even
+above the base). Max 100:
+
+| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **Average** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+---
+
+## Technical Details
+
+### Core mechanisms
+
+- **SSD expert offload** — expert weights stream from storage on demand;
+  peak memory is bounded by the active set, not the parameter count.
+- **Prerouter** — a trained head predicts expert routing one step ahead, so
+  expert loads overlap the forward pass instead of stalling it (**up to
+  +59%** decode throughput; the gain grows with storage latency, model
+  size, and routed width *K*).
+- **Recover-LoRA** — the int4 base is frozen and LoRA adapters are trained
+  by distillation from the FP teacher, recovering most of the 4-bit
+  quantization loss. Adapters stay unmerged: one read-only base serves
+  multiple adapter sets.
+
+### Design
+
+- **transformers-style usage** — `AutoModel` / `AutoConfig` / `AutoEngine`
+  resolve the tier from the model name.
+- **Backend isolation** — all MLX code lives under
+  `src/edge0/backends/mlx/`; the core logic (model specs, prerouter,
+  streaming expert pool, server) depends only on the backend facade
+  (`backends/base.py`), so a new backend implements the same facade
+  (`backends/cuda/` is a reserved slot) with zero changes to core code.
+- **Adapters as safetensors** — LoRA and prerouter weights are
+  `.safetensors` files with provenance metadata (source, version, owner
+  layers), resolved from the model directory or the gitignored `artifacts/`
+  fallback.
+- **Model + adapters in one directory** — a model directory holds the base
+  checkpoint and that model's adapters; upgrading adapters swaps adapter
+  files only — the base stays read-only and is never merged.
+
+### Package layout
+
+```
+src/edge0/
+├── backends/mlx/                  # MLX backend (isolation boundary; cuda/ reserved)
+├── engine/  models/  moe/  prerouter/  streaming/
+├── adapters/  attention/  server/
+└── cli.py  registry.py  sampling.py  config.py
+```
+
+---
+
+## Documentation
+
+- [Architecture](../docs/architecture.md)
+- [Attention](../docs/attention.md) / [MoE](../docs/moe.md) / [SSD streaming](../docs/streaming.md) / [prerouter](../docs/prerouter.md)
+- [Adding a model](../docs/adding-a-model.md)
+- [edge0-35b](../docs/models/edge0-35b.md) / [edge0-8b](../docs/models/edge0-8b.md)
+- Technical report: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](../paper/main.pdf))
+
+## Tests
+
+```bash
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q   # real-weight generation
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+```
diff --git a/python/README_es.md b/python/README_es.md
new file mode 100644
index 0000000..5276f72
--- /dev/null
+++ b/python/README_es.md
@@ -0,0 +1,215 @@
+# edge0 — Framework de Python
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | Español | [Français](README_fr.md)
+
+Este documento está dirigido a **desarrolladores que instalan y ejecutan el
+framework de Python de edge0 desde el código fuente**. Cubre: inicio rápido
+(instalación → modelos → ejecución), rendimiento y calidad medidos, y el
+diseño detrás del stack.
+
+edge0 se publica como **monorepo** —
+[`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — cuyo nivel superior
+contiene la documentación compartida (`docs/`) y los subproyectos de
+plataforma (`python/` = este framework, además de `macos/`, `ios/`,
+`android/`, `windows/`). El framework de Python es la implementación de
+referencia de la receta edge0 — **SSD expert offload + Recover-LoRA +
+predicción de enrutamiento del prerouter** — que ejecuta grandes modelos
+sparse-MoE en Apple Silicon mediante MLX, con una memoria pico limitada por el
+conjunto de expertos *activos* en lugar del número de parámetros.
+
+Atribución de terceros: consulta [`NOTICE`](../NOTICE) en la raíz del
+repositorio. Fichas de modelo:
+[Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+· [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+Todos los comandos siguientes se ejecutan desde este directorio (`python/`).
+
+---
+
+## Inicio rápido
+
+### Requisitos previos
+
+| componente | requisito |
+|---|---|
+| SO / hardware | macOS en Apple Silicon (M1/M2/M3/M4) — el backend MLX es exclusivo de Apple Silicon; un backend CUDA está en la hoja de ruta |
+| Python | 3.10+ (se recomienda 3.12) |
+| MLX | `mlx==0.30.6` / `mlx-metal==0.30.6` con `mlx-lm==0.31.0` (consulta `pyproject.toml`) |
+| Memoria | ~2.9 GB activos pico (edge0-35b), ~1.0 GB (edge0-8b), contextos cortos |
+| Disco | ~23 GB (edge0-35b) / ~4.2 GB (edge0-8b); los pesos de los expertos se mapean con mmap y se leen bajo demanda, no se cargan en RAM por adelantado |
+
+> Una salida ilegible o mezclada de idiomas en Apple A18 / A18 Pro indica un
+> `mlx` antiguo: `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
+> ([#8](https://github.com/Edge0-AI/Edge0/issues/8)).
+
+### Instalación
+
+```bash
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+### Modelos
+
+Ambos niveles están publicados en Hugging Face y ModelScope; cada repositorio
+incluye el checkpoint base y los adaptadores entrenados de LoRA + prerouter en
+**un solo directorio**, de modo que una única descarga es un modelo listo para
+ejecutar:
+
+```bash
+# the repo's helper (defaults to the two published tiers):
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b  --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview  --local-dir models/edge0-8b
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+```
+
+Apunta edge0 a un nivel mediante una variable de entorno, o pasa el directorio
+directamente (el nivel se detecta automáticamente desde `config.json`):
+
+```bash
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+### Ejecución
+
+```bash
+edge0 demo edge0-8b                 # one-shot generation demo
+edge0 chat edge0-8b --prompt "Explain streaming inference in one sentence."
+edge0 serve edge0-8b                # OpenAI-compatible server on http://127.0.0.1:8000
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-8b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+```
+
+`python -m edge0 ...` es equivalente a `edge0 ...`.
+
+### API de Python
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`examples/demo.py` es el mismo recorrido mínimo (`edge0 demo` ejecuta
+exactamente esta ruta).
+
+---
+
+## Rendimiento
+
+Medido con `examples/bench.py` (prefill de un prompt de 3.3k tokens → 10 pasos
+de calentamiento muestreados → 200 tokens de decode muestreados y
+cronometrados, 2 ejecuciones por nivel):
+
+| Nivel | Velocidad de decode | Throughput de prefill (frío / caliente)* | Memoria activa pico | Máquina de pruebas |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*Frío = primera petición tras iniciar el proceso (los pesos de expertos entran
+por fallos de página desde el SSD); caliente = peticiones posteriores
+(residentes en la caché de páginas).*
+
+```bash
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+```
+
+## Calidad
+
+Ejecutado con [OpenCompass](https://github.com/open-compass/opencompass) bajo
+configuraciones idénticas para los modelos edge0 (int4 + adaptadores
+entrenados + prerouter) y las bases fp16 originales. El pipeline edge0 pierde
+poco: **3.9 puntos de media para edge0-35b, 2.8 para edge0-8b** (MMLU-Pro
+incluso supera a la base). Máximo 100:
+
+| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **Media** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+---
+
+## Detalles técnicos
+
+### Mecanismos principales
+
+- **SSD expert offload** — los pesos de los expertos se transmiten desde el
+  almacenamiento bajo demanda; la memoria pico está limitada por el conjunto
+  activo, no por el número de parámetros.
+- **Prerouter** — una cabeza entrenada predice el enrutamiento de expertos un
+  paso por delante, de modo que las cargas de expertos se solapan con el
+  forward pass en lugar de bloquearlo (**hasta +59%** de throughput de decode;
+  la ganancia crece con la latencia del almacenamiento, el tamaño del modelo y
+  el ancho de enrutamiento *K*).
+- **Recover-LoRA** — la base int4 está congelada y los adaptadores LoRA se
+  entrenan por destilación desde el teacher FP, recuperando la mayor parte de
+  la pérdida de cuantización a 4 bits. Los adaptadores permanecen sin
+  fusionar: una única base de solo lectura sirve a varios conjuntos de
+  adaptadores.
+
+### Diseño
+
+- **Uso al estilo transformers** — `AutoModel` / `AutoConfig` / `AutoEngine`
+  resuelven el nivel a partir del nombre del modelo.
+- **Aislamiento de backends** — todo el código MLX vive en
+  `src/edge0/backends/mlx/`; la lógica principal (especificaciones de modelos,
+  prerouter, pool de expertos en streaming, servidor) depende solo de la
+  fachada del backend (`backends/base.py`), de modo que un nuevo backend
+  implementa la misma fachada (`backends/cuda/` es un espacio reservado) sin
+  ningún cambio en el código principal.
+- **Adaptadores como safetensors** — los pesos de LoRA y del prerouter son
+  archivos `.safetensors` con metadatos de procedencia (fuente, versión, capas
+  propietarias), resueltos desde el directorio del modelo o la alternativa
+  `artifacts/` ignorada por git.
+- **Modelo + adaptadores en un solo directorio** — un directorio de modelo
+  contiene el checkpoint base y los adaptadores de ese modelo; actualizar los
+  adaptadores solo reemplaza los archivos de adaptadores — la base permanece
+  de solo lectura y nunca se fusiona.
+
+### Estructura del paquete
+
+```
+src/edge0/
+├── backends/mlx/                  # MLX backend (isolation boundary; cuda/ reserved)
+├── engine/  models/  moe/  prerouter/  streaming/
+├── adapters/  attention/  server/
+└── cli.py  registry.py  sampling.py  config.py
+```
+
+---
+
+## Documentación
+
+- [Arquitectura](../docs/architecture.md)
+- [Atención](../docs/attention.md) / [MoE](../docs/moe.md) / [Streaming SSD](../docs/streaming.md) / [prerouter](../docs/prerouter.md)
+- [Añadir un modelo](../docs/adding-a-model.md)
+- [edge0-35b](../docs/models/edge0-35b.md) / [edge0-8b](../docs/models/edge0-8b.md)
+- Informe técnico: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](../paper/main.pdf))
+
+## Pruebas
+
+```bash
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q   # real-weight generation
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+```
diff --git a/python/README_fr.md b/python/README_fr.md
new file mode 100644
index 0000000..85f51d4
--- /dev/null
+++ b/python/README_fr.md
@@ -0,0 +1,219 @@
+# edge0 — Framework Python
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | Français
+
+Ce document s'adresse aux **développeurs qui installent et exécutent le
+framework Python edge0 depuis les sources**. Il couvre : le démarrage
+rapide (installation → modèles → exécution), les performances et la
+qualité mesurées, ainsi que la conception de la stack.
+
+edge0 est publié sous forme de **monorepo** —
+[`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — dont le niveau
+supérieur contient la documentation partagée (`docs/`) et les
+sous-projets par plateforme (`python/` = ce framework, ainsi que
+`macos/`, `ios/`, `android/`, `windows/`). Le framework Python est
+l'implémentation de référence de la recette edge0 — **SSD expert offload
++ Recover-LoRA + prédiction de routage prerouter** — qui exécute de
+grands modèles MoE creux (sparse-MoE) sur Apple Silicon via MLX, avec une
+mémoire de pointe limitée par l'ensemble *actif* d'experts plutôt que par
+le nombre de paramètres.
+
+Attribution tierce : voir [`NOTICE`](../NOTICE) à la racine du dépôt.
+Cartes de modèles : [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+· [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+Toutes les commandes ci-dessous s'exécutent depuis ce répertoire
+(`python/`).
+
+---
+
+## Démarrage rapide
+
+### Prérequis
+
+| composant | exigence |
+|---|---|
+| OS / matériel | macOS sur Apple Silicon (M1/M2/M3/M4) — le backend MLX est réservé à Apple Silicon ; un backend CUDA est sur la feuille de route |
+| Python | 3.10+ (3.12 recommandé) |
+| MLX | `mlx==0.30.6` / `mlx-metal==0.30.6` avec `mlx-lm==0.31.0` (voir `pyproject.toml`) |
+| Mémoire | ~2,9 GB actifs de pointe (edge0-35b), ~1,0 GB (edge0-8b), contextes courts |
+| Disque | ~23 GB (edge0-35b) / ~4,2 GB (edge0-8b) ; les poids des experts sont mmappés et lus à la demande, pas chargés en RAM au démarrage |
+
+> Une sortie illisible ou mêlant plusieurs langues sur Apple A18 / A18 Pro
+> indique une version plus ancienne de `mlx` :
+> `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
+> ([#8](https://github.com/Edge0-AI/Edge0/issues/8)).
+
+### Installation
+
+```bash
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+### Modèles
+
+Les deux gammes sont publiées sur Hugging Face et ModelScope ; chaque
+dépôt regroupe le checkpoint de base et les adaptateurs LoRA + prerouter
+entraînés dans **un seul répertoire**, de sorte qu'un unique
+téléchargement constitue un modèle prêt à l'emploi :
+
+```bash
+# the repo's helper (defaults to the two published tiers):
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b  --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview  --local-dir models/edge0-8b
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+```
+
+Pointez edge0 vers une gamme via une variable d'environnement, ou passez
+directement le répertoire (la gamme est détectée automatiquement depuis
+`config.json`) :
+
+```bash
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+### Exécution
+
+```bash
+edge0 demo edge0-8b                 # one-shot generation demo
+edge0 chat edge0-8b --prompt "Explain streaming inference in one sentence."
+edge0 serve edge0-8b                # OpenAI-compatible server on http://127.0.0.1:8000
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-8b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+```
+
+`python -m edge0 ...` est équivalent à `edge0 ...`.
+
+### API Python
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`examples/demo.py` est le même parcours minimal (`edge0 demo` exécute
+exactement ce chemin).
+
+---
+
+## Performances
+
+Mesurées avec `examples/bench.py` (prefill d'un prompt de 3,3 k tokens →
+10 pas d'échauffement échantillonnés → 200 tokens de decode
+échantillonnés et chronométrés, 2 exécutions par gamme) :
+
+| Gamme | Vitesse de decode | Débit de prefill (à froid / à chaud)* | Mémoire active de pointe | Machine de test |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*À froid = première requête après le démarrage du processus (les poids des
+experts arrivent depuis le SSD par défaut de page) ; à chaud = requêtes
+suivantes (page cache résident).*
+
+```bash
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+```
+
+## Qualité
+
+Exécutée avec [OpenCompass](https://github.com/open-compass/opencompass)
+dans des conditions identiques pour les modèles edge0 (int4 + adaptateurs
+entraînés + prerouter) et les bases fp16 d'origine. Le pipeline edge0 perd
+peu : **3,9 points en moyenne pour edge0-35b, 2,8 pour edge0-8b**
+(MMLU-Pro est même au-dessus de la base). Maximum 100 :
+
+| Benchmark | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **Moyenne** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+---
+
+## Détails techniques
+
+### Mécanismes fondamentaux
+
+- **SSD expert offload** — les poids des experts sont lus en streaming
+  depuis le stockage à la demande ; la mémoire de pointe est limitée par
+  l'ensemble actif, et non par le nombre de paramètres.
+- **Prerouter** — une tête entraînée prédit le routage des experts un pas
+  à l'avance, de sorte que les chargements d'experts chevauchent la passe
+  avant au lieu de la bloquer (**jusqu'à +59 %** de débit de decode ; le
+  gain augmente avec la latence du stockage, la taille du modèle et la
+  largeur routée *K*).
+- **Recover-LoRA** — la base int4 est gelée et les adaptateurs LoRA sont
+  entraînés par distillation depuis l'enseignant FP, récupérant
+  l'essentiel de la perte de quantification 4-bit. Les adaptateurs restent
+  non fusionnés : une base unique en lecture seule sert plusieurs jeux
+  d'adaptateurs.
+
+### Conception
+
+- **Utilisation façon transformers** — `AutoModel` / `AutoConfig` / `AutoEngine`
+  résolvent la gamme à partir du nom du modèle.
+- **Isolation du backend** — tout le code MLX se trouve sous
+  `src/edge0/backends/mlx/` ; la logique centrale (spécifications des
+  modèles, prerouter, pool d'experts en streaming, serveur) ne dépend que
+  de la façade du backend (`backends/base.py`), de sorte qu'un nouveau
+  backend implémente la même façade (`backends/cuda/` est un emplacement
+  réservé) sans aucune modification du code central.
+- **Adaptateurs au format safetensors** — les poids LoRA et prerouter sont
+  des fichiers `.safetensors` avec métadonnées de provenance (source,
+  version, couches propriétaires), résolus depuis le répertoire du modèle
+  ou le cache de secours `artifacts/` ignoré par git.
+- **Modèle + adaptateurs dans un même répertoire** — un répertoire de
+  modèle contient le checkpoint de base et les adaptateurs de ce modèle ;
+  la mise à niveau des adaptateurs ne remplace que les fichiers
+  d'adaptateurs — la base reste en lecture seule et n'est jamais
+  fusionnée.
+
+### Structure du paquet
+
+```
+src/edge0/
+├── backends/mlx/                  # MLX backend (isolation boundary; cuda/ reserved)
+├── engine/  models/  moe/  prerouter/  streaming/
+├── adapters/  attention/  server/
+└── cli.py  registry.py  sampling.py  config.py
+```
+
+---
+
+## Documentation
+
+- [Architecture](../docs/architecture.md)
+- [Attention](../docs/attention.md) / [MoE](../docs/moe.md) / [Streaming SSD](../docs/streaming.md) / [prerouter](../docs/prerouter.md)
+- [Ajouter un modèle](../docs/adding-a-model.md)
+- [edge0-35b](../docs/models/edge0-35b.md) / [edge0-8b](../docs/models/edge0-8b.md)
+- Rapport technique : [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](../paper/main.pdf))
+
+## Tests
+
+```bash
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q   # real-weight generation
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+```
diff --git a/python/README_ja.md b/python/README_ja.md
new file mode 100644
index 0000000..8f8fa79
--- /dev/null
+++ b/python/README_ja.md
@@ -0,0 +1,165 @@
+# edge0 — Python フレームワーク
+
+[English](README.md) | [中文](README_zh.md) | 日本語 | [Español](README_es.md) | [Français](README_fr.md)
+
+このドキュメントは、**edge0 Python フレームワークをソースからインストールして実行する開発者**向けのものです。内容は、クイックスタート(インストール → モデル → 実行)、実測のパフォーマンスと品質、そしてスタックの背後にある設計です。
+
+edge0 は**モノレポ**として公開されています — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — 最上位階層には共有ドキュメント(`docs/`)とプラットフォームのサブプロジェクト(`python/` = 本フレームワーク、ほかに `macos/`、`ios/`、`android/`、`windows/`)が含まれます。Python フレームワークは edge0 レシピ — **SSD expert offload + Recover-LoRA + prerouter ルーティング予測** — のリファレンス実装であり、MLX を介して Apple Silicon 上で大規模スパース MoE モデルを実行し、ピークメモリはパラメータ数ではなく*アクティブ*なエキスパート集合によって上限が決まります。
+
+サードパーティの帰属表示: リポジトリルートの [`NOTICE`](../NOTICE) を参照してください。モデルカード: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)。以下のすべてのコマンドはこのディレクトリ(`python/`)から実行します。
+
+---
+
+## クイックスタート
+
+### 前提条件
+
+| コンポーネント | 要件 |
+|---|---|
+| OS / ハードウェア | Apple Silicon(M1/M2/M3/M4)搭載の macOS — MLX バックエンドは Apple Silicon 専用;CUDA バックエンドはロードマップ上 |
+| Python | 3.10+(3.12 推奨) |
+| MLX | `mlx==0.30.6` / `mlx-metal==0.30.6` と `mlx-lm==0.31.0`(`pyproject.toml` を参照) |
+| メモリ | ピークアクティブ約 2.9 GB(edge0-35b)、約 1.0 GB(edge0-8b)、短いコンテキストの場合 |
+| ディスク | 約 23 GB(edge0-35b)/ 約 4.2 GB(edge0-8b);エキスパートの重みは mmap されてオンデマンドで読み取られ、事前に RAM へロードされない |
+
+> Apple A18 / A18 Pro で文字化けした言語混在の出力が発生する場合、古い
+> `mlx` が原因です: `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
+> ([#8](https://github.com/Edge0-AI/Edge0/issues/8))。
+
+### インストール
+
+```bash
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+### モデル
+
+両方のティアは Hugging Face と ModelScope で公開されています。各リポジトリは、ベースチェックポイントと訓練済み LoRA + prerouter アダプタを**1 つのディレクトリ**にまとめており、1 回のダウンロードですぐ実行できるモデルが揃います:
+
+```bash
+# the repo's helper (defaults to the two published tiers):
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b  --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview  --local-dir models/edge0-8b
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+```
+
+環境変数で edge0 にティアを指定するか、ディレクトリを直接渡します(ティアは `config.json` から自動検出されます):
+
+```bash
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+### 実行
+
+```bash
+edge0 demo edge0-8b                 # one-shot generation demo
+edge0 chat edge0-8b --prompt "Explain streaming inference in one sentence."
+edge0 serve edge0-8b                # OpenAI-compatible server on http://127.0.0.1:8000
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-8b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+```
+
+`python -m edge0 ...` は `edge0 ...` と等価です。
+
+### Python API
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`examples/demo.py` も同じ最小限のウォークスルーです(`edge0 demo` はまさにこのパスを実行します)。
+
+---
+
+## パフォーマンス
+
+`examples/bench.py` で測定(3.3k トークンのプロンプトのプリフィル → サンプリングによるウォームアップ 10 ステップ → 計測付きのサンプリングデコード 200 トークン、ティアごとに 2 回実行):
+
+| ティア | デコード速度 | プリフィルスループット(コールド / ウォーム)* | ピークアクティブメモリ | テストマシン |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*コールド = プロセス起動後の最初のリクエスト(エキスパートの重みが SSD からフォールトインする);ウォーム = 以降のリクエスト(ページキャッシュに常駐)。*
+
+```bash
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+```
+
+## 品質
+
+edge0 モデル(int4 + 訓練済みアダプタ + prerouter)とオリジナルの fp16 ベースの両方について、同一の設定で [OpenCompass](https://github.com/open-compass/opencompass) を使って実行しました。edge0 パイプラインの損失は小さく、**edge0-35b で平均 3.9 ポイント、edge0-8b で 2.8 ポイント**です(MMLU-Pro はベースを上回ります)。満点 100:
+
+| ベンチマーク | edge0-35b (int4) | Qwen3.6-35B-A3B (fp16) | edge0-8b (int4) | Ling 3.0 tiny (fp16) |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **平均** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+---
+
+## 技術詳細
+
+### コア機構
+
+- **SSD expert offload** — エキスパートの重みは必要に応じてストレージからストリーミングされます。ピークメモリはパラメータ数ではなく、アクティブな集合によって上限が決まります。
+- **Prerouter** — 訓練済みヘッドがエキスパートルーティングを 1 ステップ先に予測するため、エキスパートのロードがフォワードパスを停滞させるのではなくオーバーラップします(**最大 +59%** のデコードスループット。ストレージレイテンシ、モデルサイズ、ルーティング幅 *K* が増すほど効果は大きくなります)。
+- **Recover-LoRA** — int4 のベースは凍結され、FP 教師からの蒸留によって LoRA アダプタが訓練されることで、4-bit 量子化損失の大部分を回復します。アダプタはマージされません。1 つの読み取り専用ベースが複数のアダプタセットに対応します。
+
+### 設計
+
+- **transformers スタイルの使い方** — `AutoModel` / `AutoConfig` / `AutoEngine` がモデル名からティアを解決します。
+- **バックエンド分離** — すべての MLX コードは `src/edge0/backends/mlx/` 配下に置かれます。コアロジック(モデル仕様、prerouter、ストリーミングエキスパートプール、サーバー)はバックエンドファサード(`backends/base.py`)にのみ依存するため、新しいバックエンドは同じファサードを実装するだけで済み(`backends/cuda/` は予約スロット)、コアコードの変更はゼロです。
+- **safetensors 形式のアダプタ** — LoRA と prerouter の重みは、来歴メタデータ(source、version、owner layers)付きの `.safetensors` ファイルであり、モデルディレクトリまたは gitignore 済みの `artifacts/` フォールバックから解決されます。
+- **モデル + アダプタを 1 つのディレクトリに** — モデルディレクトリには、ベースチェックポイントとそのモデルのアダプタが格納されます。アダプタのアップグレードはアダプタファイルの交換のみ — ベースは読み取り専用のままで、決してマージされません。
+
+### パッケージレイアウト
+
+```
+src/edge0/
+├── backends/mlx/                  # MLX backend (isolation boundary; cuda/ reserved)
+├── engine/  models/  moe/  prerouter/  streaming/
+├── adapters/  attention/  server/
+└── cli.py  registry.py  sampling.py  config.py
+```
+
+---
+
+## ドキュメント
+
+- [アーキテクチャ](../docs/architecture.md)
+- [アテンション](../docs/attention.md) / [MoE](../docs/moe.md) / [SSD ストリーミング](../docs/streaming.md) / [prerouter](../docs/prerouter.md)
+- [モデルの追加](../docs/adding-a-model.md)
+- [edge0-35b](../docs/models/edge0-35b.md) / [edge0-8b](../docs/models/edge0-8b.md)
+- 技術レポート: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063)([PDF](../paper/main.pdf))
+
+## テスト
+
+```bash
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q   # real-weight generation
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+```
diff --git a/python/README_zh.md b/python/README_zh.md
new file mode 100644
index 0000000..85d620f
--- /dev/null
+++ b/python/README_zh.md
@@ -0,0 +1,197 @@
+# edge0 — Python 框架
+
+[English](README.md) | 中文 | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+本文档面向**从源码安装并运行 edge0 Python 框架的开发者**。内容包括：
+快速开始（安装 → 模型 → 运行）、性能与质量实测，以及技术栈背后的设计。
+
+edge0 以 **monorepo** 形式发布 ——
+[`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) —— 顶层目录存放共享
+文档（`docs/`）与各平台子项目（`python/` = 本框架，另有 `macos/`、`ios/`、
+`android/`、`windows/`）。Python 框架是 edge0 配方的参考实现 —— **SSD
+专家 offload + Recover-LoRA + prerouter 路由预判** —— 通过 MLX 在
+Apple Silicon 上运行大型稀疏 MoE 模型，峰值内存由*激活*专家集而非参数量
+决定。
+
+第三方声明：见仓库根目录的 [`NOTICE`](../NOTICE)。模型卡：
+[Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview)
+· [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)。
+下文所有命令均在本目录（`python/`）下运行。
+
+---
+
+## 快速开始
+
+### 环境要求
+
+| 组件 | 要求 |
+|---|---|
+| 系统 / 硬件 | Apple Silicon 上的 macOS（M1/M2/M3/M4）—— MLX 后端仅支持 Apple Silicon；CUDA 后端在路线图中 |
+| Python | 3.10+（推荐 3.12） |
+| MLX | `mlx==0.30.6` / `mlx-metal==0.30.6`，搭配 `mlx-lm==0.31.0`（见 `pyproject.toml`） |
+| 内存 | 短上下文下峰值激活内存约 2.9 GB（edge0-35b）/ 约 1.0 GB（edge0-8b） |
+| 磁盘 | 约 23 GB（edge0-35b）/ 约 4.2 GB（edge0-8b）；专家权重经 mmap 按需读取，不会一次性载入内存 |
+
+> Apple A18 / A18 Pro 上输出乱码、语言混杂，说明 `mlx` 版本过旧：
+> `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'`
+> （[#8](https://github.com/Edge0-AI/Edge0/issues/8)）。
+
+### 安装
+
+```bash
+python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]'
+```
+
+### 模型
+
+两个档位均发布在 Hugging Face 与 ModelScope；每个仓库把基模 checkpoint
+与训练好的 LoRA + prerouter 适配器打包在**同一目录**，一次下载即为
+可运行的模型：
+
+```bash
+# the repo's helper (defaults to the two published tiers):
+.venv/bin/python scripts/fetch_models.py --tier edge0-8b  --target-dir models
+.venv/bin/python scripts/fetch_models.py --tier edge0-35b --target-dir models
+
+# or directly with the CLI:
+.venv/bin/huggingface-cli download Edge0/Edge0-8B-A1B-preview  --local-dir models/edge0-8b
+.venv/bin/huggingface-cli download Edge0/Edge0-35B-A3B-preview --local-dir models/edge0-35b
+```
+
+通过环境变量让 edge0 指向某个档位，或直接传入目录（档位会从
+`config.json` 自动识别）：
+
+```bash
+export EDGE0_8B_MODEL=$PWD/models/edge0-8b
+```
+
+### 运行
+
+```bash
+edge0 demo edge0-8b                 # one-shot generation demo
+edge0 chat edge0-8b --prompt "Explain streaming inference in one sentence."
+edge0 serve edge0-8b                # OpenAI-compatible server on http://127.0.0.1:8000
+```
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"edge0-8b","messages":[{"role":"user","content":"Hello!"}],"max_tokens":32}'
+```
+
+`python -m edge0 ...` 等价于 `edge0 ...`。
+
+### Python API
+
+```python
+from edge0 import AutoEngine
+from edge0.server.chat import ChatMessage, ChatRequest, ChatSession
+
+engine = AutoEngine.from_pretrained("/path/to/model")  # tier auto-detected
+req = ChatRequest(
+    model=engine.name,
+    messages=[ChatMessage(role="user", content="Hello!")],
+    max_tokens=64,
+)
+tokens, meta = ChatSession(engine, req).run()
+print(engine._tok.decode(tokens))
+engine.close()   # release mmaps / expert cache
+```
+
+`examples/demo.py` 就是这条最小路径（`edge0 demo` 内部运行的正是这条
+路径）。
+
+---
+
+## 性能实测
+
+使用 `examples/bench.py` 实测（3.3k token prompt prefill → 10 步采样
+warmup → 200 token 计时解码段，每档 2 轮）：
+
+| 档位 | 解码速度 | Prefill 吞吐（冷/热）* | 峰值 active 内存 | 测试机器 |
+|---|---|---|---|---|
+| `edge0-35b` | 14.9–17.7 tok/s | 113 / 140 tok/s | 2.9 GiB | Mac mini M4 Pro, 24 GB |
+| `edge0-8b` | 23.9–25.3 tok/s | 500 / 1428 tok/s | 1.0 GiB | Mac mini M4 Pro, 24 GB |
+
+*冷 = 进程启动后首请求（专家权重从 SSD 逐页换入）；热 = 后续请求
+（页缓存常驻）。*
+
+```bash
+python examples/bench.py edge0-8b     # via $EDGE0_8B_MODEL
+python examples/bench.py edge0-35b    # via $EDGE0_35B_MODEL
+```
+
+## 质量
+
+使用 [OpenCompass](https://github.com/open-compass/opencompass) 在完全相同
+的设置下，对 edge0 模型（int4 + 训练适配器 + prerouter）与原 fp16 基座
+模型测得。edge0 管线的损失很小：**edge0-35b 平均仅落后 3.9 分、
+edge0-8b 落后 2.8 分**（MMLU-Pro 甚至反超基座）。满分 100：
+
+| 评测集 | edge0-35b（int4） | Qwen3.6-35B-A3B（fp16） | edge0-8b（int4） | Ling 3.0 tiny（fp16） |
+|---|---:|---:|---:|---:|
+| AIME 2026 | 86.6 | 92.7 | 63.3 | 73.3 |
+| HumanEval | 90.9 | 95.1 | 91.5 | 92.7 |
+| GPQA-Diamond | 79.8 | 81.8 | 70.7 | 71.2 |
+| MMLU-Pro | 81.0 | 84.6 | 70.1 | 65.8 |
+| IFBench | 57.9 | 61.7 | 53.9 | 60.6 |
+| **平均** | **79.2** | **83.2** | **69.9** | **72.7** |
+
+---
+
+## 技术细节
+
+### 核心机制
+
+- **SSD 专家 offload** —— 专家权重按需从存储流式加载；峰值内存由
+  激活集而非参数量决定。
+- **prerouter** —— 训练头提前一步预测专家路由，专家装载与前向计算重叠
+  而非阻塞（解码吞吐**最高 +59%**；收益随存储延迟、模型规模与路由宽度
+  *K* 增大）。
+- **Recover-LoRA** —— 冻结 int4 基模，用 FP teacher 蒸馏训练 LoRA
+  适配器，在 4bit 下恢复绝大部分量化损失。适配器保持不合并：一份只读
+  基模服务多套适配器。
+
+### 设计
+
+- **像 transformers 一样使用** —— `AutoModel` / `AutoConfig` /
+  `AutoEngine` 按模型名自动解析档位。
+- **后端隔离** —— 全部 MLX 代码收在 `src/edge0/backends/mlx/`；核心
+  逻辑（模型 spec、prerouter、流式专家池、server）只依赖后端门面
+  （`backends/base.py`），新增后端实现同一门面即可（`backends/cuda/`
+  为预留插槽），核心代码零改动。
+- **适配器统一为 safetensors** —— LoRA 与 prerouter 权重均为带来源
+  元数据（source、version、owner 层）的 `.safetensors` 文件，可从模型
+  目录或 gitignored 的 `artifacts/` 回退目录解析。
+- **模型 + 适配器同目录布局** —— 一个模型目录同时存放基模 checkpoint
+  与该模型的适配器；升级适配器只换适配器文件 —— 基模保持只读、永不
+  合并。
+
+### 包结构
+
+```
+src/edge0/
+├── backends/mlx/                  # MLX backend (isolation boundary; cuda/ reserved)
+├── engine/  models/  moe/  prerouter/  streaming/
+├── adapters/  attention/  server/
+└── cli.py  registry.py  sampling.py  config.py
+```
+
+---
+
+## 文档
+
+- [架构总览](../docs/architecture.md)
+- [注意力抽象](../docs/attention.md) / [MoE 抽象](../docs/moe.md) / [SSD 流式](../docs/streaming.md) / [prerouter](../docs/prerouter.md)
+- [如何接入新模型](../docs/adding-a-model.md)
+- [edge0-35b](../docs/models/edge0-35b.md) / [edge0-8b](../docs/models/edge0-8b.md)
+- 技术报告：[The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063)（[PDF](../paper/main.pdf)）
+
+## 测试
+
+```bash
+pytest                 # unit tests (no real weights)
+EDGE0_8B_MODEL=/path/to/edge0-8b pytest -m slow -q   # real-weight generation
+.venv/bin/python scripts/e2e_smoke.py \
+  --qwen-dir /path/to/edge0-35b --ling-dir /path/to/edge0-8b
+```
diff --git a/examples/bench.py b/python/examples/bench.py
similarity index 100%
rename from examples/bench.py
rename to python/examples/bench.py
diff --git a/examples/demo.py b/python/examples/demo.py
similarity index 81%
rename from examples/demo.py
rename to python/examples/demo.py
index 686700f..236e9b6 100644
--- a/examples/demo.py
+++ b/python/examples/demo.py
@@ -40,6 +40,18 @@ def main() -> int:
                          "already run that path)")
     args = ap.parse_args()
 
+    # ``--model <tier>`` resolves through the tier env var (same contract
+    # as ``edge0 demo``); ``--model-dir`` always wins.
+    if not args.model_dir and args.model:
+        import os
+
+        from edge0.cli import TIER_ENV
+        env = TIER_ENV.get(args.model)
+        args.model_dir = os.environ.get(env, "") if env else ""
+        if not args.model_dir:
+            ap.error(f"--model {args.model} needs ${env or '?'} set to the "
+                     f"checkpoint directory (or pass --model-dir)")
+
     kw = ({"prerouter": None} if args.no_prerouter
           else demo_kwargs(args.model_dir, args.model))
     engine = AutoEngine.from_pretrained(args.model_dir, name=args.model, **kw)
diff --git a/pyproject.toml b/python/pyproject.toml
similarity index 100%
rename from pyproject.toml
rename to python/pyproject.toml
diff --git a/scripts/e2e_smoke.py b/python/scripts/e2e_smoke.py
similarity index 100%
rename from scripts/e2e_smoke.py
rename to python/scripts/e2e_smoke.py
diff --git a/scripts/fetch_models.py b/python/scripts/fetch_models.py
similarity index 100%
rename from scripts/fetch_models.py
rename to python/scripts/fetch_models.py
diff --git a/scripts/generate_example.py b/python/scripts/generate_example.py
similarity index 100%
rename from scripts/generate_example.py
rename to python/scripts/generate_example.py
diff --git a/scripts/strip_vision_weights.py b/python/scripts/strip_vision_weights.py
similarity index 100%
rename from scripts/strip_vision_weights.py
rename to python/scripts/strip_vision_weights.py
diff --git a/scripts/upload_hf.py b/python/scripts/upload_hf.py
similarity index 100%
rename from scripts/upload_hf.py
rename to python/scripts/upload_hf.py
diff --git a/scripts/verify_alignment.py b/python/scripts/verify_alignment.py
similarity index 100%
rename from scripts/verify_alignment.py
rename to python/scripts/verify_alignment.py
diff --git a/src/edge0/__init__.py b/python/src/edge0/__init__.py
similarity index 100%
rename from src/edge0/__init__.py
rename to python/src/edge0/__init__.py
diff --git a/src/edge0/__main__.py b/python/src/edge0/__main__.py
similarity index 100%
rename from src/edge0/__main__.py
rename to python/src/edge0/__main__.py
diff --git a/src/edge0/adapters/__init__.py b/python/src/edge0/adapters/__init__.py
similarity index 100%
rename from src/edge0/adapters/__init__.py
rename to python/src/edge0/adapters/__init__.py
diff --git a/src/edge0/adapters/lora.py b/python/src/edge0/adapters/lora.py
similarity index 96%
rename from src/edge0/adapters/lora.py
rename to python/src/edge0/adapters/lora.py
index b13bf32..1103dd8 100644
--- a/src/edge0/adapters/lora.py
+++ b/python/src/edge0/adapters/lora.py
@@ -16,9 +16,7 @@
 
 where ``<target>`` is the module path inside the loaded model (e.g.
 ``model.layers.3.attention.q_proj`` or
-``language_model.model.layers.0.linear_attn.in_proj_qkv``).  The legacy
-npz files used the same A/B naming, so conversion is a pure format
-migration (see ``scripts/convert_adapters_legacy.py``).
+``language_model.model.layers.0.linear_attn.in_proj_qkv``).
 """
 
 from __future__ import annotations
diff --git a/src/edge0/attention/spec.py b/python/src/edge0/attention/spec.py
similarity index 100%
rename from src/edge0/attention/spec.py
rename to python/src/edge0/attention/spec.py
diff --git a/src/edge0/backends/__init__.py b/python/src/edge0/backends/__init__.py
similarity index 100%
rename from src/edge0/backends/__init__.py
rename to python/src/edge0/backends/__init__.py
diff --git a/src/edge0/backends/base.py b/python/src/edge0/backends/base.py
similarity index 100%
rename from src/edge0/backends/base.py
rename to python/src/edge0/backends/base.py
diff --git a/src/edge0/backends/mlx/__init__.py b/python/src/edge0/backends/mlx/__init__.py
similarity index 100%
rename from src/edge0/backends/mlx/__init__.py
rename to python/src/edge0/backends/mlx/__init__.py
diff --git a/src/edge0/backends/mlx/_impl/__init__.py b/python/src/edge0/backends/mlx/_impl/__init__.py
similarity index 100%
rename from src/edge0/backends/mlx/_impl/__init__.py
rename to python/src/edge0/backends/mlx/_impl/__init__.py
diff --git a/src/edge0/backends/mlx/_impl/bailing_hybrid.py b/python/src/edge0/backends/mlx/_impl/bailing_hybrid.py
similarity index 100%
rename from src/edge0/backends/mlx/_impl/bailing_hybrid.py
rename to python/src/edge0/backends/mlx/_impl/bailing_hybrid.py
diff --git a/src/edge0/backends/mlx/_impl/qwen3_5.py b/python/src/edge0/backends/mlx/_impl/qwen3_5.py
similarity index 100%
rename from src/edge0/backends/mlx/_impl/qwen3_5.py
rename to python/src/edge0/backends/mlx/_impl/qwen3_5.py
diff --git a/src/edge0/backends/mlx/_impl/qwen3_5_moe.py b/python/src/edge0/backends/mlx/_impl/qwen3_5_moe.py
similarity index 100%
rename from src/edge0/backends/mlx/_impl/qwen3_5_moe.py
rename to python/src/edge0/backends/mlx/_impl/qwen3_5_moe.py
diff --git a/src/edge0/backends/mlx/_impl/qwen3_next.py b/python/src/edge0/backends/mlx/_impl/qwen3_next.py
similarity index 100%
rename from src/edge0/backends/mlx/_impl/qwen3_next.py
rename to python/src/edge0/backends/mlx/_impl/qwen3_next.py
diff --git a/src/edge0/backends/mlx/backend.py b/python/src/edge0/backends/mlx/backend.py
similarity index 100%
rename from src/edge0/backends/mlx/backend.py
rename to python/src/edge0/backends/mlx/backend.py
diff --git a/src/edge0/backends/mlx/io.py b/python/src/edge0/backends/mlx/io.py
similarity index 100%
rename from src/edge0/backends/mlx/io.py
rename to python/src/edge0/backends/mlx/io.py
diff --git a/src/edge0/backends/mlx/quant.py b/python/src/edge0/backends/mlx/quant.py
similarity index 100%
rename from src/edge0/backends/mlx/quant.py
rename to python/src/edge0/backends/mlx/quant.py
diff --git a/src/edge0/cli.py b/python/src/edge0/cli.py
similarity index 94%
rename from src/edge0/cli.py
rename to python/src/edge0/cli.py
index 97555d8..097442c 100644
--- a/src/edge0/cli.py
+++ b/python/src/edge0/cli.py
@@ -6,7 +6,6 @@
     edge0 serve       start the HTTP server (one model, queued generations)
     edge0 chat        one-shot prompt -> answer on the terminal
     edge0 models      list registered tiers and their default profiles
-    edge0 convert-adapters   one-shot legacy npz -> safetensors migration
 """
 
 from __future__ import annotations
@@ -208,18 +207,6 @@ def cmd_serve(args) -> int:
     return 0
 
 
-def cmd_convert(args) -> int:
-    import runpy
-    sys.path.insert(0, str(__import__("pathlib").Path(__file__).resolve()
-                           .parents[1] / "scripts"))
-    runpy.run_path(
-        str(__import__("pathlib").Path(__file__).resolve().parents[1]
-            / "scripts" / "convert_adapters_legacy.py"),
-        run_name="__main__",
-    )
-    return 0
-
-
 def _add_engine_flags(p) -> None:
     """Engine-construction flags shared by demo / chat / serve."""
     p.add_argument("--no-prerouter", action="store_true")
@@ -288,10 +275,6 @@ def _build_parser() -> argparse.ArgumentParser:
     _add_engine_flags(p)
     p.set_defaults(fn=cmd_serve)
 
-    p = sub.add_parser("convert-adapters",
-                       help="one-shot legacy npz -> safetensors migration")
-    p.set_defaults(fn=cmd_convert)
-
     return ap
 
 
diff --git a/src/edge0/config.py b/python/src/edge0/config.py
similarity index 100%
rename from src/edge0/config.py
rename to python/src/edge0/config.py
diff --git a/src/edge0/engine/__init__.py b/python/src/edge0/engine/__init__.py
similarity index 100%
rename from src/edge0/engine/__init__.py
rename to python/src/edge0/engine/__init__.py
diff --git a/src/edge0/engine/base.py b/python/src/edge0/engine/base.py
similarity index 100%
rename from src/edge0/engine/base.py
rename to python/src/edge0/engine/base.py
diff --git a/src/edge0/engine/hooks.py b/python/src/edge0/engine/hooks.py
similarity index 100%
rename from src/edge0/engine/hooks.py
rename to python/src/edge0/engine/hooks.py
diff --git a/src/edge0/engine/ling.py b/python/src/edge0/engine/ling.py
similarity index 100%
rename from src/edge0/engine/ling.py
rename to python/src/edge0/engine/ling.py
diff --git a/src/edge0/engine/qwen.py b/python/src/edge0/engine/qwen.py
similarity index 100%
rename from src/edge0/engine/qwen.py
rename to python/src/edge0/engine/qwen.py
diff --git a/src/edge0/models/__init__.py b/python/src/edge0/models/__init__.py
similarity index 100%
rename from src/edge0/models/__init__.py
rename to python/src/edge0/models/__init__.py
diff --git a/src/edge0/models/base.py b/python/src/edge0/models/base.py
similarity index 95%
rename from src/edge0/models/base.py
rename to python/src/edge0/models/base.py
index d3dead5..d6ede85 100644
--- a/src/edge0/models/base.py
+++ b/python/src/edge0/models/base.py
@@ -19,16 +19,16 @@
 from edge0.prerouter.spec import PrerouterSpec
 from edge0.streaming.options import LayerOptions
 
-# Repository-root artifacts (gitignored; produced once by
-# scripts/convert_adapters_legacy.py from the training npz exports, then
-# the source npz are discarded).
+# Project-root artifacts (gitignored): an optional local cache of adapter
+# safetensors. Adapters normally ship inside the model directory; this is
+# only the fallback location artifact() returns when they are not there.
 ARTIFACTS_DIR = Path(__file__).resolve().parents[3] / "artifacts"
 
 
 def artifact(name: str, model_dir: str | None = None) -> str:
     """Path of one adapter file: the model's own directory first
     (model and its adapters side by side in one
-    directory), falling back to ``artifacts/`` (conversion cache)."""
+    directory), falling back to ``artifacts/`` (optional local cache)."""
     if model_dir:
         cand = Path(model_dir) / name
         if cand.is_file():
diff --git a/src/edge0/models/edge0_35b/__init__.py b/python/src/edge0/models/edge0_35b/__init__.py
similarity index 97%
rename from src/edge0/models/edge0_35b/__init__.py
rename to python/src/edge0/models/edge0_35b/__init__.py
index 3f27819..0e5bd40 100644
--- a/src/edge0/models/edge0_35b/__init__.py
+++ b/python/src/edge0/models/edge0_35b/__init__.py
@@ -10,7 +10,8 @@
 * trained prerouter: 33 heads (owners 6..38), start_layer 7,
   hidden 512, fp16, consumed directly by the patched MoE block at
   decode (patch_call=True, cross-token staged decode, K=4).
-* serving: port 8085; acceptance ≈13 tok/s, peak active ≈3.3 GB.
+* serving: port 8000 (the shared CLI default); acceptance ≈13 tok/s,
+  peak active ≈3.3 GB.
 """
 
 from __future__ import annotations
@@ -64,7 +65,6 @@ def _defaults(cls, model_dir: str) -> "Qwen35Config":
             hot_window=4,
             intra_staging=False,
             prefetch_history=True,
-            port=8085,
             target_tok_s=13.0,
             peak_active_mem_mb=3400.0,
         )
diff --git a/src/edge0/models/edge0_8b/__init__.py b/python/src/edge0/models/edge0_8b/__init__.py
similarity index 97%
rename from src/edge0/models/edge0_8b/__init__.py
rename to python/src/edge0/models/edge0_8b/__init__.py
index 7fcc724..5453580 100644
--- a/src/edge0/models/edge0_8b/__init__.py
+++ b/python/src/edge0/models/edge0_8b/__init__.py
@@ -11,7 +11,8 @@
   (explicit owners 1..22), start_layer 1, hidden 512, fp16,
   feature_topk "executed", consumed inside ``BailingSparseMoE``
   from ``prerouter_cache`` logits (patch_call=False).
-* serving: port 8083; acceptance ≈33 tok/s, peak active ≈1.4 GB.
+* serving: port 8000 (the shared CLI default); acceptance ≈33 tok/s,
+  peak active ≈1.4 GB.
 """
 
 from __future__ import annotations
@@ -72,7 +73,6 @@ def _defaults(cls, model_dir: str) -> "Ling8BConfig":
             hot_window=1,
             intra_staging=False,
             prefetch_history=True,
-            port=8083,
             target_tok_s=33.0,
             peak_active_mem_mb=1400.0,
         )
diff --git a/src/edge0/moe/__init__.py b/python/src/edge0/moe/__init__.py
similarity index 100%
rename from src/edge0/moe/__init__.py
rename to python/src/edge0/moe/__init__.py
diff --git a/src/edge0/moe/routing.py b/python/src/edge0/moe/routing.py
similarity index 100%
rename from src/edge0/moe/routing.py
rename to python/src/edge0/moe/routing.py
diff --git a/src/edge0/moe/spec.py b/python/src/edge0/moe/spec.py
similarity index 100%
rename from src/edge0/moe/spec.py
rename to python/src/edge0/moe/spec.py
diff --git a/src/edge0/prerouter/__init__.py b/python/src/edge0/prerouter/__init__.py
similarity index 100%
rename from src/edge0/prerouter/__init__.py
rename to python/src/edge0/prerouter/__init__.py
diff --git a/src/edge0/prerouter/heads.py b/python/src/edge0/prerouter/heads.py
similarity index 100%
rename from src/edge0/prerouter/heads.py
rename to python/src/edge0/prerouter/heads.py
diff --git a/src/edge0/prerouter/install.py b/python/src/edge0/prerouter/install.py
similarity index 100%
rename from src/edge0/prerouter/install.py
rename to python/src/edge0/prerouter/install.py
diff --git a/src/edge0/prerouter/spec.py b/python/src/edge0/prerouter/spec.py
similarity index 100%
rename from src/edge0/prerouter/spec.py
rename to python/src/edge0/prerouter/spec.py
diff --git a/src/edge0/prerouter/stager.py b/python/src/edge0/prerouter/stager.py
similarity index 100%
rename from src/edge0/prerouter/stager.py
rename to python/src/edge0/prerouter/stager.py
diff --git a/src/edge0/prerouter/state.py b/python/src/edge0/prerouter/state.py
similarity index 100%
rename from src/edge0/prerouter/state.py
rename to python/src/edge0/prerouter/state.py
diff --git a/src/edge0/registry.py b/python/src/edge0/registry.py
similarity index 100%
rename from src/edge0/registry.py
rename to python/src/edge0/registry.py
diff --git a/src/edge0/sampling.py b/python/src/edge0/sampling.py
similarity index 100%
rename from src/edge0/sampling.py
rename to python/src/edge0/sampling.py
diff --git a/src/edge0/server/__init__.py b/python/src/edge0/server/__init__.py
similarity index 100%
rename from src/edge0/server/__init__.py
rename to python/src/edge0/server/__init__.py
diff --git a/src/edge0/server/app.py b/python/src/edge0/server/app.py
similarity index 100%
rename from src/edge0/server/app.py
rename to python/src/edge0/server/app.py
diff --git a/src/edge0/server/chat.py b/python/src/edge0/server/chat.py
similarity index 100%
rename from src/edge0/server/chat.py
rename to python/src/edge0/server/chat.py
diff --git a/src/edge0/streaming/__init__.py b/python/src/edge0/streaming/__init__.py
similarity index 100%
rename from src/edge0/streaming/__init__.py
rename to python/src/edge0/streaming/__init__.py
diff --git a/src/edge0/streaming/cache.py b/python/src/edge0/streaming/cache.py
similarity index 100%
rename from src/edge0/streaming/cache.py
rename to python/src/edge0/streaming/cache.py
diff --git a/src/edge0/streaming/install.py b/python/src/edge0/streaming/install.py
similarity index 100%
rename from src/edge0/streaming/install.py
rename to python/src/edge0/streaming/install.py
diff --git a/src/edge0/streaming/layer.py b/python/src/edge0/streaming/layer.py
similarity index 100%
rename from src/edge0/streaming/layer.py
rename to python/src/edge0/streaming/layer.py
diff --git a/src/edge0/streaming/mmap.py b/python/src/edge0/streaming/mmap.py
similarity index 100%
rename from src/edge0/streaming/mmap.py
rename to python/src/edge0/streaming/mmap.py
diff --git a/src/edge0/streaming/options.py b/python/src/edge0/streaming/options.py
similarity index 100%
rename from src/edge0/streaming/options.py
rename to python/src/edge0/streaming/options.py
diff --git a/tests/test_cli_flags.py b/python/tests/test_cli_flags.py
similarity index 100%
rename from tests/test_cli_flags.py
rename to python/tests/test_cli_flags.py
diff --git a/tests/test_e2e_slow.py b/python/tests/test_e2e_slow.py
similarity index 100%
rename from tests/test_e2e_slow.py
rename to python/tests/test_e2e_slow.py
diff --git a/tests/test_moe_spec.py b/python/tests/test_moe_spec.py
similarity index 100%
rename from tests/test_moe_spec.py
rename to python/tests/test_moe_spec.py
diff --git a/tests/test_prefill_hook.py b/python/tests/test_prefill_hook.py
similarity index 100%
rename from tests/test_prefill_hook.py
rename to python/tests/test_prefill_hook.py
diff --git a/tests/test_registry.py b/python/tests/test_registry.py
similarity index 97%
rename from tests/test_registry.py
rename to python/tests/test_registry.py
index 9cf7932..d3a61d9 100644
--- a/tests/test_registry.py
+++ b/python/tests/test_registry.py
@@ -65,7 +65,7 @@ def test_qwen35_profile():
         name="edge0-35b", history_slots=True).history_slots is True
     assert cfg.gen.temperature == 0.7
     assert 248046 in cfg.gen.eos_ids
-    assert cfg.port == 8085
+    assert cfg.port == 8000  # unified serving port (CLI default)
 
 
 def test_ling8b_profile():
@@ -82,7 +82,7 @@ def test_ling8b_profile():
     assert cfg.prerouter.patch_call is False
     assert cfg.gen.repetition_penalty == 1.1
     assert 156895 in cfg.gen.eos_ids
-    assert cfg.port == 8083
+    assert cfg.port == 8000  # unified serving port (CLI default)
 
 
 def test_override_and_reject():
diff --git a/tests/test_repo_hygiene.py b/python/tests/test_repo_hygiene.py
similarity index 68%
rename from tests/test_repo_hygiene.py
rename to python/tests/test_repo_hygiene.py
index c8ef1d0..9ea9b7c 100644
--- a/tests/test_repo_hygiene.py
+++ b/python/tests/test_repo_hygiene.py
@@ -15,14 +15,18 @@
 
 import pathlib
 import re
+import subprocess
 
-ROOT = pathlib.Path(__file__).resolve().parents[1]
+# Repo root (this file lives at <root>/python/tests/): the hygiene guards
+# scan the whole multi-platform repo, not just the Python subproject.
+ROOT = pathlib.Path(__file__).resolve().parents[2]
 
 _EXCLUDE_DIRS = {".git", ".venv", ".pytest_cache", "__pycache__", ".agents",
-                 ".codex"}
+                 ".codex", "vendor", "wt", ".edge0", "models", "node_modules",
+                 ".gradle", "build-dl", "target", "dist", "build"}
 _EXCLUDE_PARTS = {"egg-info"}
 _EXCLUDE_SUFFIXES = {".pyc", ".so", ".dylib", ".bin", ".safetensors", ".npz",
-                     ".png", ".jpg", ".jpeg", ".gif", ".ico"}
+                     ".png", ".jpg", ".jpeg", ".gif", ".ico", ".pdf"}
 
 # NOTE: patterns are assembled with + so this file never contains the
 # literal string it is looking for.
@@ -51,7 +55,35 @@
 ]
 
 
+def _git_ls_files() -> list[str] | None:
+    """Repo-relative paths of the files that are (or would be) committed:
+    tracked files plus untracked-but-not-gitignored ones.  Materialized
+    build areas (``vendor/``, ``wt/``, ``models/``, ``.venv/``) are
+    gitignored, so this naturally excludes them — the hygiene guards check
+    *shipped* code, not downloaded third-party dependencies.  Returns None
+    if git is unavailable, so the caller can fall back to a directory walk.
+    """
+    try:
+        out = subprocess.run(
+            ["git", "ls-files", "--cached", "--others", "--exclude-standard",
+             "-z"],
+            cwd=ROOT, capture_output=True, check=True,
+        ).stdout
+    except (OSError, subprocess.CalledProcessError):
+        return None
+    return [p for p in out.decode("utf-8", "replace").split("\0") if p]
+
+
 def _iter_text_files():
+    shipped = _git_ls_files()
+    if shipped is not None:
+        for rel in shipped:
+            path = ROOT / rel
+            if path.suffix in _EXCLUDE_SUFFIXES or not path.is_file():
+                continue
+            yield path
+        return
+    # Fallback (no git): walk the tree, skipping vcs / build / vendor dirs.
     for path in ROOT.rglob("*"):
         rel = path.relative_to(ROOT)
         if any(part in _EXCLUDE_DIRS for part in rel.parts):
@@ -85,7 +117,7 @@ def test_no_hardcoded_local_paths():
 
 def test_mlx_imports_stay_inside_the_backend():
     bad = []
-    for path in (ROOT / "src" / "edge0").rglob("*.py"):
+    for path in (ROOT / "python" / "src" / "edge0").rglob("*.py"):
         rel = path.relative_to(ROOT)
         text = path.read_text(encoding="utf-8")
         if _MLX_IMPORT.search(text) and "backends/mlx" not in rel.as_posix():
diff --git a/tests/test_sampling.py b/python/tests/test_sampling.py
similarity index 100%
rename from tests/test_sampling.py
rename to python/tests/test_sampling.py
diff --git a/tests/test_server.py b/python/tests/test_server.py
similarity index 100%
rename from tests/test_server.py
rename to python/tests/test_server.py
diff --git a/tests/test_streaming_math.py b/python/tests/test_streaming_math.py
similarity index 100%
rename from tests/test_streaming_math.py
rename to python/tests/test_streaming_math.py
diff --git a/scripts/convert_adapters_legacy.py b/scripts/convert_adapters_legacy.py
deleted file mode 100644
index 09c1b96..0000000
--- a/scripts/convert_adapters_legacy.py
+++ /dev/null
@@ -1,159 +0,0 @@
-#!/usr/bin/env python3
-"""One-shot legacy adapter converter: npz -> safetensors.
-
-The training pipelines exported adapters as ``.npz`` (qwen round-7 and
-the training side).  edge0 consumes safetensors ONLY — this script converts each
-legacy pair once into ``artifacts/`` (gitignored), writing the
-conversion provenance (model, kind, K, r/alpha, source md5, format
-version) into the safetensors ``__metadata__``.  The npz sources are
-NOT deleted (see the conversion note in README); they are only needed
-to regenerate the artifacts.
-
-Usage:
-    python scripts/convert_adapters_legacy.py [--force]
-
-Idempotent: existing artifacts are skipped unless ``--force``.
-"""
-
-from __future__ import annotations
-
-import argparse
-import hashlib
-import os
-import json
-import sys
-from datetime import datetime, timezone
-from pathlib import Path
-
-import numpy as np
-
-ROOT = Path(__file__).resolve().parents[1]
-ARTIFACTS = ROOT / "artifacts"
-
-# (npz path, tier, kind) — every legacy source this script understands.
-# Training-side npz exports discovered per machine; pass --npz-dir (or
-# edit this list) to point at your own exports.
-NPZ_DIRS = [
-    Path(os.environ.get("EDGE0_NPZ_DIR", "/path/to/npz/exports")),
-]
-SOURCES = [
-    ("prerouter_<family>.npz", "edge0-35b", "prerouter"),
-    ("lora_<family>.npz", "edge0-35b", "lora"),
-    ("prerouter_<family>.npz", "edge0-8b", "prerouter"),
-    ("lora_<family>.npz", "edge0-8b", "lora"),
-]
-
-
-def _find_source(name: str) -> Path | None:
-    for d in NPZ_DIRS:
-        p = d / name
-        if p.is_file():
-            return p
-    return None
-
-
-def md5(path: Path) -> str:
-    h = hashlib.md5()
-    with open(path, "rb") as f:
-        for chunk in iter(lambda: f.read(1 << 20), b""):
-            h.update(chunk)
-    return h.hexdigest()
-
-
-def rewrite_key(key: str, tier: str, kind: str) -> str:
-    """Map a legacy npz key to the edge0 artifact key.
-
-    prerouter heads are normalized to ``layers.<N>.<part>.weight``;
-    qwen lora gains the ``language_model.`` prefix (the checkpoint
-    namespace the adapter is applied into).
-    """
-    if kind == "prerouter":
-        # Some families use layers.N.*; others namespace the layer index
-        # with a prefix segment (<ns>.<N>.<part>.weight).  Both are
-        # normalized to layers.<N>.<part>.weight.
-        parts = key.split(".")
-        if (len(parts) >= 4 and parts[0] != "layers"
-                and parts[1].isdigit()):
-            return f"layers.{parts[1]}.{parts[2]}.{parts[3]}"
-        return key
-    if kind == "lora" and tier == "edge0-35b" and key.startswith("model.layers."):
-        return "language_model." + key
-    return key
-
-
-def convert(source: Path, tier: str, kind: str, force: bool = False) -> Path:
-    name = {
-        ("edge0-35b", "prerouter"): "prerouter_edge0_35b_k4.safetensors",
-        ("edge0-35b", "lora"): "lora_edge0_35b_k4.safetensors",
-        ("edge0-8b", "prerouter"): "prerouter_edge0_8b.safetensors",
-        ("edge0-8b", "lora"): "lora_edge0_8b.safetensors",
-    }[(tier, kind)]
-    out = ARTIFACTS / name
-    if out.exists() and not force:
-        print(f"skip  {out.name} (exists)")
-        return out
-
-    data = np.load(source)
-    tensors = {}
-    for key in data.files:
-        tensors[rewrite_key(key, tier, kind)] = data[key]
-
-    # provenance: rank from the A matrix, alpha from the training run
-    # (32 for both shipped pairs), K from the tier profile.
-    meta = {
-        "model": tier,
-        "kind": kind,
-        "K": "4" if tier == "edge0-35b" else "8",
-        "r": "16", "alpha": "32",
-        "source": str(source),
-        "source_md5": md5(source),
-        "converted": datetime.now(timezone.utc).isoformat(timespec="seconds"),
-        "format_version": "1",
-    }
-    if kind == "prerouter":
-        owners = sorted({int(k.split(".")[1])
-                         for k in tensors if k.startswith("layers.")})
-        meta["owners"] = json.dumps(owners)
-
-    from safetensors.numpy import save_file as save_np
-    save_np(tensors, str(out), metadata={"__metadata__": json.dumps(meta)})
-    print(f"write {out.name} ({len(tensors)} tensors, "
-          f"{out.stat().st_size >> 20} MiB)")
-    return out
-
-
-def main(argv: list[str] | None = None) -> int:
-    ap = argparse.ArgumentParser(description=__doc__)
-    ap.add_argument("--force", action="store_true",
-                    help="reconvert even if the artifact exists")
-    ap.add_argument("--npz-dir", default=None,
-                    help="directory holding the legacy npz exports "
-                         "(default: $EDGE0_NPZ_DIR)")
-    args = ap.parse_args(argv)
-
-    if args.npz_dir:
-        NPZ_DIRS.insert(0, Path(args.npz_dir))
-    ARTIFACTS.mkdir(parents=True, exist_ok=True)
-    missing = []
-    resolved = []
-    for name, tier, kind in SOURCES:
-        p = _find_source(name)
-        if p is None:
-            missing.append(name)
-        else:
-            resolved.append((p, tier, kind))
-    if missing:
-        print("missing legacy sources (skip their conversion):")
-        for m in missing:
-            print("  -", m)
-    for path, tier, kind in resolved:
-        try:
-            convert(path, tier, kind, force=args.force)
-        except Exception as exc:  # noqa: BLE001 — report and continue
-            print(f"FAIL {path.name}: {exc}", file=sys.stderr)
-            return 1
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
diff --git a/scripts/patch-green.ps1 b/scripts/patch-green.ps1
new file mode 100644
index 0000000..c2744f0
--- /dev/null
+++ b/scripts/patch-green.ps1
@@ -0,0 +1,60 @@
+# edge0/scripts/patch-green.ps1 — patch-band `git am` gate (final arbiter after an upstream refresh; seconds, no compile)
+# The vendor tree is submodule-free: materialized on demand by cloning upstream at the pinned
+# commit from vendor.llama.pin (repo root). Set EDGE0_LLAMA_URL to clone from a mirror.
+# Usage: pwsh -File patch-green.ps1 [-Bands common,windows] [-Pin <sha>] [-Tree <expected-tree-hash>] [-Keep]
+#        (-Pin defaults to vendor.llama.pin)
+# Exit codes: 0 = GREEN (all bands replayed clean [+ tree parity]), 1 = RED.
+param(
+    [string]$Bands = "common,windows",
+    [string]$Pin   = "",
+    [string]$Tree  = "",
+    [switch]$Keep
+)
+$ErrorActionPreference = "Stop"
+$DEPOT  = Split-Path -Parent (Split-Path -Parent $MyInvocation.MyCommand.Path)   # = edge0 root
+$VENDOR = Join-Path $DEPOT "vendor/llama.cpp"
+$WT     = Join-Path $DEPOT "wt/greencheck"
+if (-not $Pin) {
+    if (-not (Test-Path (Join-Path $DEPOT "vendor.llama.pin"))) { throw "vendor.llama.pin missing at depot root (or pass -Pin)" }
+    $Pin = (Get-Content (Join-Path $DEPOT "vendor.llama.pin") -Raw).Trim().Split(" ")[0]
+}
+if (-not (Test-Path (Join-Path $VENDOR ".git"))) {
+    Write-Host "[gate] materializing vendor: clone llama.cpp at pinned $Pin -> $VENDOR" -ForegroundColor Cyan
+    New-Item -ItemType Directory -Force (Split-Path $VENDOR) | Out-Null
+    $LLAMA_URL = if ($env:EDGE0_LLAMA_URL) { $env:EDGE0_LLAMA_URL } else { "https://github.com/ggml-org/llama.cpp" }
+    git clone --filter=blob:none $LLAMA_URL $VENDOR
+    if ($LASTEXITCODE -ne 0) { throw "clone failed ($LLAMA_URL)" }
+    git -C $VENDOR checkout --detach $Pin
+    if ($LASTEXITCODE -ne 0) { throw "pinned commit $Pin not found" }
+}
+
+Push-Location $VENDOR
+try {
+    if (Test-Path $WT) { Pop-Location; git -C $VENDOR worktree remove --force $WT | Out-Null; Push-Location $VENDOR }
+    git worktree add --detach $WT $Pin | Out-Null
+    Push-Location $WT
+    $n = 0; $total = 0
+    foreach ($band in $Bands.Split(",")) {
+        $files = Get-ChildItem (Join-Path $DEPOT "patches/llama.cpp/$band/*.patch") | Sort-Object Name
+        foreach ($f in $files) {
+            $total++
+            git am --3way $f.FullName *> $null
+            if ($LASTEXITCODE -ne 0) {
+                Write-Host "RED: git am failed at $band\$($f.Name)" -ForegroundColor Red
+                git am --abort *> $null
+                exit 1
+            }
+            $n++
+        }
+    }
+    $actual = (git rev-parse "HEAD^{tree}").Trim()
+    Write-Host "applied=$n/$total  tree=$actual"
+    if ($Tree -and $actual -ne $Tree) {
+        Write-Host "RED: tree parity failed -- expected $Tree, got $actual" -ForegroundColor Red
+        exit 1
+    }
+    Write-Host "GREEN: band replay clean$(if ($Tree) { ' + tree parity' })" -ForegroundColor Green
+} finally {
+    Pop-Location
+    if (-not $Keep) { git -C $VENDOR worktree remove --force $WT *> $null }
+}
diff --git a/vendor.llama.pin b/vendor.llama.pin
new file mode 100644
index 0000000..145981a
--- /dev/null
+++ b/vendor.llama.pin
@@ -0,0 +1 @@
+7ab4ee7baad2d920464cbacfad4f4b07cf111fd2
diff --git a/windows/.gitignore b/windows/.gitignore
new file mode 100644
index 0000000..abddc33
--- /dev/null
+++ b/windows/.gitignore
@@ -0,0 +1,9 @@
+app/node_modules/
+app/dist/
+app/src-tauri/target/
+app/src-tauri/gen/
+.edge0/
+wt/
+/models/
+benchmarks/
+*.log
diff --git a/windows/NOTICE b/windows/NOTICE
new file mode 100644
index 0000000..7b9933f
--- /dev/null
+++ b/windows/NOTICE
@@ -0,0 +1,21 @@
+NOTICE — edge0 / windows (release tree)
+========================================
+
+This file records third-party material contained in or referenced by this subproject.
+Licensing of the distribution is governed by the LICENSE file at the monorepo root.
+
+1. ggml-org/llama.cpp (MIT) — https://github.com/ggml-org/llama.cpp
+   The engine is the pristine upstream checkout pinned at commit 7ab4ee7 (tag b11100),
+   located at the monorepo root vendor/llama.cpp. All modifications ship as replayable
+   patch bands under patches/llama.cpp/ (never committed into the vendor tree).
+
+2. Rust crates (see app/src-tauri/Cargo.lock for exact versions/licenses):
+   tauri 2, serde, reqwest (rustls), sha2, hex, windows-sys — each under its own license.
+
+3. JS/npm packages (see app/package.json + package-lock.json): React, Vite, Tailwind,
+   streamdown/markdown rendering stack — each under its own license.
+
+4. Model weights (downloaded at runtime, not part of this repository):
+   Edge0/Edge0-8B-A1B-preview and Edge0/Edge0-35B-A3B-preview — published in MLX format,
+   converted on-device to GGUF Q4_1 by the tools/ converter chain. Model-card license
+   applies to the weights and applies to the conversion outputs unchanged.
diff --git a/windows/README.md b/windows/README.md
new file mode 100644
index 0000000..3eda58a
--- /dev/null
+++ b/windows/README.md
@@ -0,0 +1,184 @@
+# edge0-windows — Release Build Guide
+
+English | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+This document is for **developers building or evaluating the Windows app from source**. It covers: quick start (build → run → test), measured performance on the reference machine, and the technical choices behind the stack.
+
+edge0 is published as a **monorepo** — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — whose top level holds the shared engine supply (`vendor.llama.pin` + the scripts-materialized `vendor/llama.cpp`, `patches/llama.cpp/` band-sets) and the platform subprojects (`windows/` = this app, `android/` = the companion). Inference runs on pinned upstream llama.cpp, patched as a replayable patch set, with dense compute on Vulkan and MoE expert weights served from CPU.
+
+Third-party attribution: see `NOTICE` in this directory. Model cards: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+
+---
+
+## 1. Quick Start
+
+### 1.1 Prerequisites
+
+| toolchain | requirement |
+|---|---|
+| OS | Windows 10 / 11 x64 |
+| Visual Studio 2022 | C++ workload (MSVC) |
+| CMake | ≥ 3.21 |
+| Vulkan SDK | for `ggml-vulkan` (any AMD/NVIDIA/Intel dGPU; iGPU works but slower) |
+| Rust | stable, with `cargo` (Tauri 2 shell) |
+| Node.js | LTS (`npm`) |
+| Python | 3.10+ with `numpy` (on-device converter + benches) |
+| Disk / RAM | ≥ 30 GB free disk; RAM: 8 GB+ runs 8B, 16 GB+ runs 35B (paging-bound), the reference machine is 48 GB |
+
+### 1.2 Build the engine (llama.cpp + patches + edge0 serving code)
+
+```powershell
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/windows
+pwsh -File scripts/vendor-build.ps1        # first build ≈ 10–20 min (compiles llama.cpp + Vulkan kernels)
+```
+
+The script resolves the engine depot in two ways, then materializes the supply:
+
+- **monorepo clone (normal path)** → the parent directory carries `vendor.llama.pin` + `patches/llama.cpp/{common,windows}`; the script replays the 8 patches (`git am --3way`) into an isolated build worktree at `../wt/win` (gitignored build area) — the vendor tree itself is **never patched in place**;
+- `EDGE0_DEPOT` env → point at an existing depot elsewhere.
+
+The llama.cpp tree is **not** a submodule: on first run the script clones upstream (shallow of blobs) into `../vendor/llama.cpp` and detaches the commit pinned in `vendor.llama.pin`; afterwards that directory is just a gitignored checkout you can refresh with `git fetch`. Set `EDGE0_LLAMA_URL` to clone from a mirror.
+
+`-AssembleOnly` runs everything except compilation (seconds; verifies the patch replay still applies and hashes match). Output:
+
+```
+wt\win\build-vk\bin\Release\llama-server.exe   (+ llama.dll, ggml*.dll)
+```
+
+### 1.3 Build the app (Tauri installer + portable exe)
+
+```powershell
+cd app
+npm install
+npx tauri build
+```
+
+Artifacts:
+
+```
+app\src-tauri\target\release\edge0-app.exe                              ← portable, no install
+app\src-tauri\target\release\bundle\nsis\edge0_0.1.0_x64-setup.exe      ← NSIS installer
+```
+
+The shell locates the engine via `EDGE0_BIN_DIR` (default: the depot engine build dir above — see `app/README.md` env table for the full override surface). Set it if your engine build lives elsewhere.
+
+### 1.4 Models
+
+Nothing to pre-install: the app downloads from HuggingFace on first use, then converts on-device (one-time, ~2–6 min for 8B). Per-file `sha256` is verified against a published manifest and resumable via HTTP Range.
+
+Manual seed (offline machines): place the MLX repo files under
+
+```
+~\.edge0\models\edge0-8b\        (chat_template.jinja, model*.safetensors, lora_edge0_8b.safetensors, tokenizer…)
+~\.edge0\models\edge0-35b\       (…sharded safetensors, lora_edge0_35b.safetensors…)
+```
+
+The converter output is `models\edge0-<tier>-gguf\edge0-<tier>.gguf` + LoRA adapter; conversion is idempotent (sha-gated). Home dir is `EDGE0_HOME` (default `~\.edge0`).
+
+### 1.5 Run the app
+
+Launch `edge0-app.exe` (or install the NSIS package). Typical flow:
+
+1. **Models page** — pick 8B (fast, ~5 GB) or 35B (~21 GB); *Download* → auto *Convert* → *Load*.
+2. **Chat page** — streamed Markdown answers; the engine runs as a supervised child process and is killed with the app (Windows Job Object, `KILL_ON_JOB_CLOSE`).
+3. **Doctor** (bottom of Models page) — one-click health check: engine binary, model files, disk, env.
+4. Troubleshooting: engine/convert logs land in `~\.edge0\logs\`.
+
+Engine API (the app uses it internally, you can point any OpenAI-compatible client at it): `http://127.0.0.1:<port>/v1/chat/completions`, loopback-only. Load params (process contract): `-ngl 99 -cmoe --ctx-size 8192 --flash-attn on --pool-mb <tier×RAM clamp> --mem-budget-mb …`.
+
+### 1.6 Test it
+
+```powershell
+cd app\src-tauri
+cargo test                                   # fast suite: resume/probe logic against a local fake server
+
+# end-to-end, real model (downloads 4.5 GB if not seeded):
+cargo test --test real8b -- --ignored --nocapture
+
+# throughput reproduction (expects the table in §2, ±day-to-day drift):
+python ..\..\tools\r3_bench.py --tier 8b
+python ..\..\tools\r3_bench.py --tier 35b
+```
+
+`r3_bench.py` resolves paths with the same contract as the app: the engine from the depot build dir
+(`EDGE0_BIN_DIR` override honored), the converted GGUF from the app's model home (`EDGE0_HOME`,
+default `~\.edge0\models`) or a repo-local `models/` when present; results land in `benchmarks/r3/`.
+
+```powershell
+# patch-replay gate (no compile, seconds):
+pwsh ..\..\scripts\vendor-build.ps1 -AssembleOnly
+```
+
+---
+
+## 2. Performance
+
+Reference machine: **Intel i7-14700K · AMD Radeon RX 9070 GRE (Vulkan) · 48 GB DDR5 · Windows 11**, steady-state (`-ngl 99 -cmoe`), bench medians of 3 reps (`tools/r3_bench.py`, ~240-token mixed-language prompt, native context).
+
+| model | prefill (tok/s) | decode (tok/s) | prefill of 240-tok prompt |
+|---|---:|---:|---:|
+| edge0-8B-A1B | ~212 | ~44.8 | ~1.1 s |
+| edge0-35B-A3B | ~70 | ~27.7 | ~3.7 s |
+
+Context to keep when reading these numbers:
+
+- Decode is the *token-generation* rate after the prompt is processed; the A1B / A3B active-parameter counts explain why a 35B model decodes at 27+ tok/s.
+- Short prompts are much faster than the 240-token column above: the app's first live request on 8B showed **~0.28 s to first token** for a 17-token prompt (engine log). The app runs with `--ctx-size 8192`; benches above use native context (131k / 262k).
+- These are warm-RAM measurements on the 48 GB reference machine (a test bed, not a product assumption). On the 16 GB / 8 GB product tiers, experts are served by mmap page-faults and decode drops with memory pressure; a process-level memory cap (`--mem-budget-mb`) is applied automatically from physical RAM to keep the OS working set sane (it measurably *helps* under pressure).
+- Same-condition A/B against unpatched vanilla llama.cpp (vendor tree with the patch bands skipped) shows the patch set costs nothing (8B: 41.5 vs 39.5 decode; 35B: 27.2 vs 27.4 — within machine drift).
+- Expect day-to-day drift of a few percent on this class of machine; compare like-for-like, same day, same process if you benchmark.
+
+---
+
+## 3. Technical Details
+
+### 3.1 Architecture
+
+```mermaid
+flowchart LR
+  subgraph ondevice ["On-device, first run"]
+    A[MLX safetensors<br/>int4 g64] -->|repack, bit-exact<br/>no requant| B[GGUF Q4_1<br/>edge0-tier.gguf]
+    A2[MLX LoRA] -->|fuse + adapt| B2[GGUF adapter]
+  end
+  B --> C[("pinned llama.cpp fork<br/>7ab4ee7 + 8 patches")]
+  B2 --> C
+  C -->|Vulkan: dense attn / GDN / norms| D[(GPU)]
+  C -->|"-cmoe: expert weights"| E[(CPU, mmap +<br/>L1 advisory pool)]
+  C --> F[llama-server<br/>127.0.0.1 OpenAI API]
+  F --> G[Tauri 2 shell<br/>supervisor + chat UI]
+```
+
+```mermaid
+flowchart TB
+  subgraph build ["Source-of-truth layout (open source)"]
+    U[github ggml-org/llama.cpp] -->|detached pin, never edited| V[vendor/llama.cpp @7ab4ee7]
+    P[patches/llama.cpp/common+windows] -->|git am --3way into per-platform worktree| W[wt/win]
+    V --> W
+    S[serve/*.cc edge0 code] -->|CMake glob patch #3| W
+  end
+```
+
+### 3.2 Repo map
+```
+serve/                 edge0 C++ pieces compiled into libllama (advisory prefetch router, memory budget)
+tools/                 MLX→GGUF converter chain (repack + LoRA adapter + catalog) and the throughput bench
+app/                   Tauri 2 shell (src-tauri/ Rust, src/ React) — see app/README.md
+scripts/               vendor-build.ps1 — engine assembly: patch replay into an isolated worktree
+```
+
+The engine changes live at the monorepo root, not in a fork: `vendor/llama.cpp` is the pinned
+pristine upstream checkout (never patched in place) and `patches/llama.cpp/` carries the 8
+hook-point patches + band README (`common/`, `windows/` here; `android/` for the companion app).
+`windows/` deliberately ships **no** nested `patches/` copy — one source of truth, no drift.
+
+### 3.3 Known gaps (as of this release)
+
+- Installer is **unsigned** (SmartScreen warning on first run); no auto-updater yet.
+- Engine is not yet bundled inside the installer — portable builds assume the engine build dir; bundling is on the roadmap.
+- The `--pool-mb` L1 pool is env-dormant on Windows pending the final memory-tier A/B; the app clamps `--mem-budget-mb` regardless.
+- 35B on ≤16 GB machines is functional but paging-bound; documented floor is 8 GB *working with* the memory cap, at reduced tok/s.
+
+---
+
+*Issues and bench results against other hardware welcome — attach `~\.edge0\logs\engine-*.log` and the `doctor` page screenshot.*
diff --git a/windows/README_es.md b/windows/README_es.md
new file mode 100644
index 0000000..b543da1
--- /dev/null
+++ b/windows/README_es.md
@@ -0,0 +1,179 @@
+# edge0-windows — Guía de compilación de release
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | Español | [Français](README_fr.md)
+
+Este documento está dirigido a **desarrolladores que compilan o evalúan la aplicación de Windows desde el código fuente**. Cubre: inicio rápido (compilar → ejecutar → probar), rendimiento medido en la máquina de referencia y las decisiones técnicas detrás del stack.
+
+edge0 se publica como **monorepo** — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — cuyo nivel superior contiene el suministro compartido del motor (`vendor.llama.pin` + el `vendor/llama.cpp` materializado por los scripts, los conjuntos de bandas de `patches/llama.cpp/`) y los subproyectos de plataforma (`windows/` = esta aplicación, `android/` = el complemento). La inferencia se ejecuta sobre una versión fijada (pin) de llama.cpp upstream, parcheada como un conjunto de parches reejecutable, con cómputo denso en Vulkan y pesos de expertos MoE servidos desde la CPU.
+
+Atribución de terceros: consulta `NOTICE` en este directorio. Fichas de modelo: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+
+---
+
+## 1. Inicio rápido
+
+### 1.1 Requisitos previos
+
+| toolchain | requisito |
+|---|---|
+| SO | Windows 10 / 11 x64 |
+| Visual Studio 2022 | carga de trabajo C++ (MSVC) |
+| CMake | ≥ 3.21 |
+| Vulkan SDK | para `ggml-vulkan` (cualquier dGPU AMD/NVIDIA/Intel; la iGPU funciona pero es más lenta) |
+| Rust | stable, con `cargo` (shell de Tauri 2) |
+| Node.js | LTS (`npm`) |
+| Python | 3.10+ con `numpy` (conversor en el dispositivo + benchmarks) |
+| Disco / RAM | ≥ 30 GB de disco libre; RAM: 8 GB+ ejecuta 8B, 16 GB+ ejecuta 35B (limitado por paginación), la máquina de referencia tiene 48 GB |
+
+### 1.2 Compilar el motor (llama.cpp + parches + código de servicio de edge0)
+
+```powershell
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/windows
+pwsh -File scripts/vendor-build.ps1        # first build ≈ 10–20 min (compiles llama.cpp + Vulkan kernels)
+```
+
+El script resuelve el depósito del motor de dos formas y después materializa el suministro:
+
+- **clon del monorepo (ruta normal)** → el directorio padre contiene `vendor.llama.pin` + `patches/llama.cpp/{common,windows}`; el script reejecuta los 8 parches (`git am --3way`) en un worktree de compilación aislado en `../wt/win` (área de compilación ignorada por git) — el árbol del vendor en sí **nunca** se parchea in situ;
+- variable de entorno `EDGE0_DEPOT` → apunta a un depósito existente en otra ubicación.
+
+El árbol de llama.cpp **no** es un submódulo: en la primera ejecución el script clona upstream (shallow de blobs) en `../vendor/llama.cpp` y hace detach del commit fijado en `vendor.llama.pin`; después, ese directorio es simplemente un checkout ignorado por git que puedes refrescar con `git fetch`. Establece `EDGE0_LLAMA_URL` para clonar desde un mirror.
+
+`-AssembleOnly` ejecuta todo menos la compilación (segundos; verifica que la reejecución de los parches sigue aplicando y que los hashes coinciden). Salida:
+
+```
+wt\win\build-vk\bin\Release\llama-server.exe   (+ llama.dll, ggml*.dll)
+```
+
+### 1.3 Compilar la aplicación (instalador de Tauri + exe portable)
+
+```powershell
+cd app
+npm install
+npx tauri build
+```
+
+Artefactos:
+
+```
+app\src-tauri\target\release\edge0-app.exe                              ← portable, no install
+app\src-tauri\target\release\bundle\nsis\edge0_0.1.0_x64-setup.exe      ← NSIS installer
+```
+
+El shell localiza el motor mediante `EDGE0_BIN_DIR` (por defecto: el directorio de compilación del motor del depósito indicado arriba — consulta la tabla de variables de entorno en `app/README.md` para ver todas las opciones de sobrescritura). Establécelo si tu compilación del motor está en otra ubicación.
+
+### 1.4 Modelos
+
+Nada que preinstalar: la aplicación descarga desde HuggingFace en el primer uso y después convierte en el dispositivo (una sola vez, ~2–6 min para 8B). El `sha256` de cada archivo se verifica contra un manifiesto publicado y la descarga se puede reanudar mediante HTTP Range.
+
+Precarga manual (máquinas sin conexión): coloca los archivos del repositorio MLX en
+
+```
+~\.edge0\models\edge0-8b\        (chat_template.jinja, model*.safetensors, lora_edge0_8b.safetensors, tokenizer…)
+~\.edge0\models\edge0-35b\       (…sharded safetensors, lora_edge0_35b.safetensors…)
+```
+
+La salida del conversor es `models\edge0-<tier>-gguf\edge0-<tier>.gguf` + adaptador LoRA; la conversión es idempotente (controlada por sha). El directorio principal es `EDGE0_HOME` (por defecto `~\.edge0`).
+
+### 1.5 Ejecutar la aplicación
+
+Inicia `edge0-app.exe` (o instala el paquete NSIS). Flujo típico:
+
+1. **Página Modelos** — elige 8B (rápido, ~5 GB) o 35B (~21 GB); *Descargar* → *Convertir* automáticamente → *Cargar*.
+2. **Página Chat** — respuestas en Markdown transmitidas en streaming; el motor se ejecuta como un proceso hijo supervisado y se termina junto con la aplicación (Windows Job Object, `KILL_ON_JOB_CLOSE`).
+3. **Doctor** (al final de la página Modelos) — chequeo de salud en un clic: binario del motor, archivos de modelos, disco, entorno.
+4. Solución de problemas: los logs del motor/conversión quedan en `~\.edge0\logs\`.
+
+API del motor (la aplicación la usa internamente, puedes apuntar a ella cualquier cliente compatible con OpenAI): `http://127.0.0.1:<port>/v1/chat/completions`, solo loopback. Parámetros de carga (contrato del proceso): `-ngl 99 -cmoe --ctx-size 8192 --flash-attn on --pool-mb <tier×RAM clamp> --mem-budget-mb …`.
+
+### 1.6 Probarla
+
+```powershell
+cd app\src-tauri
+cargo test                                   # fast suite: resume/probe logic against a local fake server
+
+# end-to-end, real model (downloads 4.5 GB if not seeded):
+cargo test --test real8b -- --ignored --nocapture
+
+# throughput reproduction (expects the table in §2, ±day-to-day drift):
+python ..\..\tools\r3_bench.py --tier 8b
+python ..\..\tools\r3_bench.py --tier 35b
+```
+
+`r3_bench.py` resuelve las rutas con el mismo contrato que la aplicación: el motor desde el directorio de compilación del depósito (se respeta la sobrescritura `EDGE0_BIN_DIR`), el GGUF convertido desde el directorio principal de modelos de la aplicación (`EDGE0_HOME`, por defecto `~\.edge0\models`) o un `models/` local del repositorio si existe; los resultados quedan en `benchmarks/r3/`.
+
+```powershell
+# patch-replay gate (no compile, seconds):
+pwsh ..\..\scripts\vendor-build.ps1 -AssembleOnly
+```
+
+---
+
+## 2. Rendimiento
+
+Máquina de referencia: **Intel i7-14700K · AMD Radeon RX 9070 GRE (Vulkan) · 48 GB DDR5 · Windows 11**, estado estacionario (`-ngl 99 -cmoe`), medianas de benchmarks de 3 repeticiones (`tools/r3_bench.py`, prompt mixto de idiomas de ~240 tokens, contexto nativo).
+
+| modelo | prefill (tok/s) | decode (tok/s) | prefill de un prompt de 240 tokens |
+|---|---:|---:|---:|
+| edge0-8B-A1B | ~212 | ~44.8 | ~1.1 s |
+| edge0-35B-A3B | ~70 | ~27.7 | ~3.7 s |
+
+Contexto a tener en cuenta al leer estas cifras:
+
+- Decode es la tasa de *generación de tokens* después de procesar el prompt; el número de parámetros activos de A1B / A3B explica por qué un modelo de 35B decodifica a más de 27 tok/s.
+- Los prompts cortos son mucho más rápidos que la columna de 240 tokens de arriba: la primera petición en vivo de la aplicación con 8B mostró **~0.28 s hasta el primer token** para un prompt de 17 tokens (log del motor). La aplicación se ejecuta con `--ctx-size 8192`; los benchmarks anteriores usan contexto nativo (131k / 262k).
+- Son mediciones con RAM caliente en la máquina de referencia de 48 GB (un banco de pruebas, no un supuesto del producto). En los niveles de producto de 16 GB / 8 GB, los expertos se sirven mediante fallos de página mmap y el decode cae con la presión de memoria; se aplica automáticamente un límite de memoria a nivel de proceso (`--mem-budget-mb`) a partir de la RAM física para mantener el conjunto de trabajo del SO bajo control (de forma medible *ayuda* bajo presión).
+- Una comparación A/B en las mismas condiciones frente a llama.cpp vanilla sin parchear (árbol del vendor omitiendo las bandas de parches) muestra que el conjunto de parches no cuesta nada (8B: 41.5 vs 39.5 en decode; 35B: 27.2 vs 27.4 — dentro de la deriva de la máquina).
+- Espera una deriva diaria de unos pocos puntos porcentuales en esta clase de máquina; si haces benchmarks, compara en igualdad de condiciones, el mismo día y el mismo proceso.
+
+---
+
+## 3. Detalles técnicos
+
+### 3.1 Arquitectura
+
+```mermaid
+flowchart LR
+  subgraph ondevice ["On-device, first run"]
+    A[MLX safetensors<br/>int4 g64] -->|repack, bit-exact<br/>no requant| B[GGUF Q4_1<br/>edge0-tier.gguf]
+    A2[MLX LoRA] -->|fuse + adapt| B2[GGUF adapter]
+  end
+  B --> C[("pinned llama.cpp fork<br/>7ab4ee7 + 8 patches")]
+  B2 --> C
+  C -->|Vulkan: dense attn / GDN / norms| D[(GPU)]
+  C -->|"-cmoe: expert weights"| E[(CPU, mmap +<br/>L1 advisory pool)]
+  C --> F[llama-server<br/>127.0.0.1 OpenAI API]
+  F --> G[Tauri 2 shell<br/>supervisor + chat UI]
+```
+
+```mermaid
+flowchart TB
+  subgraph build ["Source-of-truth layout (open source)"]
+    U[github ggml-org/llama.cpp] -->|detached pin, never edited| V[vendor/llama.cpp @7ab4ee7]
+    P[patches/llama.cpp/common+windows] -->|git am --3way into per-platform worktree| W[wt/win]
+    V --> W
+    S[serve/*.cc edge0 code] -->|CMake glob patch #3| W
+  end
+```
+
+### 3.2 Mapa del repositorio
+```
+serve/                 edge0 C++ pieces compiled into libllama (advisory prefetch router, memory budget)
+tools/                 MLX→GGUF converter chain (repack + LoRA adapter + catalog) and the throughput bench
+app/                   Tauri 2 shell (src-tauri/ Rust, src/ React) — see app/README.md
+scripts/               vendor-build.ps1 — engine assembly: patch replay into an isolated worktree
+```
+
+Los cambios del motor viven en la raíz del monorepo, no en un fork: `vendor/llama.cpp` es el checkout upstream prístino fijado (nunca parcheado in situ) y `patches/llama.cpp/` contiene los 8 parches de puntos de anclaje + el README de bandas (`common/`, `windows/` aquí; `android/` para la aplicación complementaria). `windows/` no incluye deliberadamente **ninguna** copia anidada de `patches/` — una única fuente de verdad, sin deriva.
+
+### 3.3 Limitaciones conocidas (en esta versión)
+
+- El instalador está **sin firmar** (advertencia de SmartScreen en el primer arranque); aún no hay actualizador automático.
+- El motor todavía no se incluye dentro del instalador — las compilaciones portables suponen el directorio de compilación del motor; la inclusión está en la hoja de ruta.
+- El pool L1 de `--pool-mb` está desactivado por variable de entorno en Windows a la espera del A/B definitivo por niveles de memoria; la aplicación limita `--mem-budget-mb` de todos modos.
+- 35B en máquinas de ≤16 GB es funcional pero está limitado por paginación; el mínimo documentado es 8 GB *funcionando con* el límite de memoria, a tok/s reducidos.
+
+---
+
+*Son bienvenidos los issues y los resultados de benchmarks con otro hardware — adjunta `~\.edge0\logs\engine-*.log` y una captura de la página `doctor`.*
diff --git a/windows/README_fr.md b/windows/README_fr.md
new file mode 100644
index 0000000..f875624
--- /dev/null
+++ b/windows/README_fr.md
@@ -0,0 +1,183 @@
+# edge0-windows — Guide de compilation de la release
+
+[English](README.md) | [中文](README_zh.md) | [日本語](README_ja.md) | [Español](README_es.md) | Français
+
+Ce document s'adresse aux **développeurs qui compilent ou évaluent l'application Windows depuis les sources**. Il couvre : le démarrage rapide (compilation → exécution → test), les performances mesurées sur la machine de référence et les choix techniques derrière la stack.
+
+edge0 est publié sous forme de **monorepo** — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — dont le niveau supérieur contient l'approvisionnement partagé du moteur (`vendor.llama.pin` + le `vendor/llama.cpp` matérialisé par les scripts, les jeux de bandes `patches/llama.cpp/`) et les sous-projets par plateforme (`windows/` = cette application, `android/` = le compagnon). L'inférence s'exécute sur une version épinglée de llama.cpp upstream, patchée sous forme d'un jeu de patchs rejouable, avec le calcul dense sur Vulkan et les poids d'experts MoE servis depuis le CPU.
+
+Attribution tierce : voir `NOTICE` dans ce répertoire. Cartes de modèles : [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview).
+
+---
+
+## 1. Démarrage rapide
+
+### 1.1 Prérequis
+
+| chaîne d'outils | exigence |
+|---|---|
+| OS | Windows 10 / 11 x64 |
+| Visual Studio 2022 | charge de travail C++ (MSVC) |
+| CMake | ≥ 3.21 |
+| Vulkan SDK | pour `ggml-vulkan` (tout dGPU AMD/NVIDIA/Intel ; l'iGPU fonctionne mais plus lentement) |
+| Rust | stable, avec `cargo` (shell Tauri 2) |
+| Node.js | LTS (`npm`) |
+| Python | 3.10+ avec `numpy` (convertisseur sur l'appareil + benches) |
+| Disque / RAM | ≥ 30 GB de disque libre ; RAM : 8 GB+ pour le 8B, 16 GB+ pour le 35B (limité par la pagination), la machine de référence est à 48 GB |
+
+### 1.2 Compiler le moteur (llama.cpp + patchs + code de serving edge0)
+
+```powershell
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/windows
+pwsh -File scripts/vendor-build.ps1        # first build ≈ 10–20 min (compiles llama.cpp + Vulkan kernels)
+```
+
+Le script résout le dépôt du moteur de deux manières, puis matérialise l'approvisionnement :
+
+- **clone du monorepo (chemin normal)** → le répertoire parent porte `vendor.llama.pin` + `patches/llama.cpp/{common,windows}` ; le script rejoue les 8 patchs (`git am --3way`) dans un worktree de compilation isolé dans `../wt/win` (zone de compilation ignorée par git) — l'arborescence vendor elle-même n'est **jamais patchée en place** ;
+- variable d'environnement `EDGE0_DEPOT` → pointe vers un dépôt existant ailleurs.
+
+L'arborescence llama.cpp n'est **pas** un sous-module : au premier lancement, le script clone upstream (shallow en blobs) dans `../vendor/llama.cpp` et détache le commit épinglé dans `vendor.llama.pin` ; ensuite, ce répertoire n'est qu'un checkout ignoré par git que vous pouvez rafraîchir avec `git fetch`. Définissez `EDGE0_LLAMA_URL` pour cloner depuis un miroir.
+
+`-AssembleOnly` exécute tout sauf la compilation (quelques secondes ; vérifie que le rejeu des patchs s'applique toujours et que les hashes correspondent). Sortie :
+
+```
+wt\win\build-vk\bin\Release\llama-server.exe   (+ llama.dll, ggml*.dll)
+```
+
+### 1.3 Compiler l'application (installeur Tauri + exe portable)
+
+```powershell
+cd app
+npm install
+npx tauri build
+```
+
+Artéfacts :
+
+```
+app\src-tauri\target\release\edge0-app.exe                              ← portable, no install
+app\src-tauri\target\release\bundle\nsis\edge0_0.1.0_x64-setup.exe      ← NSIS installer
+```
+
+Le shell localise le moteur via `EDGE0_BIN_DIR` (par défaut : le répertoire de compilation du moteur du dépôt ci-dessus — voir le tableau des variables d'environnement de `app/README.md` pour toutes les options de substitution). Définissez-le si votre compilation du moteur se trouve ailleurs.
+
+### 1.4 Modèles
+
+Rien à préinstaller : l'application télécharge depuis HuggingFace à la première utilisation, puis convertit sur l'appareil (une seule fois, ~2–6 min pour le 8B). Le `sha256` par fichier est vérifié contre un manifeste publié et le téléchargement est reprenable via HTTP Range.
+
+Amorçage manuel (machines hors ligne) : placez les fichiers du dépôt MLX sous
+
+```
+~\.edge0\models\edge0-8b\        (chat_template.jinja, model*.safetensors, lora_edge0_8b.safetensors, tokenizer…)
+~\.edge0\models\edge0-35b\       (…sharded safetensors, lora_edge0_35b.safetensors…)
+```
+
+La sortie du convertisseur est `models\edge0-<tier>-gguf\edge0-<tier>.gguf` + l'adaptateur LoRA ; la conversion est idempotente (contrôlée par sha). Le répertoire personnel est `EDGE0_HOME` (`~\.edge0` par défaut).
+
+### 1.5 Lancer l'application
+
+Lancez `edge0-app.exe` (ou installez le paquet NSIS). Flux typique :
+
+1. **Page Models** — choisissez 8B (rapide, ~5 GB) ou 35B (~21 GB) ; *Download* → *Convert* automatique → *Load*.
+2. **Page Chat** — réponses Markdown en streaming ; le moteur s'exécute comme un processus enfant supervisé et est tué avec l'application (Windows Job Object, `KILL_ON_JOB_CLOSE`).
+3. **Doctor** (en bas de la page Models) — vérification de santé en un clic : binaire du moteur, fichiers de modèles, disque, environnement.
+4. Dépannage : les journaux du moteur/de conversion se trouvent dans `~\.edge0\logs\`.
+
+API du moteur (l'application l'utilise en interne, vous pouvez pointer n'importe quel client compatible OpenAI dessus) : `http://127.0.0.1:<port>/v1/chat/completions`, loopback uniquement. Paramètres de chargement (contrat du processus) : `-ngl 99 -cmoe --ctx-size 8192 --flash-attn on --pool-mb <tier×RAM clamp> --mem-budget-mb …`.
+
+### 1.6 Tester
+
+```powershell
+cd app\src-tauri
+cargo test                                   # fast suite: resume/probe logic against a local fake server
+
+# end-to-end, real model (downloads 4.5 GB if not seeded):
+cargo test --test real8b -- --ignored --nocapture
+
+# throughput reproduction (expects the table in §2, ±day-to-day drift):
+python ..\..\tools\r3_bench.py --tier 8b
+python ..\..\tools\r3_bench.py --tier 35b
+```
+
+`r3_bench.py` résout les chemins avec le même contrat que l'application : le moteur depuis le répertoire de compilation du dépôt
+(substitution `EDGE0_BIN_DIR` respectée), le GGUF converti depuis le répertoire personnel des modèles de l'application (`EDGE0_HOME`,
+`~\.edge0\models` par défaut) ou un `models/` local au dépôt s'il existe ; les résultats arrivent dans `benchmarks/r3/`.
+
+```powershell
+# patch-replay gate (no compile, seconds):
+pwsh ..\..\scripts\vendor-build.ps1 -AssembleOnly
+```
+
+---
+
+## 2. Performances
+
+Machine de référence : **Intel i7-14700K · AMD Radeon RX 9070 GRE (Vulkan) · 48 GB DDR5 · Windows 11**, régime stable (`-ngl 99 -cmoe`), médianes de bench sur 3 répétitions (`tools/r3_bench.py`, prompt multilingue d'environ 240 tokens, contexte natif).
+
+| modèle | prefill (tok/s) | decode (tok/s) | prefill d'un prompt de 240 tokens |
+|---|---:|---:|---:|
+| edge0-8B-A1B | ~212 | ~44.8 | ~1.1 s |
+| edge0-35B-A3B | ~70 | ~27.7 | ~3.7 s |
+
+Contexte à garder en tête en lisant ces chiffres :
+
+- Le decode est le taux de *génération de tokens* après traitement du prompt ; le nombre de paramètres actifs A1B / A3B explique pourquoi un modèle 35B décode à plus de 27 tok/s.
+- Les prompts courts sont bien plus rapides que la colonne 240 tokens ci-dessus : la première requête réelle de l'application sur le 8B a montré **~0,28 s avant le premier token** pour un prompt de 17 tokens (journal du moteur). L'application s'exécute avec `--ctx-size 8192` ; les benches ci-dessus utilisent le contexte natif (131k / 262k).
+- Ce sont des mesures à chaud en RAM sur la machine de référence de 48 GB (un banc de test, pas une hypothèse produit). Sur les gammes produit 16 GB / 8 GB, les experts sont servis par des défauts de page mmap et le decode chute avec la pression mémoire ; une limite mémoire au niveau du processus (`--mem-budget-mb`) est appliquée automatiquement à partir de la RAM physique pour garder le jeu de travail de l'OS raisonnable (elle *aide* de façon mesurable sous pression).
+- Un A/B dans les mêmes conditions contre llama.cpp vanilla non patché (arborescence vendor avec les bandes de patchs ignorées) montre que le jeu de patchs ne coûte rien (8B : 41,5 contre 39,5 en decode ; 35B : 27,2 contre 27,4 — dans la dérive de la machine).
+- Attendez-vous à une dérive de quelques pourcents d'un jour à l'autre sur cette classe de machine ; comparez à conditions égales, le même jour, le même processus, si vous benchmarkez.
+
+---
+
+## 3. Détails techniques
+
+### 3.1 Architecture
+
+```mermaid
+flowchart LR
+  subgraph ondevice ["On-device, first run"]
+    A[MLX safetensors<br/>int4 g64] -->|repack, bit-exact<br/>no requant| B[GGUF Q4_1<br/>edge0-tier.gguf]
+    A2[MLX LoRA] -->|fuse + adapt| B2[GGUF adapter]
+  end
+  B --> C[("pinned llama.cpp fork<br/>7ab4ee7 + 8 patches")]
+  B2 --> C
+  C -->|Vulkan: dense attn / GDN / norms| D[(GPU)]
+  C -->|"-cmoe: expert weights"| E[(CPU, mmap +<br/>L1 advisory pool)]
+  C --> F[llama-server<br/>127.0.0.1 OpenAI API]
+  F --> G[Tauri 2 shell<br/>supervisor + chat UI]
+```
+
+```mermaid
+flowchart TB
+  subgraph build ["Source-of-truth layout (open source)"]
+    U[github ggml-org/llama.cpp] -->|detached pin, never edited| V[vendor/llama.cpp @7ab4ee7]
+    P[patches/llama.cpp/common+windows] -->|git am --3way into per-platform worktree| W[wt/win]
+    V --> W
+    S[serve/*.cc edge0 code] -->|CMake glob patch #3| W
+  end
+```
+
+### 3.2 Structure du dépôt
+```
+serve/                 edge0 C++ pieces compiled into libllama (advisory prefetch router, memory budget)
+tools/                 MLX→GGUF converter chain (repack + LoRA adapter + catalog) and the throughput bench
+app/                   Tauri 2 shell (src-tauri/ Rust, src/ React) — see app/README.md
+scripts/               vendor-build.ps1 — engine assembly: patch replay into an isolated worktree
+```
+
+Les modifications du moteur se trouvent à la racine du monorepo, pas dans un fork : `vendor/llama.cpp` est le checkout upstream épinglé d'origine (jamais patché en place) et `patches/llama.cpp/` porte les 8
+patchs de points d'accroche + le README de bande (`common/`, `windows/` ici ; `android/` pour l'application compagne).
+`windows/` n'embarque volontairement **aucune** copie imbriquée de `patches/` — une seule source de vérité, aucune dérive.
+
+### 3.3 Lacunes connues (à la date de cette release)
+
+- L'installeur n'est **pas signé** (avertissement SmartScreen au premier lancement) ; pas encore de mise à jour automatique.
+- Le moteur n'est pas encore intégré à l'installeur — les versions portables supposent le répertoire de compilation du moteur ; l'intégration est sur la feuille de route.
+- Le pool L1 `--pool-mb` reste en veille sur Windows dans l'attente de l'A/B final sur les niveaux de mémoire ; l'application plafonne `--mem-budget-mb` quoi qu'il en soit.
+- Le 35B sur des machines ≤ 16 GB est fonctionnel mais limité par la pagination ; le plancher documenté est 8 GB *avec* le plafond mémoire, à un débit de tok/s réduit.
+
+---
+
+*Les issues et résultats de bench sur d'autres matériels sont les bienvenus — joignez `~\.edge0\logs\engine-*.log` et la capture d'écran de la page `doctor`.*
diff --git a/windows/README_ja.md b/windows/README_ja.md
new file mode 100644
index 0000000..fe2660b
--- /dev/null
+++ b/windows/README_ja.md
@@ -0,0 +1,179 @@
+# edge0-windows — リリースビルドガイド
+
+[English](README.md) | [中文](README_zh.md) | 日本語 | [Español](README_es.md) | [Français](README_fr.md)
+
+このドキュメントは、**ソースから Windows アプリをビルドまたは評価する開発者**向けのものです。内容は、クイックスタート(ビルド → 実行 → テスト)、リファレンスマシンでの実測パフォーマンス、そしてスタックの背後にある技術的選択です。
+
+edge0 は**モノレポ**として公開されています — [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) — 最上位階層には共有のエンジンサプライ(`vendor.llama.pin` + スクリプトが実体化する `vendor/llama.cpp`、`patches/llama.cpp/` のバンドセット)とプラットフォームのサブプロジェクト(`windows/` = 本アプリ、`android/` = コンパニオン)が含まれます。推論はピン留めされた上流の llama.cpp 上で実行され、リプレイ可能なパッチセットとしてパッチが適用され、dense 計算は Vulkan 上、MoE エキスパート重みは CPU から供給されます。
+
+サードパーティの帰属表示: このディレクトリの `NOTICE` を参照してください。モデルカード: [Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)。
+
+---
+
+## 1. クイックスタート
+
+### 1.1 前提条件
+
+| ツールチェーン | 要件 |
+|---|---|
+| OS | Windows 10 / 11 x64 |
+| Visual Studio 2022 | C++ ワークロード(MSVC) |
+| CMake | ≥ 3.21 |
+| Vulkan SDK | `ggml-vulkan` 用(AMD/NVIDIA/Intel の dGPU ならどれでも;iGPU でも動作するがより遅い) |
+| Rust | stable、`cargo` 付き(Tauri 2 シェル) |
+| Node.js | LTS(`npm`) |
+| Python | 3.10+ と `numpy`(デバイス上コンバータ + ベンチ) |
+| ディスク / RAM | 30 GB 以上の空きディスク;RAM: 8 GB+ で 8B、16 GB+ で 35B が動作(ページング律速);リファレンスマシンは 48 GB |
+
+### 1.2 エンジンのビルド(llama.cpp + パッチ + edge0 サービングコード)
+
+```powershell
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/windows
+pwsh -File scripts/vendor-build.ps1        # first build ≈ 10–20 min (compiles llama.cpp + Vulkan kernels)
+```
+
+スクリプトはエンジンデポを 2 つの方法で解決し、サプライを実体化します:
+
+- **モノレポのクローン(通常のパス)** → 親ディレクトリが `vendor.llama.pin` + `patches/llama.cpp/{common,windows}` を持ちます。スクリプトは 8 つのパッチ(`git am --3way`)を `../wt/win` の隔離されたビルドワークツリーにリプレイします(gitignore されたビルド領域)— ベンダーツリー自体は**決してその場でパッチされません**;
+- `EDGE0_DEPOT` 環境変数 → 他の場所にある既存のデポを指定します。
+
+llama.cpp ツリーはサブモジュールでは**ありません**: 初回実行時、スクリプトは上流を(blobs の shallow クローンで)`../vendor/llama.cpp` にクローンし、`vendor.llama.pin` にピン留めされたコミットを detach します。以降、そのディレクトリは `git fetch` で更新できる、gitignore された単なるチェックアウトです。ミラーからクローンするには `EDGE0_LLAMA_URL` を設定します。
+
+`-AssembleOnly` はコンパイル以外のすべてを実行します(数秒;パッチのリプレイが今も適用可能でハッシュが一致することを検証します)。出力:
+
+```
+wt\win\build-vk\bin\Release\llama-server.exe   (+ llama.dll, ggml*.dll)
+```
+
+### 1.3 アプリのビルド(Tauri インストーラ + ポータブル exe)
+
+```powershell
+cd app
+npm install
+npx tauri build
+```
+
+成果物:
+
+```
+app\src-tauri\target\release\edge0-app.exe                              ← portable, no install
+app\src-tauri\target\release\bundle\nsis\edge0_0.1.0_x64-setup.exe      ← NSIS installer
+```
+
+シェルは `EDGE0_BIN_DIR` 経由でエンジンを特定します(デフォルト: 上記のデポのエンジンビルドディレクトリ — 完全なオーバーライド一覧は `app/README.md` の環境変数テーブルを参照)。エンジンビルドが他の場所にある場合は設定してください。
+
+### 1.4 モデル
+
+事前インストールは不要です: アプリは初回使用時に HuggingFace からダウンロードし、その後デバイス上で変換します(一度きり;8B で約 2〜6 分)。ファイルごとの `sha256` が公開マニフェストと照合して検証され、HTTP Range 経由で再開可能です。
+
+手動シード(オフラインマシン): MLX リポジトリのファイルを以下に配置します
+
+```
+~\.edge0\models\edge0-8b\        (chat_template.jinja, model*.safetensors, lora_edge0_8b.safetensors, tokenizer…)
+~\.edge0\models\edge0-35b\       (…sharded safetensors, lora_edge0_35b.safetensors…)
+```
+
+コンバータの出力は `models\edge0-<tier>-gguf\edge0-<tier>.gguf` + LoRA アダプタです。変換は冪等です(sha ゲート付き)。ホームディレクトリは `EDGE0_HOME`(デフォルト `~\.edge0`)です。
+
+### 1.5 アプリの実行
+
+`edge0-app.exe` を起動します(または NSIS パッケージをインストール)。典型的なフロー:
+
+1. **Models ページ** — 8B(高速、約 5 GB)または 35B(約 21 GB)を選択;*Download* → 自動で *Convert* → *Load*。
+2. **Chat ページ** — ストリーミングされる Markdown の回答。エンジンは監督付きの子プロセスとして実行され、アプリとともに kill されます(Windows Job Object、`KILL_ON_JOB_CLOSE`)。
+3. **Doctor**(Models ページの下部)— ワンクリックのヘルスチェック: エンジンバイナリ、モデルファイル、ディスク、環境変数。
+4. トラブルシューティング: エンジン/変換のログは `~\.edge0\logs\` に出力されます。
+
+エンジン API(アプリが内部的に使用しています;任意の OpenAI 互換クライアントから指定できます): `http://127.0.0.1:<port>/v1/chat/completions`、ループバック専用。ロードパラメータ(プロセスコントラクト): `-ngl 99 -cmoe --ctx-size 8192 --flash-attn on --pool-mb <tier×RAM clamp> --mem-budget-mb …`。
+
+### 1.6 テスト
+
+```powershell
+cd app\src-tauri
+cargo test                                   # fast suite: resume/probe logic against a local fake server
+
+# end-to-end, real model (downloads 4.5 GB if not seeded):
+cargo test --test real8b -- --ignored --nocapture
+
+# throughput reproduction (expects the table in §2, ±day-to-day drift):
+python ..\..\tools\r3_bench.py --tier 8b
+python ..\..\tools\r3_bench.py --tier 35b
+```
+
+`r3_bench.py` はアプリと同じコントラクトでパスを解決します: エンジンはデポのビルドディレクトリから(`EDGE0_BIN_DIR` のオーバーライドを尊重)、変換済み GGUF はアプリのモデルホーム(`EDGE0_HOME`、デフォルト `~\.edge0\models`)から、または存在する場合はリポジトリローカルの `models/` から。結果は `benchmarks/r3/` に出力されます。
+
+```powershell
+# patch-replay gate (no compile, seconds):
+pwsh ..\..\scripts\vendor-build.ps1 -AssembleOnly
+```
+
+---
+
+## 2. パフォーマンス
+
+リファレンスマシン: **Intel i7-14700K · AMD Radeon RX 9070 GRE(Vulkan)· 48 GB DDR5 · Windows 11**、定常状態(`-ngl 99 -cmoe`)、3 回実行のベンチ中央値(`tools/r3_bench.py`、約 240 トークンの言語混在プロンプト、ネイティブコンテキスト)。
+
+| モデル | プリフィル (tok/s) | デコード (tok/s) | 240 トークンプロンプトのプリフィル |
+|---|---:|---:|---:|
+| edge0-8B-A1B | ~212 | ~44.8 | ~1.1 s |
+| edge0-35B-A3B | ~70 | ~27.7 | ~3.7 s |
+
+これらの数値を読む際に念頭に置くべきコンテキスト:
+
+- デコードはプロンプト処理後の*トークン生成*レートです。A1B / A3B のアクティブパラメータ数が、35B モデルが 27+ tok/s でデコードできる理由を説明します。
+- 短いプロンプトは上記の 240 トークンの列よりはるかに高速です: 8B でのアプリの最初のライブリクエストは、17 トークンのプロンプトで**最初のトークンまで約 0.28 s** を示しました(エンジンログ)。アプリは `--ctx-size 8192` で実行されます;上記のベンチはネイティブコンテキスト(131k / 262k)を使用しています。
+- これらは 48 GB リファレンスマシンでのウォーム RAM の測定値です(テストベッドであり、製品としての前提ではありません)。16 GB / 8 GB の製品ティアでは、エキスパートは mmap のページフォルトによって供給され、メモリプレッシャーに応じてデコードが低下します。プロセスレベルのメモリ上限(`--mem-budget-mb`)が物理 RAM から自動的に適用され、OS のワーキングセットを健全に保ちます(プレッシャー下では測定可能なほど*役立ちます*)。
+- パッチ未適用の素の llama.cpp(パッチバンドをスキップしたベンダーツリー)との同一条件での A/B 比較は、パッチセットのコストがゼロであることを示しています(8B: デコード 41.5 vs 39.5;35B: 27.2 vs 27.4 — マシンのドリフトの範囲内)。
+- このクラスのマシンでは数パーセントの日ごとのドリフトが予想されます。ベンチマークを行う場合は、同一条件で、同じ日に、同じプロセスで比較してください。
+
+---
+
+## 3. 技術詳細
+
+### 3.1 アーキテクチャ
+
+```mermaid
+flowchart LR
+  subgraph ondevice ["On-device, first run"]
+    A[MLX safetensors<br/>int4 g64] -->|repack, bit-exact<br/>no requant| B[GGUF Q4_1<br/>edge0-tier.gguf]
+    A2[MLX LoRA] -->|fuse + adapt| B2[GGUF adapter]
+  end
+  B --> C[("pinned llama.cpp fork<br/>7ab4ee7 + 8 patches")]
+  B2 --> C
+  C -->|Vulkan: dense attn / GDN / norms| D[(GPU)]
+  C -->|"-cmoe: expert weights"| E[(CPU, mmap +<br/>L1 advisory pool)]
+  C --> F[llama-server<br/>127.0.0.1 OpenAI API]
+  F --> G[Tauri 2 shell<br/>supervisor + chat UI]
+```
+
+```mermaid
+flowchart TB
+  subgraph build ["Source-of-truth layout (open source)"]
+    U[github ggml-org/llama.cpp] -->|detached pin, never edited| V[vendor/llama.cpp @7ab4ee7]
+    P[patches/llama.cpp/common+windows] -->|git am --3way into per-platform worktree| W[wt/win]
+    V --> W
+    S[serve/*.cc edge0 code] -->|CMake glob patch #3| W
+  end
+```
+
+### 3.2 リポジトリマップ
+```
+serve/                 edge0 C++ pieces compiled into libllama (advisory prefetch router, memory budget)
+tools/                 MLX→GGUF converter chain (repack + LoRA adapter + catalog) and the throughput bench
+app/                   Tauri 2 shell (src-tauri/ Rust, src/ React) — see app/README.md
+scripts/               vendor-build.ps1 — engine assembly: patch replay into an isolated worktree
+```
+
+エンジンの変更はフォークではなくモノレポのルートにあります: `vendor/llama.cpp` はピン留めされた手付かずの上流チェックアウト(決してその場でパッチされない)であり、`patches/llama.cpp/` が 8 つのフックポイントパッチ + バンド README(`common/` と `windows/` はここ;コンパニオンアプリ用は `android/`)を保持します。`windows/` は意図的にネストした `patches/` のコピーを**一切**同梱しません — 信頼できる単一のソース、ドリフトなし。
+
+### 3.3 既知のギャップ(本リリース時点)
+
+- インストーラは**未署名**です(初回実行時に SmartScreen の警告)。自動アップデータはまだありません。
+- エンジンはまだインストーラ内にバンドルされていません — ポータブルビルドはエンジンビルドディレクトリを前提とします。バンドルはロードマップ上です。
+- `--pool-mb` の L1 プールは、最終的なメモリティア A/B を待つ間、Windows では環境変数レベルで休眠状態です。いずれにせよアプリは `--mem-budget-mb` をクランプします。
+- 16 GB 以下のマシンでの 35B は動作しますが、ページング律速です。ドキュメント化された下限は 8 GB で、メモリ上限と*組み合わせる*ことで動作し、tok/s は低下します。
+
+---
+
+*他のハードウェアでの issue やベンチ結果を歓迎します — `~\.edge0\logs\engine-*.log` と `doctor` ページのスクリーンショットを添付してください。*
diff --git a/windows/README_zh.md b/windows/README_zh.md
new file mode 100644
index 0000000..ec44c38
--- /dev/null
+++ b/windows/README_zh.md
@@ -0,0 +1,179 @@
+# edge0-windows — 发布构建指南
+
+[English](README.md) | 中文 | [日本語](README_ja.md) | [Español](README_es.md) | [Français](README_fr.md)
+
+本文档面向**从源码构建或评估 Windows App 的开发者**。内容包括：快速开始（构建 → 运行 → 测试）、参考机器上的性能实测，以及技术栈背后的技术选型。
+
+edge0 以 **monorepo** 形式发布 —— [`Edge0-AI/edge0`](https://github.com/Edge0-AI/edge0) —— 顶层目录存放共享的引擎供给（`vendor.llama.pin` + 由脚本物化的 `vendor/llama.cpp`、`patches/llama.cpp/` 补丁 band 集）与各平台子项目（`windows/` = 本 App，`android/` = 伙伴项目）。推理运行在固定版本（pinned）的上游 llama.cpp 上，以可重放补丁集的方式打补丁，稠密计算走 Vulkan，MoE 专家权重由 CPU 供给。
+
+第三方声明：见本目录下的 `NOTICE`。模型卡：[Edge0/Edge0-8B-A1B-preview](https://huggingface.co/Edge0/Edge0-8B-A1B-preview) · [Edge0/Edge0-35B-A3B-preview](https://huggingface.co/Edge0/Edge0-35B-A3B-preview)。
+
+---
+
+## 1. 快速开始
+
+### 1.1 环境要求
+
+| 工具链 | 要求 |
+|---|---|
+| 系统 | Windows 10 / 11 x64 |
+| Visual Studio 2022 | C++ 工作负载（MSVC） |
+| CMake | ≥ 3.21 |
+| Vulkan SDK | 用于 `ggml-vulkan`（任意 AMD/NVIDIA/Intel 独显；核显可用但更慢） |
+| Rust | stable 版，带 `cargo`（Tauri 2 外壳） |
+| Node.js | LTS（`npm`） |
+| Python | 3.10+ 且带 `numpy`（端上转换器 + 基准测试） |
+| 磁盘 / 内存 | ≥ 30 GB 可用磁盘；内存：8 GB+ 可运行 8B，16 GB+ 可运行 35B（受分页限制），参考机器为 48 GB |
+
+### 1.2 构建引擎（llama.cpp + 补丁 + edge0 serving 代码）
+
+```powershell
+git clone https://github.com/Edge0-AI/edge0
+cd edge0/windows
+pwsh -File scripts/vendor-build.ps1        # first build ≈ 10–20 min (compiles llama.cpp + Vulkan kernels)
+```
+
+脚本以两种方式解析引擎 depot，然后物化供给：
+
+- **monorepo 克隆（常规路径）** → 父目录带有 `vendor.llama.pin` + `patches/llama.cpp/{common,windows}`；脚本把 8 个补丁（`git am --3way`）重放到 `../wt/win` 的隔离构建 worktree（gitignored 构建区）—— vendor 树本身**绝不就地打补丁**；
+- `EDGE0_DEPOT` 环境变量 → 指向别处已有的 depot。
+
+llama.cpp 树**不是** submodule：首次运行时脚本会把上游克隆（blob 浅克隆）到 `../vendor/llama.cpp`，并 detach 到 `vendor.llama.pin` 固定的提交；此后该目录只是一个 gitignored 的 checkout，可用 `git fetch` 刷新。设置 `EDGE0_LLAMA_URL` 可从镜像克隆。
+
+`-AssembleOnly` 运行除编译外的全部步骤（数秒；验证补丁重放仍可应用且哈希匹配）。输出：
+
+```
+wt\win\build-vk\bin\Release\llama-server.exe   (+ llama.dll, ggml*.dll)
+```
+
+### 1.3 构建 App（Tauri 安装包 + 便携版 exe）
+
+```powershell
+cd app
+npm install
+npx tauri build
+```
+
+产物：
+
+```
+app\src-tauri\target\release\edge0-app.exe                              ← portable, no install
+app\src-tauri\target\release\bundle\nsis\edge0_0.1.0_x64-setup.exe      ← NSIS installer
+```
+
+外壳通过 `EDGE0_BIN_DIR` 定位引擎（默认：上文的 depot 引擎构建目录 —— 完整的覆盖项见 `app/README.md` 的环境变量表）。若引擎构建在别处，请设置该变量。
+
+### 1.4 模型
+
+无需预装任何东西：App 首次使用时从 HuggingFace 下载，然后在端上转换（一次性，8B 约 2–6 分钟）。逐文件 `sha256` 会对照已发布清单校验，并支持通过 HTTP Range 断点续传。
+
+手动预置（离线机器）：把 MLX 仓库文件放在
+
+```
+~\.edge0\models\edge0-8b\        (chat_template.jinja, model*.safetensors, lora_edge0_8b.safetensors, tokenizer…)
+~\.edge0\models\edge0-35b\       (…sharded safetensors, lora_edge0_35b.safetensors…)
+```
+
+转换器输出为 `models\edge0-<tier>-gguf\edge0-<tier>.gguf` + LoRA 适配器；转换是幂等的（sha 门控）。主目录为 `EDGE0_HOME`（默认 `~\.edge0`）。
+
+### 1.5 运行 App
+
+启动 `edge0-app.exe`（或安装 NSIS 包）。典型流程：
+
+1. **Models 页面** —— 选择 8B（快，约 5 GB）或 35B（约 21 GB）；*Download* → 自动 *Convert* → *Load*。
+2. **Chat 页面** —— 流式 Markdown 回答；引擎作为受监督的子进程运行，并随 App 一起被终止（Windows Job Object，`KILL_ON_JOB_CLOSE`）。
+3. **Doctor**（Models 页面底部）—— 一键健康检查：引擎二进制、模型文件、磁盘、环境。
+4. 排障：引擎/转换日志位于 `~\.edge0\logs\`。
+
+引擎 API（App 内部使用，任何 OpenAI 兼容客户端都可指向它）：`http://127.0.0.1:<port>/v1/chat/completions`，仅限回环。加载参数（进程契约）：`-ngl 99 -cmoe --ctx-size 8192 --flash-attn on --pool-mb <tier×RAM clamp> --mem-budget-mb …`。
+
+### 1.6 测试
+
+```powershell
+cd app\src-tauri
+cargo test                                   # fast suite: resume/probe logic against a local fake server
+
+# end-to-end, real model (downloads 4.5 GB if not seeded):
+cargo test --test real8b -- --ignored --nocapture
+
+# throughput reproduction (expects the table in §2, ±day-to-day drift):
+python ..\..\tools\r3_bench.py --tier 8b
+python ..\..\tools\r3_bench.py --tier 35b
+```
+
+`r3_bench.py` 以与 App 相同的契约解析路径：引擎取自 depot 构建目录（支持 `EDGE0_BIN_DIR` 覆盖），转换后的 GGUF 取自 App 的模型主目录（`EDGE0_HOME`，默认 `~\.edge0\models`）或存在时的仓库本地 `models/`；结果落在 `benchmarks/r3/`。
+
+```powershell
+# patch-replay gate (no compile, seconds):
+pwsh ..\..\scripts\vendor-build.ps1 -AssembleOnly
+```
+
+---
+
+## 2. 性能实测
+
+参考机器：**Intel i7-14700K · AMD Radeon RX 9070 GRE（Vulkan）· 48 GB DDR5 · Windows 11**，稳态（`-ngl 99 -cmoe`），3 轮基准取中位数（`tools/r3_bench.py`，约 240 token 混合语言 prompt，原生上下文）。
+
+| 模型 | prefill（tok/s） | decode（tok/s） | 240 token prompt 的 prefill |
+|---|---:|---:|---:|
+| edge0-8B-A1B | ~212 | ~44.8 | ~1.1 s |
+| edge0-35B-A3B | ~70 | ~27.7 | ~3.7 s |
+
+阅读这些数字时需要了解的背景：
+
+- decode 是 prompt 处理完成后的 *token 生成*速率；A1B / A3B 的激活参数量解释了为什么 35B 模型能以 27+ tok/s 解码。
+- 短 prompt 远快于上表的 240 token 列：App 对 8B 的首个真实请求，17 token 的 prompt **首 token 约 0.28 s**（引擎日志）。App 以 `--ctx-size 8192` 运行；上述基准使用原生上下文（131k / 262k）。
+- 这些是 48 GB 参考机器上的热内存实测（它是测试台，不代表产品假设）。在 16 GB / 8 GB 的产品档位上，专家经 mmap 缺页供给，解码速度随内存压力下降；进程级内存上限（`--mem-budget-mb`）会依据物理内存自动施加，以保持操作系统工作集合理（在压力下它确有可测量的*帮助*）。
+- 与未打补丁的原版 llama.cpp（跳过补丁 band 的 vendor 树）在同条件下 A/B 对比表明补丁集零开销（8B：decode 41.5 对 39.5；35B：27.2 对 27.4 —— 均在机器波动范围内）。
+- 这一级别的机器日间波动可达百分之几；如要跑基准，请在同一天、同一进程内做同类对比。
+
+---
+
+## 3. 技术细节
+
+### 3.1 架构
+
+```mermaid
+flowchart LR
+  subgraph ondevice ["On-device, first run"]
+    A[MLX safetensors<br/>int4 g64] -->|repack, bit-exact<br/>no requant| B[GGUF Q4_1<br/>edge0-tier.gguf]
+    A2[MLX LoRA] -->|fuse + adapt| B2[GGUF adapter]
+  end
+  B --> C[("pinned llama.cpp fork<br/>7ab4ee7 + 8 patches")]
+  B2 --> C
+  C -->|Vulkan: dense attn / GDN / norms| D[(GPU)]
+  C -->|"-cmoe: expert weights"| E[(CPU, mmap +<br/>L1 advisory pool)]
+  C --> F[llama-server<br/>127.0.0.1 OpenAI API]
+  F --> G[Tauri 2 shell<br/>supervisor + chat UI]
+```
+
+```mermaid
+flowchart TB
+  subgraph build ["Source-of-truth layout (open source)"]
+    U[github ggml-org/llama.cpp] -->|detached pin, never edited| V[vendor/llama.cpp @7ab4ee7]
+    P[patches/llama.cpp/common+windows] -->|git am --3way into per-platform worktree| W[wt/win]
+    V --> W
+    S[serve/*.cc edge0 code] -->|CMake glob patch #3| W
+  end
+```
+
+### 3.2 仓库结构
+```
+serve/                 edge0 C++ pieces compiled into libllama (advisory prefetch router, memory budget)
+tools/                 MLX→GGUF converter chain (repack + LoRA adapter + catalog) and the throughput bench
+app/                   Tauri 2 shell (src-tauri/ Rust, src/ React) — see app/README.md
+scripts/               vendor-build.ps1 — engine assembly: patch replay into an isolated worktree
+```
+
+引擎改动位于 monorepo 根目录，而非某个 fork：`vendor/llama.cpp` 是固定版本的原始上游 checkout（绝不就地打补丁），`patches/llama.cpp/` 承载 8 个 hook-point 补丁 + band README（此处为 `common/`、`windows/`；`android/` 供伙伴 App 使用）。`windows/` 刻意**不**携带嵌套的 `patches/` 副本 —— 单一事实来源，杜绝漂移。
+
+### 3.3 已知缺口（截至本版本）
+
+- 安装包**未签名**（首次运行会有 SmartScreen 警告）；尚无自动更新。
+- 引擎尚未打包进安装包 —— 便携版构建假定引擎构建目录存在；打包在路线图中。
+- `--pool-mb` L1 池在 Windows 上处于环境休眠状态，等待最终的内存档位 A/B；无论如何 App 都会钳制 `--mem-budget-mb`。
+- ≤16 GB 机器上的 35B 可用但受分页限制；文档化的下限是 8 GB *配合*内存上限工作，tok/s 有所降低。
+
+---
+
+*欢迎提交 issue 与其他硬件上的基准结果 —— 请附上 `~\.edge0\logs\engine-*.log` 与 `doctor` 页面截图。*
diff --git a/windows/app/README.md b/windows/app/README.md
new file mode 100644
index 0000000..3ba3e1f
--- /dev/null
+++ b/windows/app/README.md
@@ -0,0 +1,60 @@
+# edge0 Windows app — Tauri 2 shell
+
+The desktop shell: model download (into `~/.edge0`) → on-device MLX→GGUF conversion → supervised
+`llama-server` child process → streamed Markdown chat.
+
+## Dev run
+
+```powershell
+# 0) prerequisites: node + cargo installed; engine built (see ../scripts/vendor-build.ps1)
+cd app
+npm install
+npx tauri dev                 # window (vite:1420 + rust shell)
+```
+
+Headless end-to-end (isolated test home, re-runnable — Range-resume keeps it idempotent):
+
+```powershell
+cd app\src-tauri
+cargo test                                         # fast suite: fake-source Range resume tri-state
+cargo test --test real8b -- --ignored --nocapture   # full chain: 4.5 GB download -> convert -> engine -> chat
+```
+
+## Environment overrides (every path is overridable — test isolation is a first-class discipline)
+
+| env | default | purpose |
+|---|---|---|
+| `EDGE0_HOME` | `%USERPROFILE%\.edge0` | app home (models/state/logs/tmp) |
+| `EDGE0_REPO` | this repo checkout | where `tools/` conversion scripts live |
+| `EDGE0_PY` | `python` | interpreter for the converter chain |
+| `EDGE0_BIN_DIR` | depot engine build dir (`…/wt/win/build-vk/bin/Release`) | directory containing `llama-server.exe` |
+| `EDGE0_PHYS_MEM_GB` | detected (conservative) | override for the memory-tier clamp |
+| `EDGE0_TEST_SOURCE` | `modelscope` | source pin for headless tests (`huggingface` also supported) |
+
+## Home layout (`~/.edge0`)
+
+```
+models/edge0-<tier>/           download target (dir name is a hard contract of the converter;
+                               per-file sha256 verified before atomic rename)
+models/edge0-<tier>-gguf/      conversion output (gguf + adapter + manifest.json idempotency gate)
+state/downloads/<t>.json       resume checkpoints   state/verified-<t>.json per-file sha ledger
+state/models.json              installed-model registry
+logs/engine-*.log, logs/convert-*.log              engine & converter output
+tmp/<t>.<key>.part             in-flight download parts
+```
+
+## Proof-of-life telemetry
+
+We require positive proof that a mechanism is actually active before trusting a number
+(measured culture: past incidents taught us silent no-ops look identical to success):
+
+- after loading 35B, the engine log must show `POOL2 init` with the resolver line, or the UI
+  marks the pool as not-in-effect;
+- headless tests assert three things together: phase=ready + registered artifact sha + chat
+  returning real Markdown structure (tables/code blocks).
+
+## Packaging notes
+
+The installer currently ships the shell only; the engine and the Python converter are expected
+on the build/dev machine (env-overridable paths above). Bundling engine+converter into the
+installer, code signing, and an auto-updater are on the roadmap.
diff --git a/windows/app/index.html b/windows/app/index.html
new file mode 100644
index 0000000..b144af3
--- /dev/null
+++ b/windows/app/index.html
@@ -0,0 +1,12 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>edge0</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/main.tsx"></script>
+  </body>
+</html>
diff --git a/windows/app/package-lock.json b/windows/app/package-lock.json
new file mode 100644
index 0000000..9ea6713
--- /dev/null
+++ b/windows/app/package-lock.json
@@ -0,0 +1,4030 @@
+{
+  "name": "edge0-app-ui",
+  "version": "0.1.0",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {
+    "": {
+      "name": "edge0-app-ui",
+      "version": "0.1.0",
+      "dependencies": {
+        "@tauri-apps/api": "^2.11.1",
+        "lucide-react": "^1.46.0",
+        "react": "^19.3.0",
+        "react-dom": "^19.3.0",
+        "streamdown": "^2.6.0"
+      },
+      "devDependencies": {
+        "@tailwindcss/typography": "^0.5.20",
+        "@tailwindcss/vite": "^4.3.3",
+        "@tauri-apps/cli": "^2.11.4",
+        "@types/react": "^19.3.0",
+        "@types/react-dom": "^19.3.0",
+        "@vitejs/plugin-react": "^6.1.1",
+        "tailwindcss": "^4.3.3",
+        "typescript": "^7.0.2",
+        "vite": "^8.3.0"
+      }
+    },
+    "node_modules/@jridgewell/gen-mapping": {
+      "version": "0.3.13",
+      "resolved": "https://registry.npmmirror.com/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz",
+      "integrity": "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@jridgewell/sourcemap-codec": "^1.5.0",
+        "@jridgewell/trace-mapping": "^0.3.24"
+      }
+    },
+    "node_modules/@jridgewell/remapping": {
+      "version": "2.3.5",
+      "resolved": "https://registry.npmmirror.com/@jridgewell/remapping/-/remapping-2.3.5.tgz",
+      "integrity": "sha512-LI9u/+laYG4Ds1TDKSJW2YPrIlcVYOwi2fUC6xB43lueCjgxV4lffOCZCtYFiH6TNOX+tQKXx97T4IKHbhyHEQ==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@jridgewell/gen-mapping": "^0.3.5",
+        "@jridgewell/trace-mapping": "^0.3.24"
+      }
+    },
+    "node_modules/@jridgewell/resolve-uri": {
+      "version": "3.1.2",
+      "resolved": "https://registry.npmmirror.com/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz",
+      "integrity": "sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==",
+      "dev": true,
+      "license": "MIT",
+      "engines": {
+        "node": ">=6.0.0"
+      }
+    },
+    "node_modules/@jridgewell/sourcemap-codec": {
+      "version": "1.6.0",
+      "resolved": "https://registry.npmmirror.com/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.6.0.tgz",
+      "integrity": "sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==",
+      "dev": true,
+      "license": "MIT"
+    },
+    "node_modules/@jridgewell/trace-mapping": {
+      "version": "0.3.31",
+      "resolved": "https://registry.npmmirror.com/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz",
+      "integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@jridgewell/resolve-uri": "^3.1.0",
+        "@jridgewell/sourcemap-codec": "^1.4.14"
+      }
+    },
+    "node_modules/@oxc-project/types": {
+      "version": "0.151.0",
+      "resolved": "https://registry.npmmirror.com/@oxc-project/types/-/types-0.151.0.tgz",
+      "integrity": "sha512-J1yXrIlNDZVzE3ada310xeAw7nH8yCAyLPuUIsjKatFPmfn5bS1oW+cM+QsGOtVWd5nhSpbwZWx/rue+r5Z+PA==",
+      "dev": true,
+      "license": "MIT",
+      "funding": {
+        "url": "https://github.com/sponsors/oxc-project"
+      }
+    },
+    "node_modules/@rolldown/binding-android-arm-eabi": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-android-arm-eabi/-/binding-android-arm-eabi-1.2.11.tgz",
+      "integrity": "sha512-A5kXfGKvKWWZE0TtPrfsvT+q4Y5d1QG8gGUzpYjGydM+fARM9MuX90PrXYXe0XbsDVgyxxNzHo6giCj90bsFNw==",
+      "cpu": [
+        "arm"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "android"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-android-arm64": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-android-arm64/-/binding-android-arm64-1.2.11.tgz",
+      "integrity": "sha512-z6cTycz+iJ4PVkuL4HHW4DfTfoeU/2nqYYuSOrTmH7yHK5Y0LCOnA03V4ZNxavyVaU1oOqUgIg2klN/s+USGOA==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "android"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-darwin-arm64": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.2.11.tgz",
+      "integrity": "sha512-jShvqNtP6vDC6/A5JOAzbVV+DkgHqhl/ScVCJEbt+TUY6QYz7YnXcrg3sLtFBniro0f/Ld50ZwCWA6f7KYD1nQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-darwin-x64": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.2.11.tgz",
+      "integrity": "sha512-f2i2xiNWq1Z1l2++q2fuhZRdLAT3aqxD6vRNm1RAxpUoBcdqNB3C0s1Bt+K+PbEx2F5F4gQp6hqKkphCY/xF9w==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-freebsd-x64": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.2.11.tgz",
+      "integrity": "sha512-4Ir5FSOKIAMr4r0kExpt1s3bMgzJU3rA45AYOHtQpls0oNeqcYBKrWMlckrYH4KCfGLfkfn1tN1dmZPMVsdXow==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "freebsd"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-linux-arm-gnueabihf": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.2.11.tgz",
+      "integrity": "sha512-/gnRDM+39BROzAN/k1OZjDPnDMcZxB/0EUxKjONO5yVkNEvlsoMDrxGNKgZi/ttFriS2gwlDNzB65pvNbFOXIQ==",
+      "cpu": [
+        "arm"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-linux-arm64-gnu": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.2.11.tgz",
+      "integrity": "sha512-PFaK8HwvAHbaKbBcDNQihjMKYvFnA5hiENx/l5tphTDz1E0WFp32l0A7aq7lyUwGsRw/xSrNIy/gIK4thrSCrw==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-linux-arm64-musl": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.2.11.tgz",
+      "integrity": "sha512-AskzJUIKRLPxkruR1wLKewGbOw+EYfU/9lOrBFj4AFrEA8hPpKFnODWNu2WLaNs0QNkEb9QIJufmVZZIL/bJlg==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-linux-ppc64-gnu": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.2.11.tgz",
+      "integrity": "sha512-qlUGAheh2yh8afH7QBgx0PrRHN85hKnNd78x8MeMhXivuevgd8vgf6/CstOzmNKY/lLTHvNTrPy98cLnAugzJw==",
+      "cpu": [
+        "ppc64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-linux-s390x-gnu": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.2.11.tgz",
+      "integrity": "sha512-secpEad+0vCbSfn8upFySkDskv+bGPk3THSDS9Y89yc4rb4kzqHp8Dmyd9BkQW4SnhNXBZCl/6CrO//hZahNJQ==",
+      "cpu": [
+        "s390x"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-linux-x64-gnu": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.2.11.tgz",
+      "integrity": "sha512-mOVBT3dPpkWm8XBWPmU4bf+U6dYDLeMo/9ojUmis4N0L5uu10qra5vOyngZ7/PSdoE4G9KvRt4bloRxNjLas7A==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-linux-x64-musl": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.2.11.tgz",
+      "integrity": "sha512-Is78i9A8Ui4SqcxUwFJ9uMmjDn58IbVTjFWYdQestFEgeuEmHMLGNriXnVJKkwG2YiZjw8cP0zCTyDMdDGtOOg==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-openharmony-arm64": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.2.11.tgz",
+      "integrity": "sha512-dUCXneZ87INUMyQ0D+C0HrEBNUPNXHaPmU5GTjyKTJEiussw9Kaj5Ln8UztPe4epV/ffvgNBEadksdYhmW6xJA==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "openharmony"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-win32-arm64-msvc": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.2.11.tgz",
+      "integrity": "sha512-jByxb6qfd+bH1xUd0qnfFnb17i9sWBPY2tOavJ0l3tdr3OTu+Kvtm8cd/JV5nFt657b1VqGltxg9olOEfofXWw==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/binding-win32-x64-msvc": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.2.11.tgz",
+      "integrity": "sha512-/PzKqzAJ03i19oy2ItPvyvaVjOjBCNnfaJs8yvUdGBKmiESgnrJSQ2awd81QzFbbnAmu7YO9ZnJrDCb9VSJPRA==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      }
+    },
+    "node_modules/@rolldown/pluginutils": {
+      "version": "1.0.1",
+      "resolved": "https://registry.npmmirror.com/@rolldown/pluginutils/-/pluginutils-1.0.1.tgz",
+      "integrity": "sha512-2j9bGt5Jh8hj+vPtgzPtl72j0yRxHAyumoo6TNfAjsLB04UtpSvPbPcDcBMxz7n+9CYB0c1GxQFxYRg2jimqGw==",
+      "dev": true,
+      "license": "MIT"
+    },
+    "node_modules/@tailwindcss/node": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/node/-/node-4.3.3.tgz",
+      "integrity": "sha512-/T8IKEsf9VTU6tLjgC7+sv2mOPtQxzE2jMw7u4Tt40Tx+QSZxpzh95/H6cMKoja9XuW7iMdLJYBB0o9G1CaAgg==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@jridgewell/remapping": "^2.3.5",
+        "enhanced-resolve": "^5.24.1",
+        "jiti": "^2.7.0",
+        "lightningcss": "1.32.0",
+        "magic-string": "^0.30.21",
+        "source-map-js": "^1.2.1",
+        "tailwindcss": "4.3.3"
+      }
+    },
+    "node_modules/@tailwindcss/oxide": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide/-/oxide-4.3.3.tgz",
+      "integrity": "sha512-krXjAikiaFSPaK/FkAQT5UTx3VormQaiZ5hBFlJZ9UFQGB/rwg1MZIhHAG9smMQRTdyJxP6Qt5MwMtdyU5FWrA==",
+      "dev": true,
+      "license": "MIT",
+      "engines": {
+        "node": ">= 20"
+      },
+      "optionalDependencies": {
+        "@tailwindcss/oxide-android-arm64": "4.3.3",
+        "@tailwindcss/oxide-darwin-arm64": "4.3.3",
+        "@tailwindcss/oxide-darwin-x64": "4.3.3",
+        "@tailwindcss/oxide-freebsd-x64": "4.3.3",
+        "@tailwindcss/oxide-linux-arm-gnueabihf": "4.3.3",
+        "@tailwindcss/oxide-linux-arm64-gnu": "4.3.3",
+        "@tailwindcss/oxide-linux-arm64-musl": "4.3.3",
+        "@tailwindcss/oxide-linux-x64-gnu": "4.3.3",
+        "@tailwindcss/oxide-linux-x64-musl": "4.3.3",
+        "@tailwindcss/oxide-wasm32-wasi": "4.3.3",
+        "@tailwindcss/oxide-win32-arm64-msvc": "4.3.3",
+        "@tailwindcss/oxide-win32-x64-msvc": "4.3.3"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-android-arm64": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-android-arm64/-/oxide-android-arm64-4.3.3.tgz",
+      "integrity": "sha512-Y85A2gmPSkl5Ve5qR86GL4HT509cFqQh1aes9p3sSkyTPwt0Pppf3GkwGe4JPACcRYjgJIEhQgM6dBClnr0NYw==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "android"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-darwin-arm64": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-darwin-arm64/-/oxide-darwin-arm64-4.3.3.tgz",
+      "integrity": "sha512-BiaWatpBcERQFDlOjRDpIVXuFK5PJez5SA4JMg6VYZdBYU+qKfV/vqjcIs+IYmtitf1xYQZTwXvU/8y4lfZUGw==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-darwin-x64": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-darwin-x64/-/oxide-darwin-x64-4.3.3.tgz",
+      "integrity": "sha512-fAeUqfV5ndhxRwai8cXGzdLvul9utWOmeTkv69unv4ZXixjn61Z+p9lCWdwOwA3TYboG3BwdVuN/RDjhBRl0mw==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-freebsd-x64": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-freebsd-x64/-/oxide-freebsd-x64-4.3.3.tgz",
+      "integrity": "sha512-iyf5bV6+wnAlflVeEy7R25dupxTNECZN5QMI0qNT6eT+EgaGdZcKhGkr5SdoaWiLJ3spLqIY9VCeSGrwmtg4kw==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "freebsd"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-linux-arm-gnueabihf": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-linux-arm-gnueabihf/-/oxide-linux-arm-gnueabihf-4.3.3.tgz",
+      "integrity": "sha512-aAYUprJAJQWWbRrPvtjdroZ56Md+JM8pMiopS6xGEwDfLhqj+2ver2p4nU4Mb3CRqcMmNBjo8KkUgcxhkzVQGQ==",
+      "cpu": [
+        "arm"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-linux-arm64-gnu": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-linux-arm64-gnu/-/oxide-linux-arm64-gnu-4.3.3.tgz",
+      "integrity": "sha512-nDxldcEENOxZRzC2uu9jrutZdAAQtb+8WWDCSnWL1zvBk1+FN+x6MtDViPB5AJMfttVCUhehGWus3XBPgatM/w==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-linux-arm64-musl": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-linux-arm64-musl/-/oxide-linux-arm64-musl-4.3.3.tgz",
+      "integrity": "sha512-Md44bD6veX/PC5iyF8cDVnw4HBIANZepRZZ7a8DQOvkfo5WUBwcp6iAuCUz23u+4SUkhJlD3eL7hNdW8ezd/kA==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-linux-x64-gnu": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-linux-x64-gnu/-/oxide-linux-x64-gnu-4.3.3.tgz",
+      "integrity": "sha512-tx7us1muwOKAKWao2v/GaafFeQboE6aj88vC6ziN2NCGcRm8gWUhwjzg+YdVB1e4boAtdtma4L43onunI6NS4w==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-linux-x64-musl": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-linux-x64-musl/-/oxide-linux-x64-musl-4.3.3.tgz",
+      "integrity": "sha512-SJxX60smvHgasZoBy11dX6YRjXJFovwWBoedhbQPOBzgFWBHGB+TVPWB9BxzR7TTxU8FQZAI2AyiNCMzFm8Img==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-wasm32-wasi": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-wasm32-wasi/-/oxide-wasm32-wasi-4.3.3.tgz",
+      "integrity": "sha512-jx1+rPhY/5Ympkktd656HBWEBLxP7dH06losBLjjf5vgCODXvi9KhtftWcMIwTFIDqBr7cRnQkdLnAG+IOlGvQ==",
+      "bundleDependencies": [
+        "@napi-rs/wasm-runtime",
+        "@emnapi/core",
+        "@emnapi/runtime",
+        "@tybys/wasm-util",
+        "@emnapi/wasi-threads",
+        "tslib"
+      ],
+      "cpu": [
+        "wasm32"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "dependencies": {
+        "@emnapi/core": "^1.11.1",
+        "@emnapi/runtime": "^1.11.1",
+        "@emnapi/wasi-threads": "^1.2.2",
+        "@napi-rs/wasm-runtime": "^1.1.4",
+        "@tybys/wasm-util": "^0.10.2",
+        "tslib": "^2.8.1"
+      },
+      "engines": {
+        "node": ">=14.0.0"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-win32-arm64-msvc": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.3.3.tgz",
+      "integrity": "sha512-3rc292Ca2ceK6Ulcc/bAVnTs/3nDtoPhyEKlgPv+yQJQi/JS/AMJlqzxvlDacL1nekbrcf6bTqp/jV4qgnPxNQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/oxide-win32-x64-msvc": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/oxide-win32-x64-msvc/-/oxide-win32-x64-msvc-4.3.3.tgz",
+      "integrity": "sha512-yJ0pwIVc/nYeGoV02WtsN8KYyLQv7kyI2wDnkezyJlGGjkd4QLwDGAwl47YpPJeuI0M0ObaXGSPjvWDPeTPggw==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/@tailwindcss/typography": {
+      "version": "0.5.20",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/typography/-/typography-0.5.20.tgz",
+      "integrity": "sha512-hwbzQuNUfcPvbegQFatVPl/MY/tcM9KLl963hQ5laJKPh81TEZ1+dNG9PirGvcaDBkp+BCshExAyKVPW91dozw==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "postcss-selector-parser": "6.0.10"
+      },
+      "peerDependencies": {
+        "tailwindcss": ">=3.0.0 || >=4.0.0 || insiders"
+      }
+    },
+    "node_modules/@tailwindcss/vite": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/@tailwindcss/vite/-/vite-4.3.3.tgz",
+      "integrity": "sha512-yYU8cogLeSh/ms2jh8Fj7jaba/EWa7Ja6GoUqYZaraEuCI5YS6ms6ObZgjjedm+jm6XZjdNRWBpPP6Z86oOxcw==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@tailwindcss/node": "4.3.3",
+        "@tailwindcss/oxide": "4.3.3",
+        "tailwindcss": "4.3.3"
+      },
+      "peerDependencies": {
+        "vite": "^5.2.0 || ^6 || ^7 || ^8"
+      }
+    },
+    "node_modules/@tauri-apps/api": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/api/-/api-2.12.0.tgz",
+      "integrity": "sha512-fUSxFX8VABYe9W9K5ROdEaKVkNUFRcMSRofjruvaxztXA+I0XnQHcYHF8m/a6iExRUxsUf0HsFWhPBuH9IgzJw==",
+      "license": "Apache-2.0 OR MIT",
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/tauri"
+      }
+    },
+    "node_modules/@tauri-apps/cli": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli/-/cli-2.12.0.tgz",
+      "integrity": "sha512-GIShnAR4C16kSTA8NO02fo8fOw/I5z5sAUEB5S2HlPHFnBkbihL4c9a/MpkxOm1+q5S+tFd3X4QUjHWxfyGzpQ==",
+      "dev": true,
+      "license": "Apache-2.0 OR MIT",
+      "bin": {
+        "tauri": "tauri.js"
+      },
+      "engines": {
+        "node": ">= 10"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/tauri"
+      },
+      "optionalDependencies": {
+        "@tauri-apps/cli-darwin-arm64": "2.12.0",
+        "@tauri-apps/cli-darwin-x64": "2.12.0",
+        "@tauri-apps/cli-linux-arm-gnueabihf": "2.12.0",
+        "@tauri-apps/cli-linux-arm64-gnu": "2.12.0",
+        "@tauri-apps/cli-linux-arm64-musl": "2.12.0",
+        "@tauri-apps/cli-linux-riscv64-gnu": "2.12.0",
+        "@tauri-apps/cli-linux-x64-gnu": "2.12.0",
+        "@tauri-apps/cli-linux-x64-musl": "2.12.0",
+        "@tauri-apps/cli-win32-arm64-msvc": "2.12.0",
+        "@tauri-apps/cli-win32-ia32-msvc": "2.12.0",
+        "@tauri-apps/cli-win32-x64-msvc": "2.12.0"
+      }
+    },
+    "node_modules/@tauri-apps/cli-darwin-arm64": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-darwin-arm64/-/cli-darwin-arm64-2.12.0.tgz",
+      "integrity": "sha512-v5J+jISihLGtih1+3yWagbsL+5dbLZORu4vaSI0vhY0Heb+MAW+T/1s1RyKxrQrwEPGBG5kqwfDDTnSbmVkuqA==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-darwin-x64": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-darwin-x64/-/cli-darwin-x64-2.12.0.tgz",
+      "integrity": "sha512-g27I07pfW/dMS8vcNqB9lamgKF0TqLlcy8e7t9qDGsB1iaeyjJWb2TgBF4tM9ZwLqPAZUXbsz28oXljO12WC4g==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-linux-arm-gnueabihf": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-linux-arm-gnueabihf/-/cli-linux-arm-gnueabihf-2.12.0.tgz",
+      "integrity": "sha512-gXyQ710fNc99YuxFxNyeX9PLSrIpbtgBBj7GQo9tSeRT6fIK40q3/gm7X5jqZo7mWLR2m0R/2ALd3vAWo0Z32Q==",
+      "cpu": [
+        "arm"
+      ],
+      "dev": true,
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-linux-arm64-gnu": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-linux-arm64-gnu/-/cli-linux-arm64-gnu-2.12.0.tgz",
+      "integrity": "sha512-H9cyXuq9WgxoHh2NZ+R+oN+W78x5UwaBDEp0/fmqUUkaBoioCuOZErM9xKiIV6bpNAhmw1RoWZ49u37q/vaiAg==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-linux-arm64-musl": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-linux-arm64-musl/-/cli-linux-arm64-musl-2.12.0.tgz",
+      "integrity": "sha512-5tFuk3JKMBa1M/iN4gN32OYlSfhY3aqb5PO3zCugxaT+sfecyMKpc3EM9VDbhmwWBE9QD4Tfj7wGhlqx2yy2pQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-linux-riscv64-gnu": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-linux-riscv64-gnu/-/cli-linux-riscv64-gnu-2.12.0.tgz",
+      "integrity": "sha512-LOrMyV+nb7pKq8P0BZtTBQuyzm1ORjn16R9MlupmxuT5LqxtuYJlw3VMwbaeXCToJ8jbP5u9JnkfmbqDlvETXw==",
+      "cpu": [
+        "riscv64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-linux-x64-gnu": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-linux-x64-gnu/-/cli-linux-x64-gnu-2.12.0.tgz",
+      "integrity": "sha512-aMIkI+X9qvD2+k80QZB3tuOYT41/cEheikC510CwAIZ31KLXORsRASYYzM4mi1/58ICmtzSLFhCC08keir7COw==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-linux-x64-musl": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-linux-x64-musl/-/cli-linux-x64-musl-2.12.0.tgz",
+      "integrity": "sha512-9gl3NhUbfuGhSKjeGCypKSdtoezOQRcDUrkhTWtLQEjnajets6peO6NadgI2DF0OzbiiFOzrnnIGMwBS+Ap6Iw==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-win32-arm64-msvc": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-win32-arm64-msvc/-/cli-win32-arm64-msvc-2.12.0.tgz",
+      "integrity": "sha512-1W+v86wsYLuyBDN0iDessD4CD9o+yNyO0m8Dp3+/vfjiZrJ/v7R7/aTRBeikaq2tLGgxeag0gxnFggBNBGw91A==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-win32-ia32-msvc": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-win32-ia32-msvc/-/cli-win32-ia32-msvc-2.12.0.tgz",
+      "integrity": "sha512-SsBjvq4f7YaBGqcWoE3cycoItPf0SrQeqFqWWgfxoeZWAIz01E31LDSeZ2p/etJiIY0R9oi045chkFfb0O7WGw==",
+      "cpu": [
+        "ia32"
+      ],
+      "dev": true,
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@tauri-apps/cli-win32-x64-msvc": {
+      "version": "2.12.0",
+      "resolved": "https://registry.npmmirror.com/@tauri-apps/cli-win32-x64-msvc/-/cli-win32-x64-msvc-2.12.0.tgz",
+      "integrity": "sha512-9zzDKYUH4TKKBgRdDdVQX8fiRO5gJ8hD1nGsHU9DOHG5Jpm3arjbu9Xg1frR20EusG3Sl8X0dppoO/4EoUXUGg==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0 OR MIT",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 10"
+      }
+    },
+    "node_modules/@types/debug": {
+      "version": "4.1.13",
+      "resolved": "https://registry.npmmirror.com/@types/debug/-/debug-4.1.13.tgz",
+      "integrity": "sha512-KSVgmQmzMwPlmtljOomayoR89W4FynCAi3E8PPs7vmDVPe84hT+vGPKkJfThkmXs0x0jAaa9U8uW8bbfyS2fWw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/ms": "*"
+      }
+    },
+    "node_modules/@types/estree": {
+      "version": "1.0.9",
+      "resolved": "https://registry.npmmirror.com/@types/estree/-/estree-1.0.9.tgz",
+      "integrity": "sha512-GhdPgy1el4/ImP05X05Uw4cw2/M93BCUmnEvWZNStlCzEKME4Fkk+YpoA5OiHNQmoS7Cafb8Xa3Pya8m1Qrzeg==",
+      "license": "MIT"
+    },
+    "node_modules/@types/estree-jsx": {
+      "version": "1.0.5",
+      "resolved": "https://registry.npmmirror.com/@types/estree-jsx/-/estree-jsx-1.0.5.tgz",
+      "integrity": "sha512-52CcUVNFyfb1A2ALocQw/Dd1BQFNmSdkuC3BkZ6iqhdMfQz7JWOFRuJFloOzjk+6WijU56m9oKXFAXc7o3Towg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/estree": "*"
+      }
+    },
+    "node_modules/@types/hast": {
+      "version": "3.0.5",
+      "resolved": "https://registry.npmmirror.com/@types/hast/-/hast-3.0.5.tgz",
+      "integrity": "sha512-rp/ezSWaD1m44dPKICGhiskI13nVr7qTloFwDa/IYkhhf5nzwP+zIQcIJh3WIFSBOy/H1PzB40jPjMDksN4F+g==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "*"
+      }
+    },
+    "node_modules/@types/mdast": {
+      "version": "4.0.4",
+      "resolved": "https://registry.npmmirror.com/@types/mdast/-/mdast-4.0.4.tgz",
+      "integrity": "sha512-kGaNbPh1k7AFzgpud/gMdvIm5xuECykRR+JnWKQno9TAXVa6WIVCGTPvYGekIDL4uwCZQSYbUxNBSb1aUo79oA==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "*"
+      }
+    },
+    "node_modules/@types/ms": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/@types/ms/-/ms-2.1.0.tgz",
+      "integrity": "sha512-GsCCIZDE/p3i96vtEqx+7dBUGXrc7zeSK3wwPHIaRThS+9OhWIXRqzs4d6k1SVU8g91DrNRWxWUGhp5KXQb2VA==",
+      "license": "MIT"
+    },
+    "node_modules/@types/react": {
+      "version": "19.3.0",
+      "resolved": "https://registry.npmmirror.com/@types/react/-/react-19.3.0.tgz",
+      "integrity": "sha512-N0rFCuH9YoxG9/m61l9MfpJKfmLOVU0em7ipIz6TRgSSkvReLB9vL85GB+yr8Bs5leqpvg96JSwF4ZS1s4viQg==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "csstype": "^3.2.2"
+      }
+    },
+    "node_modules/@types/react-dom": {
+      "version": "19.3.0",
+      "resolved": "https://registry.npmmirror.com/@types/react-dom/-/react-dom-19.3.0.tgz",
+      "integrity": "sha512-ZI7bU42mZXXKHn/qNLEw2IrbiINU7X5+vfgdixBHkCNpYWXjKgfQ/P+uyGb5CjOLB9UcnTeg3rylQtV2hym44Q==",
+      "dev": true,
+      "license": "MIT",
+      "peerDependencies": {
+        "@types/react": "^19.3.0"
+      }
+    },
+    "node_modules/@types/unist": {
+      "version": "3.0.3",
+      "resolved": "https://registry.npmmirror.com/@types/unist/-/unist-3.0.3.tgz",
+      "integrity": "sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q==",
+      "license": "MIT"
+    },
+    "node_modules/@typescript/typescript-aix-ppc64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-aix-ppc64/-/typescript-aix-ppc64-7.0.2.tgz",
+      "integrity": "sha512-MTKKkWB7p/0E9xi1d1tHtZ5PiLkGEMIq88pK2CubZjOsLtYTLqhgIgi6zepFa+9GHZ6h05NMCkQxGKiPXMxXtQ==",
+      "cpu": [
+        "ppc64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "aix"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-darwin-arm64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-darwin-arm64/-/typescript-darwin-arm64-7.0.2.tgz",
+      "integrity": "sha512-gowzar9MwS/aRWp6f3a4KUqzRjAZjOsmGNCM6LcTgXum+dBfgsBVMN+AgvOCCbguXyick6LJhpBszxMebJ8syA==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-darwin-x64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-darwin-x64/-/typescript-darwin-x64-7.0.2.tgz",
+      "integrity": "sha512-SZ9xZInqApNlNGc9s0W1VSsktYSOe9cFqNOIqmN1Gs8SmkjKZYFt017G4VwPxASInODuAdbTW7sXiFUf893RgA==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-freebsd-arm64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-freebsd-arm64/-/typescript-freebsd-arm64-7.0.2.tgz",
+      "integrity": "sha512-W5NH4y/J0plIIS5b2xvTEkU7JFxyqdMAOgf+Ilhl0vHQXKO5dZoxd+C/jEtq56c4F3wk71RB4BMRQ2XdI+bwYQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "freebsd"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-freebsd-x64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-freebsd-x64/-/typescript-freebsd-x64-7.0.2.tgz",
+      "integrity": "sha512-UMGDx5sTpzNw3WiPebH7l90IWfJggEd+egHt/q6p7/Cm3zqoV7VxkGXt+3DxPIw8CcmvAB0j3sVVfbhX+M4Tpw==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "freebsd"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-arm": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-arm/-/typescript-linux-arm-7.0.2.tgz",
+      "integrity": "sha512-gffT3xPz9sR7j/YJExkyPntrI0P2EP9XbOyWzth2/Gs0RstK+90RBcO0ncXoXy/beYll1SXw846Nf2zdnEz0QQ==",
+      "cpu": [
+        "arm"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-arm64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-arm64/-/typescript-linux-arm64-7.0.2.tgz",
+      "integrity": "sha512-Qh4eU4/y3yDjnfjjyPYihMj5/ODIlmt+Bzu17OI+fiSRDW57QmU5SiN63exPRNJPKUzcc1INa1NXdrJ+MqHjUQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-loong64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-loong64/-/typescript-linux-loong64-7.0.2.tgz",
+      "integrity": "sha512-uEHck9i8hoAzXPiYRib1O7miOnz23SxIeVl6F4LXox+qov1K35jHcEW6VHKvZI+pyvl7fZEP4MCU5LYvIq1GuQ==",
+      "cpu": [
+        "loong64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-mips64el": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-mips64el/-/typescript-linux-mips64el-7.0.2.tgz",
+      "integrity": "sha512-R4KvAMnE43W5Qeqb0Ly56O3mWMWIAgsMyz36DCaycd5nbg/9kzm0liw3JocfRqyJY0KPmzFjbswozXyW0DnIYA==",
+      "cpu": [
+        "mips64el"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-ppc64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-ppc64/-/typescript-linux-ppc64-7.0.2.tgz",
+      "integrity": "sha512-DORx5b3sd/4S7eayxm4FQv+A7CrkUIGRaHiwI8oiHTAI1fAPWhF4J0vAlkC8biAlHSVVwxMQ3tjZ2/DVbnQiiA==",
+      "cpu": [
+        "ppc64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-riscv64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-riscv64/-/typescript-linux-riscv64-7.0.2.tgz",
+      "integrity": "sha512-wf0jqEDOjrPRnKwYRyyJDRo11KMbvMFrU+q4zqKyChODBzvlkbhNQfKvLxQCcwTpdDaXSHZTVuh0JoCrKCUMHQ==",
+      "cpu": [
+        "riscv64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-s390x": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-s390x/-/typescript-linux-s390x-7.0.2.tgz",
+      "integrity": "sha512-IkwJc3L7yhytWd/ewjyxNDfOmswCm9GWMJT/ue/dU4aZNbwZeYAetq42VyLmsmSjvoX7z74X6ZaYCtzAr0EuGw==",
+      "cpu": [
+        "s390x"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-linux-x64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-linux-x64/-/typescript-linux-x64-7.0.2.tgz",
+      "integrity": "sha512-EYdf2cNg7rgCWJnxCdJ+F3V39O8ihb37eHAu1LK8oAFizgTQbPOK7zHHXbPt8rX24COqODXeI3sIf0fCXG7H/A==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-netbsd-arm64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-netbsd-arm64/-/typescript-netbsd-arm64-7.0.2.tgz",
+      "integrity": "sha512-+polYF4MF04aPpO5FTkHran9yUQDSXqy5GiSDKpsll5jy3l3+g9QLhpf39T+ePtefhXLOGrLl0QIjkQP6VnelA==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "netbsd"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-netbsd-x64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-netbsd-x64/-/typescript-netbsd-x64-7.0.2.tgz",
+      "integrity": "sha512-8YIT0EHM/3dq10ZOVF/A7pc/YSMtbcecct4rWtexrnSCHOPcpC2KTLXfTCR6vDpnSiY12heNb1GiN/wu+T/FyA==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "netbsd"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-openbsd-arm64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-openbsd-arm64/-/typescript-openbsd-arm64-7.0.2.tgz",
+      "integrity": "sha512-APT8+ClYnuYm1u9+kgGXoMj2VzWzcymwh2gNSQVySHfkRDGOTVkoWLjCmOQSaO+PoqQ57B0flRp9SA+7GnnkzQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "openbsd"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-openbsd-x64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-openbsd-x64/-/typescript-openbsd-x64-7.0.2.tgz",
+      "integrity": "sha512-yX7s+Q0Dln0Dt9tEzZsAjXXR/+ytBM7AlglaqyeMPxQszJ1JhlJdZ6jLA+IzldHtflX81em7lDao1xXu+aRRkg==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "openbsd"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-sunos-x64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-sunos-x64/-/typescript-sunos-x64-7.0.2.tgz",
+      "integrity": "sha512-dLJDGaLZ1D4HPQn62u1n8mBDkJREwMsAkCdkwd4Ieqw+x3TUyTsqY0YiBCtE6H6OzzgGk3iuZ3vFWRS+E8/d1g==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "sunos"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-win32-arm64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-win32-arm64/-/typescript-win32-arm64-7.0.2.tgz",
+      "integrity": "sha512-Gyl1Vy6OsWesLzmq+EP0Fb7b4Nid5232AvcA2SFcdYreldpNtYFFofPjnt62y9hQy7VTaZp65ICJjuAQRaVcIQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@typescript/typescript-win32-x64": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/@typescript/typescript-win32-x64/-/typescript-win32-x64-7.0.2.tgz",
+      "integrity": "sha512-0BQ3HkAHHlKLSp1qRvf3SUhGpGsDuhB/jgFw75guyqbxJqEaS0Cw/VFO8i2nHglJUzQCRtMMR/IBAKE3ETMC4g==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "Apache-2.0",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">=16.20.0"
+      }
+    },
+    "node_modules/@ungap/structured-clone": {
+      "version": "1.4.0",
+      "resolved": "https://registry.npmmirror.com/@ungap/structured-clone/-/structured-clone-1.4.0.tgz",
+      "integrity": "sha512-1mEZtMKPM09vDmQt5y7YvmN2+DFTP7Tg0EWXdic8/C6VRnpb33e4ghisCIE3WZjsE2N8mf+QV1Zqh7ZFYLWInQ==",
+      "license": "ISC"
+    },
+    "node_modules/@vitejs/plugin-react": {
+      "version": "6.1.1",
+      "resolved": "https://registry.npmmirror.com/@vitejs/plugin-react/-/plugin-react-6.1.1.tgz",
+      "integrity": "sha512-yxLaQV9gkhS8ezJqCM6+ndU7mDY6gqAg75NQ+0IjwEI8IYOmQCgkRwHKVSfWXW076DsqMo0Dk+0FK1U+M5RgFw==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@rolldown/pluginutils": "^1.0.1"
+      },
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      },
+      "peerDependencies": {
+        "@rolldown/plugin-babel": "^0.1.7 || ^0.2.0",
+        "babel-plugin-react-compiler": "^1.0.0",
+        "oxc-transform-react": "^0.145.0",
+        "vite": "^8.0.0"
+      },
+      "peerDependenciesMeta": {
+        "@rolldown/plugin-babel": {
+          "optional": true
+        },
+        "babel-plugin-react-compiler": {
+          "optional": true
+        },
+        "oxc-transform-react": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/bail": {
+      "version": "2.0.2",
+      "resolved": "https://registry.npmmirror.com/bail/-/bail-2.0.2.tgz",
+      "integrity": "sha512-0xO6mYd7JB2YesxDKplafRpsiOzPt9V02ddPCLbY1xYGPOX24NTyN50qnUxgCPcSoYMhKpAuBTjQoRZCAkUDRw==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/ccount": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/ccount/-/ccount-2.0.1.tgz",
+      "integrity": "sha512-eyrF0jiFpY+3drT6383f1qhkbGsLSifNAjA61IUjZjmLCWjItY6LB9ft9YhoDgwfmclB2zhu51Lc7+95b8NRAg==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/character-entities": {
+      "version": "2.0.2",
+      "resolved": "https://registry.npmmirror.com/character-entities/-/character-entities-2.0.2.tgz",
+      "integrity": "sha512-shx7oQ0Awen/BRIdkjkvz54PnEEI/EjwXDSIZp86/KKdbafHh1Df/RYGBhn4hbe2+uKC9FnT5UCEdyPz3ai9hQ==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/character-entities-html4": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/character-entities-html4/-/character-entities-html4-2.1.0.tgz",
+      "integrity": "sha512-1v7fgQRj6hnSwFpq1Eu0ynr/CDEw0rXo2B61qXrLNdHZmPKgb7fqS1a2JwF0rISo9q77jDI8VMEHoApn8qDoZA==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/character-entities-legacy": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmmirror.com/character-entities-legacy/-/character-entities-legacy-3.0.0.tgz",
+      "integrity": "sha512-RpPp0asT/6ufRm//AJVwpViZbGM/MkjQFxJccQRHmISF/22NBtsHqAWmL+/pmkPWoIUJdWyeVleTl1wydHATVQ==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/character-reference-invalid": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/character-reference-invalid/-/character-reference-invalid-2.0.1.tgz",
+      "integrity": "sha512-iBZ4F4wRbyORVsu0jPV7gXkOsGYjGHPmAyv+HiHG8gi5PtC9KI2j1+v8/tlibRvjoWX027ypmG/n0HtO5t7unw==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/clsx": {
+      "version": "2.1.1",
+      "resolved": "https://registry.npmmirror.com/clsx/-/clsx-2.1.1.tgz",
+      "integrity": "sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=6"
+      }
+    },
+    "node_modules/comma-separated-tokens": {
+      "version": "2.0.3",
+      "resolved": "https://registry.npmmirror.com/comma-separated-tokens/-/comma-separated-tokens-2.0.3.tgz",
+      "integrity": "sha512-Fu4hJdvzeylCfQPp9SGWidpzrMs7tTrlu6Vb8XGaRGck8QSNZJJp538Wrb60Lax4fPwR64ViY468OIUTbRlGZg==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/cssesc": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmmirror.com/cssesc/-/cssesc-3.0.0.tgz",
+      "integrity": "sha512-/Tb/JcjK111nNScGob5MNtsntNM1aCNUDipB/TkwZFhyDrrE47SOx/18wF2bbjgc3ZzCSKW1T5nt5EbFoAz/Vg==",
+      "dev": true,
+      "license": "MIT",
+      "bin": {
+        "cssesc": "bin/cssesc"
+      },
+      "engines": {
+        "node": ">=4"
+      }
+    },
+    "node_modules/csstype": {
+      "version": "3.2.3",
+      "resolved": "https://registry.npmmirror.com/csstype/-/csstype-3.2.3.tgz",
+      "integrity": "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==",
+      "dev": true,
+      "license": "MIT"
+    },
+    "node_modules/debug": {
+      "version": "4.4.3",
+      "resolved": "https://registry.npmmirror.com/debug/-/debug-4.4.3.tgz",
+      "integrity": "sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==",
+      "license": "MIT",
+      "dependencies": {
+        "ms": "^2.1.3"
+      },
+      "engines": {
+        "node": ">=6.0"
+      },
+      "peerDependenciesMeta": {
+        "supports-color": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/decode-named-character-reference": {
+      "version": "1.3.0",
+      "resolved": "https://registry.npmmirror.com/decode-named-character-reference/-/decode-named-character-reference-1.3.0.tgz",
+      "integrity": "sha512-GtpQYB283KrPp6nRw50q3U9/VfOutZOe103qlN7BPP6Ad27xYnOIWv4lPzo8HCAL+mMZofJ9KEy30fq6MfaK6Q==",
+      "license": "MIT",
+      "dependencies": {
+        "character-entities": "^2.0.0"
+      },
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/dequal": {
+      "version": "2.0.3",
+      "resolved": "https://registry.npmmirror.com/dequal/-/dequal-2.0.3.tgz",
+      "integrity": "sha512-0je+qPKHEMohvfRTCEo3CrPG6cAzAYgmzKyxRiYSSDkS6eGJdyVJm7WaYA5ECaAD9wLB2T4EEeymA5aFVcYXCA==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=6"
+      }
+    },
+    "node_modules/detect-libc": {
+      "version": "2.1.2",
+      "resolved": "https://registry.npmmirror.com/detect-libc/-/detect-libc-2.1.2.tgz",
+      "integrity": "sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==",
+      "dev": true,
+      "license": "Apache-2.0",
+      "engines": {
+        "node": ">=8"
+      }
+    },
+    "node_modules/devlop": {
+      "version": "1.1.0",
+      "resolved": "https://registry.npmmirror.com/devlop/-/devlop-1.1.0.tgz",
+      "integrity": "sha512-RWmIqhcFf1lRYBvNmr7qTNuyCt/7/ns2jbpp1+PalgE/rDQcBT0fioSMUpJ93irlUhC5hrg4cYqe6U+0ImW0rA==",
+      "license": "MIT",
+      "dependencies": {
+        "dequal": "^2.0.0"
+      },
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/enhanced-resolve": {
+      "version": "5.25.1",
+      "resolved": "https://registry.npmmirror.com/enhanced-resolve/-/enhanced-resolve-5.25.1.tgz",
+      "integrity": "sha512-nGXts5znJzmWPu+mIE9izCOzdg63oJca2mDzGWWTth7sr4aCToKcoyFVBQwN75Ij5Pf6p510EwkTqViTRzDV+w==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "graceful-fs": "^4.2.4",
+        "tapable": "^2.3.3"
+      },
+      "engines": {
+        "node": ">=10.13.0"
+      }
+    },
+    "node_modules/entities": {
+      "version": "6.0.1",
+      "resolved": "https://registry.npmmirror.com/entities/-/entities-6.0.1.tgz",
+      "integrity": "sha512-aN97NXWF6AWBTahfVOIrB/NShkzi5H7F9r1s9mD3cDj4Ko5f2qhhVoYMibXF7GlLveb/D2ioWay8lxI97Ven3g==",
+      "license": "BSD-2-Clause",
+      "engines": {
+        "node": ">=0.12"
+      },
+      "funding": {
+        "url": "https://github.com/fb55/entities?sponsor=1"
+      }
+    },
+    "node_modules/escape-string-regexp": {
+      "version": "5.0.0",
+      "resolved": "https://registry.npmmirror.com/escape-string-regexp/-/escape-string-regexp-5.0.0.tgz",
+      "integrity": "sha512-/veY75JbMK4j1yjvuUxuVsiS/hr/4iHs9FTT6cgTexxdE0Ly/glccBAkloH/DofkjRbZU3bnoj38mOmhkZ0lHw==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=12"
+      },
+      "funding": {
+        "url": "https://github.com/sponsors/sindresorhus"
+      }
+    },
+    "node_modules/estree-util-is-identifier-name": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmmirror.com/estree-util-is-identifier-name/-/estree-util-is-identifier-name-3.0.0.tgz",
+      "integrity": "sha512-hFtqIDZTIUZ9BXLb8y4pYGyk6+wekIivNVTcmvk8NoOh+VeRn5y6cEHzbURrWbfp1fIqdVipilzj+lfaadNZmg==",
+      "license": "MIT",
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/extend": {
+      "version": "3.0.2",
+      "resolved": "https://registry.npmmirror.com/extend/-/extend-3.0.2.tgz",
+      "integrity": "sha512-fjquC59cD7CyW6urNXK0FBufkZcoiGG80wTuPujX590cB5Ttln20E2UB4S/WARVqhXffZl2LNgS+gQdPIIim/g==",
+      "license": "MIT"
+    },
+    "node_modules/fdir": {
+      "version": "6.5.0",
+      "resolved": "https://registry.npmmirror.com/fdir/-/fdir-6.5.0.tgz",
+      "integrity": "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==",
+      "dev": true,
+      "license": "MIT",
+      "engines": {
+        "node": ">=12.0.0"
+      },
+      "peerDependencies": {
+        "picomatch": "^3 || ^4"
+      },
+      "peerDependenciesMeta": {
+        "picomatch": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/fsevents": {
+      "version": "2.3.3",
+      "resolved": "https://registry.npmmirror.com/fsevents/-/fsevents-2.3.3.tgz",
+      "integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==",
+      "dev": true,
+      "hasInstallScript": true,
+      "license": "MIT",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": "^8.16.0 || ^10.6.0 || >=11.0.0"
+      }
+    },
+    "node_modules/graceful-fs": {
+      "version": "4.2.11",
+      "resolved": "https://registry.npmmirror.com/graceful-fs/-/graceful-fs-4.2.11.tgz",
+      "integrity": "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==",
+      "dev": true,
+      "license": "ISC"
+    },
+    "node_modules/hast-util-from-parse5": {
+      "version": "8.0.3",
+      "resolved": "https://registry.npmmirror.com/hast-util-from-parse5/-/hast-util-from-parse5-8.0.3.tgz",
+      "integrity": "sha512-3kxEVkEKt0zvcZ3hCRYI8rqrgwtlIOFMWkbclACvjlDw8Li9S2hk/d51OI0nr/gIpdMHNepwgOKqZ/sy0Clpyg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "@types/unist": "^3.0.0",
+        "devlop": "^1.0.0",
+        "hastscript": "^9.0.0",
+        "property-information": "^7.0.0",
+        "vfile": "^6.0.0",
+        "vfile-location": "^5.0.0",
+        "web-namespaces": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/hast-util-parse-selector": {
+      "version": "4.0.0",
+      "resolved": "https://registry.npmmirror.com/hast-util-parse-selector/-/hast-util-parse-selector-4.0.0.tgz",
+      "integrity": "sha512-wkQCkSYoOGCRKERFWcxMVMOcYE2K1AaNLU8DXS9arxnLOUEWbOXKXiJUNzEpqZ3JOKpnha3jkFrumEjVliDe7A==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/hast-util-raw": {
+      "version": "9.1.0",
+      "resolved": "https://registry.npmmirror.com/hast-util-raw/-/hast-util-raw-9.1.0.tgz",
+      "integrity": "sha512-Y8/SBAHkZGoNkpzqqfCldijcuUKh7/su31kEBp67cFY09Wy0mTRgtsLYsiIxMJxlu0f6AA5SUTbDR8K0rxnbUw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "@types/unist": "^3.0.0",
+        "@ungap/structured-clone": "^1.0.0",
+        "hast-util-from-parse5": "^8.0.0",
+        "hast-util-to-parse5": "^8.0.0",
+        "html-void-elements": "^3.0.0",
+        "mdast-util-to-hast": "^13.0.0",
+        "parse5": "^7.0.0",
+        "unist-util-position": "^5.0.0",
+        "unist-util-visit": "^5.0.0",
+        "vfile": "^6.0.0",
+        "web-namespaces": "^2.0.0",
+        "zwitch": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/hast-util-sanitize": {
+      "version": "5.0.2",
+      "resolved": "https://registry.npmmirror.com/hast-util-sanitize/-/hast-util-sanitize-5.0.2.tgz",
+      "integrity": "sha512-3yTWghByc50aGS7JlGhk61SPenfE/p1oaFeNwkOOyrscaOkMGrcW9+Cy/QAIOBpZxP1yqDIzFMR0+Np0i0+usg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "@ungap/structured-clone": "^1.0.0",
+        "unist-util-position": "^5.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/hast-util-to-jsx-runtime": {
+      "version": "2.3.6",
+      "resolved": "https://registry.npmmirror.com/hast-util-to-jsx-runtime/-/hast-util-to-jsx-runtime-2.3.6.tgz",
+      "integrity": "sha512-zl6s8LwNyo1P9uw+XJGvZtdFF1GdAkOg8ujOw+4Pyb76874fLps4ueHXDhXWdk6YHQ6OgUtinliG7RsYvCbbBg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/estree": "^1.0.0",
+        "@types/hast": "^3.0.0",
+        "@types/unist": "^3.0.0",
+        "comma-separated-tokens": "^2.0.0",
+        "devlop": "^1.0.0",
+        "estree-util-is-identifier-name": "^3.0.0",
+        "hast-util-whitespace": "^3.0.0",
+        "mdast-util-mdx-expression": "^2.0.0",
+        "mdast-util-mdx-jsx": "^3.0.0",
+        "mdast-util-mdxjs-esm": "^2.0.0",
+        "property-information": "^7.0.0",
+        "space-separated-tokens": "^2.0.0",
+        "style-to-js": "^1.0.0",
+        "unist-util-position": "^5.0.0",
+        "vfile-message": "^4.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/hast-util-to-parse5": {
+      "version": "8.0.1",
+      "resolved": "https://registry.npmmirror.com/hast-util-to-parse5/-/hast-util-to-parse5-8.0.1.tgz",
+      "integrity": "sha512-MlWT6Pjt4CG9lFCjiz4BH7l9wmrMkfkJYCxFwKQic8+RTZgWPuWxwAfjJElsXkex7DJjfSJsQIt931ilUgmwdA==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "comma-separated-tokens": "^2.0.0",
+        "devlop": "^1.0.0",
+        "property-information": "^7.0.0",
+        "space-separated-tokens": "^2.0.0",
+        "web-namespaces": "^2.0.0",
+        "zwitch": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/hast-util-whitespace": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmmirror.com/hast-util-whitespace/-/hast-util-whitespace-3.0.0.tgz",
+      "integrity": "sha512-88JUN06ipLwsnv+dVn+OIYOvAuvBMy/Qoi6O7mQHxdPXpjy+Cd6xRkWwux7DKO+4sYILtLBRIKgsdpS2gQc7qw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/hastscript": {
+      "version": "9.0.1",
+      "resolved": "https://registry.npmmirror.com/hastscript/-/hastscript-9.0.1.tgz",
+      "integrity": "sha512-g7df9rMFX/SPi34tyGCyUBREQoKkapwdY/T04Qn9TDWfHhAYt4/I0gMVirzK5wEzeUqIjEB+LXC/ypb7Aqno5w==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "comma-separated-tokens": "^2.0.0",
+        "hast-util-parse-selector": "^4.0.0",
+        "property-information": "^7.0.0",
+        "space-separated-tokens": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/html-url-attributes": {
+      "version": "3.0.1",
+      "resolved": "https://registry.npmmirror.com/html-url-attributes/-/html-url-attributes-3.0.1.tgz",
+      "integrity": "sha512-ol6UPyBWqsrO6EJySPz2O7ZSr856WDrEzM5zMqp+FJJLGMW35cLYmmZnl0vztAZxRUoNZJFTCohfjuIJ8I4QBQ==",
+      "license": "MIT",
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/html-void-elements": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmmirror.com/html-void-elements/-/html-void-elements-3.0.0.tgz",
+      "integrity": "sha512-bEqo66MRXsUGxWHV5IP0PUiAWwoEjba4VCzg0LjFJBpchPaTfyfCKTG6bc5F8ucKec3q5y6qOdGyYTSBEvhCrg==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/inline-style-parser": {
+      "version": "0.2.7",
+      "resolved": "https://registry.npmmirror.com/inline-style-parser/-/inline-style-parser-0.2.7.tgz",
+      "integrity": "sha512-Nb2ctOyNR8DqQoR0OwRG95uNWIC0C1lCgf5Naz5H6Ji72KZ8OcFZLz2P5sNgwlyoJ8Yif11oMuYs5pBQa86csA==",
+      "license": "MIT"
+    },
+    "node_modules/is-alphabetical": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/is-alphabetical/-/is-alphabetical-2.0.1.tgz",
+      "integrity": "sha512-FWyyY60MeTNyeSRpkM2Iry0G9hpr7/9kD40mD/cGQEuilcZYS4okz8SN2Q6rLCJ8gbCt6fN+rC+6tMGS99LaxQ==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/is-alphanumerical": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/is-alphanumerical/-/is-alphanumerical-2.0.1.tgz",
+      "integrity": "sha512-hmbYhX/9MUMF5uh7tOXyK/n0ZvWpad5caBA17GsC6vyuCqaWliRG5K1qS9inmUhEMaOBIW7/whAnSwveW/LtZw==",
+      "license": "MIT",
+      "dependencies": {
+        "is-alphabetical": "^2.0.0",
+        "is-decimal": "^2.0.0"
+      },
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/is-decimal": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/is-decimal/-/is-decimal-2.0.1.tgz",
+      "integrity": "sha512-AAB9hiomQs5DXWcRB1rqsxGUstbRroFOPPVAomNk/3XHR5JyEZChOyTWe2oayKnsSsr/kcGqF+z6yuH6HHpN0A==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/is-hexadecimal": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/is-hexadecimal/-/is-hexadecimal-2.0.1.tgz",
+      "integrity": "sha512-DgZQp241c8oO6cA1SbTEWiXeoxV42vlcJxgH+B3hi1AiqqKruZR3ZGF8In3fj4+/y/7rHvlOZLZtgJ/4ttYGZg==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/is-plain-obj": {
+      "version": "4.1.0",
+      "resolved": "https://registry.npmmirror.com/is-plain-obj/-/is-plain-obj-4.1.0.tgz",
+      "integrity": "sha512-+Pgi+vMuUNkJyExiMBt5IlFoMyKnr5zhJ4Uspz58WOhBF5QoIZkFyNHIbBAtHwzVAgk5RtndVNsDRN61/mmDqg==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=12"
+      },
+      "funding": {
+        "url": "https://github.com/sponsors/sindresorhus"
+      }
+    },
+    "node_modules/jiti": {
+      "version": "2.7.0",
+      "resolved": "https://registry.npmmirror.com/jiti/-/jiti-2.7.0.tgz",
+      "integrity": "sha512-AC/7JofJvZGrrneWNaEnJeOLUx+JlGt7tNa0wZiRPT4MY1wmfKjt2+6O2p2uz2+skll8OZZmJMNqeke7kKbNgQ==",
+      "dev": true,
+      "license": "MIT",
+      "bin": {
+        "jiti": "lib/jiti-cli.mjs"
+      }
+    },
+    "node_modules/lightningcss": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss/-/lightningcss-1.32.0.tgz",
+      "integrity": "sha512-NXYBzinNrblfraPGyrbPoD19C1h9lfI/1mzgWYvXUTe414Gz/X1FD2XBZSZM7rRTrMA8JL3OtAaGifrIKhQ5yQ==",
+      "dev": true,
+      "license": "MPL-2.0",
+      "dependencies": {
+        "detect-libc": "^2.0.3"
+      },
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      },
+      "optionalDependencies": {
+        "lightningcss-android-arm64": "1.32.0",
+        "lightningcss-darwin-arm64": "1.32.0",
+        "lightningcss-darwin-x64": "1.32.0",
+        "lightningcss-freebsd-x64": "1.32.0",
+        "lightningcss-linux-arm-gnueabihf": "1.32.0",
+        "lightningcss-linux-arm64-gnu": "1.32.0",
+        "lightningcss-linux-arm64-musl": "1.32.0",
+        "lightningcss-linux-x64-gnu": "1.32.0",
+        "lightningcss-linux-x64-musl": "1.32.0",
+        "lightningcss-win32-arm64-msvc": "1.32.0",
+        "lightningcss-win32-x64-msvc": "1.32.0"
+      }
+    },
+    "node_modules/lightningcss-android-arm64": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-android-arm64/-/lightningcss-android-arm64-1.32.0.tgz",
+      "integrity": "sha512-YK7/ClTt4kAK0vo6w3X+Pnm0D2cf2vPHbhOXdoNti1Ga0al1P4TBZhwjATvjNwLEBCnKvjJc2jQgHXH0NEwlAg==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "android"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-darwin-arm64": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-darwin-arm64/-/lightningcss-darwin-arm64-1.32.0.tgz",
+      "integrity": "sha512-RzeG9Ju5bag2Bv1/lwlVJvBE3q6TtXskdZLLCyfg5pt+HLz9BqlICO7LZM7VHNTTn/5PRhHFBSjk5lc4cmscPQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-darwin-x64": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-darwin-x64/-/lightningcss-darwin-x64-1.32.0.tgz",
+      "integrity": "sha512-U+QsBp2m/s2wqpUYT/6wnlagdZbtZdndSmut/NJqlCcMLTWp5muCrID+K5UJ6jqD2BFshejCYXniPDbNh73V8w==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-freebsd-x64": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-freebsd-x64/-/lightningcss-freebsd-x64-1.32.0.tgz",
+      "integrity": "sha512-JCTigedEksZk3tHTTthnMdVfGf61Fky8Ji2E4YjUTEQX14xiy/lTzXnu1vwiZe3bYe0q+SpsSH/CTeDXK6WHig==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "freebsd"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-linux-arm-gnueabihf": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-arm-gnueabihf/-/lightningcss-linux-arm-gnueabihf-1.32.0.tgz",
+      "integrity": "sha512-x6rnnpRa2GL0zQOkt6rts3YDPzduLpWvwAF6EMhXFVZXD4tPrBkEFqzGowzCsIWsPjqSK+tyNEODUBXeeVHSkw==",
+      "cpu": [
+        "arm"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-linux-arm64-gnu": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-arm64-gnu/-/lightningcss-linux-arm64-gnu-1.32.0.tgz",
+      "integrity": "sha512-0nnMyoyOLRJXfbMOilaSRcLH3Jw5z9HDNGfT/gwCPgaDjnx0i8w7vBzFLFR1f6CMLKF8gVbebmkUN3fa/kQJpQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-linux-arm64-musl": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-arm64-musl/-/lightningcss-linux-arm64-musl-1.32.0.tgz",
+      "integrity": "sha512-UpQkoenr4UJEzgVIYpI80lDFvRmPVg6oqboNHfoH4CQIfNA+HOrZ7Mo7KZP02dC6LjghPQJeBsvXhJod/wnIBg==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-linux-x64-gnu": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-x64-gnu/-/lightningcss-linux-x64-gnu-1.32.0.tgz",
+      "integrity": "sha512-V7Qr52IhZmdKPVr+Vtw8o+WLsQJYCTd8loIfpDaMRWGUZfBOYEJeyJIkqGIDMZPwPx24pUMfwSxxI8phr/MbOA==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-linux-x64-musl": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-x64-musl/-/lightningcss-linux-x64-musl-1.32.0.tgz",
+      "integrity": "sha512-bYcLp+Vb0awsiXg/80uCRezCYHNg1/l3mt0gzHnWV9XP1W5sKa5/TCdGWaR/zBM2PeF/HbsQv/j2URNOiVuxWg==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-win32-arm64-msvc": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-win32-arm64-msvc/-/lightningcss-win32-arm64-msvc-1.32.0.tgz",
+      "integrity": "sha512-8SbC8BR40pS6baCM8sbtYDSwEVQd4JlFTOlaD3gWGHfThTcABnNDBda6eTZeqbofalIJhFx0qKzgHJmcPTnGdw==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/lightningcss-win32-x64-msvc": {
+      "version": "1.32.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-win32-x64-msvc/-/lightningcss-win32-x64-msvc-1.32.0.tgz",
+      "integrity": "sha512-Amq9B/SoZYdDi1kFrojnoqPLxYhQ4Wo5XiL8EVJrVsB8ARoC1PWW6VGtT0WKCemjy8aC+louJnjS7U18x3b06Q==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/longest-streak": {
+      "version": "3.1.0",
+      "resolved": "https://registry.npmmirror.com/longest-streak/-/longest-streak-3.1.0.tgz",
+      "integrity": "sha512-9Ri+o0JYgehTaVBBDoMqIl8GXtbWg711O3srftcHhZ0dqnETqLaoIK0x17fUw9rFSlK/0NlsKe0Ahhyl5pXE2g==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/lucide-react": {
+      "version": "1.48.0",
+      "resolved": "https://registry.npmmirror.com/lucide-react/-/lucide-react-1.48.0.tgz",
+      "integrity": "sha512-R0CIKY/fXiC6y9xRBADgsK+VW2p/pcTJOMhLZf1T+uG+vJUvyUR02nGOgmIIC7MPkiNK4Ox8QDnbqF8rJYWPZQ==",
+      "license": "ISC",
+      "peerDependencies": {
+        "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0"
+      }
+    },
+    "node_modules/magic-string": {
+      "version": "0.30.21",
+      "resolved": "https://registry.npmmirror.com/magic-string/-/magic-string-0.30.21.tgz",
+      "integrity": "sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@jridgewell/sourcemap-codec": "^1.5.5"
+      }
+    },
+    "node_modules/markdown-table": {
+      "version": "3.0.4",
+      "resolved": "https://registry.npmmirror.com/markdown-table/-/markdown-table-3.0.4.tgz",
+      "integrity": "sha512-wiYz4+JrLyb/DqW2hkFJxP7Vd7JuTDm77fvbM8VfEQdmSMqcImWeeRbHwZjBjIFki/VaMK2BhFi7oUUZeM5bqw==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/marked": {
+      "version": "17.0.6",
+      "resolved": "https://registry.npmmirror.com/marked/-/marked-17.0.6.tgz",
+      "integrity": "sha512-gB0gkNafnonOw0obSTEGZTT86IuhILt2Wfx0mWH/1Au83kybTayroZ/V6nS25mN7u8ASy+5fMhgB3XPNrOZdmA==",
+      "license": "MIT",
+      "bin": {
+        "marked": "bin/marked.js"
+      },
+      "engines": {
+        "node": ">= 20"
+      }
+    },
+    "node_modules/mdast-util-find-and-replace": {
+      "version": "3.0.2",
+      "resolved": "https://registry.npmmirror.com/mdast-util-find-and-replace/-/mdast-util-find-and-replace-3.0.2.tgz",
+      "integrity": "sha512-Tmd1Vg/m3Xz43afeNxDIhWRtFZgM2VLyaf4vSTYwudTyeuTneoL3qtWMA5jeLyz/O1vDJmmV4QuScFCA2tBPwg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "escape-string-regexp": "^5.0.0",
+        "unist-util-is": "^6.0.0",
+        "unist-util-visit-parents": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-from-markdown": {
+      "version": "2.0.3",
+      "resolved": "https://registry.npmmirror.com/mdast-util-from-markdown/-/mdast-util-from-markdown-2.0.3.tgz",
+      "integrity": "sha512-W4mAWTvSlKvf8L6J+VN9yLSqQ9AOAAvHuoDAmPkz4dHf553m5gVj2ejadHJhoJmcmxEnOv6Pa8XJhpxE93kb8Q==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "@types/unist": "^3.0.0",
+        "decode-named-character-reference": "^1.0.0",
+        "devlop": "^1.0.0",
+        "mdast-util-to-string": "^4.0.0",
+        "micromark": "^4.0.0",
+        "micromark-util-decode-numeric-character-reference": "^2.0.0",
+        "micromark-util-decode-string": "^2.0.0",
+        "micromark-util-normalize-identifier": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0",
+        "unist-util-stringify-position": "^4.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-gfm": {
+      "version": "3.1.0",
+      "resolved": "https://registry.npmmirror.com/mdast-util-gfm/-/mdast-util-gfm-3.1.0.tgz",
+      "integrity": "sha512-0ulfdQOM3ysHhCJ1p06l0b0VKlhU0wuQs3thxZQagjcjPrlFRqY215uZGHHJan9GEAXd9MbfPjFJz+qMkVR6zQ==",
+      "license": "MIT",
+      "dependencies": {
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-gfm-autolink-literal": "^2.0.0",
+        "mdast-util-gfm-footnote": "^2.0.0",
+        "mdast-util-gfm-strikethrough": "^2.0.0",
+        "mdast-util-gfm-table": "^2.0.0",
+        "mdast-util-gfm-task-list-item": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-gfm-autolink-literal": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/mdast-util-gfm-autolink-literal/-/mdast-util-gfm-autolink-literal-2.0.1.tgz",
+      "integrity": "sha512-5HVP2MKaP6L+G6YaxPNjuL0BPrq9orG3TsrZ9YXbA3vDw/ACI4MEsnoDpn6ZNm7GnZgtAcONJyPhOP8tNJQavQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "ccount": "^2.0.0",
+        "devlop": "^1.0.0",
+        "mdast-util-find-and-replace": "^3.0.0",
+        "micromark-util-character": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-gfm-footnote": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/mdast-util-gfm-footnote/-/mdast-util-gfm-footnote-2.1.0.tgz",
+      "integrity": "sha512-sqpDWlsHn7Ac9GNZQMeUzPQSMzR6Wv0WKRNvQRg0KqHh02fpTz69Qc1QSseNX29bhz1ROIyNyxExfawVKTm1GQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "devlop": "^1.1.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0",
+        "micromark-util-normalize-identifier": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-gfm-strikethrough": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/mdast-util-gfm-strikethrough/-/mdast-util-gfm-strikethrough-2.0.1.tgz",
+      "integrity": "sha512-OuJHqvr455pwu2OaOrir7dbzqs4jlWKOtlu9L6GjcIjjQwNyw02VntQsr/Ek6/A13vgnXfUhBekWRUP1OkNivw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-gfm-table": {
+      "version": "2.0.0",
+      "resolved": "https://registry.npmmirror.com/mdast-util-gfm-table/-/mdast-util-gfm-table-2.0.0.tgz",
+      "integrity": "sha512-78UEvebzz/rJIxLvE7ZtDd/vIQ0RHv+3Mh5DR96p7cS7HsBhYIICDBCu8csTNWNO6tBWfqXPWekRuj2FNOGOZg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "devlop": "^1.0.0",
+        "markdown-table": "^3.0.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-gfm-task-list-item": {
+      "version": "2.0.0",
+      "resolved": "https://registry.npmmirror.com/mdast-util-gfm-task-list-item/-/mdast-util-gfm-task-list-item-2.0.0.tgz",
+      "integrity": "sha512-IrtvNvjxC1o06taBAVJznEnkiHxLFTzgonUdy8hzFVeDun0uTjxxrRGVaNFqkU1wJR3RBPEfsxmU6jDWPofrTQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "devlop": "^1.0.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-mdx-expression": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/mdast-util-mdx-expression/-/mdast-util-mdx-expression-2.0.1.tgz",
+      "integrity": "sha512-J6f+9hUp+ldTZqKRSg7Vw5V6MqjATc+3E4gf3CFNcuZNWD8XdyI6zQ8GqH7f8169MM6P7hMBRDVGnn7oHB9kXQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/estree-jsx": "^1.0.0",
+        "@types/hast": "^3.0.0",
+        "@types/mdast": "^4.0.0",
+        "devlop": "^1.0.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-mdx-jsx": {
+      "version": "3.2.0",
+      "resolved": "https://registry.npmmirror.com/mdast-util-mdx-jsx/-/mdast-util-mdx-jsx-3.2.0.tgz",
+      "integrity": "sha512-lj/z8v0r6ZtsN/cGNNtemmmfoLAFZnjMbNyLzBafjzikOM+glrjNHPlf6lQDOTccj9n5b0PPihEBbhneMyGs1Q==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/estree-jsx": "^1.0.0",
+        "@types/hast": "^3.0.0",
+        "@types/mdast": "^4.0.0",
+        "@types/unist": "^3.0.0",
+        "ccount": "^2.0.0",
+        "devlop": "^1.1.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0",
+        "parse-entities": "^4.0.0",
+        "stringify-entities": "^4.0.0",
+        "unist-util-stringify-position": "^4.0.0",
+        "vfile-message": "^4.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-mdxjs-esm": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/mdast-util-mdxjs-esm/-/mdast-util-mdxjs-esm-2.0.1.tgz",
+      "integrity": "sha512-EcmOpxsZ96CvlP03NghtH1EsLtr0n9Tm4lPUJUBccV9RwUOneqSycg19n5HGzCf+10LozMRSObtVr3ee1WoHtg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/estree-jsx": "^1.0.0",
+        "@types/hast": "^3.0.0",
+        "@types/mdast": "^4.0.0",
+        "devlop": "^1.0.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "mdast-util-to-markdown": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-phrasing": {
+      "version": "4.1.0",
+      "resolved": "https://registry.npmmirror.com/mdast-util-phrasing/-/mdast-util-phrasing-4.1.0.tgz",
+      "integrity": "sha512-TqICwyvJJpBwvGAMZjj4J2n0X8QWp21b9l0o7eXyVJ25YNWYbJDVIyD1bZXE6WtV6RmKJVYmQAKWa0zWOABz2w==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "unist-util-is": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-to-hast": {
+      "version": "13.2.1",
+      "resolved": "https://registry.npmmirror.com/mdast-util-to-hast/-/mdast-util-to-hast-13.2.1.tgz",
+      "integrity": "sha512-cctsq2wp5vTsLIcaymblUriiTcZd0CwWtCbLvrOzYCDZoWyMNV8sZ7krj09FSnsiJi3WVsHLM4k6Dq/yaPyCXA==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "@types/mdast": "^4.0.0",
+        "@ungap/structured-clone": "^1.0.0",
+        "devlop": "^1.0.0",
+        "micromark-util-sanitize-uri": "^2.0.0",
+        "trim-lines": "^3.0.0",
+        "unist-util-position": "^5.0.0",
+        "unist-util-visit": "^5.0.0",
+        "vfile": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-to-markdown": {
+      "version": "2.1.3",
+      "resolved": "https://registry.npmmirror.com/mdast-util-to-markdown/-/mdast-util-to-markdown-2.1.3.tgz",
+      "integrity": "sha512-wgyJtgUkUcdU7zci7uuwc/tzoAhm0TswWhaXuvnWolqug+8jY7bgNkBBL0dYBnHgc/tHL7lCzRdPsdUxUKKBqw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "@types/unist": "^3.0.0",
+        "longest-streak": "^3.0.0",
+        "mdast-util-phrasing": "^4.0.0",
+        "mdast-util-to-string": "^4.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-classify-character": "^2.0.0",
+        "micromark-util-decode-string": "^2.0.0",
+        "micromark-util-html-tag-name": "^2.0.0",
+        "unist-util-visit": "^5.0.0",
+        "zwitch": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/mdast-util-to-string": {
+      "version": "4.0.0",
+      "resolved": "https://registry.npmmirror.com/mdast-util-to-string/-/mdast-util-to-string-4.0.0.tgz",
+      "integrity": "sha512-0H44vDimn51F0YwvxSJSm0eCDOJTRlmN0R1yBh4HLj9wiV1Dn0QoXGbvFAWj2hSItVTlCmBF1hqKlIyUBVFLPg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark": {
+      "version": "4.0.3",
+      "resolved": "https://registry.npmmirror.com/micromark/-/micromark-4.0.3.tgz",
+      "integrity": "sha512-oGYfQzHSG5dOMovQcJ3fyTmZlWAWpi0XA0sJwJs+i6OT88o1+Jtw/8z0CmdowSLxGhhFK89rQ/oXph/wN02PNw==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "@types/debug": "^4.0.0",
+        "debug": "^4.0.0",
+        "decode-named-character-reference": "^1.0.0",
+        "devlop": "^1.0.0",
+        "micromark-core-commonmark": "^2.0.0",
+        "micromark-factory-space": "^2.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-chunked": "^2.0.0",
+        "micromark-util-combine-extensions": "^2.0.0",
+        "micromark-util-decode-numeric-character-reference": "^2.0.0",
+        "micromark-util-edit-map": "^1.0.0",
+        "micromark-util-encode": "^2.0.0",
+        "micromark-util-normalize-identifier": "^2.0.0",
+        "micromark-util-resolve-all": "^2.0.0",
+        "micromark-util-sanitize-uri": "^2.0.0",
+        "micromark-util-subtokenize": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-core-commonmark": {
+      "version": "2.0.4",
+      "resolved": "https://registry.npmmirror.com/micromark-core-commonmark/-/micromark-core-commonmark-2.0.4.tgz",
+      "integrity": "sha512-wxEeE8v8XVvOrxn1TZj74qYhAtTQsSqufHVS3uNVyT1MupXxhyaHk8/9xDNPh9BmL77QNDgWTkZ3RsYHB2ry7w==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "decode-named-character-reference": "^1.0.0",
+        "devlop": "^1.0.0",
+        "micromark-factory-destination": "^2.0.0",
+        "micromark-factory-label": "^2.0.0",
+        "micromark-factory-space": "^2.1.0",
+        "micromark-factory-title": "^2.0.0",
+        "micromark-factory-whitespace": "^2.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-chunked": "^2.0.0",
+        "micromark-util-classify-character": "^2.0.0",
+        "micromark-util-edit-map": "^1.0.0",
+        "micromark-util-html-tag-name": "^2.0.0",
+        "micromark-util-normalize-identifier": "^2.0.0",
+        "micromark-util-resolve-all": "^2.0.0",
+        "micromark-util-subtokenize": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.3"
+      }
+    },
+    "node_modules/micromark-extension-gfm": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmmirror.com/micromark-extension-gfm/-/micromark-extension-gfm-3.0.0.tgz",
+      "integrity": "sha512-vsKArQsicm7t0z2GugkCKtZehqUm31oeGBV/KVSorWSy8ZlNAv7ytjFhvaryUiCUJYqs+NoE6AFhpQvBTM6Q4w==",
+      "license": "MIT",
+      "dependencies": {
+        "micromark-extension-gfm-autolink-literal": "^2.0.0",
+        "micromark-extension-gfm-footnote": "^2.0.0",
+        "micromark-extension-gfm-strikethrough": "^2.0.0",
+        "micromark-extension-gfm-table": "^2.0.0",
+        "micromark-extension-gfm-tagfilter": "^2.0.0",
+        "micromark-extension-gfm-task-list-item": "^2.0.0",
+        "micromark-util-combine-extensions": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark-extension-gfm-autolink-literal": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/micromark-extension-gfm-autolink-literal/-/micromark-extension-gfm-autolink-literal-2.1.0.tgz",
+      "integrity": "sha512-oOg7knzhicgQ3t4QCjCWgTmfNhvQbDDnJeVu9v81r7NltNCVmhPy1fJRX27pISafdjL+SVc4d3l48Gb6pbRypw==",
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-sanitize-uri": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark-extension-gfm-footnote": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/micromark-extension-gfm-footnote/-/micromark-extension-gfm-footnote-2.1.0.tgz",
+      "integrity": "sha512-/yPhxI1ntnDNsiHtzLKYnE3vf9JZ6cAisqVDauhp4CEHxlb4uoOTxOCJ+9s51bIB8U1N1FJ1RXOKTIlD5B/gqw==",
+      "license": "MIT",
+      "dependencies": {
+        "devlop": "^1.0.0",
+        "micromark-core-commonmark": "^2.0.0",
+        "micromark-factory-space": "^2.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-normalize-identifier": "^2.0.0",
+        "micromark-util-sanitize-uri": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark-extension-gfm-strikethrough": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/micromark-extension-gfm-strikethrough/-/micromark-extension-gfm-strikethrough-2.1.0.tgz",
+      "integrity": "sha512-ADVjpOOkjz1hhkZLlBiYA9cR2Anf8F4HqZUO6e5eDcPQd0Txw5fxLzzxnEkSkfnD0wziSGiv7sYhk/ktvbf1uw==",
+      "license": "MIT",
+      "dependencies": {
+        "devlop": "^1.0.0",
+        "micromark-util-chunked": "^2.0.0",
+        "micromark-util-classify-character": "^2.0.0",
+        "micromark-util-resolve-all": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark-extension-gfm-table": {
+      "version": "2.1.2",
+      "resolved": "https://registry.npmmirror.com/micromark-extension-gfm-table/-/micromark-extension-gfm-table-2.1.2.tgz",
+      "integrity": "sha512-pRzm4kDTu0MjlmBkxmS9yYhw60nncfcEwu9NNdPFSQEFXS95ZKyIIyTSHu/o3ReBUrLKYEq+7YaXCRn/bPB4MA==",
+      "license": "MIT",
+      "dependencies": {
+        "devlop": "^1.0.0",
+        "micromark-factory-space": "^2.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark-extension-gfm-tagfilter": {
+      "version": "2.0.0",
+      "resolved": "https://registry.npmmirror.com/micromark-extension-gfm-tagfilter/-/micromark-extension-gfm-tagfilter-2.0.0.tgz",
+      "integrity": "sha512-xHlTOmuCSotIA8TW1mDIM6X2O1SiX5P9IuDtqGonFhEK0qgRI4yeC6vMxEV2dgyr2TiD+2PQ10o+cOhdVAcwfg==",
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-types": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark-extension-gfm-task-list-item": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/micromark-extension-gfm-task-list-item/-/micromark-extension-gfm-task-list-item-2.1.0.tgz",
+      "integrity": "sha512-qIBZhqxqI6fjLDYFTBIa4eivDMnP+OZqsNwmQ3xNLE4Cxwc+zfQEfbs6tzAo2Hjq+bh6q5F+Z8/cksrLFYWQQw==",
+      "license": "MIT",
+      "dependencies": {
+        "devlop": "^1.0.0",
+        "micromark-factory-space": "^2.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/micromark-factory-destination": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-factory-destination/-/micromark-factory-destination-2.0.1.tgz",
+      "integrity": "sha512-Xe6rDdJlkmbFRExpTOmRj9N3MaWmbAgdpSrBQvCFqhezUn4AHqJHbaEnfbVYYiexVSs//tqOdY/DxhjdCiJnIA==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-factory-label": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-factory-label/-/micromark-factory-label-2.0.1.tgz",
+      "integrity": "sha512-VFMekyQExqIW7xIChcXn4ok29YE3rnuyveW3wZQWWqF4Nv9Wk5rgJ99KzPvHjkmPXF93FXIbBp6YdW3t71/7Vg==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "devlop": "^1.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-factory-space": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/micromark-factory-space/-/micromark-factory-space-2.1.0.tgz",
+      "integrity": "sha512-fS8hnLIjnjvdQIj39Geug8wWsR0HrYZ43KShKxNfwT7t2LOHo/LbWZEzEaSOjwtjdCsjoE6syHhonEzW0zv0+Q==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-factory-title": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-factory-title/-/micromark-factory-title-2.0.1.tgz",
+      "integrity": "sha512-5bZ+3CjhAd9eChYTHsjy6TGxpOFSKgKKJPJxr293jTbfry2KDoWkhBb6TcPVB4NmzaPhMs1Frm9AZH7OD4Cjzw==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-factory-space": "^2.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-factory-whitespace": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-factory-whitespace/-/micromark-factory-whitespace-2.0.1.tgz",
+      "integrity": "sha512-Ob0nuZ3PKt/n0hORHyvoD9uZhr+Za8sFoP+OnMcnWK5lngSzALgQYKMr9RJVOWLqQYuyn6ulqGWSXdwf6F80lQ==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-factory-space": "^2.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-character": {
+      "version": "2.1.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-character/-/micromark-util-character-2.1.1.tgz",
+      "integrity": "sha512-wv8tdUTJ3thSFFFJKtpYKOYiGP2+v96Hvk4Tu8KpCAsTMs6yi+nVmGh1syvSCsaxz45J6Jbw+9DD6g97+NV67Q==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-chunked": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-chunked/-/micromark-util-chunked-2.0.1.tgz",
+      "integrity": "sha512-QUNFEOPELfmvv+4xiNg2sRYeS/P84pTW0TCgP5zc9FpXetHY0ab7SxKyAQCNCc1eK0459uoLI1y5oO5Vc1dbhA==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-symbol": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-classify-character": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-classify-character/-/micromark-util-classify-character-2.0.1.tgz",
+      "integrity": "sha512-K0kHzM6afW/MbeWYWLjoHQv1sgg2Q9EccHEDzSkxiP/EaagNzCm7T/WMKZ3rjMbvIpvBiZgwR3dKMygtA4mG1Q==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-combine-extensions": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-combine-extensions/-/micromark-util-combine-extensions-2.0.1.tgz",
+      "integrity": "sha512-OnAnH8Ujmy59JcyZw8JSbK9cGpdVY44NKgSM7E9Eh7DiLS2E9RNQf0dONaGDzEG9yjEl5hcqeIsj4hfRkLH/Bg==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-chunked": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-decode-numeric-character-reference": {
+      "version": "2.0.2",
+      "resolved": "https://registry.npmmirror.com/micromark-util-decode-numeric-character-reference/-/micromark-util-decode-numeric-character-reference-2.0.2.tgz",
+      "integrity": "sha512-ccUbYk6CwVdkmCQMyr64dXz42EfHGkPQlBj5p7YVGzq8I7CtjXZJrubAYezf7Rp+bjPseiROqe7G6foFd+lEuw==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-symbol": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-decode-string": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-decode-string/-/micromark-util-decode-string-2.0.1.tgz",
+      "integrity": "sha512-nDV/77Fj6eH1ynwscYTOsbK7rR//Uj0bZXBwJZRfaLEJ1iGBR6kIfNmlNqaqJf649EP0F3NWNdeJi03elllNUQ==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "decode-named-character-reference": "^1.0.0",
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-decode-numeric-character-reference": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-edit-map": {
+      "version": "1.0.0",
+      "resolved": "https://registry.npmmirror.com/micromark-util-edit-map/-/micromark-util-edit-map-1.0.0.tgz",
+      "integrity": "sha512-Pa2ljlsEL6sVwFaYeyrOLSYbQt73JvGbPOYFq+9AElXiSnfI5Q4465RRZ1sHv7lDjDOnGRDDaqPXqClqWK/q1Q==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-encode": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-encode/-/micromark-util-encode-2.0.1.tgz",
+      "integrity": "sha512-c3cVx2y4KqUnwopcO9b/SCdo2O67LwJJ/UyqGfbigahfegL9myoEFoDYZgkT7f36T0bLrM9hZTAaAyH+PCAXjw==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT"
+    },
+    "node_modules/micromark-util-html-tag-name": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-html-tag-name/-/micromark-util-html-tag-name-2.0.1.tgz",
+      "integrity": "sha512-2cNEiYDhCWKI+Gs9T0Tiysk136SnR13hhO8yW6BGNyhOC4qYFnwF1nKfD3HFAIXA5c45RrIG1ub11GiXeYd1xA==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT"
+    },
+    "node_modules/micromark-util-normalize-identifier": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-normalize-identifier/-/micromark-util-normalize-identifier-2.0.1.tgz",
+      "integrity": "sha512-sxPqmo70LyARJs0w2UclACPUUEqltCkJ6PhKdMIDuJ3gSf/Q+/GIe3WKl0Ijb/GyH9lOpUkRAO2wp0GVkLvS9Q==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-symbol": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-resolve-all": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-resolve-all/-/micromark-util-resolve-all-2.0.1.tgz",
+      "integrity": "sha512-VdQyxFWFT2/FGJgwQnJYbe1jjQoNTS4RjglmSjTUlpUMa95Htx9NHeYW4rGDJzbjvCsl9eLjMQwGeElsqmzcHg==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-sanitize-uri": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-sanitize-uri/-/micromark-util-sanitize-uri-2.0.1.tgz",
+      "integrity": "sha512-9N9IomZ/YuGGZZmQec1MbgxtlgougxTodVwDzzEouPKo3qFWvymFHWcnDi2vzV1ff6kas9ucW+o3yzJK9YB1AQ==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "micromark-util-character": "^2.0.0",
+        "micromark-util-encode": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-subtokenize": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmmirror.com/micromark-util-subtokenize/-/micromark-util-subtokenize-2.1.0.tgz",
+      "integrity": "sha512-XQLu552iSctvnEcgXw6+Sx75GflAPNED1qx7eBJ+wydBb2KCbRZe+NwvIEEMM83uml1+2WSXpBAcp9IUCgCYWA==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "devlop": "^1.0.0",
+        "micromark-util-chunked": "^2.0.0",
+        "micromark-util-symbol": "^2.0.0",
+        "micromark-util-types": "^2.0.0"
+      }
+    },
+    "node_modules/micromark-util-symbol": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/micromark-util-symbol/-/micromark-util-symbol-2.0.1.tgz",
+      "integrity": "sha512-vs5t8Apaud9N28kgCrRUdEed4UJ+wWNvicHLPxCa9ENlYuAY31M0ETy5y1vA33YoNPDFTghEbnh6efaE8h4x0Q==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT"
+    },
+    "node_modules/micromark-util-types": {
+      "version": "2.0.3",
+      "resolved": "https://registry.npmmirror.com/micromark-util-types/-/micromark-util-types-2.0.3.tgz",
+      "integrity": "sha512-oxB2Ik03hI0gv+VNn9tnh1t1YEe9MDPptViAEgfdf3YQHsn0pzGTgCdlSCJXcwhqm8phaEuM7zeEu3QQzVBrPg==",
+      "funding": [
+        {
+          "type": "GitHub Sponsors",
+          "url": "https://github.com/sponsors/unifiedjs"
+        },
+        {
+          "type": "OpenCollective",
+          "url": "https://opencollective.com/unified"
+        }
+      ],
+      "license": "MIT"
+    },
+    "node_modules/ms": {
+      "version": "2.1.3",
+      "resolved": "https://registry.npmmirror.com/ms/-/ms-2.1.3.tgz",
+      "integrity": "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==",
+      "license": "MIT"
+    },
+    "node_modules/nanoid": {
+      "version": "3.3.19",
+      "resolved": "https://registry.npmmirror.com/nanoid/-/nanoid-3.3.19.tgz",
+      "integrity": "sha512-Y2tUNy4ouw6tq5oDSKeQYGOyhkUBhNOcGV/02KC+6kd9eDGqdZd++mjMiIDilrBYvjEnCYvVtsuHCuP+okSfug==",
+      "dev": true,
+      "funding": [
+        {
+          "type": "github",
+          "url": "https://github.com/sponsors/ai"
+        }
+      ],
+      "license": "MIT",
+      "bin": {
+        "nanoid": "bin/nanoid.cjs"
+      },
+      "engines": {
+        "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1"
+      }
+    },
+    "node_modules/parse-entities": {
+      "version": "4.0.2",
+      "resolved": "https://registry.npmmirror.com/parse-entities/-/parse-entities-4.0.2.tgz",
+      "integrity": "sha512-GG2AQYWoLgL877gQIKeRPGO1xF9+eG1ujIb5soS5gPvLQ1y2o8FL90w2QWNdf9I361Mpp7726c+lj3U0qK1uGw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^2.0.0",
+        "character-entities-legacy": "^3.0.0",
+        "character-reference-invalid": "^2.0.0",
+        "decode-named-character-reference": "^1.0.0",
+        "is-alphanumerical": "^2.0.0",
+        "is-decimal": "^2.0.0",
+        "is-hexadecimal": "^2.0.0"
+      },
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/parse-entities/node_modules/@types/unist": {
+      "version": "2.0.11",
+      "resolved": "https://registry.npmmirror.com/@types/unist/-/unist-2.0.11.tgz",
+      "integrity": "sha512-CmBKiL6NNo/OqgmMn95Fk9Whlp2mtvIv+KNpQKN2F4SjvrEesubTRWGYSg+BnWZOnlCaSTU1sMpsBOzgbYhnsA==",
+      "license": "MIT"
+    },
+    "node_modules/parse5": {
+      "version": "7.3.0",
+      "resolved": "https://registry.npmmirror.com/parse5/-/parse5-7.3.0.tgz",
+      "integrity": "sha512-IInvU7fabl34qmi9gY8XOVxhYyMyuH2xUNpb2q8/Y+7552KlejkRvqvD19nMoUW/uQGGbqNpA6Tufu5FL5BZgw==",
+      "license": "MIT",
+      "dependencies": {
+        "entities": "^6.0.0"
+      },
+      "funding": {
+        "url": "https://github.com/inikulin/parse5?sponsor=1"
+      }
+    },
+    "node_modules/picocolors": {
+      "version": "1.1.1",
+      "resolved": "https://registry.npmmirror.com/picocolors/-/picocolors-1.1.1.tgz",
+      "integrity": "sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==",
+      "dev": true,
+      "license": "ISC"
+    },
+    "node_modules/picomatch": {
+      "version": "4.0.7",
+      "resolved": "https://registry.npmmirror.com/picomatch/-/picomatch-4.0.7.tgz",
+      "integrity": "sha512-qcJu88Q2IWqJsDD529JKMdwGm/dvInW4HvQnRwiH9JtihJvzGOscDtHE3x1pBKeUOTysQ8kVmLnJ2kJu7yhcGA==",
+      "dev": true,
+      "license": "MIT",
+      "engines": {
+        "node": ">=12"
+      },
+      "funding": {
+        "url": "https://github.com/sponsors/jonschlinkert"
+      }
+    },
+    "node_modules/postcss": {
+      "version": "8.5.28",
+      "resolved": "https://registry.npmmirror.com/postcss/-/postcss-8.5.28.tgz",
+      "integrity": "sha512-RRuzqDtt5Y9h3quz5hWhK+TPnsmVs6WwSU6LkJMeY4HstUEDuYTG8UJSdawMRzmzAtV+KEoG8N3Qg2qLy5vM/A==",
+      "dev": true,
+      "funding": [
+        {
+          "type": "opencollective",
+          "url": "https://opencollective.com/postcss/"
+        },
+        {
+          "type": "tidelift",
+          "url": "https://tidelift.com/funding/github/npm/postcss"
+        },
+        {
+          "type": "github",
+          "url": "https://github.com/sponsors/ai"
+        }
+      ],
+      "license": "MIT",
+      "dependencies": {
+        "nanoid": "^3.3.18",
+        "picocolors": "^1.1.1",
+        "source-map-js": "^1.2.1"
+      },
+      "engines": {
+        "node": "^10 || ^12 || >=14"
+      }
+    },
+    "node_modules/postcss-selector-parser": {
+      "version": "6.0.10",
+      "resolved": "https://registry.npmmirror.com/postcss-selector-parser/-/postcss-selector-parser-6.0.10.tgz",
+      "integrity": "sha512-IQ7TZdoaqbT+LCpShg46jnZVlhWD2w6iQYAcYXfHARZ7X1t/UGhhceQDs5X0cGqKvYlHNOuv7Oa1xmb0oQuA3w==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "cssesc": "^3.0.0",
+        "util-deprecate": "^1.0.2"
+      },
+      "engines": {
+        "node": ">=4"
+      }
+    },
+    "node_modules/property-information": {
+      "version": "7.2.0",
+      "resolved": "https://registry.npmmirror.com/property-information/-/property-information-7.2.0.tgz",
+      "integrity": "sha512-IAtzIB6sUiWaJYrX9smp3V46pBGbBeLFRGdh25kg1334VcBlD8HzhPeNIWQH9zhGmo2itIe25EHt9dQP7G5hmg==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/react": {
+      "version": "19.3.0",
+      "resolved": "https://registry.npmmirror.com/react/-/react-19.3.0.tgz",
+      "integrity": "sha512-E8LUcbtBWt20bbl2YoHfx4ZDBdxVTfOKtCZn9cDSJ4l6/nuoApcpIBcj47t2wZoVX8g2ZHuMHbiShgCR1T5Sog==",
+      "license": "MIT",
+      "engines": {
+        "node": ">=0.10.0"
+      }
+    },
+    "node_modules/react-dom": {
+      "version": "19.3.0",
+      "resolved": "https://registry.npmmirror.com/react-dom/-/react-dom-19.3.0.tgz",
+      "integrity": "sha512-JDk8dgif51OjFoDE70+OT9ICyYr+69HlmihNwp1+Nsfbna3t5sIiCa9ZJktDmQ4/1b/rn26hIAR2uYXDMr5r0Q==",
+      "license": "MIT",
+      "dependencies": {
+        "scheduler": "^0.28.0"
+      },
+      "peerDependencies": {
+        "react": "^19.3.0"
+      }
+    },
+    "node_modules/rehype-harden": {
+      "version": "1.1.8",
+      "resolved": "https://registry.npmmirror.com/rehype-harden/-/rehype-harden-1.1.8.tgz",
+      "integrity": "sha512-Qn7vR1xrf6fZCrkm9TDWi/AB4ylrHy+jqsNm1EHOAmbARYA6gsnVJBq/sdBh6kmT4NEZxH5vgIjrscefJAOXcw==",
+      "license": "MIT",
+      "dependencies": {
+        "unist-util-visit": "^5.0.0"
+      }
+    },
+    "node_modules/rehype-raw": {
+      "version": "7.0.0",
+      "resolved": "https://registry.npmmirror.com/rehype-raw/-/rehype-raw-7.0.0.tgz",
+      "integrity": "sha512-/aE8hCfKlQeA8LmyeyQvQF3eBiLRGNlfBJEvWH7ivp9sBqs7TNqBL5X3v157rM4IFETqDnIOO+z5M/biZbo9Ww==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "hast-util-raw": "^9.0.0",
+        "vfile": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/rehype-sanitize": {
+      "version": "6.0.0",
+      "resolved": "https://registry.npmmirror.com/rehype-sanitize/-/rehype-sanitize-6.0.0.tgz",
+      "integrity": "sha512-CsnhKNsyI8Tub6L4sm5ZFsme4puGfc6pYylvXo1AeqaGbjOYyzNv3qZPwvs0oMJ39eryyeOdmxwUIo94IpEhqg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "hast-util-sanitize": "^5.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/remark-gfm": {
+      "version": "4.0.1",
+      "resolved": "https://registry.npmmirror.com/remark-gfm/-/remark-gfm-4.0.1.tgz",
+      "integrity": "sha512-1quofZ2RQ9EWdeN34S79+KExV1764+wCUGop5CPL1WGdD0ocPpu91lzPGbwWMECpEpd42kJGQwzRfyov9j4yNg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "mdast-util-gfm": "^3.0.0",
+        "micromark-extension-gfm": "^3.0.0",
+        "remark-parse": "^11.0.0",
+        "remark-stringify": "^11.0.0",
+        "unified": "^11.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/remark-parse": {
+      "version": "11.0.0",
+      "resolved": "https://registry.npmmirror.com/remark-parse/-/remark-parse-11.0.0.tgz",
+      "integrity": "sha512-FCxlKLNGknS5ba/1lmpYijMUzX2esxW5xQqjWxw2eHFfS2MSdaHVINFmhjo+qN1WhZhNimq0dZATN9pH0IDrpA==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "mdast-util-from-markdown": "^2.0.0",
+        "micromark-util-types": "^2.0.0",
+        "unified": "^11.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/remark-rehype": {
+      "version": "11.1.2",
+      "resolved": "https://registry.npmmirror.com/remark-rehype/-/remark-rehype-11.1.2.tgz",
+      "integrity": "sha512-Dh7l57ianaEoIpzbp0PC9UKAdCSVklD8E5Rpw7ETfbTl3FqcOOgq5q2LVDhgGCkaBv7p24JXikPdvhhmHvKMsw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/hast": "^3.0.0",
+        "@types/mdast": "^4.0.0",
+        "mdast-util-to-hast": "^13.0.0",
+        "unified": "^11.0.0",
+        "vfile": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/remark-stringify": {
+      "version": "11.0.0",
+      "resolved": "https://registry.npmmirror.com/remark-stringify/-/remark-stringify-11.0.0.tgz",
+      "integrity": "sha512-1OSmLd3awB/t8qdoEOMazZkNsfVTeY4fTsgzcQFdXNq8ToTN4ZGwrMnlda4K6smTFKD+GRV6O48i6Z4iKgPPpw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/mdast": "^4.0.0",
+        "mdast-util-to-markdown": "^2.0.0",
+        "unified": "^11.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/remend": {
+      "version": "1.3.1",
+      "resolved": "https://registry.npmmirror.com/remend/-/remend-1.3.1.tgz",
+      "integrity": "sha512-N3DiY5qbRPoa5vkxn1oDLMyXOVTeo6Hp+XOj6SIqJAYUgLS0Q587gILPMom/qm86AQ/ZrcOdwEIzCz8V3J0nxQ==",
+      "license": "Apache-2.0"
+    },
+    "node_modules/rolldown": {
+      "version": "1.2.11",
+      "resolved": "https://registry.npmmirror.com/rolldown/-/rolldown-1.2.11.tgz",
+      "integrity": "sha512-qpSwIyz0jHQq5qXBTNxFmE6664rJ7O+4TvPFOiOaBSrz8IOHc1koKKSqTM2H6u1UG1+TveuC6vaDHKXFOvb1Kw==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "@oxc-project/types": "=0.151.0",
+        "@rolldown/pluginutils": "^1.0.0"
+      },
+      "bin": {
+        "rolldown": "bin/cli.mjs"
+      },
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      },
+      "optionalDependencies": {
+        "@rolldown/binding-android-arm-eabi": "1.2.11",
+        "@rolldown/binding-android-arm64": "1.2.11",
+        "@rolldown/binding-darwin-arm64": "1.2.11",
+        "@rolldown/binding-darwin-x64": "1.2.11",
+        "@rolldown/binding-freebsd-x64": "1.2.11",
+        "@rolldown/binding-linux-arm-gnueabihf": "1.2.11",
+        "@rolldown/binding-linux-arm64-gnu": "1.2.11",
+        "@rolldown/binding-linux-arm64-musl": "1.2.11",
+        "@rolldown/binding-linux-ppc64-gnu": "1.2.11",
+        "@rolldown/binding-linux-s390x-gnu": "1.2.11",
+        "@rolldown/binding-linux-x64-gnu": "1.2.11",
+        "@rolldown/binding-linux-x64-musl": "1.2.11",
+        "@rolldown/binding-openharmony-arm64": "1.2.11",
+        "@rolldown/binding-win32-arm64-msvc": "1.2.11",
+        "@rolldown/binding-win32-x64-msvc": "1.2.11"
+      }
+    },
+    "node_modules/scheduler": {
+      "version": "0.28.0",
+      "resolved": "https://registry.npmmirror.com/scheduler/-/scheduler-0.28.0.tgz",
+      "integrity": "sha512-juorfCmIkIw8tT+p5BXSm6PJjQF/ycEYmKyzURCIt/RaZIhL+PulbQ9Yu2z1HdOJDdqDTlxA1+xKBmHXJsczAw==",
+      "license": "MIT"
+    },
+    "node_modules/source-map-js": {
+      "version": "1.2.1",
+      "resolved": "https://registry.npmmirror.com/source-map-js/-/source-map-js-1.2.1.tgz",
+      "integrity": "sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==",
+      "dev": true,
+      "license": "BSD-3-Clause",
+      "engines": {
+        "node": ">=0.10.0"
+      }
+    },
+    "node_modules/space-separated-tokens": {
+      "version": "2.0.2",
+      "resolved": "https://registry.npmmirror.com/space-separated-tokens/-/space-separated-tokens-2.0.2.tgz",
+      "integrity": "sha512-PEGlAwrG8yXGXRjW32fGbg66JAlOAwbObuqVoJpv/mRgoWDQfgH1wDPvtzWyUSNAXBGSk8h755YDbbcEy3SH2Q==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/streamdown": {
+      "version": "2.6.0",
+      "resolved": "https://registry.npmmirror.com/streamdown/-/streamdown-2.6.0.tgz",
+      "integrity": "sha512-nQZVUn4GvB2R5SAlDNph20iKZeW+RM4gv6G1H3ACypEnmQf8PgTHZ/1Ta2OACRwQ2coK7aW5AeIAa7Ls5zk+RA==",
+      "license": "Apache-2.0",
+      "dependencies": {
+        "clsx": "^2.1.1",
+        "hast-util-to-jsx-runtime": "^2.3.6",
+        "html-url-attributes": "^3.0.1",
+        "marked": "^17.0.1",
+        "rehype-harden": "^1.1.8",
+        "rehype-raw": "^7.0.0",
+        "rehype-sanitize": "^6.0.0",
+        "remark-gfm": "^4.0.1",
+        "remark-parse": "^11.0.0",
+        "remark-rehype": "^11.1.2",
+        "remend": "1.3.1",
+        "tailwind-merge": "^3.6.0",
+        "unified": "^11.0.5",
+        "unist-util-visit": "^5.1.0",
+        "unist-util-visit-parents": "^6.0.2"
+      },
+      "peerDependencies": {
+        "react": "^18.0.0 || ^19.0.0",
+        "react-dom": "^18.0.0 || ^19.0.0"
+      }
+    },
+    "node_modules/stringify-entities": {
+      "version": "4.0.4",
+      "resolved": "https://registry.npmmirror.com/stringify-entities/-/stringify-entities-4.0.4.tgz",
+      "integrity": "sha512-IwfBptatlO+QCJUo19AqvrPNqlVMpW9YEL2LIVY+Rpv2qsjCGxaDLNRgeGsQWJhfItebuJhsGSLjaBbNSQ+ieg==",
+      "license": "MIT",
+      "dependencies": {
+        "character-entities-html4": "^2.0.0",
+        "character-entities-legacy": "^3.0.0"
+      },
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/style-to-js": {
+      "version": "1.1.21",
+      "resolved": "https://registry.npmmirror.com/style-to-js/-/style-to-js-1.1.21.tgz",
+      "integrity": "sha512-RjQetxJrrUJLQPHbLku6U/ocGtzyjbJMP9lCNK7Ag0CNh690nSH8woqWH9u16nMjYBAok+i7JO1NP2pOy8IsPQ==",
+      "license": "MIT",
+      "dependencies": {
+        "style-to-object": "1.0.14"
+      }
+    },
+    "node_modules/style-to-object": {
+      "version": "1.0.14",
+      "resolved": "https://registry.npmmirror.com/style-to-object/-/style-to-object-1.0.14.tgz",
+      "integrity": "sha512-LIN7rULI0jBscWQYaSswptyderlarFkjQ+t79nzty8tcIAceVomEVlLzH5VP4Cmsv6MtKhs7qaAiwlcp+Mgaxw==",
+      "license": "MIT",
+      "dependencies": {
+        "inline-style-parser": "0.2.7"
+      }
+    },
+    "node_modules/tailwind-merge": {
+      "version": "3.7.0",
+      "resolved": "https://registry.npmmirror.com/tailwind-merge/-/tailwind-merge-3.7.0.tgz",
+      "integrity": "sha512-XPPUyAc+cvspz3lHTcR/QgPfW2A0lv/xQNIjX3HGhLR+Nq2lHaLq5MtTesHn8GUr3W3DguT2KT5x3NVgRtYwmA==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/dcastil"
+      }
+    },
+    "node_modules/tailwindcss": {
+      "version": "4.3.3",
+      "resolved": "https://registry.npmmirror.com/tailwindcss/-/tailwindcss-4.3.3.tgz",
+      "integrity": "sha512-gOhV3P7ufE62QDGg1zVaTgCR+EtPv92k2nIhVcVKcLmxT1sUBsQGhnZj175j+MqRt4zLF7ic+sCYjfhxMxj7YQ==",
+      "dev": true,
+      "license": "MIT"
+    },
+    "node_modules/tapable": {
+      "version": "2.3.3",
+      "resolved": "https://registry.npmmirror.com/tapable/-/tapable-2.3.3.tgz",
+      "integrity": "sha512-uxc/zpqFg6x7C8vOE7lh6Lbda8eEL9zmVm/PLeTPBRhh1xCgdWaQ+J1CUieGpIfm2HdtsUpRv+HshiasBMcc6A==",
+      "dev": true,
+      "license": "MIT",
+      "engines": {
+        "node": ">=6"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/webpack"
+      }
+    },
+    "node_modules/tinyglobby": {
+      "version": "0.2.17",
+      "resolved": "https://registry.npmmirror.com/tinyglobby/-/tinyglobby-0.2.17.tgz",
+      "integrity": "sha512-wXR/dYpcqKmfWpEdZjiKJOwCNFndD0DMnrW/cYjVGttEkBfVgcLFHoNrlj47mjOVic9yyNu65alsgF4NQyTa2g==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "fdir": "^6.5.0",
+        "picomatch": "^4.0.4"
+      },
+      "engines": {
+        "node": ">=12.0.0"
+      },
+      "funding": {
+        "url": "https://github.com/sponsors/SuperchupuDev"
+      }
+    },
+    "node_modules/trim-lines": {
+      "version": "3.0.1",
+      "resolved": "https://registry.npmmirror.com/trim-lines/-/trim-lines-3.0.1.tgz",
+      "integrity": "sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/trough": {
+      "version": "2.2.0",
+      "resolved": "https://registry.npmmirror.com/trough/-/trough-2.2.0.tgz",
+      "integrity": "sha512-tmMpK00BjZiUyVyvrBK7knerNgmgvcV/KLVyuma/SC+TQN167GrMRciANTz09+k3zW8L8t60jWO1GpfkZdjTaw==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/typescript": {
+      "version": "7.0.2",
+      "resolved": "https://registry.npmmirror.com/typescript/-/typescript-7.0.2.tgz",
+      "integrity": "sha512-8FYau96o3NKOhbjKi/qNvG/W5jhzxkbdm5sj9AbZ/5T5sWqn3hJgLfGx27sRKZWTvyzCP8dLRBTf5tBTSRVUNA==",
+      "dev": true,
+      "license": "Apache-2.0",
+      "bin": {
+        "tsc": "bin/tsc"
+      },
+      "engines": {
+        "node": ">=16.20.0"
+      },
+      "optionalDependencies": {
+        "@typescript/typescript-aix-ppc64": "7.0.2",
+        "@typescript/typescript-darwin-arm64": "7.0.2",
+        "@typescript/typescript-darwin-x64": "7.0.2",
+        "@typescript/typescript-freebsd-arm64": "7.0.2",
+        "@typescript/typescript-freebsd-x64": "7.0.2",
+        "@typescript/typescript-linux-arm": "7.0.2",
+        "@typescript/typescript-linux-arm64": "7.0.2",
+        "@typescript/typescript-linux-loong64": "7.0.2",
+        "@typescript/typescript-linux-mips64el": "7.0.2",
+        "@typescript/typescript-linux-ppc64": "7.0.2",
+        "@typescript/typescript-linux-riscv64": "7.0.2",
+        "@typescript/typescript-linux-s390x": "7.0.2",
+        "@typescript/typescript-linux-x64": "7.0.2",
+        "@typescript/typescript-netbsd-arm64": "7.0.2",
+        "@typescript/typescript-netbsd-x64": "7.0.2",
+        "@typescript/typescript-openbsd-arm64": "7.0.2",
+        "@typescript/typescript-openbsd-x64": "7.0.2",
+        "@typescript/typescript-sunos-x64": "7.0.2",
+        "@typescript/typescript-win32-arm64": "7.0.2",
+        "@typescript/typescript-win32-x64": "7.0.2"
+      }
+    },
+    "node_modules/unified": {
+      "version": "11.0.5",
+      "resolved": "https://registry.npmmirror.com/unified/-/unified-11.0.5.tgz",
+      "integrity": "sha512-xKvGhPWw3k84Qjh8bI3ZeJjqnyadK+GEFtazSfZv/rKeTkTjOJho6mFqh2SM96iIcZokxiOpg78GazTSg8+KHA==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0",
+        "bail": "^2.0.0",
+        "devlop": "^1.0.0",
+        "extend": "^3.0.0",
+        "is-plain-obj": "^4.0.0",
+        "trough": "^2.0.0",
+        "vfile": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/unist-util-is": {
+      "version": "6.0.1",
+      "resolved": "https://registry.npmmirror.com/unist-util-is/-/unist-util-is-6.0.1.tgz",
+      "integrity": "sha512-LsiILbtBETkDz8I9p1dQ0uyRUWuaQzd/cuEeS1hoRSyW5E5XGmTzlwY1OrNzzakGowI9Dr/I8HVaw4hTtnxy8g==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/unist-util-position": {
+      "version": "5.0.0",
+      "resolved": "https://registry.npmmirror.com/unist-util-position/-/unist-util-position-5.0.0.tgz",
+      "integrity": "sha512-fucsC7HjXvkB5R3kTCO7kUjRdrS0BJt3M/FPxmHMBOm8JQi2BsHAHFsy27E0EolP8rp0NzXsJ+jNPyDWvOJZPA==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/unist-util-stringify-position": {
+      "version": "4.0.0",
+      "resolved": "https://registry.npmmirror.com/unist-util-stringify-position/-/unist-util-stringify-position-4.0.0.tgz",
+      "integrity": "sha512-0ASV06AAoKCDkS2+xw5RXJywruurpbC4JZSm7nr7MOt1ojAzvyyaO+UxZf18j8FCF6kmzCZKcAgN/yu2gm2XgQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/unist-util-visit": {
+      "version": "5.1.0",
+      "resolved": "https://registry.npmmirror.com/unist-util-visit/-/unist-util-visit-5.1.0.tgz",
+      "integrity": "sha512-m+vIdyeCOpdr/QeQCu2EzxX/ohgS8KbnPDgFni4dQsfSCtpz8UqDyY5GjRru8PDKuYn7Fq19j1CQ+nJSsGKOzg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0",
+        "unist-util-is": "^6.0.0",
+        "unist-util-visit-parents": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/unist-util-visit-parents": {
+      "version": "6.0.2",
+      "resolved": "https://registry.npmmirror.com/unist-util-visit-parents/-/unist-util-visit-parents-6.0.2.tgz",
+      "integrity": "sha512-goh1s1TBrqSqukSc8wrjwWhL0hiJxgA8m4kFxGlQ+8FYQ3C/m11FcTs4YYem7V664AhHVvgoQLk890Ssdsr2IQ==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0",
+        "unist-util-is": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/util-deprecate": {
+      "version": "1.0.2",
+      "resolved": "https://registry.npmmirror.com/util-deprecate/-/util-deprecate-1.0.2.tgz",
+      "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==",
+      "dev": true,
+      "license": "MIT"
+    },
+    "node_modules/vfile": {
+      "version": "6.0.3",
+      "resolved": "https://registry.npmmirror.com/vfile/-/vfile-6.0.3.tgz",
+      "integrity": "sha512-KzIbH/9tXat2u30jf+smMwFCsno4wHVdNmzFyL+T/L3UGqqk6JKfVqOFOZEpZSHADH1k40ab6NUIXZq422ov3Q==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0",
+        "vfile-message": "^4.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/vfile-location": {
+      "version": "5.0.3",
+      "resolved": "https://registry.npmmirror.com/vfile-location/-/vfile-location-5.0.3.tgz",
+      "integrity": "sha512-5yXvWDEgqeiYiBe1lbxYF7UMAIm/IcopxMHrMQDq3nvKcjPKIhZklUKL+AE7J7uApI4kwe2snsK+eI6UTj9EHg==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0",
+        "vfile": "^6.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/vfile-message": {
+      "version": "4.0.3",
+      "resolved": "https://registry.npmmirror.com/vfile-message/-/vfile-message-4.0.3.tgz",
+      "integrity": "sha512-QTHzsGd1EhbZs4AsQ20JX1rC3cOlt/IWJruk893DfLRr57lcnOeMaWG4K0JrRta4mIJZKth2Au3mM3u03/JWKw==",
+      "license": "MIT",
+      "dependencies": {
+        "@types/unist": "^3.0.0",
+        "unist-util-stringify-position": "^4.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/unified"
+      }
+    },
+    "node_modules/vite": {
+      "version": "8.3.1",
+      "resolved": "https://registry.npmmirror.com/vite/-/vite-8.3.1.tgz",
+      "integrity": "sha512-/bvH9E9tmCXRGp2uXY3WbOldqpTwFkbha/8ANaEQ6VkxhH60KyqLwgZq6lG2y+4uT55x9+9eUHMpQ7uGnOCKjA==",
+      "dev": true,
+      "license": "MIT",
+      "dependencies": {
+        "lightningcss": "^1.33.0",
+        "picomatch": "^4.0.7",
+        "postcss": "^8.5.28",
+        "rolldown": "~1.2.9",
+        "tinyglobby": "^0.2.17"
+      },
+      "bin": {
+        "vite": "bin/vite.js"
+      },
+      "engines": {
+        "node": "^20.19.0 || >=22.12.0"
+      },
+      "funding": {
+        "url": "https://github.com/vitejs/vite?sponsor=1"
+      },
+      "optionalDependencies": {
+        "fsevents": "~2.3.3"
+      },
+      "peerDependencies": {
+        "@types/node": "^20.19.0 || >=22.12.0",
+        "@vitejs/devtools": "^0.7.1",
+        "esbuild": "^0.27.0 || ^0.28.0",
+        "jiti": ">=1.21.0",
+        "less": "^4.0.0",
+        "sass": "^1.70.0",
+        "sass-embedded": "^1.70.0",
+        "stylus": ">=0.54.8",
+        "sugarss": "^5.0.0",
+        "terser": "^5.16.0",
+        "tsx": "^4.8.1",
+        "yaml": "^2.4.2"
+      },
+      "peerDependenciesMeta": {
+        "@types/node": {
+          "optional": true
+        },
+        "@vitejs/devtools": {
+          "optional": true
+        },
+        "esbuild": {
+          "optional": true
+        },
+        "jiti": {
+          "optional": true
+        },
+        "less": {
+          "optional": true
+        },
+        "sass": {
+          "optional": true
+        },
+        "sass-embedded": {
+          "optional": true
+        },
+        "stylus": {
+          "optional": true
+        },
+        "sugarss": {
+          "optional": true
+        },
+        "terser": {
+          "optional": true
+        },
+        "tsx": {
+          "optional": true
+        },
+        "yaml": {
+          "optional": true
+        }
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss/-/lightningcss-1.33.0.tgz",
+      "integrity": "sha512-WkUDrojuJs0xkgGf2udWxa3yGBRxPtxUkB79i6aCZLRgc7PM8fZe9TosfPDcvEpQZbuFASnHYmRLBLUbmLOIIA==",
+      "dev": true,
+      "license": "MPL-2.0",
+      "dependencies": {
+        "detect-libc": "^2.0.3"
+      },
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      },
+      "optionalDependencies": {
+        "lightningcss-android-arm64": "1.33.0",
+        "lightningcss-darwin-arm64": "1.33.0",
+        "lightningcss-darwin-x64": "1.33.0",
+        "lightningcss-freebsd-x64": "1.33.0",
+        "lightningcss-linux-arm-gnueabihf": "1.33.0",
+        "lightningcss-linux-arm64-gnu": "1.33.0",
+        "lightningcss-linux-arm64-musl": "1.33.0",
+        "lightningcss-linux-x64-gnu": "1.33.0",
+        "lightningcss-linux-x64-musl": "1.33.0",
+        "lightningcss-win32-arm64-msvc": "1.33.0",
+        "lightningcss-win32-x64-msvc": "1.33.0"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-android-arm64": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-android-arm64/-/lightningcss-android-arm64-1.33.0.tgz",
+      "integrity": "sha512-gEpRTalKdosp4Bb8qWtc2iOgE5SeIHlpS1up9bFq2wAyYhl1UdTObYiHe98zEM9SQvSoqQZ1IQD0JNpg3Ml5pg==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "android"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-darwin-arm64": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-darwin-arm64/-/lightningcss-darwin-arm64-1.33.0.tgz",
+      "integrity": "sha512-Sciaz8eenNTKn9b3t7+xr0ipTp9YxKQY4npwQ3mrRuL0BAVHBLyZxofhaKBAVtzmtRZ/zTyo0/to4B1uWG/Djg==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-darwin-x64": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-darwin-x64/-/lightningcss-darwin-x64-1.33.0.tgz",
+      "integrity": "sha512-Z5UPAxzrjlWNNyGy6i65cJzzvgJ5D3T6wMvs+gWpY9d7qRhANrxqAp6LhxIgZhWEw18RfJTGcRxjuLIBr+m8XQ==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "darwin"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-freebsd-x64": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-freebsd-x64/-/lightningcss-freebsd-x64-1.33.0.tgz",
+      "integrity": "sha512-QQM/Ti/hQajJwCY+RiWuCZ9sdtI/XQk7nDK5vC8kkdwixezOlDgvDx7+RT+QjK6FcFT4MpsuoBnHIo/O3StRRg==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "freebsd"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-linux-arm-gnueabihf": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-arm-gnueabihf/-/lightningcss-linux-arm-gnueabihf-1.33.0.tgz",
+      "integrity": "sha512-N7FVBe6iS24MlM6R/4RBTxGhQheZGs7tiQ9U32UtF75NzP5Q7xWPRqLBCKxlRQRk3rY1jCIPLzx7WzOhuUIRLQ==",
+      "cpu": [
+        "arm"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-linux-arm64-gnu": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-arm64-gnu/-/lightningcss-linux-arm64-gnu-1.33.0.tgz",
+      "integrity": "sha512-j2v/itmy4HlNxlc6voKXYgBqNi0Ng2LShg4z7GufpEgs05P+2suBVyi9I6YHq5uoVFx9ETin3eCEhLVyXGQnKg==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-linux-arm64-musl": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-arm64-musl/-/lightningcss-linux-arm64-musl-1.33.0.tgz",
+      "integrity": "sha512-yiO5ROMuYQgXbC60yjZU5CYSFZGKXL0HFATXt9mHJn1+zW55oCtMI9NfcVhYLMFDL7gV7oBPon/EmMMGg2OvtQ==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-linux-x64-gnu": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-x64-gnu/-/lightningcss-linux-x64-gnu-1.33.0.tgz",
+      "integrity": "sha512-ar+Ju7LmcN0Jo4FpL4hpFybwNG9/3A/Br5KW2n2jyODg3MEZXaDYADdemoNS+BDNfMgKvylJLj4S5tyRActuAg==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "glibc"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-linux-x64-musl": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-linux-x64-musl/-/lightningcss-linux-x64-musl-1.33.0.tgz",
+      "integrity": "sha512-RYiYbkokw0trfKqqzfF55lginwEPrD3OJDfTuJzFs1MK6iFnDenaz1fqLLtX4ITG3OktJQXOeTaw1awrBAlZPw==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "libc": [
+        "musl"
+      ],
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "linux"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-win32-arm64-msvc": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-win32-arm64-msvc/-/lightningcss-win32-arm64-msvc-1.33.0.tgz",
+      "integrity": "sha512-1K+MPfLSFVpphzpdbfkhlWk6wBrTObBzS2T6db10PNOZgR9GoVsAWzwNyuhUYYbTp23j+4RrncfujZ4uAzXvwA==",
+      "cpu": [
+        "arm64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/vite/node_modules/lightningcss-win32-x64-msvc": {
+      "version": "1.33.0",
+      "resolved": "https://registry.npmmirror.com/lightningcss-win32-x64-msvc/-/lightningcss-win32-x64-msvc-1.33.0.tgz",
+      "integrity": "sha512-OlEICDx/Xl0FqSp4bry8zFnCvGpig3Gl4gCquvYwHuqJKEC1+n9NgDniFvqHGmMv1ZkqDJrDqKKSykTDX+ehuA==",
+      "cpu": [
+        "x64"
+      ],
+      "dev": true,
+      "license": "MPL-2.0",
+      "optional": true,
+      "os": [
+        "win32"
+      ],
+      "engines": {
+        "node": ">= 12.0.0"
+      },
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/parcel"
+      }
+    },
+    "node_modules/web-namespaces": {
+      "version": "2.0.1",
+      "resolved": "https://registry.npmmirror.com/web-namespaces/-/web-namespaces-2.0.1.tgz",
+      "integrity": "sha512-bKr1DkiNa2krS7qxNtdrtHAmzuYGFQLiQ13TsorsdT6ULTkPLKuu5+GsFpDlg6JFjUTwX2DyhMPG2be8uPrqsQ==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    },
+    "node_modules/zwitch": {
+      "version": "2.0.4",
+      "resolved": "https://registry.npmmirror.com/zwitch/-/zwitch-2.0.4.tgz",
+      "integrity": "sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/wooorm"
+      }
+    }
+  }
+}
diff --git a/windows/app/package.json b/windows/app/package.json
new file mode 100644
index 0000000..ad64f1b
--- /dev/null
+++ b/windows/app/package.json
@@ -0,0 +1,30 @@
+{
+  "name": "edge0-app-ui",
+  "private": true,
+  "version": "0.1.0",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "tsc --noEmit && vite build",
+    "typecheck": "tsc --noEmit",
+    "tauri": "tauri"
+  },
+  "dependencies": {
+    "@tauri-apps/api": "^2.11.1",
+    "lucide-react": "^1.46.0",
+    "react": "^19.3.0",
+    "react-dom": "^19.3.0",
+    "streamdown": "^2.6.0"
+  },
+  "devDependencies": {
+    "@tailwindcss/typography": "^0.5.20",
+    "@tailwindcss/vite": "^4.3.3",
+    "@tauri-apps/cli": "^2.11.4",
+    "@types/react": "^19.3.0",
+    "@types/react-dom": "^19.3.0",
+    "@vitejs/plugin-react": "^6.1.1",
+    "tailwindcss": "^4.3.3",
+    "typescript": "^7.0.2",
+    "vite": "^8.3.0"
+  }
+}
diff --git a/windows/app/src-tauri/Cargo.lock b/windows/app/src-tauri/Cargo.lock
new file mode 100644
index 0000000..06c343f
--- /dev/null
+++ b/windows/app/src-tauri/Cargo.lock
@@ -0,0 +1,4722 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "adler2"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
+
+[[package]]
+name = "aho-corasick"
+version = "1.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "alloc-no-stdlib"
+version = "3.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c2fb6cfd47bf496ff64095c20eaba0c201404ee38714d4142fcfa1dc334fcc7a"
+
+[[package]]
+name = "alloc-stdlib"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b5c1865780388bfa186411ab5f247819487fc4864c6e9c3106611fa347586e1"
+dependencies = [
+ "alloc-no-stdlib",
+]
+
+[[package]]
+name = "android_system_properties"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "anyhow"
+version = "1.0.104"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
+
+[[package]]
+name = "atk"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "241b621213072e993be4f6f3a9e4b45f65b7e6faad43001be957184b7bb1824b"
+dependencies = [
+ "atk-sys",
+ "glib",
+ "libc",
+]
+
+[[package]]
+name = "atk-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c5e48b684b0ca77d2bbadeef17424c2ea3c897d44d566a1617e7e8f30614d086"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "atomic-waker"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
+
+[[package]]
+name = "autocfg"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
+
+[[package]]
+name = "base64"
+version = "0.21.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567"
+
+[[package]]
+name = "base64"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
+
+[[package]]
+name = "base64"
+version = "0.23.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5"
+
+[[package]]
+name = "bit-set"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
+dependencies = [
+ "bit-vec",
+]
+
+[[package]]
+name = "bit-vec"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
+
+[[package]]
+name = "bitflags"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
+
+[[package]]
+name = "bitflags"
+version = "2.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "block-buffer"
+version = "0.10.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
+dependencies = [
+ "generic-array",
+]
+
+[[package]]
+name = "block2"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5"
+dependencies = [
+ "objc2",
+]
+
+[[package]]
+name = "brotli"
+version = "9.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8b851b75c23ca7873623d612fe49bd1989aeb03d08fb9432187eb253d3d4c6b"
+dependencies = [
+ "alloc-no-stdlib",
+ "alloc-stdlib",
+ "brotli-decompressor",
+]
+
+[[package]]
+name = "brotli-decompressor"
+version = "6.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "941cd9bd4ddab83cb46fa5a2d428f1c857b24ac78cb876cf7beb710840934bd7"
+dependencies = [
+ "alloc-no-stdlib",
+ "alloc-stdlib",
+]
+
+[[package]]
+name = "bs58"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4"
+dependencies = [
+ "tinyvec",
+]
+
+[[package]]
+name = "bumpalo"
+version = "3.20.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649"
+
+[[package]]
+name = "bytemuck"
+version = "1.25.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797"
+
+[[package]]
+name = "byteorder"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
+
+[[package]]
+name = "byteorder-lite"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495"
+
+[[package]]
+name = "bytes"
+version = "1.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "cairo-rs"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ca26ef0159422fb77631dc9d17b102f253b876fe1586b03b803e63a309b4ee2"
+dependencies = [
+ "bitflags 2.13.2",
+ "cairo-sys-rs",
+ "glib",
+ "libc",
+ "once_cell",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "cairo-sys-rs"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "685c9fa8e590b8b3d678873528d83411db17242a73fccaed827770ea0fedda51"
+dependencies = [
+ "glib-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "camino"
+version = "1.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbbad30e4b4c14a39e3cc8aed085a12a327257c316619c93581e017bc52be591"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "cargo-platform"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e35af189006b9c0f00a064685c727031e3ed2d8020f7ba284d78cc2671bd36ea"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "cargo_metadata"
+version = "0.19.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dd5eb614ed4c27c5d706420e4320fbe3216ab31fa1c33cd8246ac36dae4479ba"
+dependencies = [
+ "camino",
+ "cargo-platform",
+ "semver",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.21",
+]
+
+[[package]]
+name = "cargo_toml"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "82f4b26e751e711a5302649417f2da046dce6391b2ea30a4820f37462314f0b9"
+dependencies = [
+ "semver",
+ "serde",
+ "toml 1.1.6+spec-1.1.0",
+]
+
+[[package]]
+name = "cc"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f360145194ee8e21db5ee7f3fcd4fe52210864c75c985dae33218202c8bbe040"
+dependencies = [
+ "find-msvc-tools",
+ "shlex",
+]
+
+[[package]]
+name = "cesu8"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c"
+
+[[package]]
+name = "cfb"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a347dcabdae9c31b0825fd6a8bed285ec9c2acb89c47827126d52fa4f59cece3"
+dependencies = [
+ "fnv",
+ "uuid",
+ "web-time",
+]
+
+[[package]]
+name = "cfg-expr"
+version = "0.15.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02"
+dependencies = [
+ "smallvec",
+ "target-lexicon",
+]
+
+[[package]]
+name = "cfg-if"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600"
+
+[[package]]
+name = "cfg_aliases"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527"
+
+[[package]]
+name = "chacha20"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06"
+dependencies = [
+ "cfg-if",
+ "cpufeatures 0.3.1",
+ "rand_core",
+]
+
+[[package]]
+name = "chrono"
+version = "0.4.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327"
+dependencies = [
+ "iana-time-zone",
+ "num-traits",
+ "serde",
+ "windows-link",
+]
+
+[[package]]
+name = "combine"
+version = "4.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e"
+dependencies = [
+ "bytes",
+ "memchr",
+]
+
+[[package]]
+name = "cookie"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87"
+dependencies = [
+ "time",
+ "version_check",
+]
+
+[[package]]
+name = "core-foundation"
+version = "0.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6"
+dependencies = [
+ "core-foundation-sys",
+ "libc",
+]
+
+[[package]]
+name = "core-foundation-sys"
+version = "0.8.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
+
+[[package]]
+name = "core-graphics"
+version = "0.25.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "064badf302c3194842cf2c5d61f56cc88e54a759313879cdf03abdd27d0c3b97"
+dependencies = [
+ "bitflags 2.13.2",
+ "core-foundation",
+ "core-graphics-types",
+ "foreign-types",
+ "libc",
+]
+
+[[package]]
+name = "core-graphics-types"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb"
+dependencies = [
+ "bitflags 2.13.2",
+ "core-foundation",
+ "libc",
+]
+
+[[package]]
+name = "cpufeatures"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "cpufeatures"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "crc32fast"
+version = "1.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "01a7799fd6b852db0e61728dde9a204c423b44d689dbd432522543614b490e78"
+dependencies = [
+ "cfg-if",
+]
+
+[[package]]
+name = "crossbeam-channel"
+version = "0.5.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6"
+
+[[package]]
+name = "crypto-common"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a"
+dependencies = [
+ "generic-array",
+ "typenum",
+]
+
+[[package]]
+name = "cssparser"
+version = "0.37.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8c9cdaae01d5ed7882b04d795e7f752f46ff52d2fa3b50a20d28c464510bba98"
+dependencies = [
+ "cssparser-macros",
+ "dtoa-short",
+ "itoa",
+ "phf",
+ "smallvec",
+]
+
+[[package]]
+name = "cssparser-macros"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d045de693cb712d0b22c6a64be5b953f67b3ce00ab5ad3dd5d8b441886ab8e1a"
+dependencies = [
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "ctor"
+version = "1.0.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "914a755b7c2d4af2bdcff7ce1739e2db9a1b81a9b07123d8015786ae03c0980d"
+
+[[package]]
+name = "darling"
+version = "0.24.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed17f5901b6630b993ca003def43f2f8ef4014fc13b047b57aad617ff32bc2ec"
+dependencies = [
+ "darling_core",
+ "darling_macro",
+]
+
+[[package]]
+name = "darling_core"
+version = "0.24.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6837e2cf7485aaae18f86181d2f0e9a7ed297a025e220aeabf63fdebd3a2ddff"
+dependencies = [
+ "ident_case",
+ "proc-macro2",
+ "quote",
+ "strsim",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "darling_macro"
+version = "0.24.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2ac7135c3ef02b2f7833bbeb1be5ba7f966dcde8a87c6b87f65a778d71a02785"
+dependencies = [
+ "darling_core",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "dbus"
+version = "0.9.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3ab69f03cc8c4340c9c8e315114e1658e6775a9b16a04357973aa21cec22b32e"
+dependencies = [
+ "libc",
+ "libdbus-sys",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "defmt"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1"
+dependencies = [
+ "bitflags 1.3.2",
+ "defmt-macros",
+]
+
+[[package]]
+name = "defmt-macros"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8"
+dependencies = [
+ "defmt-parser",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "defmt-parser"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e"
+dependencies = [
+ "thiserror 2.0.21",
+]
+
+[[package]]
+name = "deranged"
+version = "0.5.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "derive_more"
+version = "2.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134"
+dependencies = [
+ "derive_more-impl",
+]
+
+[[package]]
+name = "derive_more-impl"
+version = "2.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "rustc_version",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "digest"
+version = "0.10.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
+dependencies = [
+ "block-buffer",
+ "crypto-common",
+]
+
+[[package]]
+name = "dirs"
+version = "7.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8d57d423b3c82e89b9a24ca3091fee61f456a26edbd28d26c65906f4bc1dcd8f"
+dependencies = [
+ "dirs-sys",
+]
+
+[[package]]
+name = "dirs-sys"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab"
+dependencies = [
+ "libc",
+ "option-ext",
+ "redox_users",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "dispatch2"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "libc",
+ "objc2",
+]
+
+[[package]]
+name = "displaydoc"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "dlopen2"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e2c5bd4158e66d1e215c49b837e11d62f3267b30c92f1d171c4d3105e3dc4d4"
+dependencies = [
+ "dlopen2_derive",
+ "libc",
+ "once_cell",
+ "winapi",
+]
+
+[[package]]
+name = "dlopen2_derive"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fbbb781877580993a8707ec48672673ec7b81eeba04cfd2310bd28c08e47c8f"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "dom_query"
+version = "0.28.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fac5fca71e65e94cc718a6e2af65d6e0f9c6027751c2aa562fbb5087fda639bc"
+dependencies = [
+ "bit-set",
+ "cssparser",
+ "foldhash",
+ "html5ever",
+ "precomputed-hash",
+ "selectors",
+ "tendril",
+]
+
+[[package]]
+name = "dpi"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8b14ccef22fc6f5a8f4d7d768562a182c04ce9a3b3157b91390b52ddfdf1a76"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "dtoa"
+version = "1.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4c3cf4824e2d5f025c7b531afcb2325364084a16806f6d47fbc1f5fbd9960590"
+
+[[package]]
+name = "dtoa-short"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd1511a7b6a56299bd043a9c167a6d2bfb37bf84a6dfceaba651168adfb43c87"
+dependencies = [
+ "dtoa",
+]
+
+[[package]]
+name = "dunce"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
+
+[[package]]
+name = "dyn-clone"
+version = "1.0.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
+
+[[package]]
+name = "edge0-app"
+version = "0.1.0"
+dependencies = [
+ "hex",
+ "reqwest 0.12.28",
+ "serde",
+ "serde_json",
+ "sha2",
+ "tauri",
+ "tauri-build",
+ "windows-sys 0.59.0",
+]
+
+[[package]]
+name = "embed-resource"
+version = "3.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbfdaacccebec3b28e4866b8973543c7647797db5ada1bdab552e48fe665fbbd"
+dependencies = [
+ "cc",
+ "memchr",
+ "rustc_version",
+ "toml 1.1.6+spec-1.1.0",
+ "vswhom",
+ "winreg",
+]
+
+[[package]]
+name = "embed_plist"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ef6b89e5b37196644d8796de5268852ff179b44e96276cf4290264843743bb7"
+
+[[package]]
+name = "equivalent"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
+
+[[package]]
+name = "erased-serde"
+version = "0.4.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2add8a07dd6a8d93ff627029c51de145e12686fbc36ecb298ac22e74cf02dec"
+dependencies = [
+ "serde",
+ "serde_core",
+ "typeid",
+]
+
+[[package]]
+name = "fastrand"
+version = "2.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223"
+
+[[package]]
+name = "fdeflate"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c"
+dependencies = [
+ "simd-adler32",
+]
+
+[[package]]
+name = "field-offset"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38e2275cc4e4fc009b0669731a1e5ab7ebf11f469eaede2bab9309a5b4d6057f"
+dependencies = [
+ "memoffset",
+ "rustc_version",
+]
+
+[[package]]
+name = "find-msvc-tools"
+version = "0.1.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aedcfb3409746eddb02b9e19ebda1c3394f759a152e48ee875a0844d1b955484"
+
+[[package]]
+name = "flate2"
+version = "1.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide 0.9.1",
+ "zlib-rs",
+]
+
+[[package]]
+name = "fnv"
+version = "1.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
+
+[[package]]
+name = "foldhash"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
+
+[[package]]
+name = "foreign-types"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965"
+dependencies = [
+ "foreign-types-macros",
+ "foreign-types-shared",
+]
+
+[[package]]
+name = "foreign-types-macros"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ea5190182e6915eb873ddbc16e23b711b6eb1f9c00a0d0a3a91b5f6228475225"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "foreign-types-shared"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b"
+
+[[package]]
+name = "form_urlencoded"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf"
+dependencies = [
+ "percent-encoding",
+]
+
+[[package]]
+name = "futures-channel"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4"
+dependencies = [
+ "futures-core",
+ "futures-sink",
+]
+
+[[package]]
+name = "futures-core"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e"
+
+[[package]]
+name = "futures-executor"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432"
+dependencies = [
+ "futures-core",
+ "futures-task",
+ "futures-util",
+]
+
+[[package]]
+name = "futures-io"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed"
+
+[[package]]
+name = "futures-macro"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "futures-sink"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d"
+
+[[package]]
+name = "futures-task"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd"
+
+[[package]]
+name = "futures-util"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc"
+dependencies = [
+ "futures-core",
+ "futures-io",
+ "futures-macro",
+ "futures-sink",
+ "futures-task",
+ "memchr",
+ "pin-project-lite",
+ "slab",
+]
+
+[[package]]
+name = "gdk"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d9f245958c627ac99d8e529166f9823fb3b838d1d41fd2b297af3075093c2691"
+dependencies = [
+ "cairo-rs",
+ "gdk-pixbuf",
+ "gdk-sys",
+ "gio",
+ "glib",
+ "libc",
+ "pango",
+]
+
+[[package]]
+name = "gdk-pixbuf"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "50e1f5f1b0bfb830d6ccc8066d18db35c487b1b2b1e8589b5dfe9f07e8defaec"
+dependencies = [
+ "gdk-pixbuf-sys",
+ "gio",
+ "glib",
+ "libc",
+ "once_cell",
+]
+
+[[package]]
+name = "gdk-pixbuf-sys"
+version = "0.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f9839ea644ed9c97a34d129ad56d38a25e6756f99f3a88e15cd39c20629caf7"
+dependencies = [
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "gdk-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c2d13f38594ac1e66619e188c6d5a1adb98d11b2fcf7894fc416ad76aa2f3f7"
+dependencies = [
+ "cairo-sys-rs",
+ "gdk-pixbuf-sys",
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "pango-sys",
+ "pkg-config",
+ "system-deps",
+]
+
+[[package]]
+name = "gdkwayland-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "140071d506d223f7572b9f09b5e155afbd77428cd5cc7af8f2694c41d98dfe69"
+dependencies = [
+ "gdk-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "pkg-config",
+ "system-deps",
+]
+
+[[package]]
+name = "gdkx11"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3caa00e14351bebbc8183b3c36690327eb77c49abc2268dd4bd36b856db3fbfe"
+dependencies = [
+ "gdk",
+ "gdkx11-sys",
+ "gio",
+ "glib",
+ "libc",
+ "x11",
+]
+
+[[package]]
+name = "gdkx11-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e2e7445fe01ac26f11601db260dd8608fe172514eb63b3b5e261ea6b0f4428d"
+dependencies = [
+ "gdk-sys",
+ "glib-sys",
+ "libc",
+ "system-deps",
+ "x11",
+]
+
+[[package]]
+name = "generic-array"
+version = "0.14.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
+dependencies = [
+ "typenum",
+ "version_check",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
+dependencies = [
+ "cfg-if",
+ "js-sys",
+ "libc",
+ "wasi",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "r-efi 5.3.0",
+ "wasip2",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099"
+dependencies = [
+ "cfg-if",
+ "js-sys",
+ "libc",
+ "r-efi 6.0.0",
+ "rand_core",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "gio"
+version = "0.18.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d4fc8f532f87b79cbc51a79748f16a6828fb784be93145a322fa14d06d354c73"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-io",
+ "futures-util",
+ "gio-sys",
+ "glib",
+ "libc",
+ "once_cell",
+ "pin-project-lite",
+ "smallvec",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "gio-sys"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37566df850baf5e4cb0dfb78af2e4b9898d817ed9263d1090a2df958c64737d2"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+ "winapi",
+]
+
+[[package]]
+name = "glib"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "233daaf6e83ae6a12a52055f568f9d7cf4671dabb78ff9560ab6da230ce00ee5"
+dependencies = [
+ "bitflags 2.13.2",
+ "futures-channel",
+ "futures-core",
+ "futures-executor",
+ "futures-task",
+ "futures-util",
+ "gio-sys",
+ "glib-macros",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "memchr",
+ "once_cell",
+ "smallvec",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "glib-macros"
+version = "0.18.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0bb0228f477c0900c880fd78c8759b95c7636dbd7842707f49e132378aa2acdc"
+dependencies = [
+ "heck 0.4.1",
+ "proc-macro-crate 2.0.2",
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "glib-sys"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "063ce2eb6a8d0ea93d2bf8ba1957e78dbab6be1c2220dd3daca57d5a9d869898"
+dependencies = [
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "glob"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b"
+
+[[package]]
+name = "gobject-sys"
+version = "0.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0850127b514d1c4a4654ead6dedadb18198999985908e6ffe4436f53c785ce44"
+dependencies = [
+ "glib-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "gtk"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd56fb197bfc42bd5d2751f4f017d44ff59fbb58140c6b49f9b3b2bdab08506a"
+dependencies = [
+ "atk",
+ "cairo-rs",
+ "field-offset",
+ "futures-channel",
+ "gdk",
+ "gdk-pixbuf",
+ "gio",
+ "glib",
+ "gtk-sys",
+ "gtk3-macros",
+ "libc",
+ "pango",
+ "pkg-config",
+]
+
+[[package]]
+name = "gtk-sys"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f29a1c21c59553eb7dd40e918be54dccd60c52b049b75119d5d96ce6b624414"
+dependencies = [
+ "atk-sys",
+ "cairo-sys-rs",
+ "gdk-pixbuf-sys",
+ "gdk-sys",
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "pango-sys",
+ "system-deps",
+]
+
+[[package]]
+name = "gtk3-macros"
+version = "0.18.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52ff3c5b21f14f0736fed6dcfc0bfb4225ebf5725f3c0209edeec181e4d73e9d"
+dependencies = [
+ "proc-macro-crate 1.3.1",
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "hashbrown"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+
+[[package]]
+name = "hashbrown"
+version = "0.17.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
+
+[[package]]
+name = "heck"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "hex"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
+
+[[package]]
+name = "html5ever"
+version = "0.39.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "46a1761807faccc9a19e86944bbf40610014066306f96edcdedc2fb714bcb7b8"
+dependencies = [
+ "log",
+ "markup5ever",
+]
+
+[[package]]
+name = "http"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0"
+dependencies = [
+ "bytes",
+ "itoa",
+]
+
+[[package]]
+name = "http-body"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c"
+dependencies = [
+ "bytes",
+ "http",
+]
+
+[[package]]
+name = "http-body-util"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23169fe34a5fbcdd3f3862e78fb9b6fccd5f02a6dc6f732547005d45631ce71c"
+dependencies = [
+ "bytes",
+ "futures-core",
+ "http",
+ "http-body",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "httparse"
+version = "1.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
+
+[[package]]
+name = "hyper"
+version = "1.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "27b501faa50e7a26c3d3560ca625132f4078a17771f4810baf70475ae48cbe43"
+dependencies = [
+ "atomic-waker",
+ "bytes",
+ "futures-channel",
+ "futures-core",
+ "http",
+ "http-body",
+ "httparse",
+ "itoa",
+ "pin-project-lite",
+ "smallvec",
+ "tokio",
+ "want",
+]
+
+[[package]]
+name = "hyper-rustls"
+version = "0.27.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dfa8e654703247911e29c23fbeaa261834bd9bb74efba2f9acddc37bfb127f53"
+dependencies = [
+ "http",
+ "hyper",
+ "hyper-util",
+ "rustls",
+ "tokio",
+ "tokio-rustls",
+ "tower-service",
+ "webpki-roots",
+]
+
+[[package]]
+name = "hyper-util"
+version = "0.1.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ddc03d96684f9226b8a787cdb71488417b53ab5ea8fdb1dac946cb9431cc8bff"
+dependencies = [
+ "base64 0.23.1",
+ "bytes",
+ "futures-channel",
+ "futures-util",
+ "http",
+ "http-body",
+ "httparse",
+ "hyper",
+ "ipnet",
+ "libc",
+ "percent-encoding",
+ "pin-project-lite",
+ "socket2",
+ "tokio",
+ "tower-service",
+ "tracing",
+]
+
+[[package]]
+name = "iana-time-zone"
+version = "0.1.65"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
+dependencies = [
+ "android_system_properties",
+ "core-foundation-sys",
+ "iana-time-zone-haiku",
+ "js-sys",
+ "log",
+ "wasm-bindgen",
+ "windows-core",
+]
+
+[[package]]
+name = "iana-time-zone-haiku"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
+dependencies = [
+ "cc",
+]
+
+[[package]]
+name = "ico"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3e795dff5605e0f04bff85ca41b51a96b83e80b281e96231bcaaf1ac35103371"
+dependencies = [
+ "byteorder",
+ "png 0.17.16",
+]
+
+[[package]]
+name = "icu_collections"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa68d21081c4a05d5a901a1c62add574c77048b6a1c67be3b50ce0b60d4ca513"
+dependencies = [
+ "displaydoc",
+ "potential_utf",
+ "utf8_iter",
+ "yoke",
+ "zerofrom",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_locale_core"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb"
+dependencies = [
+ "displaydoc",
+ "litemap",
+ "tinystr",
+ "writeable",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_normalizer"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "12f9cf5f235641ed274641dd81c3f28d870e276763d0797aeeab72317b1c646f"
+dependencies = [
+ "icu_collections",
+ "icu_normalizer_data",
+ "icu_properties",
+ "icu_provider",
+ "smallvec",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_normalizer_data"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1563da1ed3e0b3bf3d74c9b85917ac9c56464d2f57242270c09c9e752f8021a0"
+
+[[package]]
+name = "icu_properties"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e7ca276ad3145661a65914e6daf131ca5120cd3dcee8f8f3214b8875184a148"
+dependencies = [
+ "displaydoc",
+ "icu_collections",
+ "icu_locale_core",
+ "icu_properties_data",
+ "icu_provider",
+ "zerotrie",
+ "zerovec",
+]
+
+[[package]]
+name = "icu_properties_data"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa"
+
+[[package]]
+name = "icu_provider"
+version = "2.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d27bbb9d3abbefac45d55f647c9de1d44aafcd1186eb91879afef17c396c3e73"
+dependencies = [
+ "displaydoc",
+ "icu_locale_core",
+ "writeable",
+ "yoke",
+ "zerofrom",
+ "zerotrie",
+ "zerovec",
+]
+
+[[package]]
+name = "ident_case"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
+
+[[package]]
+name = "idna"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de"
+dependencies = [
+ "idna_adapter",
+ "smallvec",
+ "utf8_iter",
+]
+
+[[package]]
+name = "idna_adapter"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714"
+dependencies = [
+ "icu_normalizer",
+ "icu_properties",
+]
+
+[[package]]
+name = "image"
+version = "0.25.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85ab80394333c02fe689eaf900ab500fbd0c2213da414687ebf995a65d5a6104"
+dependencies = [
+ "bytemuck",
+ "byteorder-lite",
+ "moxcms",
+ "num-traits",
+ "png 0.18.1",
+]
+
+[[package]]
+name = "indexmap"
+version = "1.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
+dependencies = [
+ "autocfg",
+ "hashbrown 0.12.3",
+ "serde",
+]
+
+[[package]]
+name = "indexmap"
+version = "2.14.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855"
+dependencies = [
+ "equivalent",
+ "hashbrown 0.17.1",
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "infer"
+version = "0.22.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4200d433cbd5178df7797c9c2e75b348b728e39631cf14520d1e2fc424201f4"
+dependencies = [
+ "cfb",
+]
+
+[[package]]
+name = "ipnet"
+version = "2.12.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "791930b43c0d5973160d90a8f3894509f2b273430f5c5c73b668636d0287c5c0"
+
+[[package]]
+name = "itoa"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
+
+[[package]]
+name = "javascriptcore-rs"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca5671e9ffce8ffba57afc24070e906da7fc4b1ba66f2cabebf61bf2ea257fcc"
+dependencies = [
+ "bitflags 1.3.2",
+ "glib",
+ "javascriptcore-rs-sys",
+]
+
+[[package]]
+name = "javascriptcore-rs-sys"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "af1be78d14ffa4b75b66df31840478fef72b51f8c2465d4ca7c194da9f7a5124"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "jiff"
+version = "0.2.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ab1baf72f08796de0260609515130699b890ac25f30e610ad894bc5856cafdb"
+dependencies = [
+ "defmt",
+ "jiff-core",
+ "jiff-static",
+ "jiff-tzdb-platform",
+ "log",
+ "portable-atomic",
+ "portable-atomic-util",
+ "serde_core",
+ "windows-link",
+]
+
+[[package]]
+name = "jiff-core"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e52fe76043ccecc9005d2305ebaadf7d7fc0cc89ca6baa10a94d6bc68c7128c"
+dependencies = [
+ "defmt",
+ "log",
+]
+
+[[package]]
+name = "jiff-static"
+version = "0.2.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "378268a1116ad67ae6228701118ac9f491d78fda38a40a1f1a9e1348de6f7212"
+dependencies = [
+ "jiff-core",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "jiff-tzdb"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "142bd39932ad231f10513df9ab62661fead8719872150b7ad02a2df79f4e141e"
+
+[[package]]
+name = "jiff-tzdb-platform"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "875a5a69ac2bab1a891711cf5eccbec1ce0341ea805560dcd90b7a2e925132e8"
+dependencies = [
+ "jiff-tzdb",
+]
+
+[[package]]
+name = "jni"
+version = "0.21.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97"
+dependencies = [
+ "cesu8",
+ "cfg-if",
+ "combine",
+ "jni-sys 0.3.1",
+ "log",
+ "thiserror 1.0.69",
+ "walkdir",
+ "windows-sys 0.45.0",
+]
+
+[[package]]
+name = "jni-sys"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258"
+dependencies = [
+ "jni-sys 0.4.1",
+]
+
+[[package]]
+name = "jni-sys"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2"
+dependencies = [
+ "jni-sys-macros",
+]
+
+[[package]]
+name = "jni-sys-macros"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264"
+dependencies = [
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "js-sys"
+version = "0.3.106"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7883d941dae510fb2d978fc3fe018c71c9e2892fd38854de3e8b92c2e5ad9cc5"
+dependencies = [
+ "cfg-if",
+ "futures-util",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "json-patch"
+version = "4.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7421438de105a0827e44fadd05377727847d717c80ce29a229f85fd04c427b72"
+dependencies = [
+ "jsonptr",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.21",
+]
+
+[[package]]
+name = "jsonptr"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a5a3cc660ba5d72bce0b3bb295bf20847ccbb40fd423f3f05b61273672e561fe"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "keyboard-types"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fbe853b403ae61a04233030ae8a79d94975281ed9770a1f9e246732b534b28d"
+dependencies = [
+ "bitflags 2.13.2",
+ "serde",
+]
+
+[[package]]
+name = "libappindicator"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03589b9607c868cc7ae54c0b2a22c8dc03dd41692d48f2d7df73615c6a95dc0a"
+dependencies = [
+ "glib",
+ "gtk",
+ "gtk-sys",
+ "libappindicator-sys",
+ "log",
+]
+
+[[package]]
+name = "libappindicator-sys"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e9ec52138abedcc58dc17a7c6c0c00a2bdb4f3427c7f63fa97fd0d859155caf"
+dependencies = [
+ "gtk-sys",
+ "libloading",
+ "once_cell",
+]
+
+[[package]]
+name = "libc"
+version = "0.2.189"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
+
+[[package]]
+name = "libdbus-sys"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "328c4789d42200f1eeec05bd86c9c13c7f091d2ba9a6ea35acdf51f31bc0f043"
+dependencies = [
+ "pkg-config",
+]
+
+[[package]]
+name = "libloading"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f"
+dependencies = [
+ "cfg-if",
+ "winapi",
+]
+
+[[package]]
+name = "libredox"
+version = "0.1.25"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61ff90caf6077a803a240f62fdbe88645a890bbca49ef8174c3cb0404362171d"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "litemap"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae"
+
+[[package]]
+name = "lock_api"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
+dependencies = [
+ "scopeguard",
+]
+
+[[package]]
+name = "log"
+version = "0.4.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6"
+
+[[package]]
+name = "lru-slab"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4050469837a6ff301cd14c1f8f24f88549e6d548f24f64e2148eb0f72cebc51f"
+
+[[package]]
+name = "markup5ever"
+version = "0.39.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7122d987ec5f704ee56f6e5b41a7d93722e9aae27ae07cafa4036c4d3f9757de"
+dependencies = [
+ "log",
+ "tendril",
+ "web_atoms",
+]
+
+[[package]]
+name = "memchr"
+version = "2.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
+
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "mime"
+version = "0.3.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
+
+[[package]]
+name = "miniz_oxide"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316"
+dependencies = [
+ "adler2",
+ "simd-adler32",
+]
+
+[[package]]
+name = "miniz_oxide"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c"
+dependencies = [
+ "adler2",
+ "simd-adler32",
+]
+
+[[package]]
+name = "mio"
+version = "1.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b18443e9c262bfe8fa82f51666e2642c53393f7e5c27b3e1aeab922cff5b9d8"
+dependencies = [
+ "libc",
+ "wasi",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "moxcms"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bb85c154ba489f01b25c0d36ae69a87e4a1c73a72631fc6c0eb6dde34a73e44b"
+dependencies = [
+ "num-traits",
+ "pxfm",
+]
+
+[[package]]
+name = "muda"
+version = "0.20.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cca139e57da4383727f189e43e66e84fa372347027fc996f9f1a007ec0a37996"
+dependencies = [
+ "crossbeam-channel",
+ "dpi",
+ "gtk",
+ "keyboard-types",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+ "once_cell",
+ "png 0.18.1",
+ "serde",
+ "thiserror 2.0.21",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "ndk"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3f42e7bbe13d351b6bead8286a43aac9534b82bd3cc43e47037f012ebfd62d4"
+dependencies = [
+ "bitflags 2.13.2",
+ "jni-sys 0.3.1",
+ "log",
+ "ndk-sys",
+ "num_enum",
+ "raw-window-handle",
+ "thiserror 1.0.69",
+]
+
+[[package]]
+name = "ndk-context"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b"
+
+[[package]]
+name = "ndk-sys"
+version = "0.6.0+11769913"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873"
+dependencies = [
+ "jni-sys 0.3.1",
+]
+
+[[package]]
+name = "new_debug_unreachable"
+version = "1.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086"
+
+[[package]]
+name = "num-conv"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441"
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "num_enum"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26"
+dependencies = [
+ "num_enum_derive",
+ "rustversion",
+]
+
+[[package]]
+name = "num_enum_derive"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8"
+dependencies = [
+ "proc-macro-crate 3.5.0",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "objc2"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f"
+dependencies = [
+ "objc2-encode",
+ "objc2-exception-helper",
+]
+
+[[package]]
+name = "objc2-app-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d49e936b501e5c5bf01fda3a9452ff86dc3ea98ad5f283e1455153142d97518c"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-foundation",
+ "objc2-quartz-core",
+]
+
+[[package]]
+name = "objc2-cloud-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73ad74d880bb43877038da939b7427bba67e9dd42004a18b809ba7d87cee241c"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-data"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b402a653efbb5e82ce4df10683b6b28027616a2715e90009947d50b8dd298fa"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-foundation"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536"
+dependencies = [
+ "bitflags 2.13.2",
+ "dispatch2",
+ "objc2",
+]
+
+[[package]]
+name = "objc2-core-graphics"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807"
+dependencies = [
+ "bitflags 2.13.2",
+ "dispatch2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-io-surface",
+]
+
+[[package]]
+name = "objc2-core-image"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5d563b38d2b97209f8e861173de434bd0214cf020e3423a52624cd1d989f006"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-location"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca347214e24bc973fc025fd0d36ebb179ff30536ed1f80252706db19ee452009"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-core-text"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0cde0dfb48d25d2b4862161a4d5fcc0e3c24367869ad306b0c9ec0073bfed92d"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+]
+
+[[package]]
+name = "objc2-encode"
+version = "4.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33"
+
+[[package]]
+name = "objc2-exception-helper"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7a1c5fbb72d7735b076bb47b578523aedc40f3c439bea6dfd595c089d79d98a"
+dependencies = [
+ "cc",
+]
+
+[[package]]
+name = "objc2-foundation"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-core-foundation",
+]
+
+[[package]]
+name = "objc2-io-surface"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-core-foundation",
+]
+
+[[package]]
+name = "objc2-quartz-core"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f"
+dependencies = [
+ "bitflags 2.13.2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-ui-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d87d638e33c06f577498cbcc50491496a3ed4246998a7fbba7ccb98b1e7eab22"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-cloud-kit",
+ "objc2-core-data",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-core-image",
+ "objc2-core-location",
+ "objc2-core-text",
+ "objc2-foundation",
+ "objc2-quartz-core",
+ "objc2-user-notifications",
+]
+
+[[package]]
+name = "objc2-user-notifications"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9df9128cbbfef73cda168416ccf7f837b62737d748333bfe9ab71c245d76613e"
+dependencies = [
+ "objc2",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "objc2-web-kit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2e5aaab980c433cf470df9d7af96a7b46a9d892d521a2cbbb2f8a4c16751e7f"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "option-ext"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
+
+[[package]]
+name = "pango"
+version = "0.18.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ca27ec1eb0457ab26f3036ea52229edbdb74dee1edd29063f5b9b010e7ebee4"
+dependencies = [
+ "gio",
+ "glib",
+ "libc",
+ "once_cell",
+ "pango-sys",
+]
+
+[[package]]
+name = "pango-sys"
+version = "0.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "436737e391a843e5933d6d9aa102cb126d501e815b83601365a948a518555dc5"
+dependencies = [
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "parking_lot"
+version = "0.12.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
+dependencies = [
+ "lock_api",
+ "parking_lot_core",
+]
+
+[[package]]
+name = "parking_lot_core"
+version = "0.9.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "redox_syscall",
+ "smallvec",
+ "windows-link",
+]
+
+[[package]]
+name = "percent-encoding"
+version = "2.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
+
+[[package]]
+name = "phf"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf"
+dependencies = [
+ "phf_macros",
+ "phf_shared",
+ "serde",
+]
+
+[[package]]
+name = "phf_codegen"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49aa7f9d80421bca176ca8dbfebe668cc7a2684708594ec9f3c0db0805d5d6e1"
+dependencies = [
+ "phf_generator",
+ "phf_shared",
+]
+
+[[package]]
+name = "phf_generator"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "135ace3a761e564ec88c03a77317a7c6b80bb7f7135ef2544dbe054243b89737"
+dependencies = [
+ "fastrand",
+ "phf_shared",
+]
+
+[[package]]
+name = "phf_macros"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "812f032b54b1e759ccd5f8b6677695d5268c588701effba24601f6932f8269ef"
+dependencies = [
+ "phf_generator",
+ "phf_shared",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "phf_shared"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266"
+dependencies = [
+ "siphasher",
+]
+
+[[package]]
+name = "pin-project-lite"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
+
+[[package]]
+name = "pkg-config"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548"
+
+[[package]]
+name = "plist"
+version = "1.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2896bade328c13f7042a297ea5ac5b0951f6cf989dea5f32c2fd98da398195cb"
+dependencies = [
+ "base64 0.23.1",
+ "indexmap 2.14.2",
+ "quick-xml",
+ "serde",
+ "time",
+]
+
+[[package]]
+name = "png"
+version = "0.17.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "82151a2fc869e011c153adc57cf2789ccb8d9906ce52c0b39a6b5697749d7526"
+dependencies = [
+ "bitflags 1.3.2",
+ "crc32fast",
+ "fdeflate",
+ "flate2",
+ "miniz_oxide 0.8.9",
+]
+
+[[package]]
+name = "png"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61"
+dependencies = [
+ "bitflags 2.13.2",
+ "crc32fast",
+ "fdeflate",
+ "flate2",
+ "miniz_oxide 0.8.9",
+]
+
+[[package]]
+name = "portable-atomic"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "potential_utf"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661"
+dependencies = [
+ "zerovec",
+]
+
+[[package]]
+name = "powerfmt"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
+
+[[package]]
+name = "precomputed-hash"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c"
+
+[[package]]
+name = "proc-macro-crate"
+version = "1.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f4c021e1093a56626774e81216a4ce732a735e5bad4868a03f3ed65ca0c3919"
+dependencies = [
+ "once_cell",
+ "toml_edit 0.19.15",
+]
+
+[[package]]
+name = "proc-macro-crate"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b00f26d3400549137f92511a46ac1cd8ce37cb5598a96d382381458b992a5d24"
+dependencies = [
+ "toml_datetime 0.6.3",
+ "toml_edit 0.20.2",
+]
+
+[[package]]
+name = "proc-macro-crate"
+version = "3.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f"
+dependencies = [
+ "toml_edit 0.25.15+spec-1.1.0",
+]
+
+[[package]]
+name = "proc-macro-error"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c"
+dependencies = [
+ "proc-macro-error-attr",
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro-error-attr"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.107"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pxfm"
+version = "0.1.30"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d55d956fa96f5ec02be2e13af0e20391a5aa83d6a074e3ad368959d0fab299ea"
+
+[[package]]
+name = "quick-xml"
+version = "0.42.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41b1177fdf999d2321d3fb46ff47159d9c1fb9ad66a4879f8c50a0b504615e9b"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "quinn"
+version = "0.11.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4051e23e9185c255a7e33ef59cdbca87a22d359052eecd22fc6b901fb37d9d11"
+dependencies = [
+ "bytes",
+ "cfg_aliases",
+ "pin-project-lite",
+ "quinn-proto",
+ "quinn-udp",
+ "rustc-hash",
+ "rustls",
+ "socket2",
+ "thiserror 2.0.21",
+ "tokio",
+ "tracing",
+ "web-time",
+]
+
+[[package]]
+name = "quinn-proto"
+version = "0.11.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9746dbde176634f4f2f1faf2404e30a31b2bc1e9cafb5329c95d8177a18c9fc"
+dependencies = [
+ "bytes",
+ "getrandom 0.4.3",
+ "lru-slab",
+ "rand",
+ "rand_pcg",
+ "ring",
+ "rustc-hash",
+ "rustls",
+ "rustls-pki-types",
+ "slab",
+ "thiserror 2.0.21",
+ "tinyvec",
+ "tracing",
+ "web-time",
+]
+
+[[package]]
+name = "quinn-udp"
+version = "0.5.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694"
+dependencies = [
+ "cfg_aliases",
+ "libc",
+ "once_cell",
+ "socket2",
+ "tracing",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "r-efi"
+version = "5.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
+
+[[package]]
+name = "r-efi"
+version = "6.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
+
+[[package]]
+name = "rand"
+version = "0.10.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af"
+dependencies = [
+ "chacha20",
+ "getrandom 0.4.3",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69"
+
+[[package]]
+name = "rand_pcg"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a"
+dependencies = [
+ "rand_core",
+]
+
+[[package]]
+name = "raw-window-handle"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"
+
+[[package]]
+name = "redox_syscall"
+version = "0.5.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
+dependencies = [
+ "bitflags 2.13.2",
+]
+
+[[package]]
+name = "redox_users"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60dc65c0ff1a7ae1294b0c67b9f14baf70b644404010370171787bfac1038fc0"
+dependencies = [
+ "libredox",
+ "thiserror 2.0.21",
+]
+
+[[package]]
+name = "ref-cast"
+version = "1.0.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3"
+dependencies = [
+ "ref-cast-impl",
+]
+
+[[package]]
+name = "ref-cast-impl"
+version = "1.0.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "regex"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
+
+[[package]]
+name = "reqwest"
+version = "0.12.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
+dependencies = [
+ "base64 0.22.1",
+ "bytes",
+ "futures-channel",
+ "futures-core",
+ "futures-util",
+ "http",
+ "http-body",
+ "http-body-util",
+ "hyper",
+ "hyper-rustls",
+ "hyper-util",
+ "js-sys",
+ "log",
+ "percent-encoding",
+ "pin-project-lite",
+ "quinn",
+ "rustls",
+ "rustls-pki-types",
+ "serde",
+ "serde_json",
+ "serde_urlencoded",
+ "sync_wrapper",
+ "tokio",
+ "tokio-rustls",
+ "tower",
+ "tower-http",
+ "tower-service",
+ "url",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "web-sys",
+ "webpki-roots",
+]
+
+[[package]]
+name = "reqwest"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029"
+dependencies = [
+ "base64 0.23.1",
+ "bytes",
+ "futures-core",
+ "futures-util",
+ "http",
+ "http-body",
+ "http-body-util",
+ "hyper",
+ "hyper-util",
+ "js-sys",
+ "log",
+ "percent-encoding",
+ "pin-project-lite",
+ "serde",
+ "serde_json",
+ "sync_wrapper",
+ "tokio",
+ "tokio-util",
+ "tower",
+ "tower-http",
+ "tower-service",
+ "url",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "wasm-streams",
+ "web-sys",
+]
+
+[[package]]
+name = "ring"
+version = "0.17.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7"
+dependencies = [
+ "cc",
+ "cfg-if",
+ "getrandom 0.2.17",
+ "libc",
+ "untrusted",
+ "windows-sys 0.52.0",
+]
+
+[[package]]
+name = "rustc-hash"
+version = "2.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d"
+
+[[package]]
+name = "rustc_version"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92"
+dependencies = [
+ "semver",
+]
+
+[[package]]
+name = "rustls"
+version = "0.23.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634"
+dependencies = [
+ "once_cell",
+ "ring",
+ "rustls-pki-types",
+ "rustls-webpki",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
+name = "rustls-pki-types"
+version = "1.15.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96"
+dependencies = [
+ "web-time",
+ "zeroize",
+]
+
+[[package]]
+name = "rustls-webpki"
+version = "0.103.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2"
+dependencies = [
+ "ring",
+ "rustls-pki-types",
+ "untrusted",
+]
+
+[[package]]
+name = "rustversion"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f"
+
+[[package]]
+name = "ryu"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
+
+[[package]]
+name = "same-file"
+version = "1.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
+dependencies = [
+ "winapi-util",
+]
+
+[[package]]
+name = "schemars"
+version = "0.8.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615"
+dependencies = [
+ "dyn-clone",
+ "indexmap 1.9.3",
+ "schemars_derive",
+ "serde",
+ "serde_json",
+ "url",
+ "uuid",
+]
+
+[[package]]
+name = "schemars"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f"
+dependencies = [
+ "dyn-clone",
+ "ref-cast",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "schemars"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a"
+dependencies = [
+ "dyn-clone",
+ "ref-cast",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "schemars_derive"
+version = "0.8.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "serde_derive_internals",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "scopeguard"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
+
+[[package]]
+name = "selectors"
+version = "0.38.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8adfa1c298912827b8a28b223b3b874357397ae706e6190acd9bf28cee99114d"
+dependencies = [
+ "bitflags 2.13.2",
+ "cssparser",
+ "derive_more",
+ "log",
+ "new_debug_unreachable",
+ "phf",
+ "phf_codegen",
+ "precomputed-hash",
+ "rustc-hash",
+ "servo_arc",
+ "smallvec",
+]
+
+[[package]]
+name = "semver"
+version = "1.0.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd"
+dependencies = [
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "serde"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
+dependencies = [
+ "serde_core",
+ "serde_derive",
+]
+
+[[package]]
+name = "serde-untagged"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9faf48a4a2d2693be24c6289dbe26552776eb7737074e6722891fadbe6c5058"
+dependencies = [
+ "erased-serde",
+ "serde",
+ "serde_core",
+ "typeid",
+]
+
+[[package]]
+name = "serde_core"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "serde_derive_internals"
+version = "0.29.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "serde_json"
+version = "1.0.151"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14"
+dependencies = [
+ "itoa",
+ "memchr",
+ "serde",
+ "serde_core",
+ "zmij",
+]
+
+[[package]]
+name = "serde_repr"
+version = "0.1.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8d3b1629de253c70a0508c3899572da79ca359fdab27c7920ff00406df418906"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "serde_spanned"
+version = "0.6.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "serde_spanned"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "serde_urlencoded"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd"
+dependencies = [
+ "form_urlencoded",
+ "itoa",
+ "ryu",
+ "serde",
+]
+
+[[package]]
+name = "serde_with"
+version = "3.24.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df9adc193c780ef8f159aee8b61e2d5801aaa555e6eb0947fe45530ec506296f"
+dependencies = [
+ "base64 0.23.1",
+ "bs58",
+ "chrono",
+ "hex",
+ "indexmap 1.9.3",
+ "indexmap 2.14.2",
+ "jiff",
+ "schemars 0.9.0",
+ "schemars 1.2.2",
+ "serde_core",
+ "serde_json",
+ "serde_with_macros",
+ "time",
+]
+
+[[package]]
+name = "serde_with_macros"
+version = "3.24.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3e17bbc68e28663bbbb90df47e058aa7eda4fb445b89fe70457bb94fbccf6e49"
+dependencies = [
+ "darling",
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "serialize-to-javascript"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04f3666a07a197cdb77cdf306c32be9b7f598d7060d50cfd4d5aa04bfd92f6c5"
+dependencies = [
+ "serde",
+ "serde_json",
+ "serialize-to-javascript-impl",
+]
+
+[[package]]
+name = "serialize-to-javascript-impl"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "772ee033c0916d670af7860b6e1ef7d658a4629a6d0b4c8c3e67f09b3765b75d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "servo_arc"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "170fb83ab34de17dc69aa7c67482b22218ddb85da56546f9bd6b929e32a05930"
+dependencies = [
+ "stable_deref_trait",
+]
+
+[[package]]
+name = "sha2"
+version = "0.10.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
+dependencies = [
+ "cfg-if",
+ "cpufeatures 0.2.17",
+ "digest",
+]
+
+[[package]]
+name = "shlex"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba"
+
+[[package]]
+name = "simd-adler32"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea"
+
+[[package]]
+name = "siphasher"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "33f4fe9184a62d842c9ef383018f3306d8ba224fd9d836f56d7288308847c256"
+
+[[package]]
+name = "slab"
+version = "0.4.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5"
+
+[[package]]
+name = "smallvec"
+version = "1.16.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9395f0f0eee849a9b707b2f06bb92a6a422090e2123bb2ef8e87a0e61892a8e"
+
+[[package]]
+name = "socket2"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4"
+dependencies = [
+ "libc",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "softbuffer"
+version = "0.4.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aac18da81ebbf05109ab275b157c22a653bb3c12cf884450179942f81bcbf6c3"
+dependencies = [
+ "bytemuck",
+ "js-sys",
+ "ndk",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-foundation",
+ "objc2-quartz-core",
+ "raw-window-handle",
+ "redox_syscall",
+ "tracing",
+ "wasm-bindgen",
+ "web-sys",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "soup3"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "471f924a40f31251afc77450e781cb26d55c0b650842efafc9c6cbd2f7cc4f9f"
+dependencies = [
+ "futures-channel",
+ "gio",
+ "glib",
+ "libc",
+ "soup3-sys",
+]
+
+[[package]]
+name = "soup3-sys"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ebe8950a680a12f24f15ebe1bf70db7af98ad242d9db43596ad3108aab86c27"
+dependencies = [
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "libc",
+ "system-deps",
+]
+
+[[package]]
+name = "stable_deref_trait"
+version = "1.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596"
+
+[[package]]
+name = "string_cache"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a18596f8c785a729f2819c0f6a7eae6ebeebdfffbfe4214ae6b087f690e31901"
+dependencies = [
+ "new_debug_unreachable",
+ "parking_lot",
+ "phf_shared",
+ "precomputed-hash",
+]
+
+[[package]]
+name = "string_cache_codegen"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "585635e46db231059f76c5849798146164652513eb9e8ab2685939dd90f29b69"
+dependencies = [
+ "phf_generator",
+ "phf_shared",
+ "proc-macro2",
+ "quote",
+]
+
+[[package]]
+name = "strsim"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
+
+[[package]]
+name = "subtle"
+version = "2.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
+
+[[package]]
+name = "swift-rs"
+version = "1.0.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e45c444e496845d3f2a351146bff59aae4975b2280238df1dfaa0c7d1846f38e"
+dependencies = [
+ "base64 0.21.7",
+ "serde",
+ "serde_json",
+]
+
+[[package]]
+name = "syn"
+version = "1.0.109"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
+dependencies = [
+ "proc-macro2",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn"
+version = "2.0.119"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn"
+version = "3.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "sync_wrapper"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263"
+dependencies = [
+ "futures-core",
+]
+
+[[package]]
+name = "synstructure"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "901704edd0dfe137f1987838ee4f259e4e063c31371bdb423f7ae38ec6f77f02"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "system-deps"
+version = "6.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a3e535eb8dded36d55ec13eddacd30dec501792ff23a0b1682c38601b8cf2349"
+dependencies = [
+ "cfg-expr",
+ "heck 0.5.0",
+ "pkg-config",
+ "toml 0.8.2",
+ "version-compare",
+]
+
+[[package]]
+name = "tao"
+version = "0.37.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f37f381f4e048e6cdf038b5705f8cf14ad108279d46eb968140a7b291aba9400"
+dependencies = [
+ "bitflags 2.13.2",
+ "block2",
+ "core-foundation",
+ "core-graphics",
+ "crossbeam-channel",
+ "dbus",
+ "dispatch2",
+ "dlopen2",
+ "dpi",
+ "gdkwayland-sys",
+ "gdkx11-sys",
+ "gtk",
+ "jni",
+ "libc",
+ "log",
+ "ndk",
+ "ndk-context",
+ "ndk-sys",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-foundation",
+ "objc2-ui-kit",
+ "parking_lot",
+ "percent-encoding",
+ "raw-window-handle",
+ "tao-macros",
+ "unicode-segmentation",
+ "url",
+ "windows",
+ "windows-core",
+ "windows-version",
+ "x11-dl",
+]
+
+[[package]]
+name = "tao-macros"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5f7eeb6d99155545da6150a1795945f16ac9c178deb2a5f2e74d776107bd5849"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
+[[package]]
+name = "tauri"
+version = "2.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb20ef8e6f97d6fe03039666cd82861f09a38b56291fa2f0f9fe140023f61c1a"
+dependencies = [
+ "anyhow",
+ "bytes",
+ "cookie",
+ "dirs",
+ "dunce",
+ "embed_plist",
+ "getrandom 0.3.4",
+ "glob",
+ "gtk",
+ "heck 0.5.0",
+ "http",
+ "image",
+ "jni",
+ "libc",
+ "log",
+ "mime",
+ "muda",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-foundation",
+ "objc2-ui-kit",
+ "objc2-web-kit",
+ "percent-encoding",
+ "plist",
+ "raw-window-handle",
+ "reqwest 0.13.5",
+ "serde",
+ "serde_json",
+ "serde_repr",
+ "serialize-to-javascript",
+ "swift-rs",
+ "tauri-build",
+ "tauri-macros",
+ "tauri-runtime",
+ "tauri-runtime-wry",
+ "tauri-utils",
+ "thiserror 2.0.21",
+ "tokio",
+ "tray-icon",
+ "url",
+ "webkit2gtk",
+ "webview2-com",
+ "window-vibrancy",
+ "windows",
+]
+
+[[package]]
+name = "tauri-build"
+version = "2.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b5ae674f48836f5dd2eeaf9e221c095bd3e78a8d0439c77ca93b48740a1b644"
+dependencies = [
+ "anyhow",
+ "cargo_toml",
+ "dirs",
+ "glob",
+ "heck 0.5.0",
+ "json-patch",
+ "schemars 0.8.22",
+ "semver",
+ "serde",
+ "serde_json",
+ "tauri-utils",
+ "tauri-winres",
+ "walkdir",
+]
+
+[[package]]
+name = "tauri-codegen"
+version = "2.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3926b73ac01df7c14f19ff872cd56e67213a87acc0000ea37ddb300f3d11a3e3"
+dependencies = [
+ "base64 0.22.1",
+ "brotli",
+ "ico",
+ "json-patch",
+ "plist",
+ "png 0.18.1",
+ "proc-macro2",
+ "quote",
+ "semver",
+ "serde",
+ "serde_json",
+ "sha2",
+ "syn 2.0.119",
+ "tauri-utils",
+ "thiserror 2.0.21",
+ "time",
+ "url",
+ "uuid",
+ "walkdir",
+]
+
+[[package]]
+name = "tauri-macros"
+version = "2.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d21d0fa2c0529972ea6ecaf7404fcf40efac098eda0beb06ca682f41e33a6aa5"
+dependencies = [
+ "heck 0.5.0",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+ "tauri-codegen",
+ "tauri-utils",
+]
+
+[[package]]
+name = "tauri-runtime"
+version = "2.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c9e6d5ef76985b32ff012a43b58df205d25a285e400ca878afcacddc5d5f3238"
+dependencies = [
+ "cookie",
+ "dpi",
+ "gtk",
+ "http",
+ "jni",
+ "objc2",
+ "objc2-ui-kit",
+ "objc2-web-kit",
+ "raw-window-handle",
+ "serde",
+ "serde_json",
+ "tauri-utils",
+ "thiserror 2.0.21",
+ "url",
+ "webkit2gtk",
+ "webview2-com",
+ "windows",
+]
+
+[[package]]
+name = "tauri-runtime-wry"
+version = "2.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa829eb69860cf1fb0bedb21d2d5f1effce3527d054a68ca1cd7355c654400a3"
+dependencies = [
+ "gtk",
+ "http",
+ "jni",
+ "log",
+ "objc2",
+ "objc2-app-kit",
+ "once_cell",
+ "percent-encoding",
+ "raw-window-handle",
+ "softbuffer",
+ "tao",
+ "tauri-runtime",
+ "tauri-utils",
+ "url",
+ "webkit2gtk",
+ "webview2-com",
+ "windows",
+ "wry",
+]
+
+[[package]]
+name = "tauri-utils"
+version = "2.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4dd257082deff79550de3409ed146104aeaf86fe38649615b04a304144939677"
+dependencies = [
+ "anyhow",
+ "brotli",
+ "cargo_metadata",
+ "ctor",
+ "dom_query",
+ "dunce",
+ "glob",
+ "http",
+ "infer",
+ "json-patch",
+ "log",
+ "memchr",
+ "phf",
+ "plist",
+ "proc-macro2",
+ "quote",
+ "regex",
+ "schemars 0.8.22",
+ "semver",
+ "serde",
+ "serde-untagged",
+ "serde_json",
+ "serde_with",
+ "swift-rs",
+ "thiserror 2.0.21",
+ "toml 1.1.6+spec-1.1.0",
+ "url",
+ "urlpattern",
+ "uuid",
+ "walkdir",
+]
+
+[[package]]
+name = "tauri-winres"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cc65d45c68858bfe420dd29e834b5d15dbecf8a07a8a16cf4d532c7b1f69d4b6"
+dependencies = [
+ "dunce",
+ "embed-resource",
+ "toml 1.1.6+spec-1.1.0",
+]
+
+[[package]]
+name = "tendril"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5fed54709c5b3a53d09bb1c113ea4f5ceafd1e772ddcb0030a82e1d56c087b08"
+dependencies = [
+ "new_debug_unreachable",
+]
+
+[[package]]
+name = "thiserror"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52"
+dependencies = [
+ "thiserror-impl 1.0.69",
+]
+
+[[package]]
+name = "thiserror"
+version = "2.0.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e"
+dependencies = [
+ "thiserror-impl 2.0.21",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "1.0.69"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "2.0.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "time"
+version = "0.3.55"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134"
+dependencies = [
+ "deranged",
+ "num-conv",
+ "powerfmt",
+ "serde_core",
+ "time-core",
+ "time-macros",
+]
+
+[[package]]
+name = "time-core"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109"
+
+[[package]]
+name = "time-macros"
+version = "0.2.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e689342a48d2ea927c87ea50cabf8594854bf940e9310208848d680d668ed85"
+dependencies = [
+ "num-conv",
+ "time-core",
+]
+
+[[package]]
+name = "tinystr"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643"
+dependencies = [
+ "displaydoc",
+ "zerovec",
+]
+
+[[package]]
+name = "tinyvec"
+version = "1.13.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd3ca314f692efd6c868f8408f53fe444634a845f96c028b97d35f6a1f79f0ee"
+
+[[package]]
+name = "tokio"
+version = "1.53.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed"
+dependencies = [
+ "bytes",
+ "libc",
+ "mio",
+ "pin-project-lite",
+ "socket2",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "tokio-rustls"
+version = "0.26.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c9cc2678c2cdd569ef8215e2afd7954ada2ae20b4fdd2c5fe6139a3b02d105db"
+dependencies = [
+ "rustls",
+ "tokio",
+]
+
+[[package]]
+name = "tokio-util"
+version = "0.7.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52"
+dependencies = [
+ "bytes",
+ "futures-core",
+ "futures-sink",
+ "pin-project-lite",
+ "tokio",
+]
+
+[[package]]
+name = "toml"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "185d8ab0dfbb35cf1399a6344d8484209c088f75f8f68230da55d48d95d43e3d"
+dependencies = [
+ "serde",
+ "serde_spanned 0.6.9",
+ "toml_datetime 0.6.3",
+ "toml_edit 0.20.2",
+]
+
+[[package]]
+name = "toml"
+version = "1.1.6+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "920602543f0911ab71da12c50d59701da54c196d1a2bf5cb4b75667f137a406a"
+dependencies = [
+ "indexmap 2.14.2",
+ "serde_core",
+ "serde_spanned 1.1.1",
+ "toml_datetime 1.1.1+spec-1.1.0",
+ "toml_parser",
+ "toml_writer",
+ "winnow 1.0.4",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7cda73e2f1397b1262d6dfdcef8aafae14d1de7748d66822d3bfeeb6d03e5e4b"
+dependencies = [
+ "serde",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "1.1.1+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.19.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421"
+dependencies = [
+ "indexmap 2.14.2",
+ "toml_datetime 0.6.3",
+ "winnow 0.5.40",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "396e4d48bbb2b7554c944bde63101b5ae446cff6ec4a24227428f15eb72ef338"
+dependencies = [
+ "indexmap 2.14.2",
+ "serde",
+ "serde_spanned 0.6.9",
+ "toml_datetime 0.6.3",
+ "winnow 0.5.40",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.25.15+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1340ea94a5856333492c9064b02c778b191dd2c853778d9609debdcdfea3a614"
+dependencies = [
+ "indexmap 2.14.2",
+ "toml_datetime 1.1.1+spec-1.1.0",
+ "toml_parser",
+ "winnow 1.0.4",
+]
+
+[[package]]
+name = "toml_parser"
+version = "1.1.3+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56"
+dependencies = [
+ "winnow 1.0.4",
+]
+
+[[package]]
+name = "toml_writer"
+version = "1.1.2+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2"
+
+[[package]]
+name = "tower"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4"
+dependencies = [
+ "futures-core",
+ "futures-util",
+ "pin-project-lite",
+ "sync_wrapper",
+ "tokio",
+ "tower-layer",
+ "tower-service",
+]
+
+[[package]]
+name = "tower-http"
+version = "0.6.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840"
+dependencies = [
+ "bitflags 2.13.2",
+ "bytes",
+ "futures-util",
+ "http",
+ "http-body",
+ "pin-project-lite",
+ "tower",
+ "tower-layer",
+ "tower-service",
+ "url",
+]
+
+[[package]]
+name = "tower-layer"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e"
+
+[[package]]
+name = "tower-service"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3"
+
+[[package]]
+name = "tracing"
+version = "0.1.44"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
+dependencies = [
+ "pin-project-lite",
+ "tracing-core",
+]
+
+[[package]]
+name = "tracing-core"
+version = "0.1.36"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a"
+dependencies = [
+ "once_cell",
+]
+
+[[package]]
+name = "tray-icon"
+version = "0.25.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2b9c52859a94554803ccd4a24b98f74148ebc73b90676d783f3490b1bff9d72"
+dependencies = [
+ "crossbeam-channel",
+ "dirs",
+ "libappindicator",
+ "muda",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-foundation",
+ "once_cell",
+ "png 0.18.1",
+ "serde",
+ "thiserror 2.0.21",
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "try-lock"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
+
+[[package]]
+name = "typeid"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c"
+
+[[package]]
+name = "typenum"
+version = "1.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.26"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954"
+
+[[package]]
+name = "unicode-segmentation"
+version = "1.13.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
+
+[[package]]
+name = "untrusted"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
+
+[[package]]
+name = "url"
+version = "2.5.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed"
+dependencies = [
+ "form_urlencoded",
+ "idna",
+ "percent-encoding",
+ "serde",
+ "serde_derive",
+]
+
+[[package]]
+name = "urlpattern"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df16f50ef4cc145211879a3867ba757076b25dfee812040dcb0658bd9ae7904b"
+dependencies = [
+ "icu_properties",
+ "regex",
+ "serde",
+ "url",
+]
+
+[[package]]
+name = "utf8_iter"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
+
+[[package]]
+name = "uuid"
+version = "1.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2ef6dac1e96601b4fb3acccccff2139741fcb757cb9a36089bf5be91cfb285ce"
+dependencies = [
+ "getrandom 0.4.3",
+ "js-sys",
+ "serde_core",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "version-compare"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03c2856837ef78f57382f06b2b8563a2f512f7185d732608fd9176cb3b8edf0e"
+
+[[package]]
+name = "version_check"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
+
+[[package]]
+name = "vswhom"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "be979b7f07507105799e854203b470ff7c78a1639e330a58f183b5fea574608b"
+dependencies = [
+ "libc",
+ "vswhom-sys",
+]
+
+[[package]]
+name = "vswhom-sys"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb067e4cbd1ff067d1df46c9194b5de0e98efd2810bbc95c5d5e5f25a3231150"
+dependencies = [
+ "cc",
+ "libc",
+]
+
+[[package]]
+name = "walkdir"
+version = "2.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
+dependencies = [
+ "same-file",
+ "winapi-util",
+]
+
+[[package]]
+name = "want"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e"
+dependencies = [
+ "try-lock",
+]
+
+[[package]]
+name = "wasi"
+version = "0.11.1+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
+
+[[package]]
+name = "wasip2"
+version = "1.0.4+wasi-0.2.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487"
+dependencies = [
+ "wit-bindgen",
+]
+
+[[package]]
+name = "wasm-bindgen"
+version = "0.2.129"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9bb54f33acc68fd454578d9820b0bde1a1a3d17aa17bb7b6595806d02886d409"
+dependencies = [
+ "cfg-if",
+ "once_cell",
+ "rustversion",
+ "wasm-bindgen-macro",
+ "wasm-bindgen-shared",
+]
+
+[[package]]
+name = "wasm-bindgen-futures"
+version = "0.4.79"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3cbab34de2d982e9b48e18d216d04c4a6f641066ff19ffb699980f591ee3610e"
+dependencies = [
+ "js-sys",
+ "tokio",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "wasm-bindgen-macro"
+version = "0.2.129"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2e29d0c35b16e224a7eeb5cd2d25e3e1968fbd65604117b44d3b789d00ee8535"
+dependencies = [
+ "quote",
+ "wasm-bindgen-macro-support",
+]
+
+[[package]]
+name = "wasm-bindgen-macro-support"
+version = "0.2.129"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f501a8bc3719dba86ef8ae4728879c08001bea749eb1333ac5b91e040e2a6b7"
+dependencies = [
+ "bumpalo",
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+ "wasm-bindgen-shared",
+]
+
+[[package]]
+name = "wasm-bindgen-shared"
+version = "0.2.129"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23f0c9c52aa7cd7d77769a4cfe2a9adb1b331f489a41d912ce14513d5ab995c6"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "wasm-streams"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb"
+dependencies = [
+ "futures-util",
+ "js-sys",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "web-sys",
+]
+
+[[package]]
+name = "web-sys"
+version = "0.3.106"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "88261b9deccee56594c11a3460c462c41f58d148598fe70ad77070126a68aba4"
+dependencies = [
+ "js-sys",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "web-time"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb"
+dependencies = [
+ "js-sys",
+ "wasm-bindgen",
+]
+
+[[package]]
+name = "web_atoms"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba8b815c1b593dc0baf78dd0f4fc8fdb2de53198fb1163738093e9a311c33fb3"
+dependencies = [
+ "phf",
+ "phf_codegen",
+ "string_cache",
+ "string_cache_codegen",
+]
+
+[[package]]
+name = "webkit2gtk"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1027150013530fb2eaf806408df88461ae4815a45c541c8975e61d6f2fc4793"
+dependencies = [
+ "bitflags 1.3.2",
+ "cairo-rs",
+ "gdk",
+ "gdk-sys",
+ "gio",
+ "gio-sys",
+ "glib",
+ "glib-sys",
+ "gobject-sys",
+ "gtk",
+ "gtk-sys",
+ "javascriptcore-rs",
+ "libc",
+ "once_cell",
+ "soup3",
+ "webkit2gtk-sys",
+]
+
+[[package]]
+name = "webkit2gtk-sys"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "916a5f65c2ef0dfe12fff695960a2ec3d4565359fdbb2e9943c974e06c734ea5"
+dependencies = [
+ "bitflags 1.3.2",
+ "cairo-sys-rs",
+ "gdk-sys",
+ "gio-sys",
+ "glib-sys",
+ "gobject-sys",
+ "gtk-sys",
+ "javascriptcore-rs-sys",
+ "libc",
+ "pkg-config",
+ "soup3-sys",
+ "system-deps",
+]
+
+[[package]]
+name = "webpki-roots"
+version = "1.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a"
+dependencies = [
+ "rustls-pki-types",
+]
+
+[[package]]
+name = "webview2-com"
+version = "0.39.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f89fca7a704cee10dcb3654c1dbb8941d1783132f1917358af75bec37a7d7e6"
+dependencies = [
+ "webview2-com-macros",
+ "webview2-com-sys",
+ "windows",
+ "windows-core",
+]
+
+[[package]]
+name = "webview2-com-macros"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67a921c1b6914c367b2b823cd4cde6f96beec77d30a939c8199bb377cf9b9b54"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "webview2-com-sys"
+version = "0.39.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3a07132775117d6065853d9d1178157b8c90e228de47129d6bce2c7edebedfb"
+dependencies = [
+ "thiserror 2.0.21",
+ "windows",
+ "windows-core",
+]
+
+[[package]]
+name = "winapi"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
+dependencies = [
+ "winapi-i686-pc-windows-gnu",
+ "winapi-x86_64-pc-windows-gnu",
+]
+
+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
+
+[[package]]
+name = "winapi-util"
+version = "0.1.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
+dependencies = [
+ "windows-sys 0.61.2",
+]
+
+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
+
+[[package]]
+name = "window-vibrancy"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "111e51caca442cafd9bab396628ac4d814880eaea4e63dfd76a12e9f342b5580"
+dependencies = [
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+ "objc2-quartz-core",
+ "raw-window-handle",
+ "windows-sys 0.61.2",
+ "windows-version",
+]
+
+[[package]]
+name = "windows"
+version = "0.62.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580"
+dependencies = [
+ "windows-collections",
+ "windows-core",
+ "windows-future",
+ "windows-numerics",
+]
+
+[[package]]
+name = "windows-collections"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610"
+dependencies = [
+ "windows-core",
+]
+
+[[package]]
+name = "windows-core"
+version = "0.62.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
+dependencies = [
+ "windows-implement",
+ "windows-interface",
+ "windows-link",
+ "windows-result",
+ "windows-strings",
+]
+
+[[package]]
+name = "windows-future"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb"
+dependencies = [
+ "windows-core",
+ "windows-link",
+ "windows-threading",
+]
+
+[[package]]
+name = "windows-implement"
+version = "0.60.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "windows-interface"
+version = "0.59.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "windows-link"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
+
+[[package]]
+name = "windows-numerics"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26"
+dependencies = [
+ "windows-core",
+ "windows-link",
+]
+
+[[package]]
+name = "windows-result"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5"
+dependencies = [
+ "windows-link",
+]
+
+[[package]]
+name = "windows-strings"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091"
+dependencies = [
+ "windows-link",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.45.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0"
+dependencies = [
+ "windows-targets 0.42.2",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.52.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
+dependencies = [
+ "windows-targets 0.52.6",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.59.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
+dependencies = [
+ "windows-targets 0.52.6",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.61.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
+dependencies = [
+ "windows-link",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071"
+dependencies = [
+ "windows_aarch64_gnullvm 0.42.2",
+ "windows_aarch64_msvc 0.42.2",
+ "windows_i686_gnu 0.42.2",
+ "windows_i686_msvc 0.42.2",
+ "windows_x86_64_gnu 0.42.2",
+ "windows_x86_64_gnullvm 0.42.2",
+ "windows_x86_64_msvc 0.42.2",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
+dependencies = [
+ "windows_aarch64_gnullvm 0.52.6",
+ "windows_aarch64_msvc 0.52.6",
+ "windows_i686_gnu 0.52.6",
+ "windows_i686_gnullvm",
+ "windows_i686_msvc 0.52.6",
+ "windows_x86_64_gnu 0.52.6",
+ "windows_x86_64_gnullvm 0.52.6",
+ "windows_x86_64_msvc 0.52.6",
+]
+
+[[package]]
+name = "windows-threading"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37"
+dependencies = [
+ "windows-link",
+]
+
+[[package]]
+name = "windows-version"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4060a1da109b9d0326b7262c8e12c84df67cc0dbc9e33cf49e01ccc2eb63631"
+dependencies = [
+ "windows-link",
+]
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8"
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
+
+[[package]]
+name = "windows_i686_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
+
+[[package]]
+name = "winnow"
+version = "0.5.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f593a95398737aeed53e489c785df13f3618e41dbcd6718c6addbf1395aa6876"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "winnow"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "winreg"
+version = "0.55.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb5a765337c50e9ec252c2069be9bf91c7df47afb103b642ba3a53bf8101be97"
+dependencies = [
+ "cfg-if",
+ "windows-sys 0.59.0",
+]
+
+[[package]]
+name = "wit-bindgen"
+version = "0.57.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
+
+[[package]]
+name = "writeable"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc"
+
+[[package]]
+name = "wry"
+version = "0.57.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a819957a01b3119af85e638a38d242af76dbc87d130dca67bfd0441072e21ff0"
+dependencies = [
+ "base64 0.22.1",
+ "block2",
+ "cookie",
+ "crossbeam-channel",
+ "dirs",
+ "dom_query",
+ "dpi",
+ "dunce",
+ "gdkx11",
+ "gtk",
+ "http",
+ "javascriptcore-rs",
+ "jni",
+ "libc",
+ "ndk",
+ "objc2",
+ "objc2-app-kit",
+ "objc2-core-foundation",
+ "objc2-foundation",
+ "objc2-ui-kit",
+ "objc2-web-kit",
+ "once_cell",
+ "percent-encoding",
+ "raw-window-handle",
+ "sha2",
+ "soup3",
+ "tao-macros",
+ "thiserror 2.0.21",
+ "url",
+ "webkit2gtk",
+ "webkit2gtk-sys",
+ "webview2-com",
+ "windows",
+ "windows-core",
+ "windows-version",
+ "x11-dl",
+]
+
+[[package]]
+name = "x11"
+version = "2.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "502da5464ccd04011667b11c435cb992822c2c0dbde1770c988480d312a0db2e"
+dependencies = [
+ "libc",
+ "pkg-config",
+]
+
+[[package]]
+name = "x11-dl"
+version = "2.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38735924fedd5314a6e548792904ed8c6de6636285cb9fec04d5b1db85c1516f"
+dependencies = [
+ "libc",
+ "once_cell",
+ "pkg-config",
+]
+
+[[package]]
+name = "yoke"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5"
+dependencies = [
+ "stable_deref_trait",
+ "yoke-derive",
+ "zerofrom",
+]
+
+[[package]]
+name = "yoke-derive"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "33811428bee40dbceb6d545e95754741d17a6aef9a4849f0fd62e2ba4f412a78"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+ "synstructure",
+]
+
+[[package]]
+name = "zerofrom"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272"
+dependencies = [
+ "zerofrom-derive",
+]
+
+[[package]]
+name = "zerofrom-derive"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f75b4683f6c7f45248d4d64056a24298c6281e0993356d7d1b4a1a962ef10d4a"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+ "synstructure",
+]
+
+[[package]]
+name = "zeroize"
+version = "1.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e"
+
+[[package]]
+name = "zerotrie"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ea269c3bd32f0a32c321907a2ae912ba6f4649bb0fc764a15627e99a7095a3f"
+dependencies = [
+ "displaydoc",
+ "yoke",
+ "zerofrom",
+]
+
+[[package]]
+name = "zerovec"
+version = "0.11.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bb0464e17806c1d976d5cba29399c7f08e516e279e2ba493f63123b5fca67dd8"
+dependencies = [
+ "yoke",
+ "zerofrom",
+ "zerovec-derive",
+]
+
+[[package]]
+name = "zerovec-derive"
+version = "0.11.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34df6fc39dbd26ddc9c10e6a2984476e13acce22e64e4487636ef494369225da"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.6",
+]
+
+[[package]]
+name = "zlib-rs"
+version = "0.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b268e58e7c693d7c271f93ffc4ba3b380412554231c85bf61ca7af91042a4112"
+
+[[package]]
+name = "zmij"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
diff --git a/windows/app/src-tauri/Cargo.toml b/windows/app/src-tauri/Cargo.toml
new file mode 100644
index 0000000..e085233
--- /dev/null
+++ b/windows/app/src-tauri/Cargo.toml
@@ -0,0 +1,31 @@
+[package]
+name = "edge0-app"
+version = "0.1.0"
+description = "edge0 Windows desktop shell (Tauri 2) — pull/convert/load/chat"
+edition = "2021"
+rust-version = "1.88"
+
+[lib]
+name = "edge0_app_lib"
+crate-type = ["staticlib", "cdylib", "rlib"]
+
+[build-dependencies]
+tauri-build = { version = "2", features = [] }
+
+[dependencies]
+tauri = { version = "2", features = ["tray-icon", "image-png"] }
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "blocking", "json"] }
+sha2 = "0.10"
+hex = "0.4"
+
+[target.'cfg(windows)'.dependencies]
+# Job Object binding for kill-on-close process governance (engine child dies with the shell).
+# Win32_Threading is required: JOBOBJECT_EXTENDED_LIMIT_INFORMATION is gated behind it.
+windows-sys = { version = "0.59", features = ["Win32_Foundation", "Win32_Security", "Win32_System_JobObjects", "Win32_System_Threading"] }
+
+[profile.release]
+codegen-units = 1
+lto = "thin"
+strip = true
diff --git a/windows/app/src-tauri/build.rs b/windows/app/src-tauri/build.rs
new file mode 100644
index 0000000..d860e1e
--- /dev/null
+++ b/windows/app/src-tauri/build.rs
@@ -0,0 +1,3 @@
+fn main() {
+    tauri_build::build()
+}
diff --git a/windows/app/src-tauri/capabilities/default.json b/windows/app/src-tauri/capabilities/default.json
new file mode 100644
index 0000000..3ae0133
--- /dev/null
+++ b/windows/app/src-tauri/capabilities/default.json
@@ -0,0 +1,7 @@
+{
+  "$schema": "../gen/schemas/desktop-schema.json",
+  "identifier": "default",
+  "description": "Main window capability — core plugins only; all app commands are native (no extra IPC surface)",
+  "windows": ["main"],
+  "permissions": ["core:default", "core:event:default", "core:window:allow-set-title"]
+}
diff --git a/windows/app/src-tauri/catalog.json b/windows/app/src-tauri/catalog.json
new file mode 100644
index 0000000..c99fd92
--- /dev/null
+++ b/windows/app/src-tauri/catalog.json
@@ -0,0 +1,198 @@
+{
+ "8b": {
+  "repo": "Edge0/Edge0-8B-A1B-preview",
+  "rev": {
+   "modelscope": "master",
+   "hf": "cadc35c7059469fad8569a63ce6c851172bbca61"
+  },
+  "files": [
+   {
+    "path": "chat_template.jinja",
+    "size": 6041,
+    "sha256": "eb6226c94ae38058f875d159f86a206b3a165828c0e7d6bda664ae14667f798a",
+    "skip": false
+   },
+   {
+    "path": "config.json",
+    "size": 3069,
+    "sha256": "5790c795d490d7b90eb04b269f5db4faad08bcc560f8da6e30b61057b1606340",
+    "skip": false
+   },
+   {
+    "path": "generation_config.json",
+    "size": 56,
+    "sha256": "64752c5973a55faf4cfc02604c7587c38b090f00013f91191f52362dcc79a4a8",
+    "skip": false
+   },
+   {
+    "path": "lora_edge0_8b.safetensors",
+    "size": 16413656,
+    "sha256": "a32ff07ba3d3b3c9811146c43a3f0e96af4ec2960ebc5072910f9a6663dc2e2e",
+    "skip": false
+   },
+   {
+    "path": "model.safetensors",
+    "size": 4508697328,
+    "sha256": "5bcde14438cfe13a547965573471e8ddc830de3a047dd2dfd38a0eb2b390d2b6",
+    "skip": false
+   },
+   {
+    "path": "prerouter_edge0_8b.safetensors",
+    "size": 38802304,
+    "sha256": "5f05d25e52ce18103d0867e5a0c79ec4e3334ad8c8906c22ec6266a95db2047c",
+    "skip": true
+   },
+   {
+    "path": "special_tokens_map.json",
+    "size": 580,
+    "sha256": "69b63b9f81044ead642d16a5fdc01bcc737dc1183746485c8397ab14d3126614",
+    "skip": false
+   },
+   {
+    "path": "tokenizer.json",
+    "size": 12205732,
+    "sha256": "40fb9d7d7795b8bd305aeff39ce9963f3f450915b9553f2938e009be9a1fed60",
+    "skip": false
+   },
+   {
+    "path": "tokenizer_config.json",
+    "size": 49671,
+    "sha256": "2456b0372956cd3e82f17e33372148b115a94970bfd4878ba5e7e60cd3204f74",
+    "skip": false
+   }
+  ],
+  "total_bytes": 4537376133,
+  "pool_mb": 2048
+ },
+ "35b": {
+  "repo": "Edge0/Edge0-35B-A3B-preview",
+  "rev": {
+   "modelscope": "master",
+   "hf": "e21098e7faa916a00f5493795f12f20497d032bb"
+  },
+  "files": [
+   {
+    "path": ".gitattributes",
+    "size": 2238,
+    "sha256": "f53d675d9d2a887034cf45c1403575f03c9878ceaabca7b10e7b6c494e1e194e",
+    "skip": true
+   },
+   {
+    "path": "20260908-223115.jpg",
+    "size": 1068426,
+    "sha256": "d7285e3c6d5f1039a8b5283031e27fd48aa0cebbb9fb57d180903481eeb8f204",
+    "skip": true
+   },
+   {
+    "path": "20260910-105854.mp4",
+    "size": 20328509,
+    "sha256": "b266d0ec149cb752581789df5846c23ca345501051ee9cc86de31415ed4d6a59",
+    "skip": true
+   },
+   {
+    "path": "README.md",
+    "size": 7012,
+    "sha256": "bba904365a53ec7b6981eb72adc42b17492a45fd304ab070869f38d171578833",
+    "skip": true
+   },
+   {
+    "path": "chat_template.jinja",
+    "size": 7764,
+    "sha256": "e84f32a23fdda27689f868aa4a1a5621f41133e51a48d7f3efcbea2839574259",
+    "skip": false
+   },
+   {
+    "path": "config.json",
+    "size": 23591,
+    "sha256": "a822a9e48b0aafbe3144ec37d4fb067e178ed96615ce6e4420b3149893cc5767",
+    "skip": false
+   },
+   {
+    "path": "configuration.json",
+    "size": 58,
+    "sha256": "c1b09db419119513247e9b8b912c4b9897106c9b20c6cada7e107d993c5435eb",
+    "skip": false
+   },
+   {
+    "path": "generation_config.json",
+    "size": 202,
+    "sha256": "e70c136c1b78ddc1fb0905bac8e733a4dc448d4f852a5dd75143fffc70be550e",
+    "skip": false
+   },
+   {
+    "path": "lora_edge0_35b.safetensors",
+    "size": 42413928,
+    "sha256": "a407f6d27bb55e39e6461bd4bce1fd90489e70c7582b68a796fb764bdb441ec1",
+    "skip": false
+   },
+   {
+    "path": "model-00001-of-00004.safetensors",
+    "size": 4395019264,
+    "sha256": "f7fff28fecb8221e53787c6339c938b113bcb94c7bbc60d09a8826eba08ef1ab",
+    "skip": false
+   },
+   {
+    "path": "model-00002-of-00004.safetensors",
+    "size": 5368472749,
+    "sha256": "31dcdb1c49eebdb1505bd14e3cb33f9cf900bd2546b638f2464694ae763a033f",
+    "skip": false
+   },
+   {
+    "path": "model-00003-of-00004.safetensors",
+    "size": 5368324139,
+    "sha256": "3e66de06a1f03dade16a612a368cfce4a4c9caa4efd7d28185454384082cec03",
+    "skip": false
+   },
+   {
+    "path": "model-00004-of-00004.safetensors",
+    "size": 4377211365,
+    "sha256": "a5d0cf03519c26f8b506df6b0ba60526e5c08c8cea22d0c21ce92950e58a5422",
+    "skip": false
+   },
+   {
+    "path": "model.safetensors.index.json",
+    "size": 180182,
+    "sha256": "e359cd95565275b923fed8df019dd8e869d9af5dc4a621b4114ab7a8e6c534d1",
+    "skip": false
+   },
+   {
+    "path": "preprocessor_config.json",
+    "size": 390,
+    "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516",
+    "skip": false
+   },
+   {
+    "path": "prerouter_edge0_35b.safetensors",
+    "size": 138422000,
+    "sha256": "93d07af08e1f5c8d174e941016740cb4d2c143c6b90f33581237679eca9b006f",
+    "skip": true
+   },
+   {
+    "path": "processor_config.json",
+    "size": 1312,
+    "sha256": "51c5f33646392e0bf1ed3e96be02504da47406c21c6cb86925d8908d1a6c4da7",
+    "skip": false
+   },
+   {
+    "path": "tokenizer.json",
+    "size": 19989343,
+    "sha256": "87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4",
+    "skip": false
+   },
+   {
+    "path": "tokenizer_config.json",
+    "size": 1139,
+    "sha256": "e98f1901ac6f0adff67b1d540bfa0c36ac1a0cf59eb72ed78146ef89aafa1182",
+    "skip": false
+   },
+   {
+    "path": "vocab.json",
+    "size": 6722759,
+    "sha256": "ce99b4cb2983d118806ce0a8b777a35b093e2000a503ebde25853284c9dfa003",
+    "skip": false
+   }
+  ],
+  "total_bytes": 19578368185,
+  "pool_mb": 4096
+ }
+}
\ No newline at end of file
diff --git a/windows/app/src-tauri/icons/128x128.png b/windows/app/src-tauri/icons/128x128.png
new file mode 100644
index 0000000..8881e31
Binary files /dev/null and b/windows/app/src-tauri/icons/128x128.png differ
diff --git a/windows/app/src-tauri/icons/128x128@2x.png b/windows/app/src-tauri/icons/128x128@2x.png
new file mode 100644
index 0000000..cede665
Binary files /dev/null and b/windows/app/src-tauri/icons/128x128@2x.png differ
diff --git a/windows/app/src-tauri/icons/32x32.png b/windows/app/src-tauri/icons/32x32.png
new file mode 100644
index 0000000..91b2b70
Binary files /dev/null and b/windows/app/src-tauri/icons/32x32.png differ
diff --git a/windows/app/src-tauri/icons/64x64.png b/windows/app/src-tauri/icons/64x64.png
new file mode 100644
index 0000000..27b33c4
Binary files /dev/null and b/windows/app/src-tauri/icons/64x64.png differ
diff --git a/windows/app/src-tauri/icons/Square107x107Logo.png b/windows/app/src-tauri/icons/Square107x107Logo.png
new file mode 100644
index 0000000..1786101
Binary files /dev/null and b/windows/app/src-tauri/icons/Square107x107Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square142x142Logo.png b/windows/app/src-tauri/icons/Square142x142Logo.png
new file mode 100644
index 0000000..a8d9be0
Binary files /dev/null and b/windows/app/src-tauri/icons/Square142x142Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square150x150Logo.png b/windows/app/src-tauri/icons/Square150x150Logo.png
new file mode 100644
index 0000000..6ce53f6
Binary files /dev/null and b/windows/app/src-tauri/icons/Square150x150Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square284x284Logo.png b/windows/app/src-tauri/icons/Square284x284Logo.png
new file mode 100644
index 0000000..978b900
Binary files /dev/null and b/windows/app/src-tauri/icons/Square284x284Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square30x30Logo.png b/windows/app/src-tauri/icons/Square30x30Logo.png
new file mode 100644
index 0000000..8196a25
Binary files /dev/null and b/windows/app/src-tauri/icons/Square30x30Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square310x310Logo.png b/windows/app/src-tauri/icons/Square310x310Logo.png
new file mode 100644
index 0000000..b8b47b3
Binary files /dev/null and b/windows/app/src-tauri/icons/Square310x310Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square44x44Logo.png b/windows/app/src-tauri/icons/Square44x44Logo.png
new file mode 100644
index 0000000..519be7f
Binary files /dev/null and b/windows/app/src-tauri/icons/Square44x44Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square71x71Logo.png b/windows/app/src-tauri/icons/Square71x71Logo.png
new file mode 100644
index 0000000..90e28f4
Binary files /dev/null and b/windows/app/src-tauri/icons/Square71x71Logo.png differ
diff --git a/windows/app/src-tauri/icons/Square89x89Logo.png b/windows/app/src-tauri/icons/Square89x89Logo.png
new file mode 100644
index 0000000..744ea79
Binary files /dev/null and b/windows/app/src-tauri/icons/Square89x89Logo.png differ
diff --git a/windows/app/src-tauri/icons/StoreLogo.png b/windows/app/src-tauri/icons/StoreLogo.png
new file mode 100644
index 0000000..0b37792
Binary files /dev/null and b/windows/app/src-tauri/icons/StoreLogo.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-anydpi-v26/ic_launcher.xml b/windows/app/src-tauri/icons/android/mipmap-anydpi-v26/ic_launcher.xml
new file mode 100644
index 0000000..2ffbf24
--- /dev/null
+++ b/windows/app/src-tauri/icons/android/mipmap-anydpi-v26/ic_launcher.xml
@@ -0,0 +1,5 @@
+<?xml version="1.0" encoding="utf-8"?>
+<adaptive-icon xmlns:android="http://schemas.android.com/apk/res/android">
+  <foreground android:drawable="@mipmap/ic_launcher_foreground"/>
+  <background android:drawable="@color/ic_launcher_background"/>
+</adaptive-icon>
\ No newline at end of file
diff --git a/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher.png b/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher.png
new file mode 100644
index 0000000..59f59cc
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher_foreground.png b/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher_foreground.png
new file mode 100644
index 0000000..66abee4
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher_foreground.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher_round.png b/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher_round.png
new file mode 100644
index 0000000..e84186c
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-hdpi/ic_launcher_round.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher.png b/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher.png
new file mode 100644
index 0000000..ca629e8
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher_foreground.png b/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher_foreground.png
new file mode 100644
index 0000000..1c4cdb0
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher_foreground.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher_round.png b/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher_round.png
new file mode 100644
index 0000000..7732498
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-mdpi/ic_launcher_round.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher.png b/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher.png
new file mode 100644
index 0000000..de6a42e
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher_foreground.png b/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher_foreground.png
new file mode 100644
index 0000000..b1047a5
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher_foreground.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher_round.png b/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher_round.png
new file mode 100644
index 0000000..3e76dfe
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xhdpi/ic_launcher_round.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher.png b/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher.png
new file mode 100644
index 0000000..3e663f7
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher_foreground.png b/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher_foreground.png
new file mode 100644
index 0000000..a98192c
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher_foreground.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher_round.png b/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher_round.png
new file mode 100644
index 0000000..40dbcb9
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xxhdpi/ic_launcher_round.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher.png b/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher.png
new file mode 100644
index 0000000..a649d68
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher_foreground.png b/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher_foreground.png
new file mode 100644
index 0000000..a33a803
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher_foreground.png differ
diff --git a/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher_round.png b/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher_round.png
new file mode 100644
index 0000000..3db3bbe
Binary files /dev/null and b/windows/app/src-tauri/icons/android/mipmap-xxxhdpi/ic_launcher_round.png differ
diff --git a/windows/app/src-tauri/icons/android/values/ic_launcher_background.xml b/windows/app/src-tauri/icons/android/values/ic_launcher_background.xml
new file mode 100644
index 0000000..16cdd5a
--- /dev/null
+++ b/windows/app/src-tauri/icons/android/values/ic_launcher_background.xml
@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="utf-8"?>
+<resources>
+  <color name="ic_launcher_background">#FFFFFF</color>
+</resources>
\ No newline at end of file
diff --git a/windows/app/src-tauri/icons/edge0_logo.png b/windows/app/src-tauri/icons/edge0_logo.png
new file mode 100644
index 0000000..712a001
Binary files /dev/null and b/windows/app/src-tauri/icons/edge0_logo.png differ
diff --git a/windows/app/src-tauri/icons/icon.icns b/windows/app/src-tauri/icons/icon.icns
new file mode 100644
index 0000000..d009eda
Binary files /dev/null and b/windows/app/src-tauri/icons/icon.icns differ
diff --git a/windows/app/src-tauri/icons/icon.ico b/windows/app/src-tauri/icons/icon.ico
new file mode 100644
index 0000000..02a0484
Binary files /dev/null and b/windows/app/src-tauri/icons/icon.ico differ
diff --git a/windows/app/src-tauri/icons/icon.png b/windows/app/src-tauri/icons/icon.png
new file mode 100644
index 0000000..94b3731
Binary files /dev/null and b/windows/app/src-tauri/icons/icon.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-20x20@1x.png b/windows/app/src-tauri/icons/ios/AppIcon-20x20@1x.png
new file mode 100644
index 0000000..f793bc5
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-20x20@1x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-20x20@2x-1.png b/windows/app/src-tauri/icons/ios/AppIcon-20x20@2x-1.png
new file mode 100644
index 0000000..204e9d5
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-20x20@2x-1.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-20x20@2x.png b/windows/app/src-tauri/icons/ios/AppIcon-20x20@2x.png
new file mode 100644
index 0000000..204e9d5
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-20x20@2x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-20x20@3x.png b/windows/app/src-tauri/icons/ios/AppIcon-20x20@3x.png
new file mode 100644
index 0000000..11b0321
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-20x20@3x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-29x29@1x.png b/windows/app/src-tauri/icons/ios/AppIcon-29x29@1x.png
new file mode 100644
index 0000000..4509b5f
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-29x29@1x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-29x29@2x-1.png b/windows/app/src-tauri/icons/ios/AppIcon-29x29@2x-1.png
new file mode 100644
index 0000000..d4414d5
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-29x29@2x-1.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-29x29@2x.png b/windows/app/src-tauri/icons/ios/AppIcon-29x29@2x.png
new file mode 100644
index 0000000..d4414d5
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-29x29@2x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-29x29@3x.png b/windows/app/src-tauri/icons/ios/AppIcon-29x29@3x.png
new file mode 100644
index 0000000..bbce047
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-29x29@3x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-40x40@1x.png b/windows/app/src-tauri/icons/ios/AppIcon-40x40@1x.png
new file mode 100644
index 0000000..204e9d5
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-40x40@1x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-40x40@2x-1.png b/windows/app/src-tauri/icons/ios/AppIcon-40x40@2x-1.png
new file mode 100644
index 0000000..6050755
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-40x40@2x-1.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-40x40@2x.png b/windows/app/src-tauri/icons/ios/AppIcon-40x40@2x.png
new file mode 100644
index 0000000..6050755
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-40x40@2x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-40x40@3x.png b/windows/app/src-tauri/icons/ios/AppIcon-40x40@3x.png
new file mode 100644
index 0000000..dc58169
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-40x40@3x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-512@2x.png b/windows/app/src-tauri/icons/ios/AppIcon-512@2x.png
new file mode 100644
index 0000000..244be5e
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-512@2x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-60x60@2x.png b/windows/app/src-tauri/icons/ios/AppIcon-60x60@2x.png
new file mode 100644
index 0000000..dc58169
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-60x60@2x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-60x60@3x.png b/windows/app/src-tauri/icons/ios/AppIcon-60x60@3x.png
new file mode 100644
index 0000000..a1b7dc1
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-60x60@3x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-76x76@1x.png b/windows/app/src-tauri/icons/ios/AppIcon-76x76@1x.png
new file mode 100644
index 0000000..afbad7f
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-76x76@1x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-76x76@2x.png b/windows/app/src-tauri/icons/ios/AppIcon-76x76@2x.png
new file mode 100644
index 0000000..445dbf3
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-76x76@2x.png differ
diff --git a/windows/app/src-tauri/icons/ios/AppIcon-83.5x83.5@2x.png b/windows/app/src-tauri/icons/ios/AppIcon-83.5x83.5@2x.png
new file mode 100644
index 0000000..3357bca
Binary files /dev/null and b/windows/app/src-tauri/icons/ios/AppIcon-83.5x83.5@2x.png differ
diff --git a/windows/app/src-tauri/icons/tray-idle.png b/windows/app/src-tauri/icons/tray-idle.png
new file mode 100644
index 0000000..5f7cb03
Binary files /dev/null and b/windows/app/src-tauri/icons/tray-idle.png differ
diff --git a/windows/app/src-tauri/icons/tray-loading.png b/windows/app/src-tauri/icons/tray-loading.png
new file mode 100644
index 0000000..efd2679
Binary files /dev/null and b/windows/app/src-tauri/icons/tray-loading.png differ
diff --git a/windows/app/src-tauri/icons/tray-resident.png b/windows/app/src-tauri/icons/tray-resident.png
new file mode 100644
index 0000000..5888448
Binary files /dev/null and b/windows/app/src-tauri/icons/tray-resident.png differ
diff --git a/windows/app/src-tauri/icons/tray-throttled.png b/windows/app/src-tauri/icons/tray-throttled.png
new file mode 100644
index 0000000..8405431
Binary files /dev/null and b/windows/app/src-tauri/icons/tray-throttled.png differ
diff --git a/windows/app/src-tauri/src/catalog.rs b/windows/app/src-tauri/src/catalog.rs
new file mode 100644
index 0000000..9c57f1c
--- /dev/null
+++ b/windows/app/src-tauri/src/catalog.rs
@@ -0,0 +1,51 @@
+// catalog.rs — embedded tier table, generated by tools/make_app_catalog.py.
+// The UI must not hardcode model sizes/sha; this is the single source of truth.
+use serde::Deserialize;
+use serde_json::Value;
+use std::collections::HashMap;
+
+pub const CAT_JSON: &str = include_str!("../catalog.json");
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct CatFile {
+    pub path: String,
+    pub size: u64,
+    pub sha256: String,
+    #[serde(default)]
+    pub skip: bool,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct Rev {
+    pub modelscope: String,
+    pub hf: String,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct Tier {
+    pub repo: String,
+    pub rev: Rev,
+    pub files: Vec<CatFile>,
+    pub total_bytes: u64,
+    pub pool_mb: u32,
+}
+
+pub fn catalog() -> HashMap<String, Tier> {
+    serde_json::from_str(CAT_JSON).expect("catalog.json parse failed (rerun tools/make_app_catalog.py)")
+}
+
+pub fn catalog_value() -> Value {
+    serde_json::from_str(CAT_JSON).unwrap()
+}
+
+pub fn base_urls(tier: &Tier) -> Vec<(&'static str, String)> {
+    // Source order for the overseas release: HuggingFace first; hf-mirror and
+    // modelscope are fallbacks (mainland users configure their own proxy — network
+    // routing is not the app's job). All three still participate in the probe-speed
+    // race, so the fastest reachable source wins at runtime.
+    vec![
+        ("hf", format!("https://huggingface.co/{}/resolve/{}", tier.repo, tier.rev.hf)),
+        ("hf-mirror", format!("https://hf-mirror.com/{}/resolve/{}", tier.repo, tier.rev.hf)),
+        ("modelscope", format!("https://modelscope.cn/models/{}/resolve/{}", tier.repo, tier.rev.modelscope)),
+    ]
+}
diff --git a/windows/app/src-tauri/src/convert.rs b/windows/app/src-tauri/src/convert.rs
new file mode 100644
index 0000000..9d0ed19
--- /dev/null
+++ b/windows/app/src-tauri/src/convert.rs
@@ -0,0 +1,77 @@
+// convert.rs — on-device conversion: MLX files -> GGUF + LoRA adapter.
+// Invokes tools/convert_mlx_to_gguf.py (interpreter from EDGE0_PY, default "python")
+// with --dir <files> --out <gguf>. Idempotency and self-checks live inside the
+// converter; this wrapper streams stdout/stderr to logs/convert-<tier>.log and back
+// to the UI, then returns the sha256 of the finished artifact for registration.
+use crate::paths;
+use crate::sink::Sink;
+use serde_json::json;
+use std::io::{BufRead, BufReader, BufWriter, Write};
+use std::process::{Command, Stdio};
+use std::sync::Arc;
+
+pub fn run(sink: &Arc<dyn Sink>, tier: &str) -> Result<String, String> {
+    let repo = paths::repo_root();
+    let script = repo.join("tools").join("convert_mlx_to_gguf.py");
+    if !script.exists() {
+        return Err(format!("conversion script not found {script:?} (set EDGE0_REPO)"));
+    }
+    let py = std::env::var("EDGE0_PY").unwrap_or_else(|_| "python".into());
+    let out = paths::gguf_dir(tier);
+    let log_p = paths::logs_dir().join(format!("convert-{tier}.log"));
+    let logfile = std::fs::File::create(&log_p).map_err(|e| e.to_string())?;
+    let mut cmd = Command::new(&py);
+    crate::no_window(&mut cmd)
+        .arg(&script)
+        .arg("--dir")
+        .arg(paths::files_dir(tier))
+        .arg("--out")
+        .arg(&out)
+        .current_dir(&repo)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped());
+    let mut child = cmd.spawn()
+        .map_err(|e| format!("spawn {py} failed: {e}"))?;
+    let so = child.stdout.take().unwrap();
+    let se = child.stderr.take().unwrap();
+    let s2 = sink.clone();
+    let tc = tier.to_string();
+    std::thread::spawn(move || {
+        let mut w = BufWriter::new(logfile);
+        let a = BufReader::new(so);
+        let b = BufReader::new(se);
+        let lines = a.lines().map(|r| r.unwrap_or_default()).chain(b.lines().map(|r| r.unwrap_or_default()));
+        for l in lines {
+            let _ = writeln!(w, "{l}");
+            let _ = w.flush();
+            s2.emit("convert.progress", &json!({ "tier": tc, "line": l }));
+        }
+    });
+    let st = child.wait().map_err(|e| e.to_string())?;
+    if !st.success() {
+        return Err(format!("converter exited nonzero {}", st.code().unwrap_or(-1)));
+    }
+    let gguf = out.join(format!("edge0-{tier}.gguf"));
+    if !gguf.exists() {
+        return Err(format!("conversion finished but artifact {gguf:?} missing"));
+    }
+    Ok(sha_of(&gguf))
+}
+
+fn sha_of(p: &std::path::Path) -> String {
+    use sha2::Digest;
+    let mut f = match std::fs::File::open(p) {
+        Ok(f) => f,
+        Err(_) => return String::new(),
+    };
+    let mut h = sha2::Sha256::new();
+    let mut buf = vec![0u8; 1 << 22]; // heap-allocated: a 4 MB stack array here caused STATUS_STACK_OVERFLOW
+    loop {
+        let n = std::io::Read::read(&mut f, &mut buf).unwrap_or(0);
+        if n == 0 {
+            break;
+        }
+        h.update(&buf[..n]);
+    }
+    hex::encode(h.finalize())
+}
diff --git a/windows/app/src-tauri/src/doctor.rs b/windows/app/src-tauri/src/doctor.rs
new file mode 100644
index 0000000..811ff15
--- /dev/null
+++ b/windows/app/src-tauri/src/doctor.rs
@@ -0,0 +1,163 @@
+// doctor.rs — environment self-check. Every check reports measured facts (the UI
+// renders them as-is and never estimates). Disk/memory attributes are read through
+// plain kernel32 externs (windows-sys bindings for these tripped link issues here).
+use crate::{catalog, engine, paths};
+use serde_json::{json, Value};
+
+extern "system" {
+    fn GetDiskFreeSpaceExW(
+        dir: *const u16,
+        free_available_to_caller: *mut u64,
+        total_bytes: *mut u64,
+        total_free_bytes: *mut u64,
+    ) -> i32;
+    fn GetFileAttributesW(path: *const u16) -> u32;
+}
+
+const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x0000_0400;
+const INVALID_FILE_ATTRIBUTES: u32 = u32::MAX;
+
+fn wstr(s: &str) -> Vec<u16> {
+    s.encode_utf16().chain(std::iter::once(0)).collect()
+}
+
+pub fn disk_free_gb(path: &std::path::Path) -> Option<u64> {
+    let p = wstr(&path.to_string_lossy());
+    let mut free: u64 = 0;
+    let mut total: u64 = 0;
+    let mut total_free: u64 = 0;
+    let ok = unsafe { GetDiskFreeSpaceExW(p.as_ptr(), &mut free, &mut total, &mut total_free) };
+    if ok != 0 { Some(free / 1_000_000_000) } else { None }
+}
+
+/// Reparse point (junction/symlink) detection — the safety line of model_delete:
+/// seeded links are never physically deleted.
+pub fn is_reparse(path: &std::path::Path) -> bool {
+    let p = wstr(&path.to_string_lossy());
+    let attr = unsafe { GetFileAttributesW(p.as_ptr()) };
+    attr != INVALID_FILE_ATTRIBUTES && (attr & FILE_ATTRIBUTE_REPARSE_POINT) != 0
+}
+
+fn check(id: &str, verdict: &str, detail: impl Into<String>, code: Option<&str>, next: Option<&str>) -> Value {
+    json!({ "id": id, "verdict": verdict, "detail": detail.into(), "code": code, "next": next })
+}
+
+/// Main entry (forwarded by doctor_run in lib.rs): takes the engine state so the
+/// pool-telemetry check reads the live process, not a snapshot.
+pub fn run(state: &engine::EngineState) -> Value {
+    let mut checks: Vec<Value> = Vec::new();
+
+    // home: exists and writable (write + delete one probe file to be sure)
+    let home = paths::home();
+    let probe = home.join("state").join(".doctor-probe");
+    let writable = std::fs::create_dir_all(probe.parent().unwrap()).is_ok()
+        && std::fs::write(&probe, b"ok").is_ok()
+        && { let _ = std::fs::remove_file(&probe); true };
+    checks.push(if writable {
+        check("home", "pass", format!("{} is writable", home.display()), None, None)
+    } else {
+        check("home", "fail", format!("{} is not writable", home.display()), Some("E-IO"), None)
+    });
+
+    // free disk (thresholds from the runtime contract: download+convert workspace for 35B ~19.6 GB + 8B ~4.5 GB)
+    match disk_free_gb(&home) {
+        Some(g) if g >= 25 => checks.push(check("disk", "pass", format!("{g} GB free"), None, None)),
+        Some(g) if g >= 10 => checks.push(check("disk", "warn",
+            format!("{g} GB free: fits the 8b tier only (35b needs >= 25 GB)"), Some("E-DL-DISK"), Some("free disk space or download the 8b tier only"))),
+        Some(g) => checks.push(check("disk", "fail", format!("{g} GB free: not enough for any tier download"), Some("E-DL-DISK"), Some("free up disk space"))),
+        None => checks.push(check("disk", "warn", "cannot read disk free space", None, None)),
+    }
+
+    // physical RAM (measured via GlobalMemoryStatusEx)
+    let mem = engine::phys_mem_gb();
+    checks.push(if mem >= 14 {
+        check("memory", "pass", format!("{mem} GB: both tiers loadable (P-16G class)"), None, None)
+    } else if mem >= 8 {
+        check("memory", "warn", format!("{mem} GB: 8b tier recommended only"), Some("E-MEM-LOAD"), Some("pick 8b on the Models page"))
+    } else {
+        check("memory", "fail", format!("{mem} GB: below the 8b floor"), Some("E-MEM-LOAD"), None)
+    });
+
+    // engine binary presence + version probe
+    let exe = paths::bin_dir().join("llama-server.exe");
+    checks.push(if exe.exists() {
+        match engine::server_version() {
+            Some(v) => check("engine", "pass", v, None, None),
+            None => check("engine", "warn", format!("{} exists but version probe failed", exe.display()), None, None),
+        }
+    } else {
+        check("engine", "fail", format!("missing {}", exe.display()), Some("E-ENGINE-MISSING"), Some("set EDGE0_BIN_DIR"))
+    });
+
+    // per-tier presence (all catalog tiers): not installed = warn pointing at the
+    // Models page; registered + files present = pass
+    let models = read_models_json();
+    for tier in catalog::catalog().keys() {
+        let gguf = paths::gguf_dir(tier).join(format!("edge0-{tier}.gguf"));
+        let lora = paths::files_dir(tier).join(format!("lora_edge0_{tier}-gguf.gguf"));
+        if models.get(tier).is_some() && gguf.exists() && lora.exists() {
+            checks.push(check(&format!("model-{tier}"), "pass",
+                format!("installed (pool {} MB from table), files present", engine::pool_for(tier)), None, None));
+        } else if models.get(tier).is_some() {
+            checks.push(check(&format!("model-{tier}"), "fail",
+                "registered but files missing (record does not match disk)", Some("E-MODEL-INVALID"), Some("delete this tier and download again")));
+        } else {
+            checks.push(check(&format!("model-{tier}"), "warn", "not installed", None, Some("download from the Models page")));
+        }
+    }
+
+    // pool telemetry: engine running but no POOL2 init line in the log = warn (surfaced honestly)
+    let st = engine::status(state);
+    if st["running"].as_bool().unwrap_or(false) {
+        let log = st["log"].as_str().unwrap_or("");
+        match engine::pool_telemetry(log) {
+            Some(line) => checks.push(check("pool-telemetry", "pass", line, None, None)),
+            None => checks.push(check("pool-telemetry", "warn",
+                "engine is running but no POOL2 init telemetry line in log = pool inactive", Some("E-POOL-OFF"), None)),
+        }
+    }
+
+    let has_fail = checks.iter().any(|c| c["verdict"] == "fail");
+    let has_warn = checks.iter().any(|c| c["verdict"] == "warn");
+    json!({ "overall": if has_fail || has_warn { "warn" } else { "pass" }, "checks": checks })
+}
+
+pub fn read_models_json() -> Value {
+    std::fs::read_to_string(paths::state_dir().join("models.json"))
+        .ok()
+        .and_then(|s| serde_json::from_str(&s).ok())
+        .unwrap_or_else(|| json!({}))
+}
+
+/// Delete a tier (physical directories only): a junction/symlink seed returns
+/// E-MODEL-PROTECTED — the repo-asset safety line.
+pub fn delete(tier: &str, resident: Option<&str>) -> Result<Value, String> {
+    if resident == Some(tier) {
+        return Err(format!("E-MODEL-BUSY {tier} is resident (evict it before deleting)"));
+    }
+    let dirs = [paths::files_dir(tier), paths::gguf_dir(tier)];
+    for d in &dirs {
+        if d.exists() && is_reparse(d) {
+            return Err(format!(
+                "E-MODEL-PROTECTED {tier} is a junction seed link — the shell unlinks, never deletes targets; remove the junction in a terminal to reclaim"));
+        }
+    }
+    for d in &dirs {
+        if d.exists() {
+            std::fs::remove_dir_all(d).map_err(|e| format!("E-IO {d:?} {e}"))?;
+        }
+    }
+    // deregister + sweep leftover state files
+    let p = paths::state_dir().join("models.json");
+    if let Ok(s) = std::fs::read_to_string(&p) {
+        if let Ok(mut m) = serde_json::from_str::<Value>(&s) {
+            m.as_object_mut().map(|o| o.remove(tier));
+            let tmp = paths::state_dir().join("models.json.tmp");
+            let _ = std::fs::write(&tmp, serde_json::to_string_pretty(&m).unwrap());
+            let _ = std::fs::rename(&tmp, &p);
+        }
+    }
+    let _ = std::fs::remove_file(paths::state_dir().join(format!("verified-{tier}.json")));
+    let _ = std::fs::remove_file(paths::state_dir().join("downloads").join(format!("{tier}.json")));
+    Ok(json!({ "deleted": tier }))
+}
diff --git a/windows/app/src-tauri/src/engine.rs b/windows/app/src-tauri/src/engine.rs
new file mode 100644
index 0000000..ed95452
--- /dev/null
+++ b/windows/app/src-tauri/src/engine.rs
@@ -0,0 +1,257 @@
+// engine.rs — supervisor for the llama-server child process (the daemon's duties
+// folded into the shell). The spawn arg table is the process form of the runtime
+// contract: -ngl 99 -cmoe --ctx-size 8192 --flash-attn on --pool-mb <tier-clamped>
+// --no-webui --port <random free>. Readiness is decided by polling /health, not by
+// a fixed sleep. A crash is surfaced and not auto-restarted within the session.
+use crate::catalog;
+use crate::paths;
+use crate::sink::Sink;
+use serde_json::{json, Value};
+use std::net::TcpListener;
+use std::process::{Child, Command, Stdio};
+use std::sync::{Arc, Mutex, OnceLock};
+
+#[derive(Debug)]
+pub struct Engine {
+    pub child: Child,
+    pub port: u16,
+    pub tier: String,
+    pub pool_mb: u32,
+    pub started_at: u64, // UNIX seconds (source of truth for the service panel's uptime)
+    pub log_path: String,
+    pub job: bool, // true = attached to a KILL_ON_JOB_CLOSE job; false = only the ExitRequested fallback (recorded honestly)
+}
+
+pub type EngineState = Mutex<Option<Engine>>;
+
+// —— direct kernel32 read (no windows-sys here: this specific binding tripped link/
+//    resolution issues before, and a plain extern declaration is all we need) ——
+#[repr(C)]
+struct MemoryStatusEx {
+    dw_length: u32,
+    memory_load: u32,
+    total_phys: u64,
+    avail_phys: u64,
+    total_pagefile: u64,
+    avail_pagefile: u64,
+    total_virtual: u64,
+    avail_virtual: u64,
+    avail_ext_virtual: u64,
+}
+
+extern "system" {
+    fn GlobalMemoryStatusEx(buf: *mut MemoryStatusEx) -> i32;
+}
+
+/// Physical RAM in GB: env override first (tests pin it), else a real
+/// GlobalMemoryStatusEx read; on read failure fall back conservatively to 48 (the
+/// reference machine — clamping stays safe without this value).
+pub fn phys_mem_gb() -> u64 {
+    if let Ok(v) = std::env::var("EDGE0_PHYS_MEM_GB") {
+        if let Ok(n) = v.parse::<u64>() {
+            return n;
+        }
+    }
+    unsafe {
+        let mut m: MemoryStatusEx = std::mem::zeroed();
+        m.dw_length = std::mem::size_of::<MemoryStatusEx>() as u32;
+        if GlobalMemoryStatusEx(&mut m) != 0 && m.total_phys > 0 {
+            return (m.total_phys / 1_000_000_000).max(1);
+        }
+    }
+    48
+}
+
+/// Pool size = catalog value clamped by physical RAM (safe band ~2-4 GB; no XL tier).
+pub fn pool_for(tier: &str) -> u32 {
+    let base = catalog::catalog().get(tier).map(|t| t.pool_mb).unwrap_or(2048);
+    let gb = phys_mem_gb();
+    let cap = match () {
+        _ if gb < 10 => 1024,
+        _ if gb < 16 => 2048,
+        _ => base,
+    };
+    base.min(cap).max(512)
+}
+
+// Job Object safety net: llama-server is assigned to a job created with
+// KILL_ON_JOB_CLOSE, so whether the shell exits cleanly, crashes, or is force-killed
+// (taskkill /f), the OS closes the process handle -> closes the job -> kills the
+// child. No orphans. This uses the official windows-sys JobObjects bindings (a
+// standard kernel32 import; the struct layout also rules out the hand-rolled-FFI
+// ERROR_BAD_LENGTH class of bugs).
+use std::os::windows::io::AsRawHandle;
+use windows_sys::Win32::Foundation::CloseHandle;
+use windows_sys::Win32::System::JobObjects::{
+    AssignProcessToJobObject, CreateJobObjectW, JobObjectExtendedLimitInformation,
+    SetInformationJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+};
+
+struct JobHandle(Option<*mut core::ffi::c_void>); // windows-sys 0.59: HANDLE = *mut c_void
+// The job object is thread-agnostic, so sharing the raw handle across threads is
+// safe; we intentionally never close it during process lifetime — the OS reclaims
+// handles at exit, and that very close is what kills the child.
+unsafe impl Sync for JobHandle {}
+unsafe impl Send for JobHandle {}
+
+static KILL_JOB: OnceLock<JobHandle> = OnceLock::new();
+
+fn kill_job() -> Option<*mut core::ffi::c_void> {
+    KILL_JOB.get_or_init(|| unsafe {
+        let h = CreateJobObjectW(std::ptr::null(), std::ptr::null());
+        if h.is_null() {
+            return JobHandle(None);
+        }
+        let mut li: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
+        li.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+        let ok = SetInformationJobObject(
+            h,
+            JobObjectExtendedLimitInformation,
+            &li as *const _ as *const core::ffi::c_void,
+            std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
+        );
+        if ok == 0 {
+            CloseHandle(h);
+            return JobHandle(None);
+        }
+        JobHandle(Some(h))
+    })
+    .0
+}
+
+/// Attach the child to the kill-on-close job (best effort: assign fails if the
+/// process already belongs to another job; the caller records false honestly, and
+/// the normal path keeps the ExitRequested→stop fallback as double insurance).
+pub fn assign_kill_job(child: &Child) -> bool {
+    match kill_job() {
+        Some(job) => unsafe { AssignProcessToJobObject(job, child.as_raw_handle()) != 0 },
+        None => false,
+    }
+}
+
+fn free_port() -> Result<u16, String> {
+    TcpListener::bind("127.0.0.1:0").map(|l| l.local_addr().map(|a| a.port()).unwrap_or(0)).map_err(|e| e.to_string())
+}
+
+pub fn start(sink: &Arc<dyn Sink>, state: &EngineState, tier: &str) -> Result<Value, String> {
+    stop(state);
+    let gguf = paths::gguf_dir(tier).join(format!("edge0-{tier}.gguf"));
+    let adapter = paths::files_dir(tier).join(format!("lora_edge0_{tier}-gguf.gguf"));
+    if !gguf.exists() {
+        return Err(format!("E-MODEL-MISSING not converted yet (run the download first): {gguf:?}"));
+    }
+    let bin = paths::bin_dir();
+    let exe = bin.join("llama-server.exe");
+    if !exe.exists() {
+        return Err(format!("E-ENGINE-MISSING {exe:?} (set EDGE0_BIN_DIR)"));
+    }
+    let port = free_port()?;
+    let pool = pool_for(tier);
+    let log_p = paths::logs_dir().join(format!("engine-{tier}-{}.log", std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0)));
+    let logfile = std::fs::File::create(&log_p).map_err(|e| e.to_string())?;
+    let log2 = logfile.try_clone().map_err(|e| e.to_string())?;
+    let mut cmd = Command::new(&exe);
+    crate::no_window(&mut cmd)
+        .arg("-m").arg(&gguf)
+        .arg("--lora").arg(&adapter)
+        .args(["-ngl", "99", "-cmoe", "--ctx-size", "8192", "--flash-attn", "on", "--no-webui"])
+        .arg("--pool-mb").arg(pool.to_string())
+        .arg("--port").arg(port.to_string())
+        .current_dir(&bin) // resolve llama/ggml DLLs from the engine's own directory
+        .stdout(Stdio::from(logfile))
+        .stderr(Stdio::from(log2));
+    let child = cmd.spawn()
+        .map_err(|e| format!("spawn {exe:?} failed: {e}"))?;
+    let job = assign_kill_job(&child); // attach to the job before recording state (failure never blocks; normal path has double insurance)
+    let e = Engine {
+        child,
+        port,
+        tier: tier.into(),
+        pool_mb: pool,
+        started_at: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0),
+        log_path: log_p.to_string_lossy().to_string(),
+        job,
+    };
+    let base = format!("http://127.0.0.1:{}", e.port);
+    *state.lock().unwrap() = Some(e);
+
+    // Readiness = the /health endpoint actually answering (700 ms tick, 300 s budget)
+    let t0 = std::time::Instant::now();
+    let client = reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(2)).build().unwrap();
+    loop {
+        if client.get(format!("{base}/health")).send().map(|r| r.status().is_success()).unwrap_or(false) {
+            break;
+        }
+        let dead = state.lock().unwrap().as_mut().map(|en| en.child.try_wait().ok().flatten().is_some()).unwrap_or(true);
+        if dead {
+            *state.lock().unwrap() = None;
+            return Err(format!("E-ENGINE-DIED see {log_p:?}"));
+        }
+        if t0.elapsed().as_secs() > 300 {
+            return Err("E-ENGINE-READY-TIMEOUT".into());
+        }
+        std::thread::sleep(std::time::Duration::from_millis(700));
+    }
+    // Pool telemetry line from the engine log (absent line = pool not active; surfaced, not hidden)
+    let pool_line = std::fs::read_to_string(&log_p).ok().and_then(|s| {
+        s.lines().rev().find(|l| l.contains("POOL2 init")).map(|l| l.to_string())
+    });
+    let v = json!({ "running": true, "tier": tier, "base_url": base, "bound": "127.0.0.1",
+                    "port": port, "pool_mb": pool, "uptime_s": t0.elapsed().as_secs(),
+                    "log": log_p.to_string_lossy(), "version": server_version(), "job": job,
+                    "phys_mem_gb": phys_mem_gb(), "pool_telemetry": pool_line });
+    sink.emit("engine.status", &v);
+    Ok(v)
+}
+
+pub fn stop(state: &EngineState) {
+    if let Some(mut e) = state.lock().unwrap().take() {
+        let _ = e.child.kill();
+        let _ = e.child.wait();
+    }
+}
+
+pub fn status(state: &EngineState) -> Value {
+    let mut ts = state.lock().unwrap();
+    match ts.as_mut() {
+        Some(e) => {
+            let dead = e.child.try_wait().ok().flatten().is_some();
+            let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0);
+            json!({ "running": !dead, "tier": e.tier,
+                    "base_url": format!("http://127.0.0.1:{}", e.port),
+                    "bound": "127.0.0.1", "port": e.port, "pool_mb": e.pool_mb,
+                    "pid": e.child.id(), "uptime_s": now.saturating_sub(e.started_at),
+                    "log": e.log_path, "version": server_version(), "job": e.job,
+                    "phys_mem_gb": phys_mem_gb() })
+        }
+        None => json!({ "running": false }),
+    }
+}
+
+/// Engine binary version: the actual `llama-server --version` line, cached in a
+/// OnceLock; the UI shows the measured value and falls back to null rather than guessing.
+pub fn server_version() -> Option<String> {
+    static VER: OnceLock<Option<String>> = OnceLock::new();
+    VER.get_or_init(|| {
+        let exe = paths::bin_dir().join("llama-server.exe");
+        let mut cmd = Command::new(&exe);
+        crate::no_window(&mut cmd)
+            .arg("--version")
+            .current_dir(paths::bin_dir())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::null());
+        cmd.spawn()
+            .and_then(|c| c.wait_with_output())
+            .ok()
+            .map(|o| String::from_utf8_lossy(&o.stdout).lines().find(|l| l.starts_with("version:")).map(|l| l.trim().to_string()).unwrap_or_default())
+            .filter(|s| !s.is_empty())
+    }).clone()
+}
+
+/// Read the POOL2 init telemetry line from a log (used by doctor; a missing line =
+/// pool-not-active red flag, surfaced honestly).
+pub fn pool_telemetry(log_path: &str) -> Option<String> {
+    std::fs::read_to_string(log_path).ok().and_then(|s| {
+        s.lines().rev().find(|l| l.contains("POOL2 init")).map(|l| l.trim().to_string())
+    })
+}
diff --git a/windows/app/src-tauri/src/lib.rs b/windows/app/src-tauri/src/lib.rs
new file mode 100644
index 0000000..89cda2d
--- /dev/null
+++ b/windows/app/src-tauri/src/lib.rs
@@ -0,0 +1,163 @@
+// lib.rs — Tauri 2 shell assembly: managed state + command surface + window
+// (created in code at 1100×720, clamped to the primary monitor's work area).
+pub mod catalog;
+pub mod convert;
+pub mod doctor;
+pub mod engine;
+pub mod paths;
+pub mod pull;
+pub mod sink;
+
+use serde_json::{json, Value};
+use std::sync::{Arc, Mutex};
+use tauri::{AppHandle, Manager, WebviewUrl, WebviewWindowBuilder};
+
+/// CREATE_NO_WINDOW = 0x0800_0000. A GUI shell has no console, so every child we
+/// spawn (llama-server / python) would otherwise flash its own black console window.
+/// All shell std Command spawns route through here, so the fix lives in one place.
+pub fn no_window(cmd: &mut std::process::Command) -> &mut std::process::Command {
+    #[cfg(windows)]
+    {
+        use std::os::windows::process::CommandExt;
+        cmd.creation_flags(0x0800_0000)
+    }
+    #[cfg(not(windows))]
+    { cmd }
+}
+
+pub struct AppState {
+    pub tasks: pull::Tasks,
+    pub cancels: pull::Cancels,
+    pub engine: engine::EngineState,
+}
+
+#[tauri::command]
+fn catalog_get() -> Value {
+    catalog::catalog_value()
+}
+
+#[tauri::command]
+fn app_paths() -> Value {
+    json!({ "home": paths::home().to_string_lossy(),
+            "repo": paths::repo_root().to_string_lossy(),
+            "bin_dir": paths::bin_dir().to_string_lossy() })
+}
+
+#[tauri::command]
+fn models_installed() -> Value {
+    std::fs::read_to_string(paths::state_dir().join("models.json"))
+        .ok()
+        .and_then(|s| serde_json::from_str(&s).ok())
+        .unwrap_or_else(|| json!({}))
+}
+
+#[tauri::command]
+fn download_start(app: AppHandle, st: tauri::State<AppState>, tier: String, source: Option<String>) -> Result<Value, String> {
+    pull::start(app, &st.tasks, &st.cancels, &tier, source)
+}
+
+#[tauri::command]
+fn download_status(st: tauri::State<AppState>, tier: String) -> Value {
+    pull::status(&st.tasks, &tier)
+}
+
+#[tauri::command]
+fn download_cancel(st: tauri::State<AppState>, tier: String) {
+    pull::cancel_task(&st.cancels, &tier);
+}
+
+#[tauri::command]
+async fn model_load(app: AppHandle, tier: String) -> Result<Value, String> {
+    // engine::start polls /health with reqwest::blocking + sleep (up to 300 s); a
+    // sync command would run that on the main thread and freeze the UI. Offload to a
+    // dedicated blocking thread (same convention as pull::start's worker); the
+    // promise still resolves on readiness, so the frontend's loading semantics are
+    // unchanged. EngineState isn't Arc, so the closure borrows it via AppHandle.state().
+    tauri::async_runtime::spawn_blocking(move || {
+        let st = app.state::<AppState>();
+        let s: Arc<dyn sink::Sink> = Arc::new(sink::TauriSink(app.clone()));
+        engine::start(&s, &st.engine, &tier)
+    })
+    .await
+    .map_err(|e| format!("E-LOAD-JOIN {e}"))?
+}
+
+#[tauri::command]
+fn model_unload(st: tauri::State<AppState>) -> Value {
+    engine::stop(&st.engine);
+    json!({ "stopped": true })
+}
+
+#[tauri::command]
+fn engine_status(st: tauri::State<AppState>) -> Value {
+    engine::status(&st.engine)
+}
+
+#[tauri::command]
+async fn doctor_run(app: AppHandle) -> Result<Value, String> {
+    // Offloaded like model_load: disk-scan / recompute commands never run on the main thread.
+    tauri::async_runtime::spawn_blocking(move || {
+        let st = app.state::<AppState>();
+        doctor::run(&st.engine)
+    })
+    .await
+    .map_err(|e| format!("E-DOCTOR-JOIN {e}"))
+}
+
+#[tauri::command]
+async fn model_delete(app: AppHandle, tier: String) -> Result<Value, String> {
+    // Multi-GB deletion (seconds under HDD / AV scanning) stays off the main thread;
+    // the resident check (engine::status) happens inside the closure too.
+    tauri::async_runtime::spawn_blocking(move || {
+        let st = app.state::<AppState>();
+        let s = engine::status(&st.engine);
+        let resident = if s["running"].as_bool().unwrap_or(false) { s["tier"].as_str().map(|t| t.to_string()) } else { None };
+        doctor::delete(&tier, resident.as_deref())
+    })
+    .await
+    .map_err(|e| format!("E-DELETE-JOIN {e}"))?
+}
+
+#[cfg_attr(mobile, tauri::mobile_entry_point)]
+pub fn run() {
+    tauri::Builder::default()
+        .manage(AppState {
+            tasks: Arc::new(Mutex::new(Default::default())),
+            cancels: Arc::new(Mutex::new(Default::default())),
+            engine: Mutex::new(None),
+        })
+        .setup(|app| {
+            // Base 1100×720, clamped to the primary monitor's work area (a fixed 720
+            // on a 200%-DPI display pushes the sidebar footer under the taskbar) + centered.
+            let (mut w, mut h): (f64, f64) = (1100.0, 720.0);
+            if let Some(m) = app.primary_monitor()? {
+                let wa = m.work_area();
+                let sf = m.scale_factor();
+                w = w.min(wa.size.width as f64 / sf - 40.0);
+                h = h.min(wa.size.height as f64 / sf - 40.0);
+            }
+            let _win = WebviewWindowBuilder::new(app, "main", WebviewUrl::default())
+                .title("edge0")
+                .inner_size(w.max(480.0), h.max(420.0))
+                .center()
+                .build()?;
+            Ok(())
+        })
+        .invoke_handler(tauri::generate_handler![
+            catalog_get, app_paths, models_installed,
+            download_start, download_status, download_cancel,
+            model_load, model_unload, engine_status,
+            doctor_run, model_delete
+        ])
+        .build(tauri::generate_context!())
+        .expect("edge0 shell failed to start")
+        .run(|app, event| {
+            if let tauri::RunEvent::ExitRequested { .. } = event {
+                // Reclaim the engine child on shell exit (no leaked llama-server orphans;
+                // the Job Object is the backstop for crashes/force-kill, this is the clean path)
+                if let Some(st) = app.try_state::<AppState>() {
+                    engine::stop(&st.engine);
+                }
+            }
+        });
+}
diff --git a/windows/app/src-tauri/src/main.rs b/windows/app/src-tauri/src/main.rs
new file mode 100644
index 0000000..2376b6f
--- /dev/null
+++ b/windows/app/src-tauri/src/main.rs
@@ -0,0 +1,6 @@
+// Prevents additional console window on Windows in release, DO NOT REMOVE!!
+#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
+
+fn main() {
+    edge0_app_lib::run()
+}
diff --git a/windows/app/src-tauri/src/paths.rs b/windows/app/src-tauri/src/paths.rs
new file mode 100644
index 0000000..dc39eaf
--- /dev/null
+++ b/windows/app/src-tauri/src/paths.rs
@@ -0,0 +1,89 @@
+// paths.rs — the ~/.edge0 directory contract. Every path here is overridable via
+// env (EDGE0_HOME etc.) so a test can point the whole app at a throwaway home dir;
+// this is the isolation discipline that makes headless end-to-end tests repeatable.
+use std::path::PathBuf;
+
+pub fn home() -> PathBuf {
+    if let Ok(h) = std::env::var("EDGE0_HOME") {
+        return PathBuf::from(h);
+    }
+    dirs_like_profile().join(".edge0")
+}
+
+fn dirs_like_profile() -> PathBuf {
+    std::env::var("USERPROFILE")
+        .or_else(|_| std::env::var("HOME"))
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| PathBuf::from("."))
+}
+
+// Layout leaf names honor the converter contract: convert_mlx_to_gguf.py derives the
+// tier from the basename and expects --out to be a "<dir>-gguf" sibling, matching the
+// dev repo's models/edge0-8b{,-gguf} shape.
+pub fn models_dir(tier: &str) -> PathBuf {
+    home().join("models").join(format!("edge0-{tier}"))
+}
+pub fn files_dir(tier: &str) -> PathBuf {
+    models_dir(tier) // MLX sources land here (basename=edge0-<tier>, the converter's --dir input)
+}
+pub fn gguf_dir(tier: &str) -> PathBuf {
+    home().join("models").join(format!("edge0-{tier}-gguf"))
+}
+pub fn tmp_dir() -> PathBuf {
+    home().join("tmp")
+}
+pub fn state_dir() -> PathBuf {
+    home().join("state")
+}
+pub fn logs_dir() -> PathBuf {
+    home().join("logs")
+}
+
+pub fn ensure_dirs(tier: Option<&str>) -> std::io::Result<()> {
+    for d in [home(), tmp_dir(), state_dir(), logs_dir()] {
+        std::fs::create_dir_all(&d)?;
+    }
+    if let Some(t) = tier {
+        std::fs::create_dir_all(files_dir(t))?;
+        std::fs::create_dir_all(gguf_dir(t))?;
+    }
+    Ok(())
+}
+
+/// Repo root (dev layout: where the convert script and tools/ live); override with
+/// EDGE0_REPO. Default is derived at compile time from the crate location (this file
+/// lives at <repo>/app/src-tauri/src), so it tracks any clone without a hardcoded path.
+pub fn repo_root() -> PathBuf {
+    std::env::var("EDGE0_REPO")
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| {
+            PathBuf::from(env!("CARGO_MANIFEST_DIR")) // <repo>/app/src-tauri
+                .join("../..")                        // -> <repo> (windows)
+        })
+}
+
+/// Engine binary directory; override with EDGE0_BIN_DIR. Default is the monorepo depot
+/// build (`<repo>/../wt/win/build-vk/bin/Release`, i.e. edge0/wt/... relative to this
+/// windows/ subproject). For a sparse/bootstrap build, set EDGE0_BIN_DIR to the engine
+/// location printed by scripts/vendor-build.ps1.
+pub fn bin_dir() -> PathBuf {
+    std::env::var("EDGE0_BIN_DIR")
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| repo_root().join("../wt/win/build-vk/bin/Release"))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn repo_root_derivation_tracks_the_checkout() {
+        // No env override in unit-test runs: the derived root must actually contain
+        // the converter script (works under any clone path/name; fails only if the
+        // manifest-relative assumption breaks).
+        std::env::remove_var("EDGE0_REPO");
+        let rr = repo_root();
+        assert!(rr.join("tools").join("convert_mlx_to_gguf.py").exists(),
+                "repo_root() = {rr:?} does not look like the windows repo");
+    }
+}
diff --git a/windows/app/src-tauri/src/pull.rs b/windows/app/src-tauri/src/pull.rs
new file mode 100644
index 0000000..354ad57
--- /dev/null
+++ b/windows/app/src-tauri/src/pull.rs
@@ -0,0 +1,357 @@
+// pull.rs — direct-mode downloader. Resume trusts only the .part file's own length;
+// if a source answers a Range request with 200, we reject it (E-DL-SRC) and rotate to
+// the next source. Every file is sha256-checked against the embedded manifest before an
+// atomic rename into place. On completion the chain runs convert → register
+// (phase: probe → download → convert → ready | error | cancelled).
+use crate::catalog::{self, Tier};
+use crate::convert;
+use crate::paths;
+use crate::sink::{Sink, TauriSink};
+use reqwest::blocking::Client;
+use serde_json::{json, Value};
+use sha2::{Digest, Sha256};
+use std::io::{Read, Write};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex};
+use tauri::AppHandle;
+
+pub type Tasks = Arc<Mutex<std::collections::HashMap<String, TaskView>>>;
+pub type Cancels = Arc<Mutex<std::collections::HashMap<String, Arc<AtomicBool>>>>;
+
+#[derive(Debug, Clone)]
+pub struct FileProg {
+    pub path: String,
+    pub done: u64,
+    pub total: u64,
+    pub state: String, // pending|active|done
+}
+
+#[derive(Debug, Clone)]
+pub struct TaskView {
+    pub tier: String,
+    pub phase: String,
+    pub bytes_done: u64,
+    pub bytes_total: u64,
+    pub source: String,
+    pub error: Option<String>,
+    pub files: Vec<FileProg>,
+}
+
+impl TaskView {
+    pub fn to_json(&self) -> Value {
+        json!({
+            "tier": self.tier, "phase": self.phase,
+            "bytes_done": self.bytes_done, "bytes_total": self.bytes_total,
+            "source": self.source, "error": self.error,
+            "files": self.files.iter().map(|f| json!({"path": f.path, "done": f.done, "total": f.total, "state": f.state})).collect::<Vec<_>>(),
+        })
+    }
+}
+
+fn persist(tier: &str, v: &Value) {
+    let p = paths::state_dir().join("downloads");
+    let _ = std::fs::create_dir_all(&p);
+    let f = p.join(format!("{tier}.json"));
+    let tmp = p.join(format!("{tier}.json.tmp"));
+    if let Ok(s) = serde_json::to_string_pretty(v) {
+        let _ = std::fs::write(&tmp, s);
+        let _ = std::fs::rename(&tmp, &f);
+    }
+}
+
+fn set_phase(tasks: &Tasks, sink: &dyn Sink, tier: &str, phase: &str, err: Option<String>) {
+    let v = {
+        let mut ts = tasks.lock().unwrap();
+        let tv = ts.get_mut(tier).unwrap();
+        tv.phase = phase.into();
+        tv.error = err;
+        tv.to_json()
+    };
+    persist(tier, &v);
+    sink.emit("download.progress", &v);
+}
+
+fn verified_map(tier: &str) -> std::collections::HashMap<String, String> {
+    std::fs::read_to_string(paths::state_dir().join(format!("verified-{tier}.json")))
+        .ok()
+        .and_then(|s| serde_json::from_str(&s).ok())
+        .unwrap_or_default()
+}
+
+fn record_verified(tier: &str, path: &str, sha: &str) {
+    let p = paths::state_dir().join(format!("verified-{tier}.json"));
+    let mut m: std::collections::HashMap<String, String> =
+        std::fs::read_to_string(&p).ok().and_then(|s| serde_json::from_str(&s).ok()).unwrap_or_default();
+    m.insert(path.into(), sha.into());
+    if let Ok(s) = serde_json::to_string_pretty(&m) {
+        let _ = std::fs::write(&p, s);
+    }
+}
+
+/// Probe sources by timing a GET of config.json; return candidates ordered fastest-first.
+fn probe_ordered(client: &Client, tr: &Tier) -> Vec<(&'static str, String)> {
+    let cands = catalog::base_urls(tr);
+    let mut times: Vec<std::time::Duration> = Vec::new();
+    for (_, base) in &cands {
+        let t0 = std::time::Instant::now();
+        let ok = client
+            .get(format!("{base}/config.json"))
+            .timeout(std::time::Duration::from_secs(15))
+            .send()
+            .map(|r| r.status().is_success())
+            .unwrap_or(false);
+        times.push(if ok { t0.elapsed() } else { std::time::Duration::from_secs(9e6 as u64) });
+    }
+    let mut idx: Vec<usize> = (0..cands.len()).collect();
+    idx.sort_by_key(|i| times[*i]);
+    idx.into_iter().map(|i| cands[i].clone()).collect()
+}
+
+/// One file: Range resume + sha256 verification + atomic rename on the same volume.
+pub fn pull_one(client: &Client, base: &str, tier: &str, fc: &catalog::CatFile,
+            cancel: &AtomicBool, on_bytes: &dyn Fn(u64)) -> Result<(), String> {
+    if cancel.load(Ordering::Relaxed) {
+        return Err("E-CANCEL".into());
+    }
+    let key: String = fc
+        .path
+        .chars()
+        .map(|c| if c.is_ascii_alphanumeric() || c == '.' || c == '-' { c } else { '_' })
+        .collect();
+    let part = paths::tmp_dir().join(format!("{tier}.{key}.part"));
+    let final_p = paths::files_dir(tier).join(&fc.path);
+    if final_p.exists() && final_p.metadata().map(|m| m.len() == fc.size).unwrap_or(false) {
+        return Ok(());
+    }
+    let mut from: u64 = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0);
+    if from > fc.size {
+        // dirty resume tail beyond the manifest size: truncate to the true length
+        let f = std::fs::OpenOptions::new().write(true).open(&part).map_err(|e| e.to_string())?;
+        f.set_len(fc.size).map_err(|e| e.to_string())?;
+        from = fc.size;
+    }
+    let mut req = client
+        .get(format!("{base}/{}", fc.path))
+        .timeout(std::time::Duration::from_secs(3600));
+    if from > 0 {
+        req = req.header("Range", format!("bytes={from}-"));
+    }
+    let mut resp = req.send().map_err(|e| format!("E-DL-NET {e}"))?;
+    if from > 0 && resp.status() == reqwest::StatusCode::OK {
+        return Err("E-DL-SRC".into()); // source ignores Range semantics: refuse the full body, let the caller rotate
+    }
+    if !resp.status().is_success() {
+        return Err(format!("E-DL-NET http {}", resp.status()));
+    }
+    if from > 0 && resp.status() == reqwest::StatusCode::PARTIAL_CONTENT {
+        // Content-Range start must equal our local prefix, else discard the .part and refetch
+        if let Some(cr) = resp.headers().get("Content-Range").and_then(|v| v.to_str().ok()).map(|s| s.to_string()) {
+            if !cr.starts_with(&format!("bytes {from}-")) {
+                let _ = std::fs::remove_file(&part);
+                return Err("E-DL-SRC".into());
+            }
+        }
+    }
+    let mut file = std::fs::OpenOptions::new()
+        .create(true)
+        .append(true)
+        .open(&part)
+        .map_err(|e| e.to_string())?;
+    let mut buf = vec![0u8; 1 << 20]; // heap buffer (thread stacks default to 2 MB — large buffers must leave the stack)
+    loop {
+        if cancel.load(Ordering::Relaxed) {
+            return Err("E-CANCEL".into());
+        }
+        let n = resp.read(&mut buf).map_err(|e| format!("E-DL-NET stream {e}"))?;
+        if n == 0 {
+            break;
+        }
+        file.write_all(&buf[..n]).map_err(|e| e.to_string())?;
+        on_bytes(n as u64);
+    }
+    file.flush().ok();
+    let size = std::fs::metadata(&part).map_err(|e| e.to_string())?.len();
+    if size != fc.size {
+        let _ = std::fs::remove_file(&part);
+        return Err(format!("E-DL-SIZE {size}/{}", fc.size));
+    }
+    let mut f = std::fs::File::open(&part).map_err(|e| e.to_string())?;
+    let mut h = Sha256::new();
+    let mut acc = 0u64;
+    while acc < size {
+        let n = f.read(&mut buf).map_err(|e| e.to_string())?;
+        if n == 0 {
+            break;
+        }
+        h.update(&buf[..n]);
+        acc += n as u64;
+    }
+    let got = hex::encode(h.finalize());
+    if got != fc.sha256 {
+        let _ = std::fs::remove_file(&part);
+        return Err(format!("E-DL-SHA {}", fc.path));
+    }
+    std::fs::rename(&part, &final_p).map_err(|e| e.to_string())?; // atomic rename on the same volume
+    record_verified(tier, &fc.path, &fc.sha256);
+    Ok(())
+}
+
+pub fn start(app: AppHandle, tasks: &Tasks, cancels: &Cancels,
+             tier: &str, source_hint: Option<String>) -> Result<Value, String> {
+    let cat = catalog::catalog();
+    let tr = cat.get(tier).ok_or_else(|| format!("no such tier: {tier}"))?.clone();
+    {
+        let mut ts = tasks.lock().unwrap();
+        if let Some(v) = ts.get(tier) {
+            if matches!(v.phase.as_str(), "probe" | "download" | "convert") {
+                return Ok(json!({ "already_running": true }));
+            }
+        }
+        let verif = verified_map(tier);
+        let mut pre_done = 0u64;
+        let files = tr
+            .files
+            .iter()
+            .filter(|f| !f.skip)
+            .map(|f| {
+                let done = verif.contains_key(&f.path) || paths::files_dir(tier).join(&f.path).metadata().map(|m| m.len() == f.size).unwrap_or(false);
+                if done {
+                    pre_done += f.size;
+                }
+                FileProg { path: f.path.clone(), done: if done { f.size } else { 0 }, total: f.size, state: if done { "done".into() } else { "pending".into() } }
+            })
+            .collect();
+        ts.insert(tier.to_string(), TaskView {
+            tier: tier.into(), phase: "probe".into(), bytes_done: pre_done,
+            bytes_total: tr.total_bytes, source: String::new(), error: None, files,
+        });
+    }
+    let cancel = Arc::new(AtomicBool::new(false));
+    cancels.lock().unwrap().insert(tier.to_string(), cancel.clone());
+    let sink: Arc<dyn Sink> = Arc::new(TauriSink(app.clone()));
+    let (tasks_c, canc_c, tier_s) = (tasks.clone(), cancels.clone(), tier.to_string());
+    std::thread::Builder::new()
+        .name(format!("pull-{tier}"))
+        .spawn(move || {
+            run(sink, tasks_c, canc_c, tier_s, tr, cancel, source_hint);
+        })
+        .map_err(|e| e.to_string())
+    .map(|_| json!({ "started": true }))
+}
+
+/// The full download + convert + register chain, callable headlessly (tests use this
+/// directly). An Err return means the error phase was already set on the task view.
+pub fn run(sink: Arc<dyn Sink>, tasks: Tasks, cancels: Cancels, tier: String, tr: Tier,
+           cancel: Arc<AtomicBool>, source_hint: Option<String>) {
+    let _ = paths::ensure_dirs(Some(&tier));
+    let client = Client::builder().user_agent("edge0-app/0.1").build().unwrap();
+    let ordered: Vec<(&'static str, String)> = match source_hint {
+        Some(h) => {
+            let mut c = catalog::base_urls(&tr);
+            if let Some(i) = c.iter().position(|(n, _)| *n == h.as_str()) {
+                c.rotate_left(i);
+            }
+            c
+        }
+        None => probe_ordered(&client, &tr),
+    };
+    {
+        let mut ts = tasks.lock().unwrap();
+        ts.get_mut(&tier).unwrap().source = ordered[0].0.to_string();
+    }
+    set_phase(&tasks, &*sink,&tier, "download", None);
+
+    let todo: Vec<catalog::CatFile> = tr.files.iter().filter(|f| !f.skip).cloned().collect();
+    let mut src_i = 0usize;
+    for fc in todo {
+        loop {
+            let (name, base) = ordered[src_i].clone();
+            let (sink_c, tasks_c, tier_c, path_c) = (sink.clone(), tasks.clone(), tier.clone(), fc.path.clone());
+            let r = pull_one(&client, &base, &tier, &fc, &cancel, &|n| {
+                let v = {
+                    let mut ts = tasks_c.lock().unwrap();
+                    let tv = ts.get_mut(&tier_c).unwrap();
+                    tv.bytes_done = (tv.bytes_done + n).min(tv.bytes_total);
+                    if let Some(f) = tv.files.iter_mut().find(|x| x.path == path_c) {
+                        f.done = (f.done + n).min(f.total);
+                        if f.state == "pending" {
+                            f.state = "active".into();
+                        }
+                    }
+                    tv.to_json()
+                };
+                sink_c.emit("download.progress", &v);
+            });
+            match r {
+                Ok(()) => {
+                    {
+                        let mut ts = tasks.lock().unwrap();
+                        let tv = ts.get_mut(&tier).unwrap();
+                        if let Some(f) = tv.files.iter_mut().find(|x| x.path == fc.path) {
+                            f.state = "done".into();
+                            f.done = f.total;
+                            tv.bytes_done = tv.files.iter().map(|x| x.done).sum();
+                        }
+                    }
+                    persist(&tier, &tasks.lock().unwrap().get(&tier).unwrap().to_json());
+                    break;
+                }
+                Err(e) if e == "E-CANCEL" => {
+                    set_phase(&tasks, &*sink,&tier, "cancelled", None);
+                    cancels.lock().unwrap().remove(&tier);
+                    return;
+                }
+                Err(e) if (e.starts_with("E-DL-SRC") || e.starts_with("E-DL-NET")) && src_i + 1 < ordered.len() => {
+                    src_i += 1; // rotate to the next source and retry this file
+                    {
+                        let mut ts = tasks.lock().unwrap();
+                        ts.get_mut(&tier).unwrap().source = ordered[src_i].0.to_string();
+                    }
+                    continue;
+                }
+                Err(e) => {
+                    set_phase(&tasks, &*sink,&tier, "error", Some(format!("{e} @{} via {name}", fc.path)));
+                    cancels.lock().unwrap().remove(&tier);
+                    return;
+                }
+            }
+        }
+    }
+
+    set_phase(&tasks, &*sink,&tier, "convert", None);
+    match convert::run(&sink, &tier) {
+        Ok(r3_sha) => {
+            register_ready(&tier, &tr, r3_sha);
+            set_phase(&tasks, &*sink,&tier, "ready", None);
+        }
+        Err(e) => set_phase(&tasks, &*sink, &tier, "error", Some(format!("E-CONVERT {e}"))),
+    }
+    cancels.lock().unwrap().remove(&tier);
+}
+
+fn register_ready(tier: &str, tr: &Tier, r3_sha: String) {
+    let p = paths::state_dir().join("models.json");
+    let mut m: Value = std::fs::read_to_string(&p)
+        .ok()
+        .and_then(|s| serde_json::from_str(&s).ok())
+        .unwrap_or_else(|| json!({}));
+    m[tier] = json!({
+        "installed_at": std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0),
+        "repo": tr.repo, "rev_hf": tr.rev.hf,
+        "gguf_dir": paths::gguf_dir(tier).to_string_lossy(),
+        "files_dir": paths::files_dir(tier).to_string_lossy(),
+        "r3_sha256": r3_sha,
+        "pool_mb": tr.pool_mb,
+    });
+    let _ = std::fs::write(&p, serde_json::to_string_pretty(&m).unwrap());
+}
+
+pub fn status(tasks: &Tasks, tier: &str) -> Value {
+    tasks.lock().unwrap().get(tier).map(|v| v.to_json()).unwrap_or_else(|| json!({ "tier": tier, "phase": "idle" }))
+}
+
+pub fn cancel_task(cancels: &Cancels, tier: &str) {
+    if let Some(c) = cancels.lock().unwrap().get(tier) {
+        c.store(true, Ordering::Relaxed);
+    }
+}
diff --git a/windows/app/src-tauri/src/sink.rs b/windows/app/src-tauri/src/sink.rs
new file mode 100644
index 0000000..a2711ac
--- /dev/null
+++ b/windows/app/src-tauri/src/sink.rs
@@ -0,0 +1,21 @@
+// sink.rs — event-output abstraction. Running in the shell it is TauriSink(AppHandle);
+// in headless tests it is NoopSink. pull/convert/engine depend on this trait instead of
+// the window handle directly, which is what lets them run without a GUI.
+use serde_json::Value;
+use tauri::{AppHandle, Emitter};
+
+pub trait Sink: Send + Sync {
+    fn emit(&self, event: &str, payload: &Value);
+}
+
+pub struct NoopSink;
+impl Sink for NoopSink {
+    fn emit(&self, _event: &str, _payload: &Value) {}
+}
+
+pub struct TauriSink(pub AppHandle);
+impl Sink for TauriSink {
+    fn emit(&self, event: &str, payload: &Value) {
+        let _ = self.0.emit(event, payload);
+    }
+}
diff --git a/windows/app/src-tauri/tauri.conf.json b/windows/app/src-tauri/tauri.conf.json
new file mode 100644
index 0000000..145edfd
--- /dev/null
+++ b/windows/app/src-tauri/tauri.conf.json
@@ -0,0 +1,23 @@
+{
+  "$schema": "https://schema.tauri.app/config/2",
+  "productName": "edge0",
+  "version": "0.1.0",
+  "identifier": "app.edge0.client-win",
+  "build": {
+    "beforeDevCommand": "npm run dev",
+    "devUrl": "http://localhost:1420",
+    "beforeBuildCommand": "npm run build",
+    "frontendDist": "../dist"
+  },
+  "app": {
+    "windows": [],
+    "security": {
+      "csp": "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data: blob:; font-src 'self' data:; connect-src 'self' ipc: http://ipc.localhost http://127.0.0.1:*"
+    }
+  },
+  "bundle": {
+    "active": true,
+    "targets": ["nsis"],
+    "icon": ["icons/icon.ico", "icons/icon.png"]
+  }
+}
diff --git a/windows/app/src-tauri/tests/doctor.rs b/windows/app/src-tauri/tests/doctor.rs
new file mode 100644
index 0000000..dcfc6d5
--- /dev/null
+++ b/windows/app/src-tauri/tests/doctor.rs
@@ -0,0 +1,82 @@
+// tests/doctor.rs — doctor check contract + model_delete protection lines (isolated home;
+// tests serialized via a process lock to avoid env races).
+use edge0_app_lib::{doctor, paths};
+use serde_json::json;
+use std::sync::Once;
+
+static INIT: Once = Once::new();
+// cargo test runs threads in parallel: both tests share the isolated home/env,
+// so serialize them with a mutex (poison-tolerant).
+static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
+
+fn iso_home() -> std::path::PathBuf {
+    INIT.call_once(|| {
+        let h = std::env::temp_dir().join(format!("edge0-doctor-itest-{}", std::process::id()));
+        let _ = std::fs::remove_dir_all(&h);
+        std::fs::create_dir_all(h.join("state")).unwrap();
+        std::env::set_var("EDGE0_HOME", &h);
+        std::env::set_var("EDGE0_BIN_DIR", &h); // empty dir ⇒ engine=fail anchors the assertion
+        std::env::set_var("EDGE0_PHYS_MEM_GB", "8");
+    });
+    paths::home()
+}
+
+#[test]
+fn doctor_contract_shape() {
+    let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
+    let h = iso_home();
+    std::fs::create_dir_all(h.join("models")).unwrap();
+    // register a fake "installed" tier without files ⇒ the model-8b fail / E-MODEL-INVALID path
+    std::fs::write(h.join("state/models.json"), json!({ "8b": { "pool_mb": 2048 } }).to_string()).unwrap();
+
+    let r = doctor::run(&std::sync::Mutex::new(None));
+    assert_eq!(r["overall"], "warn", "overall=warn whenever any fail/warn exists");
+    let find = |id: &str| r["checks"].as_array().expect("checks is an array").iter().find(|c| c["id"] == id).cloned();
+
+    assert_eq!(find("home").unwrap()["verdict"], "pass", "isolated home is writable");
+    assert!(doctor::disk_free_gb(&h).is_some(), "disk free space reads as a real value");
+    let eng = find("engine").unwrap();
+    assert_eq!(eng["verdict"], "fail");
+    assert_eq!(eng["code"], "E-ENGINE-MISSING");
+    assert_eq!(find("memory").unwrap()["code"], "E-MEM-LOAD", "8 GB ⇒ warn, 8b tier only");
+    let m8 = find("model-8b").unwrap();
+    assert_eq!(m8["verdict"], "fail", "registered but files missing = red INVALID flag");
+    assert_eq!(m8["code"], "E-MODEL-INVALID");
+    assert_eq!(find("model-35b").unwrap()["verdict"], "warn", "not installed = warn, not fail");
+}
+
+#[test]
+fn delete_protection_and_idempotent() {
+    let _g = LOCK.lock().unwrap_or_else(|e| e.into_inner());
+    let h = iso_home();
+    // busy: deleting the resident tier is refused
+    let e = doctor::delete("8b", Some("8b"));
+    assert!(e.unwrap_err().starts_with("E-MODEL-BUSY"));
+    // real directories: physical delete + deregistration
+    let f = paths::files_dir("t-del");
+    let g = paths::gguf_dir("t-del");
+    std::fs::create_dir_all(&f).unwrap();
+    std::fs::create_dir_all(&g).unwrap();
+    std::fs::write(f.join("x.bin"), b"hi").unwrap();
+    let mp = h.join("state/models.json");
+    std::fs::write(&mp, json!({ "t-del": {} }).to_string()).unwrap();
+    doctor::delete("t-del", None).unwrap();
+    assert!(!f.exists() && !g.exists());
+    let m: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&mp).unwrap()).unwrap();
+    assert!(m.get("t-del").is_none(), "tier removed from models.json");
+    // idempotent: deleting an absent tier is Ok
+    assert!(doctor::delete("t-none", None).is_ok());
+    // junction refuses deletion (seed-link protection): create a junction at
+    // files_dir(jtest) ⇒ E-MODEL-PROTECTED
+    let j = paths::files_dir("jtest");
+    let out = std::process::Command::new("cmd")
+        .args(["/c", "mklink", "/J", &j.to_string_lossy(), &paths::home().to_string_lossy()])
+        .output();
+    if out.map(|o| o.status.success()).unwrap_or(false) && j.exists() {
+        let err = doctor::delete("jtest", None).unwrap_err();
+        assert!(err.starts_with("E-MODEL-PROTECTED"), "{err}");
+        assert!(j.exists(), "the junction survives the refused deletion");
+        let _ = std::process::Command::new("cmd").args(["/c", "rmdir", &j.to_string_lossy()]).output();
+    }
+    // if mklink fails (no privilege) this leg skips — the real seed-link path is covered by manual runs.
+}
diff --git a/windows/app/src-tauri/tests/job.rs b/windows/app/src-tauri/tests/job.rs
new file mode 100644
index 0000000..74ee296
--- /dev/null
+++ b/windows/app/src-tauri/tests/job.rs
@@ -0,0 +1,84 @@
+// tests/job.rs — Job Object gate: (1) production assign_kill_job returns true and really
+// attaches; (2) closing the job handle makes the OS kill the member process
+// (KILL_ON_JOB_CLOSE semantics, same windows-sys structs as production); (3) a real
+// llama-server gets attached end-to-end (ignored test, needs the model).
+use edge0_app_lib::engine::assign_kill_job;
+use std::os::windows::io::AsRawHandle;
+use std::process::{Command, Stdio};
+use windows_sys::Win32::Foundation::CloseHandle;
+use windows_sys::Win32::System::JobObjects::{
+    AssignProcessToJobObject, CreateJobObjectW, JobObjectExtendedLimitInformation,
+    SetInformationJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+};
+
+fn spawn_linger() -> std::process::Child {
+    Command::new("ping")
+        .args(["-n", "60", "127.0.0.1"])
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()
+        .expect("ping spawn")
+}
+
+#[test]
+fn assign_kill_job_true() {
+    let mut child = spawn_linger();
+    let ok = assign_kill_job(&child);
+    let _ = child.kill();
+    let _ = child.wait();
+    assert!(ok, "job create + SetInformation(KILL_ON_JOB_CLOSE) + Assign all succeeded");
+}
+
+#[test]
+fn job_close_kills_child() {
+    unsafe {
+        let mut child = spawn_linger();
+        let job = CreateJobObjectW(std::ptr::null(), std::ptr::null());
+        assert!(!job.is_null(), "CreateJobObjectW");
+        let mut li: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
+        li.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+        assert_ne!(
+            SetInformationJobObject(
+                job,
+                JobObjectExtendedLimitInformation,
+                &li as *const _ as *const core::ffi::c_void,
+                std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
+            ),
+            0,
+            "SetInformationJobObject"
+        );
+        assert_ne!(AssignProcessToJobObject(job, child.as_raw_handle() as _), 0, "Assign");
+        assert_ne!(CloseHandle(job), 0, "CloseHandle(job)");
+
+        // closing the job must make the kernel kill the member ping (a 60 s process should die within seconds; 10 s ceiling)
+        let t0 = std::time::Instant::now();
+        loop {
+            if child.try_wait().expect("try_wait").is_some() {
+                break;
+            }
+            assert!(t0.elapsed().as_secs() < 10, "KILL_ON_JOB_CLOSE did not fire (child alive after 10 s)");
+            std::thread::sleep(std::time::Duration::from_millis(100));
+        }
+    }
+}
+
+/// Real end-to-end (needs an 8b model under ~/.edge0 + llama-server at EDGE0_BIN_DIR):
+/// engine::start launches the actual llama-server, asserts job:true (the real process
+/// really attached), checks status echoes it, then stop() brings running=false.
+#[test]
+#[ignore = "needs the real 8b model and llama-server; run manually: cargo test --test job -- --ignored"]
+fn engine_start_attaches_real_llama_to_job() {
+    use edge0_app_lib::{engine, sink::NoopSink, sink::Sink};
+    use std::sync::{Arc, Mutex};
+    let sink: Arc<dyn Sink> = Arc::new(NoopSink);
+    let state: engine::EngineState = Mutex::new(None);
+    let v = engine::start(&sink, &state, "8b").expect("real 8b chain start");
+    assert_eq!(v["running"], serde_json::json!(true));
+    assert_eq!(v["job"], serde_json::json!(true), "the real llama-server must be attached to the kill-on-close job");
+    let st = engine::status(&state);
+    let pid = st["pid"].as_u64().expect("pid present");
+    assert_eq!(st["job"], serde_json::json!(true), "status must echo the job flag too");
+    engine::stop(&state);
+    assert_eq!(engine::status(&state)["running"], serde_json::json!(false));
+    println!("REAL JOB CHAIN GREEN pid={pid}");
+}
diff --git a/windows/app/src-tauri/tests/pull.rs b/windows/app/src-tauri/tests/pull.rs
new file mode 100644
index 0000000..7e1f628
--- /dev/null
+++ b/windows/app/src-tauri/tests/pull.rs
@@ -0,0 +1,181 @@
+// tests/pull.rs — fast gate: a fake HTTP server plays Range tri-state + sha mismatch +
+// dirty-tail truncation. Whole run isolated in a temp EDGE0_HOME (unique tier names
+// per case so tests never tread on each other).
+use edge0_app_lib::catalog::CatFile;
+use edge0_app_lib::paths;
+use edge0_app_lib::pull::pull_one;
+use sha2::{Digest, Sha256};
+use std::io::{Read, Write};
+use std::net::TcpListener;
+use std::sync::atomic::AtomicBool;
+use std::sync::Once;
+
+static INIT: Once = Once::new();
+
+fn test_home() -> std::path::PathBuf {
+    INIT.call_once(|| {
+        let h = std::env::temp_dir().join(format!("edge0-itest-{}", std::process::id()));
+        let _ = std::fs::remove_dir_all(&h);
+        std::env::set_var("EDGE0_HOME", &h);
+    });
+    paths::home()
+}
+
+fn sha(bytes: &[u8]) -> String {
+    hex::encode(Sha256::digest(bytes))
+}
+
+fn pattern(n: usize) -> Vec<u8> {
+    (0..n).map(|i| ((i * 31 + 7) % 251) as u8).collect()
+}
+
+/// Fake source: mode = range (honors Range) | rude (answers Range with 200 full body) | corrupt (shifted content)
+fn serve(content: Vec<u8>, mode: &'static str) -> String {
+    let l = TcpListener::bind("127.0.0.1:0").unwrap();
+    let port = l.local_addr().unwrap().port();
+    std::thread::spawn(move || {
+        for stream in l.incoming() {
+            let mut s = match stream { Ok(s) => s, Err(_) => continue };
+            let mut buf = [0u8; 4096];
+            let mut head = String::new();
+            while let Ok(n) = s.read(&mut buf) {
+                if n == 0 {
+                    break;
+                }
+                head.push_str(&String::from_utf8_lossy(&buf[..n]));
+                if head.contains("\r\n\r\n") {
+                    break;
+                }
+            }
+            let want: Option<u64> = match mode {
+                _ if !head.contains("Range:") => None,
+                _ => head
+                    .find("bytes=")
+                    .and_then(|i| head[i + 6..].split('-').next())
+                    .and_then(|s| s.parse().ok()),
+            };
+            let mut body = content.clone();
+            if mode == "corrupt" {
+                for b in body.iter_mut() {
+                    *b = b.wrapping_add(1);
+                }
+            }
+            let resp_head = match (mode, want) {
+                ("rude", Some(_)) => format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()),
+                (_, Some(f)) if f <= body.len() as u64 => {
+                    let tail = &body[f as usize..];
+                    let end = f + tail.len() as u64 - 1;
+                    format!("HTTP/1.1 206 Partial Content\r\nContent-Range: bytes {f}-{end}/{}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
+                        content.len(), tail.len())
+                }
+                _ => format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()),
+            };
+            let send_body = match (mode, want) {
+                (_, Some(f)) if f <= body.len() as u64 && mode != "rude" => body[f as usize..].to_vec(),
+                _ => body.clone(),
+            };
+            let _ = s.write_all(resp_head.as_bytes());
+            let _ = s.write_all(&send_body);
+            let _ = s.flush();
+        }
+    });
+    format!("http://127.0.0.1:{port}")
+}
+
+fn fc(path: &str, bytes: &[u8]) -> CatFile {
+    CatFile { path: path.into(), size: bytes.len() as u64, sha256: sha(bytes), skip: false }
+}
+
+fn cancel() -> AtomicBool {
+    AtomicBool::new(false)
+}
+
+#[test]
+fn pull_fresh_full() {
+    let _h = test_home();
+    let _ = paths::ensure_dirs(Some("t1"));
+    let data = pattern(5000);
+    let c = reqwest::blocking::Client::new();
+    let base = serve(data.clone(), "range");
+    pull_one(&c, &base, "t1", &fc("a.bin", &data), &cancel(), &|_| {}).unwrap();
+    assert_eq!(std::fs::read(paths::files_dir("t1").join("a.bin")).unwrap(), data);
+}
+
+#[test]
+fn pull_resume_206() {
+    let _h = test_home();
+    let h = test_home();
+    let _ = paths::ensure_dirs(Some("t2"));
+    let data = pattern(10_000);
+    std::fs::write(h.join("tmp/t2.a.bin.part"), &data[..4096]).unwrap();
+    let c = reqwest::blocking::Client::new();
+    let base = serve(data.clone(), "range");
+    let got_n = std::sync::atomic::AtomicU64::new(0);
+    pull_one(&c, &base, "t2", &fc("a.bin", &data), &cancel(), &|n| {
+        got_n.fetch_add(n, std::sync::atomic::Ordering::Relaxed);
+    })
+    .unwrap();
+    assert_eq!(std::fs::read(paths::files_dir("t2").join("a.bin")).unwrap(), data);
+    assert_eq!(got_n.load(std::sync::atomic::Ordering::Relaxed), 10_000 - 4096, "resume must fetch only the missing tail");
+}
+
+#[test]
+fn pull_resume_dirty_tail_bad_sha() {
+    let h = test_home();
+    let _ = paths::ensure_dirs(Some("t3"));
+    let data = pattern(2000);
+    std::fs::write(h.join("tmp/t3.a.bin.part"), vec![9u8; 5000]).unwrap(); // oversized dirty tail, wrong content
+    let c = reqwest::blocking::Client::new();
+    let base = serve(data.clone(), "range");
+    let r = pull_one(&c, &base, "t3", &fc("a.bin", &data), &cancel(), &|_| {});
+    assert!(matches!(&r, Err(e) if e.starts_with("E-DL-SHA") || e.starts_with("E-DL-SRC")), "{r:?}");
+}
+
+#[test]
+fn pull_rude_source_rejected() {
+    let h = test_home();
+    let _ = paths::ensure_dirs(Some("t4"));
+    let data = pattern(3000);
+    std::fs::write(h.join("tmp/t4.a.bin.part"), &data[..1000]).unwrap();
+    let c = reqwest::blocking::Client::new();
+    let base = serve(data.clone(), "rude");
+    let r = pull_one(&c, &base, "t4", &fc("a.bin", &data), &cancel(), &|_| {});
+    assert_eq!(r.unwrap_err(), "E-DL-SRC");
+}
+
+#[test]
+fn pull_corrupt_content_sha_red() {
+    let _h = test_home();
+    let _ = paths::ensure_dirs(Some("t5"));
+    let data = pattern(2048);
+    let c = reqwest::blocking::Client::new();
+    let base = serve(data.clone(), "corrupt");
+    let r = pull_one(&c, &base, "t5", &fc("a.bin", &data), &cancel(), &|_| {});
+    assert!(r.unwrap_err().starts_with("E-DL-SHA"));
+}
+
+#[test]
+fn catalog_contract() {
+    let cat = edge0_app_lib::catalog::catalog();
+    let t8 = cat.get("8b").expect("8b tier in catalog");
+    assert!(t8.files.iter().any(|f| f.skip && f.path.contains("prerouter")));
+    assert!(!t8.files.iter().any(|f| !f.skip && f.path.contains("prerouter")));
+    assert_eq!(t8.total_bytes, t8.files.iter().filter(|f| !f.skip).map(|f| f.size).sum::<u64>());
+    let t35 = cat.get("35b").expect("35b tier in catalog");
+    assert!(t35.files.iter().any(|f| f.skip && (f.path.ends_with(".jpg") || f.path.ends_with(".mp4"))));
+    assert_eq!(t35.rev.hf.len(), 40, "35b HF revision pinned to a full commit sha");
+    // source contract (overseas release decision): hf first, hf-mirror/modelscope as fallbacks; hf url tail = commit sha
+    let urls = edge0_app_lib::catalog::base_urls(t35);
+    let names: Vec<&str> = urls.iter().map(|(n, _)| *n).collect();
+    assert_eq!(names, vec!["hf", "hf-mirror", "modelscope"], "source order = huggingface first");
+    assert!(urls[0].1.starts_with("https://huggingface.co/") && urls[0].1.ends_with(&t35.rev.hf), "hf url carries the commit sha");
+    assert!(urls[2].1.contains(&t35.rev.modelscope), "modelscope url carries its revision");
+}
+
+#[test]
+fn pool_clamp_contract() {
+    for t in ["8b", "35b", "nonexist"] {
+        let p = edge0_app_lib::engine::pool_for(t);
+        assert!((512..=4096).contains(&p), "{t} → {p} outside the safe band (contract forbids XL)");
+    }
+}
diff --git a/windows/app/src-tauri/tests/real8b.rs b/windows/app/src-tauri/tests/real8b.rs
new file mode 100644
index 0000000..d1a06ae
--- /dev/null
+++ b/windows/app/src-tauri/tests/real8b.rs
@@ -0,0 +1,78 @@
+// tests/real8b.rs — real end-to-end pipeline (EDGE0_TEST_TIER=8b|35b: download → convert
+// → engine → chat → assertions). Manual tier: cargo test --test real8b -- --ignored
+// --nocapture. The isolated EDGE0_HOME uses a fixed name and is not wiped at startup —
+// rerunning the same command resumes idempotently. For 35b the POOL2 telemetry line is
+// REQUIRED (phantom defense: never claim pool-backed numbers without the live proof).
+use edge0_app_lib::{catalog, engine, paths, pull, sink::NoopSink};
+use serde_json::{json, Value};
+use std::collections::HashMap;
+use std::sync::atomic::AtomicBool;
+use std::sync::{Arc, Mutex};
+
+#[test]
+#[ignore]
+fn real8b_pipeline() {
+    let tier = std::env::var("EDGE0_TEST_TIER").unwrap_or_else(|_| "8b".into());
+    let h = std::env::temp_dir().join(format!("edge0-real-{}", tier));
+    std::env::set_var("EDGE0_HOME", &h);
+    paths::ensure_dirs(Some(&tier)).unwrap();
+
+    let mut cat = catalog::catalog();
+    let tr = cat.remove(&tier).unwrap_or_else(|| panic!("{tier} tier missing from catalog"));
+    let tasks: pull::Tasks = Arc::new(Mutex::new(HashMap::new()));
+    tasks.lock().unwrap().insert(tier.clone(), pull::TaskView {
+        tier: tier.clone(), phase: "probe".into(), bytes_done: 0, bytes_total: tr.total_bytes,
+        source: String::new(), error: None,
+        files: tr.files.iter().filter(|f| !f.skip)
+            .map(|f| pull::FileProg { path: f.path.clone(), done: 0, total: f.size, state: "pending".into() })
+            .collect(),
+    });
+    let cancels: pull::Cancels = Arc::new(Mutex::new(HashMap::new()));
+    let sink: Arc<dyn edge0_app_lib::sink::Sink> = Arc::new(NoopSink);
+
+    let src = std::env::var("EDGE0_TEST_SOURCE").unwrap_or_else(|_| "modelscope".into());
+    pull::run(sink.clone(), tasks.clone(), cancels, tier.clone(), tr, Arc::new(AtomicBool::new(false)), Some(src));
+
+    let v = tasks.lock().unwrap().get(&tier).unwrap().to_json();
+    println!("pipeline phase={}", v["phase"]);
+    assert_eq!(v["phase"].as_str(), Some("ready"), "full chain should end ready: {v}");
+
+    let gguf = paths::gguf_dir(&tier).join(format!("edge0-{tier}.gguf"));
+    let adapter = paths::files_dir(&tier).join(format!("lora_edge0_{tier}-gguf.gguf"));
+    assert!(gguf.exists() && adapter.exists(), "conversion artifacts present");
+
+    // engine + chat
+    let state: engine::EngineState = Mutex::new(None);
+    let info = engine::start(&sink, &state, &tier).expect("engine start");
+    println!("engine: {info}");
+    if tier == "35b" {
+        let tel = info["pool_telemetry"].as_str().unwrap_or_default().to_string();
+        assert!(tel.contains("tt=120") && tel.contains("resolver=on"),
+            "phantom defense: 35b must show live pool telemetry, never claim pool-backed numbers without it: {tel:?}");
+        assert_eq!(info["pool_mb"].as_u64(), Some(4096), "16 GB-class pool table value = 4096");
+    }
+    let base = info["base_url"].as_str().unwrap().to_string();
+    let client = reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(600)).build().unwrap();
+    // prompt asks for exact markdown artifacts (table + code fence); assertions below
+    // check those tokens. Chinese on purpose: this fixture also exercises the model's native language.
+    let resp: Value = client
+        .post(format!("{base}/v1/chat/completions"))
+        .json(&json!({"messages": [{"role": "user",
+            "content": "请原样输出：①markdown 表格（两列 算式/结果，两行 1+1 与 2+2）②python 代码块 print(42)。除此之外不要任何文字"}],
+            "max_tokens": 400, "temperature": 0.0}))
+        .send().unwrap().json().unwrap();
+    let msg = &resp["choices"][0]["message"];
+    let content = msg["content"].as_str().unwrap_or_default().to_string();
+    let reasoning = msg["reasoning_content"].as_str().unwrap_or_default().to_string();
+    engine::stop(&state); // reclaim the engine before asserting (a panic must not leak the process)
+    println!("resp: content={}B reasoning={}B", content.len(), reasoning.len());
+    let joined = format!("{content}{reasoning}");
+    assert!(joined.contains('|') && joined.contains("```"), "markdown artifacts (table + code fence) missing");
+    println!("reply head: {}", &joined[..joined.len().min(300)]);
+
+    // ledger cross-check
+    let m: Value = serde_json::from_str(&std::fs::read_to_string(paths::state_dir().join("models.json")).unwrap()).unwrap();
+    let sha = m[&tier]["r3_sha256"].as_str().unwrap_or_default().to_string();
+    println!("registered r3 sha256={sha}");
+    assert_eq!(sha.len(), 64);
+}
diff --git a/windows/app/src/api/client.ts b/windows/app/src/api/client.ts
new file mode 100644
index 0000000..27a3f83
--- /dev/null
+++ b/windows/app/src/api/client.ts
@@ -0,0 +1,140 @@
+// api/client.ts — shell IPC wrapper + one-pulse stores. Two-layer architecture:
+// data sources are invoke commands + Tauri events only; the frontend never
+// invents polling and never fabricates estimates.
+import { invoke } from "@tauri-apps/api/core";
+import { listen } from "@tauri-apps/api/event";
+
+export type EngStatus = {
+  running: boolean; tier?: string; base_url?: string; bound?: string; port?: number;
+  pool_mb?: number; uptime_s?: number; log?: string; version?: string | null;
+  phys_mem_gb?: number; pool_telemetry?: string | null; pid?: number; job?: boolean;
+};
+export type FileProg = { path: string; done: number; total: number; state: string };
+export type TaskView = {
+  tier: string; phase: string; bytes_done: number; bytes_total: number;
+  source: string; error: string | null; files?: FileProg[];
+};
+export type CatTier = { repo: string; rev: { modelscope: string; hf: string }; total_bytes: number; pool_mb: number; files: { path: string; size: number; skip?: boolean }[] };
+
+class Sig {
+  private ls = new Set<() => void>();
+  private v = 0;
+  sub = (fn: () => void): (() => void) => {
+    this.ls.add(fn);
+    return () => { this.ls.delete(fn); };
+  };
+  version = () => this.v;
+  emit() {
+    this.v++;
+    for (const f of [...this.ls]) f();
+  }
+}
+
+export const engineStore = {
+  status: null as EngStatus | null,
+  loadingTier: null as string | null, // in-flight model_load marker (one source of truth for the "loading" light)
+  sig: new Sig(),
+};
+
+export const downloadStore = {
+  tasks: {} as Record<string, TaskView>,
+  sig: new Sig(),
+  apply(v: TaskView) {
+    this.tasks = { ...this.tasks, [v.tier]: v };
+    this.sig.emit();
+  },
+};
+
+export const shellStore = {
+  catalog: {} as Record<string, CatTier>,
+  paths: { home: "", repo: "", bin_dir: "" },
+  installed: {} as Record<string, unknown>,
+  sig: new Sig(),
+};
+
+/** Ring buffer of request log rows (service-page data = frontend-measured: model / latency / tokens / status; since app start, cleared on exit — labeled honestly). */
+export type ReqRow = { ts: number; model: string; status: "done" | "cancelled" | "error"; latencyMs: number; tokens: number | null; code: string | null };
+export const reqLog = {
+  rows: [] as ReqRow[],
+  sig: new Sig(),
+  add(r: ReqRow) {
+    this.rows = [r, ...this.rows].slice(0, 200);
+    this.sig.emit();
+  },
+};
+
+let bootOnce: Promise<void> | null = null;
+
+/** Idempotent boot: initial fetch + event wiring (shared by main.tsx and first-screen components; repeat calls share one Promise). */
+export function boot(): Promise<void> {
+  if (!bootOnce) bootOnce = bootNow();
+  return bootOnce;
+}
+
+async function bootNow(): Promise<void> {
+  try {
+    shellStore.catalog = await invoke<Record<string, CatTier>>("catalog_get");
+    shellStore.paths = await invoke("app_paths");
+    await refreshInstalled();
+    engineStore.status = await invoke<EngStatus>("engine_status");
+    for (const tier of Object.keys(shellStore.catalog)) {
+      downloadStore.tasks[tier] = await invoke<TaskView>("download_status", { tier });
+    }
+  } catch { /* shell not ready yet: stores stay null, UI renders its waiting shape */ }
+  void listen<EngStatus>("engine.status", (e) => {
+    engineStore.status = e.payload;
+    engineStore.sig.emit();
+  });
+  void listen<TaskView>("download.progress", (e) => downloadStore.apply(e.payload));
+  shellStore.sig.emit();
+  engineStore.sig.emit();
+  downloadStore.sig.emit();
+}
+
+export async function refreshInstalled(): Promise<void> {
+  shellStore.installed = await invoke<Record<string, unknown>>("models_installed");
+  shellStore.sig.emit();
+}
+
+/** Models-page refresh button: refresh re-pulls the sources of truth; it is not a poll. */
+export async function refreshAll(): Promise<void> {
+  try {
+    await refreshInstalled();
+    engineStore.status = await invoke<EngStatus>("engine_status");
+    for (const tier of Object.keys(shellStore.catalog)) {
+      downloadStore.tasks[tier] = await invoke<TaskView>("download_status", { tier });
+    }
+  } finally {
+    engineStore.sig.emit();
+    downloadStore.sig.emit();
+  }
+}
+
+export async function loadModel(tier: string): Promise<EngStatus> {
+  engineStore.loadingTier = tier;
+  engineStore.sig.emit();
+  try {
+    const s = await invoke<EngStatus>("model_load", { tier });
+    engineStore.status = s;
+    return s;
+  } finally {
+    engineStore.loadingTier = null;
+    engineStore.sig.emit();
+  }
+}
+
+export async function unloadModel(): Promise<void> {
+  await invoke("model_unload");
+  engineStore.status = await invoke<EngStatus>("engine_status");
+  engineStore.sig.emit();
+}
+
+export const deleteModel = (tier: string): Promise<{ deleted: string }> => invoke("model_delete", { tier });
+export const startDownload = (tier: string, source: string | null): Promise<unknown> => invoke("download_start", { tier, source });
+export const cancelDownload = (tier: string): Promise<void> => invoke("download_cancel", { tier });
+export const doctorRun = (): Promise<{ overall: string; checks: DoctorCheck[] }> => invoke("doctor_run");
+export type DoctorCheck = { id: string; verdict: string; detail: string; code?: string | null; next?: string | null };
+
+export function installedTiers(): string[] {
+  return Object.keys(shellStore.installed);
+}
diff --git a/windows/app/src/chat/buckets.ts b/windows/app/src/chat/buckets.ts
new file mode 100644
index 0000000..e0740ee
--- /dev/null
+++ b/windows/app/src/chat/buckets.ts
@@ -0,0 +1,31 @@
+// Chat-history time bucketing: pure function, unit-testable; Today / Yesterday / Previous 7 days / Earlier.
+export interface Bucket<T> {
+  key: "today" | "yesterday" | "week" | "earlier";
+  items: T[];
+}
+
+const DAY_MS = 24 * 60 * 60 * 1000;
+
+function localDayStart(ms: number): number {
+  const d = new Date(ms);
+  d.setHours(0, 0, 0, 0);
+  return d.getTime();
+}
+
+export function bucketThreads<T extends { updatedAt: number }>(
+  threads: T[],
+  nowMs: number,
+): Bucket<T>[] {
+  const today = localDayStart(nowMs);
+  const buckets: Record<Bucket<T>["key"], T[]> = { today: [], yesterday: [], week: [], earlier: [] };
+  for (const th of threads) {
+    const day = localDayStart(th.updatedAt);
+    if (day >= today) buckets.today.push(th);
+    else if (day >= today - DAY_MS) buckets.yesterday.push(th);
+    else if (day >= today - 6 * DAY_MS) buckets.week.push(th);
+    else buckets.earlier.push(th);
+  }
+  return (["today", "yesterday", "week", "earlier"] as const)
+    .filter((k) => buckets[k].length > 0)
+    .map((k) => ({ key: k, items: buckets[k] }));
+}
diff --git a/windows/app/src/chat/modelStatus.ts b/windows/app/src/chat/modelStatus.ts
new file mode 100644
index 0000000..fab2dea
--- /dev/null
+++ b/windows/app/src/chat/modelStatus.ts
@@ -0,0 +1,23 @@
+// chat/modelStatus.ts — status-light derivation. Inputs are shell snapshots only:
+// the installed list + engineStore.status + loadingTier. No timers, no guessing:
+// the light reflects facts alone.
+import type { EngStatus } from "../api/client";
+
+export type ModelLight = "resident" | "loading" | "installed" | "absent";
+
+export function deriveLight(
+  status: EngStatus | null,
+  loadingTier: string | null,
+  installed: string[],
+  tier: string,
+): ModelLight {
+  if (!installed.includes(tier)) return "absent";
+  if (status?.running && status.tier === tier) return "resident";
+  if (loadingTier === tier) return "loading";
+  return "installed";
+}
+
+/** Picker candidates = installed tiers (catalog/download browsing belongs to the models page). */
+export function candidateModels(installed: string[]): string[] {
+  return installed;
+}
diff --git a/windows/app/src/chat/scroll.ts b/windows/app/src/chat/scroll.ts
new file mode 100644
index 0000000..396e498
--- /dev/null
+++ b/windows/app/src/chat/scroll.ts
@@ -0,0 +1,13 @@
+// Follow-bottom predicate: pin to the bottom on streaming increments only while the user
+// is near the bottom (prevents the "scroll lock" bug that drags a scrolled-up reader back).
+export const STICK_BOTTOM_THRESHOLD = 40;
+
+export interface ScrollGeometry {
+  scrollHeight: number;
+  scrollTop: number;
+  clientHeight: number;
+}
+
+export function isAtBottom(el: ScrollGeometry, threshold = STICK_BOTTOM_THRESHOLD): boolean {
+  return el.scrollHeight - el.scrollTop - el.clientHeight < threshold;
+}
diff --git a/windows/app/src/chat/useThreadChat.ts b/windows/app/src/chat/useThreadChat.ts
new file mode 100644
index 0000000..a28e422
--- /dev/null
+++ b/windows/app/src/chat/useThreadChat.ts
@@ -0,0 +1,192 @@
+// chat/useThreadChat.ts — per-thread generation hook. No AI SDK dependency: talks
+// straight to the loopback llama-server OpenAI SSE endpoint. The final chunk's
+// usage+timings are recorded as measured (missing keys render nothing; zero estimates).
+import { useCallback, useRef, useState } from "react";
+import { engineStore, reqLog } from "../api/client";
+import { threadStore, type MsgUsage, type ThreadMsg, type ThreadRow } from "../state/threads";
+import { currentTemperature } from "../lib/generation";
+import { storedSystemPrompt } from "../lib/systemPrompt";
+import { codeOf } from "../lib/t";
+
+export type ThreadChatApi = {
+  messages: ThreadMsg[];
+  busy: boolean;
+  send: (text: string) => Promise<void>;
+  stop: () => void;
+  regenerate: () => Promise<void>;
+  editResend: (msgId: string, text: string) => Promise<void>;
+  wireJson: () => { role: string; content: string }[]; // the exact wire form used for token metering
+};
+
+export function toWire(history: ThreadMsg[]): { role: string; content: string }[] {
+  const out: { role: string; content: string }[] = [];
+  const sp = storedSystemPrompt();
+  if (sp) out.push({ role: "system", content: sp });
+  for (const m of history) {
+    if (m.role === "user") out.push({ role: "user", content: m.content });
+    else if (m.content.trim()) out.push({ role: "assistant", content: m.content });
+  }
+  return out;
+}
+
+const uid = () => `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}`;
+
+export function useThreadChat(thread: ThreadRow, model: string | null): ThreadChatApi {
+  const [busy, setBusy] = useState(false);
+  const abortRef = useRef<AbortController | null>(null);
+
+  const patchLast = useCallback(
+    (patch: Partial<ThreadMsg>) => {
+      threadStore.patchMessages(thread.id, (ms) =>
+        ms.length ? ms.map((m, i) => (i === ms.length - 1 ? { ...m, ...patch } : m)) : ms,
+      );
+    },
+    [thread.id],
+  );
+
+  const run = useCallback(
+    async (history: ThreadMsg[]) => {
+      const base = engineStore.status?.running ? engineStore.status.base_url : undefined;
+      if (!base) {
+        patchLast({ meta: { status: "error", code: "E-NOT-LOADED" } });
+        return;
+      }
+      setBusy(true);
+      const t0 = Date.now();
+      const ac = new AbortController();
+      abortRef.current = ac;
+      let acc = "";
+      let reason = "";
+      let usage: MsgUsage | undefined;
+      let perf: { decode_tps?: number; prefill_tps?: number; prompt_ms?: number } | undefined;
+      try {
+        const resp = await fetch(`${base}/v1/chat/completions`, {
+          method: "POST",
+          signal: ac.signal,
+          headers: { "Content-Type": "application/json" },
+          body: JSON.stringify({
+            model: model ?? "edge0",
+            messages: toWire(history),
+            max_tokens: 1024,
+            temperature: currentTemperature(),
+            stream: true,
+            stream_options: { include_usage: true }, // final-chunk usage comes straight from the engine
+            // enable_thinking:false => answer text streams directly (server-side normalization point)
+            chat_template_kwargs: { enable_thinking: false },
+          }),
+        });
+        if (!resp.ok || !resp.body) throw new Error(`E-ENGINE HTTP ${resp.status}`);
+        const reader = resp.body.getReader();
+        const dec = new TextDecoder();
+        let buf = "";
+        for (;;) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          buf += dec.decode(value, { stream: true });
+          const lines = buf.split("\n");
+          buf = lines.pop() ?? "";
+          for (const line of lines) {
+            const tr = line.trim();
+            if (!tr.startsWith("data:")) continue;
+            const payload = tr.slice(5).trim();
+            if (payload === "[DONE]") continue;
+            try {
+              const d = JSON.parse(payload);
+              const delta = d.choices?.[0]?.delta ?? {};
+              acc += delta.content ?? "";
+              reason += delta.reasoning_content ?? "";
+              if (d.usage) {
+                usage = {
+                  prompt_tokens: d.usage.prompt_tokens ?? 0,
+                  completion_tokens: d.usage.completion_tokens ?? 0,
+                  total_tokens: d.usage.total_tokens ?? 0,
+                };
+              }
+              if (d.timings && typeof d.timings === "object") {
+                const nums: Record<string, number> = {};
+                for (const [k, v] of Object.entries(d.timings)) if (typeof v === "number") nums[k] = v;
+                perf = {
+                  decode_tps: nums.predicted_per_second,
+                  prefill_tps: nums.prompt_per_second,
+                  prompt_ms: nums.prompt_ms,
+                };
+              }
+              patchLast({ content: acc, reasoning: reason, streaming: true });
+            } catch {
+              /* tolerate a split (partial) JSON line; next chunk completes it */
+            }
+          }
+        }
+        patchLast({ content: acc, reasoning: reason, streaming: false, meta: { status: "done", usage, perf } });
+        reqLog.add({ ts: Date.now(), model: model ?? "edge0", status: "done", latencyMs: Date.now() - t0, tokens: usage?.completion_tokens ?? null, code: null });
+      } catch (e) {
+        const aborted = (e as Error).name === "AbortError";
+        patchLast({
+          content: acc, reasoning: reason, streaming: false,
+          meta: aborted
+            ? { status: "aborted", usage, perf }
+            : { status: "error", code: codeOf(String(e)) ?? "E-ENGINE" },
+        });
+        reqLog.add({
+          ts: Date.now(), model: model ?? "edge0",
+          status: aborted ? "cancelled" : "error",
+          latencyMs: Date.now() - t0, tokens: null, code: aborted ? null : codeOf(String(e)),
+        });
+      } finally {
+        setBusy(false);
+        abortRef.current = null;
+      }
+    },
+    [model, patchLast, thread.id],
+  );
+
+  const send = useCallback(
+    async (text: string) => {
+      const history: ThreadMsg[] = [
+        ...threadStore.messages(thread.id),
+        { id: uid(), role: "user", content: text, created_at: Date.now() },
+      ];
+      threadStore.patchMessages(thread.id, (ms) => [
+        ...ms,
+        history[history.length - 1],
+        { id: uid(), role: "assistant", content: "", created_at: Date.now(), streaming: true },
+      ]);
+      await run(history);
+    },
+    [run, thread.id],
+  );
+
+  const regenerate = useCallback(async () => {
+    const msgs = threadStore.messages(thread.id);
+    const lastA = msgs.map((m) => m.role).lastIndexOf("assistant");
+    if (lastA < 0) return;
+    const history = msgs.slice(0, lastA); // drop the previous answer
+    threadStore.patchMessages(thread.id, () => [
+      ...history,
+      { id: uid(), role: "assistant", content: "", created_at: Date.now(), streaming: true },
+    ]);
+    await run(history);
+  }, [run, thread.id]);
+
+  const editResend = useCallback(
+    async (msgId: string, text: string) => {
+      const msgs = threadStore.messages(thread.id);
+      const idx = msgs.findIndex((m) => m.id === msgId);
+      if (idx < 0) return;
+      const edited: ThreadMsg = { id: uid(), role: "user", content: text, created_at: Date.now() };
+      const history = [...msgs.slice(0, idx), edited];
+      threadStore.patchMessages(thread.id, () => [
+        ...history,
+        { id: uid(), role: "assistant", content: "", created_at: Date.now(), streaming: true },
+      ]);
+      await run(history);
+    },
+    [run, thread.id],
+  );
+
+  const stop = useCallback(() => abortRef.current?.abort(), []);
+
+  const wireJson = useCallback(() => toWire(threadStore.messages(thread.id)), [thread.id]);
+
+  return { messages: threadStore.messages(thread.id), busy, send, stop, regenerate, editResend, wireJson };
+}
diff --git a/windows/app/src/components/AppShell.tsx b/windows/app/src/components/AppShell.tsx
new file mode 100644
index 0000000..0e1e34b
--- /dev/null
+++ b/windows/app/src/components/AppShell.tsx
@@ -0,0 +1,180 @@
+// AppShell — sidebar + full-width main area (Atomic-Chat style three-segment layout).
+// Hash routing; state comes from the single-pulse stores in api/client.
+// autoLoad intent chain: download ready -> issue one load -> on success create a thread and enter chat.
+import { BookOpen, Menu, Plus, Server, Settings, Sparkles, X } from "lucide-react";
+import { useEffect, useState, useSyncExternalStore, type ReactNode } from "react";
+
+import {
+  downloadStore,
+  engineStore,
+  loadModel,
+  refreshInstalled,
+  shellStore,
+} from "../api/client";
+import { autoLoad } from "../lib/autoLoad";
+import { t } from "../lib/t";
+import { bucketThreads } from "../chat/buckets";
+import { threadStore } from "../state/threads";
+import { SidebarThreadItem } from "./SidebarThreadItem";
+
+const go = (path: string) => {
+  window.location.hash = `#${path}`;
+};
+
+export function AppShell({ view, children }: { view: string; children: ReactNode }) {
+  const [sidebarOpen, setSidebarOpen] = useState(false);
+  const threadsVer = useSyncExternalStore(threadStore.subscribe, () => threadStore.version());
+  const engVer = useSyncExternalStore(engineStore.sig.sub, engineStore.sig.version);
+  const dlVer = useSyncExternalStore(downloadStore.sig.sub, downloadStore.sig.version);
+  const shellVer = useSyncExternalStore(shellStore.sig.sub, shellStore.sig.version);
+  void threadsVer; void engVer; void dlVer; void shellVer; // pulse-driven rerenders; data read straight from stores
+
+  // autoLoad one-shot intent chain: ready -> explicitly issue one load -> on success create thread and enter chat.
+  useEffect(() => {
+    const check = () => {
+      const cur = autoLoad.current;
+      if (!cur) return;
+      if (cur.phase === "download") {
+        const tv = downloadStore.tasks[cur.tier];
+        if (tv && (tv.phase === "error" || tv.phase === "cancelled")) autoLoad.cancel();
+        else if (tv?.phase === "ready") {
+          const tier = autoLoad.claimCompleted(cur.tier);
+          if (tier) {
+            void (async () => {
+              await refreshInstalled();
+              try {
+                await loadModel(tier);
+              } catch {
+                autoLoad.cancel();
+              }
+            })();
+          }
+        }
+      } else {
+        const st = engineStore.status;
+        if (st?.running && st.tier === cur.tier && autoLoad.consume(cur.tier)) {
+          const row = threadStore.create(t("chat.newThread"), cur.tier);
+          threadStore.setActive(row.id);
+          go("/"); // entering a session means returning to the chat page
+        } else if (!st?.running && engineStore.loadingTier === null) {
+          autoLoad.cancel(); // load failed (engine not running and not loading)
+        }
+      }
+    };
+    const off1 = downloadStore.sig.sub(check);
+    const off2 = engineStore.sig.sub(check);
+    const off3 = autoLoad.subscribe(check);
+    check();
+    return () => {
+      off1();
+      off2();
+      off3();
+    };
+  }, []);
+
+  const threads = threadStore.threads();
+  const activeId = threadStore.activeId;
+  const buckets = bucketThreads(threads, Date.now());
+  const st = engineStore.status;
+  const running = !!st?.running;
+  const serviceLine = running
+    ? `:${st?.port} · ${t("sidebar.resident")} ${st?.tier}`
+    : st
+      ? t("sidebar.idle")
+      : t("sidebar.noService");
+
+  function newChat() {
+    const row = threadStore.create(t("chat.newThread"), running ? (st?.tier ?? null) : null);
+    threadStore.setActive(row.id);
+    go("/"); // "New chat" clicked from a non-chat page must land back on the chat page
+    setSidebarOpen(false);
+  }
+
+  const navCls = (active: boolean) =>
+    "flex items-center gap-2 rounded-lg px-3 py-2 text-sm " +
+    (active ? "bg-surface2 text-fg" : "text-muted hover:bg-surface2 hover:text-fg");
+
+  return (
+    <div className="e0-app-shell text-fg">
+      {sidebarOpen && (
+        <button type="button" aria-label={t("nav.closeSidebar")} className="e0-sidebar-backdrop" onClick={() => setSidebarOpen(false)} />
+      )}
+      <aside data-testid="app-sidebar" className={"e0-sidebar " + (sidebarOpen ? "e0-sidebar-open" : "")} data-open={sidebarOpen ? "true" : "false"}>
+        <div className="flex items-center justify-between px-4 pb-2 pt-4">
+          <div className="flex items-center gap-2">
+            <span className="grid size-8 place-items-center rounded-xl bg-fg text-surface shadow-sm">
+              <Sparkles size={16} aria-hidden />
+            </span>
+            <div>
+              <span className="block text-sm font-bold tracking-tight">{t("app.name")}</span>
+              <span className="block text-[10px] font-medium uppercase tracking-[0.16em] text-muted">{t("app.localLabel")}</span>
+            </div>
+          </div>
+          <button type="button" aria-label={t("nav.closeSidebar")} className="e0-icon-btn md:hidden" onClick={() => setSidebarOpen(false)}>
+            <X size={16} aria-hidden />
+          </button>
+        </div>
+        <div className="px-3 pb-3">
+          <button type="button" onClick={newChat} className="e0-btn e0-btn-primary w-full">
+            <Plus size={16} aria-hidden />
+            {t("sidebar.newChat")}
+          </button>
+        </div>
+        <nav aria-label={t("sidebar.navigation")} className="space-y-1 px-2 pb-2">
+          <a href="#/models" aria-label={t("nav.models")} data-testid="nav-models" className={navCls(view === "/models")} onClick={() => setSidebarOpen(false)}>
+            <BookOpen size={16} aria-hidden />
+            {t("nav.models")}
+          </a>
+          <a href="#/service" aria-label={t("nav.service")} data-testid="nav-service" className={navCls(view === "/service")} onClick={() => setSidebarOpen(false)}>
+            <Server size={16} aria-hidden />
+            {t("nav.service")}
+          </a>
+        </nav>
+        <nav aria-label={t("sidebar.history")} className="min-h-0 flex-1 overflow-y-auto px-2">
+          <h2 className="px-3 pb-2 pt-2 text-[10px] font-bold uppercase tracking-[0.16em] text-muted">{t("sidebar.history")}</h2>
+          {buckets.map((b) => (
+            <section key={b.key} className="mb-2">
+              <h2 className="px-2 pb-1 pt-2 text-[11px] font-semibold uppercase tracking-wide text-muted">
+                {t(`sidebar.bucket.${b.key}`)}
+              </h2>
+              <ul className="space-y-0.5">
+                {b.items.map((th) => (
+                  <SidebarThreadItem
+                    key={th.id}
+                    thread={th}
+                    active={th.id === activeId}
+                    onSelect={() => {
+                      threadStore.setActive(th.id);
+                      setSidebarOpen(false);
+                      go("/");
+                    }}
+                    onRename={(title) => threadStore.rename(th.id, title)}
+                    onDelete={() => threadStore.remove(th.id)}
+                  />
+                ))}
+              </ul>
+            </section>
+          ))}
+          {threads.length === 0 && <p className="px-3 py-3 text-xs leading-relaxed text-muted">{t("chat.noThreads")}</p>}
+        </nav>
+        <div className="e0-divider mx-3" />
+        <div className="flex items-center gap-2 px-3 py-3 text-xs text-muted">
+          <span aria-label={serviceLine} className={"inline-block size-2 shrink-0 rounded-full " + (running ? "bg-ok" : st ? "bg-warn" : "bg-muted")} />
+          <span className="min-w-0 flex-1 truncate" data-testid="service-summary">{serviceLine}</span>
+          <a href="#/settings" aria-label={t("nav.settings")} className="e0-icon-btn ml-auto" onClick={() => setSidebarOpen(false)}>
+            <Settings size={16} aria-hidden />
+          </a>
+        </div>
+      </aside>
+      <main data-testid="app-main" className="e0-main-frame">
+        <div className="e0-mobile-toggle absolute left-4 top-4 z-10">
+          <button type="button" aria-label={t("nav.openSidebar")} title={t("nav.openSidebar")}
+            className="e0-icon-btn border border-line bg-surface shadow-sm" onClick={() => setSidebarOpen(true)}>
+            <Menu size={17} aria-hidden />
+          </button>
+        </div>
+        <div className="h-full min-h-0 overflow-y-auto">{children}</div>
+      </main>
+    </div>
+  );
+}
diff --git a/windows/app/src/components/ShellControls.tsx b/windows/app/src/components/ShellControls.tsx
new file mode 100644
index 0000000..64b8213
--- /dev/null
+++ b/windows/app/src/components/ShellControls.tsx
@@ -0,0 +1,51 @@
+// ShellControls — theme + font-size control groups (no radix: native data-state keeps the same shape).
+import { useEffect, useState } from "react";
+import { t } from "../lib/t";
+import {
+  applyFont,
+  applyTheme,
+  storedFont,
+  storedTheme,
+  type FontScale,
+  type Theme,
+} from "../lib/theme";
+
+export function ShellControls() {
+  const [theme, setTheme] = useState<Theme>(storedTheme());
+  const [font, setFont] = useState<FontScale>(storedFont());
+
+  useEffect(() => {
+    applyTheme(theme);
+  }, [theme]);
+  useEffect(() => {
+    applyFont(font);
+  }, [font]);
+  // When following the system preference, re-resolve on system changes.
+  useEffect(() => {
+    const mq = window.matchMedia("(prefers-color-scheme: dark)");
+    const onChange = () => theme === "system" && applyTheme("system");
+    mq.addEventListener("change", onChange);
+    return () => mq.removeEventListener("change", onChange);
+  }, [theme]);
+
+  return (
+    <div className="flex flex-wrap items-center gap-3 border-b border-line pb-3" data-testid="shell-controls">
+      <div className="e0-toggle-group" role="group" aria-label={t("theme.system")}>
+        {(["light", "dark", "system"] as const).map((v) => (
+          <button key={v} type="button" data-state={v === theme ? "on" : "off"} className="e0-toggle"
+            aria-pressed={v === theme} onClick={() => setTheme(v)}>
+            {t(`theme.${v}`)}
+          </button>
+        ))}
+      </div>
+      <div className="e0-toggle-group" role="group" aria-label={t("font.md")}>
+        {(["sm", "md", "lg"] as const).map((v) => (
+          <button key={v} type="button" data-state={v === font ? "on" : "off"} className="e0-toggle"
+            aria-pressed={v === font} onClick={() => setFont(v)}>
+            {t(`font.${v}`)}
+          </button>
+        ))}
+      </div>
+    </div>
+  );
+}
diff --git a/windows/app/src/components/SidebarThreadItem.tsx b/windows/app/src/components/SidebarThreadItem.tsx
new file mode 100644
index 0000000..4a82b3a
--- /dev/null
+++ b/windows/app/src/components/SidebarThreadItem.tsx
@@ -0,0 +1,83 @@
+// SidebarThreadItem — rename/delete actions revealed on hover/focus.
+import { Pencil, Trash2 } from "lucide-react";
+import { useState } from "react";
+import { t } from "../lib/t";
+
+export function SidebarThreadItem({
+  thread,
+  active,
+  onSelect,
+  onRename,
+  onDelete,
+}: {
+  thread: { id: string; data: { title: string } };
+  active: boolean;
+  onSelect: () => void;
+  onRename: (title: string) => void;
+  onDelete: () => void;
+}) {
+  const [editing, setEditing] = useState(false);
+  const [draft, setDraft] = useState("");
+
+  if (editing) {
+    return (
+      <li>
+        <form
+          onSubmit={(e) => {
+            e.preventDefault();
+            onRename(draft.trim() || thread.data.title);
+            setEditing(false);
+          }}
+        >
+          <input
+            autoFocus
+            value={draft}
+            onChange={(e) => setDraft(e.target.value)}
+            onBlur={() => setEditing(false)}
+            className="w-full rounded-lg border border-accent bg-surface px-3 py-2 text-sm outline-none"
+          />
+        </form>
+      </li>
+    );
+  }
+
+  return (
+    <li
+      className="group flex items-center rounded-lg transition-colors hover:bg-surface2/80"
+      data-active={active || undefined}
+    >
+      <button
+        type="button"
+        onClick={onSelect}
+        className={
+          "min-w-0 flex-1 truncate rounded-md px-2.5 py-1.5 text-left text-sm " +
+          (active ? "bg-surface2 font-semibold text-fg shadow-sm" : "text-fg/85")
+        }
+        title={thread.data.title}
+      >
+        {thread.data.title}
+      </button>
+      <span className="mr-1 hidden shrink-0 items-center gap-0.5 group-hover:flex group-focus-within:flex">
+        <button
+          type="button"
+          aria-label={t("chat.rename")}
+          onClick={() => {
+            setEditing(true);
+            setDraft(thread.data.title);
+          }}
+          className="e0-icon-btn !size-7 !rounded-md"
+        >
+          <Pencil size={13} />
+        </button>
+        <button
+          type="button"
+          aria-label={t("chat.deleteThread")}
+          onClick={onDelete}
+          className="e0-icon-btn !size-7 !rounded-md hover:!text-err"
+        >
+          <Trash2 size={13} />
+        </button>
+      </span>
+    </li>
+  );
+}
diff --git a/windows/app/src/components/chat/ComposerBar.tsx b/windows/app/src/components/chat/ComposerBar.tsx
new file mode 100644
index 0000000..455ba15
--- /dev/null
+++ b/windows/app/src/components/chat/ComposerBar.tsx
@@ -0,0 +1,81 @@
+// ComposerBar — floating rounded composer: Enter sends / Shift+Enter newlines,
+// the button becomes Stop while generating, ModelPill + GenPill on the left.
+import { ArrowUp, Square } from "lucide-react";
+import { useState } from "react";
+
+import type { ThreadChatApi } from "../../chat/useThreadChat";
+import { t } from "../../lib/t";
+import { GenPill } from "./GenPill";
+import { ModelPill } from "./ModelPill";
+
+export function ComposerBar({
+  chat,
+  model,
+  onModelChange,
+}: {
+  chat: ThreadChatApi;
+  model: string;
+  onModelChange: (id: string) => void;
+}) {
+  const [text, setText] = useState("");
+  const [sending, setSending] = useState(false);
+
+  async function submit() {
+    const trimmed = text.trim();
+    if (!trimmed || sending) return;
+    setSending(true);
+    setText("");
+    try {
+      await chat.send(trimmed);
+    } finally {
+      setSending(false);
+    }
+  }
+
+  return (
+    <div className="e0-composer" data-testid="composer">
+      <div className="e0-composer-shell">
+        <textarea
+          value={text}
+          onChange={(e) => setText(e.target.value)}
+          onKeyDown={(e) => {
+            if (e.key === "Enter" && !e.shiftKey) {
+              e.preventDefault();
+              void submit();
+            }
+          }}
+          rows={2}
+          placeholder={t("chat.placeholder")}
+          className="e0-composer-textarea w-full resize-none bg-transparent text-sm outline-none placeholder:text-muted"
+        />
+        <div className="e0-composer-toolbar">
+          <ModelPill model={model} onModelChange={onModelChange} wireJson={chat.wireJson} />
+          <GenPill />
+          <span className="min-w-0 flex-1 truncate text-xs text-muted">{!model ? t("chat.pickModelFirst") : ""}</span>
+          {chat.busy ? (
+            <button
+              type="button"
+              aria-label={t("chat.stop")}
+              title={t("chat.stop")}
+              onClick={chat.stop}
+              className="e0-send-btn bg-fg text-surface"
+            >
+              <Square size={14} fill="currentColor" />
+            </button>
+          ) : (
+            <button
+              type="button"
+              aria-label={t("chat.send")}
+              title={t("chat.send")}
+              onClick={() => void submit()}
+              disabled={!model || !text.trim()}
+              className="e0-send-btn disabled:opacity-40"
+            >
+              <ArrowUp size={16} />
+            </button>
+          )}
+        </div>
+      </div>
+    </div>
+  );
+}
diff --git a/windows/app/src/components/chat/GenPill.tsx b/windows/app/src/components/chat/GenPill.tsx
new file mode 100644
index 0000000..652f015
--- /dev/null
+++ b/windows/app/src/components/chat/GenPill.tsx
@@ -0,0 +1,35 @@
+// GenPill — three generation styles, switched in place on the composer toolbar; the next message reads it live.
+import { useState } from "react";
+import { t } from "../../lib/t";
+import { GEN_TEMPERATURE, setGen, storedGen, type GenPreset } from "../../lib/generation";
+
+const PRESETS: GenPreset[] = ["precise", "balanced", "creative"];
+
+export function GenPill() {
+  const [gen, setGenState] = useState<GenPreset>(() => storedGen());
+  return (
+    <div
+      className="flex shrink-0 items-center gap-0.5 rounded-full border border-line bg-surface2 px-1 py-0.5"
+      data-testid="gen-pill"
+      title={t("gen.hint")}
+    >
+      {PRESETS.map((p) => (
+        <button
+          key={p}
+          type="button"
+          data-testid={`gen-${p}`}
+          data-state={gen === p ? "on" : "off"}
+          aria-pressed={gen === p}
+          onClick={() => {
+            setGen(p);
+            setGenState(p);
+          }}
+          className="rounded-full px-2 py-0.5 text-[11px] text-muted transition-colors hover:text-fg data-[state=on]:bg-accent data-[state=on]:text-accentink"
+          title={`${t(`gen.${p}`)} · ${GEN_TEMPERATURE[p]}`}
+        >
+          {t(`gen.${p}`)}
+        </button>
+      ))}
+    </div>
+  );
+}
diff --git a/windows/app/src/components/chat/MessageItem.tsx b/windows/app/src/components/chat/MessageItem.tsx
new file mode 100644
index 0000000..4658eba
--- /dev/null
+++ b/windows/app/src/components/chat/MessageItem.tsx
@@ -0,0 +1,193 @@
+// MessageItem — one rendered message: user bubble (right, editable + resend), assistant
+// full-width with a drawn avatar, action bar on hover, collapsible reasoning, and a
+// measured usage/perf line that renders nothing when keys are missing.
+import { Streamdown } from "streamdown";
+import { Brain, Check, Copy, Pencil, RotateCcw } from "lucide-react";
+import { useState } from "react";
+
+import type { MsgPerf, ThreadMsg } from "../../state/threads";
+import { codeOf, errKey, t } from "../../lib/t";
+
+export function MessageItem({
+  message,
+  streaming,
+  isLastAssistant,
+  busy,
+  onRegenerate,
+  onEditResend,
+}: {
+  message: ThreadMsg;
+  streaming: boolean;
+  isLastAssistant: boolean;
+  busy: boolean;
+  onRegenerate: () => void;
+  onEditResend: (newText: string) => void;
+}) {
+  const [editing, setEditing] = useState(false);
+  const [draft, setDraft] = useState("");
+  const [copied, setCopied] = useState<"text" | "reasoning" | null>(null);
+  const isUser = message.role === "user";
+  const text = message.content;
+  const reasoning = message.reasoning ?? "";
+  const meta = message.meta ?? {};
+  const live = streaming && isLastAssistant && !isUser;
+  const isAborted = meta.status === "aborted";
+  const isError = meta.status === "error";
+
+  async function copy(value: string, which: "text" | "reasoning") {
+    try {
+      await navigator.clipboard.writeText(value);
+      setCopied(which);
+      setTimeout(() => setCopied(null), 1200);
+    } catch {
+      /* clipboard denied in webview: fail silently, never block the UI */
+    }
+  }
+
+  if (isUser) {
+    return (
+      <article data-msg-role="user" className="group/msg flex flex-col items-end py-2.5">
+        {editing ? (
+          <form
+            onSubmit={(e) => {
+              e.preventDefault();
+              setEditing(false);
+              if (draft.trim()) onEditResend(draft);
+            }}
+            className="w-full max-w-[88%]"
+          >
+            <textarea
+              autoFocus
+              value={draft}
+              onChange={(e) => setDraft(e.target.value)}
+              rows={3}
+              className="w-full resize-y rounded-2xl border border-accent/50 bg-surface2 p-3 text-sm outline-none"
+            />
+            <div className="flex justify-end gap-2 pt-1.5">
+              <button type="submit" className="e0-btn e0-btn-primary">{t("chat.resend")}</button>
+              <button type="button" onClick={() => setEditing(false)} className="e0-btn e0-btn-ghost">{t("common.cancel")}</button>
+            </div>
+          </form>
+        ) : (
+          <>
+            <div className="max-w-[88%] whitespace-pre-wrap break-words rounded-2xl rounded-br-md bg-accent px-4 py-3 text-sm leading-relaxed text-accentink shadow-sm">
+              {text}
+            </div>
+            <div className="mt-0.5 flex h-6 gap-1 opacity-0 transition-opacity group-hover/msg:opacity-100 focus-within:opacity-100 [@media(hover:none)]:opacity-100">
+              {text && (
+                <IconAction label={t("chat.copy")} onClick={() => void copy(text, "text")}>
+                  {copied === "text" ? <Check size={13} /> : <Copy size={13} />}
+                </IconAction>
+              )}
+              {!busy && (
+                <IconAction label={t("chat.editResend")} onClick={() => { setEditing(true); setDraft(text); }}>
+                  <Pencil size={13} />
+                </IconAction>
+              )}
+            </div>
+          </>
+        )}
+      </article>
+    );
+  }
+
+  return (
+    <article data-msg-role="assistant" className="group/msg flex gap-3 border-t border-line/40 py-5 first:border-t-0">
+      <Avatar />
+      <div className="min-w-0 flex-1">
+        {reasoning && (
+          <details className="mb-3 max-w-full rounded-xl border border-line bg-surface2/70 px-3 py-2" data-testid="reasoning-block">
+            <summary className="flex cursor-pointer select-none items-center gap-1.5 text-xs font-medium text-muted">
+              <Brain size={13} aria-hidden />
+              {t("chat.reasoning")}
+            </summary>
+            <div className="mt-1.5 grid gap-1.5">
+              <div className="whitespace-pre-wrap text-xs leading-relaxed text-muted">{reasoning}</div>
+              <button type="button" onClick={() => void copy(reasoning, "reasoning")} className="flex w-fit items-center gap-1 text-xs text-muted hover:text-fg">
+                {copied === "reasoning" ? <Check size={12} /> : <Copy size={12} />}
+                {t("chat.copyReasoning")}
+              </button>
+            </div>
+          </details>
+        )}
+
+        <div
+          className="e0-prose prose prose-sm max-w-none break-words text-sm [&>*:first-child]:mt-0 [&>*:last-child]:mb-0"
+          data-testid="message-body"
+        >
+          <Streamdown mode={live ? "streaming" : "static"}>{text}</Streamdown>
+          {live && !text && <span className="inline-block h-4 w-1 animate-pulse bg-fg align-middle" />}
+        </div>
+
+        {isError && (
+          <p className="mt-1.5 flex items-center gap-2 rounded-lg bg-err/10 px-2.5 py-1.5 text-xs text-err" role="alert">
+            <span>{t(errKey(meta.code))}</span>
+            {meta.code && <code className="rounded bg-err/10 px-1">{meta.code}</code>}
+          </p>
+        )}
+        {isAborted && <p className="mt-1 text-xs text-warn">{t("chat.aborted")}</p>}
+
+        {/* per-turn performance metrics: shown as soon as complete, stays; renders nothing when all fields are absent (no-fake-values rule) */}
+        {!live && <PerfLine perf={meta.perf} />}
+
+        <div className="mt-0.5 flex h-6 items-center gap-1 opacity-0 transition-opacity group-hover/msg:opacity-100 focus-within:opacity-100 [@media(hover:none)]:opacity-100">
+          {text && (
+            <IconAction label={t("chat.copy")} onClick={() => void copy(text, "text")}>
+              {copied === "text" ? <Check size={13} /> : <Copy size={13} />}
+            </IconAction>
+          )}
+          {isLastAssistant && !busy && (
+            <IconAction label={t("chat.regenerate")} onClick={onRegenerate}>
+              <RotateCcw size={13} />
+            </IconAction>
+          )}
+          {meta.usage && (
+            <span className="ml-auto text-[11px] tabular-nums text-muted" data-testid="usage-line">
+              {t("chat.usage", {
+                p: meta.usage.prompt_tokens,
+                c: meta.usage.completion_tokens,
+                t: meta.usage.total_tokens,
+              })}
+            </span>
+          )}
+        </div>
+      </div>
+    </article>
+  );
+}
+
+function IconAction({ label, onClick, children }: { label: string; onClick: () => void; children: React.ReactNode }) {
+  return (
+    <button type="button" aria-label={label} title={label} onClick={onClick} className="e0-icon-btn !size-7 !rounded-md">
+      {children}
+    </button>
+  );
+}
+
+/** Per-turn performance line (timings measured in the final stream chunk): renders nothing when no value is present. */
+function PerfLine({ perf }: { perf: MsgPerf | undefined }) {
+  const p = (perf ?? {}) as Record<string, number | undefined>;
+  const parts: string[] = [];
+  if (typeof p.decode_tps === "number") parts.push(t("chat.perf.decode", { v: p.decode_tps.toFixed(1) }));
+  if (typeof p.prefill_tps === "number") parts.push(t("chat.perf.prefill", { v: p.prefill_tps.toFixed(1) }));
+  if (typeof p.prompt_ms === "number") parts.push(t("chat.perf.prompt", { ms: Math.round(p.prompt_ms) }));
+  if (parts.length === 0) return null;
+  return (
+    <p className="mt-1 flex flex-wrap gap-x-2 text-[11px] tabular-nums text-muted" data-testid="perf-line">
+      {parts.map((s, i) => (
+        <span key={i}>{s}</span>
+      ))}
+    </p>
+  );
+}
+
+/** Drawn monochrome avatar (brand mark). */
+function Avatar() {
+  return (
+    <span aria-hidden className="mt-0.5 grid h-7 w-7 shrink-0 place-content-center rounded-full border border-line bg-surface2 text-fg">
+      <svg width="14" height="14" viewBox="0 0 14 14" fill="none">
+        <path d="M7 1l1.6 4.4L13 7l-4.4 1.6L7 13 5.4 8.6 1 7l4.4-1.6L7 1z" fill="currentColor" />
+      </svg>
+    </span>
+  );
+}
diff --git a/windows/app/src/components/chat/MessageList.tsx b/windows/app/src/components/chat/MessageList.tsx
new file mode 100644
index 0000000..b12cdad
--- /dev/null
+++ b/windows/app/src/components/chat/MessageList.tsx
@@ -0,0 +1,28 @@
+// MessageList — full-width list + "waiting for first chunk" placeholder (all threads live locally; no pagination).
+import type { ThreadChatApi } from "../../chat/useThreadChat";
+import { t } from "../../lib/t";
+import { MessageItem } from "./MessageItem";
+
+export function MessageList({ chat }: { chat: ThreadChatApi }) {
+  const lastAssistantIdx = chat.messages.map((m) => m.role).lastIndexOf("assistant");
+  return (
+    <div data-testid="message-list">
+      {chat.messages.map((m, i) => (
+        <MessageItem
+          key={m.id ?? i}
+          message={m}
+          streaming={chat.busy}
+          isLastAssistant={i === lastAssistantIdx}
+          busy={chat.busy}
+          onRegenerate={() => void chat.regenerate()}
+          onEditResend={(next) => void chat.editResend(m.id, next)}
+        />
+      ))}
+      {chat.busy && chat.messages[chat.messages.length - 1]?.role !== "assistant" && (
+        <p className="px-3 text-xs text-muted" data-testid="waiting-first-chunk">
+          {t("chat.waitingFirstChunk")}
+        </p>
+      )}
+    </div>
+  );
+}
diff --git a/windows/app/src/components/chat/ModelPill.tsx b/windows/app/src/components/chat/ModelPill.tsx
new file mode 100644
index 0000000..e9b8540
--- /dev/null
+++ b/windows/app/src/components/chat/ModelPill.tsx
@@ -0,0 +1,166 @@
+// ModelPill — model pill at the composer's left end. Data comes from the shell:
+// candidates = installed list, light = pure modelStatus function, token metering =
+// measured via /apply-template -> /tokenize (unavailable = unknown, never estimated).
+import { ChevronDown, Loader2 } from "lucide-react";
+import { useEffect, useRef, useState, useSyncExternalStore } from "react";
+
+import { engineStore, installedTiers, loadModel, shellStore } from "../../api/client";
+import { threadStore } from "../../state/threads";
+import { candidateModels, deriveLight, type ModelLight } from "../../chat/modelStatus";
+import { codeOf, errKey, t } from "../../lib/t";
+
+function LightDot({ light }: { light: ModelLight }) {
+  const cls =
+    light === "resident" ? "bg-ok" : light === "loading" ? "bg-warn animate-pulse" : light === "installed" ? "bg-muted" : "bg-line";
+  return <span data-testid="model-light" data-light={light} className={"inline-block h-2 w-2 shrink-0 rounded-full " + cls} aria-hidden />;
+}
+
+export function ModelPill({
+  model,
+  onModelChange,
+  wireJson,
+}: {
+  model: string;
+  onModelChange: (id: string) => void;
+  wireJson: () => { role: string; content: string }[];
+}) {
+  const [, force] = useState(0);
+  const [open, setOpen] = useState(false);
+  const [loadError, setLoadError] = useState<string | null>(null);
+  const [contextTokens, setContextTokens] = useState<number | null>(null);
+  const debounceRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const wireRef = useRef(wireJson);
+  wireRef.current = wireJson;
+  const rootRef = useRef<HTMLDivElement>(null);
+
+  useSyncExternalStore(engineStore.sig.sub, engineStore.sig.version);
+  useSyncExternalStore(shellStore.sig.sub, shellStore.sig.version);
+  useEffect(() => engineStore.sig.sub(() => force((n) => n + 1)), []);
+
+  useEffect(() => {
+    if (!open) return;
+    const onDoc = (e: MouseEvent) => {
+      if (!rootRef.current?.contains(e.target as Node)) setOpen(false);
+    };
+    document.addEventListener("mousedown", onDoc);
+    return () => document.removeEventListener("mousedown", onDoc);
+  }, [open]);
+
+  const st = engineStore.status;
+  const installed = installedTiers();
+  const candidates = candidateModels(installed);
+  const light = model ? deriveLight(st, engineStore.loadingTier, installed, model) : "absent";
+
+  // metering: on message/model/port change, debounce 400 ms then measure via the two-hop call.
+  const msgVer = useSyncExternalStore(
+    threadStore.subscribe,
+    () => threadStore.version(),
+  );
+  useEffect(() => {
+    if (debounceRef.current) clearTimeout(debounceRef.current);
+    debounceRef.current = setTimeout(() => void measure(), 400);
+    return () => {
+      if (debounceRef.current) clearTimeout(debounceRef.current);
+    };
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [msgVer, model, st?.port]);
+
+  async function measure() {
+    const base = engineStore.status?.running ? engineStore.status.base_url : undefined;
+    if (!base || !model) {
+      setContextTokens(null);
+      return;
+    }
+    try {
+      const r1 = await fetch(`${base}/apply-template`, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({ messages: wireRef.current(), add_generation_prompt: true }),
+      });
+      if (!r1.ok) throw new Error(`HTTP ${r1.status}`);
+      const j1 = (await r1.json()) as { prompt?: string };
+      if (typeof j1.prompt !== "string") throw new Error("no prompt");
+      const r2 = await fetch(`${base}/tokenize`, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({ content: j1.prompt }),
+      });
+      if (!r2.ok) throw new Error(`HTTP ${r2.status}`);
+      const j2 = (await r2.json()) as { tokens?: unknown[] };
+      setContextTokens(Array.isArray(j2.tokens) ? j2.tokens.length : null);
+    } catch {
+      setContextTokens(null); // unavailable = unknown: no estimate, no fake zero
+    }
+  }
+
+  async function load(id: string) {
+    setLoadError(null);
+    try {
+      await loadModel(id);
+    } catch (e) {
+      setLoadError(t(errKey(codeOf(String(e))))) + (codeOf(String(e)) ? ` · ${codeOf(String(e))}` : "");
+    }
+  }
+
+  return (
+    <div ref={rootRef} className="relative" data-testid="model-pill">
+      <button
+        type="button"
+        aria-haspopup="listbox"
+        aria-expanded={open}
+        onClick={() => setOpen((o) => !o)}
+        className="flex min-w-0 max-w-[18rem] items-center gap-1.5 rounded-full border border-line bg-surface2 px-2.5 py-1.5 text-xs font-medium text-fg hover:border-accent/50 hover:bg-surface"
+      >
+        <LightDot light={light} />
+        <span className="max-w-[10rem] truncate">{model || t("chat.noModel")}</span>
+        <span className="text-muted" data-testid="context-meter">
+          {contextTokens === null ? t("chat.contextUnknown") : t("chat.contextUsed", { n: contextTokens })}
+        </span>
+        <ChevronDown size={12} aria-hidden />
+      </button>
+
+      {open && (
+        <div role="listbox" className="absolute bottom-full left-0 z-20 mb-2 w-72 rounded-2xl border border-line bg-surface p-2 shadow-lg">
+          {candidates.length === 0 && <p className="px-2 py-1.5 text-sm text-muted">{t("chat.noModel")}</p>}
+          {candidates.map((id) => {
+            const l = deriveLight(st, engineStore.loadingTier, installed, id);
+            return (
+              <div key={id} className="flex items-center gap-2 rounded-lg px-2 py-1.5 text-sm hover:bg-surface2">
+                <button
+                  type="button"
+                  role="option"
+                  aria-selected={id === model}
+                  onClick={() => {
+                    onModelChange(id);
+                    setOpen(false);
+                  }}
+                  className="flex min-w-0 flex-1 items-center gap-2 text-left"
+                >
+                  <LightDot light={l} />
+                  <span className="min-w-0 flex-1 truncate">{id}</span>
+                  <span className="text-xs text-muted">{t(`chat.light.${l}`)}</span>
+                </button>
+                {(l === "installed" || l === "loading") && id === model && (
+                  <button
+                    type="button"
+                    onClick={() => void load(id)}
+                    disabled={l === "loading"}
+                    className="flex items-center gap-1 rounded-full border border-line px-2 py-0.5 text-xs text-muted hover:text-fg disabled:opacity-60"
+                  >
+                    {l === "loading" && <Loader2 size={11} className="animate-spin" aria-hidden />}
+                    {t("chat.loadModel")}
+                  </button>
+                )}
+              </div>
+            );
+          })}
+          {loadError && (
+            <p className="mt-1 rounded-lg bg-err/10 px-2 py-1.5 text-xs text-err" role="alert">
+              {loadError}
+            </p>
+          )}
+        </div>
+      )}
+    </div>
+  );
+}
diff --git a/windows/app/src/lib/autoLoad.ts b/windows/app/src/lib/autoLoad.ts
new file mode 100644
index 0000000..2468ea5
--- /dev/null
+++ b/windows/app/src/lib/autoLoad.ts
@@ -0,0 +1,65 @@
+// One-shot intent token for "Download & enable" (in memory; consumed on use).
+// This module is the ONLY place that may auto-chain a finished download into a load —
+// every other load path requires an explicit user click.
+type Phase = "download" | "load";
+
+interface Token {
+  tier: string;
+  taskId: string; // download tasks are keyed by tier (shell-side pull.rs)
+  phase: Phase;
+}
+
+type Listener = () => void;
+
+class AutoLoadStore {
+  private token: Token | null = null;
+  private listeners = new Set<Listener>();
+
+  get current(): Token | null {
+    return this.token;
+  }
+
+  subscribe(fn: Listener): () => void {
+    this.listeners.add(fn);
+    return () => {
+      this.listeners.delete(fn);
+    };
+  }
+
+  private emit(): void {
+    for (const l of this.listeners) l();
+  }
+
+  /** User pressed "Download & enable" (the only legal entry point). */
+  begin(tier: string, taskId: string): void {
+    this.token = { tier, taskId, phase: "download" };
+    this.emit();
+  }
+
+  /** download ready for the same task: move to the load phase (returns tier). One-shot: never revives. */
+  claimCompleted(tier: string): string | null {
+    if (!this.token || this.token.phase !== "download" || this.token.taskId !== tier) return null;
+    const t = this.token.tier;
+    this.token = { ...this.token, phase: "load" };
+    this.emit();
+    return t;
+  }
+
+  /** load succeeded: consume the token (returns tier for thread creation); later duplicate ready events are no-ops. */
+  consume(tier: string): boolean {
+    if (!this.token || this.token.phase !== "load" || this.token.tier !== tier) return false;
+    this.token = null;
+    this.emit();
+    return true;
+  }
+
+  /** failure / abandon / close: invalidate immediately (never persisted, so a restart naturally drops it). */
+  cancel(): void {
+    if (this.token) {
+      this.token = null;
+      this.emit();
+    }
+  }
+}
+
+export const autoLoad = new AutoLoadStore();
diff --git a/windows/app/src/lib/byteFormat.ts b/windows/app/src/lib/byteFormat.ts
new file mode 100644
index 0000000..b04f1c8
--- /dev/null
+++ b/windows/app/src/lib/byteFormat.ts
@@ -0,0 +1,17 @@
+// Byte / rate / ETA display formatting — unit conversion and rounding only; never invents numbers the data source didn't provide.
+export function humanBytes(n: number): string {
+  if (n >= 1e12) return `${(n / 1e12).toFixed(2)} TB`;
+  if (n >= 1e9) return `${(n / 1e9).toFixed(2)} GB`;
+  if (n >= 1e6) return `${(n / 1e6).toFixed(2)} MB`;
+  if (n >= 1e3) return `${(n / 1e3).toFixed(1)} KB`;
+  return `${n} B`;
+}
+
+export function humanRate(bytesPerSec: number): string {
+  return `${humanBytes(bytesPerSec)}/s`;
+}
+
+/** ETA display granularity: round to 30 s to avoid jitter; returns minutes (0.5 steps allowed). */
+export function etaMinutes(etaS: number): number {
+  return Math.max(0.5, (Math.round(etaS / 30) * 30) / 60);
+}
diff --git a/windows/app/src/lib/generation.ts b/windows/app/src/lib/generation.ts
new file mode 100644
index 0000000..3b8721b
--- /dev/null
+++ b/windows/app/src/lib/generation.ts
@@ -0,0 +1,26 @@
+// Generation presets: three temperature tiers applied globally, default "Precise" 0.2.
+// UI is the GenPill on the composer toolbar; the transport reads the current value on
+// every send (tier changes take effect from the next message).
+export type GenPreset = "precise" | "balanced" | "creative";
+
+const GEN_KEY = "edge0.gen";
+
+export const GEN_TEMPERATURE: Record<GenPreset, number> = {
+  precise: 0.2,
+  balanced: 0.7,
+  creative: 1.0,
+};
+
+export function storedGen(): GenPreset {
+  const v = localStorage.getItem(GEN_KEY);
+  return v === "balanced" || v === "creative" ? v : "precise";
+}
+
+export function setGen(p: GenPreset): void {
+  localStorage.setItem(GEN_KEY, p);
+}
+
+/** Current temperature (read live on every request). */
+export function currentTemperature(): number {
+  return GEN_TEMPERATURE[storedGen()];
+}
diff --git a/windows/app/src/lib/systemPrompt.ts b/windows/app/src/lib/systemPrompt.ts
new file mode 100644
index 0000000..8ca2dab
--- /dev/null
+++ b/windows/app/src/lib/systemPrompt.ts
@@ -0,0 +1,22 @@
+// Optional system prompt: empty string = not sent.
+export const DEFAULT_SYSTEM_PROMPT = "You're edge0, an on-device AI assistant.";
+
+const SYSTEM_PROMPT_KEY = "edge0.system-prompt";
+const SYSTEM_PROMPT_EVENT = "edge0:system-prompt";
+
+export function storedSystemPrompt(): string | null {
+  const stored = localStorage.getItem(SYSTEM_PROMPT_KEY);
+  if (stored === null) return DEFAULT_SYSTEM_PROMPT;
+  const trimmed = stored.trim();
+  return trimmed.length > 0 ? stored : null;
+}
+
+export function setSystemPrompt(value: string): void {
+  localStorage.setItem(SYSTEM_PROMPT_KEY, value);
+  window.dispatchEvent(new Event(SYSTEM_PROMPT_EVENT));
+}
+
+export function subscribeSystemPrompt(listener: () => void): () => void {
+  window.addEventListener(SYSTEM_PROMPT_EVENT, listener);
+  return () => window.removeEventListener(SYSTEM_PROMPT_EVENT, listener);
+}
diff --git a/windows/app/src/lib/t.ts b/windows/app/src/lib/t.ts
new file mode 100644
index 0000000..a323cd7
--- /dev/null
+++ b/windows/app/src/lib/t.ts
@@ -0,0 +1,183 @@
+// lib/t.ts — single source for UI strings: **English-only** (product decision for the
+// overseas release: no language switching). Dot-separated keys with {{var}} interpolation;
+// a missing key echoes the key itself (honest surface, never swallowed silently).
+type Table = Record<string, unknown>;
+
+const en: Table = {
+  app: { name: "edge0", localLabel: "local inference" },
+  platform: { note: "Windows · local llama.cpp engine" },
+  theme: { light: "Light", dark: "Dark", system: "System" },
+  font: { sm: "Small", md: "Default", lg: "Large" },
+  nav: {
+    models: "Models", service: "Service", settings: "Settings",
+    openSidebar: "Open sidebar", closeSidebar: "Close sidebar",
+  },
+  common: { cancel: "Cancel" },
+  sidebar: {
+    newChat: "New chat", history: "Chat history", navigation: "Navigation",
+    resident: "resident", idle: "idle", noService: "No service",
+    scopeShell: "managed in-app",
+    bucket: { today: "Today", yesterday: "Yesterday", week: "Previous 7 days", earlier: "Earlier" },
+  },
+  chat: {
+    newThread: "New chat", rename: "Rename", deleteThread: "Delete chat", noThreads: "No chats yet",
+    placeholder: "Type a message, Enter to send", send: "Send", stop: "Stop", regenerate: "Regenerate",
+    editResend: "Edit & resend", resend: "Resend", copy: "Copy", copyReasoning: "Copy thinking",
+    reasoning: "Thinking", aborted: "Response interrupted", saveFailed: "Not saved (conversation store write failed)",
+    usage: "tokens in {{p}}/out {{c}}/total {{t}}",
+    perf: { decode: "{{v}} tok/s decode", prefill: "{{v}} tok/s prefill", prompt: "prefill {{ms}} ms", mem: "memory {{gi}} GiB" },
+    model: "Model", noModel: "(no installed model)",
+    light: { resident: "resident", loading: "loading", installed: "idle", absent: "unknown" },
+    loadModel: "Load",
+    contextUsed: "context {{n}} tok", contextUnknown: "context: unknown",
+    greeting: "What would you like to explore?", pickOrCreate: "Select or create a chat",
+    serviceShell: "local engine", noInstalledModel: "No models to chat with yet — download one first",
+    goDownload: "Go to downloads",
+    waitingFirstChunk: "Waiting for the model…", pickModelFirst: "Pick a model first",
+    residentForever: "Resident until manually evicted",
+    err: {
+      notload: "Model not loaded — press Load, or download it",
+      busy: "Service busy, retry shortly", cancel: "Generation cancelled",
+      worker: "Engine process crashed — restart it from the Service page and retry",
+      param: "Request parameters rejected — check the model selection",
+      missing: "Model not installed — download it first",
+      memload: "Not enough memory to load — use a smaller tier",
+      version: "Engine binary missing — set EDGE0_BIN_DIR or start from the Service page",
+      pool: "Prefetch pool not active (POOL2 telemetry line missing) — performance claims downgraded",
+      protected: "This tier is a seed link (junction) — the shell never deletes through links",
+      generic: "Request failed",
+    },
+  },
+  settings: {
+    title: "Settings", kicker: "Tune edge0 appearance and local service access",
+    appearance: "Appearance", appearanceNote: "Theme and font size apply across the app immediately.",
+    systemPrompt: "System prompt", systemPromptNote: "Optional instructions sent before each chat. Leave blank to disable.",
+    systemPromptEnabled: "System prompt enabled", systemPromptDisabled: "No system prompt will be sent",
+    systemPromptReset: "Use default",
+    service: "Service", serviceNote: "View a compact summary; full controls, logs, and diagnostics live in the service console.",
+    serviceUnavailable: "Service not connected", advancedNote: "Request logs and environment checks live in the service console.",
+    openConsole: "Open service console",
+  },
+  gen: {
+    precise: "Precise", balanced: "Balanced", creative: "Creative",
+    hint: "Generation style: sampling temperature (Precise 0.2 / Balanced 0.7 / Creative 1.0), applied from the next message.",
+  },
+  models: {
+    title: "Models & Downloads", kicker: "Manage installed tiers, residency, and downloads",
+    refresh: "Refresh", loading: "Reading catalog…", retry: "Retry",
+    stale: "Catalog manifest expired (mirrors unreachable): browsing OK, downloads paused",
+    rev: "Rev", bytes: "Size", files: "Files", source: "Source", singleSource: "single source",
+    license: "License", notes: "Model card", pool: "Prefetch pool", repo: "Source repo",
+    state: { absent: "Not installed", installed: "Installed", update: "Installed" },
+    download: "Download", sourceAuto: "Auto source (speed probe)", sourceMs: "ModelScope", sourceHf: "Hugging Face", sourceHfMirror: "HF mirror",
+    load: "Enable", unload: "Evict", loadingModel: "Loading…",
+    delete: "Delete", deleteConfirm: "Delete this tier? Both ~/.edge0 directories are removed physically (seed links are refused).",
+    loadFailed: "Load failed: {{msg}}",
+    pause: "Cancel", cancelTask: "Cancel task", deleteTask: "Delete task",
+    sourceSwitched: "source switched", eta: "~{{n}} min left", estimating: "estimating speed…", verifying: "verifying…",
+    filesProgress: "{{done}}/{{total}} files",
+    phase: {
+      idle: "Idle", probe: "Probing sources", download: "Downloading", convert: "On-device converting",
+      ready: "Done", error: "Failed", cancelled: "Cancelled",
+    },
+    err: {
+      net: "Mirror or source temporarily unreachable — check network and retry",
+      src: "Source rejects resume semantics or unreachable — retry later or switch network",
+      disk: "Not enough disk space — free space and retry, or pick a smaller tier",
+      hash: "File hash mismatch with manifest — delete the task and retry",
+      size: "File size mismatch with manifest — delete the task and retry",
+      convert: "On-device conversion failed — check tools/convert_mlx_to_gguf.py, then retry",
+      protected: "This tier is a seed link (junction) — the shell unlinks, never deletes targets",
+      busy: "Tier is resident — evict it before deleting",
+      invalid: "Model files do not match the install record — delete the tier and download again",
+      generic: "Request failed",
+    },
+  },
+  welcome: {
+    title: "Get started with edge0 on this PC",
+    body: "Recommended: {{tier}} ({{bytes}}, stored in {{path}}) — inference runs entirely on-device.",
+    bytesUnknown: "size unknown", pathUnknown: "local app directory",
+    downloadEnable: "Download & enable", skip: "Skip for now",
+    errors: { disk: "Not enough disk space", memory: "Not enough memory to load the recommended tier", engine: "The local engine is unavailable", generic: "First launch did not complete" },
+  },
+  service: {
+    title: "Local API Service", kicker: "Manage engine lifecycle, access, and runtime diagnostics",
+    sections: {
+      lifecycle: "Lifecycle & ownership", lifecycleNote: "Start or stop the engine and inspect port, version, and uptime.",
+      logs: "Request log", logsNote: "Shows the generation-request window since this app started.",
+      doctor: "Environment checks", doctorNote: "Shows shell-side measured findings and next actions.",
+    },
+    panel: {
+      title: "Service Panel", start: "Start service", stop: "Stop service",
+      busy: "Service operation in progress… (restart ends in-flight requests)", actionFailed: "The service action could not be completed.",
+      port: "Listening", bind: "Bind address", scope: "Ownership", uptime: "Uptime",
+      version: "Engine version", home: "Storage location", pool: "Prefetch pool", tier: "Resident tier",
+      governance: "Process governance", governanceJob: "Job Object armed (OS reclaims the engine when the app dies)",
+      governanceSoft: "Normal-exit fallback only (job attach failed — stated honestly)",
+      token: "Access token", noCredential: "No credential needed (loopback is exempt)",
+      notOwned: "This service was not started by this app; stop it via its owner",
+      lan: "Allow LAN", lanOn: "On", lanOff: "Off", lanLater: "Phase 2: LAN access is not implemented yet",
+      idle: "Engine not running",
+    },
+    log: {
+      title: "Request Log", count: "{{n}} requests since start",
+      note: "Window = generation requests in this app session (in memory; cleared on exit)",
+      empty: "No requests since this session started",
+      th: { status: "Status", model: "Model", latency: "Latency", tokens: "tokens", source: "Source" },
+      status: { done: "done", cancelled: "cancelled", error: "error" },
+      source: { loopback: "local" },
+    },
+    doctor: {
+      title: "Environment Check", loading: "Reading check results…", failed: "Cannot fetch environment findings", retry: "Retry",
+      next: "Next", goModels: "Go to Models",
+      overall: { pass: "All passing", warn: "Warnings" },
+      verdict: { pass: "pass", warn: "warn", fail: "fail" },
+    },
+  },
+};
+
+function lookup(key: string): string {
+  let cur: unknown = en;
+  for (const seg of key.split(".")) {
+    if (cur == null || typeof cur !== "object") return key;
+    cur = (cur as Table)[seg];
+  }
+  return typeof cur === "string" ? cur : key;
+}
+
+/** t("sidebar.bucket.today") / t("chat.usage", { p, c, t }) — with {{var}} interpolation. */
+export function t(key: string, vars?: Record<string, string | number>): string {
+  let s = lookup(key);
+  if (vars) {
+    for (const [k, v] of Object.entries(vars)) s = s.replaceAll(`{{${k}}}`, String(v));
+  }
+  return s;
+}
+
+/** Error code → human-readable key (this repo's code set is small: a flat table). */
+export function errKey(code: string | null | undefined): string {
+  if (!code) return "chat.err.generic";
+  if (code.startsWith("E-DL-NET")) return "models.err.net";
+  if (code.startsWith("E-DL-SRC")) return "models.err.src";
+  if (code.startsWith("E-DL-DISK")) return "models.err.disk";
+  if (code.startsWith("E-DL-SHA")) return "models.err.hash";
+  if (code.startsWith("E-DL-SIZE")) return "models.err.size";
+  if (code.startsWith("E-CONVERT")) return "models.err.convert";
+  if (code.startsWith("E-MODEL-PROTECTED")) return "models.err.protected";
+  if (code.startsWith("E-MODEL-BUSY")) return "models.err.busy";
+  if (code.startsWith("E-MODEL-INVALID")) return "models.err.invalid";
+  if (code.startsWith("E-MODEL-MISSING")) return "chat.err.missing";
+  if (code.startsWith("E-ENGINE-MISSING")) return "chat.err.version";
+  if (code.startsWith("E-ENGINE-DIED")) return "chat.err.worker";
+  if (code.startsWith("E-POOL")) return "chat.err.pool";
+  if (code.startsWith("E-MEM")) return "chat.err.memload";
+  if (code.startsWith("E-DISK")) return "models.err.disk";
+  return "chat.err.generic";
+}
+
+/** Extract the code from an error string (shell-side format: "E-XXX human text"; take the leading token). */
+export function codeOf(err: string | null | undefined): string | null {
+  if (!err) return null;
+  const m = err.match(/E-[A-Z-]+/);
+  return m ? m[0] : null;
+}
diff --git a/windows/app/src/lib/theme.ts b/windows/app/src/lib/theme.ts
new file mode 100644
index 0000000..31038a1
--- /dev/null
+++ b/windows/app/src/lib/theme.ts
@@ -0,0 +1,31 @@
+// Dark/light theme + font size: the shell is the single source of truth (system-follow mode resolves in JS, then persists to data-theme).
+export type Theme = "light" | "dark" | "system";
+export type FontScale = "sm" | "md" | "lg";
+
+const THEME_KEY = "edge0.token…heme";
+const FONT_KEY = "edge0.token…font";
+
+export function resolveTheme(t: Theme): "light" | "dark" {
+  if (t !== "system") return t;
+  return window.matchMedia("(prefers-color-scheme: dark)").matches ? "dark" : "light";
+}
+
+export function storedTheme(): Theme {
+  const v = localStorage.getItem(THEME_KEY);
+  return v === "light" || v === "dark" || v === "system" ? v : "system";
+}
+
+export function storedFont(): FontScale {
+  const v = localStorage.getItem(FONT_KEY);
+  return v === "sm" || v === "md" || v === "lg" ? v : "md";
+}
+
+export function applyTheme(t: Theme): void {
+  localStorage.setItem(THEME_KEY, t);
+  document.documentElement.dataset.theme = resolveTheme(t);
+}
+
+export function applyFont(f: FontScale): void {
+  localStorage.setItem(FONT_KEY, f);
+  document.documentElement.dataset.font = f;
+}
diff --git a/windows/app/src/main.tsx b/windows/app/src/main.tsx
new file mode 100644
index 0000000..18a8927
--- /dev/null
+++ b/windows/app/src/main.tsx
@@ -0,0 +1,52 @@
+// main.tsx — frontend entry. Hash routing (no router dependency):
+// `/` = chat (default landing), `/models`, `/service`, `/settings`; legacy `/chat` normalizes to `/`.
+import React, { useEffect, useState } from "react";
+import ReactDOM from "react-dom/client";
+
+import { boot } from "./api/client";
+import { AppShell } from "./components/AppShell";
+import { applyFont, applyTheme, storedFont, storedTheme } from "./lib/theme";
+import { ChatPage } from "./routes/chat";
+import { ModelsPage } from "./routes/models";
+import { ServicePage } from "./routes/service";
+import { SettingsPage } from "./routes/settings";
+import "./styles.css";
+
+// The shell is the single source of truth: boot applies data-theme/data-font before first paint (no flash of wrong theme).
+applyTheme(storedTheme());
+applyFont(storedFont());
+
+function normalize(hash: string): string {
+  const p = (hash || "").replace(/^#/, "") || "/";
+  return p === "/chat" ? "/" : p; // normalize the legacy entry; keep a single canonical route
+}
+
+function useHashRoute(): string {
+  const [path, setPath] = useState(() => normalize(window.location.hash));
+  useEffect(() => {
+    const on = () => setPath(normalize(window.location.hash));
+    window.addEventListener("hashchange", on);
+    return () => window.removeEventListener("hashchange", on);
+  }, []);
+  return path;
+}
+
+function App() {
+  const path = useHashRoute();
+  useEffect(() => {
+    void boot();
+  }, []);
+  const view = path === "/models" ? path : path === "/service" ? path : path === "/settings" ? path : "/";
+  return (
+    <AppShell view={view}>
+      {view === "/models" && <ModelsPage />}
+      {view === "/service" && <ServicePage />}
+      {view === "/settings" && <SettingsPage />}
+      {view === "/" && <ChatPage />}
+    </AppShell>
+  );
+}
+
+ReactDOM.createRoot(document.getElementById("root")!).render(
+  <React.StrictMode><App /></React.StrictMode>
+);
diff --git a/windows/app/src/models/cards.tsx b/windows/app/src/models/cards.tsx
new file mode 100644
index 0000000..72e7dfe
--- /dev/null
+++ b/windows/app/src/models/cards.tsx
@@ -0,0 +1,95 @@
+// models/cards.tsx — in-flight/failed download task cards (one implementation shared by
+// the catalog page and the welcome card). Speed is frontend-measured: sliding-window
+// deltas over the shell's progress byte stream (the shell sends no rate field).
+import { useEffect, useRef, useState, useSyncExternalStore } from "react";
+
+import { cancelDownload, downloadStore, startDownload, type TaskView } from "../api/client";
+import { etaMinutes, humanBytes, humanRate } from "../lib/byteFormat";
+import { codeOf, errKey, t } from "../lib/t";
+
+export function DownloadCardForTier({ tier }: { tier: string }) {
+  useSyncExternalStore(downloadStore.sig.sub, downloadStore.sig.version);
+  const task = downloadStore.tasks[tier];
+  if (!task || task.phase === "idle") return null;
+  return <DownloadCardBody task={task} />;
+}
+
+export function DownloadCardBody({ task }: { task: TaskView }) {
+  const pct = task.bytes_total > 0 ? Math.min(100, (task.bytes_done / task.bytes_total) * 100) : 0;
+  const downloading = task.phase === "download";
+  const rateRef = useRef<{ b: number; t: number } | null>(null);
+  const [rate, setRate] = useState<number | null>(null);
+
+  useEffect(() => {
+    if (!downloading) {
+      setRate(null);
+      rateRef.current = null;
+      return;
+    }
+    const now = Date.now();
+    const prev = rateRef.current;
+    rateRef.current = { b: task.bytes_done, t: now };
+    if (prev && now - prev.t > 500) {
+      const db = task.bytes_done - prev.b;
+      if (db >= 0) setRate(db / ((now - prev.t) / 1000));
+    }
+  }, [task.bytes_done, downloading]);
+
+  const failed = task.phase === "error" && task.error;
+  const showEta = downloading && rate !== null && rate > 0;
+  const filesDone = task.files ? task.files.filter((f) => f.state === "done").length : 0;
+
+  return (
+    <div data-testid={`dl-${task.tier}`} className="space-y-1 rounded-lg bg-surface2/50 p-3 text-left text-xs">
+      <div className="flex items-center justify-between">
+        <span data-testid={`dl-phase-${task.tier}`}>{t(`models.phase.${task.phase}`)}</span>
+        <span className="text-muted tabular-nums">
+          {humanBytes(task.bytes_done)} / {humanBytes(task.bytes_total)}
+          {downloading && rate !== null && ` · ${humanRate(rate)}`}
+        </span>
+      </div>
+      <div
+        className="h-1.5 w-full overflow-hidden rounded bg-line"
+        role="progressbar"
+        aria-valuenow={Math.round(pct)}
+        data-testid={`dl-bar-${task.tier}`}
+      >
+        <div className={"h-full " + (failed ? "bg-err" : "bg-accent")} style={{ width: `${pct}%` }} />
+      </div>
+      <div className="flex items-center justify-between gap-3">
+        <span data-testid={`dl-eta-${task.tier}`} className="min-w-0 flex-1 truncate text-muted">
+          {failed
+            ? `${t(errKey(codeOf(task.error)))} · ${codeOf(task.error) ?? "E-DL"}`
+            : showEta
+              ? t("models.eta", { n: etaMinutes((task.bytes_total - task.bytes_done) / (rate as number)) })
+              : downloading
+                ? t("models.estimating")
+                : task.files?.length
+                  ? t("models.filesProgress", { done: filesDone, total: task.files.length })
+                  : ""}
+        </span>
+        <div className="flex shrink-0 items-center gap-2">
+          {failed ? (
+            <button
+              type="button"
+              data-testid={`dl-retry-${task.tier}`}
+              onClick={() => void startDownload(task.tier, null).catch(() => undefined)}
+              className="e0-btn e0-btn-secondary"
+            >
+              {t("models.retry")}
+            </button>
+          ) : downloading ? (
+            <button
+              type="button"
+              data-testid={`dl-cancel-${task.tier}`}
+              onClick={() => void cancelDownload(task.tier)}
+              className="rounded-full border border-line px-2 py-0.5 text-muted hover:text-err"
+            >
+              {t("models.cancelTask")}
+            </button>
+          ) : null}
+        </div>
+      </div>
+    </div>
+  );
+}
diff --git a/windows/app/src/routes/chat.tsx b/windows/app/src/routes/chat.tsx
new file mode 100644
index 0000000..85a0a4d
--- /dev/null
+++ b/windows/app/src/routes/chat.tsx
@@ -0,0 +1,103 @@
+// routes/chat.tsx — chat main area: centered max-width message column + floating bottom
+// composer; no selected thread = big-title empty state; nothing installed = WelcomeCard;
+// the boot snapshot decides the waiting shape.
+import { useCallback, useEffect, useRef, useState, useSyncExternalStore } from "react";
+
+import { engineStore, installedTiers, shellStore } from "../api/client";
+import { useThreadChat } from "../chat/useThreadChat";
+import { isAtBottom } from "../chat/scroll";
+import { t } from "../lib/t";
+import { storedSystemPrompt, subscribeSystemPrompt } from "../lib/systemPrompt";
+import { threadStore, type ThreadRow } from "../state/threads";
+import { ComposerBar } from "../components/chat/ComposerBar";
+import { MessageList } from "../components/chat/MessageList";
+import { WelcomeCard } from "./welcome";
+
+export function ChatPage() {
+  useSyncExternalStore(threadStore.subscribe, () => threadStore.version());
+  useSyncExternalStore(engineStore.sig.sub, engineStore.sig.version);
+  useSyncExternalStore(shellStore.sig.sub, shellStore.sig.version);
+
+  const st = engineStore.status;
+  const active = threadStore.threads().find((x) => x.id === threadStore.activeId) ?? null;
+  const hasInstalled = installedTiers().length > 0;
+
+  if (!st) {
+    return (
+      <div className="e0-chat-empty text-sm text-muted">
+        <span className="mx-auto grid size-10 place-items-center rounded-2xl bg-surface2 text-fg">
+          <span className="size-2 rounded-full bg-warn animate-pulse" aria-hidden />
+        </span>
+        <p className="font-medium text-fg">{t("service.panel.idle")}</p>
+        <p className="text-xs">{t("platform.note")}</p>
+      </div>
+    );
+  }
+  if (!hasInstalled) return <WelcomeCard />;
+  if (!active) {
+    return (
+      <div className="e0-chat-empty">
+        <span className="mx-auto grid size-12 place-items-center rounded-2xl bg-fg text-surface shadow-sm">
+          <span className="text-lg font-semibold" aria-hidden>+</span>
+        </span>
+        <h2 className="mt-1 text-3xl font-semibold tracking-tight">{t("chat.greeting")}</h2>
+        <p className="text-sm text-muted">{t("chat.pickOrCreate")}</p>
+      </div>
+    );
+  }
+  return <ChatPane key={active.id} thread={active} />;
+}
+
+function ChatPane({ thread }: { thread: ThreadRow }) {
+  const st = engineStore.status;
+  const [model, setModel] = useState(thread.data.model ?? (st?.running ? st.tier ?? "" : ""));
+  useSyncExternalStore(subscribeSystemPrompt, () => storedSystemPrompt(), () => storedSystemPrompt());
+  const scrollerRef = useRef<HTMLDivElement>(null);
+  // Follow-bottom: pin to the bottom on streaming increments only while the user is already near the bottom.
+  const stickBottom = useRef(true);
+
+  const chat = useThreadChat(thread, model || null);
+
+  const onModelChange = useCallback(
+    (id: string) => {
+      setModel(id);
+      threadStore.setModel(thread.id, id);
+    },
+    [thread.id],
+  );
+
+  useEffect(() => {
+    const last = chat.messages[chat.messages.length - 1];
+    // right after the user sends: always jump to bottom (clear intent)
+    if (last?.role === "user") stickBottom.current = true;
+    const el = scrollerRef.current;
+    if (el && stickBottom.current) el.scrollTo({ top: el.scrollHeight });
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [chat.messages]);
+
+  const onScrollerScroll = () => {
+    const el = scrollerRef.current;
+    if (el) stickBottom.current = isAtBottom(el);
+  };
+
+  return (
+    <div className="flex h-full min-h-0 flex-col">
+      <header className="e0-chat-header">
+        <div className="min-w-0">
+          <div className="e0-chat-header-title">{thread.data.title}</div>
+          <div className="e0-chat-header-meta">
+            <span className={"size-2 rounded-full " + (st?.running ? "bg-ok" : "bg-muted")} aria-hidden />
+            {t("chat.serviceShell")}
+          </div>
+        </div>
+        <span className="hidden text-xs text-muted sm:inline">{model || t("chat.noModel")}</span>
+      </header>
+      <div ref={scrollerRef} onScroll={onScrollerScroll} className="min-h-0 flex-1 overflow-y-auto">
+        <div className="e0-chat-column space-y-1 px-5 pb-4 pt-7 sm:px-6">
+          <MessageList chat={chat} />
+        </div>
+      </div>
+      <ComposerBar chat={chat} model={model} onModelChange={onModelChange} />
+    </div>
+  );
+}
diff --git a/windows/app/src/routes/models.tsx b/windows/app/src/routes/models.tsx
new file mode 100644
index 0000000..5e75482
--- /dev/null
+++ b/windows/app/src/routes/models.tsx
@@ -0,0 +1,204 @@
+// routes/models.tsx — model catalog + download management. Tiers come from the embedded
+// catalog table; local facts = models_installed + engine status + folded download tasks.
+// The UI never estimates.
+import { BookOpen, Download, Power, RefreshCw, Trash2 } from "lucide-react";
+import { useState, useSyncExternalStore } from "react";
+
+import {
+  deleteModel,
+  downloadStore,
+  engineStore,
+  loadModel,
+  refreshAll,
+  shellStore,
+  startDownload,
+  unloadModel,
+} from "../api/client";
+import { deriveLight } from "../chat/modelStatus";
+import { humanBytes } from "../lib/byteFormat";
+import { codeOf, errKey, t } from "../lib/t";
+import type { CatTier } from "../api/client";
+import { DownloadCardBody } from "../models/cards";
+
+export function ModelsPage() {
+  const [, force] = useState(0);
+  useSyncExternalStore(shellStore.sig.sub, shellStore.sig.version);
+  useSyncExternalStore(downloadStore.sig.sub, downloadStore.sig.version);
+  useSyncExternalStore(engineStore.sig.sub, engineStore.sig.version);
+  void force;
+
+  const tiers = Object.entries(shellStore.catalog);
+  const ready = tiers.length > 0;
+
+  return (
+    <section className="e0-page space-y-4" data-testid="models-page">
+      <div className="e0-page-header">
+        <div>
+          <h2 className="e0-page-title flex items-center gap-2">
+            <BookOpen size={18} aria-hidden />
+            {t("models.title")}
+          </h2>
+          <p className="e0-page-kicker">{t("models.kicker")}</p>
+        </div>
+        <button type="button" data-testid="models-refresh" onClick={() => void refreshAll()} className="e0-btn e0-btn-secondary">
+          <RefreshCw size={12} aria-hidden />
+          {t("models.refresh")}
+        </button>
+      </div>
+
+      {!ready && <p className="text-sm text-muted">{t("models.loading")}</p>}
+
+      {tiers.map(([tier, rec]) => (
+        <TierCard key={tier} tier={tier} rec={rec} />
+      ))}
+    </section>
+  );
+}
+
+function TierCard({ tier, rec }: { tier: string; rec: CatTier }) {
+  const [source, setSource] = useState("");
+  const [opError, setOpError] = useState<string | null>(null);
+  const [confirmDelete, setConfirmDelete] = useState(false);
+
+  const st = engineStore.status;
+  const installed = shellStore.installed[tier] !== undefined;
+  const task = downloadStore.tasks[tier];
+  const inFlight = task && (task.phase === "probe" || task.phase === "download" || task.phase === "convert");
+  const light = deriveLight(st, engineStore.loadingTier, Object.keys(shellStore.installed), tier);
+  const stateLabel = installed ? "installed" : "absent";
+  const resident = light === "resident";
+  const filesCount = rec.files.filter((f) => !f.skip).length;
+
+  async function run(fn: () => Promise<unknown>) {
+    setOpError(null);
+    setConfirmDelete(false);
+    try {
+      await fn();
+      await refreshAll();
+    } catch (e) {
+      const s = String(e);
+      setOpError(`${t(errKey(codeOf(s)))}${codeOf(s) ? ` · ${codeOf(s)}` : ""}`);
+      await refreshAll().catch(() => undefined);
+    }
+  }
+
+  return (
+    <div data-testid={`tier-${tier}`} className="e0-section flex flex-col gap-3 p-5">
+      <div className="flex items-baseline justify-between gap-2">
+        <div>
+          <h3 className="text-base font-semibold">{tier}</h3>
+          <p className="text-xs text-muted" data-testid={`tier-state-${tier}`}>{t(`models.state.${stateLabel}`)}</p>
+        </div>
+        <span data-testid={`light-${tier}`} data-light={light} className={"e0-status " + (light === "resident" ? "e0-status-ok" : light === "loading" ? "e0-status-warn" : "")}>
+          <span className={"size-1.5 rounded-full " + (light === "resident" ? "bg-ok" : light === "loading" ? "bg-warn" : "bg-muted")} aria-hidden />
+          {t(`chat.light.${light}`)}
+        </span>
+      </div>
+
+      <dl className="order-3 grid grid-cols-2 gap-x-4 gap-y-1 border-t border-line/60 pt-3 text-xs tabular-nums">
+        <Field k={t("models.rev")} v={rec.rev.hf.slice(0, 12)} />
+        <Field k={t("models.bytes")} v={humanBytes(rec.total_bytes)} />
+        <Field k={t("models.files")} v={String(filesCount)} />
+        <Field k={t("models.pool")} v={`${rec.pool_mb} MB`} />
+        <Field k={t("models.source")} v="Hugging Face (primary) + mirrors" />
+        <Field k={t("models.repo")} v={rec.repo} />
+      </dl>
+
+      {inFlight && task && (
+        <div className="order-2">
+          <DownloadCardBody task={task} />
+        </div>
+      )}
+      {!inFlight && task && (task.phase === "error" || task.phase === "cancelled") && (
+        <div className="order-2">
+          <DownloadCardBody task={task} />
+        </div>
+      )}
+
+      {opError && (
+        <p data-testid={`tier-operr-${tier}`} role="alert" className="e0-status e0-status-err order-2 w-full justify-start rounded-lg">
+          {opError}
+        </p>
+      )}
+
+      {!installed && !inFlight && (
+        <div className="order-1 flex flex-wrap items-center gap-2">
+          <select
+            data-testid={`source-select-${tier}`}
+            value={source}
+            onChange={(e) => setSource(e.target.value)}
+            className="rounded-lg border border-line bg-surface px-2.5 py-1.5 text-xs"
+          >
+            <option value="">{t("models.sourceAuto")}</option>
+            <option value="hf">{t("models.sourceHf")}</option>
+            <option value="hf-mirror">{t("models.sourceHfMirror")}</option>
+            <option value="modelscope">{t("models.sourceMs")}</option>
+          </select>
+          <button
+            type="button"
+            data-testid={`download-${tier}`}
+            onClick={() => void run(() => startDownload(tier, source || null))}
+            className="e0-btn e0-btn-primary"
+          >
+            <Download size={12} aria-hidden />
+            {t("models.download")}
+          </button>
+        </div>
+      )}
+
+      {installed && (
+        <div className="order-1 flex flex-wrap items-center gap-2">
+          {!inFlight && light !== "resident" && (
+            <button
+              type="button"
+              data-testid={`load-${tier}`}
+              disabled={light === "loading"}
+              onClick={() => void run(() => loadModel(tier))}
+              className="e0-btn e0-btn-secondary"
+            >
+              <Power size={12} aria-hidden />
+              {light === "loading" ? t("models.loadingModel") : t("models.load")}
+            </button>
+          )}
+          {resident && (
+            <button type="button" data-testid={`unload-${tier}`} onClick={() => void run(() => unloadModel())} className="e0-btn e0-btn-secondary">
+              {t("models.unload")}
+            </button>
+          )}
+          {resident && <span className="text-xs text-muted">{t("chat.residentForever")}</span>}
+          {confirmDelete ? (
+            <span className="flex items-center gap-1 text-xs">
+              <span>{t("models.deleteConfirm")}</span>
+              <button type="button" data-testid={`delete-yes-${tier}`} onClick={() => void run(() => deleteModel(tier))} className="e0-btn bg-err text-white">
+                {t("models.delete")}
+              </button>
+              <button type="button" data-testid={`delete-no-${tier}`} onClick={() => setConfirmDelete(false)} className="e0-btn e0-btn-secondary">
+                {t("common.cancel")}
+              </button>
+            </span>
+          ) : (
+            <button
+              type="button"
+              data-testid={`delete-${tier}`}
+              disabled={!!inFlight}
+              onClick={() => setConfirmDelete(true)}
+              className="e0-btn e0-btn-danger"
+            >
+              <Trash2 size={12} aria-hidden />
+              {t("models.delete")}
+            </button>
+          )}
+        </div>
+      )}
+    </div>
+  );
+}
+
+function Field({ k, v }: { k: string; v: string }) {
+  return (
+    <>
+      <dt className="text-muted">{k}</dt>
+      <dd className="text-fg truncate" title={v}>{v}</dd>
+    </>
+  );
+}
diff --git a/windows/app/src/routes/service.tsx b/windows/app/src/routes/service.tsx
new file mode 100644
index 0000000..3456581
--- /dev/null
+++ b/windows/app/src/routes/service.tsx
@@ -0,0 +1,46 @@
+// routes/service.tsx — local API service screen with three landmarks: lifecycle /
+// request log / environment checks. Adds no new control paths or endpoints — a skin
+// over existing shell capabilities.
+import { Server } from "lucide-react";
+
+import { DoctorSection } from "../service/doctor";
+import { ServicePanel } from "../service/panel";
+import { RequestLogSection } from "../service/reqlog";
+import { t } from "../lib/t";
+
+export function ServicePage() {
+  return (
+    <div className="e0-page space-y-5" data-testid="service-page">
+      <header className="e0-page-header">
+        <div>
+          <h2 className="e0-page-title flex items-center gap-2">
+            <Server size={18} aria-hidden />
+            {t("service.title")}
+          </h2>
+          <p className="e0-page-kicker">{t("service.kicker")}</p>
+        </div>
+      </header>
+      <section aria-labelledby="service-lifecycle-title" data-testid="service-lifecycle-section">
+        <div className="mb-2 px-1">
+          <h3 id="service-lifecycle-title" className="e0-section-title">{t("service.sections.lifecycle")}</h3>
+          <p className="e0-section-note">{t("service.sections.lifecycleNote")}</p>
+        </div>
+        <ServicePanel />
+      </section>
+      <section aria-labelledby="service-log-title" data-testid="service-log-section">
+        <div className="mb-2 px-1">
+          <h3 id="service-log-title" className="e0-section-title">{t("service.sections.logs")}</h3>
+          <p className="e0-section-note">{t("service.sections.logsNote")}</p>
+        </div>
+        <RequestLogSection />
+      </section>
+      <section aria-labelledby="service-doctor-title" data-testid="service-doctor-section">
+        <div className="mb-2 px-1">
+          <h3 id="service-doctor-title" className="e0-section-title">{t("service.sections.doctor")}</h3>
+          <p className="e0-section-note">{t("service.sections.doctorNote")}</p>
+        </div>
+        <DoctorSection />
+      </section>
+    </div>
+  );
+}
diff --git a/windows/app/src/routes/settings.tsx b/windows/app/src/routes/settings.tsx
new file mode 100644
index 0000000..40ed212
--- /dev/null
+++ b/windows/app/src/routes/settings.tsx
@@ -0,0 +1,96 @@
+// routes/settings.tsx — settings page: appearance group reuses ShellControls, a system
+// prompt group, and a service block carrying the summary + console entry.
+import { Settings2, SlidersHorizontal, Server, RotateCcw } from "lucide-react";
+import { useState, useSyncExternalStore } from "react";
+
+import { engineStore } from "../api/client";
+import { ShellControls } from "../components/ShellControls";
+import { DEFAULT_SYSTEM_PROMPT, setSystemPrompt, storedSystemPrompt } from "../lib/systemPrompt";
+import { t } from "../lib/t";
+
+export function SettingsPage() {
+  const [systemPrompt, setPrompt] = useState(() => storedSystemPrompt() ?? "");
+  useSyncExternalStore(engineStore.sig.sub, engineStore.sig.version); // service-block pulse
+  const s = engineStore.status;
+  return (
+    <div className="e0-page space-y-5">
+      <header className="e0-page-header">
+        <div>
+          <h2 className="e0-page-title flex items-center gap-2">
+            <Settings2 size={19} aria-hidden />
+            {t("settings.title")}
+          </h2>
+          <p className="e0-page-kicker">{t("settings.kicker")}</p>
+        </div>
+      </header>
+      <section className="e0-section overflow-hidden" aria-labelledby="settings-appearance-title">
+        <div className="e0-section-header">
+          <div>
+            <h3 id="settings-appearance-title" className="e0-section-title">{t("settings.appearance")}</h3>
+            <p className="e0-section-note">{t("settings.appearanceNote")}</p>
+          </div>
+          <SlidersHorizontal size={17} className="text-muted" aria-hidden />
+        </div>
+        <div className="px-5 pb-5">
+          <ShellControls />
+        </div>
+      </section>
+      <section className="e0-section overflow-hidden" aria-labelledby="settings-prompt-title">
+        <div className="e0-section-header">
+          <div>
+            <h3 id="settings-prompt-title" className="e0-section-title">{t("settings.systemPrompt")}</h3>
+            <p className="e0-section-note">{t("settings.systemPromptNote")}</p>
+          </div>
+        </div>
+        <div className="border-t border-line/60 px-5 pb-5 pt-4">
+          <textarea
+            value={systemPrompt}
+            onChange={(e) => {
+              setPrompt(e.target.value);
+              setSystemPrompt(e.target.value);
+            }}
+            placeholder={DEFAULT_SYSTEM_PROMPT}
+            rows={4}
+            aria-label={t("settings.systemPrompt")}
+            className="w-full resize-y rounded-xl border border-line bg-surface2 p-3 text-sm leading-relaxed outline-none focus:border-accent"
+          />
+          <div className="mt-3 flex items-center justify-between gap-3">
+            <span className="text-xs text-muted">
+              {systemPrompt.trim() ? t("settings.systemPromptEnabled") : t("settings.systemPromptDisabled")}
+            </span>
+            <button
+              type="button"
+              className="e0-btn e0-btn-secondary"
+              onClick={() => {
+                setPrompt(DEFAULT_SYSTEM_PROMPT);
+                setSystemPrompt(DEFAULT_SYSTEM_PROMPT);
+              }}
+            >
+              <RotateCcw size={13} aria-hidden />
+              {t("settings.systemPromptReset")}
+            </button>
+          </div>
+        </div>
+      </section>
+      <section className="e0-section overflow-hidden" aria-labelledby="settings-service-title">
+        <div className="e0-section-header">
+          <div>
+            <h3 id="settings-service-title" className="e0-section-title">{t("settings.service")}</h3>
+            <p className="e0-section-note">{t("settings.serviceNote")}</p>
+          </div>
+          <Server size={17} className="text-muted" aria-hidden />
+        </div>
+        <div className="flex items-center gap-3 border-t border-line/60 px-5 py-4">
+          <span className={"size-2 rounded-full " + (s?.running ? "bg-ok" : "bg-muted")} aria-hidden />
+          <span className="min-w-0 flex-1 truncate text-sm text-muted">
+            {s?.running ? `${s.bound}:${s.port} · ${s.tier}` : t("settings.serviceUnavailable")}
+          </span>
+          <a href="#/service" data-testid="service-console-link" className="e0-btn e0-btn-secondary shrink-0">
+            {t("settings.openConsole")}
+          </a>
+        </div>
+      </section>
+      <p className="px-1 text-xs text-muted">{t("settings.advancedNote")}</p>
+    </div>
+  );
+}
diff --git a/windows/app/src/routes/welcome.tsx b/windows/app/src/routes/welcome.tsx
new file mode 100644
index 0000000..e32ae9e
--- /dev/null
+++ b/windows/app/src/routes/welcome.tsx
@@ -0,0 +1,111 @@
+// routes/welcome.tsx — first-run guide card (shown in the chat area when nothing is
+// installed; completes the loop within three clicks). "Download & enable" = one explicit
+// intent: create the task + arm the one-shot token (auto-load on completion lives in the
+// AppShell chain).
+import { Sparkles, Download } from "lucide-react";
+import { useState, useSyncExternalStore } from "react";
+
+import { downloadStore, shellStore, startDownload } from "../api/client";
+import { autoLoad } from "../lib/autoLoad";
+import { humanBytes } from "../lib/byteFormat";
+import { codeOf, errKey, t } from "../lib/t";
+import { DownloadCardBody } from "../models/cards";
+
+export function WelcomeCard() {
+  useSyncExternalStore(downloadStore.sig.sub, downloadStore.sig.version);
+  useSyncExternalStore(shellStore.sig.sub, shellStore.sig.version);
+
+  const [dismissed, setDismissed] = useState(false);
+  const [err, setErr] = useState<string | null>(null);
+  const [busy, setBusy] = useState(false);
+
+  const cat = shellStore.catalog;
+  const tier = cat["8b"] ? "8b" : Object.keys(cat)[0] ?? null;
+  const rec = tier ? cat[tier] : null;
+  const task = tier ? downloadStore.tasks[tier] : undefined;
+  const path = shellStore.paths.home || null;
+
+  if (dismissed) {
+    return (
+      <div className="e0-chat-empty">
+        <p className="text-lg font-medium">{t("chat.noInstalledModel")}</p>
+        <a
+          href="#/models"
+          className="justify-self-center rounded-full border border-line px-4 py-1.5 text-sm text-muted hover:bg-surface2 hover:text-fg"
+        >
+          {t("chat.goDownload")}
+        </a>
+      </div>
+    );
+  }
+
+  async function downloadAndEnable() {
+    if (!tier) return;
+    setErr(null);
+    setBusy(true);
+    try {
+      await startDownload(tier, null);
+      autoLoad.begin(tier, tier); // task key = tier (shell pull.rs)
+    } catch (e) {
+      setErr(String(e));
+    } finally {
+      setBusy(false);
+    }
+  }
+
+  return (
+    <div className="e0-chat-empty px-6" data-testid="welcome-card">
+      <div className="e0-chat-column mx-auto space-y-5 text-center">
+        <span className="mx-auto grid size-12 place-items-center rounded-2xl bg-fg text-surface shadow-sm">
+          <Sparkles size={20} aria-hidden />
+        </span>
+        <div>
+          <h2 className="text-3xl font-semibold tracking-tight">{t("welcome.title")}</h2>
+          <p className="mx-auto mt-2 max-w-lg text-sm leading-relaxed text-muted">
+            {t("welcome.body", {
+              tier: tier ?? "8b",
+              bytes: rec ? humanBytes(rec.total_bytes) : t("welcome.bytesUnknown"),
+              path: path ?? t("welcome.pathUnknown"),
+            })}
+          </p>
+        </div>
+        {task && task.phase !== "idle" ? (
+          <div className="mx-auto w-full max-w-lg text-left">
+            <DownloadCardBody task={task} />
+          </div>
+        ) : (
+          <button
+            type="button"
+            data-testid="welcome-go"
+            disabled={busy || !rec}
+            onClick={() => void downloadAndEnable()}
+            className="e0-btn e0-btn-primary mx-auto min-h-10 px-5 text-sm"
+          >
+            <Download size={15} aria-hidden />
+            {t("welcome.downloadEnable")}
+          </button>
+        )}
+        {err && (
+          <p data-testid="welcome-err" role="alert" className="e0-status e0-status-err mx-auto max-w-lg justify-start rounded-xl px-3 py-2 text-left">
+            <span>
+              <span className="block font-semibold">{t("welcome.errors.generic")}</span>
+              <span className="mt-0.5 block font-normal">
+                {t(errKey(codeOf(err)))} · {codeOf(err) ?? ""}
+              </span>
+            </span>
+          </p>
+        )}
+        <div>
+          <button
+            type="button"
+            data-testid="welcome-skip"
+            onClick={() => setDismissed(true)}
+            className="text-xs text-muted underline hover:text-fg"
+          >
+            {t("welcome.skip")}
+          </button>
+        </div>
+      </div>
+    </div>
+  );
+}
diff --git a/windows/app/src/service/doctor.tsx b/windows/app/src/service/doctor.tsx
new file mode 100644
index 0000000..0f235b4
--- /dev/null
+++ b/windows/app/src/service/doctor.tsx
@@ -0,0 +1,99 @@
+// service/doctor.tsx — environment-check skin. Discipline: all verdicts are shell-side
+// measurements (doctor_run); the UI only maps strings to styles — no thresholds, no
+// re-judgement; on failure it keeps no stale data and never fakes a conclusion.
+import { Stethoscope } from "lucide-react";
+import { useCallback, useEffect, useState } from "react";
+
+import { doctorRun, type DoctorCheck } from "../api/client";
+import { t } from "../lib/t";
+
+interface DoctorReport {
+  overall: string;
+  checks: DoctorCheck[];
+}
+
+/** Whitelisted `code` -> action routing (the skin only navigates, never judges). */
+const ACTION_ROUTES: Record<string, { to: string; labelKey: string }> = {
+  "E-DL-DISK": { to: "/models", labelKey: "service.doctor.goModels" },
+  "E-MODEL-INVALID": { to: "/models", labelKey: "service.doctor.goModels" },
+  "E-MODEL-MISSING": { to: "/models", labelKey: "service.doctor.goModels" },
+};
+
+export function DoctorSection() {
+  const [report, setReport] = useState<DoctorReport | null>(null);
+  const [failed, setFailed] = useState(false);
+  const [loading, setLoading] = useState(true);
+
+  const fetchDoctor = useCallback(async () => {
+    setLoading(true);
+    setFailed(false);
+    try {
+      setReport(await doctorRun());
+    } catch {
+      setReport(null); // failure state must not keep stale data
+      setFailed(true);
+    } finally {
+      setLoading(false);
+    }
+  }, []);
+
+  useEffect(() => {
+    void fetchDoctor();
+  }, [fetchDoctor]);
+
+  const verdictTone = (v: string) => (v === "pass" ? "ok" : v === "warn" ? "warn" : "err");
+  const cls = (tone: string) =>
+    tone === "ok" ? "bg-accent/15 text-accent" : tone === "warn" ? "bg-warn/15 text-warn" : "bg-err/15 text-err";
+
+  return (
+    <section className="e0-section overflow-hidden p-5" data-testid="svc-doctor">
+      <div className="mb-2 flex items-center justify-between">
+        <h3 className="flex items-center gap-2 text-sm font-semibold">
+          <Stethoscope size={14} aria-hidden />
+          {t("service.doctor.title")}
+        </h3>
+        {report && (
+          <span data-testid="doctor-overall" className={"rounded-full px-2 py-0.5 text-xs " + cls(verdictTone(report.overall))}>
+            {t(`service.doctor.overall.${report.overall}`)}
+          </span>
+        )}
+      </div>
+
+      {loading && !report && <p className="text-xs text-muted">{t("service.doctor.loading")}</p>}
+
+      {failed && (
+        <p data-testid="doctor-failed" className="text-xs text-muted">
+          {t("service.doctor.failed")}
+          <button type="button" data-testid="doctor-retry" onClick={() => void fetchDoctor()} className="ml-2 underline">
+            {t("service.doctor.retry")}
+          </button>
+        </p>
+      )}
+
+      {report && (
+        <ul className="space-y-1">
+          {report.checks.map((c) => {
+            const action = c.code ? ACTION_ROUTES[c.code] : undefined;
+            return (
+              <li key={c.id} data-testid={`doctor-item-${c.id}`} className="flex items-baseline gap-2 border-t border-line/60 py-1.5 text-xs first:border-t-0">
+                <span className={"w-14 shrink-0 rounded-full px-1.5 py-0.5 text-center " + cls(verdictTone(c.verdict))}>
+                  {t(`service.doctor.verdict.${c.verdict}`)}
+                </span>
+                <span data-testid={`doctor-detail-${c.id}`} className="min-w-0 flex-1">
+                  <span className="text-muted">{c.id}</span> · {c.detail}
+                  {c.code && <span className="text-muted"> · {c.code}</span>}
+                  {c.next && <span className="block text-muted">{t("service.doctor.next")}: {c.next}</span>}
+                </span>
+                {action && (
+                  <a href={`#${action.to}`} data-testid={`doctor-action-${c.id}`} className="e0-btn e0-btn-secondary shrink-0">
+                    {t(action.labelKey)}
+                  </a>
+                )}
+              </li>
+            );
+          })}
+        </ul>
+      )}
+    </section>
+  );
+}
diff --git a/windows/app/src/service/panel.tsx b/windows/app/src/service/panel.tsx
new file mode 100644
index 0000000..4f6eeb4
--- /dev/null
+++ b/windows/app/src/service/panel.tsx
@@ -0,0 +1,112 @@
+// service/panel.tsx — lifecycle panel. Skin boundaries: forwards existing commands only,
+// reads measured fields only. This shell has no daemon token: loopback is access-exempt
+// and stated as such. LAN = phase-2 placeholder row (kept visible, disabled + honest
+// title).
+import { useState, useSyncExternalStore } from "react";
+
+import {
+  engineStore,
+  installedTiers,
+  loadModel,
+  shellStore,
+  unloadModel,
+} from "../api/client";
+import { codeOf, errKey, t } from "../lib/t";
+
+export function ServicePanel() {
+  useSyncExternalStore(engineStore.sig.sub, engineStore.sig.version);
+  useSyncExternalStore(shellStore.sig.sub, shellStore.sig.version);
+  const [busy, setBusy] = useState<string | null>(null);
+  const [error, setError] = useState<string | null>(null);
+
+  const st = engineStore.status;
+  const running = !!st?.running;
+  const loading = engineStore.loadingTier !== null;
+  const tiers = installedTiers();
+  const pick = tiers.includes("8b") ? "8b" : tiers[0] ?? null;
+
+  async function act(name: string, fn: () => Promise<unknown>) {
+    setBusy(name);
+    setError(null);
+    try {
+      await fn();
+    } catch (e) {
+      const s = String(e);
+      setError(`${t(errKey(codeOf(s)))}${codeOf(s) ? ` · ${codeOf(s)}` : ""}`);
+    } finally {
+      setBusy(null);
+    }
+  }
+
+  const inFlight = busy !== null || loading;
+
+  return (
+    <section className="e0-section overflow-hidden p-5" data-testid="svc-panel">
+      <div className="mb-2 flex items-center justify-between">
+        <h3 className="text-sm font-semibold">{t("service.panel.title")}</h3>
+        <div className="flex items-center gap-2">
+          {running ? (
+            <button type="button" data-testid="svc-stop" disabled={inFlight} onClick={() => void act("stop", () => unloadModel())} className="e0-btn e0-btn-secondary">
+              {t("service.panel.stop")}
+            </button>
+          ) : (
+            <button
+              type="button"
+              data-testid="svc-start"
+              disabled={inFlight || !pick}
+              title={pick ? undefined : t("chat.noInstalledModel")}
+              onClick={() => pick && void act("start", () => loadModel(pick))}
+              className="rounded-full bg-accent px-3 py-1 text-xs font-medium text-accentink hover:opacity-90 disabled:opacity-50"
+            >
+              {loading ? t("models.loadingModel") : t("service.panel.start")}
+            </button>
+          )}
+        </div>
+      </div>
+
+      {inFlight && (
+        <p data-testid="svc-busy" className="mb-2 text-xs text-muted">{t("service.panel.busy")}</p>
+      )}
+      {error && (
+        <p data-testid="svc-error" role="alert" className="mb-2 rounded-lg bg-err/10 px-3 py-2 text-xs text-err">{error}</p>
+      )}
+
+      <dl className="grid grid-cols-[auto_1fr] gap-x-4 gap-y-1 text-xs" data-testid="svc-fields">
+        <dt className="text-muted">{t("service.panel.port")}</dt>
+        <dd data-testid="svc-port" className="tabular-nums">{st?.port ? `${st.bound}:${st.port}` : "—"}</dd>
+        <dt className="text-muted">{t("service.panel.bind")}</dt>
+        <dd data-testid="svc-bind">{st?.bound ?? "—"}</dd>
+        <dt className="text-muted">{t("service.panel.scope")}</dt>
+        <dd data-testid="svc-scope">{running ? t("sidebar.scopeShell") : "—"}</dd>
+        <dt className="text-muted">{t("service.panel.tier")}</dt>
+        <dd data-testid="svc-tier">{running ? st?.tier ?? "—" : "—"}</dd>
+        <dt className="text-muted">{t("service.panel.pool")}</dt>
+        <dd data-testid="svc-pool" className="tabular-nums">{running && st?.pool_mb ? `${st.pool_mb} MB` : "—"}</dd>
+        <dt className="text-muted">{t("service.panel.uptime")}</dt>
+        <dd data-testid="svc-uptime" className="tabular-nums">{running ? `${st?.uptime_s ?? 0}s` : "—"}</dd>
+        <dt className="text-muted">{t("service.panel.version")}</dt>
+        <dd data-testid="svc-version">{st?.version ?? "—"}</dd>
+        <dt className="text-muted">{t("service.panel.governance")}</dt>
+        <dd data-testid="svc-governance" className={running && !st?.job ? "text-warn" : ""}>
+          {running ? (st?.job ? t("service.panel.governanceJob") : t("service.panel.governanceSoft")) : "—"}
+        </dd>
+        <dt className="text-muted">{t("service.panel.home")}</dt>
+        <dd data-testid="svc-home" className="break-all">{shellStore.paths.home || "—"}</dd>
+      </dl>
+
+      {/* token row: loopback direct access needs no credential — the fact, stated plainly (no mask, no copy button) */}
+      <div className="mt-3 flex items-center gap-2 text-xs" data-testid="svc-token-row">
+        <span className="text-muted">{t("service.panel.token")}</span>
+        <span data-testid="svc-token-none">{t("service.panel.noCredential")}</span>
+      </div>
+
+      {/* LAN toggle: phase-2 placeholder — row kept visible, disabled + honest title */}
+      <div className="mt-3 flex items-center justify-between border-t border-line pt-3 text-xs" data-testid="svc-lan-row">
+        <span className="text-muted">{t("service.panel.lan")}</span>
+        <button type="button" role="switch" aria-checked={false} data-testid="svc-lan-switch" disabled title={t("service.panel.lanLater")} className="e0-btn e0-btn-secondary disabled:opacity-50">
+          {t("service.panel.lanOff")}
+        </button>
+      </div>
+    </section>
+  );
+}
diff --git a/windows/app/src/service/reqlog.tsx b/windows/app/src/service/reqlog.tsx
new file mode 100644
index 0000000..d23bd15
--- /dev/null
+++ b/windows/app/src/service/reqlog.tsx
@@ -0,0 +1,70 @@
+// service/reqlog.tsx — request log table. Data source = frontend-measured generation
+// requests (useThreadChat writes into the reqLog ring): since app start, in memory,
+// cleared on exit — labeled honestly in the section header.
+import { useSyncExternalStore } from "react";
+
+import { reqLog, type ReqRow } from "../api/client";
+import { t } from "../lib/t";
+
+export function statusTone(status: ReqRow["status"]): "ok" | "warn" | "err" {
+  if (status === "done") return "ok";
+  if (status === "error") return "err";
+  return "warn";
+}
+
+export function RequestLogSection() {
+  useSyncExternalStore(reqLog.sig.sub, reqLog.sig.version);
+  const rows = reqLog.rows;
+  return (
+    <section className="e0-section overflow-hidden p-5" data-testid="svc-log">
+      <div className="mb-1 flex items-baseline justify-between">
+        <h3 className="text-sm font-semibold">{t("service.log.title")}</h3>
+        <span data-testid="svc-log-count" className="text-xs text-muted tabular-nums">
+          {t("service.log.count", { n: rows.length })}
+        </span>
+      </div>
+      <p data-testid="svc-log-note" className="mb-2 text-xs text-muted">{t("service.log.note")}</p>
+      {rows.length === 0 ? (
+        <p data-testid="svc-log-empty" className="text-xs text-muted">{t("service.log.empty")}</p>
+      ) : (
+        <table className="w-full text-xs">
+          <thead>
+            <tr className="text-left text-muted">
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.status")}</th>
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.model")}</th>
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.latency")}</th>
+              <th className="py-1 pr-2 font-normal">{t("service.log.th.tokens")}</th>
+              <th className="py-1 font-normal">{t("service.log.th.source")}</th>
+            </tr>
+          </thead>
+          <tbody>
+            {rows.map((r, i) => {
+              const tone = statusTone(r.status);
+              return (
+                <tr key={`${r.ts}-${i}`} data-testid={`svc-log-row-${i}`} className="border-t border-line/60">
+                  <td className="py-1 pr-2">
+                    <span className="flex items-center gap-1">
+                      <span
+                        className={"inline-block size-2 rounded-full " + (tone === "ok" ? "bg-accent" : tone === "err" ? "bg-err" : "bg-warn")}
+                        aria-hidden
+                      />
+                      {t(`service.log.status.${r.status}`)}
+                      {r.code && <span className="text-muted">{r.code}</span>}
+                    </span>
+                    <span className="ml-1 text-[10px] text-muted tabular-nums">
+                      {new Date(r.ts).toLocaleTimeString([], { hour: "2-digit", minute: "2-digit", second: "2-digit" })}
+                    </span>
+                  </td>
+                  <td className="py-1 pr-2">{r.model || "—"}</td>
+                  <td className="py-1 pr-2 tabular-nums">{r.latencyMs}ms</td>
+                  <td className="py-1 pr-2 tabular-nums">{r.tokens ?? "—"}</td>
+                  <td className="py-1">{t("service.log.source.loopback")}</td>
+                </tr>
+              );
+            })}
+          </tbody>
+        </table>
+      )}
+    </section>
+  );
+}
diff --git a/windows/app/src/state/threads.ts b/windows/app/src/state/threads.ts
new file mode 100644
index 0000000..8625e9d
--- /dev/null
+++ b/windows/app/src/state/threads.ts
@@ -0,0 +1,144 @@
+// state/threads.ts — localStorage-backed thread store (message/row shape mirrors a
+// typical chat thread store). A SQLite thread store would be a phase-2 drop-in:
+// this layer can be replaced wholesale without touching the UI.
+export type MsgUsage = { prompt_tokens: number; completion_tokens: number; total_tokens: number };
+export type MsgPerf = { decode_tps?: number; prefill_tps?: number; prompt_ms?: number };
+export type MsgMeta = {
+  status?: "done" | "aborted" | "error";
+  usage?: MsgUsage; perf?: MsgPerf; code?: string | null;
+};
+export type ThreadMsg = {
+  id: string; role: "user" | "assistant"; content: string;
+  reasoning?: string; created_at: number; meta?: MsgMeta;
+  streaming?: boolean; // in-session transient (stripped before persisting; always false on load)
+};
+export type ThreadRow = {
+  id: string;
+  data: { title: string; model: string | null };
+  updatedAt: number;
+  messages: ThreadMsg[];
+};
+
+const STORE_KEY = "edge0.threads.v2";
+const ACTIVE_KEY = "edge0.active-thread";
+
+function uid(): string {
+  return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}`;
+}
+
+class ThreadStore {
+  private rows: ThreadRow[] = [];
+  private ver = 0;
+  activeId: string | null = null;
+  private listeners = new Set<() => void>();
+  private persistTimer: ReturnType<typeof setTimeout> | null = null;
+  private loaded = false;
+
+  load(): void {
+    if (this.loaded) return;
+    this.loaded = true;
+    try {
+      this.rows = (JSON.parse(localStorage.getItem(STORE_KEY) ?? "[]") as ThreadRow[])
+        .map((r) => ({
+          ...r,
+          // converge crash residue: a trailing assistant message without terminal meta means generation was interrupted -> mark aborted
+          messages: (() => {
+            const ms = r.messages.map((m) => ({ ...m, streaming: false }));
+            const last = ms[ms.length - 1];
+            if (last && last.role === "assistant" && !last.meta?.status) {
+              ms[ms.length - 1] = { ...last, meta: { ...last.meta, status: "aborted" } };
+            }
+            return ms;
+          })(),
+        }));
+    } catch {
+      this.rows = []; // drop a corrupt cache (never crash startup)
+    }
+    const a = localStorage.getItem(ACTIVE_KEY);
+    this.activeId = a && this.rows.some((r) => r.id === a) ? a : null;
+  }
+
+  subscribe = (cb: () => void): (() => void) => {
+    this.listeners.add(cb);
+    return () => { this.listeners.delete(cb); };
+  };
+
+  /** snapshot = version number (stable source for useSyncExternalStore; get rows via snapshot()). */
+  version(): number {
+    return this.ver;
+  }
+
+  threads(): ThreadRow[] {
+    return this.rows;
+  }
+
+  private emit(): void {
+    this.ver++;
+    for (const l of [...this.listeners]) l();
+    if (this.persistTimer) clearTimeout(this.persistTimer);
+    this.persistTimer = setTimeout(() => {
+      try {
+        localStorage.setItem(STORE_KEY, JSON.stringify(
+          this.rows.map((r) => ({ ...r, messages: r.messages.map((m) => {
+            const { streaming: _s, ...rest } = m;
+            return rest;
+          }) })),
+        ));
+      } catch { /* quota exceeded: skip the write (UI must not crash) */ }
+    }, 400);
+  }
+
+  /** refresh() compat shim kept from the original API shape (this implementation is always in sync). */
+  async refresh(): Promise<void> {
+    this.load();
+    this.emit();
+  }
+
+  create(title: string, model: string | null): ThreadRow {
+    this.load();
+    const row: ThreadRow = { id: uid(), data: { title, model }, updatedAt: Date.now(), messages: [] };
+    this.rows = [row, ...this.rows];
+    this.activeId = row.id;
+    localStorage.setItem(ACTIVE_KEY, row.id);
+    this.emit();
+    return row;
+  }
+
+  rename(id: string, title: string): void {
+    this.rows = this.rows.map((r) => (r.id === id ? { ...r, data: { ...r.data, title }, updatedAt: Date.now() } : r));
+    this.emit();
+  }
+
+  remove(id: string): void {
+    this.rows = this.rows.filter((r) => r.id !== id);
+    if (this.activeId === id) this.setActive(null);
+    this.emit();
+  }
+
+  setActive(id: string | null): void {
+    this.activeId = id;
+    if (id) localStorage.setItem(ACTIVE_KEY, id);
+    else localStorage.removeItem(ACTIVE_KEY);
+    this.emit();
+  }
+
+  setModel(id: string, model: string): void {
+    this.rows = this.rows.map((r) => (r.id === id ? { ...r, data: { ...r.data, model } } : r));
+    this.emit();
+  }
+
+  messages(id: string): ThreadMsg[] {
+    return this.rows.find((r) => r.id === id)?.messages ?? [];
+  }
+
+  /** Single entry point for message-list mutations (updatedAt written along). */
+  patchMessages(id: string, fn: (msgs: ThreadMsg[]) => ThreadMsg[]): void {
+    this.rows = this.rows.map((r) =>
+      r.id === id ? { ...r, updatedAt: Date.now(), messages: fn(r.messages) } : r,
+    );
+    this.emit();
+  }
+}
+
+export const threadStore = new ThreadStore();
+threadStore.load();
diff --git a/windows/app/src/styles.css b/windows/app/src/styles.css
new file mode 100644
index 0000000..06d8b35
--- /dev/null
+++ b/windows/app/src/styles.css
@@ -0,0 +1,485 @@
+@import "tailwindcss";
+@plugin "@tailwindcss/typography";
+
+/* Dark mode does not use prefers-color-scheme: the shell is the single source of truth
+   (in "follow system" mode, JS resolves it and writes data-theme). */
+@custom-variant dark (&:where([data-theme="dark"], [data-theme="dark"] *));
+
+:root {
+  --e0-app: #f3f4f6;
+  --e0-surface: #ffffff;
+  --e0-surface-2: #f7f7f8;
+  --e0-sidebar: #eceef1;
+  --e0-text: #181a1f;
+  --e0-muted: #70747d;
+  --e0-border: #dfe2e7;
+  --e0-accent: #2f6feb;
+  --e0-accent-ink: #ffffff;
+  --e0-ok: #168052;
+  --e0-warn: #a66700;
+  --e0-err: #c53636;
+  --e0-focus: #2f6feb;
+  --e0-shadow: 0 12px 36px rgba(29, 34, 43, 0.08);
+  --e0-radius-lg: 18px;
+  --e0-radius-md: 12px;
+}
+[data-theme="dark"] {
+  --e0-app: #111317;
+  --e0-surface: #181b20;
+  --e0-surface-2: #20242a;
+  --e0-sidebar: #1b1e24;
+  --e0-text: #f1f3f5;
+  --e0-muted: #a2a8b2;
+  --e0-border: #30353d;
+  --e0-accent: #6ea3ff;
+  --e0-accent-ink: #111317;
+  --e0-ok: #55d69a;
+  --e0-warn: #efb456;
+  --e0-err: #ff8585;
+  --e0-focus: #8bb5ff;
+  --e0-shadow: 0 16px 42px rgba(0, 0, 0, 0.28);
+}
+
+/* Three font-size tiers: the whole tree scales via rem. */
+html[data-font="sm"] { font-size: 14px; }
+html[data-font="md"] { font-size: 16px; }
+html[data-font="lg"] { font-size: 18px; }
+
+@theme inline {
+  --color-app: var(--e0-app);
+  --color-surface: var(--e0-surface);
+  --color-surface2: var(--e0-surface-2);
+  --color-sidebar: var(--e0-sidebar);
+  --color-fg: var(--e0-text);
+  --color-muted: var(--e0-muted);
+  --color-line: var(--e0-border);
+  --color-accent: var(--e0-accent);
+  --color-accentink: var(--e0-accent-ink);
+  --color-ok: var(--e0-ok);
+  --color-warn: var(--e0-warn);
+  --color-err: var(--e0-err);
+}
+
+body {
+  @apply bg-app text-fg antialiased;
+  margin: 0;
+  min-width: 320px;
+  font-family: ui-sans-serif, system-ui, "Segoe UI", "Microsoft YaHei", sans-serif;
+}
+
+button,
+input,
+textarea,
+select {
+  font: inherit;
+}
+
+button,
+a,
+input,
+textarea,
+select {
+  -webkit-tap-highlight-color: transparent;
+}
+
+button:focus-visible,
+a:focus-visible,
+input:focus-visible,
+textarea:focus-visible,
+select:focus-visible {
+  outline: 2px solid var(--e0-focus);
+  outline-offset: 2px;
+}
+
+/* The composer focus state is expressed by the shell's focus-within; suppress the inner
+   textarea's square outline. Kept unlayered: unlayered author styles beat Tailwind utilities. */
+.e0-composer-shell textarea:focus-visible {
+  outline: none;
+}
+
+@layer components {
+  .e0-app-shell {
+    display: flex;
+    gap: 12px;
+    height: 100vh;
+    min-height: 0;
+    max-height: 100vh;
+    padding: 12px;
+    background: var(--e0-app);
+    box-sizing: border-box;
+    overflow: hidden;
+  }
+
+  .e0-sidebar {
+    position: relative;
+    z-index: 30;
+    display: flex;
+    width: 258px;
+    height: 100%;
+    flex: 0 0 258px;
+    flex-direction: column;
+    overflow: hidden;
+    border: 1px solid var(--e0-border);
+    border-radius: var(--e0-radius-lg);
+    background: var(--e0-sidebar);
+    box-shadow: var(--e0-shadow);
+  }
+
+  .e0-main-frame {
+    position: relative;
+    height: 100%;
+    min-width: 0;
+    min-height: 0;
+    flex: 1;
+    overflow: hidden;
+    border: 1px solid var(--e0-border);
+    border-radius: var(--e0-radius-lg);
+    background: var(--e0-surface);
+    box-shadow: var(--e0-shadow);
+  }
+
+  .e0-page {
+    width: min(100%, 980px);
+    margin: 0 auto;
+    padding: 34px clamp(20px, 4vw, 56px) 44px;
+  }
+
+  .e0-page-header {
+    display: flex;
+    align-items: flex-start;
+    justify-content: space-between;
+    gap: 16px;
+    margin-bottom: 28px;
+  }
+
+  .e0-page-title {
+    margin: 0;
+    font-size: 1.35rem;
+    font-weight: 650;
+    letter-spacing: -0.01em;
+  }
+
+  .e0-page-kicker {
+    margin: 5px 0 0;
+    color: var(--e0-muted);
+    font-size: 0.8rem;
+    line-height: 1.5;
+  }
+
+  .e0-section {
+    border: 1px solid var(--e0-border);
+    border-radius: var(--e0-radius-md);
+    background: var(--e0-surface);
+    box-shadow: 0 4px 16px rgba(29, 34, 43, 0.04);
+  }
+
+  .e0-section-header {
+    display: flex;
+    align-items: flex-start;
+    justify-content: space-between;
+    gap: 16px;
+    padding: 18px 20px 12px;
+  }
+
+  .e0-section-title {
+    margin: 0;
+    font-size: 0.9rem;
+    font-weight: 650;
+  }
+
+  .e0-section-note {
+    margin: 4px 0 0;
+    color: var(--e0-muted);
+    font-size: 0.75rem;
+    line-height: 1.5;
+  }
+
+  .e0-btn {
+    display: inline-flex;
+    min-height: 32px;
+    align-items: center;
+    justify-content: center;
+    gap: 7px;
+    border: 1px solid transparent;
+    border-radius: 999px;
+    padding: 0 13px;
+    font-size: 0.75rem;
+    font-weight: 600;
+    transition: background-color 140ms ease, border-color 140ms ease, color 140ms ease, opacity 140ms ease;
+  }
+
+  .e0-btn:disabled {
+    cursor: not-allowed;
+    opacity: 0.48;
+  }
+
+  .e0-btn-primary {
+    background: var(--e0-accent);
+    color: var(--e0-accent-ink);
+  }
+
+  .e0-btn-primary:hover:not(:disabled) {
+    filter: brightness(0.96);
+  }
+
+  .e0-btn-secondary {
+    border-color: var(--e0-border);
+    background: var(--e0-surface);
+    color: var(--e0-text);
+  }
+
+  .e0-btn-secondary:hover:not(:disabled),
+  .e0-btn-ghost:hover:not(:disabled) {
+    background: var(--e0-surface-2);
+  }
+
+  .e0-btn-ghost {
+    color: var(--e0-muted);
+  }
+
+  .e0-btn-danger {
+    border-color: color-mix(in srgb, var(--e0-err) 35%, var(--e0-border));
+    color: var(--e0-err);
+  }
+
+  .e0-icon-btn {
+    display: inline-grid;
+    width: 32px;
+    height: 32px;
+    place-items: center;
+    border-radius: 10px;
+    color: var(--e0-muted);
+    transition: background-color 140ms ease, color 140ms ease;
+  }
+
+  .e0-icon-btn:hover {
+    background: color-mix(in srgb, var(--e0-text) 8%, transparent);
+    color: var(--e0-text);
+  }
+
+  .e0-status {
+    display: inline-flex;
+    min-height: 24px;
+    align-items: center;
+    gap: 6px;
+    border-radius: 999px;
+    padding: 0 9px;
+    font-size: 0.7rem;
+    font-weight: 600;
+  }
+
+  .e0-status-ok { background: color-mix(in srgb, var(--e0-ok) 13%, transparent); color: var(--e0-ok); }
+  .e0-status-warn { background: color-mix(in srgb, var(--e0-warn) 14%, transparent); color: var(--e0-warn); }
+  .e0-status-err { background: color-mix(in srgb, var(--e0-err) 13%, transparent); color: var(--e0-err); }
+
+  .e0-toggle-group {
+    display: inline-flex;
+    gap: 2px;
+    border: 1px solid var(--e0-border);
+    border-radius: 999px;
+    padding: 3px;
+    background: var(--e0-surface-2);
+  }
+
+  .e0-toggle {
+    border-radius: 999px;
+    padding: 5px 9px;
+    color: var(--e0-muted);
+    font-size: 0.7rem;
+    font-weight: 600;
+  }
+
+  .e0-toggle[data-state="on"] {
+    background: var(--e0-accent);
+    color: var(--e0-accent-ink);
+  }
+
+  .e0-divider {
+    border-top: 1px solid color-mix(in srgb, var(--e0-border) 72%, transparent);
+  }
+
+  .e0-mobile-toggle {
+    display: none;
+  }
+
+  .e0-chat-header {
+    display: flex;
+    min-height: 58px;
+    align-items: center;
+    justify-content: space-between;
+    gap: 16px;
+    border-bottom: 1px solid color-mix(in srgb, var(--e0-border) 72%, transparent);
+    padding: 0 28px;
+  }
+
+  .e0-chat-header-title {
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+    font-size: 0.8rem;
+    font-weight: 650;
+  }
+
+  .e0-chat-header-meta {
+    display: inline-flex;
+    align-items: center;
+    gap: 7px;
+    color: var(--e0-muted);
+    font-size: 0.7rem;
+  }
+
+  .e0-chat-empty {
+    display: grid;
+    min-height: 100%;
+    place-content: center;
+    gap: 8px;
+    padding: 32px;
+    text-align: center;
+  }
+
+  .e0-chat-column {
+    width: min(100%, 760px);
+    margin: 0 auto;
+  }
+
+  .e0-composer {
+    padding: 12px 22px 20px;
+  }
+
+  .e0-composer-shell {
+    width: min(100%, 760px);
+    margin: 0 auto;
+    overflow: visible;
+    border: 1px solid var(--e0-border);
+    border-radius: 20px;
+    background: var(--e0-surface);
+    box-shadow: 0 10px 28px rgba(29, 34, 43, 0.09);
+    transition: border-color 140ms ease, box-shadow 140ms ease;
+  }
+
+  .e0-composer-shell:focus-within {
+    border-color: color-mix(in srgb, var(--e0-accent) 55%, var(--e0-border));
+    box-shadow: 0 10px 30px color-mix(in srgb, var(--e0-accent) 13%, transparent);
+  }
+
+  .e0-composer-textarea {
+    min-height: 82px;
+    max-height: 220px;
+    padding: 17px 18px 8px;
+    line-height: 1.5;
+  }
+
+  .e0-composer-toolbar {
+    display: flex;
+    min-width: 0;
+    align-items: center;
+    gap: 8px;
+    padding: 6px 10px 10px;
+  }
+
+  .e0-send-btn {
+    display: inline-grid;
+    width: 34px;
+    height: 34px;
+    flex: 0 0 auto;
+    place-items: center;
+    border-radius: 999px;
+    background: var(--e0-accent);
+    color: var(--e0-accent-ink);
+    transition: transform 140ms ease, opacity 140ms ease;
+  }
+
+  .e0-send-btn:hover:not(:disabled) {
+    transform: translateY(-1px);
+  }
+
+  .e0-send-btn:disabled {
+    cursor: not-allowed;
+    opacity: 0.38;
+  }
+
+  /* Prose colors follow the design tokens (the typography plugin's default gray is too bright on the near-black background) */
+  .e0-prose {
+    --tw-prose-body: var(--e0-text);
+    --tw-prose-headings: var(--e0-text);
+    --tw-prose-bold: var(--e0-text);
+    --tw-prose-links: var(--e0-accent);
+    --tw-prose-code: var(--e0-text);
+    --tw-prose-quotes: var(--e0-text);
+    --tw-prose-quote-borders: var(--e0-border);
+    --tw-prose-bullets: var(--e0-muted);
+    --tw-prose-hr: var(--e0-border);
+    --tw-prose-th-borders: var(--e0-border);
+    --tw-prose-td-borders: var(--e0-border);
+  }
+}
+
+@media (max-width: 760px) {
+  .e0-app-shell {
+    display: block;
+    padding: 0;
+  }
+
+  .e0-sidebar {
+    position: fixed;
+    inset: 10px auto 10px 10px;
+    width: min(280px, calc(100vw - 28px));
+    transform: translateX(calc(-100% - 20px));
+    transition: transform 180ms ease;
+  }
+
+  .e0-sidebar.e0-sidebar-open {
+    transform: translateX(0);
+  }
+
+  .e0-sidebar-backdrop {
+    position: fixed;
+    z-index: 20;
+    inset: 0;
+    background: rgba(13, 17, 23, 0.34);
+  }
+
+  .e0-main-frame {
+    height: 100vh;
+    min-height: 100vh;
+    border: 0;
+    border-radius: 0;
+  }
+
+  .e0-mobile-toggle {
+    display: inline-grid;
+  }
+
+  .e0-page {
+    padding: 26px 16px 32px;
+  }
+
+  .e0-page-header {
+    align-items: center;
+    margin-bottom: 22px;
+  }
+
+  .e0-chat-header {
+    min-height: 54px;
+    padding: 0 16px 0 58px;
+  }
+
+  .e0-chat-column {
+    width: 100%;
+  }
+
+  .e0-composer {
+    padding: 10px 12px 14px;
+  }
+}
+
+@media (prefers-reduced-motion: reduce) {
+  *,
+  *::before,
+  *::after {
+    scroll-behavior: auto !important;
+    transition-duration: 0.01ms !important;
+    animation-duration: 0.01ms !important;
+  }
+}
diff --git a/windows/app/src/vite-env.d.ts b/windows/app/src/vite-env.d.ts
new file mode 100644
index 0000000..39fc679
--- /dev/null
+++ b/windows/app/src/vite-env.d.ts
@@ -0,0 +1,3 @@
+/// <reference types="vite/client" />
+
+declare module "*.css";
diff --git a/windows/app/tsconfig.json b/windows/app/tsconfig.json
new file mode 100644
index 0000000..08ff715
--- /dev/null
+++ b/windows/app/tsconfig.json
@@ -0,0 +1,15 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "lib": ["ES2022", "DOM", "DOM.Iterable"],
+    "module": "ESNext",
+    "moduleResolution": "bundler",
+    "jsx": "react-jsx",
+    "strict": true,
+    "noEmit": true,
+    "skipLibCheck": true,
+    "isolatedModules": true,
+    "resolveJsonModule": true
+  },
+  "include": ["src"]
+}
diff --git a/windows/app/vite.config.ts b/windows/app/vite.config.ts
new file mode 100644
index 0000000..0dd5c55
--- /dev/null
+++ b/windows/app/vite.config.ts
@@ -0,0 +1,12 @@
+import react from "@vitejs/plugin-react";
+import tailwindcss from "@tailwindcss/vite";
+import { defineConfig } from "vite";
+
+// Tauri expects a fixed dev port on loopback; fail fast if it is taken.
+export default defineConfig({
+  plugins: [react(), tailwindcss()],
+  clearScreen: false,
+  server: { port: 1420, strictPort: true, host: "127.0.0.1" },
+  envPrefix: ["VITE_", "TAURI_"],
+  build: { target: "es2022", sourcemap: false },
+});
diff --git a/windows/scripts/vendor-build.ps1 b/windows/scripts/vendor-build.ps1
new file mode 100644
index 0000000..0d55340
--- /dev/null
+++ b/windows/scripts/vendor-build.ps1
@@ -0,0 +1,88 @@
+# scripts/vendor-build.ps1 — llama engine assembly for the Windows product (monorepo layout)
+# The pristine upstream tree (../vendor/llama.cpp, pinned commit) is NEVER patched in place:
+# this script replays the patch bands into an isolated build worktree (../wt/win), copies the
+# edge0 serving pieces from serve/ into src/edge0 (patch #3's CMake glob compiles them into
+# the llama target), then configures and builds with Vulkan.
+# The vendor tree is submodule-free: it is MATERIALIZED on first run by cloning upstream at
+# the pinned commit listed in vendor.llama.pin (repo root); afterwards it is just a checkout.
+# Depot resolution (first match wins):
+#   A. $env:EDGE0_DEPOT                                   explicit override
+#   B. parent directory containing vendor.llama.pin       monorepo checkout (the normal path)
+#   If neither matches, stop and explain — run from a full monorepo clone, or point
+#   EDGE0_DEPOT at a depot that carries vendor.llama.pin.
+# Set EDGE0_LLAMA_URL to clone from a mirror instead of GitHub.
+# Usage: pwsh -File scripts/vendor-build.ps1 [-BuildDir build-vk] [-AssembleOnly]
+#   -AssembleOnly = stages 1-3 only (checks + patch replay + serve copy), no compile.
+param(
+    [string]$BuildDir = "build-vk",
+    [switch]$Clean,
+    [switch]$AssembleOnly
+)
+$ErrorActionPreference = "Stop"
+$ROOT = Split-Path -Parent (Split-Path -Parent $MyInvocation.MyCommand.Path)
+
+$DEPOT = if ($env:EDGE0_DEPOT) { $env:EDGE0_DEPOT } else {
+    $parent = Split-Path -Parent $ROOT
+    if (Test-Path (Join-Path $parent "vendor.llama.pin")) { $parent } else { $null }
+}
+if (-not $DEPOT) { throw "no depot: set EDGE0_DEPOT to a tree carrying vendor.llama.pin (the monorepo root), or run inside a full monorepo clone" }
+$PINFILE = Join-Path $DEPOT "vendor.llama.pin"
+if (-not (Test-Path $PINFILE)) { throw "vendor.llama.pin missing in depot $DEPOT" }
+$PIN = (Get-Content $PINFILE -Raw).Trim().Split(" ")[0]   # pinned upstream commit (tag b11100) — ledger: patches/llama.cpp/README.md
+$VENDOR = Join-Path $DEPOT "vendor/llama.cpp"
+$WT     = Join-Path $DEPOT "wt/win"
+# Patch bands resolve from the depot root ONLY — one source of truth. The old in-repo
+# windows/patches/ mirror was retired 2026-09-30 (a drifting copy; depot carries the bands).
+$PDIR = Join-Path $DEPOT "patches/llama.cpp"
+if (-not (Test-Path (Join-Path $PDIR "common"))) { throw "patch bands missing at $PDIR/common — run inside the monorepo or point EDGE0_DEPOT at a depot carrying patches/llama.cpp" }
+$BANDS = @("common", "windows") | ForEach-Object { Join-Path $PDIR "$_/*.patch" }
+
+# Materialize the pinned tree when absent (submodule-free supply).
+if (-not (Test-Path (Join-Path $VENDOR ".git"))) {
+    Write-Host "[0/4] materializing vendor: clone llama.cpp at pinned $PIN -> $VENDOR" -ForegroundColor Cyan
+    New-Item -ItemType Directory -Force (Split-Path $VENDOR) | Out-Null
+    $LLAMA_URL = if ($env:EDGE0_LLAMA_URL) { $env:EDGE0_LLAMA_URL } else { "https://github.com/ggml-org/llama.cpp" }
+    git clone --filter=blob:none $LLAMA_URL $VENDOR
+    if ($LASTEXITCODE -ne 0) { throw "clone failed ($LLAMA_URL) — network? or set EDGE0_LLAMA_URL / EDGE0_DEPOT" }
+    git -C $VENDOR checkout --detach $PIN
+    if ($LASTEXITCODE -ne 0) { throw "pinned commit $PIN not found upstream (refresh the patch ledger)" }
+}
+
+Push-Location $VENDOR
+try {
+    Write-Host "[1/4] checking pristine vendor tree (must sit exactly at the pin, zero local changes)" -ForegroundColor Cyan
+    $head = (git rev-parse HEAD).Trim()
+    if (-not $head.StartsWith($PIN)) { throw "vendor HEAD=$head is not the pinned $PIN — the pristine tree must never be patched in place; restore with: git checkout --detach $PIN" }
+    if ((git status --porcelain).Length -gt 0) { throw "vendor worktree is dirty — the pristine source of truth tolerates no local edits" }
+
+    Write-Host "[2/4] (re)creating build worktree $WT and replaying patch bands" -ForegroundColor Cyan
+    if (-not (Test-Path $WT)) { git worktree add --detach $WT $PIN | Out-Null }
+    # Unconditional reset before replay = idempotent
+    git -C $WT reset --hard $PIN 2>&1 | Out-Null
+    git -C $WT clean -fd src/edge0 tools/llama-bench 2>&1 | Out-Null
+    Set-Location $WT
+    foreach ($pat in $BANDS) {
+        Get-ChildItem $pat | Sort-Object Name | ForEach-Object {
+            Write-Host "  am $($_.Name)"
+            git am --3way $_.FullName
+            if ($LASTEXITCODE -ne 0) { throw "git am failed at $($_.Name) (upstream drift or band conflict — three-way by hand, see patches/llama.cpp/README.md)" }
+        }
+    }
+
+    Write-Host "[3/4] copying serve/ pieces -> src/edge0/ (compiled in via patch #3 glob)" -ForegroundColor Cyan
+    New-Item -ItemType Directory -Force (Join-Path $WT "src/edge0") | Out-Null
+    Copy-Item (Join-Path $ROOT "serve/*.cc") (Join-Path $WT "src/edge0/") -Force -ErrorAction SilentlyContinue
+    Copy-Item (Join-Path $ROOT "serve/*.h")  (Join-Path $WT "src/edge0/") -Force -ErrorAction SilentlyContinue
+    if ($AssembleOnly) {
+        Write-Host "OK(assemble-only): worktree=$WT — patches replayed, serving pieces copied" -ForegroundColor Green
+        return
+    }
+
+    Write-Host "[4/4] configure + build (Release/Vulkan)" -ForegroundColor Cyan
+    $bd = Join-Path $WT $BuildDir
+    cmake -S $WT -B $bd -DGGML_VULKAN=ON -DCMAKE_BUILD_TYPE=Release
+    cmake --build $bd --config Release
+    Write-Host "OK: $bd/bin/Release/llama-server.exe" -ForegroundColor Green
+} finally {
+    Pop-Location
+}
diff --git a/windows/serve/mem_budget.cc b/windows/serve/mem_budget.cc
new file mode 100644
index 0000000..bc96226
--- /dev/null
+++ b/windows/serve/mem_budget.cc
@@ -0,0 +1,38 @@
+// serve/mem_budget.cc — process memory budget (hard working-set cap on Windows).
+// Semantics: budget = N MB ⇒ resident set is capped at N; beyond that the OS evicts
+// our pages (mmap'd model-file pages first, so the next touch faults a real re-read).
+// This is the direct implementation of the low-RAM product tiers (8 GB / 16 GB machines).
+#include "mem_budget.h"
+
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+
+#include <cstdio>
+#include <cstdlib>
+
+E0_MB_API bool edge0_apply_mem_budget(int mb) {
+    if (mb <= 0) {                                        // debug-compat env fallback
+        const char * e = getenv("E0_MEM_BUDGET_MB");
+        if (e && *e) mb = (int) atol(e);
+    }
+    if (mb <= 0) return false;
+
+    const SIZE_T budget = (SIZE_T) mb * 1024 * 1024;
+    SIZE_T cur_min = 0, cur_max = 0; DWORD cur_flags = 0;
+    SIZE_T want_min = budget / 8;                 // soft floor = 12.5%: a residency base so the OS doesn't evict everything
+    if (GetProcessWorkingSetSizeEx(GetCurrentProcess(), &cur_min, &cur_max, &cur_flags) && cur_max) {
+        if (want_min > cur_min) want_min = cur_min < budget ? cur_min : budget;  // never raise the original minimum
+    }
+    // max working set = hard cap: the OS trims immediately on overflow (unlike the
+    // default ~1.5 TB soft limit, which only trims under global memory pressure)
+    if (!SetProcessWorkingSetSizeEx(GetCurrentProcess(), want_min, budget,
+                                    QUOTA_LIMITS_HARDWS_MIN_DISABLE | QUOTA_LIMITS_HARDWS_MAX_ENABLE)) {
+        fprintf(stderr, "[mem-budget] SetProcessWorkingSetSizeEx failed err=%lu (budget %dMB NOT applied)\n",
+                GetLastError(), mb);
+        return false;
+    }
+    fprintf(stderr, "[mem-budget] applied: max working set = %d MB (min %llu MB, hard-trim on)\n",
+            mb, (unsigned long long) (want_min / (1024 * 1024)));
+    fflush(stderr);
+    return true;
+}
diff --git a/windows/serve/mem_budget.h b/windows/serve/mem_budget.h
new file mode 100644
index 0000000..c03627a
--- /dev/null
+++ b/windows/serve/mem_budget.h
@@ -0,0 +1,14 @@
+// serve/mem_budget.h — process memory budget (first-class product feature).
+// Independent of the prefetch router: it only caps physical residency; no prediction.
+#pragma once
+#ifdef _WIN32
+#  define E0_MB_API __declspec(dllexport)
+#else
+#  define E0_MB_API __attribute__((visibility("default")))
+#endif
+
+// Primary channel: CLI flag (--mem-budget-mb → mb > 0); when mb <= 0, falls back to
+// the E0_MEM_BUDGET_MB env var (debug compat). mb > 0 pins this process's maximum
+// working set to that size — the OS then evicts pages beyond the cap, so cold model
+// reads genuinely occur. Returns whether the cap was applied.
+E0_MB_API bool edge0_apply_mem_budget(int mb);
diff --git a/windows/serve/prefetch.cc b/windows/serve/prefetch.cc
new file mode 100644
index 0000000..6aeec2c
--- /dev/null
+++ b/windows/serve/prefetch.cc
@@ -0,0 +1,854 @@
+// serve/prefetch.cc — advisory expert-prefetch router + self-managed L1 pool + accounting.
+// Depends only on llama public/near-public APIs (declared inline, linked from llama.lib)
+// and parses GGUF/safetensors headers by hand; llama internal headers are not touched.
+// The real routing gate is never modified: predictions only warm the page cache. A wrong
+// prediction costs latency, never quality.
+// Three-way overlap accounting (printed as [pref-trace]):
+//   head   = pred@t-1 ∩ teacher@t        teacher = host-side recompute of gate·m_in (gate is F32 in our GGUF)
+//   sticky = teacher@t-1 ∩ teacher@t     (the free win the OS page-cache LRU already provides)
+//   random = uniform sample of K ∩ teacher@t   (popularity floor)
+#include "prefetch.h"
+
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+#include <psapi.h>
+
+#include <atomic>
+#include <chrono>
+#include <cmath>
+#include <condition_variable>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <map>
+#include <mutex>
+#include <string>
+#include <thread>
+#include <vector>
+
+// Near-public staging API (src/llama-ext.h). NOTE: that header has **no extern "C"
+// wrapper** — its symbols are C++-decorated, so you must include the real header rather
+// than hand-write declarations. This unit is globbed into the llama target by the patch
+// band (patches/llama.cpp/common #3), which puts src/ on the PRIVATE include path;
+// llama.h is reachable via ../include PUBLIC.
+#include "llama.h"
+#include "llama-ext.h"
+
+
+namespace {
+
+// Pool final telemetry (defined in the pool section below) — fwd-declared for ~Pref.
+static void e0_pool_trace_final();
+
+// ── Shapes (pinned for the 35B GGUF: E=256, H=2048, feat=2560, K=4, owners 6..38 → consumers 7..39) ──
+constexpr int   EK = 256, HH = 2048, FE = 2560, KDIM = 4, NH = 33;
+constexpr int   OWNER0 = 6;                       // first owner
+std::vector<int> consumers;                        // 7..39
+
+// ── GGUF/safetensors header parsing ──
+std::string rd_str(FILE * f) {
+    uint64_t n = 0;
+    if (fread(&n, 8, 1, f) != 1) return {};
+    std::string s((size_t) n, '\0');
+    if (n && fread(s.data(), 1, (size_t) n, f) != n) return {};
+    return s;
+}
+struct GGUFType { uint32_t t; };
+size_t align_up(size_t x, size_t a) { return (x + a - 1) / a * a; }
+// GGUF KV value types: 0u8 1i8 2u16 3i16 4u32 5i32 6f32 7bool 8str 9arr 10u64 11i64 12f64
+size_t scalar_size(uint32_t t) { static const size_t s[] = {1,1,2,2,4,4,4,1,0,0,8,8,8}; return t <= 12 ? s[t] : 0; }
+void skip_val(FILE * f, uint32_t t) {
+    if (t == 8) { rd_str(f); return; }
+    if (t == 9) {
+        uint32_t et; uint64_t n;
+        if (fread(&et, 4, 1, f) != 1 || fread(&n, 8, 1, f) != 1) return;
+        for (uint64_t i = 0; i < n; ++i) skip_val(f, et);
+        return;
+    }
+    size_t s = scalar_size(t);
+    if (s) fseek(f, (long) s, SEEK_CUR);
+}
+// Tensor dtype uses the GGML_TYPE enum: 0=F32 1=F16 2=BF16 3=Q4_1 (NOT the GGUF KV enum — don't mix).
+bool   is_q4_1(uint32_t t) { return t == 3; }
+size_t plain_esz(uint32_t t) { return t == 0 ? 4 : (t == 1 || t == 2) ? 2 : 0; }
+
+struct TT {
+    std::string name;
+    uint32_t type = 0;
+    uint64_t dim[4] = {1, 1, 1, 1};
+    uint64_t off = 0;
+};
+size_t tt_bytes(const TT & x) {
+    size_t n = 1;
+    for (int i = 0; i < 4; ++i) n *= (size_t) x.dim[i];
+    if (is_q4_1(x.type)) return (size_t) x.dim[1] * x.dim[2] * x.dim[3] * ((size_t) x.dim[0] / 32) * 20;
+    return n * plain_esz(x.type);
+}
+
+// IEEE half → float (subnormals included)
+float f16_to_f32(uint16_t h) {
+    uint32_t sign = (uint32_t) (h & 0x8000) << 16;
+    uint32_t e = (h >> 10) & 0x1F, m = h & 0x3FF;
+    uint32_t bits;
+    if (e == 0) {
+        if (m == 0) bits = sign;
+        else {                                        // subnormal: renormalize
+            e = 127 - 15 + 1;
+            while (!(m & 0x400)) { m <<= 1; e--; }
+            m &= 0x3FF;
+            bits = sign | (e << 23) | (m << 13);
+        }
+    } else if (e == 0x1F) {
+        bits = sign | 0x7F800000u | (m << 13);
+    } else {
+        bits = sign | ((e + 127 - 15) << 23) | (m << 13);
+    }
+    float r;
+    memcpy(&r, &bits, 4);
+    return r;
+}
+
+// safetensors: { "<key>": {"dtype":"F16","shape":[..],"data_offsets":[a,b]} }
+struct SEnt { long long a = -1, b = -1; std::vector<long long> shape; };
+bool st_load(const char * path, std::map<std::string, SEnt> & out) {
+    FILE * f = fopen(path, "rb");
+    if (!f) return false;
+    uint64_t hn = 0;
+    if (fread(&hn, 8, 1, f) != 1) { fclose(f); return false; }
+    std::string hdr((size_t) hn, '\0');
+    bool ok = fread(hdr.data(), 1, (size_t) hn, f) == hn;
+    fclose(f);
+    if (!ok) return false;
+    size_t i = 0;
+    while ((i = hdr.find("\"layers.", i)) != std::string::npos) {
+        size_t ke = hdr.find("\":", i + 1);
+        size_t vo = hdr.find('{', ke);
+        size_t vc = ke == std::string::npos || vo == std::string::npos ? std::string::npos : hdr.find('}', vo);
+        if (vc == std::string::npos) break;
+        std::string key = hdr.substr(i + 1, ke - i - 1);
+        std::string o = hdr.substr(vo, vc - vo + 1);
+        SEnt e;
+        size_t doff = o.find("\"data_offsets\"");
+        size_t lb = doff == std::string::npos ? std::string::npos : o.find('[', doff);
+        size_t rb = lb == std::string::npos ? std::string::npos : o.find(']', lb);
+        if (rb == std::string::npos) { i = vc; continue; }
+        e.a = strtoll(o.c_str() + lb + 1, nullptr, 10);
+        const char * p = o.c_str() + lb + 1;
+        while (*p && *p != ',') p++;
+        if (*p == ',') e.b = strtoll(p + 1, nullptr, 10);
+        size_t sp = o.find("\"shape\""), slb = sp == std::string::npos ? std::string::npos : o.find('[', sp),
+               srb = slb == std::string::npos ? std::string::npos : o.find(']', slb);
+        if (srb != std::string::npos) {
+            std::string nums = o.substr(slb + 1, srb - slb - 1);
+            const char * q = nums.c_str();
+            while (*q) {
+                char * r; long long v = strtoll(q, &r, 10);
+                if (r != q) e.shape.push_back(v);
+                q = r; while (*q == ',' || *q == ' ') q++;
+            }
+        }
+        out[key] = e;
+        i = vc;
+    }
+    return !out.empty();
+}
+
+bool grab_f16(FILE * f, uint64_t data_base, const SEnt & e, std::vector<float> & dst) {
+    size_t n = 1;
+    for (auto d : e.shape) n *= (size_t) d;
+    dst.resize(n);
+    if (_fseeki64(f, (long long) (data_base + e.a), SEEK_SET) != 0) return false;
+    std::vector<uint16_t> raw(n);
+    if (fread(raw.data(), 2, n, f) != n) return false;
+    for (size_t i = 0; i < n; ++i) dst[i] = f16_to_f32(raw[i]);
+    return true;
+}
+
+// ── Runtime state ──
+struct Head { std::vector<float> fc1, fc2, lin; };   // [512][2560] / [256][512] / [256][2560]
+struct Hist { std::vector<float> cur, prev; };        // owner-layer two-step onehot (E)
+
+struct Pref {
+    bool active = false;
+    std::map<int, Head> heads;                        // owner → head triple
+    std::map<int, std::vector<float>> gate;           // consumer → [E][H]
+    struct EXP { std::vector<TT> t; };               // consumer → up/gate/down exps metadata
+    std::map<int, EXP> exps;
+    std::map<int, Hist> hist;                         // layers 6..39 (39 keeps teacher onehot only)
+    std::map<int, std::vector<int>> pred_prev;        // consumer → previous-step head top-K
+    std::map<int, std::vector<int>> exec_prev;        // consumer → previous-step teacher (sticky)
+    uint64_t data_base = 0;
+    HANDLE fh = INVALID_HANDLE_VALUE, mh = NULL;
+    void * view = nullptr;
+    bool do_prefetch = false;
+    int  trace_every = 32;
+    int  stride = 1;          // full head/teacher recompute period (E0_PREF_STRIDE); skipped steps only re-issue prefetch
+    bool async = false;       // E0_PREF_ASYNC=0 falls back to synchronous (tax-splitting A/B)
+    // Predicted-range VirtualLock: pinned bytes are not trimmed while under budget;
+    // lock failure once the cap is hit = natural ceiling. NOTE: VirtualLock normally
+    // requires SeLockMemoryPrivilege, absent on consumer machines — treat this as best-effort.
+    bool   do_pin = false;
+    size_t pin_cap = 3ull << 30, pin_bytes = 0;
+    std::map<int, std::vector<std::pair<size_t, size_t>>> pinned;   // consumer → currently pinned (file_off,len)
+    // Self-managed L1 over the mmap view (E0_POOL): re-touch at step boundaries, in-table
+    // age-based LRU; "evict" = drop from table (stop re-touching; reference-bit aging lets
+    // the OS reclaim to standby — a free soft layer, no privileged APIs involved).
+    int  pool = 0;                                    // 0=legacy advisory 1=sticky 2=sticky+observation
+    struct Pin { size_t off, len; int age; uint8_t reg; };
+    std::map<int, std::vector<Pin>> pins;             // consumer → pin table (age = steps since last hit)
+    size_t reg_bytes[2] = {0, 0}, reg_cap[2] = {0, 0};   // 0=hot zone (teacher) 1=observation zone (head)
+    long pool_hit = 0, pool_tot = 0, lock_fail = 0, evicts = 0, shrinks = 0;
+    // Async: m_in row-snapshot mailbox ((NH+1)×H floats, contiguous); worker runs full steps.
+    std::vector<float> snap;
+    std::atomic<bool> busy{false}, quit{false};
+    std::mutex wmtx;
+    std::condition_variable cv;
+    std::thread worker;
+    long hits = 0, sticky = 0, rnd = 0, total = 0, steps = 0;
+    long long pool_mb = 0;                        // private-page pool hot-zone MB (>0 enables; E0_POOL_MB)
+    uint32_t lcg = 2463534242u;
+    uint32_t next() { lcg ^= lcg << 13; lcg ^= lcg >> 17; lcg ^= lcg << 5; return lcg; }
+    ~Pref() {
+        quit = true; cv.notify_all();
+        if (worker.joinable()) worker.join();
+        if (!total) return;
+        fprintf(stderr, "\n[pref-trace] head   overlap = %ld/%ld = %.1f%%\n", hits, total, 100.0 * hits / total);
+        fprintf(stderr, "[pref-trace] sticky overlap = %ld/%ld = %.1f%%\n", sticky, total, 100.0 * sticky / total);
+        fprintf(stderr, "[pref-trace] random overlap = %ld/%ld = %.1f%%\n", rnd, total, 100.0 * rnd / total);
+        fprintf(stderr, "[pref-trace] head−sticky = %.1f pp | head−random = %.1f pp\n",
+                100.0 * (hits - sticky) / total, 100.0 * (hits - rnd) / total);
+        e0_pool_trace_final();
+        if (pool)
+            fprintf(stderr, "[pref-trace] POOL(%s) hit = %ld/%ld = %.1f%% | pins hot=%lluMB obs=%lluMB (cap %llu/%llu) evict=%ld lockfail=%ld shrink=%ld\n",
+                    pool == 2 ? "sticky+obs" : "sticky", pool_hit, pool_tot, pool_tot ? 100.0 * pool_hit / pool_tot : 0.0,
+                    (unsigned long long) (reg_bytes[0] >> 20), (unsigned long long) (reg_bytes[1] >> 20),
+                    (unsigned long long) (reg_cap[0] >> 20), (unsigned long long) (reg_cap[1] >> 20),
+                    evicts, lock_fail, shrinks);
+        fflush(stderr);
+    }
+};
+Pref g;
+
+float gelu_erf(float x) { return 0.5f * x * (1.f + erff(x * 0.7071067811865475f)); }
+float dotf(const float * __restrict a, const float * __restrict b, size_t n) {
+    float s = 0.f;
+    for (size_t i = 0; i < n; ++i) s += a[i] * b[i];   // /arch:AVX2 auto-FMA reduction; the bottleneck is weight streaming, not ALU
+    return s;
+}
+void topk(const float * v, int n, int k, std::vector<int> & out) {
+    out.clear();
+    std::vector<char> used(n, 0);
+    for (int j = 0; j < k; ++j) {
+        int best = -1;
+        for (int i = 0; i < n; ++i) if (!used[i] && (best < 0 || v[i] > v[best])) best = i;
+        if (best < 0) break;
+        used[best] = 1; out.push_back(best);
+    }
+}
+int ovlp(const std::vector<int> & a, const std::vector<int> & b) {
+    int n = 0;
+    for (int x : a) for (int y : b) if (x == y) { ++n; break; }
+    return n;
+}
+
+}  // namespace
+
+// ══ Private-page expert pool (E0_POOL_MB): L1 with fill-once slots ══
+// Design: decode reads expert bytes **directly from pool pages** via the expert-base
+// resolver (patch surface #7), bypassing the view page-fault path entirely — a hit is a
+// plain read with no kernel involvement. Slots fill once and never recycle (no eviction
+// ⇒ no read-vs-overwrite race, correctness by construction); once the budget is
+// exhausted, new experts pass through to the mmap view = graceful degradation to the
+// un-pooled behavior. Three acquisition channels:
+//   demand = synchronous whole-block ReadFile on the compute thread (first touch)
+//   hot    = worker prefill from the teacher set
+//   obs    = dedicated quota for head predictions
+// Whole-block reads are also fragmentation-proof (1×1MB sequential vs a 256×4K fault
+// storm — an order of magnitude apart on consumer SSDs).
+static_assert(sizeof(LONG) == 4, "slot state atoms");
+namespace {
+
+struct PoolTT {                       // one 3D experts tensor (40 slots measured across all layers)
+    uint8_t * arena = nullptr;        // MEM_RESERVE, MEM_COMMIT on demand
+    volatile LONG * st = nullptr;     // slot state: 0=empty 1=filling 2=ready 3=not-poolable
+    CRITICAL_SECTION * cvlk = nullptr;// per-slot pair: CS lock + CV (filler owns it, readers sleep until woken)
+    CONDITION_VARIABLE * cv = nullptr;
+    size_t per = 0;                   // bytes per expert
+    uint64_t foff = 0;                // file offset of expert 0
+};
+static volatile LONG g_p_sc_ok = 0, g_p_sc_bad = 0;      // byte-level self-check (fill vs view memcmp sampling)
+static bool g_pool_sc = false;                           // E0_POOL_SELFCHECK=1 enables
+static void e0_pool_selfcheck(uint8_t * slot, uint64_t foff, size_t per) {
+    if (!g_pool_sc || !g.view) return;
+    const uint8_t * v = (const uint8_t *) g.view + foff;
+    int bad = 0;
+    for (size_t i = 0; i < per; i += 4096) if (memcmp(slot + i, v + i, 64) != 0) { bad = 1; break; }
+    if (bad) InterlockedIncrement(&g_p_sc_bad); else InterlockedIncrement(&g_p_sc_ok);
+}
+// ⚠ The tensor *pointer* changes every forward (same tensor registered multiple times,
+// measured) ⇒ key by tensor NAME (table built at init, read-only afterwards → lock-free).
+std::map<std::string, PoolTT *> g_ptt_by_name;                  // GGUF name → PTT
+std::map<int, std::vector<PoolTT *>> g_ptt_by_c;                // consumer → 3 PTTs (for prefill)
+std::mutex g_pmtx;                                              // fill/accounting lock (fills are ms-scale, not a spin hotspot)
+size_t g_pool_cap[2] = {0, 0};                                  // 0=hot 1=obs (E0_POOL_MB/E0_POOL_OBS_MB)
+size_t g_pool_bytes[2] = {0, 0};
+volatile LONG g_p_calls = 0;
+bool g_pool_on = false;
+volatile LONG g_p_hits = 0, g_p_miss_d = 0, g_p_bypass = 0, g_p_waits = 0, g_p_fills = 0;
+
+// Expert-base resolver hook from the patch band (exported by ggml.dll, C decoration).
+typedef const char * (*ggml_edge0_mmid_base_fn)(const struct ggml_tensor * src0, int64_t id, size_t nb02);
+extern "C" __declspec(dllimport) void ggml_edge0_set_mmid_resolver(ggml_edge0_mmid_base_fn fn);
+
+// Resolver callback (patch surface #7, runs on ggml compute threads): returns the base
+// address of expert `id`'s data.
+extern "C" static const char * e0_mmid_base_res(const ggml_tensor * s, int64_t id, size_t nb02) {
+    LONG c = InterlockedIncrement(&g_p_calls);
+    if (c <= 2) { fprintf(stderr, "[pref-trace] POOL2 call#%ld name=%s nb02=%zu\n", c, s->name ? s->name : "?", nb02); fflush(stderr); }
+    const char * orig = (const char *) s->data + (size_t) id * nb02;
+    auto it = g_ptt_by_name.find(s->name ? s->name : "");
+    if (it == g_ptt_by_name.end() || !it->second->st || it->second->per != nb02) {
+        LONG b = InterlockedIncrement(&g_p_bypass);
+        if (b <= 6)
+            fprintf(stderr, "[pref-trace] POOL2 miss#%ld name=%s nb02=%zu in_table=%d table=%zu\n",
+                    b, s->name ? s->name : "?", nb02, (int) (it != g_ptt_by_name.end()), g_ptt_by_name.size());
+        return orig;
+    }
+    PoolTT * pt = it->second;
+    for (;;) {
+        LONG st = pt->st[id];
+        if (st == 1) {                                        // another thread is filling: sleep until woken (no spinning)
+            InterlockedIncrement(&g_p_waits);
+            CRITICAL_SECTION & lk = pt->cvlk[id];
+            EnterCriticalSection(&lk);
+            if (pt->st[id] == 1) SleepConditionVariableCS(&pt->cv[id], &lk, 2000);
+            LeaveCriticalSection(&lk);
+            continue;
+        }
+        if (st == 2) {
+            LONG h = InterlockedIncrement(&g_p_hits);
+            if (h <= 3 || (h & 0x3FFF) == 0)
+                fprintf(stderr, "[pref-trace] POOL2 hit=%ld fillD=%ld bypass=%ld wait=%ld sc=%ld/%ld | hot=%lluMB obs=%lluMB\n",
+                        h, g_p_miss_d, g_p_bypass, g_p_waits, (long) g_p_sc_ok, (long) g_p_sc_bad,
+                        (unsigned long long) (g_pool_bytes[0] >> 20), (unsigned long long) (g_pool_bytes[1] >> 20));
+            fflush(stderr);
+            return (const char *) (pt->arena + (size_t) id * pt->per);
+        }
+        if (st == 3) { InterlockedIncrement(&g_p_bypass); return orig; }
+        {                                                                   // claim the fill: CAS + exclusive CV section
+            CRITICAL_SECTION & lk = pt->cvlk[id];
+            if (InterlockedCompareExchange(&pt->st[id], 1, 0) != 0) continue;
+            EnterCriticalSection(&lk);
+        }
+        LONG fs = InterlockedIncrement(&g_p_fills);
+        if (fs <= 3 || (fs & 0x3FF) == 0) {
+            fprintf(stderr, "[pref-trace] POOL2 fill#%ld name=%s id=%lld per=%zu hits=%ld\n",
+                    fs, s->name ? s->name : "?", (long long) id, pt->per, g_p_hits);
+            fflush(stderr);
+        }
+        // Demand fill on this thread (hot budget).
+        bool ok = false;
+        {
+            std::lock_guard<std::mutex> lk(g_pmtx);
+            if (g_pool_bytes[0] + pt->per <= g_pool_cap[0]) {
+                uint8_t * dst = pt->arena + (size_t) id * pt->per;
+                if (VirtualAlloc(dst, pt->per, MEM_COMMIT, PAGE_READWRITE)) {
+                    uint64_t fo = pt->foff + (uint64_t) id * pt->per;
+                    LARGE_INTEGER off; off.QuadPart = (LONGLONG) fo;
+                    ok = SetFilePointerEx(g.fh, off, nullptr, FILE_BEGIN)
+                         && ReadFile(g.fh, dst, (DWORD) pt->per, nullptr, nullptr);
+                    if (ok) { g_pool_bytes[0] += pt->per; e0_pool_selfcheck(dst, fo, pt->per); }
+                }
+            }
+            if (ok) InterlockedIncrement(&g_p_miss_d);
+        }
+        pt->st[id] = ok ? 2 : 3;
+        WakeAllConditionVariable(&pt->cv[id]);
+        LeaveCriticalSection(&pt->cvlk[id]);
+        return ok ? (const char *) (pt->arena + (size_t) id * pt->per) : orig;
+    }
+}
+
+// Worker prefill (teacher=hot / head=obs): called outside the resolver path; same per-slot CAS pattern.
+static void e0_pool_prefill(int c, const std::vector<int> & ids, int reg) {
+    auto it = g_ptt_by_c.find(c);
+    if (it == g_ptt_by_c.end()) return;
+    for (PoolTT * pt : it->second)
+        for (int id : ids) {
+            if (pt->st[id] != 0) continue;
+            CRITICAL_SECTION & lk = pt->cvlk[id];
+            if (InterlockedCompareExchange(&pt->st[id], 1, 0) != 0) continue;
+            EnterCriticalSection(&lk);
+            bool ok = false;
+            {
+                std::lock_guard<std::mutex> lk(g_pmtx);
+                if (g_pool_bytes[reg] + pt->per <= g_pool_cap[reg]) {
+                    uint8_t * dst = pt->arena + (size_t) id * pt->per;
+                    if (VirtualAlloc(dst, pt->per, MEM_COMMIT, PAGE_READWRITE)) {
+                        LARGE_INTEGER off = { (DWORD) (pt->foff + (uint64_t) id * pt->per),
+                                              (LONG) ((pt->foff + (uint64_t) id * pt->per) >> 32) };
+                        ok = SetFilePointerEx(g.fh, off, nullptr, FILE_BEGIN)
+                             && ReadFile(g.fh, dst, (DWORD) pt->per, nullptr, nullptr);
+                        if (ok) { g_pool_bytes[reg] += pt->per; e0_pool_selfcheck(dst, pt->foff + (uint64_t) id * pt->per, pt->per); }
+                    }
+                }
+            }
+            pt->st[id] = ok ? 2 : 3;
+            WakeAllConditionVariable(&pt->cv[id]);
+            LeaveCriticalSection(&lk);
+        }
+}
+
+static void e0_pool_trace_final() {
+    if (!g_pool_on) return;
+    fprintf(stderr, "[pref-trace] POOL2 hit=%ld fillD=%ld bypass=%ld wait=%ld sc=%ld/%ld | hot=%lluMB obs=%lluMB\n",
+            g_p_hits, g_p_miss_d, g_p_bypass, g_p_waits, (long) g_p_sc_ok, (long) g_p_sc_bad,
+            (unsigned long long) (g_pool_bytes[0] >> 20), (unsigned long long) (g_pool_bytes[1] >> 20));
+}
+
+static void e0_pool_scan(const std::map<std::string, TT> & tinfo) {   // pool all *_exps tensors (independent of consumer shapes)
+    for (auto & kv : tinfo) {
+        const std::string & nm = kv.first;
+        const TT & tt = kv.second;
+        auto suff = [&](const char * sfx) -> size_t {           // exact suffix match (returns length on hit)
+            size_t len = strlen(sfx), k = nm.rfind(sfx);
+            return (k != std::string::npos && k + len == nm.size()) ? len : 0;
+        };
+        size_t l1 = suff(".ffn_up_exps.weight"), l2 = suff(".ffn_gate_exps.weight"), l3 = suff(".ffn_down_exps.weight");
+        size_t sl = l1 ? l1 : (l2 ? l2 : l3);
+        if (!sl || tt.dim[2] < 2 || tt.dim[2] > 4096) continue;
+        const size_t per = tt_bytes(tt) / (size_t) tt.dim[2];
+        if (!per) continue;
+        PoolTT * e = new PoolTT();
+        e->per = per;
+        e->foff = g.data_base + tt.off;
+        e->arena = (uint8_t *) VirtualAlloc(nullptr, per * (size_t) tt.dim[2], MEM_RESERVE, PAGE_READWRITE);
+        e->st = (volatile LONG *) VirtualAlloc(nullptr, sizeof(LONG) * (size_t) tt.dim[2], MEM_COMMIT, PAGE_READWRITE);
+        e->cvlk = new CRITICAL_SECTION[(size_t) tt.dim[2]];
+        e->cv = new CONDITION_VARIABLE[(size_t) tt.dim[2]];
+        for (size_t z = 0; z < (size_t) tt.dim[2]; ++z) { InitializeCriticalSection(&e->cvlk[z]); InitializeConditionVariable(&e->cv[z]); }
+        if (!e->arena || !e->st) { delete[] e->cvlk; delete[] e->cv; e->cvlk = nullptr; e->cv = nullptr; delete e; continue; }
+        g_ptt_by_name[nm] = e;
+    }
+    g_pool_on = !g_ptt_by_name.empty();
+    if (g_pool_on) ggml_edge0_set_mmid_resolver(e0_mmid_base_res);   // attach to patch surface #7
+    { int k = 0; for (auto & q : g_ptt_by_name) if (k++ < 3) fprintf(stderr, "[pref-trace] POOL2 key%d=%s per=%zu\n", k, q.first.c_str(), q.second->per); }
+    fprintf(stderr, "[pref-trace] POOL2 init: tt=%zu caps hot=%lluMB obs=%lluMB resolver=%s\n",
+            g_ptt_by_name.size(), (unsigned long long) (g_pool_cap[0] >> 20),
+            (unsigned long long) (g_pool_cap[1] >> 20), g_pool_on ? "on" : "FAIL");
+    fflush(stderr);
+}
+
+}  // namespace
+
+static void enable_capture(llama_context * ctx) {
+    for (int L = OWNER0; L <= OWNER0 + NH; ++L) llama_set_embeddings_layer_inp(ctx, (uint32_t) L, true);
+}
+
+// ── Pool re-touch ladder: deterministic accounting (zero prediction = sticky; + observation zone = sticky/obs) ──
+// Mechanism: VirtualLock needs SeLockMemoryPrivilege, which consumer machines lack ⇒ pinning
+// is NOT the product mechanism. What works unprivileged = demand touch + working-set hard cap
+// + the OS clock's natural LRU:
+//   "pin"   = worker volatile-reads the range in 4K steps (demands pages into the WS; any
+//             trimmed page's fault tax is paid inside the worker, off the critical path);
+//   "evict" = drop from the table only (stop re-touching → reference-bit ages → OS trims it
+//             into standby = free soft layer);
+//   "cap"   = at init, SetProcessWorkingSetSizeEx(MAX = baseline + hot + obs + headroom, HARDWS_MAX_ENABLE).
+// Hot zone = re-touch of the previous step's teacher ids (same experts in a row = zero-fault
+// hit); observation zone = head-guess quota (a bad guess only ages its own pages — the hot
+// zone re-touches every step and never ages = contamination ceiling). Low-memory notification
+// ⇒ shrink the tables (hot to 75%, observation to empty).
+static volatile char g_touch_sink;                  // prevent DCE (every page is really read)
+static void e0_touch(const void * view, size_t off, size_t len) {
+    const volatile char * p = (const volatile char *) view + off;
+    char s = 0;
+    for (size_t i = 0; i < len; i += 4096) s += p[i];
+    g_touch_sink = s;
+}
+static void e0_pool_evict(int R) {                  // evict the oldest entry in region R (age==0 entries are exempt)
+    int wc = -1, wi = -1, ba = 0;
+    for (auto & pc : g.pins)
+        for (int k = 0; k < (int) pc.second.size(); ++k) {
+            auto & p = pc.second[k];
+            if (p.reg == R && p.age > ba) { ba = p.age; wc = pc.first; wi = k; }
+        }
+    if (wc < 0) return;
+    g.reg_bytes[R] -= g.pins[wc][wi].len; ++g.evicts;
+    g.pins[wc].erase(g.pins[wc].begin() + wi);      // dropping from the table = stop re-touching; physical reclaim is the OS clock's job
+}
+static void e0_issue_pool() {
+    if (!g.view) return;
+    const int nreg = g.pool == 2 ? 2 : 1;
+    std::map<int, std::vector<std::pair<size_t, size_t>>> want[2];   // 0=hot(teacher) 1=observation(head)
+    auto ranges = [&](const std::map<int, std::vector<int>> & src, int R) {
+        for (auto & pc : src) {
+            auto & v = want[R][pc.first];
+            for (const TT & tt : g.exps[pc.first].t) {
+                const size_t per = tt_bytes(tt) / EK;
+                for (int e : pc.second)
+                    v.push_back({ (size_t) g.data_base + tt.off + (size_t) e * per, per });
+            }
+        }
+    };
+    ranges(g.exec_prev, 0);
+    if (nreg > 1) {
+        std::map<int, std::vector<int>> preds;
+        { std::lock_guard<std::mutex> lk(g.wmtx); preds = g.pred_prev; }
+        ranges(preds, 1);
+    }
+    for (auto & kv : g.pins) for (auto & p : kv.second) ++p.age;     // uniform step-rate aging
+    std::vector<std::pair<size_t, size_t>> touch;                     // this step's demand set (≤ table cap)
+    for (int R = 0; R < nreg; ++R) {
+        for (auto & wv : want[R]) {
+            auto & cur = g.pins[wv.first];
+            std::vector<char> hit(wv.second.size(), 0);
+            for (size_t i = 0; i < wv.second.size(); ++i)
+                for (auto & p : cur)
+                    if (p.reg == R && p.off == wv.second[i].first) { p.age = 0; hit[i] = 1; if (!R) ++g.pool_hit; break; }
+            if (!R) g.pool_tot += (long) wv.second.size();           // denominator = all hot-zone demand this step
+            for (size_t i = 0; i < wv.second.size(); ++i) {
+                if (hit[i]) continue;
+                auto & w = wv.second[i];
+                while (g.reg_bytes[R] + w.second > g.reg_cap[R]) {
+                    size_t b = g.reg_bytes[R];
+                    e0_pool_evict(R);
+                    if (g.reg_bytes[R] == b) { ++g.lock_fail; break; }  // nothing evictable ⇒ hard ceiling (refused-pin counter)
+                }
+                if (g.reg_bytes[R] + w.second > g.reg_cap[R]) { ++g.lock_fail; continue; }
+                cur.push_back({ w.first, w.second, 0, (uint8_t) R });
+                g.reg_bytes[R] += w.second;
+            }
+        }
+    }
+    for (auto & pc : g.pins) for (auto & p : pc.second) e0_touch(g.view, p.off, p.len);  // re-touch the whole table
+    typedef HANDLE(WINAPI * QMRN)(LPCWSTR, void *);                   // memory-pressure cooperation (table-side fast shrink)
+    typedef BOOL(WINAPI * GMRN_T)(HANDLE, PBOOL);
+    static QMRN qmrn = (QMRN) GetProcAddress(GetModuleHandleA("kernel32"), "QueryMemoryResourceNotification");
+    static GMRN_T gmrn = (GMRN_T) GetProcAddress(GetModuleHandleA("kernel32"), "GetMemoryResourceNotification");
+    static HANDLE hn = qmrn ? qmrn(L"MemoryLow", nullptr) : nullptr;
+    if (hn && gmrn) {
+        BOOL low = FALSE;
+        if (gmrn(hn, &low) && low) {
+            ++g.shrinks;
+            for (int R = 0; R < nreg; ++R)
+                while (g.reg_bytes[R] && (R || g.reg_bytes[R] * 4 > g.reg_cap[R] * 3)) {
+                    size_t b = g.reg_bytes[R];
+                    e0_pool_evict(R);
+                    if (g.reg_bytes[R] == b) break;
+                }
+        }
+    }
+}
+
+// ── Advisory issue: PVM soft warm-up + VirtualLock hard retention (lock failure once the
+// budget is hit = natural ceiling). Callable from the worker thread or the main thread
+// (synchronous / stride-skipped); pred_prev is copied under a short wmtx lock.
+static void e0_issue() {
+    if (g.pool) { e0_issue_pool(); return; }          // pool accounting takes over (legacy pin/prefetch path retired)
+    if (!g.view || (!g.do_prefetch && !g.do_pin)) return;
+    std::map<int, std::vector<int>> preds;
+    { std::lock_guard<std::mutex> lk(g.wmtx); preds = g.pred_prev; }
+    static WIN32_MEMORY_RANGE_ENTRY win[4 * 3 * NH];
+    int nw = 0;
+    for (auto & pc : preds) {
+        const int c = pc.first;
+        std::vector<std::pair<size_t, size_t>> want;
+        for (const TT & tt : g.exps[c].t) {
+            const size_t per = tt_bytes(tt) / EK;
+            for (int e : pc.second) want.push_back({ (size_t) g.data_base + tt.off + (size_t) e * per, per });
+        }
+        if (g.do_pin) {
+            auto & cur = g.pinned[c];
+            for (auto & r : cur) {                                   // release stale pins
+                bool keep = false;
+                for (auto & w : want) if (w.first == r.first) { keep = true; break; }
+                if (!keep) { VirtualUnlock((char *) g.view + r.first, r.second); g.pin_bytes -= r.second; }
+            }
+            for (auto & w : want) {                                  // pin new predictions (skipped once capped)
+                bool have = false;
+                for (auto & r : cur) if (r.first == w.first) { have = true; break; }
+                if (have || g.pin_bytes + w.second > g.pin_cap) continue;
+                if (VirtualLock((char *) g.view + w.first, w.second)) g.pin_bytes += w.second;
+            }
+            cur = want;   // remember current pin state (entries skipped by the cap unlock harmlessly next round)
+        }
+        if (g.do_prefetch)
+            for (auto & w : want) {
+                if (nw == 4 * 3 * NH) break;
+                win[nw].VirtualAddress = (PVOID) ((char *) g.view + w.first);
+                win[nw].NumberOfBytes  = w.second;
+                ++nw;
+            }
+    }
+    if (nw) {
+        static decltype(&PrefetchVirtualMemory) pvm =
+            (decltype(&PrefetchVirtualMemory)) GetProcAddress(GetModuleHandleA("kernel32"), "PrefetchVirtualMemory");
+        if (pvm) pvm(GetCurrentProcess(), nw, win, 0);
+    }
+}
+
+// ── Full-step computation: (1) teacher (2) three-way accounting (3) head → preds + roll history (4) issue ──
+// rows = snapshot buffer, row order = layers OWNER0..OWNER0+NH, HH floats per row.
+// This function has no ctx dependency ⇒ it can run on the worker thread (keeps the ~90ms
+// head GEMV off the critical path).
+static void e0_full_step(const float * rows) {
+    std::map<int, std::vector<int>> teacher;
+    for (int c : consumers) {
+        const float * x = rows + (size_t) (c - OWNER0) * HH;
+        std::vector<float> tl(EK);
+        const float * gw = g.gate[c].data();
+        for (int e = 0; e < EK; ++e) tl[e] = dotf(gw + (size_t) e * HH, x, HH);
+        topk(tl.data(), EK, KDIM, teacher[c]);
+    }
+    // Three-way settlement: pred copy (short lock) vs teacher @ this snapshot.
+    std::map<int, std::vector<int>> preds_copy;
+    { std::lock_guard<std::mutex> lk(g.wmtx); preds_copy = g.pred_prev; }
+    for (int c : consumers) {
+        auto pp = preds_copy.find(c);
+        if (pp != preds_copy.end()) {
+            g.hits += ovlp(pp->second, teacher[c]); g.total += KDIM;
+            auto sp = g.exec_prev.find(c);
+            if (sp != g.exec_prev.end()) g.sticky += ovlp(sp->second, teacher[c]);
+            std::vector<int> r;
+            for (int tries = 0; tries < 64 && (int) r.size() < KDIM; ++tries) {
+                int v = (int) (g.next() % (uint32_t) EK);
+                bool dup = false; for (int y : r) if (y == v) { dup = true; break; }
+                if (!dup) r.push_back(v);
+            }
+            g.rnd += ovlp(r, teacher[c]);
+        }
+        g.exec_prev[c] = teacher[c];
+    }
+    if (g_pool_on) for (auto & kv : g.exec_prev) e0_pool_prefill(kv.first, kv.second, 0);   // teacher → hot zone
+    // head: owner o features = [m_in|cur_oh|prev_oh] → new pred (consumer = o+1)
+    std::map<int, std::vector<int>> new_preds;
+    std::vector<float> feat(FE), act(512), lg(EK);
+    for (int o = OWNER0; o < OWNER0 + NH; ++o) {
+        const int c = o + 1;
+        auto & h = g.hist[o];
+        memcpy(feat.data(),          rows + (size_t) (o - OWNER0) * HH, HH * 4);
+        memcpy(feat.data() + HH,     h.cur.data(),  EK * 4);
+        memcpy(feat.data() + HH + EK, h.prev.data(), EK * 4);
+        const Head & hd = g.heads[o];
+        for (int j = 0; j < 512; ++j) act[j] = gelu_erf(dotf(hd.fc1.data() + (size_t) j * FE, feat.data(), FE));
+        for (int e = 0; e < EK; ++e)
+            lg[e] = dotf(hd.fc2.data() + (size_t) e * 512, act.data(), 512)
+                  + dotf(hd.lin.data() + (size_t) e * FE, feat.data(), FE);
+        topk(lg.data(), EK, KDIM, new_preds[c]);
+    }
+    if (g_pool_on && g_pool_cap[1]) for (auto & kv : new_preds) e0_pool_prefill(kv.first, kv.second, 1);  // head → observation zone
+    { std::lock_guard<std::mutex> lk(g.wmtx); g.pred_prev = std::move(new_preds); }
+    // roll the onehot history (teacher→cur, cur→prev)
+    for (auto & kv : g.hist) {
+        auto t = teacher.find(kv.first);
+        if (t == teacher.end()) continue;
+        kv.second.prev = kv.second.cur;
+        kv.second.cur.assign(EK, 0.f);
+        for (int e : t->second) kv.second.cur[e] = 1.f;
+    }
+    e0_issue();
+    if (g.trace_every && g.steps % g.trace_every == 0 && g.total) {
+        fprintf(stderr, "[pref-trace] step %ld: head %.1f%% sticky %.1f%% rand %.1f%% pinned=%lluMB (n=%ld)\n",
+                g.steps, 100.0 * g.hits / g.total, 100.0 * g.sticky / g.total,
+                100.0 * g.rnd / g.total, (unsigned long long) (g.pin_bytes >> 20), g.total);
+        if (g.pool)
+            fprintf(stderr, "[pref-trace]  POOL hit=%.1f%% (n=%ld) pins %llu+%lluMB evict=%ld fail=%ld shrink=%ld\n",
+                    100.0 * g.pool_hit / (g.pool_tot ? g.pool_tot : 1), g.pool_tot,
+                    (unsigned long long) (g.reg_bytes[0] >> 20), (unsigned long long) (g.reg_bytes[1] >> 20),
+                    g.evicts, g.lock_fail, g.shrinks);
+        fflush(stderr);
+    }
+}
+
+static void e0_worker() {
+    for (;;) {
+        {
+            std::unique_lock<std::mutex> lk(g.wmtx);
+            g.cv.wait(lk, [&] { return g.quit.load() || g.busy.load(); });
+            if (g.quit) return;
+        }
+        if (g.busy.load()) {
+            e0_full_step(g.snap.data());
+            g.busy = false;
+        }
+    }
+}
+
+bool edge0_pref_init(llama_model * model, llama_context * ctx, int pool_mb, const char * gguf_path) {
+    (void) model;
+    if (g.active) { if (!g.gate.empty()) enable_capture(ctx); return true; }   // hot reload (only head mode needs capture)
+    const char * side = getenv("E0_PREROUTER");
+    const char * gg   = (gguf_path && *gguf_path) ? gguf_path : getenv("E0_GGUF");   // primary CLI channel = model path
+    if (pool_mb > 0) g.pool_mb = pool_mb;                                        // --pool-mb formal CLI
+    else if (const char * pmb = getenv("E0_POOL_MB")) g.pool_mb = atoll(pmb);    // env fallback (harness/debug); runnable without the sidecar
+    const bool headson = side && *side;
+    if ((!headson && !g.pool_mb) || !gg || !*gg) return false;
+
+    if (consumers.empty()) for (int c = OWNER0 + 1; c <= OWNER0 + NH; ++c) consumers.push_back(c);  // 7..39
+
+    // ── GGUF header: skip-scan KVs (for general.alignment) + build the tensor table ──
+    FILE * f = fopen(gg, "rb");
+    if (!f) return false;
+    struct { uint32_t magic, ver; uint64_t nt, nkv; } h;
+    if (fread(&h, sizeof h, 1, f) != 1 || memcmp(&h.magic, "GGUF", 4)) { fclose(f); return false; }
+    size_t alignment = 32;
+    for (uint64_t i = 0; i < h.nkv; ++i) {
+        std::string k = rd_str(f);
+        uint32_t t; if (fread(&t, 4, 1, f) != 1) { fclose(f); return false; }
+        if (k == "general.alignment" && t == 4) fread(&alignment, 4, 1, f);
+        else skip_val(f, t);
+    }
+    std::map<std::string, TT> tinfo;
+    for (uint64_t i = 0; i < h.nt; ++i) {
+        TT x; x.name = rd_str(f);
+        uint32_t nd; fread(&nd, 4, 1, f);
+        for (uint32_t d = 0; d < nd && d < 4; ++d) fread(&x.dim[d], 8, 1, f);
+        fread(&x.type, 4, 1, f); fread(&x.off, 8, 1, f);
+        tinfo[x.name] = x;
+    }
+    g.data_base = align_up((uint64_t) ftell(f), alignment);
+
+    // Register exps metadata (needed by both pool and head) + copy the F32 gate weights host-side (head mode only).
+    for (int c : consumers) {
+        for (const char * suf : {"ffn_up_exps", "ffn_gate_exps", "ffn_down_exps"}) {
+            auto et = tinfo.find("blk." + std::to_string(c) + "." + suf + ".weight");
+            if (et != tinfo.end() && et->second.dim[2] == (uint64_t) EK) g.exps[c].t.push_back(et->second);
+        }
+        if (!headson) continue;
+        auto it = tinfo.find("blk." + std::to_string(c) + ".ffn_gate_inp.weight");
+        if (it == tinfo.end() || it->second.type != 0) {
+            fclose(f); fprintf(stderr, "[pref-trace] gate missing/bad type c=%d => dormant\n", c); return false;
+        }
+        const TT & gt = it->second;
+        size_t nb = tt_bytes(gt);
+        std::vector<float> buf(nb / 4);
+        _fseeki64(f, (long long) (g.data_base + gt.off), SEEK_SET);
+        if (fread(buf.data(), 1, nb, f) != nb) { fclose(f); return false; }
+        g.gate[c] = std::move(buf);
+    }
+    fclose(f);
+
+    // sidecar head weights (head mode only)
+    if (headson) {
+    std::map<std::string, SEnt> st;
+    if (!st_load(side, st)) { fprintf(stderr, "[pref-trace] sidecar parse failed => dormant\n"); return false; }
+    FILE * sf = fopen(side, "rb");
+    if (!sf) return false;
+    uint64_t shn = 0; fread(&shn, 8, 1, sf);
+    uint64_t sdata = 8 + shn;
+    for (int o = OWNER0; o < OWNER0 + NH; ++o) {
+        Head hd;
+        auto key = [&](const char * w) { return "layers." + std::to_string(o) + "." + w + ".weight"; };
+        bool ok = grab_f16(sf, sdata, st[key("fc1")], hd.fc1)
+               && grab_f16(sf, sdata, st[key("fc2")], hd.fc2)
+               && grab_f16(sf, sdata, st[key("linear_init")], hd.lin);
+        if (!ok) { fclose(sf); fprintf(stderr, "[pref-trace] sidecar missing layer o=%d => dormant\n", o); return false; }
+        g.heads[o] = std::move(hd);
+    }
+    fclose(sf);
+    }   // headson
+
+    // Dedicated read handle (pool fills via g.fh) + a self-owned read-only mapping (PVM anchor; head mode only).
+    if (headson || g.pool_mb) {
+        g.fh = CreateFileA(gg, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
+                           NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
+        if (g.fh != INVALID_HANDLE_VALUE && headson) {
+            g.mh = CreateFileMappingA(g.fh, NULL, PAGE_READONLY, 0, 0, NULL);
+            if (g.mh) g.view = MapViewOfFile(g.mh, FILE_MAP_READ, 0, 0, 0);
+        }
+    }
+    g.do_prefetch = getenv("E0_PREFETCH") != nullptr;
+    if (const char * tn = getenv("E0_PREF_TRACE_N")) g.trace_every = atoi(tn);
+    if (const char * sn = getenv("E0_PREF_STRIDE")) g.stride = atoi(sn) > 0 ? atoi(sn) : 1;
+    const char * an = getenv("E0_PREF_ASYNC");
+    g.async = !(an && *an == '0');                     // async is the default; =0 falls back to synchronous (tax isolation)
+    g.do_pin = getenv("E0_PREF_PIN") != nullptr;
+    if (const char * pm = getenv("E0_PIN_MAX_MB")) g.pin_cap = (size_t) atoll(pm) << 20;
+    if (const char * pl = getenv("E0_POOL")) {                      // sticky=zone1  obs=zone1+2
+        if (!strcmp(pl, "obs")) g.pool = 2;
+        else if (*pl && strcmp(pl, "0")) g.pool = 1;
+    }
+    if (g.pool) {
+        g.reg_cap[0] = g.pin_cap;                                   // hot zone = E0_PIN_MAX_MB (default 3G)
+        const char * om = getenv("E0_POOL_OBS_MB");
+        g.reg_cap[1] = (size_t) (om ? atoll(om) : 1024) << 20;      // observation zone default 1G
+        g.do_pin = false;                                           // the advisory pin path does not run in parallel
+        // Working-set hard cap = startup baseline WS + pool quotas + headroom (the residency contract without SeLockMemory).
+        PROCESS_MEMORY_COUNTERS_EX pmc = {};
+        pmc.cb = sizeof(pmc);
+        if (K32GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS *) &pmc, sizeof(pmc))) {
+            const char * hm = getenv("E0_POOL_WSHEAD_MB");
+            size_t head = (size_t) (hm ? atoll(hm) : 1024) << 20;
+            SIZE_T wmax = (SIZE_T) pmc.WorkingSetSize + g.reg_cap[0] + g.reg_cap[1] + head;
+            SIZE_T qmin = 0, qmax = 0; DWORD qf = 0;      // the Ex variant rejects min=-1 (error 87) ⇒ read current quotas and keep them
+            GetProcessWorkingSetSizeEx(GetCurrentProcess(), &qmin, &qmax, &qf);
+            BOOL r = SetProcessWorkingSetSizeEx(GetCurrentProcess(), qmin, wmax,
+                                                QUOTA_LIMITS_HARDWS_MIN_DISABLE | QUOTA_LIMITS_HARDWS_MAX_ENABLE);
+            fprintf(stderr, "[pref-trace] POOL ws-cap: base=%lluMB max=%lluMB set=%d err=%lu\n",
+                    (unsigned long long) (pmc.WorkingSetSize >> 20), (unsigned long long) (wmax >> 20),
+                    (int) r, r ? 0UL : GetLastError());
+        }
+    }
+    g.snap.assign((size_t) (NH + 1) * HH, 0.f);
+
+    // Open the layer-input capture channel (same staging API as speculative decode; head mode only).
+    if (headson) {
+        enable_capture(ctx);
+        for (int L = OWNER0; L <= OWNER0 + NH; ++L) g.hist[L] = Hist{ std::vector<float>(EK, 0.f), std::vector<float>(EK, 0.f) };
+    }
+    if (g.pool_mb) {
+        g_pool_cap[0] = (size_t) g.pool_mb << 20;
+        if (const char * pom = getenv("E0_POOL_OBS_MB")) g_pool_cap[1] = (size_t) atoll(pom) << 20;
+        g_pool_sc = getenv("E0_POOL_SELFCHECK") != nullptr;
+        if (g_pool_sc && !g.view) {      // self-check needs the view: map it even in pool-only mode
+            g.mh = CreateFileMappingA(g.fh, NULL, PAGE_READONLY, 0, 0, NULL);
+            if (g.mh) g.view = MapViewOfFile(g.mh, FILE_MAP_READ, 0, 0, 0);
+        }
+        e0_pool_scan(tinfo);
+        for (int c : consumers)
+            for (const TT & tt : g.exps[c].t) {
+                auto it = g_ptt_by_name.find(tt.name);
+                if (it != g_ptt_by_name.end()) g_ptt_by_c[c].push_back(it->second);
+            }
+    }
+    g.active = headson || g_pool_on;
+    if (g.active && !headson) { g.async = false; }        // pool-only mode: no worker (demand is self-sufficient)
+    if (g.active && headson && g.async) g.worker = std::thread(e0_worker);
+    fprintf(stderr, "[pref-trace] INIT ok: heads=%zu gates=%zu exps=%zu map=%s prefetch=%s pin=%s(async=%d) stride=%d K=%d E=%d pool=%d(hot=%lluMB obs=%lluMB)\n",
+            g.heads.size(), g.gate.size(), g.exps.size(), g.view ? "ok" : "FAIL",
+            g.do_prefetch ? "on" : "off", g.do_pin ? "on" : "off", (int) g.async, g.stride, KDIM, EK,
+            (int) g_pool_on, (unsigned long long) (g_pool_cap[0] >> 20), (unsigned long long) (g_pool_cap[1] >> 20));
+    return g.active;
+}
+
+void edge0_pref_on_step(llama_context * ctx, int64_t last_row) {
+    if (!g.active) return;
+    if (g.gate.empty()) return;                 // pool-only mode: no capture channel, nothing to account — demand is autonomous
+    llama_synchronize(ctx);
+    ++g.steps;
+
+    if (g.stride > 1 && (g.steps % g.stride) != 0) {   // downsampled: reuse existing preds, just re-issue (µs scale)
+        if (!g.async && !g.pool) e0_issue();            // async=worker already issued on cadence; pool=skipped steps neither age nor evict
+        return;
+    }
+    if (g.async) {
+        if (g.busy.load()) return;                     // worker busy: drop this snapshot (preds already valid, lossless)
+        float * dst = g.snap.data();
+        for (int L = OWNER0; L <= OWNER0 + NH; ++L, dst += HH) {
+            float * buf = llama_get_embeddings_layer_inp(ctx, (uint32_t) L);
+            if (!buf) return;                          // channel not ready
+            memcpy(dst, buf + last_row * HH, (size_t) HH * 4);
+        }
+        { std::lock_guard<std::mutex> lk(g.wmtx); g.busy = true; }
+        g.cv.notify_one();
+        return;
+    }
+    // Synchronous path: snapshot + full step inline (reference/fallback).
+    float * dst = g.snap.data();
+    for (int L = OWNER0; L <= OWNER0 + NH; ++L, dst += HH) {
+        float * buf = llama_get_embeddings_layer_inp(ctx, (uint32_t) L);
+        if (!buf) return;
+        memcpy(dst, buf + last_row * HH, (size_t) HH * 4);
+    }
+    e0_full_step(g.snap.data());
+}
diff --git a/windows/serve/prefetch.h b/windows/serve/prefetch.h
new file mode 100644
index 0000000..3fe7e87
--- /dev/null
+++ b/windows/serve/prefetch.h
@@ -0,0 +1,50 @@
+// serve/prefetch.h — edge0 advisory expert-prefetch router.
+// Self-contained unit compiled into the llama target; depends only on llama
+// public/near-public APIs (llama.h + the src/llama-ext.h layer-input channel,
+// same route as speculative.cpp), linked into llama-server.
+// The real routing gate is never touched: the router head predicts, for step t+1,
+// the expert set of consumer layers from the owner layer's input (m_in) and a
+// two-step history, then issues prefetches for the matching GGUF byte ranges.
+// A miss costs latency only, never quality.
+//
+// Env gates (all dormant by default):
+//   E0_PREROUTER=<path to prerouter_edge0_*.safetensors>   enable head + advisory prefetch
+//   E0_PREFETCH=0                                          account only, no prefetch (A/B split)
+//   E0_PREF_TRACE_N=<k>                                    print [pref-trace] every k steps (default 32)
+//   E0_POOL_MB=<n>                                         hot-zone size MB of the private-page
+//                                                          L1 pool (served through patch #7's
+//                                                          expert-base resolver; can run without
+//                                                          the sidecar as a demand pool)
+//   E0_POOL_OBS_MB=<n>                                     observation-zone MB (meaningful only with heads on)
+//   E0_POOL=sticky|obs                                     route MUL_MAT_ID reads through the L1 pool
+//                                                          instead of plain mmap (sticky = re-touch the
+//                                                          teacher set at step boundaries; obs = add the
+//                                                          head-driven observation quota on top)
+//   E0_PIN_MAX_MB=<n>                                      hot-zone byte cap (default 3072)
+//   E0_POOL_WSHEAD_MB=<n>                                  headroom for the working-set hard cap
+//                                                          (cap = base + hot + obs + headroom; default 1024)
+// [pref-trace] prints head/sticky/random accounting = proof the mechanism is live.
+#pragma once
+
+#include <cstdint>
+
+// This unit is compiled into the llama target; when llama is built as a DLL it
+// must be exported so llama-server-impl can link against it.
+#ifdef _WIN32
+#  define E0_PREF_API __declspec(dllexport)
+#else
+#  define E0_PREF_API __attribute__((visibility("default")))
+#endif
+
+struct llama_model;
+struct llama_context;
+
+// Called once after server model+ctx are ready. Returns false = not enabled or
+// shape mismatch (stays dormant). pool_mb/gguf_path = formal CLI channels
+// (--pool-mb / model path); <= 0 / NULL falls back to the env channel (debug/harness).
+E0_PREF_API bool edge0_pref_init(llama_model * model, llama_context * ctx,
+                                 int pool_mb, const char * gguf_path);
+
+// Called after every llama_decode; last_row = the row index of this ubatch's last token
+// in the layer-input capture buffer.
+E0_PREF_API void edge0_pref_on_step(llama_context * ctx, int64_t last_row);
diff --git a/windows/tools/convert_mlx_to_gguf.py b/windows/tools/convert_mlx_to_gguf.py
new file mode 100644
index 0000000..94e7041
--- /dev/null
+++ b/windows/tools/convert_mlx_to_gguf.py
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+# -*- coding: utf8 -*-
+"""convert_mlx_to_gguf.py — single entry for on-device MLX→GGUF conversion.
+
+The app ships MLX artifacts only; this script (distributed with the app) runs the
+full conversion on the user's machine in three steps:
+  1) runtime GGUF generation (35b = repack_r3 / 8b = repack_r3_8b, incl. their
+     built-in full-tensor golden comparison gates)
+  2) LoRA adapter generation (lora_mlx_to_gguf, incl. name/shape cross-check gate)
+  3) on-device self-check + manifest (GGUF header/tensor-count revalidation +
+     sha256 of everything + converter/engine version lock-step triple)
+
+Idempotency: if manifest.json exists and every artifact sha matches ⇒ skip
+(resume granularity = whole steps).
+Failure: nonzero exit + partial artifacts kept + error recorded; MLX sources are
+never deleted (cleanup policy belongs to the app).
+Dependencies: python3 + numpy only (reuses the three tools/ scripts; gguf-py is
+not assumed present, so the header revalidation is hand-rolled on stdlib).
+
+Usage: python tools/convert_mlx_to_gguf.py --dir models/edge0-8b [--out models/edge0-8b-gguf]
+       [--compare-against <dir>]  # determinism proof: artifact shas compared 1:1 against same-named files in <dir>
+"""
+import argparse, hashlib, json, os, re, subprocess, sys, time
+
+TOOLS = os.path.dirname(os.path.abspath(__file__))
+CONVERTER_REV = subprocess.run(["git", "-C", os.path.dirname(TOOLS), "rev-parse", "--short", "HEAD"],
+                               capture_output=True, text=True).stdout.strip() or "unknown"
+ENGINE_REF = "edge0/lora-bailingmoe3@a49736d+"          # version lock-step: minimum compatible engine (after patch band 5)
+
+
+def sha256(path, blk=1 << 22):
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(blk), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def check_gguf_head(path, min_tensors):
+    """stdlib GGUF header revalidation: magic/version/n_tensors + KV section walkable + file size consistent."""
+    with open(path, "rb") as f:
+        if f.read(4) != b"GGUF":
+            raise ValueError(f"{path}: not GGUF")
+        ver = int.from_bytes(f.read(4), "little")
+        n_tensors = int.from_bytes(f.read(8), "little")
+        n_kv = int.from_bytes(f.read(8), "little")
+        if ver < 2 or n_tensors < min_tensors:
+            raise ValueError(f"{path}: ver={ver} n_tensors={n_tensors}")
+        def rd_s():
+            n = int.from_bytes(f.read(8), "little"); s = f.read(n); return s.decode("utf8")
+        def skip(t):
+            if t == 8: rd_s()
+            elif t == 9:
+                et = int.from_bytes(f.read(4), "little"); n = int.from_bytes(f.read(8), "little")
+                for _ in range(n): skip(et)
+            else:
+                f.read({0:1,1:1,2:2,3:2,4:4,5:4,6:4,7:1,10:8,11:8,12:8}[t])
+        for _ in range(n_kv):
+            rd_s(); t = int.from_bytes(f.read(4), "little"); skip(t)
+        total = 0
+        for _ in range(n_tensors):
+            nm = rd_s()
+            nd = int.from_bytes(f.read(4), "little")
+            dims = [int.from_bytes(f.read(8), "little") for _ in range(nd)]
+            dt = int.from_bytes(f.read(4), "little")
+            off = int.from_bytes(f.read(8), "little")
+            n = 1
+            for d in dims: n *= d
+            tsz = {0:4, 1:2, 2:2}.get(dt)
+            if tsz is None:                              # quantized block type (r3: 3 = Q4_1)
+                assert dt == 3, f"{nm}: unknown dtype {dt}"
+                assert dims[0] % 32 == 0, f"{nm}: Q4_1 ne0 % 32"
+                n = n // 32 * 20
+            total = max(total, off + n)
+        f.seek(0, 2)
+        if f.tell() < total:
+            raise ValueError(f"{path}: data section shorter than tensor table ({f.tell()}<{total})")
+        return n_tensors, n_kv
+
+
+def run_step(argv, tag, logf):
+    print(f"[{tag}] $ {' '.join(argv)}", flush=True)
+    r = subprocess.run(argv, stdout=logf, stderr=subprocess.STDOUT)
+    if r.returncode != 0:
+        raise SystemExit(f"[{tag}] FAILED rc={r.returncode} (see {logf.name})")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dir", required=True)
+    ap.add_argument("--out", default=None, help="default: <dir>-gguf")
+    ap.add_argument("--compare-against", default=None)
+    a = ap.parse_args()
+    d = os.path.normpath(a.dir)
+    out = a.out or (d + "-gguf")
+    stem = os.path.basename(d)
+    m = re.match(r"edge0[-_](35b|8b)", stem)
+    assert m, f"cannot detect tier from directory name: {stem}"
+    tier = m.group(1)
+    if a.out and os.path.normpath(a.out) != os.path.normpath(d + "-gguf"):
+        print("WARNING: --out deviates from the <dir>-gguf convention ⇒ the lora cross-check gate still reads the base from the conventional path (use the default layout on-device)", flush=True)
+    os.makedirs(out, exist_ok=True)
+    log = open(os.path.join(out, "convert.log"), "ab")
+    arts = {"gguf": os.path.join(out, f"edge0-{tier}.gguf"),
+            "lora": os.path.join(out, f"lora_edge0_{tier}-gguf.gguf")}
+    src_lora = os.path.join(d, f"lora_{stem.replace('-', '_')}-gguf.gguf")
+    manifest_p = os.path.join(out, "manifest.json")
+
+    # ---- idempotency gate ----
+    if os.path.isfile(manifest_p):
+        try:
+            old = json.load(open(manifest_p, encoding="utf8"))
+            ok = all(os.path.isfile(p) and sha256(p) == old["artifacts"][k] for k, p in arts.items())
+            if ok:
+                print("manifest hit and all artifact shas match ⇒ already converted, skipping (idempotent).", flush=True)
+                return 0
+        except Exception:
+            pass
+
+    t0 = time.time()
+    # determinism baselines captured BEFORE any overwrite (the lora original lives in the MLX dir and is regenerated by the conversion)
+    ref_base = {}
+    if a.compare_against:
+        ref = os.path.normpath(a.compare_against)
+        for k, p in arts.items():
+            for cand in (os.path.join(ref, os.path.basename(p)), os.path.join(d, os.path.basename(p))):
+                if os.path.isfile(cand):
+                    ref_base[k] = (cand, sha256(cand))
+                    break
+    # ---- step 1: runtime GGUF ----
+    if tier == "35b":
+        run_step([sys.executable, os.path.join(TOOLS, "repack_r3.py"), "--dir", d, "--out", out,
+                  "--verify", "all"], "r3-35b", log)
+        n_min = 700
+    else:
+        run_step([sys.executable, os.path.join(TOOLS, "repack_r3_8b.py"), "--dir", d, "--out", out],
+                 "r3-8b", log)
+        n_min = 500
+    # ---- step 2: adapter ----
+    run_step([sys.executable, os.path.join(TOOLS, "lora_mlx_to_gguf.py"), "--dir", d], "lora", log)
+    assert os.path.isfile(src_lora), f"adapter not produced: {src_lora}"
+    if os.path.abspath(src_lora) != os.path.abspath(arts["lora"]):
+        import shutil; shutil.copyfile(src_lora, arts["lora"])
+
+    # ---- step 3: on-device self-check ----
+    nt, nk = check_gguf_head(arts["gguf"], n_min)
+    print(f"[self-check] {os.path.basename(arts['gguf'])}: tensors={nt} kv={nk} OK", flush=True)
+    nta, _ = check_gguf_head(arts["lora"], 100)
+    print(f"[self-check] adapter tensors={nta} OK", flush=True)
+
+    man = {
+        "model": stem, "tier": tier,
+        "produced_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "elapsed_s": round(time.time() - t0, 1),
+        "converter": {"path": "tools/convert_mlx_to_gguf.py", "repo_rev": CONVERTER_REV},
+        "engine_min": ENGINE_REF,
+        "inputs_sha256": {k: sha256(os.path.join(d, k)) for k in sorted(os.listdir(d))
+                          if k.endswith((".safetensors",)) },
+        "artifacts": {k: sha256(p) for k, p in arts.items()},
+        "sizes": {k: os.path.getsize(p) for k, p in arts.items()},
+    }
+    if a.compare_against:
+        det = {k: (k in ref_base and ref_base[k][1] == man["artifacts"][k]) for k in arts}
+        man["determinism_vs"] = {"baselines": {k: ref_base[k][0] for k in ref_base}, **det}
+        assert all(det.get(k) for k in arts), f"determinism mismatch (artifact sha differs from baseline): {det}"
+        print(f"[determinism] vs baselines: ALL MATCH ✔", flush=True)
+    json.dump(man, open(manifest_p, "w", encoding="utf8"), ensure_ascii=False, indent=1)
+    print(f"DONE {tier}: gguf={man['sizes']['gguf']/2**30:.2f}GiB lora={man['sizes']['lora']/2**20:.1f}MiB "
+          f"took {man['elapsed_s']}s → {manifest_p}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/windows/tools/lora_mlx_to_gguf.py b/windows/tools/lora_mlx_to_gguf.py
new file mode 100644
index 0000000..226febf
--- /dev/null
+++ b/windows/tools/lora_mlx_to_gguf.py
@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""lora_mlx_to_gguf.py — edge0 LoRA (MLX safetensors, peft shape) → llama-adapter GGUF.
+
+Basis (verified against fork @7ab4ee7):
+  naming law: adapter tensor names must == base weight names + ".lora_a/.lora_b"
+              (llama-adapter.cpp:333 hard lookup)
+  shape law:  A declared dims = reversed(torch [r,K]) = ggml ne [K,r]; B = [r,N]
+              (:365-370 checks; torch row-major bytes untouched, only the dims
+              declaration flips)
+  scale law:  runtime scale = adapter.lora.alpha ? s·alpha/rank : s
+              (llama-adapter.h:53, rank = B.ne0)
+              source-of-truth scale = alpha/r = 32/16 = 2.0 ⇒ write
+              adapter.lora.alpha=32.0 so plain `--lora FILE` (default s=1.0) is the same expression
+  ssm_ba special tensor: the model side interleaves+merges in_proj_b/a
+              (qwen3next.cpp:421-438) and the two source LoRAs have different A matrices
+              ⇒ block-diagonal merge: A=[32,2048] (top 16 rows = A_b, bottom 16 = A_a),
+              B=[64,32] with rows placed per the interleave; rank=32 ⇒ default
+              scale=32/32=1.0 while the source needs 2.0 per path ⇒ **B entries ×2
+              compensation** (linear value scaling, NOT requantization).
+  arch law:   general.architecture must match the base (:212 hard compare).
+"""
+import argparse, os, re, struct, sys
+import numpy as np
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from repack_mlx_to_gguf import StIndex, GgufWriter, KV, T_F16, bf16_to_f32
+from repack_r3 import ssm_ba_perm
+
+RENAME = [
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.linear_attn\.in_proj_qkv$"), r"blk.\1.attn_qkv.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.linear_attn\.in_proj_z$"), r"blk.\1.attn_gate.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.linear_attn\.out_proj$"), r"blk.\1.ssm_out.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.self_attn\.q_proj$"), r"blk.\1.attn_q.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.self_attn\.k_proj$"), r"blk.\1.attn_k.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.self_attn\.v_proj$"), r"blk.\1.attn_v.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.self_attn\.o_proj$"), r"blk.\1.attn_output.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.mlp\.shared_expert\.gate_proj$"), r"blk.\1.ffn_gate_shexp.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.mlp\.shared_expert\.up_proj$"), r"blk.\1.ffn_up_shexp.weight"),
+    (re.compile(r"^language_model\.model\.layers\.(\d+)\.mlp\.shared_expert\.down_proj$"), r"blk.\1.ffn_down_shexp.weight"),
+]
+BA = re.compile(r"^language_model\.model\.layers\.(\d+)\.linear_attn\.in_proj_([ab])$")
+
+
+def f16(idx, name):
+    a = idx.load(name)
+    return np.ascontiguousarray(a, dtype=np.float16)
+
+
+RENAME_8B = [
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.q_proj$"), r"blk.\1.attn_q.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.k_proj$"), r"blk.\1.attn_k.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.v_proj$"), r"blk.\1.attn_v.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.o_proj$"), r"blk.\1.attn_output.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.dense$"), r"blk.\1.attn_output.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.f_proj$"), r"blk.\1.ssm_f_a.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.b_proj$"), r"blk.\1.ssm_beta.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.q_a_proj$"), r"blk.\1.attn_q_a.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.q_b_proj$"), r"blk.\1.attn_q_b.weight"),
+    # g_proj routes by layer type (KDA→ssm_g_a / MLA→attn_gate), decided via FULL() in main
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.g_proj$"), r"G"),
+    # kv_b_proj targets a 3D split tensor (unimplemented upstream-side); never touched by --only experiments; plain conversions skip it and report
+    (re.compile(r"^model\.layers\.(\d+)\.attention\.kv_b_proj$"), r"SKIP:kv_b_3d"),
+    (re.compile(r"^model\.layers\.(\d+)\.mlp\.gate_proj$"), r"blk.\1.ffn_gate.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.mlp\.up_proj$"), r"blk.\1.ffn_up.weight"),
+    (re.compile(r"^model\.layers\.(\d+)\.mlp\.down_proj$"), r"blk.\1.ffn_down.weight"),
+]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dir", required=True)
+    ap.add_argument("--alpha", type=float, default=32.0)   # = alpha baked in the source artifact; plain --lora (default scale 1.0) then reproduces the same expression
+    ap.add_argument("--only", default=None, help="generate only families whose base name matches this regex (single-variable experiments)")
+    ap.add_argument("--out-suffix", default="-gguf", help="infix in the output name")
+    a = ap.parse_args()
+    stem = os.path.basename(a.dir.rstrip("\\/"))
+    is8 = "8b" in stem
+    sfx = stem.replace("-", "_")          # artifact names use underscores: lora_edge0_35b.safetensors
+    lora_path = os.path.join(a.dir, f"lora_{sfx}.safetensors")
+    out = os.path.join(a.dir, f"lora_{sfx}{a.out_suffix}.gguf")
+    idx = StIndex([lora_path])
+    cfg = __import__("json").load(open(os.path.join(a.dir, "config.json"), encoding="utf8"))
+    tc = cfg.get("text_config", cfg)
+    n_v = tc.get("linear_num_value_heads", 32)
+    n_k = tc.get("linear_num_key_heads", 16)
+    only = re.compile(a.only) if a.only else None
+    skipped = []
+
+    simple, ba_pairs = {}, {}
+    for name in idx.t:
+        base = re.sub(r"\.lora_[AB]$", "", name)
+        if only and not only.search(base):
+            skipped.append(base)
+            continue
+        ab = "a" if name.endswith(".lora_A") else "b"
+        if not is8:
+            m = BA.match(base)
+            if m:
+                ba_pairs.setdefault((int(m.group(1)), m.group(2)), {})[ab] = name
+                continue
+        for rx, tpl in (RENAME_8B if is8 else RENAME):
+            mm = rx.match(base)
+            if mm:
+                if tpl.startswith("SKIP:"):
+                    skipped.append(base + " [3D split tensor unimplemented]")
+                    break
+                if tpl == "G":                  # 8b g_proj routes by layer type
+                    il = int(mm.group(1))
+                    dst = (f"blk.{il}.attn_gate.weight" if (il + 1) % 4 == 0
+                           else f"blk.{il}.ssm_g_a.weight")
+                else:
+                    dst = rx.sub(tpl, base)
+                simple.setdefault(dst, {})[ab] = name
+                break
+        else:
+            if "pregate" not in base and "mlp.gate" not in base:
+                raise SystemExit(f"unmapped LoRA family: {base}")
+    if skipped:
+        print("skip:", sorted(set(skipped))[:6], f"({len(skipped)} entries)")
+
+    entries = []   # (model_name, A_f16 [r,K] torch, B_f16 [N,r] torch)
+    for mname, d in sorted(simple.items()):
+        A = idx.load(d["a"]); B = idx.load(d["b"])
+        assert A.shape[0] == B.shape[1], (mname, A.shape, B.shape)
+        entries.append((mname, A, B))
+    # ssm_ba block-diagonal merge (with alpha/rank compensation: B ×(32/alpha·alpha/16→2))
+    comp = 2.0     # target scale per path = 2.0; merged rank=32 ⇒ get_scale=alpha·s/32=1.0 ⇒ B needs ×2
+    for il in sorted({k[0] for k in ba_pairs}):
+        db, da = ba_pairs[(il, "b")], ba_pairs[(il, "a")]
+        Ab, Bb = idx.load(db["a"]), idx.load(db["b"])
+        Aa, Ba = idx.load(da["a"]), idx.load(da["b"])
+        r = Ab.shape[0]
+        A = np.concatenate([Ab, Aa], 0)                       # [2r, K]
+        B = np.zeros((2 * n_v, 2 * r), dtype=np.float32)
+        for out_row, (tag, j) in enumerate(ssm_ba_perm(n_k, n_v)):
+            src_mat, src_row = (Bb, j) if tag == "b" else (Ba, j)
+            cols = slice(0, r) if tag == "b" else slice(r, 2 * r)
+            B[out_row, cols] = np.asarray(src_mat[src_row], dtype=np.float32) * comp
+        entries.append((f"blk.{il}.ssm_ba.weight", A, B.astype(np.float16)))
+
+    w = GgufWriter(out)
+    w.kv_s("general.architecture", "bailingmoe3" if is8 else "qwen3next")
+    w.kv_s("general.type", "adapter")
+    w.kv_s("general.name", f"{stem}-lora (repack of lora_{stem}.safetensors)")
+    w.kv_s("adapter.type", "lora")
+    w.kv_s("adapter.name", f"edge0-{stem}")
+    w.kv("adapter.lora.alpha", KV["f32"], struct.pack("<f", a.alpha))
+    for mname, A, B in entries:
+        w.register(mname + ".lora_a", T_F16, (A.shape[1], A.shape[0]), A.size * 2)
+        w.register(mname + ".lora_b", T_F16, (B.shape[1], B.shape[0]), B.size * 2)
+    w.finalize()
+    for mname, A, B in entries:
+        w.write(np.ascontiguousarray(A, np.float16).tobytes()); w.align_pad()
+        w.write(np.ascontiguousarray(B, np.float16).tobytes()); w.align_pad()
+    w.close(); idx.close()
+    print(f"adapter: {out}  tensors={2*len(entries)}  pairs={len(entries)} (incl. ssm_ba merge ×{comp} compensation)")
+    verify_against_base(out, "edge0-8b" if is8 else "edge0-35b")
+
+
+def verify_against_base(adapter_path, base_stem):
+    """Cross-check gate: right after writing, verify (1) adapter names ⊆ base names
+    (2) the upstream shape three-law (llama-adapter.cpp:358-371: a.ne0=base.ne0 ∧
+    b.ne1=base.ne1 ∧ a.ne1=b.ne0; token_embd flip is a special case).
+    Any miss/shape violation = hard throw, preventing the 'loads without error but
+    never enters the graph' failure mode from recurring."""
+    # gguf-py comes from the pinned vendored upstream (monorepo layout or bootstrap depot);
+    # standalone checkouts fall back to an installed `gguf` package on sys.path.
+    _here = os.path.dirname(os.path.abspath(__file__))
+    _repo = os.path.dirname(_here)
+    for _p in [os.path.join(x, "vendor", "llama.cpp", "gguf-py") for x in
+               ([os.environ["EDGE0_DEPOT"]] if os.environ.get("EDGE0_DEPOT") else [])
+               + [os.path.join(_repo, os.pardir), os.path.join(_repo, ".edge0")]]:
+        if os.path.isdir(_p):
+            sys.path.insert(0, _p)
+            break
+    from gguf import GGUFReader
+    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    base_path = os.path.join(root, "models", f"{base_stem}-gguf", f"{base_stem}.gguf")
+    bt = {str(t.name): [int(x) for x in t.shape]
+          for t in GGUFReader(base_path).tensors}
+    ad = GGUFReader(adapter_path)
+    A = {str(t.name)[:-7]: [int(x) for x in t.shape] for t in ad.tensors if str(t.name).endswith(".lora_a")}
+    B = {str(t.name)[:-7]: [int(x) for x in t.shape] for t in ad.tensors if str(t.name).endswith(".lora_b")}
+    miss, shape_bad = [], []
+    for stem in sorted(A.keys() & B.keys()):
+        b = bt.get(stem)
+        if b is None:
+            miss.append(stem); continue
+        a, bb = A[stem], B[stem]
+        if stem.startswith("token_embd"):
+            ok = b[0] == bb[1] and b[1] == a[1] and a[1] == bb[0]   # token_embd flip special case
+        else:
+            ok = b[0] == a[0] and b[1] == bb[1] and a[1] == bb[0]
+        if not ok:
+            shape_bad.append((stem, f"base={b} a={a} b={bb}"))
+    if miss or shape_bad:
+        raise SystemExit(f"cross-check gate RED: miss={len(miss)} {miss[:5]} shape_bad={len(shape_bad)} {shape_bad[:5]}")
+    print(f"cross-check gate GREEN: {len(A.keys() & B.keys())} pairs all hit base {os.path.basename(base_path)} + shape three-law satisfied")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/windows/tools/make_app_catalog.py b/windows/tools/make_app_catalog.py
new file mode 100644
index 0000000..86f74d7
--- /dev/null
+++ b/windows/tools/make_app_catalog.py
@@ -0,0 +1,38 @@
+# -*- coding: utf8 -*-
+"""tools/make_app_catalog.py — generator for the app's embedded tier table.
+Reads the two per-tier manifests (8b=path+revision / 35b=name+rev-missing) and
+normalizes them into catalog.json:
+  {tier: {repo, rev:{modelscope:"master", hf:<sha>}, base:{modelscope,hf,hf_mirror},
+          files:[{path,size,sha256,skip}], total_bytes, pool_mb:{35b:4096|8b:2048}}}
+skip set = tensors the GGUF route never consumes (media/prerouter/.gitattributes/README).
+Idempotent, safe to re-run."""
+import json
+import os
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+HF_SHA = {"8b": None, "35b": "e21098e7faa916a00f5493795f12f20497d032bb"}  # 8b falls back to manifest.revision
+SKIP = {".gitattributes", "README.md"}  # + per-tier media/prerouter added by the predicate below
+
+cat = {}
+for tier in ("8b", "35b"):
+    m = json.load(open(os.path.join(ROOT, "models", f"fetch-manifest-edge0-{tier}.json"), encoding="utf8"))
+    rev_hf = m.get("revision") or HF_SHA[tier]
+    files, total = [], 0
+    for f in m["files"]:
+        p = f.get("path") or f.get("name")
+        skip = p in SKIP or "prerouter" in p or p.endswith((".jpg", ".mp4"))
+        files.append(dict(path=p, size=f["size"], sha256=f["sha256"], skip=bool(skip)))
+        if not skip:
+            total += f["size"]
+    cat[tier] = dict(repo=m["repo"], rev=dict(modelscope="master", hf=rev_hf),
+                     files=files, total_bytes=total,
+                     pool_mb=(2048 if tier == "8b" else 4096))
+
+out = os.path.join(ROOT, "app", "src-tauri", "catalog.json")
+os.makedirs(os.path.dirname(out), exist_ok=True)
+json.dump(cat, open(out, "w", encoding="utf8"), ensure_ascii=False, indent=1)
+for t, c in cat.items():
+    print(t, c["repo"], "rev_hf=", c["rev"]["hf"][:12],
+          "pull=", sum(1 for f in c["files"] if not f["skip"]), "/", len(c["files"]),
+          f"bytes={c['total_bytes']:,}")
diff --git a/windows/tools/r3_bench.py b/windows/tools/r3_bench.py
new file mode 100644
index 0000000..79b4eb6
--- /dev/null
+++ b/windows/tools/r3_bench.py
@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""r3_bench.py — throughput measurement on the resident tier (48 GB, warm cache).
+
+Each run: 3 reps — long warmup case (a 512-token warmup rep is discarded first) →
+256-token greedy decode; parse prompt/eval tok/s from llama-completion; 5s idle
+between reps. Absolute values are only meaningful within the same machine and
+window when compared against the reference — cross-window comparison is always
+suspect (measure on the same day / same process, or treat numbers as shape-only).
+
+Usage: python tools/r3_bench.py --tier 35b [--reps 3]
+"""
+import argparse, json, os, re, subprocess, sys, time
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def resolve_bin_dir():
+    """Engine dir — same contract as the app (paths.rs::bin_dir): EDGE0_BIN_DIR override,
+    else monorepo depot build <repo>/../wt/win/..., else bootstrap <repo>/.edge0/wt/win/..."""
+    env = os.environ.get("EDGE0_BIN_DIR")
+    if env:
+        return env
+    for c in (os.path.join(ROOT, "..", "wt", "win", "build-vk", "bin", "Release"),
+              os.path.join(ROOT, ".edge0", "wt", "win", "build-vk", "bin", "Release")):
+        if os.path.isfile(os.path.join(c, "llama-completion.exe")):
+            return c
+    return os.path.join(ROOT, "..", "wt", "win", "build-vk", "bin", "Release")
+
+
+SERVER_BIN = resolve_bin_dir()
+
+
+def models_root(tier):
+    """Where the GGUF lives — repo-local models/ wins when present (dev layout), else the
+    app's model home (EDGE0_HOME, default ~/.edge0) where the download+convert step lands it."""
+    home = os.environ.get("EDGE0_HOME") or os.path.join(os.path.expanduser("~"), ".edge0")
+    roots = [os.path.join(ROOT, "models"), os.path.join(home, "models")]
+    for r in roots:
+        if os.path.isfile(os.path.join(r, f"edge0-{tier}-gguf", f"edge0-{tier}.gguf")):
+            return r
+    raise SystemExit(f"no edge0-{tier}.gguf under {roots[0]} or {roots[1]} "
+                     "— load/convert once in the app, or run tools/convert_mlx_to_gguf.py first")
+
+
+PROMPT = ("混合专家模型（MoE）通过稀疏激活把多个专家网络与门控路由组合起来：" * 12)
+PAT = {
+    "prompt": re.compile(r"prompt eval time =.*?,\s*([\d.]+) tokens per second"),
+    "eval": re.compile(r"(?<!prompt )eval time =.*?,\s*([\d.]+) tokens per second"),
+    "total_ms": re.compile(r"total time =\s*([\d.]+) ms"),
+    "load_ms": re.compile(r"load time =\s*([\d.]+) ms"),
+    "prompt_ms": re.compile(r"prompt eval time =\s*([\d.]+) ms"),
+}
+
+
+def one(gguf, lora, n, tier):
+    cmd = [os.path.join(SERVER_BIN, "llama-completion.exe"), "-m", gguf,
+           "-ngl", "99", "-cmoe", "--temp", "0", "-n", str(n), "--seed", "7",
+           "-no-cnv", "--no-display-prompt", "-p", PROMPT]
+    if lora:
+        cmd[3:3] = ["--lora", lora]
+    t0 = time.time()
+    r = subprocess.run(cmd, capture_output=True, text=True, encoding="utf8", errors="replace")
+    err = r.stderr or ""
+    out = {"wall_s": round(time.time() - t0, 1)}
+    for k, pat in PAT.items():
+        m = pat.search(err)
+        if m:
+            out[k] = float(m.group(1))
+    ev = re.search(r"n_ctx = (\d+)", err)
+    out["n_ctx"] = int(ev.group(1)) if ev else None
+    if "error" in err.lower() and "eval time" not in err:
+        out["error"] = [l for l in err.splitlines() if "error" in l.lower()][:2]
+    if "total_ms" in out and "prompt_ms" in out:
+        dec_ms = out["total_ms"] - out["prompt_ms"]        # total = generation end-to-end, excludes load (verified)
+        out["decode_tps_derived"] = round(n * 1000 / dec_ms, 2) if dec_ms > 0 else None
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--tier", choices=["35b", "8b"], required=True)
+    ap.add_argument("--reps", type=int, default=3)
+    a = ap.parse_args()
+    exe = os.path.join(SERVER_BIN, "llama-completion.exe")
+    if not os.path.isfile(exe):
+        raise SystemExit(f"llama-completion.exe not found in {SERVER_BIN} "
+                         "— build the engine (scripts/vendor-build.ps1) or set EDGE0_BIN_DIR")
+    md = models_root(a.tier)
+    gguf = os.path.join(md, f"edge0-{a.tier}-gguf", f"edge0-{a.tier}.gguf")
+    lora = os.path.join(md, f"edge0-{a.tier}",
+                        f"lora_edge0_{a.tier}-gguf.gguf")
+    rows = []
+    warm = one(gguf, lora, 64, a.tier)          # warmup (page cache + mmap fault-in), discarded
+    print(f"warm(discard): {warm}", flush=True)
+    for i in range(a.reps):
+        time.sleep(5)
+        m = one(gguf, lora, 256, a.tier)
+        m["rep"] = i
+        rows.append(m)
+        print(f"rep{i}: {m}", flush=True)
+    ok = [r for r in rows if "eval" in r]
+    evs = sorted(r["eval"] for r in ok)
+    prs = sorted(r["prompt"] for r in ok if "prompt" in r)
+    summ = dict(tier=a.tier, reps=len(ok), decode_tps_med=evs[len(evs) // 2] if evs else None,
+                decode_tps_min=evs[0] if evs else None, decode_tps_max=evs[-1] if evs else None,
+                prefill_tps_med=prs[len(prs) // 2] if prs else None, rows=rows,
+                note="48 GB tier, -cmoe -ngl99 with LoRA; compare against the same-machine reference")
+    outp = os.path.join(ROOT, "benchmarks", "r3", f"bench_{a.tier}.json")
+    os.makedirs(os.path.dirname(outp), exist_ok=True)
+    json.dump(summ, open(outp, "w", encoding="utf8"), ensure_ascii=False, indent=1)
+    print("BENCH", json.dumps({k: v for k, v in summ.items() if k != "rows"}, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/windows/tools/repack_r3.py b/windows/tools/repack_r3.py
new file mode 100644
index 0000000..c00907c
--- /dev/null
+++ b/windows/tools/repack_r3.py
@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""repack_r3.py — 35b → GGUF with llama.cpp canonical naming (the runtime artifact).
+
+Built on the lossless repack core (tools/repack_mlx_to_gguf.py, proven bit-exact),
+this script does three things and nothing else:
+  A. rename: MLX full names → qwen3next registered names
+     (per qwen3next.cpp:31-129 / llama-arch.cpp:434-518)
+  B. value transforms (each with its source of truth):
+     - ssm_a   = -exp(A_log)                          conversion/qwen.py:395-396, qwen3next.cpp:452
+     - ssm_ba  = in_proj_b/a row interleave [b,b,a,a]×16 groups   qwen3next.cpp:421-446 (v-heads contiguous within a group)
+     - conv1d  = [C,4,1] squeeze last axis → dims flipped to ne{4,C}
+               (row-major bytes unchanged, pure dims rewrite)
+     - everything else stored directly: norm (+1 is already baked into the source
+       of truth — NEVER add +1 a second time), gate int8→F32, BF16→F32
+  C. metadata + tokenizer:
+     expert_used_count=4, rope.dimension_count=64, recurrent_layers explicit array,
+     no MTP/vision/pregate tensors (an unregistered tensor = hard throw at load time),
+     tokens padded to 248320, pre=qwen2, add_bos=false, eos=im_end(248046),
+     byte_fallback=true
+
+Gates: tensor count / name-set identity against the registry; ba permutation
+reversible; Q4_1 sampled pack two-sided sha. Output: <stem>.gguf.
+"""
+import argparse, json, math, os, re, struct, sys, time
+import numpy as np
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from repack_mlx_to_gguf import (StIndex, GgufWriter, KV, T_F32, T_F16, T_Q4_1, ALIGN, QK41,
+                                DT_Q41, mlx_q_bytes, pack_q41_rows, deq_mlx_side, deq_gguf_side,
+                                sha16)
+import repack_mlx_to_gguf as core
+
+P = "language_model.model.layers.{}."
+EMB, LMH, ONORM = "language_model.model.embed_tokens", "language_model.lm_head", "language_model.model.norm"
+R3VER = "r3-1.0"
+
+
+def full_attn_layer(i, interval):
+    return (i + 1) % interval == 0
+
+
+def build_name_table(nl, interval):
+    """→ list[(r3_name, kind, src_key_or_pair)]"""
+    t = [("token_embd.weight", "q41", EMB + ".weight"),
+         ("output.weight", "q41", LMH + ".weight"),
+         ("output_norm.weight", "f32raw", ONORM + ".weight")]
+    for i in range(nl):
+        b = P.format(i)
+        rb = f"blk.{i}."
+        t += [(rb + "attn_norm.weight", "f32raw", b + "input_layernorm.weight"),
+              (rb + "post_attention_norm.weight", "f32raw", b + "post_attention_layernorm.weight")]
+        if full_attn_layer(i, interval):
+            t += [(rb + "attn_q.weight", "q41", b + "self_attn.q_proj.weight"),
+                  (rb + "attn_k.weight", "q41", b + "self_attn.k_proj.weight"),
+                  (rb + "attn_v.weight", "q41", b + "self_attn.v_proj.weight"),
+                  (rb + "attn_q_norm.weight", "f32raw", b + "self_attn.q_norm.weight"),
+                  (rb + "attn_k_norm.weight", "f32raw", b + "self_attn.k_norm.weight"),
+                  (rb + "attn_output.weight", "q41", b + "self_attn.o_proj.weight")]
+        else:
+            la = b + "linear_attn."
+            t += [(rb + "attn_qkv.weight", "q41", la + "in_proj_qkv.weight"),
+                  (rb + "attn_gate.weight", "q41", la + "in_proj_z.weight"),
+                  (rb + "ssm_ba.weight", "ba", (la + "in_proj_b.weight", la + "in_proj_a.weight")),
+                  (rb + "ssm_conv1d.weight", "conv", la + "conv1d.weight"),
+                  (rb + "ssm_dt.bias", "f32raw", la + "dt_bias"),
+                  (rb + "ssm_a", "sa", la + "A_log"),
+                  (rb + "ssm_norm.weight", "f32raw", la + "norm.weight"),
+                  (rb + "ssm_out.weight", "q41", la + "out_proj.weight")]
+        m = b + "mlp."
+        t += [(rb + "ffn_gate_inp.weight", "f32i8", m + "gate.weight"),
+              (rb + "ffn_gate_exps.weight", "q41", m + "switch_mlp.gate_proj.weight"),
+              (rb + "ffn_up_exps.weight", "q41", m + "switch_mlp.up_proj.weight"),
+              (rb + "ffn_down_exps.weight", "q41", m + "switch_mlp.down_proj.weight"),
+              (rb + "ffn_gate_shexp.weight", "q41", m + "shared_expert.gate_proj.weight"),
+              (rb + "ffn_up_shexp.weight", "q41", m + "shared_expert.up_proj.weight"),
+              (rb + "ffn_down_shexp.weight", "q41", m + "shared_expert.down_proj.weight"),
+              (rb + "ffn_gate_inp_shexp.weight", "f32i8_1d", m + "shared_expert_gate.weight")]
+    return t
+
+
+def ssm_ba_perm(n_k, n_v):
+    """group g row order = b[2g],b[2g+1],a[2g],a[2g+1] (v-heads contiguous within a group, qwen3next.cpp:421-446)."""
+    order = []
+    for g in range(n_k):
+        order += [("b", 2 * g), ("b", 2 * g + 1), ("a", 2 * g), ("a", 2 * g + 1)]
+    return order
+
+
+def load_ba(idx, kb, ka, n_k, n_v):
+    """two source tensors (int4 g64, rows=n_v) → interleaved (2n_v, K) q/s/b arrays."""
+    out = {}
+    for tag, key in (("b", kb), ("a", ka)):
+        w = idx.load(key); s = idx.load(key[:-7] + ".scales"); b = idx.load(key[:-7] + ".biases")
+        out[tag] = (mlx_q_bytes(w.reshape(n_v, -1)).astype(np.float32), s.reshape(n_v, -1), b.reshape(n_v, -1))
+    K = out["b"][0].shape[1]
+    q = np.empty((2 * n_v, K), np.uint8); S = np.empty((2 * n_v, K // 64), np.float32); Bs = np.empty_like(S)
+    for r, (tag, row) in enumerate(ssm_ba_perm(n_k, n_v)):
+        q[r] = out[tag][0][row].astype(np.uint8); S[r] = out[tag][1][row]; Bs[r] = out[tag][2][row]
+    return q, S, Bs, K
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dir", required=True); ap.add_argument("--out", required=True)
+    ap.add_argument("--verify", default="all")
+    ap.add_argument("--dry", action="store_true")
+    a = ap.parse_args()
+    cfg = json.load(open(os.path.join(a.dir, "config.json"), encoding="utf8"))["text_config"]
+    n_k, n_v = cfg["linear_num_key_heads"], cfg["linear_num_value_heads"]
+    nl, intv = cfg["num_hidden_layers"], cfg["full_attention_interval"]
+    table = build_name_table(nl, intv)
+    names = [t[0] for t in table]
+    assert len(names) == len(set(names)), "name-set collision"
+    if a.dry:
+        print(f"r3 registry entries={len(table)}  (3+40*2+10*6+30*8+40*8=703 expected)")
+        return
+    idx = StIndex([os.path.join(a.dir, f) for f in sorted(os.listdir(a.dir))
+                   if re.fullmatch(r"model(-\d+-of-\d+)?\.safetensors", f)])
+    os.makedirs(a.out, exist_ok=True)
+    stem = os.path.basename(a.dir.rstrip("\\/"))
+    gp = os.path.join(a.out, f"{stem}.gguf")
+    t0, audit = time.time(), dict(tool=R3VER, tier=stem, mode="r3", gates=[], errors=[])
+    w = GgufWriter(gp)
+
+    # ---- metadata (required read set + overrides) ----
+    kv = w.kv
+    w.kv_s("general.architecture", "qwen3next")
+    w.kv_s("general.name", "edge0-35b")
+    A = "qwen3next."      # this fork keys every KV with the arch prefix (gguf-py "{arch}.X" verified; the old "llm." spelling is wrong and will not load)
+    for k, ty, v in [
+        (A + "block_count", KV["u32"], nl), (A + "context_length", KV["u32"], cfg["max_position_embeddings"]),
+        (A + "embedding_length", KV["u32"], cfg["hidden_size"]), (A + "feed_forward_length", KV["u32"], cfg["moe_intermediate_size"]),
+        (A + "attention.head_count", KV["u32"], cfg["num_attention_heads"]),
+        (A + "attention.head_count_kv", KV["u32"], cfg["num_key_value_heads"]),
+        (A + "attention.key_length", KV["u32"], cfg["head_dim"]), (A + "attention.value_length", KV["u32"], cfg["head_dim"]),
+        (A + "attention.layer_norm_rms_epsilon", KV["f32"], cfg["rms_norm_eps"]),
+        (A + "rope.freq_base", KV["f32"], cfg["rope_parameters"]["rope_theta"]),
+        (A + "expert_count", KV["u32"], cfg["num_experts"]), (A + "expert_used_count", KV["u32"], 4),
+        (A + "expert_feed_forward_length", KV["u32"], cfg["moe_intermediate_size"]),
+        (A + "expert_shared_feed_forward_length", KV["u32"], cfg["shared_expert_intermediate_size"]),
+        (A + "ssm.conv_kernel", KV["u32"], cfg["linear_conv_kernel_dim"]),
+        (A + "ssm.state_size", KV["u32"], cfg["linear_key_head_dim"]),
+        (A + "ssm.group_count", KV["u32"], n_k), (A + "ssm.time_step_rank", KV["u32"], n_v),
+        (A + "ssm.inner_size", KV["u32"], cfg["linear_value_head_dim"] * n_v),
+        (A + "rope.dimension_count", KV["u32"], int(cfg["head_dim"] * cfg["partial_rotary_factor"])),
+        (A + "full_attention_interval", KV["u32"], intv),
+    ]:
+        kv(k, ty, struct.pack("<I", v) if ty in (KV["u32"], KV["i32"]) else struct.pack("<f", v))
+    kv("qwen3next.attention.recurrent_layers", KV["arr"],
+       struct.pack("<I", 7) + struct.pack("<Q", nl) + b"".join(b"\x00" if full_attn_layer(i, intv) else b"\x01" for i in range(nl)))
+
+    # ---- tokenizer (required: tokens/merges/model; vocab padded to 248320) ----
+    tkj = json.load(open(os.path.join(a.dir, "tokenizer.json"), encoding="utf8"))
+    vocab, merges = tkj["model"]["vocab"], tkj["model"]["merges"]
+    toks = [None] * cfg["vocab_size"]
+    for s, i in vocab.items():
+        toks[i] = s
+    for at in tkj["added_tokens"]:
+        toks[at["id"]] = at["content"]
+    for i in range(len(toks)):
+        if toks[i] is None:
+            toks[i] = f"<0x{255 - (i % 256):02X}>"      # placeholder (high ids the source of truth never produces)
+    tt = [1] * len(toks)
+    for at in tkj["added_tokens"]:
+        tt[at["id"]] = 3
+    def arr(key, et, elems):
+        kv(key, KV["arr"], struct.pack("<I", et) + struct.pack("<Q", len(elems)) + b"".join(elems))
+    w.kv_s("tokenizer.ggml.model", "gpt2")
+    arr("tokenizer.ggml.tokens", KV["str"], [w._str(s) for s in toks])
+    arr("tokenizer.ggml.merges", KV["str"], [w._str(" ".join(m)) for m in merges])
+    arr("tokenizer.ggml.scores", KV["f32"], [struct.pack("<f", 0.0) for _ in toks])
+    arr("tokenizer.ggml.token_type", KV["i32"], [struct.pack("<i", x) for x in tt])
+    w.kv("tokenizer.ggml.added_tokens_ids", KV["arr"],
+         struct.pack("<I", KV["u32"]) + struct.pack("<Q", len(tkj["added_tokens"]))
+         + b"".join(struct.pack("<I", x["id"]) for x in tkj["added_tokens"]))
+    for k, v in [("tokenizer.ggml.bos_token_id", 248044), ("tokenizer.ggml.eos_token_id", 248046),
+                 ("tokenizer.ggml.padding_token_id", 248044)]:
+        kv(k, KV["u32"], struct.pack("<I", v))
+    kv("tokenizer.ggml.add_bos_token", KV["bool"], b"\x00")
+    kv("tokenizer.ggml.add_eos_token", KV["bool"], b"\x00")
+    kv("tokenizer.ggml.byte_fallback", KV["bool"], b"\x01")
+    w.kv_s("tokenizer.ggml.pre", "qwen2")
+    w.kv_s("tokenizer.chat_template", open(os.path.join(a.dir, "chat_template.jinja"), encoding="utf8").read())
+    kv("tokenizer.ggml.eot_token_id", KV["u32"], struct.pack("<I", 248046))
+
+    # ---- tensors (two passes: A registers dims/size, B streams writes + gates) ----
+    def r3_spec(kind, src):
+        """→ (tt, dims_ne, size). Pure shape arithmetic, no data loaded."""
+        if kind == "ba":
+            K = idx.t[src[0]][4][-1] * 8
+            return T_Q4_1, (K, 2 * n_v), 2 * n_v * (K // QK41) * 20
+        shp = idx.t[src][4]; dt = idx.t[src][3]
+        if kind == "q41":
+            K = shp[-1] * 8
+            rows = int(np.prod(shp[:-1])) if len(shp) > 1 else 1
+            return T_Q4_1, (K,) + tuple(reversed(shp[:-1])), rows * (K // QK41) * 20
+        if kind == "f32i8":
+            K = shp[-1] * 4
+            rows = int(np.prod(shp[:-1])) if len(shp) > 1 else 1
+            return T_F32, (K,) + tuple(reversed(shp[:-1])), rows * K * 4
+        if kind == "f32i8_1d":
+            K = shp[-1] * 4
+            return T_F32, (K,), K * 4
+        if kind == "conv":
+            return T_F32, (shp[1], shp[0]), int(np.prod(shp)) * 4
+        if kind == "sa":
+            return T_F32, (shp[0],), shp[0] * 4
+        return T_F32, (tuple(reversed(shp)) if len(shp) > 1 else (shp[0],)), int(np.prod(shp)) * 4
+
+    for r3n, kind, src in table:
+        tt, dims, size = r3_spec(kind, src)
+        w.register(r3n, tt, dims, size)
+    w.finalize()
+
+    n_done = 0
+    for r3n, kind, src in table:
+        if kind == "q41":
+            base = src[:-7]
+            W = idx.load(src); S = idx.load(base + ".scales"); Bs = idx.load(base + ".biases")
+            lead = W.shape[:-1]; rows = int(np.prod(lead)) if lead else 1
+            Wf, Sf, Bsf = W.reshape(rows, -1), S.reshape(rows, -1), Bs.reshape(rows, -1)
+            blocks, exm = pack_q41_rows(mlx_q_bytes(Wf), Sf, Bsf, src)
+            w.write(blocks.tobytes()); w.align_pad()
+            # sentinel gate: first row depacked on both sides, sha compared; if that row contains an exempt block, downgrade to record
+            if rows:
+                a_side = deq_mlx_side(mlx_q_bytes(Wf[:1]), Sf[:1], Bsf[:1])
+                b_side = deq_gguf_side(blocks[:1])
+                same = bool(sha16(a_side.tobytes()) == sha16(b_side.tobytes()))
+                row0_ex = any(r == 0 for (r, _b) in exm)
+                st = "OK" if same else ("EXEMPT" if row0_ex else "FAIL")
+                audit["gates"].append(dict(t=r3n, sentinel=st))
+                if st == "FAIL":
+                    audit["errors"].append(f"{r3n}: sentinel row non-exempt mismatch")
+        elif kind in ("f32i8", "f32i8_1d"):
+            base = src[:-7]
+            W = idx.load(src); S = idx.load(base + ".scales"); Bs = idx.load(base + ".biases")
+            rows = int(np.prod(W.shape[:-1])) if W.ndim > 1 else 1
+            Wf, Sf, Bsf = W.reshape(rows, -1), S.reshape(rows, -1), Bs.reshape(rows, -1)
+            q8 = Wf.view(np.uint8).reshape(rows, Wf.shape[1] * 4)
+            v = deq_mlx_side(q8, Sf, Bsf)
+            w.write(np.ascontiguousarray(v, np.float32).tobytes()); w.align_pad()
+        elif kind == "f32raw":
+            x = idx.load(src)
+            w.write(np.ascontiguousarray(x, np.float32).tobytes()); w.align_pad()
+        elif kind == "sa":
+            x = idx.load(src)
+            v = -np.exp(np.asarray(x, dtype=np.float32))
+            v2 = np.array([-math.exp(float(t)) for t in x.ravel()], dtype=np.float32)  # libm second path, cross-checked against numpy
+            dv = np.abs(v - v2).max(initial=0.0)
+            audit["gates"].append(dict(t=r3n, exp_dual_dev=float(dv)))
+            if dv > 1e-5:
+                audit["errors"].append(f"{r3n}: -exp dual-path deviation {dv:.2e} > 1e-5")
+            w.write(np.ascontiguousarray(v, np.float32).tobytes()); w.align_pad()
+        elif kind == "conv":
+            x = idx.load(src)                      # [C,4,1]
+            assert x.shape[-1] == 1
+            y = np.ascontiguousarray(x.reshape(x.shape[0], x.shape[1]), dtype=np.float32)
+            w.write(y.tobytes()); w.align_pad()
+        elif kind == "ba":
+            kb, ka = src
+            q, S, Bs, K = load_ba(idx, kb, ka, n_k, n_v)
+            blocks, _ = pack_q41_rows(q, S, Bs, "ssm_ba")
+            w.write(blocks.tobytes()); w.align_pad()
+            # reversibility gate: undo the permutation and compare against the pre-interleave source q rows for exact equality
+            wb = idx.load(kb); wa = idx.load(ka)
+            qb = mlx_q_bytes(wb.reshape(n_v, -1)); qa = mlx_q_bytes(wa.reshape(n_v, -1))
+            rec = np.empty_like(q, dtype=np.uint8)
+            for r, (tag, row) in enumerate(ssm_ba_perm(n_k, n_v)):
+                rec[r] = (qb if tag == "b" else qa)[row]
+            ok = bool((rec == q).all())
+            audit["gates"].append(dict(t=r3n, ba_permute_reversible=ok))
+            if not ok:
+                audit["errors"].append(f"{r3n}: ba permutation not reversible")
+        n_done += 1
+        if n_done % 60 == 0 or n_done == len(table):
+            print(f"  {n_done}/{len(table)} {r3n:44s} {w.dlen/1e9:6.2f}GB {time.time()-t0:5.0f}s", flush=True)
+    w.close(); idx.close()
+    bad = [g for g in audit["gates"] if g.get("ok") is False or g.get("ba_permute_reversible") is False]
+    audit["n_tensors"] = n_done; audit["elapsed_s"] = round(time.time() - t0, 1)
+    v = "GREEN" if not bad and not audit["errors"] and n_done == len(table) else "RED"
+    json.dump(dict(verdict=f"R3-artifact {v}", audit=audit, size=os.path.getsize(gp)),
+              open(os.path.join(a.out, f"{stem}-audit.json"), "w"), ensure_ascii=False, indent=1)
+    print(f"\nGGUF {gp} {os.path.getsize(gp)/1e9:.2f}GB tensors={n_done}/{len(table)} ⇒ R3-artifact {v}"
+          f" (errors={audit['errors']})")
+    return 0 if v == "GREEN" else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)
diff --git a/windows/tools/repack_r3_8b.py b/windows/tools/repack_r3_8b.py
new file mode 100644
index 0000000..0a373ea
--- /dev/null
+++ b/windows/tools/repack_r3_8b.py
@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""repack_r3_8b.py — 8b (BailingMoeV3/Ling-3) → bailingmoe3 canonical-naming GGUF.
+
+The contract is the full reference table (conversion/bailingmoe3.py +
+llama-arch.cpp + bailingmoe3.cpp):
+  name strings (llama-arch.cpp table as-is): attn_q_a/attn_q_a_norm/attn_q_b/
+    attn_kv_a_mqa/attn_kv_a_norm/attn_k_b/attn_v_b/attn_gate/
+    ssm_conv1d_{q,k,v}/ssm_a(no .weight)/ssm_dt(.bias)/ssm_beta/ssm_f_a/ssm_g_a/
+    ssm_norm/exp_probs_b(.bias)
+  transforms (checked line-by-line against bailingmoe3.py):
+    conv [C,1,4] → reshape(1,C,1,4) → ne{4,1,C,1}, bytes unchanged, F32
+    ssm_a = torch.exp(A_log).reshape(-1,1)   ← NOTE +exp here (35b uses −exp; do not mix them up)
+    dt_bias [2048] elementwise → ssm_dt.bias stored directly as F32
+    kv_b_proj → split + transpose (view(16,256,512) → [128 noope | 128 v] → k transpose(1,2))
+      k_b/v_b are stored as **F32, not Q4_1**: the transpose scatters the quantization
+      grouping axis (dim0=128 spans 128 original rows), so lossless repack is
+      physically impossible; 6 layers x 8MB = 48MB cost, negligible.
+    f_proj→ssm_f_a; g_proj→(KDA)ssm_g_a /(MLA)attn_gate; dense(MLA)→attn_output
+    gate BF16→F32 stored directly (8b gate is NOT quantized, unlike 35b int8);
+    expert_bias→exp_probs_b.bias F32
+    all norms stored directly (bailing_hybrid sanitize has no +1 and no moveaxis;
+    the source of truth is the final value multiplied directly)
+  metadata: head_count_kv is a 24-element array ((i+1)%4==0 →1 else 0, bailingmoe3.py:55-56);
+    key_length=576(=kv_lora+qk_rope) value_length=128 key_length_mla=192 value_length_mla=128;
+    expert_gating_func=2(SIGMOID, llama-hparams.h:19) expert_weights_scale=2.5
+    (config routed_scaling_factor) expert_weights_norm=true; kda.head_dim=128
+    kda.safe_gate=true kda.gate_lower_bound=-5; leading_dense_block_count=1;
+    rope.dimension_count=64; **tokenizer.ggml.pre="bailingmoe2"** (verified by hashing
+    tokenizer.json against the fork, NOT qwen2); architecture string
+    general.architecture=bailingmoe3.
+pregate x66 stripped (sidecar feeds the prefetch router). Gates: tensor-count
+identity / sentinel sha / kv_b transform reversibility spot-check / exp dual-path.
+"""
+import argparse, json, math, os, re, struct, sys, time
+import numpy as np
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from repack_mlx_to_gguf import (StIndex, GgufWriter, KV, T_F32, T_F16, T_Q4_1, QK41, DT_Q41,
+                                mlx_q_bytes, pack_q41_rows, deq_mlx_side, deq_gguf_side, sha16)
+
+FULL = lambda i: (i + 1) % 4 == 0            # MLA layer (bailingmoe3.py:39-42 same rule)
+NL, E, KDA_DIM, HEADS = 24, 128, 2048, 16
+R3VER = "r3-8b-1.0"
+
+
+def q41_src(idx, name):
+    W = idx.load(name); base = name[:-7]
+    return W, idx.load(base + ".scales"), idx.load(base + ".biases")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dir", default="models/edge0-8b")
+    ap.add_argument("--out", default="models/edge0-8b-gguf")
+    a = ap.parse_args()
+    idx = StIndex([os.path.join(a.dir, f) for f in sorted(os.listdir(a.dir))
+                   if re.fullmatch(r"model(-\d+-of-\d+)?\.safetensors", f)])
+    cfg = json.load(open(os.path.join(a.dir, "config.json"), encoding="utf8"))
+    tkj = json.load(open(os.path.join(a.dir, "tokenizer.json"), encoding="utf8"))
+    A = "bailingmoe3."
+    P = "model.layers.{}."
+    ATT = P + "attention."
+    MLP = P + "mlp."
+
+    # ---- tensor table: (dst_name, kind, src); kind ∈ q41/f32raw/conv/sa/kvb ----
+    table = [("token_embd.weight", "q41", "model.word_embeddings.weight"),
+             ("output.weight", "q41", "lm_head.weight"),
+             ("output_norm.weight", "f32raw", "model.norm.weight")]
+    for i in range(NL):
+        at_, mp = ATT.format(i), MLP.format(i)
+        table += [(f"blk.{i}.attn_norm.weight", "f32raw", P.format(i) + "input_layernorm.weight"),
+                  (f"blk.{i}.ffn_norm.weight", "f32raw", P.format(i) + "post_attention_layernorm.weight")]
+        if FULL(i):
+            table += [(f"blk.{i}.attn_q_a.weight", "q41", at_ + "q_a_proj.weight"),
+                      (f"blk.{i}.attn_q_a_norm.weight", "f32raw", at_ + "q_a_layernorm.weight"),
+                      (f"blk.{i}.attn_q_b.weight", "q41", at_ + "q_b_proj.weight"),
+                      (f"blk.{i}.attn_kv_a_mqa.weight", "q41", at_ + "kv_a_proj_with_mqa.weight"),
+                      (f"blk.{i}.attn_kv_a_norm.weight", "f32raw", at_ + "kv_a_layernorm.weight"),
+                      (f"blk.{i}.attn_k_b.weight", "kvb_k", at_ + "kv_b_proj.weight"),
+                      (f"blk.{i}.attn_v_b.weight", "kvb_v", at_ + "kv_b_proj.weight"),
+                      (f"blk.{i}.attn_gate.weight", "q41", at_ + "g_proj.weight"),
+                      (f"blk.{i}.attn_output.weight", "q41", at_ + "dense.weight")]
+        else:
+            table += [(f"blk.{i}.attn_q.weight", "q41", at_ + "q_proj.weight"),
+                      (f"blk.{i}.attn_k.weight", "q41", at_ + "k_proj.weight"),
+                      (f"blk.{i}.attn_v.weight", "q41", at_ + "v_proj.weight"),
+                      (f"blk.{i}.ssm_conv1d_q.weight", "conv", at_ + "q_conv1d.weight"),
+                      (f"blk.{i}.ssm_conv1d_k.weight", "conv", at_ + "k_conv1d.weight"),
+                      (f"blk.{i}.ssm_conv1d_v.weight", "conv", at_ + "v_conv1d.weight"),
+                      (f"blk.{i}.ssm_a", "sa", at_ + "A_log"),
+                      (f"blk.{i}.ssm_dt.bias", "f32raw", at_ + "dt_bias"),
+                      (f"blk.{i}.ssm_beta.weight", "q41", at_ + "b_proj.weight"),
+                      (f"blk.{i}.ssm_f_a.weight", "q41", at_ + "f_proj.weight"),
+                      (f"blk.{i}.ssm_g_a.weight", "q41", at_ + "g_proj.weight"),
+                      (f"blk.{i}.ssm_norm.weight", "f32raw", at_ + "o_norm.weight"),
+                      (f"blk.{i}.attn_output.weight", "q41", at_ + "o_proj.weight")]
+        if i == 0:
+            table += [(f"blk.0.ffn_gate.weight", "q41", mp + "gate_proj.weight"),
+                      (f"blk.0.ffn_up.weight", "q41", mp + "up_proj.weight"),
+                      (f"blk.0.ffn_down.weight", "q41", mp + "down_proj.weight")]
+        else:
+            table += [(f"blk.{i}.ffn_gate_inp.weight", "f32raw", mp + "gate.weight"),
+                      (f"blk.{i}.exp_probs_b.bias", "f32raw", mp + "gate.expert_bias"),
+                      (f"blk.{i}.ffn_gate_exps.weight", "q41", mp + "experts.gate_proj.weight"),
+                      (f"blk.{i}.ffn_up_exps.weight", "q41", mp + "experts.up_proj.weight"),
+                      (f"blk.{i}.ffn_down_exps.weight", "q41", mp + "experts.down_proj.weight"),
+                      (f"blk.{i}.ffn_gate_shexp.weight", "q41", mp + "shared_experts.gate_proj.weight"),
+                      (f"blk.{i}.ffn_up_shexp.weight", "q41", mp + "shared_experts.up_proj.weight"),
+                      (f"blk.{i}.ffn_down_shexp.weight", "q41", mp + "shared_experts.down_proj.weight")]
+    miss = [t for t in table if t[2] not in idx.t]
+    assert not miss, f"source tensor missing: {miss[:5]}"
+    names = [t[0] for t in table]
+    assert len(names) == len(set(names))
+
+    # ---- metadata ----
+    os.makedirs(a.out, exist_ok=True)
+    stem = "edge0-8b"
+    gp = os.path.join(a.out, f"{stem}.gguf")
+    w = GgufWriter(gp)
+    w.kv_s("general.architecture", "bailingmoe3")
+    w.kv_s("general.type", "model")
+    w.kv_s("general.name", "edge0-8b")
+    w.kv("general.alignment", KV["u32"], struct.pack("<I", 32))
+    u = lambda v: struct.pack("<I", v)
+    f32 = lambda v: struct.pack("<f", v)
+    b1 = lambda v: b"\x01" if v else b"\x00"
+    for k, ty, v in [
+        (A + "block_count", KV["u32"], u(NL)), (A + "context_length", KV["u32"], u(cfg["max_position_embeddings"])),
+        (A + "embedding_length", KV["u32"], u(cfg["hidden_size"])), (A + "feed_forward_length", KV["u32"], u(cfg["intermediate_size"])),
+        (A + "attention.head_count", KV["u32"], u(HEADS)),
+        (A + "attention.key_length", KV["u32"], u(cfg["kv_lora_rank"] + cfg["qk_rope_head_dim"])),
+        (A + "attention.value_length", KV["u32"], u(cfg["head_dim"])),
+        (A + "attention.key_length_mla", KV["u32"], u(cfg["qk_head_dim"])),
+        (A + "attention.value_length_mla", KV["u32"], u(cfg["v_head_dim"])),
+        (A + "attention.q_lora_rank", KV["u32"], u(cfg["q_lora_rank"])),
+        (A + "attention.kv_lora_rank", KV["u32"], u(cfg["kv_lora_rank"])),
+        (A + "rope.freq_base", KV["f32"], f32(float(cfg["rope_theta"]))),
+        (A + "rope.dimension_count", KV["u32"], u(cfg["qk_rope_head_dim"])),
+        (A + "attention.layer_norm_rms_epsilon", KV["f32"], f32(cfg["rms_norm_eps"])),
+        (A + "vocab_size", KV["u32"], u(cfg["vocab_size"])),
+        (A + "ssm.conv_kernel", KV["u32"], u(cfg["short_conv_kernel_size"])),
+        (A + "kda.head_dim", KV["u32"], u(cfg["head_dim"])),
+        (A + "expert_count", KV["u32"], u(E)), (A + "expert_used_count", KV["u32"], u(cfg["num_experts_per_tok"])),
+        (A + "expert_group_count", KV["u32"], u(cfg["n_group"])), (A + "expert_group_used_count", KV["u32"], u(cfg["topk_group"])),
+        (A + "expert_gating_func", KV["u32"], u(2)),
+        (A + "expert_weights_scale", KV["f32"], f32(cfg["routed_scaling_factor"])),
+        (A + "expert_feed_forward_length", KV["u32"], u(cfg["moe_intermediate_size"])),
+        (A + "expert_shared_feed_forward_length", KV["u32"], u(cfg.get("shared_expert_intermediate_size", cfg["moe_intermediate_size"]))),
+        (A + "expert_shared_count", KV["u32"], u(1)),
+        (A + "leading_dense_block_count", KV["u32"], u(cfg["first_k_dense_replace"])),
+    ]:
+        w.kv(k, ty, v)
+    w.kv(A + "kda.safe_gate", KV["bool"], b1(True))
+    w.kv(A + "kda.gate_lower_bound", KV["f32"], f32(cfg["kda_lower_bound"]))
+    w.kv(A + "expert_weights_norm", KV["bool"], b1(cfg.get("norm_topk_prob", True)))
+    w.kv(A + "attention.head_count_kv", KV["arr"], struct.pack("<I", KV["u32"]) + struct.pack("<Q", NL)
+         + b"".join(u(1 if FULL(i) else 0) for i in range(NL)))
+    w.kv_s("tokenizer.ggml.model", "gpt2")
+    w.kv_s("tokenizer.ggml.pre", "bailingmoe2")
+
+    # tokenizer arrays
+    vocab = tkj["model"]["vocab"]; toks = [None] * cfg["vocab_size"]
+    for s, i2 in vocab.items():
+        toks[i2] = s
+    for at in tkj["added_tokens"]:
+        toks[at["id"]] = at["content"]
+    tt = [1] * len(toks)
+    for at in tkj["added_tokens"]:
+        tt[at["id"]] = 3
+    n_pad = sum(1 for x in toks if x is None)
+    for j, x in enumerate(toks):
+        if x is None:
+            toks[j] = f"<0x{j:06X}>"
+
+    def arr(key, et, elems):
+        w.kv(key, KV["arr"], struct.pack("<I", et) + struct.pack("<Q", len(elems)) + b"".join(elems))
+    arr("tokenizer.ggml.tokens", KV["str"], [w._str(s) for s in toks])
+    arr("tokenizer.ggml.merges", KV["str"],
+        [w._str(" ".join(m) if isinstance(m, list) else m) for m in tkj["model"]["merges"]])
+    arr("tokenizer.ggml.scores", KV["f32"], [f32(0.0) for _ in toks])
+    arr("tokenizer.ggml.token_type", KV["i32"], [struct.pack("<i", x) for x in tt])
+    w.kv("tokenizer.ggml.bos_token_id", KV["u32"], u(156891))
+    w.kv("tokenizer.ggml.eos_token_id", KV["u32"], u(cfg.get("eos_token_id", 156895)))
+    w.kv("tokenizer.ggml.padding_token_id", KV["u32"], u(cfg.get("pad_token_id", 156892)))
+    w.kv("tokenizer.ggml.add_bos_token", KV["bool"], b1(False))
+    w.kv("tokenizer.ggml.add_eos_token", KV["bool"], b1(False))
+    w.kv("tokenizer.ggml.byte_fallback", KV["bool"], b1(True))
+    ct = os.path.join(a.dir, "chat_template.jinja")
+    if os.path.isfile(ct):
+        w.kv_s("tokenizer.chat_template", open(ct, encoding="utf8").read())
+
+    # ---- two-pass tensors ----
+    def spec(kind, src):
+        shp = idx.t[src][4]
+        if kind == "kvb_k":
+            return T_F32, (cfg["qk_nope_head_dim"], cfg["kv_lora_rank"], HEADS), HEADS * cfg["kv_lora_rank"] * cfg["qk_nope_head_dim"] * 4
+        if kind == "kvb_v":
+            return T_F32, (cfg["kv_lora_rank"], cfg["v_head_dim"], HEADS), HEADS * cfg["kv_lora_rank"] * cfg["v_head_dim"] * 4
+        if kind == "conv":
+            return T_F32, (shp[2], 1, shp[0], 1), int(np.prod(shp)) * 4
+        if kind == "sa":
+            return T_F32, (1, shp[0]), shp[0] * 4          # bailingmoe3.cpp:89 create {1, n_head}
+        if kind == "f32raw":
+            return T_F32, (tuple(reversed(shp)) if len(shp) > 1 else (shp[0],)), int(np.prod(shp)) * 4
+        K = shp[-1] * 8
+        rows = int(np.prod(shp[:-1])) if len(shp) > 1 else 1
+        return T_Q4_1, (K,) + tuple(reversed(shp[:-1])), rows * (K // QK41) * 20
+
+    for dst, kind, src in table:
+        tt_, dims, size = spec(kind, src)
+        w.register(dst, tt_, dims, size)
+    ds = w.finalize()
+    audit = dict(tool=R3VER, gates=[], errors=[])
+    t0 = time.time()
+    for n, (dst, kind, src) in enumerate(table):
+        if kind == "q41":
+            W, S, Bs = q41_src(idx, src)
+            rows = int(np.prod(W.shape[:-1])) if W.ndim > 1 else 1
+            Wf, Sf, Bsf = W.reshape(rows, -1), S.reshape(rows, -1), Bs.reshape(rows, -1)
+            blocks, exm = pack_q41_rows(mlx_q_bytes(Wf), Sf, Bsf, src)
+            w.write(blocks.tobytes()); w.align_pad()
+            a_side = deq_mlx_side(mlx_q_bytes(Wf[:1]), Sf[:1], Bsf[:1])
+            same = sha16(a_side.tobytes()) == sha16(deq_gguf_side(blocks[:1]).tobytes())
+            st = "OK" if same else ("EXEMPT" if any(r == 0 for r, _ in exm) else "FAIL")
+            audit["gates"].append((dst, st))
+            if st == "FAIL":
+                audit["errors"].append(dst)
+        elif kind == "kvb_k" or kind == "kvb_v":
+            W, S, Bs = q41_src(idx, src)
+            K = W.shape[-1] * 8
+            q = mlx_q_bytes(W.reshape(4096, -1))
+            v = deq_mlx_side(q, S.reshape(4096, -1), Bs.reshape(4096, -1))   # (4096,512)
+            kv = v.reshape(HEADS, 256, K)
+            kt = np.ascontiguousarray(kv[:, :128, :].transpose(0, 2, 1), np.float32)   # (16,512,128)
+            vv = np.ascontiguousarray(kv[:, 128:, :], np.float32)                       # (16,128,512)
+            payload = kt.tobytes() if kind == "kvb_k" else vv.tobytes()
+            w.write(payload); w.align_pad()
+        elif kind == "conv":
+            x = idx.load(src)                          # [C,1,4]
+            y = np.ascontiguousarray(x.reshape(1, x.shape[0], 1, x.shape[2]), np.float32)
+            w.write(y.tobytes()); w.align_pad()
+        elif kind == "sa":
+            x = idx.load(src)
+            v = np.exp(np.asarray(x, dtype=np.float32))
+            v2 = np.array([math.exp(float(t)) for t in x.ravel()], np.float32)
+            dv = float(np.abs(v - v2).max(initial=0.0))
+            if dv > 1e-5:
+                audit["errors"].append(f"{dst} exp dev {dv}")
+            y = np.ascontiguousarray(v.reshape(-1, 1), np.float32)
+            w.write(y.tobytes()); w.align_pad()         # ne {1,16} via reversed((16,1))
+        else:  # f32raw
+            x = idx.load(src)
+            w.write(np.ascontiguousarray(x, np.float32).tobytes()); w.align_pad()
+        if (n + 1) % 100 == 0:
+            print(f"  {n+1}/{len(table)} {dst} {time.time()-t0:.0f}s", flush=True)
+    w.close(); idx.close()
+    ng = len(audit["gates"]); bad = [g for g in audit["gates"] if g[1] != "OK"]
+    audit.update(n_tensors=len(table), pad_tokens=n_pad, elapsed_s=round(time.time() - t0, 1))
+    json.dump(dict(audit=audit, size=os.path.getsize(gp),
+                   verdict="R3-8b-artifact " + ("GREEN" if not bad and not audit["errors"] else "RED")),
+              open(os.path.join(a.out, f"{stem}-audit.json"), "w"), ensure_ascii=False, indent=1)
+    print(f"GGUF {gp} {os.path.getsize(gp)/1e9:.2f}GB tensors={len(table)} gate-mismatch={len(bad)} errors={audit['errors'][:5]}")
+    print("sa dims check:", spec("sa", "model.layers.0.attention.A_log")[1])
+    return 0 if not bad and not audit["errors"] else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)